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

* fix(datasets): stop frame errors being treated as shard exhaustion in StreamingLeRobotDataset

StreamingLeRobotDataset.__iter__ caught every RuntimeError and treated all as exhausted shard. Real errors like video decode failure made each shard get dropped on the first frame, so iteration ended while yeilding zero frames with no errors.

Shard exahustion is StopIteration raised from make_frame generator, which python converts to RuntimeError with StopIteration as __cause__. Added check to tell StopIteration from everything else, consuming real shard exhaustion while re-raising everything else.

Added a test that injects a decode failure and asserts iteration raises instead of returning on empty stream.

Fixes #4066

* refactor(datasets): exception streaming

---------

Co-authored-by: Mohit Yadav <mohitydv09@gmail.com>

src/lerobot/datasets/streaming_dataset.py

@@ -58,6 +58,10 @@ class LookAheadError(Exception):
     pass
 
 
+class _ShardExhaustedError(Exception):
+    """Raised when a streaming dataset shard has no more items."""
+
+
 class Backtrackable[T]:
     """
     Wrap any iterator/iterable so you can step back up to `history` items
@@ -178,7 +182,7 @@ class Backtrackable[T]:
         """
         Check if we can go back `steps` items without raising an IndexError.
         """
-        return steps <= len(self._back_buf) + self._cursor
+        return steps < len(self._back_buf) + self._cursor
 
     def can_peek_ahead(self, steps: int = 1) -> bool:
         """
@@ -422,10 +426,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                     else:
                         frames_buffer.append(frame)
                     break  # random shard sampled, switch shard
-            except (
+            except _ShardExhaustedError:
-                RuntimeError,
-                StopIteration,
-            ):  # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7
                 del idx_to_backtrack_dataset[shard_key]  # Remove exhausted shard, onto another shard
 
         # Once shards are all exhausted, shuffle the buffer and yield the remaining frames
@@ -503,7 +504,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
 
     def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
         """Makes a frame starting from a dataset iterator"""
+        try:
             item = next(dataset_iterator)
+        except StopIteration as e:
+            # Translate exhaustion here, before PEP 479 turns it into an indistinguishable RuntimeError.
+            raise _ShardExhaustedError from e
         item = item_to_torch(item)
 
         updates = []  # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features)

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/bi_so_follower/bi_so_follower.py

@@ -62,6 +62,7 @@ class BiSOFollower(BimanualMixin, Robot):
             position_i_coefficient=config.left_arm_config.position_i_coefficient,
             position_d_coefficient=config.left_arm_config.position_d_coefficient,
             use_degrees=config.left_arm_config.use_degrees,
+            num_read_retries=config.left_arm_config.num_read_retries,
             cameras=left_arm_cameras,
         )
 
@@ -75,6 +76,7 @@ class BiSOFollower(BimanualMixin, Robot):
             position_i_coefficient=config.right_arm_config.position_i_coefficient,
             position_d_coefficient=config.right_arm_config.position_d_coefficient,
             use_degrees=config.right_arm_config.use_degrees,
+            num_read_retries=config.right_arm_config.num_read_retries,
             cameras=config.right_arm_config.cameras,
         )
 

src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py

@@ -44,12 +44,16 @@ class BiSOLeader(BimanualMixin, Teleoperator):
             id=f"{config.id}_left" if config.id else None,
             calibration_dir=config.calibration_dir,
             port=config.left_arm_config.port,
+            use_degrees=config.left_arm_config.use_degrees,
+            num_read_retries=config.left_arm_config.num_read_retries,
         )
 
         right_arm_config = SOLeaderTeleopConfig(
             id=f"{config.id}_right" if config.id else None,
             calibration_dir=config.calibration_dir,
             port=config.right_arm_config.port,
+            use_degrees=config.right_arm_config.use_degrees,
+            num_read_retries=config.right_arm_config.num_read_retries,
         )
 
         self.left_arm = SOLeader(left_arm_config)

tests/datasets/test_streaming.py

@@ -252,6 +252,54 @@ def test_frames_order_with_shards(tmp_path, lerobot_dataset_factory, shuffle):
             assert frames_match
 
 
+def test_iter_raises_on_frame_error(tmp_path, lerobot_dataset_factory, monkeypatch):
+    """Video decode failures must propagate instead of being silently ignored."""
+    ds_num_frames = 20
+    ds_num_episodes = 2
+    buffer_size = 10
+
+    local_path = tmp_path / "test"
+    repo_id = f"{DUMMY_REPO_ID}"
+
+    lerobot_dataset_factory(
+        root=local_path,
+        repo_id=repo_id,
+        total_episodes=ds_num_episodes,
+        total_frames=ds_num_frames,
+    )
+
+    streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=buffer_size)
+
+    def broken_video_decode(*args, **kwargs):
+        raise RuntimeError("Could not load libtorchcodec")
+
+    monkeypatch.setattr(streaming_dataset_module, "decode_video_frames_torchcodec", broken_video_decode)
+
+    with pytest.raises(RuntimeError, match="libtorchcodec"):
+        next(iter(streaming_ds))
+
+
+def test_iter_raises_on_nested_generator_error(tmp_path, lerobot_dataset_factory, monkeypatch):
+    """PEP 479 errors below frame construction must not be mistaken for shard exhaustion."""
+    local_path = tmp_path / "test"
+    repo_id = DUMMY_REPO_ID
+
+    lerobot_dataset_factory(root=local_path, repo_id=repo_id, total_episodes=2, total_frames=20)
+    streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=10)
+
+    def broken_frame_generator():
+        raise StopIteration("decoder internal failure")
+        yield
+
+    def broken_video_decode(*args, **kwargs):
+        return next(broken_frame_generator())
+
+    monkeypatch.setattr(streaming_dataset_module, "decode_video_frames_torchcodec", broken_video_decode)
+
+    with pytest.raises(RuntimeError, match="generator raised StopIteration"):
+        next(iter(streaming_ds))
+
+
 @pytest.mark.parametrize(
     "state_deltas, action_deltas",
     [