File size: 4,758 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """Server for serving a policy over websockets.
Adapted from https://github.com/robo-arena/roboarena/
"""
import asyncio
import dataclasses
import logging
import traceback
from openpi_client.base_policy import BasePolicy
from openpi_client import msgpack_numpy
import websockets.asyncio.server
import websockets.frames
@dataclasses.dataclass
class PolicyServerConfig:
# Resolution that images get resized to client-side, None means no resizing.
# It's beneficial to resize images to the desired resolution client-side for faster communication.
image_resolution: tuple[int, int] | None = (224, 224)
# Whether or not wrist camera image(s) should be sent.
needs_wrist_camera: bool = True
# Number of external cameras to send.
n_external_cameras: int = 1 # can be in [0, 1, 2]
# Whether or not stereo camera image(s) should be sent.
needs_stereo_camera: bool = False
# Whether or not the unique eval session id should be sent (e.g. for policies that want to keep track of history).
needs_session_id: bool = False
# Which action space to use.
action_space: str = "joint_position" # can be in ["joint_position", "joint_velocity", "cartesian_position", "cartesian_velocity"]
class WebsocketPolicyServer:
"""
Serves a policy using the websocket protocol.
Interface:
Observation:
- observation/wrist_image_left: (H, W, 3) if needs_wrist_camera is True
- observation/wrist_image_right: (H, W, 3) if needs_wrist_camera is True and needs_stereo_camera is True
- observation/exterior_image_{i}_left: (H, W, 3) if n_external_cameras >= 1
- observation/exterior_image_{i}_right: (H, W, 3) if needs_stereo_camera is True
- session_id: (1,) if needs_session_id is True
- observation/joint_position: (7,)
- observation/cartesian_position: (6,)
- observation/gripper_position: (1,)
- prompt: str, the natural language task instruction for the policy
Action:
- action: (N, 8,) or (N, 7,): either 7 movement actions (for joint action spaces) or 6 (for cartesian) plus one dimension for gripper position
--> all N actions will get executed on the robot before the server is queried again
"""
def __init__(
self,
policy: BasePolicy,
server_config: PolicyServerConfig,
host: str = "0.0.0.0",
port: int = 8000,
) -> None:
self._policy = policy
self._server_config = server_config
self._host = host
self._port = port
logging.getLogger("websockets.server").setLevel(logging.INFO)
def serve_forever(self) -> None:
asyncio.run(self.run())
async def run(self):
async with websockets.asyncio.server.serve(
self._handler,
self._host,
self._port,
compression=None,
max_size=None,
) as server:
await server.serve_forever()
async def _handler(self, websocket: websockets.asyncio.server.ServerConnection):
logging.info(f"Connection from {websocket.remote_address} opened")
packer = msgpack_numpy.Packer()
# Send server config to client to configure what gets sent to server.
await websocket.send(packer.pack(dataclasses.asdict(self._server_config)))
while True:
try:
obs = msgpack_numpy.unpackb(await websocket.recv())
endpoint = obs["endpoint"]
del obs["endpoint"]
if endpoint == "reset":
self._policy.reset(obs)
to_return = "reset successful"
else:
action = self._policy.infer(obs)
to_return = packer.pack(action)
await websocket.send(to_return)
except websockets.ConnectionClosed:
logging.info(f"Connection from {websocket.remote_address} closed")
break
except Exception:
await websocket.send(traceback.format_exc())
await websocket.close(
code=websockets.frames.CloseCode.INTERNAL_ERROR,
reason="Internal server error. Traceback included in previous frame.",
)
raise
if __name__ == "__main__":
import numpy as np
class DummyPolicy(BasePolicy):
def infer(self, obs):
return np.zeros((1, 8), dtype=np.float32)
def reset(self, reset_info):
pass
logging.basicConfig(level=logging.INFO)
policy = DummyPolicy()
server = WebsocketPolicyServer(policy, PolicyServerConfig())
server.serve_forever()
|