| import logging |
| import time |
| from typing import Dict, Tuple |
|
|
| import websockets.sync.client |
| from typing_extensions import override |
|
|
| from openpi_client import base_policy as _base_policy |
| from openpi_client import msgpack_numpy |
|
|
|
|
| class WebsocketClientPolicy(_base_policy.BasePolicy): |
| """Implements the Policy interface by communicating with a server over websocket. |
| |
| See WebsocketPolicyServer for a corresponding server implementation. |
| """ |
|
|
| def __init__(self, host: str = "0.0.0.0", port: int = 8000) -> None: |
| self._uri = f"ws://{host}:{port}" |
| self._packer = msgpack_numpy.Packer() |
| self._ws, self._server_metadata = self._wait_for_server() |
|
|
| def get_server_metadata(self) -> Dict: |
| return self._server_metadata |
|
|
| def _wait_for_server(self) -> Tuple[websockets.sync.client.ClientConnection, Dict]: |
| logging.info(f"Waiting for server at {self._uri}...") |
| while True: |
| try: |
| conn = websockets.sync.client.connect(self._uri, compression=None, max_size=None) |
| metadata = msgpack_numpy.unpackb(conn.recv()) |
| return conn, metadata |
| except (ConnectionRefusedError, TimeoutError, OSError) as exc: |
| logging.info(f"Still waiting for server: {exc}") |
| time.sleep(5) |
|
|
| @override |
| def infer(self, obs: Dict) -> Dict: |
| data = self._packer.pack(obs) |
| self._ws.send(data) |
| response = self._ws.recv() |
| if isinstance(response, str): |
| |
| raise RuntimeError(f"Error in inference server:\n{response}") |
| return msgpack_numpy.unpackb(response) |
|
|
| @override |
| def reset(self) -> None: |
| pass |
|
|
