Add simple client CPU-only dry run

examples/simple_client/README.md

@@ -1,13 +1,29 @@
 # 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:
 
 ```bash
 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
+```

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))

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,)),