policy-websocket 0.1.0

WebSocket-based policy client/server for robot learning

policy-websocket

WebSocket-based policy client/server for robot learning. Provides a minimal, dependency-light interface for running policies remotely over WebSocket.

Compatible with openpi, RoboCasa, and other robot environments.

Installation

pip install policy-websocket

For development from source:

git clone https://github.com/YufengJin/policy_websocket.git && cd policy_websocket
pip install -e .

Quick Start

# Terminal 1: Start single-step server
python examples/policy_server.py --port 8000

# Terminal 2: Run client
python examples/policy_client.py --host localhost --port 8000 --steps 10

Or run the Action Chunk server (policy_server_ac.py) and integration test (examples/run_example_test.py).

Components

Class	Description
BasePolicy	Abstract base: infer(obs) -> dict, reset()
WebsocketClientPolicy	Client that sends obs to a remote server, returns actions
WebsocketPolicyServer	Server that wraps any BasePolicy and serves over WebSocket
ActionChunkBroker	Wraps chunk-returning policies to yield one action per step

Documentation

• Policy Server Setup Guide — Wrap PyTorch models, Action Chunk, RoboCasa integration
• Module Reference — Component overview, data flow, API details
• Examples — Single-step server, Action Chunk server, client, integration test
• Stress Test — 3-view RGB-D 720p throughput, FPS and loss measurement

Usage

Server (wrap your policy)

from policy_websocket import BasePolicy, WebsocketPolicyServer
import numpy as np

class MyPolicy(BasePolicy):
    def infer(self, obs):
        return {"actions": np.zeros(7)}

server = WebsocketPolicyServer(policy=MyPolicy(), host="0.0.0.0", port=8000)
server.serve_forever()

Client

from policy_websocket import WebsocketClientPolicy

policy = WebsocketClientPolicy(host="localhost", port=8000)
action_dict = policy.infer(obs_dict)
action = action_dict["actions"]

ActionChunkBroker (predict N, execute M)

from policy_websocket import BasePolicy, WebsocketPolicyServer, ActionChunkBroker

# Inner policy returns (16, action_dim) chunks
chunk_policy = MyChunkPolicy()
broker = ActionChunkBroker(policy=chunk_policy, action_horizon=8)
server = WebsocketPolicyServer(policy=broker, port=8000)

Protocol

• Transport: WebSocket
• Serialization: msgpack with NumPy array support
• Flow: Client sends obs dict → Server calls policy.infer(obs) → Returns action dict
• Health: GET /healthz returns 200 OK

Dependencies

• websockets>=11.0
• msgpack>=1.0.5
• numpy>=1.22.4,<2.0.0

License

MIT