fix(policies): vla jepa prepare model input to take index 0 and not index -1

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`.
 
-Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
+On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
-D405, do not support them and raise an error at connection time.
+exposure or gain also affects the depth stream.
 
 </hfoption>
 </hfoptions>

src/lerobot/cameras/realsense/camera_realsense.py

@@ -365,12 +365,11 @@ class RealSenseCamera(Camera):
         return self._async_read(timeout_ms=10000, read_depth=read_depth)
 
     def _get_color_sensor(self) -> "rs.sensor":
-        """Returns the dedicated "RGB Camera" sensor that controls the color stream.
+        """Returns the sensor that controls the color stream.
 
-        Manual color controls are only applied to a dedicated RGB module. Cameras
+        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
-        without one (e.g. the D405, whose color stream comes from the shared
+        separate RGB module — its color stream comes from "Stereo Module".
-        "Stereo Module") are unsupported, so we never fall back to another sensor
+        We try RGB Camera first, then fall back to Stereo Module.
-        to avoid altering the depth stream.
         """
         if self.rs_profile is None:
             raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
@@ -378,14 +377,12 @@ class RealSenseCamera(Camera):
         device = self.rs_profile.get_device()
         sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
 
-        if "RGB Camera" in sensors:
+        for name in ("RGB Camera", "Stereo Module"):
-            return sensors["RGB Camera"]
+            if name in sensors:
+                return sensors[name]
 
         available = list(sensors.keys())
-        raise RuntimeError(
+        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
-            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/common/train_utils.py

@@ -52,17 +52,6 @@ def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Pa
     return output_dir / CHECKPOINTS_DIR / step_identifier
 
 
-def should_save_checkpoint(step: int, save_freq: int, total_steps: int) -> bool:
-    """Whether a checkpoint should be saved at ``step``.
-
-    A checkpoint is saved every ``save_freq`` steps and always after the final step. A
-    non-positive ``save_freq`` disables periodic saving (only the final checkpoint is
-    written), mirroring how ``log_freq``/``eval_freq`` treat non-positive values and
-    avoiding a ``ZeroDivisionError`` from ``step % 0``.
-    """
-    return (save_freq > 0 and step % save_freq == 0) or step == total_steps
-
-
 def save_training_step(
     step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
 ) -> None:

src/lerobot/configs/train.py

@@ -119,7 +119,6 @@ class TrainPipelineConfig(HubMixin):
     tolerance_s: float = 1e-4
     save_checkpoint: bool = True
     # Checkpoint is saved every `save_freq` training iterations and after the last training step.
-    # A non-positive value disables periodic saving, keeping only the final checkpoint.
     save_freq: int = 20_000
     use_policy_training_preset: bool = True
     optimizer: OptimizerConfig | None = None

src/lerobot/datasets/dataset_writer.py

@@ -386,9 +386,6 @@ class DatasetWriter:
                 self._episodes_since_last_encoding = 0
 
         if episode_data is None:
-            # Post-save cleanup deliberately does not go through clear_episode_buffer():
-            # staging frames of video cameras must survive here — the (possibly batched)
-            # encoder still needs them and deletes them once each video is written.
             if len(self._meta.image_keys) > 0:
                 self._delete_camera_frame_dirs(self._meta.image_keys)
             self.episode_buffer = self._create_episode_buffer()

src/lerobot/policies/vla_jepa/configuration_vla_jepa.py

@@ -15,11 +15,13 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from typing import Any
 
 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+from lerobot.utils.constants import OBS_STATE
 
 
 @PreTrainedConfig.register_subclass("vla_jepa")
@@ -122,6 +124,13 @@ 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,7 +399,8 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         state = batch.get(OBS_STATE)
         if state is not None:
             if state.ndim > 2:
-                state = state[:, -1, :]
+                # deltas are forward-looking here, so index 0 is the current observation, not -1.
+                state = state[:, 0, :]
             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,20 +23,10 @@ from ..config import RobotConfig
 def lekiwi_cameras_config() -> dict[str, CameraConfig]:
     return {
         "front": OpenCVCameraConfig(
-            index_or_path="/dev/video0",
+            index_or_path="/dev/video0", fps=30, width=640, height=480, rotation=Cv2Rotation.ROTATE_180
-            fps=30,
-            width=640,
-            height=480,
-            fourcc="MJPG",
-            rotation=Cv2Rotation.ROTATE_180,
         ),
         "wrist": OpenCVCameraConfig(
-            index_or_path="/dev/video2",
+            index_or_path="/dev/video2", fps=30, width=480, height=640, rotation=Cv2Rotation.ROTATE_90
-            fps=30,
-            width=480,
-            height=640,
-            fourcc="MJPG",
-            rotation=Cv2Rotation.ROTATE_90,
         ),
     }
 

src/lerobot/robots/lekiwi/lekiwi_client.py

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

src/lerobot/robots/lekiwi/lekiwi_host.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
 import json
 import logging
 import time
@@ -43,9 +44,7 @@ class LeKiwiHost:
         self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}")
 
         self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH)
-        # CONFLATE does not support multipart messages; a 2-deep send queue keeps
+        self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1)
-        # 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
@@ -100,23 +99,19 @@ def main(cfg: LeKiwiServerConfig):
 
             last_observation = robot.get_observation()
 
-            # Send one multipart message: a JSON header frame (state + camera
+            # Encode ndarrays to base64 strings
-            # order) followed by one raw JPEG frame per camera. Raw JPEG avoids
+            for cam_key, _ in robot.cameras.items():
-            # the 33% base64 inflation of embedding binary data in JSON.
+                ret, buffer = cv2.imencode(
-            cam_keys = list(robot.cameras.keys())
+                    ".jpg", last_observation[cam_key], [int(cv2.IMWRITE_JPEG_QUALITY), 90]
-            jpeg_frames = []
-            for cam_key in cam_keys:
-                ret, jpeg = cv2.imencode(
-                    ".jpg", last_observation.pop(cam_key), [int(cv2.IMWRITE_JPEG_QUALITY), 90]
                 )
-                jpeg_frames.append(jpeg if ret else b"")
+                if ret:
-            header = {"_cams": cam_keys, **last_observation}
+                    last_observation[cam_key] = base64.b64encode(buffer).decode("utf-8")
+                else:
+                    last_observation[cam_key] = ""
 
             # Send the observation to the remote agent
             try:
-                host.zmq_observation_socket.send_multipart(
+                host.zmq_observation_socket.send_string(json.dumps(last_observation), flags=zmq.NOBLOCK)
-                    [json.dumps(header).encode()] + jpeg_frames, flags=zmq.NOBLOCK
-                )
             except zmq.Again:
                 logging.info("Dropping observation, no client connected")
 

src/lerobot/scripts/lerobot_train.py

@@ -45,7 +45,6 @@ from lerobot.common.train_utils import (
     load_training_state,
     push_checkpoint_to_hub,
     save_checkpoint,
-    should_save_checkpoint,
     update_last_checkpoint,
 )
 from lerobot.common.wandb_utils import WandBLogger
@@ -617,7 +616,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             progbar.update(1)
         train_tracker.step()
         is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
-        is_saving_step = should_save_checkpoint(step, cfg.save_freq, cfg.steps)
+        is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
         is_env_eval_step = cfg.env_eval_freq > 0 and step % cfg.env_eval_freq == 0
         is_eval_step = cfg.eval_steps > 0 and eval_dataloader is not None and step % cfg.eval_steps == 0
 

tests/cameras/test_realsense.py

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

tests/datasets/test_dataset_writer.py

@@ -236,42 +236,6 @@ def test_clear_removes_video_frame_staging_dir(tmp_path):
     assert not video_staging_dir.exists()
 
 
-def test_batched_encoding_staging_survives_save(tmp_path):
-    """The post-save clear must NOT delete video staging frames.
-
-    With ``batch_encoding_size > 1`` the frames of already-saved episodes stay
-    on disk until the batch encode runs; the encoder deletes them afterwards.
-    A blanket switch of the post-save cleanup to ``camera_keys`` (as done in the
-    discard path) would silently break batched encoding.
-    """
-    video_key = "observation.images.cam"
-    features = {
-        video_key: {
-            "dtype": "video",
-            "shape": (64, 96, 3),
-            "names": ["height", "width", "channels"],
-        },
-        "action": {"dtype": "float32", "shape": (2,), "names": None},
-    }
-    dataset = LeRobotDataset.create(
-        repo_id=DUMMY_REPO_ID,
-        fps=DEFAULT_FPS,
-        features=features,
-        root=tmp_path / "ds",
-        use_videos=True,
-        batch_encoding_size=2,
-    )
-    for _ in range(3):
-        dataset.add_frame(_make_frame(features))
-
-    staging_dir = dataset.writer._get_image_file_dir(0, video_key)
-    assert staging_dir.is_dir()
-
-    dataset.save_episode()  # first of a batch of 2: no encoding yet
-
-    assert staging_dir.is_dir() and any(staging_dir.iterdir())
-
-
 def test_finalize_is_idempotent(tmp_path):
     """Calling finalize() twice does not raise."""
     dataset = LeRobotDataset.create(

tests/utils/test_train_utils.py

@@ -30,7 +30,6 @@ from lerobot.common.train_utils import (
     save_checkpoint,
     save_training_state,
     save_training_step,
-    should_save_checkpoint,
     update_last_checkpoint,
 )
 from lerobot.utils.constants import (
@@ -51,18 +50,6 @@ def test_get_step_identifier():
     assert get_step_identifier(456789, 1_000_000) == "0456789"
 
 
-def test_should_save_checkpoint():
-    # Periodic checkpoints land on multiples of save_freq.
-    assert should_save_checkpoint(10, save_freq=10, total_steps=100) is True
-    assert should_save_checkpoint(5, save_freq=10, total_steps=100) is False
-    # The final step always saves, even when it is not a multiple of save_freq.
-    assert should_save_checkpoint(100, save_freq=30, total_steps=100) is True
-    # save_freq <= 0 disables periodic saving without raising ZeroDivisionError.
-    assert should_save_checkpoint(1, save_freq=0, total_steps=100) is False
-    assert should_save_checkpoint(100, save_freq=0, total_steps=100) is True
-    assert should_save_checkpoint(1, save_freq=-1, total_steps=100) is False
-
-
 def test_get_step_checkpoint_dir():
     output_dir = Path("/checkpoints")
     step_dir = get_step_checkpoint_dir(output_dir, 1000, 5)