chore(format): formatting code

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>

docs/source/cameras.mdx

@@ -164,8 +164,8 @@ includes the range reported by the sensor. Requesting an unsupported control als
 Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
 require `use_rgb=True`.
 
-On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
-exposure or gain also affects the depth stream.
+Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
+D405, do not support them and raise an error at connection time.
 
 </hfoption>
 </hfoptions>

src/lerobot/cameras/realsense/camera_realsense.py

@@ -365,11 +365,12 @@ class RealSenseCamera(Camera):
         return self._async_read(timeout_ms=10000, read_depth=read_depth)
 
     def _get_color_sensor(self) -> "rs.sensor":
-        """Returns the sensor that controls the color stream.
+        """Returns the dedicated "RGB Camera" sensor that controls the color stream.
 
-        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
-        separate RGB module — its color stream comes from "Stereo Module".
-        We try RGB Camera first, then fall back to Stereo Module.
+        Manual color controls are only applied to a dedicated RGB module. Cameras
+        without one (e.g. the D405, whose color stream comes from the shared
+        "Stereo Module") are unsupported, so we never fall back to another sensor
+        to avoid altering the depth stream.
         """
         if self.rs_profile is None:
             raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
@@ -377,12 +378,14 @@ class RealSenseCamera(Camera):
         device = self.rs_profile.get_device()
         sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
 
-        for name in ("RGB Camera", "Stereo Module"):
-            if name in sensors:
-                return sensors[name]
+        if "RGB Camera" in sensors:
+            return sensors["RGB Camera"]
 
         available = list(sensors.keys())
-        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
+        raise RuntimeError(
+            f"{self}: manual color controls require a dedicated 'RGB Camera' module, which this camera does not have. ",
+            f"Available sensors: {available}.",
+        )
 
     def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
         """Sets a sensor option, re-raising range errors with actionable diagnostics."""

src/lerobot/datasets/dataset_tools.py

@@ -1045,10 +1045,12 @@ def _copy_data_with_feature_changes(
                     df[feature_name] = feature_values
                 else:
                     feature_slice = values[frame_idx:end_idx]
-                    if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
+                    if feature_slice.ndim == 1:
+                        df[feature_name] = feature_slice
+                    elif feature_slice.ndim == 2 and feature_slice.shape[1] == 1:
                         df[feature_name] = feature_slice.flatten()
                     else:
-                        df[feature_name] = feature_slice
+                        df[feature_name] = list(feature_slice)
             frame_idx = end_idx
 
         # Write using the same chunk/file structure as source

src/lerobot/policies/vla_jepa/configuration_vla_jepa.py

@@ -15,13 +15,11 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any
 
 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.types import NormalizationMode
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-from lerobot.utils.constants import OBS_STATE
 
 
 @PreTrainedConfig.register_subclass("vla_jepa")
@@ -124,13 +122,6 @@ class VLAJEPAConfig(PreTrainedConfig):
         if self.robot_state_feature is not None:
             self.state_dim = self.robot_state_feature.shape[0]
 
-    def set_dataset_feature_metadata(self, dataset_features: dict[str, Any]) -> None:
-        """Add `observation.state` to `input_features` if missing, so it gets normalized."""
-        if OBS_STATE in self.input_features or OBS_STATE not in dataset_features:
-            return
-        shape = tuple(dataset_features[OBS_STATE]["shape"])
-        self.input_features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=shape)
-
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(
             lr=self.optimizer_lr,

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -399,8 +399,7 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         state = batch.get(OBS_STATE)
         if state is not None:
             if state.ndim > 2:
-                # deltas are forward-looking here, so index 0 is the current observation, not -1.
-                state = state[:, 0, :]
+                state = state[:, -1, :]
             inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float()  # [B, 1, dim]
 
         return inputs

src/lerobot/robots/lekiwi/config_lekiwi.py

@@ -23,10 +23,20 @@ from ..config import RobotConfig
 def lekiwi_cameras_config() -> dict[str, CameraConfig]:
     return {
         "front": OpenCVCameraConfig(
-            index_or_path="/dev/video0", fps=30, width=640, height=480, rotation=Cv2Rotation.ROTATE_180
+            index_or_path="/dev/video0",
+            fps=30,
+            width=640,
+            height=480,
+            fourcc="MJPG",
+            rotation=Cv2Rotation.ROTATE_180,
         ),
         "wrist": OpenCVCameraConfig(
-            index_or_path="/dev/video2", fps=30, width=480, height=640, rotation=Cv2Rotation.ROTATE_90
+            index_or_path="/dev/video2",
+            fps=30,
+            width=480,
+            height=640,
+            fourcc="MJPG",
+            rotation=Cv2Rotation.ROTATE_90,
         ),
     }
 

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")
 

tests/cameras/test_realsense.py

@@ -322,15 +322,16 @@ def test_get_color_sensor_prefers_rgb_camera():
     assert camera._get_color_sensor() is rgb
 
 
-def test_get_color_sensor_falls_back_to_stereo_module():
-    """D405 has no separate RGB module; color comes from Stereo Module."""
+def test_get_color_sensor_raises_without_dedicated_rgb_module():
+    """D405 has no separate RGB module; we refuse to touch the shared Stereo Module."""
     config = RealSenseCameraConfig(serial_number_or_name="042")
     camera = RealSenseCamera(config)
 
     stereo = _make_mock_sensor("Stereo Module")
     _attach_mock_color_sensor(camera, stereo)
 
-    assert camera._get_color_sensor() is stereo
+    with pytest.raises(RuntimeError, match="dedicated 'RGB Camera' module"):
+        camera._get_color_sensor()
 
 
 def test_get_color_sensor_raises_with_available_sensors():

tests/datasets/test_dataset_tools.py

@@ -304,40 +304,45 @@ def test_merge_empty_list(tmp_path):
         merge_datasets([], output_repo_id="merged", output_dir=tmp_path)
 
 
-def test_add_features_with_values(sample_dataset, tmp_path):
-    """Test adding a feature with pre-computed values."""
+@pytest.mark.parametrize(
+    "values, feature_shape, expected_item_shape",
+    [
+        (np.random.randn(50).astype(np.float32), (1,), ()),
+        (np.random.randn(50, 1).astype(np.float32), (1,), ()),
+        (np.random.randn(50, 7).astype(np.float32), (7,), (7,)),
+        (np.random.randn(50, 1, 4).astype(np.float32), (1, 4), (1, 4)),
+    ],
+    ids=["scalar_1d", "scalar_2d", "vector", "matrix"],
+)
+def test_add_features_with_values(sample_dataset, tmp_path, values, feature_shape, expected_item_shape):
+    """Test adding a pre-computed feature across supported per-frame shapes."""
     num_frames = sample_dataset.meta.total_frames
-    reward_values = np.random.randn(num_frames, 1).astype(np.float32)
+    assert len(values) == num_frames
 
-    feature_info = {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": None,
-    }
-    features = {
-        "reward": (reward_values, feature_info),
-    }
+    feature_info = {"dtype": "float32", "shape": feature_shape, "names": None}
+    features = {"new_feature": (values, feature_info)}
 
     with (
         patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
         patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
     ):
         mock_get_safe_version.return_value = "v3.0"
-        mock_snapshot_download.return_value = str(tmp_path / "with_reward")
+        mock_snapshot_download.return_value = str(tmp_path / "with_feature")
 
         new_dataset = add_features(
             dataset=sample_dataset,
             features=features,
-            output_dir=tmp_path / "with_reward",
+            output_dir=tmp_path / "with_feature",
         )
 
-    assert "reward" in new_dataset.meta.features
-    assert new_dataset.meta.features["reward"] == feature_info
+    assert "new_feature" in new_dataset.meta.features
+    assert new_dataset.meta.features["new_feature"] == feature_info
 
     assert len(new_dataset) == num_frames
     sample_item = new_dataset[0]
-    assert "reward" in sample_item
-    assert isinstance(sample_item["reward"], torch.Tensor)
+    assert "new_feature" in sample_item
+    assert isinstance(sample_item["new_feature"], torch.Tensor)
+    assert tuple(sample_item["new_feature"].shape) == expected_item_shape
 
 
 def test_add_features_with_callable(sample_dataset, tmp_path):