Custom policy support: almond_axol.policy SDK + run-policy policy_type custom

[shawnpatel]

Summary

Run Policy previously only ran LeRobot checkpoints. This PR adds support for running your own model: the robot sends joints and camera frames, and your code returns action chunks. It works from both the SDK and the control panel.

SDK: almond_axol.policy. It needs only numpy and websockets, not LeRobot or torch.

class MyPolicy(Policy):
    def setup(self, spec): ...          # spec: state/action names, cameras, fps, task, policy_path
    def reset(self): ...                # called at the start of each episode
    def infer(self, obs): return chunk  # obs.state / obs.joints / obs.images (RGB) / obs.task -> (T, D)
serve(MyPolicy(), port=8765)

• A plain obs -> chunk function also works.
• A policy can declare action_names and fps. The robot refuses to start on a mismatch.
• If your code raises, the error message goes back to the robot, which stops the episode and shows it.
• The wire protocol is documented (length-prefixed JSON header plus raw uint8/float32 payload over WebSocket), so a model outside Python can implement it directly.
• PolicyServer runs in the background (for tests and notebooks). PolicyClient is the robot side and can also exercise a server without a robot.

run-policy. New --policy_type custom, which connects to your server at --server_host/--server_port (default 127.0.0.1:8765). No LeRobot server is started for it, and the checkpoint fps check is skipped. policy_path becomes optional for custom runs and is passed to your server unchanged.

Only the connection to the policy is new. The existing client's gRPC-specific parts (open, reset, and installing a chunk) are split into methods, and a subclass swaps in the WebSocket client. The robot sends one observation at a time and waits for its chunk. That chunk then goes through the same code as before: chunk alignment and ensembling, the execution filter, the contact watchdog, episode control, and dataset recording.

Control panel.

• Dropdown options for a field now come from its Literal type annotation. custom therefore appears for run-policy but not for collect-dagger, which loads LeRobot in-process.
• Side effect: video_backend also becomes a dropdown.
• Policy path is required only for LeRobot policy types.
• The panel shows a note when custom is selected.
• The Settings → Inference help text now mentions custom servers.

Docs. New operations/custom-policy (how-to) and api/policy (reference, including the wire protocol). The run-policy pages, the control-panel guide, the README, and the module tree are updated.

websockets>=13,<16 is now a direct dependency. It was already installed through uvicorn[standard].

Not in this PR

• DAgger with a custom policy, since it loads its policy in-process through LeRobot.
• Configurable timeouts. They are fixed: 5 min for setup, 60 s per inference reply.
• An axol policy.check command for validating a policy server without a robot.
• The connection has no authentication, the same as the existing inference server. The docs recommend binding 127.0.0.1 or firewalling the port.

Test plan

• uv run pytest: 1591 passed, coverage 57%. The new tests/test_custom_policy.py covers:
  • protocol round trips and rejection of malformed input
  • real server/client sessions, including error passthrough, a second robot being refused, and setup refusing a session
  • the robot-side custom client against a real PolicyServer with a fake robot: chunk truncation and timestep stamping, action-layout and fps mismatches, failures that stop the run, and an unreachable server
  • the form schema's options and required fields
• tests/test_docs_examples.py checks the new API page's tables and examples against the code.
• ruff check, ruff format --check
• web: npm test, npm run lint, npm run format:check, npm run build
• On hardware: run a custom policy server on the robot and on a remote GPU box, then run episodes from the control panel (Save/Discard, contact stop, dataset recording).

🤖 Generated with Claude Code

---

Note

Medium Risk
New real-time robot control path over unauthenticated WebSocket; failures are guarded by action-schema and fps checks, but misconfigured or hostile policy servers could send bad actions if the port is exposed.

Overview
Adds custom policy support so Run Policy can drive the robot from a non-LeRobot model via a new almond_axol.policy SDK and WebSocket wire protocol (hello / reset / observation → actions).

axol run-policy --policy_type custom connects to an operator-hosted policy server (almond_axol.policy.serve) at --server_host/--server_port (default 127.0.0.1:8765) instead of spawning or using LeRobot’s gRPC PolicyServer. policy_path is optional for custom runs and forwarded verbatim; LeRobot checkpoint fps checks are skipped (the server may declare fps in handshake). AxolCustomPolicyClient reuses the same chunk aggregation, execution filtering, and episode flow as LeRobot by swapping only transport.

CLI/config: new RunPolicyType (custom only on run-policy, not collect-dagger). websockets is a direct base dependency. Control panel derives policy_type dropdown from Literal annotations, makes policy_path required only for LeRobot types, and shows a custom-policy hint. Docs/README cover the SDK, operations guide, and API reference; tests cover protocol, server/client sessions, and the robot-side client.

^{Reviewed by Cursor Bugbot for commit 9c037dc. Configure here.}

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated
axol	Ready	Preview	Sep 24, 2026 7:24pm UTC

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 9c037dc. Configure here.}

[cursor]

Slow inference aborts episode teardown

High Severity

The custom receiver stays inside PolicyClient.infer until the policy replies (up to 60s), and it still starts that wait for a leftover mailbox observation after shutdown. Episode teardown only joins workers for 5s, then closes the socket and flags a fatal error, so Save/Discard/timeout discards the episode and ends the run whenever inference is still in flight.

Additional Locations (2)

• almond_axol/cli/run_policy.py#L925-L995
• almond_axol/cli/run_policy.py#L69-L71

 

^{Reviewed by Cursor Bugbot for commit 9c037dc. Configure here.}