mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 21:19:40 +00:00
feat(lekiwi): stream observations as ZMQ multipart (-25% bandwidth) (#4088)
Observations were sent as base64-in-JSON, which inflates every camera frame 33% purely to fit binary JPEG inside a text protocol (~11 Mbps at 3 cameras x 30 Hz). Send a ZMQ multipart message instead: a JSON header frame (state + camera order) followed by one raw JPEG frame per camera. Benchmarked on hardware: -25% wire size (invariant across contention), lower and tighter latency, and no dropped frames under load where the base64 format stalled. ZMQ_CONFLATE does not support multipart, so the observation sockets use 2-deep high-water marks; the client's existing drain-to-latest loop preserves the keep-newest behavior. Breaking wire change: host and client must run the same version. Co-authored-by: Xingdong Zuo <18168681+zuoxingdong@users.noreply.github.com> Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
@@ -14,7 +14,6 @@
|
|||||||
|
|
||||||
# TODO(aliberts, Steven, Pepijn): use gRPC calls instead of zmq?
|
# TODO(aliberts, Steven, Pepijn): use gRPC calls instead of zmq?
|
||||||
|
|
||||||
import base64
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
from functools import cached_property
|
from functools import cached_property
|
||||||
@@ -134,7 +133,9 @@ class LeKiwiClient(Robot):
|
|||||||
self.zmq_observation_socket = self.zmq_context.socket(zmq.PULL)
|
self.zmq_observation_socket = self.zmq_context.socket(zmq.PULL)
|
||||||
zmq_observations_locator = f"tcp://{self.remote_ip}:{self.port_zmq_observations}"
|
zmq_observations_locator = f"tcp://{self.remote_ip}:{self.port_zmq_observations}"
|
||||||
self.zmq_observation_socket.connect(zmq_observations_locator)
|
self.zmq_observation_socket.connect(zmq_observations_locator)
|
||||||
self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1)
|
# CONFLATE does not support multipart messages; a small receive queue plus
|
||||||
|
# the existing drain-to-latest loop keeps newest-only semantics.
|
||||||
|
self.zmq_observation_socket.setsockopt(zmq.RCVHWM, 2)
|
||||||
|
|
||||||
poller = zmq.Poller()
|
poller = zmq.Poller()
|
||||||
poller.register(self.zmq_observation_socket, zmq.POLLIN)
|
poller.register(self.zmq_observation_socket, zmq.POLLIN)
|
||||||
@@ -147,8 +148,8 @@ class LeKiwiClient(Robot):
|
|||||||
def calibrate(self) -> None:
|
def calibrate(self) -> None:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def _poll_and_get_latest_message(self) -> str | None:
|
def _poll_and_get_latest_message(self) -> list[bytes] | None:
|
||||||
"""Polls the ZMQ socket for a limited time and returns the latest message string."""
|
"""Polls the ZMQ socket for a limited time and returns the latest message's frames."""
|
||||||
zmq = self._zmq
|
zmq = self._zmq
|
||||||
poller = zmq.Poller()
|
poller = zmq.Poller()
|
||||||
poller.register(self.zmq_observation_socket, zmq.POLLIN)
|
poller.register(self.zmq_observation_socket, zmq.POLLIN)
|
||||||
@@ -166,7 +167,7 @@ class LeKiwiClient(Robot):
|
|||||||
last_msg = None
|
last_msg = None
|
||||||
while True:
|
while True:
|
||||||
try:
|
try:
|
||||||
msg = self.zmq_observation_socket.recv_string(zmq.NOBLOCK)
|
msg = self.zmq_observation_socket.recv_multipart(zmq.NOBLOCK)
|
||||||
last_msg = msg
|
last_msg = msg
|
||||||
except zmq.Again:
|
except zmq.Again:
|
||||||
break
|
break
|
||||||
@@ -176,28 +177,27 @@ class LeKiwiClient(Robot):
|
|||||||
|
|
||||||
return last_msg
|
return last_msg
|
||||||
|
|
||||||
def _parse_observation_json(self, obs_string: str) -> RobotObservation | None:
|
def _parse_observation(self, frames: list[bytes]) -> RobotObservation | None:
|
||||||
"""Parses the JSON observation string."""
|
"""Parses a multipart observation: JSON header + one raw JPEG frame per camera."""
|
||||||
try:
|
try:
|
||||||
return json.loads(obs_string)
|
header = json.loads(frames[0])
|
||||||
except json.JSONDecodeError as e:
|
cam_names = header.pop("_cams")
|
||||||
logging.error(f"Error decoding JSON observation: {e}")
|
observation: RobotObservation = header
|
||||||
|
for cam_name, jpeg in zip(cam_names, frames[1:], strict=True):
|
||||||
|
observation[cam_name] = jpeg
|
||||||
|
return observation
|
||||||
|
except (json.JSONDecodeError, KeyError, ValueError) as e:
|
||||||
|
logging.error(f"Error decoding observation: {e}")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def _decode_image_from_b64(self, image_b64: str) -> np.ndarray | None:
|
def _decode_image(self, jpeg: bytes) -> np.ndarray | None:
|
||||||
"""Decodes a base64 encoded image string to an OpenCV image."""
|
"""Decodes a raw JPEG buffer to an OpenCV image."""
|
||||||
if not image_b64:
|
if not jpeg:
|
||||||
return None
|
|
||||||
try:
|
|
||||||
jpg_data = base64.b64decode(image_b64)
|
|
||||||
np_arr = np.frombuffer(jpg_data, dtype=np.uint8)
|
|
||||||
frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
|
|
||||||
if frame is None:
|
|
||||||
logging.warning("cv2.imdecode returned None for an image.")
|
|
||||||
return frame
|
|
||||||
except (TypeError, ValueError) as e:
|
|
||||||
logging.error(f"Error decoding base64 image data: {e}")
|
|
||||||
return None
|
return None
|
||||||
|
frame = cv2.imdecode(np.frombuffer(jpeg, dtype=np.uint8), cv2.IMREAD_COLOR)
|
||||||
|
if frame is None:
|
||||||
|
logging.warning("cv2.imdecode returned None for an image.")
|
||||||
|
return frame
|
||||||
|
|
||||||
def _remote_state_from_obs(
|
def _remote_state_from_obs(
|
||||||
self, observation: RobotObservation
|
self, observation: RobotObservation
|
||||||
@@ -212,10 +212,10 @@ class LeKiwiClient(Robot):
|
|||||||
|
|
||||||
# Decode images
|
# Decode images
|
||||||
current_frames: dict[str, np.ndarray] = {}
|
current_frames: dict[str, np.ndarray] = {}
|
||||||
for cam_name, image_b64 in observation.items():
|
for cam_name, jpeg in observation.items():
|
||||||
if cam_name not in self._cameras_ft:
|
if cam_name not in self._cameras_ft:
|
||||||
continue
|
continue
|
||||||
frame = self._decode_image_from_b64(image_b64)
|
frame = self._decode_image(jpeg)
|
||||||
if frame is not None:
|
if frame is not None:
|
||||||
current_frames[cam_name] = frame
|
current_frames[cam_name] = frame
|
||||||
|
|
||||||
@@ -230,15 +230,15 @@ class LeKiwiClient(Robot):
|
|||||||
If no new data arrives or decoding fails, returns the last known values.
|
If no new data arrives or decoding fails, returns the last known values.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
# 1. Get the latest message string from the socket
|
# 1. Get the latest message's frames from the socket
|
||||||
latest_message_str = self._poll_and_get_latest_message()
|
latest_frames = self._poll_and_get_latest_message()
|
||||||
|
|
||||||
# 2. If no message, return cached data
|
# 2. If no message, return cached data
|
||||||
if latest_message_str is None:
|
if latest_frames is None:
|
||||||
return self.last_frames, self.last_remote_state
|
return self.last_frames, self.last_remote_state
|
||||||
|
|
||||||
# 3. Parse the JSON message
|
# 3. Parse the multipart message
|
||||||
observation = self._parse_observation_json(latest_message_str)
|
observation = self._parse_observation(latest_frames)
|
||||||
|
|
||||||
# 4. If JSON parsing failed, return cached data
|
# 4. If JSON parsing failed, return cached data
|
||||||
if observation is None:
|
if observation is None:
|
||||||
|
|||||||
@@ -14,7 +14,6 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
import base64
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
@@ -44,7 +43,9 @@ class LeKiwiHost:
|
|||||||
self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}")
|
self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}")
|
||||||
|
|
||||||
self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH)
|
self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH)
|
||||||
self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1)
|
# CONFLATE does not support multipart messages; a 2-deep send queue keeps
|
||||||
|
# near-latest-only semantics and sheds stale observations during stalls.
|
||||||
|
self.zmq_observation_socket.setsockopt(zmq.SNDHWM, 2)
|
||||||
self.zmq_observation_socket.bind(f"tcp://*:{config.port_zmq_observations}")
|
self.zmq_observation_socket.bind(f"tcp://*:{config.port_zmq_observations}")
|
||||||
|
|
||||||
self.connection_time_s = config.connection_time_s
|
self.connection_time_s = config.connection_time_s
|
||||||
@@ -99,19 +100,23 @@ def main(cfg: LeKiwiServerConfig):
|
|||||||
|
|
||||||
last_observation = robot.get_observation()
|
last_observation = robot.get_observation()
|
||||||
|
|
||||||
# Encode ndarrays to base64 strings
|
# Send one multipart message: a JSON header frame (state + camera
|
||||||
for cam_key, _ in robot.cameras.items():
|
# order) followed by one raw JPEG frame per camera. Raw JPEG avoids
|
||||||
ret, buffer = cv2.imencode(
|
# the 33% base64 inflation of embedding binary data in JSON.
|
||||||
".jpg", last_observation[cam_key], [int(cv2.IMWRITE_JPEG_QUALITY), 90]
|
cam_keys = list(robot.cameras.keys())
|
||||||
|
jpeg_frames = []
|
||||||
|
for cam_key in cam_keys:
|
||||||
|
ret, jpeg = cv2.imencode(
|
||||||
|
".jpg", last_observation.pop(cam_key), [int(cv2.IMWRITE_JPEG_QUALITY), 90]
|
||||||
)
|
)
|
||||||
if ret:
|
jpeg_frames.append(jpeg if ret else b"")
|
||||||
last_observation[cam_key] = base64.b64encode(buffer).decode("utf-8")
|
header = {"_cams": cam_keys, **last_observation}
|
||||||
else:
|
|
||||||
last_observation[cam_key] = ""
|
|
||||||
|
|
||||||
# Send the observation to the remote agent
|
# Send the observation to the remote agent
|
||||||
try:
|
try:
|
||||||
host.zmq_observation_socket.send_string(json.dumps(last_observation), flags=zmq.NOBLOCK)
|
host.zmq_observation_socket.send_multipart(
|
||||||
|
[json.dumps(header).encode()] + jpeg_frames, flags=zmq.NOBLOCK
|
||||||
|
)
|
||||||
except zmq.Again:
|
except zmq.Again:
|
||||||
logging.info("Dropping observation, no client connected")
|
logging.info("Dropping observation, no client connected")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user