fix(train): guard checkpoint saving against save_freq=0 (#4112)

is_saving_step was the only step-frequency check without a `> 0` guard,
so save_freq=0 raised ZeroDivisionError from `step % 0` on the first
step. Route the decision through a should_save_checkpoint helper that
treats a non-positive save_freq as "save only the final checkpoint",
matching how log_freq/eval_freq handle non-positive values.

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/common/train_utils.py

@@ -52,6 +52,17 @@ 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,6 +119,7 @@ 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,6 +386,9 @@ 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,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")
 

src/lerobot/scripts/lerobot_train.py

@@ -45,6 +45,7 @@ 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
@@ -616,7 +617,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 = step % cfg.save_freq == 0 or step == cfg.steps
+        is_saving_step = should_save_checkpoint(step, cfg.save_freq, 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,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_writer.py

@@ -236,6 +236,42 @@ 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,6 +30,7 @@ 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 (
@@ -50,6 +51,18 @@ 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)