Metadata-Version: 2.4
Name: policy-websocket
Version: 0.1.0
Summary: WebSocket-based policy client/server for robot learning
Project-URL: Homepage, https://github.com/YufengJin/policy_websocket
Project-URL: Repository, https://github.com/YufengJin/policy_websocket
License-File: LICENSE
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.8
Requires-Dist: msgpack>=1.0.5
Requires-Dist: numpy<2.0.0,>=1.22.4
Requires-Dist: websockets>=11.0
Description-Content-Type: text/markdown

# 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](https://github.com/Physical-Intelligence/openpi), [RoboCasa](https://github.com/robocasa/robocasa), and other robot environments.

## Installation

```bash
pip install policy-websocket
```

For development from source:

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

## Quick Start

```bash
# 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](docs/policy_server.md) — Wrap PyTorch models, Action Chunk, RoboCasa integration
- [Module Reference](docs/policy_client.md) — Component overview, data flow, API details
- [Examples](examples/) — [Single-step server](examples/policy_server.py), [Action Chunk server](examples/policy_server_ac.py), [client](examples/policy_client.py), [integration test](examples/run_example_test.py)
- [Stress Test](scripts/stress_test.py) — 3-view RGB-D 720p throughput, FPS and loss measurement

## Usage

### Server (wrap your policy)

```python
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

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

```python
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
