From 7a9cff68b71a3b80b9368d60034ab803f6af687d Mon Sep 17 00:00:00 2001 From: Travor <3488616445@qq.com> Date: Fri, 22 May 2026 01:13:40 +0800 Subject: [PATCH] Add simple client dry-run server --- examples/simple_client/README.md | 32 +++++++- examples/simple_client/fake_policy_server.py | 78 ++++++++++++++++++++ examples/simple_client/main.py | 22 ++++++ 3 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 examples/simple_client/fake_policy_server.py diff --git a/examples/simple_client/README.md b/examples/simple_client/README.md index bc381c1d7a..fa2724eaa6 100644 --- a/examples/simple_client/README.md +++ b/examples/simple_client/README.md @@ -1,6 +1,6 @@ # Simple Client -A minimal client that sends observations to the server and prints the inference rate. +A minimal client that sends observations to a policy server and prints the inference rate. You can specify which runtime environment to use using the `--env` flag. You can see the available options by running: @@ -8,6 +8,22 @@ You can specify which runtime environment to use using the `--env` flag. You can uv run examples/simple_client/main.py --help ``` +## CPU-only dry run + +If you want to test the websocket path before downloading a checkpoint or connecting a robot, start the fake policy server in one terminal: + +```bash +uv run examples/simple_client/fake_policy_server.py --action-horizon 10 --action-dim 8 +``` + +Then send random observations from another terminal: + +```bash +uv run examples/simple_client/main.py --env DROID --num-steps 3 --show-action-chunk +``` + +This verifies that the client can connect, send an observation dictionary, receive an `actions` array, and print the action chunk shape. The fake server does not load model weights. + ## With Docker ```bash @@ -28,3 +44,17 @@ Terminal window 2: ```bash uv run scripts/serve_policy.py --env DROID ``` + +For a trained or fine-tuned checkpoint, start the real policy server with the training config and checkpoint directory: + +```bash +uv run scripts/serve_policy.py policy:checkpoint \ + --policy.config=pi05_droid \ + --policy.dir=/path/to/checkpoint +``` + +Then keep the same client command, choosing the `--env` whose random observation fields match the policy input mapping: + +```bash +uv run examples/simple_client/main.py --env DROID --num-steps 10 --show-action-chunk +``` diff --git a/examples/simple_client/fake_policy_server.py b/examples/simple_client/fake_policy_server.py new file mode 100644 index 0000000000..615632455d --- /dev/null +++ b/examples/simple_client/fake_policy_server.py @@ -0,0 +1,78 @@ +import dataclasses +import logging +import time + +import numpy as np +from openpi_client import msgpack_numpy +import tyro +import websockets.exceptions +import websockets.sync.server + + +@dataclasses.dataclass +class Args: + """Run a CPU-only fake policy server for testing the websocket client.""" + + # Host to bind the websocket server to. + host: str = "0.0.0.0" + # Port to bind the websocket server to. + port: int = 8000 + # Number of actions returned in each chunk. + action_horizon: int = 10 + # Size of each action vector. + action_dim: int = 8 + # Optional sleep that simulates model latency. + delay_ms: float = 0.0 + + +def main(args: Args) -> None: + packer = msgpack_numpy.Packer() + metadata = { + "model": "fake_policy", + "action_horizon": args.action_horizon, + "action_dim": args.action_dim, + } + + def handler(websocket: websockets.sync.server.ServerConnection) -> None: + logging.info("Connection opened") + websocket.send(packer.pack(metadata)) + request_index = 0 + + while True: + try: + obs = msgpack_numpy.unpackb(websocket.recv()) + start_time = time.monotonic() + if args.delay_ms > 0: + time.sleep(args.delay_ms / 1000) + + actions = _make_action_chunk( + action_horizon=args.action_horizon, + action_dim=args.action_dim, + request_index=request_index, + ) + response = { + "actions": actions, + "server_timing": {"infer_ms": (time.monotonic() - start_time) * 1000}, + "policy_timing": {"fake_policy_ms": args.delay_ms}, + "prompt": obs.get("prompt", ""), + } + websocket.send(packer.pack(response)) + request_index += 1 + except websockets.exceptions.ConnectionClosed: + logging.info("Connection closed") + break + + logging.info("Serving fake policy on %s:%s", args.host, args.port) + with websockets.sync.server.serve(handler, args.host, args.port, compression=None, max_size=None) as server: + server.serve_forever() + + +def _make_action_chunk(*, action_horizon: int, action_dim: int, request_index: int) -> np.ndarray: + base = np.linspace(-1.0, 1.0, action_horizon, dtype=np.float32)[:, None] + offsets = np.linspace(0.0, 0.1, action_dim, dtype=np.float32)[None, :] + return base + offsets + request_index * 0.01 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + main(tyro.cli(Args)) diff --git a/examples/simple_client/main.py b/examples/simple_client/main.py index cd7eda1090..0d6b828900 100644 --- a/examples/simple_client/main.py +++ b/examples/simple_client/main.py @@ -35,6 +35,10 @@ class Args: api_key: str | None = None # Number of steps to run the policy for. num_steps: int = 20 + # Print a compact preview of the returned action chunk on each step. + show_action_chunk: bool = False + # Maximum number of action rows to print when --show-action-chunk is enabled. + max_action_rows: int = 3 # Path to save the timings to a parquet file. (e.g., timing.parquet) timing_file: pathlib.Path | None = None # Environment to run the policy in. @@ -139,6 +143,8 @@ def main(args: Args) -> None: inference_start = time.time() action = policy.infer(obs_fn()) timing_recorder.record("client_infer_ms", 1000 * (time.time() - inference_start)) + if args.show_action_chunk: + _print_action_chunk(action, max_rows=args.max_action_rows) for key, value in action.get("server_timing", {}).items(): timing_recorder.record(f"server_{key}", value) for key, value in action.get("policy_timing", {}).items(): @@ -150,6 +156,22 @@ def main(args: Args) -> None: timing_recorder.write_parquet(args.timing_file) +def _print_action_chunk(action: dict, *, max_rows: int) -> None: + actions = action.get("actions") + if actions is None: + rich.print("[yellow]Response did not include an 'actions' field.[/yellow]") + return + + actions = np.asarray(actions) + rich.print(f"[bold]actions[/bold] shape={actions.shape}") + if actions.ndim == 0: + rich.print(actions) + return + + rows = actions[: max(1, max_rows)] + rich.print(np.array2string(rows, precision=3, suppress_small=True)) + + def _random_observation_aloha() -> dict: return { "state": np.ones((14,)),