File size: 3,306 Bytes
fbd9366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Client for communicating with a policy server.
Adapted from https://github.com/robo-arena/roboarena/
"""
import logging
import time
from typing import Dict, Tuple
import websockets.sync.client
from typing_extensions import override
from openpi_client.base_policy import BasePolicy
from openpi_client import msgpack_numpy
# The websockets library by default sends a ping every 20 seconds and
# expects a pong response within 20 seconds. However, the sever may not
# send a pong response immediately if it is busy processing a request.
# Increase the ping interval and timeout so that the client can wait
# for a longer time before closing the connection.
PING_INTERVAL_SECS = 60
PING_TIMEOUT_SECS = 600
class WebsocketClientPolicy(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}...")
try:
conn = websockets.sync.client.connect(
self._uri,
compression=None,
max_size=None,
ping_interval=PING_INTERVAL_SECS,
ping_timeout=PING_TIMEOUT_SECS,
)
metadata = msgpack_numpy.unpackb(conn.recv())
return conn, metadata
except:
logging.info("Connection to server with ws:// failed. Trying wss:// ...")
self._uri = "wss://" + self._uri.split("//")[1]
conn = websockets.sync.client.connect(
self._uri,
compression=None,
max_size=None,
ping_interval=PING_INTERVAL_SECS,
ping_timeout=PING_TIMEOUT_SECS,
)
metadata = msgpack_numpy.unpackb(conn.recv())
return conn, metadata
@override
def infer(self, obs: Dict) -> Dict: # noqa: UP006
# Notify server that we're calling the infer endpoint (as opposed to the reset endpoint)
obs["endpoint"] = "infer"
data = self._packer.pack(obs)
self._ws.send(data)
response = self._ws.recv()
if isinstance(response, str):
# we're expecting bytes; if the server sends a string, it's an error.
raise RuntimeError(f"Error in inference server:\n{response}")
return msgpack_numpy.unpackb(response)
@override
def reset(self, reset_info: Dict) -> None:
# Notify server that we're calling the reset endpoint (as opposed to the infer endpoint)
reset_info["endpoint"] = "reset"
data = self._packer.pack(reset_info)
self._ws.send(data)
response = self._ws.recv()
return response
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
client = WebsocketClientPolicy()
actions = client.infer({})
print(f"Actions received: {actions}")
client.reset({}) |