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>

src/lerobot/robots/lekiwi/lekiwi_client.py

@@ -14,7 +14,6 @@
 
 # TODO(aliberts, Steven, Pepijn): use gRPC calls instead of zmq?
 
-import base64
 import json
 import logging
 from functools import cached_property
@@ -134,7 +133,9 @@ class LeKiwiClient(Robot):
         self.zmq_observation_socket = self.zmq_context.socket(zmq.PULL)
         zmq_observations_locator = f"tcp://{self.remote_ip}:{self.port_zmq_observations}"
         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.register(self.zmq_observation_socket, zmq.POLLIN)
@@ -147,8 +148,8 @@ class LeKiwiClient(Robot):
     def calibrate(self) -> None:
         pass
 
-    def _poll_and_get_latest_message(self) -> str | None:
-        """Polls the ZMQ socket for a limited time and returns the latest message string."""
+    def _poll_and_get_latest_message(self) -> list[bytes] | None:
+        """Polls the ZMQ socket for a limited time and returns the latest message's frames."""
         zmq = self._zmq
         poller = zmq.Poller()
         poller.register(self.zmq_observation_socket, zmq.POLLIN)
@@ -166,7 +167,7 @@ class LeKiwiClient(Robot):
         last_msg = None
         while True:
             try:
-                msg = self.zmq_observation_socket.recv_string(zmq.NOBLOCK)
+                msg = self.zmq_observation_socket.recv_multipart(zmq.NOBLOCK)
                 last_msg = msg
             except zmq.Again:
                 break
@@ -176,28 +177,27 @@ class LeKiwiClient(Robot):
 
         return last_msg
 
-    def _parse_observation_json(self, obs_string: str) -> RobotObservation | None:
-        """Parses the JSON observation string."""
+    def _parse_observation(self, frames: list[bytes]) -> RobotObservation | None:
+        """Parses a multipart observation: JSON header + one raw JPEG frame per camera."""
         try:
-            return json.loads(obs_string)
-        except json.JSONDecodeError as e:
-            logging.error(f"Error decoding JSON observation: {e}")
+            header = json.loads(frames[0])
+            cam_names = header.pop("_cams")
+            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
 
-    def _decode_image_from_b64(self, image_b64: str) -> np.ndarray | None:
-        """Decodes a base64 encoded image string to an OpenCV image."""
-        if not image_b64:
-            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}")
+    def _decode_image(self, jpeg: bytes) -> np.ndarray | None:
+        """Decodes a raw JPEG buffer to an OpenCV image."""
+        if not jpeg:
             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(
         self, observation: RobotObservation
@@ -212,10 +212,10 @@ class LeKiwiClient(Robot):
 
         # Decode images
         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:
                 continue
-            frame = self._decode_image_from_b64(image_b64)
+            frame = self._decode_image(jpeg)
             if frame is not None:
                 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.
         """
 
-        # 1. Get the latest message string from the socket
-        latest_message_str = self._poll_and_get_latest_message()
+        # 1. Get the latest message's frames from the socket
+        latest_frames = self._poll_and_get_latest_message()
 
         # 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
 
-        # 3. Parse the JSON message
-        observation = self._parse_observation_json(latest_message_str)
+        # 3. Parse the multipart message
+        observation = self._parse_observation(latest_frames)
 
         # 4. If JSON parsing failed, return cached data
         if observation is None:

src/lerobot/robots/lekiwi/lekiwi_host.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import base64
 import json
 import logging
 import time
@@ -44,7 +43,9 @@ class LeKiwiHost:
         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.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.connection_time_s = config.connection_time_s
@@ -99,19 +100,23 @@ def main(cfg: LeKiwiServerConfig):
 
             last_observation = robot.get_observation()
 
-            # Encode ndarrays to base64 strings
-            for cam_key, _ in robot.cameras.items():
-                ret, buffer = cv2.imencode(
-                    ".jpg", last_observation[cam_key], [int(cv2.IMWRITE_JPEG_QUALITY), 90]
+            # Send one multipart message: a JSON header frame (state + camera
+            # order) followed by one raw JPEG frame per camera. Raw JPEG avoids
+            # the 33% base64 inflation of embedding binary data in JSON.
+            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:
-                    last_observation[cam_key] = base64.b64encode(buffer).decode("utf-8")
-                else:
-                    last_observation[cam_key] = ""
+                jpeg_frames.append(jpeg if ret else b"")
+            header = {"_cams": cam_keys, **last_observation}
 
             # Send the observation to the remote agent
             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:
                 logging.info("Dropping observation, no client connected")