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>

docs/source/robocasa.mdx

@@ -82,6 +82,8 @@ By default the env samples objects only from the `lightwheel` registry (what `--
 
 All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on.
 
+By default, each task uses the rollout horizon registered by RoboCasa. Set `--env.episode_length=<steps>` to apply the same explicit horizon to every selected task.
+
 ### Single-task evaluation (recommended for quick iteration)
 
 ```bash

docs/source/using_dataset_tools.mdx

@@ -11,10 +11,9 @@ LeRobot provides several utilities for manipulating datasets:
 3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
 4. **Add Features** - Add new features to a dataset
 5. **Remove Features** - Remove features from a dataset
-6. **Modify Tasks** - Change the natural-language task descriptions associated with episodes
-7. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
-8. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
-9. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
+6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
+7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
+8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
 
 The core implementation is in `lerobot.datasets.dataset_tools`.
 An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
@@ -90,53 +89,6 @@ lerobot-edit-dataset \
     --operation.feature_names "['observation.images.top']"
 ```
 
-#### Modify Tasks
-
-Change the natural-language task descriptions attached to episodes. This is useful for fixing typos, standardizing wording, or re-labeling episodes.
-
-> [!WARNING]
-> `modify_tasks` modifies the dataset **in-place** (updating `meta/tasks.parquet`, the `task_index` column in the data files, the `tasks` column in the episode metadata, and `total_tasks` in `meta/info.json`). The `--new_repo_id` and `--new_root` parameters are ignored for this operation.
-
-```bash
-# Set a single task for all episodes
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.new_task "Pick up the cube and place it"
-
-# Set different tasks for specific episodes
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.episode_tasks '{"0": "Task A", "1": "Task B", "2": "Task A"}'
-
-# Replace existing task strings wherever they appear
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
-
-# Combine modes in a single run
-lerobot-edit-dataset \
-    --repo_id lerobot/pusht \
-    --operation.type modify_tasks \
-    --operation.new_task "Default task" \
-    --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}' \
-    --operation.episode_tasks '{"5": "Special task for episode 5"}'
-```
-
-**Parameters:**
-
-- `new_task`: A single task string used as the default for episodes not otherwise covered.
-- `episode_tasks`: Mapping from episode index to task string.
-- `task_replacements`: Mapping from existing task strings to their replacements, applied to episodes whose current task matches a key. Every key must be an existing task in the dataset.
-
-The modes can be combined in a single run. Per episode, the task is resolved with the following precedence:
-
-`episode_tasks` > `task_replacements` > `new_task` > original task
-
-At least one of `new_task`, `episode_tasks`, or `task_replacements` must be specified. An episode that ends up with no task raises an error.
-
 #### Convert to Video
 
 Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.

examples/isaac_teleop_to_so101/common.py

@@ -44,6 +44,7 @@ from typing import Protocol
 
 import numpy as np
 
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -56,7 +57,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     EEBoundsAndSafety,
     InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS
 from lerobot.utils.robot_utils import precise_sleep
 

examples/isaac_teleop_to_so101/isaac_teleop/teleop_so101_leader_arm.py

@@ -38,7 +38,7 @@ from typing import TYPE_CHECKING
 
 import numpy as np
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 
 from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available
 from .config_isaac_teleop import SO101LeaderArmConfig

examples/isaac_teleop_to_so101/isaac_teleop/teleop_xr_controller.py

@@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 
 from .base import IsaacTeleopTeleoperator, _isaacteleop_available
 from .config_isaac_teleop import XRControllerConfig

examples/isaac_teleop_to_so101/isaac_teleop/xr_controller_processor.py

@@ -26,8 +26,8 @@ from __future__ import annotations
 from dataclasses import dataclass
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import RobotAction
 from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
-from lerobot.types import RobotAction
 from lerobot.utils.rotation import Rotation
 
 from .base import _GRIPPER_MOTOR_SCALE

examples/phone_to_so100/evaluate.py

@@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.common.control_utils import predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
@@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener

examples/phone_to_so100/record.py

@@ -16,6 +16,7 @@
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -36,7 +37,6 @@ from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.phone import Phone, PhoneConfig
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.feature_utils import combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.utils import log_say

examples/phone_to_so100/replay.py

@@ -17,6 +17,7 @@
 import time
 
 from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -27,7 +28,6 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say

examples/phone_to_so100/rollout.py

@@ -27,6 +27,7 @@ Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``.
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.configs import PreTrainedConfig
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -43,7 +44,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
 from lerobot.rollout.inference import SyncInferenceConfig
 from lerobot.rollout.strategies import BaseStrategy
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging
 

examples/phone_to_so100/teleoperate.py

@@ -15,6 +15,7 @@
 
 import time
 
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -31,7 +32,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.teleoperators.phone import Phone, PhoneConfig
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 

examples/so100_to_so100_EE/evaluate.py

@@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.common.control_utils import predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
@@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener

examples/so100_to_so100_EE/record.py

@@ -17,6 +17,7 @@
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -33,7 +34,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 )
 from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.feature_utils import combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.utils import log_say

examples/so100_to_so100_EE/replay.py

@@ -18,6 +18,7 @@
 import time
 
 from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -28,7 +29,6 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say

examples/so100_to_so100_EE/rollout.py

@@ -25,6 +25,7 @@ forward/inverse kinematics.
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.configs import PreTrainedConfig
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -41,7 +42,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
 from lerobot.rollout.inference import SyncInferenceConfig
 from lerobot.rollout.strategies import BaseStrategy
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging
 

examples/so100_to_so100_EE/teleoperate.py

@@ -16,6 +16,7 @@
 
 import time
 
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -30,7 +31,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 

src/lerobot/async_inference/policy_server.py

@@ -38,6 +38,7 @@ import draccus
 import grpc
 import torch
 
+from lerobot.lerobot_types import PolicyAction
 from lerobot.policies import get_policy_class, make_pre_post_processors
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.transport import (
@@ -45,7 +46,6 @@ from lerobot.transport import (
     services_pb2_grpc,  # type: ignore
 )
 from lerobot.transport.utils import receive_bytes_in_chunks
-from lerobot.types import PolicyAction
 
 from .configs import PolicyServerConfig
 from .constants import SUPPORTED_POLICIES

src/lerobot/common/control_utils.py

@@ -35,9 +35,9 @@ else:
 
 if TYPE_CHECKING:
     from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import PolicyAction
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.robots import Robot
-from lerobot.types import PolicyAction
 
 
 def predict_action(

src/lerobot/datasets/dataset_tools.py

@@ -1435,18 +1435,15 @@ def modify_tasks(
     dataset: LeRobotDataset,
     new_task: str | None = None,
     episode_tasks: dict[int, str] | None = None,
-    task_replacements: dict[str, str] | None = None,
 ) -> LeRobotDataset:
     """Modify tasks in a LeRobotDataset.
 
     This function allows you to either:
     1. Set a single task for the entire dataset (using `new_task`)
     2. Set specific tasks for specific episodes (using `episode_tasks`)
-    3. Replace existing task strings wherever they appear (using `task_replacements`)
 
-    Per episode, the task is resolved with precedence:
-    `episode_tasks` > `task_replacements` > `new_task` > original task. An episode that ends
-    up with no task (none of the above apply and it had no original task) raises an error.
+    You can combine both: `new_task` sets the default, and `episode_tasks` overrides
+    specific episodes.
 
     The dataset is modified in-place, updating only the task-related files:
     - meta/tasks.parquet
@@ -1456,14 +1453,11 @@ def modify_tasks(
 
     Args:
         dataset: The source LeRobotDataset to modify.
-        new_task: Default task applied to any episode not covered by `episode_tasks` or a
-            matching `task_replacements` entry.
-        episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
-            over both `task_replacements` and `new_task`.
-        task_replacements: Optional dict mapping existing task strings to new ones. Applied to
-            episodes whose current task matches a key. Every key must be an existing task.
+        new_task: A single task string to apply to all episodes. If None and episode_tasks
+            is also None, raises an error.
+        episode_tasks: Optional dict mapping episode indices to their task strings.
+            Overrides `new_task` for specific episodes.
 
-    At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
 
     Examples:
         Set a single task for all episodes:
@@ -1481,17 +1475,11 @@ def modify_tasks(
                 new_task="Default task",
                 episode_tasks={5: "Special task for episode 5"}
             )
-
-        Replace existing task strings in-place:
-            dataset = modify_tasks(
-                dataset,
-                task_replacements={"Pick up the cube": "Lift the cube"}
-            )
     """
-    if not new_task and not episode_tasks and not task_replacements:
-        raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
+    if new_task is None and episode_tasks is None:
+        raise ValueError("Must specify at least one of new_task or episode_tasks")
 
-    if episode_tasks:
+    if episode_tasks is not None:
         valid_indices = set(range(dataset.meta.total_episodes))
         invalid = set(episode_tasks.keys()) - valid_indices
         if invalid:
@@ -1501,29 +1489,19 @@ def modify_tasks(
     if dataset.meta.episodes is None:
         dataset.meta.episodes = load_episodes(dataset.root)
 
-    if task_replacements:
-        current_tasks = set(dataset.meta.tasks.index)
-        invalid_tasks = set(task_replacements) - current_tasks
-        if invalid_tasks:
-            raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
-
     # Build the mapping from episode index to task string
     episode_to_task: dict[int, str] = {}
     for ep_idx in range(dataset.meta.total_episodes):
-        original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
-        original_task = original_tasks[0] if original_tasks else None
-
         if episode_tasks and ep_idx in episode_tasks:
             episode_to_task[ep_idx] = episode_tasks[ep_idx]
-        elif task_replacements and original_task in task_replacements:
-            episode_to_task[ep_idx] = task_replacements[original_task]
-        elif new_task:
+        elif new_task is not None:
             episode_to_task[ep_idx] = new_task
-        elif original_task:
-            # Keep original task if not overridden and no default provided
-            episode_to_task[ep_idx] = original_task
         else:
-            raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
+            # Keep original task if not overridden and no default provided
+            original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
+            if not original_tasks:
+                raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
+            episode_to_task[ep_idx] = original_tasks[0]
 
     # Collect all unique tasks and create new task mapping
     unique_tasks = sorted(set(episode_to_task.values()))

src/lerobot/datasets/pipeline_features.py

@@ -17,8 +17,8 @@ from collections.abc import Sequence
 from typing import Any
 
 from lerobot.configs import PipelineFeatureType
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.processor import DataProcessorPipeline
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.feature_utils import hw_to_dataset_features
 

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 (
-                RuntimeError,
-                StopIteration,
-            ):  # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7
+            except _ShardExhaustedError:
                 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"""
-        item = next(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/envs/configs.py

@@ -507,7 +507,7 @@ class MetaworldEnv(EnvConfig):
 class RoboCasaEnv(EnvConfig):
     task: str = "CloseFridge"
     fps: int = 20
-    episode_length: int = 1000
+    episode_length: int | None = None
     obs_type: str = "pixels_agent_pos"
     render_mode: str = "rgb_array"
     camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"

src/lerobot/envs/libero.py

@@ -30,7 +30,7 @@ from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv
 
-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv, parse_camera_names
 

src/lerobot/envs/metaworld.py

@@ -25,7 +25,7 @@ import metaworld.policies as policies
 import numpy as np
 from gymnasium import spaces
 
-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv
 

src/lerobot/envs/robocasa.py

@@ -25,7 +25,7 @@ import gymnasium as gym
 import numpy as np
 from gymnasium import spaces
 
-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv, parse_camera_names
 
@@ -98,6 +98,19 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
     return names, None
 
 
+def _get_task_horizon(task: str) -> int:
+    """Return the rollout horizon registered by RoboCasa for a task."""
+    from robocasa.utils.dataset_registry_utils import get_task_horizon
+
+    try:
+        return int(get_task_horizon(task))
+    except ValueError as exc:
+        raise ValueError(
+            f"No RoboCasa horizon is registered for task '{task}'. "
+            "Set `--env.episode_length=<steps>` explicitly."
+        ) from exc
+
+
 def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
     """Split a flat (12,) action vector into a RoboCasa action dict.
 
@@ -154,7 +167,7 @@ class RoboCasaEnv(gym.Env):
 
         self.camera_name = parse_camera_names(camera_name)
 
-        self._max_episode_steps = episode_length if episode_length is not None else 1000
+        self._max_episode_steps = episode_length if episode_length is not None else _get_task_horizon(task)
 
         # Deferred — created on first reset() inside the worker subprocess
         # to avoid inheriting stale GPU/EGL contexts across fork().

src/lerobot/envs/robotwin.py

@@ -28,7 +28,7 @@ import numpy as np
 import torch
 from gymnasium import spaces
 
-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation
 from lerobot.utils.import_utils import _scipy_available
 
 from .utils import _LazyAsyncVectorEnv

src/lerobot/envs/vlabench.py

@@ -37,7 +37,7 @@ import numpy as np
 from gymnasium import spaces
 from scipy.spatial.transform import Rotation
 
-from lerobot.types import RobotObservation
+from lerobot.lerobot_types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv
 

src/lerobot/policies/eo1/processor_eo1.py

@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import TransitionKey
 from lerobot.processor import (
     ComplementaryDataProcessorStep,
     PolicyAction,
@@ -31,7 +32,6 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
-from lerobot.types import TransitionKey
 from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.import_utils import _transformers_available, require_package
 

src/lerobot/policies/evo1/processor_evo1.py

@@ -21,6 +21,7 @@ from typing import Any
 import torch
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -40,7 +41,6 @@ from lerobot.processor.converters import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     DONE,

src/lerobot/policies/factory.py

@@ -28,6 +28,7 @@ if TYPE_CHECKING:
 
 from lerobot.configs import FeatureType, PreTrainedConfig
 from lerobot.envs import EnvConfig, env_to_policy_features
+from lerobot.lerobot_types import PolicyAction
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     PolicyProcessorPipeline,
@@ -37,7 +38,6 @@ from lerobot.processor import (
     transition_to_batch,
     transition_to_policy_action,
 )
-from lerobot.types import PolicyAction
 from lerobot.utils.constants import (
     ACTION,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,

src/lerobot/policies/groot/processor_groot.py

@@ -50,6 +50,7 @@ if TYPE_CHECKING or _datasets_available:
 else:
     LeRobotDataset = None
 
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
@@ -66,7 +67,6 @@ from lerobot.processor import (
     transition_to_batch,
     transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     OBS_IMAGE,

src/lerobot/policies/molmoact2/processor_molmoact2.py

@@ -36,6 +36,7 @@ import torch
 from torch import Tensor
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -49,7 +50,6 @@ from lerobot.processor import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     OBS_IMAGES,

src/lerobot/policies/pi05/processor_pi05.py

@@ -22,6 +22,7 @@ import numpy as np
 import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     PolicyAction,
@@ -33,7 +34,6 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE
 
 from .configuration_pi05 import PI05Config

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -22,6 +22,7 @@ import numpy as np
 import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     ActionTokenizerProcessorStep,
@@ -34,7 +35,6 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE
 
 from .configuration_pi0_fast import PI0FastConfig

src/lerobot/policies/utils.py

@@ -22,7 +22,7 @@ import torch
 from torch import nn
 
 from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
-from lerobot.types import PolicyAction, RobotAction, RobotObservation
+from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame
 

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/policies/xvla/modeling_xvla.py

@@ -150,7 +150,7 @@ class XVLAModel(nn.Module):
         # Freeze or unfreeze policy transformer
         if not self.config.train_policy_transformer:
             for name, param in self.transformer.named_parameters():
-                if "soft_prompts" not in name:
+                if "soft_prompt" not in name:
                     param.requires_grad = False
 
         # Freeze or unfreeze soft prompts

src/lerobot/policies/xvla/processor_xvla.py

@@ -21,6 +21,7 @@ import numpy as np
 import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     ObservationProcessorStep,
     PolicyAction,
@@ -31,7 +32,6 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     IMAGENET_STATS,
     OBS_IMAGES,

src/lerobot/processor/__init__.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.types import (
+from lerobot.lerobot_types import (
     EnvAction,
     EnvTransition,
     PolicyAction,

src/lerobot/processor/batch_processor.py

@@ -25,7 +25,7 @@ from dataclasses import dataclass, field
 from torch import Tensor
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvTransition, PolicyAction
+from lerobot.lerobot_types import EnvTransition, PolicyAction
 from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
 from .pipeline import (

src/lerobot/processor/converters.py

@@ -23,7 +23,7 @@ from typing import Any
 import numpy as np
 import torch
 
-from lerobot.types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
+from lerobot.lerobot_types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
 from lerobot.utils.constants import ACTION, DONE, INFO, OBS_PREFIX, REWARD, TRUNCATED
 
 

src/lerobot/processor/delta_action_processor.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.types import PolicyAction, RobotAction
+from lerobot.lerobot_types import PolicyAction, RobotAction
 
 from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep
 

src/lerobot/processor/device_processor.py

@@ -25,7 +25,7 @@ from typing import Any
 import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvTransition, PolicyAction, TransitionKey
+from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
 from lerobot.utils.device_utils import get_safe_torch_device
 
 from .pipeline import ProcessorStep, ProcessorStepRegistry

src/lerobot/processor/factory.py

@@ -20,7 +20,7 @@ from typing import Any
 import torch
 
 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.types import PolicyAction, RobotAction, RobotObservation
+from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .batch_processor import AddBatchDimensionProcessorStep

src/lerobot/processor/gym_action_processor.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvAction, EnvTransition, PolicyAction, TransitionKey
+from lerobot.lerobot_types import EnvAction, EnvTransition, PolicyAction, TransitionKey
 
 from .converters import to_tensor
 from .hil_processor import TELEOP_ACTION_KEY

src/lerobot/processor/hil_processor.py

@@ -29,7 +29,7 @@ from lerobot.teleoperators.utils import TeleopEvents
 if TYPE_CHECKING:
     from lerobot.teleoperators.teleoperator import Teleoperator
 
-from lerobot.types import EnvTransition, PolicyAction, TransitionKey
+from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
 
 from .pipeline import (
     ComplementaryDataProcessorStep,

src/lerobot/processor/normalize_processor.py

@@ -25,7 +25,7 @@ import torch
 from torch import Tensor
 
 from lerobot.configs import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvTransition, PolicyAction, TransitionKey
+from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
 
 if TYPE_CHECKING:
     from lerobot.datasets import LeRobotDataset

src/lerobot/processor/pipeline.py

@@ -45,7 +45,14 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
+from lerobot.lerobot_types import (
+    EnvAction,
+    EnvTransition,
+    PolicyAction,
+    RobotAction,
+    RobotObservation,
+    TransitionKey,
+)
 from lerobot.utils.constants import HF_LEROBOT_HOME
 from lerobot.utils.hub import HubMixin
 

src/lerobot/processor/policy_robot_bridge.py

@@ -20,7 +20,7 @@ from typing import Any
 import torch
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.types import PolicyAction, RobotAction
+from lerobot.lerobot_types import PolicyAction, RobotAction
 from lerobot.utils.constants import ACTION
 
 from .pipeline import ActionProcessorStep, ProcessorStepRegistry

src/lerobot/processor/relative_action_processor.py

@@ -20,7 +20,7 @@ import torch
 from torch import Tensor
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvTransition, TransitionKey
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE
 
 from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep

src/lerobot/processor/render_messages_processor.py

@@ -23,7 +23,7 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.configs.recipe import TrainingRecipe
 from lerobot.datasets.language import LANGUAGE_EVENTS, LANGUAGE_PERSISTENT
 from lerobot.datasets.language_render import render_sample
-from lerobot.types import EnvTransition, TransitionKey
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.utils.utils import unwrap_scalar
 
 from .pipeline import ProcessorStep, ProcessorStepRegistry

src/lerobot/processor/tokenizer_processor.py

@@ -30,7 +30,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.types import EnvTransition, RobotObservation, TransitionKey
+from lerobot.lerobot_types import EnvTransition, RobotObservation, TransitionKey
 from lerobot.utils.constants import (
     ACTION_TOKEN_MASK,
     ACTION_TOKENS,

src/lerobot/rewards/robometer/compute_rabc_weights.py

@@ -57,10 +57,10 @@ import torch
 from tqdm import tqdm
 
 from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import TransitionKey
 from lerobot.rewards.robometer.configuration_robometer import RobometerConfig
 from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
 from lerobot.rewards.robometer.processor_robometer import RobometerEncoderProcessorStep
-from lerobot.types import TransitionKey
 
 DEFAULT_OUTPUT_FILENAME = "robometer_progress.parquet"
 

src/lerobot/rewards/robometer/processor_robometer.py

@@ -25,6 +25,7 @@ from PIL import Image
 from torch import Tensor
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -39,7 +40,6 @@ from lerobot.rewards.robometer.configuration_robometer import (
     RobometerConfig,
 )
 from lerobot.rewards.robometer.modeling_robometer import ROBOMETER_FEATURE_PREFIX
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_IMAGES,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,

src/lerobot/rewards/sarm/processor_sarm.py

@@ -47,6 +47,7 @@ else:
     Faker = None  # type: ignore[assignment, misc]
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -58,7 +59,6 @@ from lerobot.processor import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
-from lerobot.types import EnvTransition, PolicyAction, TransitionKey
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_sarm import SARMConfig

src/lerobot/rewards/topreward/compute_rabc_weights.py

@@ -48,10 +48,10 @@ import torch
 from tqdm import tqdm
 
 from lerobot.datasets import LeRobotDataset
+from lerobot.lerobot_types import TransitionKey
 from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig
 from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
 from lerobot.rewards.topreward.processor_topreward import TOPRewardEncoderProcessorStep
-from lerobot.types import TransitionKey
 
 DEFAULT_OUTPUT_FILENAME = "topreward_progress.parquet"
 

src/lerobot/rewards/topreward/processor_topreward.py

@@ -23,6 +23,7 @@ import torch
 from torch import Tensor
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -37,7 +38,6 @@ from lerobot.rewards.topreward.configuration_topreward import (
     DEFAULT_PROMPT_SUFFIX_TEMPLATE,
     TOPRewardConfig,
 )
-from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_IMAGES,
     OBS_PREFIX,

src/lerobot/rl/algorithms/base.py

@@ -28,7 +28,7 @@ from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_file as load_safetensors, save_file as save_safetensors
 from torch.optim import Optimizer
 
-from lerobot.types import BatchType
+from lerobot.lerobot_types import BatchType
 from lerobot.utils.hub import HubMixin
 
 from .configs import RLAlgorithmConfig, TrainingStats

src/lerobot/rl/algorithms/sac/sac_algorithm.py

@@ -26,6 +26,7 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 from torch.optim import Optimizer
 
+from lerobot.lerobot_types import BatchType
 from lerobot.policies.gaussian_actor.modeling_gaussian_actor import (
     DISCRETE_DIMENSION_INDEX,
     MLP,
@@ -35,7 +36,6 @@ from lerobot.policies.gaussian_actor.modeling_gaussian_actor import (
     orthogonal_init,
 )
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.types import BatchType
 from lerobot.utils.constants import ACTION
 from lerobot.utils.transition import move_state_dict_to_device
 

src/lerobot/rl/data_sources/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.types import BatchType
+from lerobot.lerobot_types import BatchType
 
 from .data_mixer import DataMixer, OnlineOfflineMixer
 

src/lerobot/rl/data_sources/data_mixer.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 
 import abc
 
-from lerobot.types import BatchType
+from lerobot.lerobot_types import BatchType
 
 from ..buffer import ReplayBuffer, concatenate_batch_transitions
 

src/lerobot/rl/trainer.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 from collections.abc import Iterator
 from typing import Any
 
-from lerobot.types import BatchType
+from lerobot.lerobot_types import BatchType
 
 from .algorithms.base import RLAlgorithm
 from .algorithms.configs import TrainingStats

src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 

src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 

src/lerobot/robots/bi_so_follower/bi_so_follower.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 
@@ -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/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py

@@ -23,7 +23,7 @@ import cv2
 import numpy as np
 import requests
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.errors import DeviceNotConnectedError
 

src/lerobot/robots/hope_jr/hope_jr_arm.py

@@ -19,12 +19,12 @@ import time
 from functools import cached_property
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorNormMode
 from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/hope_jr/hope_jr_hand.py

@@ -19,12 +19,12 @@ import time
 from functools import cached_property
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorNormMode
 from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/koch_follower/koch_follower.py

@@ -19,12 +19,12 @@ import time
 from functools import cached_property
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.dynamixel import (
     DynamixelMotorsBus,
     OperatingMode,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/lekiwi/lekiwi.py

@@ -23,12 +23,12 @@ from typing import Any
 import numpy as np
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/lekiwi/lekiwi_client.py

@@ -22,7 +22,7 @@ from functools import cached_property
 import cv2
 import numpy as np
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.errors import DeviceNotConnectedError

src/lerobot/robots/omx_follower/omx_follower.py

@@ -19,13 +19,13 @@ import time
 from functools import cached_property
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.dynamixel import (
     DriveMode,
     DynamixelMotorsBus,
     OperatingMode,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/openarm_follower/openarm_follower.py

@@ -20,9 +20,9 @@ from functools import cached_property
 from typing import Any
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.damiao import DamiaoMotorsBus
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/reachy2/robot_reachy2.py

@@ -19,7 +19,7 @@ import time
 from typing import TYPE_CHECKING, Any
 
 from lerobot.cameras import make_cameras_from_configs
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _reachy2_sdk_available, require_package
 
 from ..robot import Robot

src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py

@@ -21,8 +21,8 @@ from functools import cached_property
 from typing import TYPE_CHECKING
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import MotorCalibration
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.import_utils import _motorbridge_available, require_package
 

src/lerobot/robots/robot.py

@@ -18,8 +18,8 @@ from pathlib import Path
 
 import draccus
 
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import MotorCalibration
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS
 
 from .config import RobotConfig

src/lerobot/robots/so_follower/so_follower.py

@@ -19,12 +19,12 @@ import time
 from functools import cached_property
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
 import numpy as np
 
 from lerobot.cameras import make_cameras_from_configs
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _unitree_sdk_available, require_package
 
 from ..robot import Robot

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -127,12 +127,6 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
         --operation.new_task "Default task" \
         --operation.episode_tasks '{"5": "Special task for episode 5"}'
 
-Modify tasks - replace existing task strings in-place (WARNING: modifies in-place):
-    lerobot-edit-dataset \
-        --repo_id lerobot/pusht \
-        --operation.type modify_tasks \
-        --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
-
 Convert image dataset to video format and save locally:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht_image \
@@ -309,7 +303,6 @@ class RemoveFeatureConfig(OperationConfig):
 class ModifyTasksConfig(OperationConfig):
     new_task: str | None = None
     episode_tasks: dict[str, str] | None = None
-    task_replacements: dict[str, str] | None = None
 
 
 @OperationConfig.register_subclass("convert_image_to_video")
@@ -558,12 +551,9 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
 
     new_task = cfg.operation.new_task
     episode_tasks_raw = cfg.operation.episode_tasks
-    task_replacements = cfg.operation.task_replacements
 
-    if new_task is None and episode_tasks_raw is None and task_replacements is None:
-        raise ValueError(
-            "Must specify at least one of new_task, episode_tasks, or task_replacements for modify_tasks operation"
-        )
+    if new_task is None and episode_tasks_raw is None:
+        raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
 
     if cfg.new_repo_id is not None or cfg.new_root is not None:
         logging.warning(
@@ -583,14 +573,11 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
         logging.info(f"  Default task: '{new_task}'")
     if episode_tasks:
         logging.info(f"  Episode-specific tasks: {episode_tasks}")
-    if task_replacements:
-        logging.info(f"  Task replacements: {task_replacements}")
 
     modified_dataset = modify_tasks(
         dataset,
         new_task=new_task,
         episode_tasks=episode_tasks,
-        task_replacements=task_replacements,
     )
 
     logging.info(f"Dataset modified at {dataset.root}")

src/lerobot/scripts/lerobot_eval.py

@@ -82,9 +82,9 @@ from lerobot.envs import (
     make_env_pre_post_processors,
     preprocess_observation,
 )
+from lerobot.lerobot_types import PolicyAction
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.processor import PolicyProcessorPipeline
-from lerobot.types import PolicyAction
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
 from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package

src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 

src/lerobot/teleoperators/bi_openarm_mini/bi_openarm_mini.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 

src/lerobot/teleoperators/bi_rebot_102_leader/bi_rebot_102_leader.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 

src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py

@@ -17,7 +17,7 @@
 import logging
 from functools import cached_property
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 from lerobot.utils.bimanual import BimanualMixin
 from lerobot.utils.decorators import check_if_not_connected
 
@@ -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)

src/lerobot/teleoperators/gamepad/teleop_gamepad.py

@@ -21,7 +21,7 @@ from typing import Any
 
 import numpy as np
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 from lerobot.utils.decorators import check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/keyboard/teleop_keyboard.py

@@ -19,7 +19,7 @@ import time
 from queue import Queue
 from typing import Any
 
-from lerobot.types import RobotAction
+from lerobot.lerobot_types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.import_utils import _pynput_available, require_package
 from lerobot.utils.keyboard_input import pynput_can_capture

src/lerobot/teleoperators/openarm_leader/openarm_leader.py

@@ -18,9 +18,9 @@ import logging
 import time
 from typing import Any
 
+from lerobot.lerobot_types import RobotAction
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.damiao import DamiaoMotorsBus
-from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/openarm_mini/openarm_mini.py

@@ -18,12 +18,12 @@ import logging
 import time
 from typing import Any
 
+from lerobot.lerobot_types import RobotAction
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
-from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/phone/phone_processor.py

@@ -17,8 +17,8 @@
 from dataclasses import dataclass, field
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.lerobot_types import RobotAction
 from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
-from lerobot.types import RobotAction
 
 from .config_phone import PhoneOS
 

src/lerobot/teleoperators/rebot_102_leader/rebot_102_leader.py

@@ -18,8 +18,8 @@ import logging
 import time
 from typing import TYPE_CHECKING
 
+from lerobot.lerobot_types import RobotAction
 from lerobot.motors import MotorCalibration
-from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.import_utils import _motorbridge_smart_servo_available, require_package
 

src/lerobot/teleoperators/teleoperator.py

@@ -19,8 +19,8 @@ from typing import Any
 
 import draccus
 
+from lerobot.lerobot_types import RobotAction
 from lerobot.motors.motors_bus import MotorCalibration
-from lerobot.types import RobotAction
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS
 
 from .config import TeleoperatorConfig

src/lerobot/utils/foxglove_visualization.py

@@ -27,7 +27,7 @@ import time
 import cv2
 import numpy as np
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 
 from .constants import (
     ACTION,

src/lerobot/utils/rerun_visualization.py

@@ -25,7 +25,7 @@ import os
 import numpy as np
 
 from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 
 from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
 from .import_utils import require_package

src/lerobot/utils/visualization_utils.py

@@ -21,7 +21,7 @@ this module does not import ``rerun`` or ``foxglove`` (each backend imports its
 ``require_package`` guard).
 """
 
-from lerobot.types import RobotAction, RobotObservation
+from lerobot.lerobot_types import RobotAction, RobotObservation
 
 from .foxglove_visualization import init_foxglove, log_foxglove_data, shutdown_foxglove
 from .rerun_visualization import init_rerun, log_rerun_data, shutdown_rerun

tests/datasets/test_dataset_tools.py

@@ -1125,61 +1125,9 @@ def test_modify_tasks_default_with_overrides(sample_dataset):
             assert ep_data["tasks"][0] == default_task
 
 
-def test_modify_tasks_replacements(sample_dataset):
-    """Test replacing task strings based on their current values."""
-    modified_dataset = modify_tasks(
-        sample_dataset,
-        task_replacements={
-            "task_0": "Pick the cube",
-            "task_1": "Place the cube",
-        },
-    )
-
-    assert len(modified_dataset.meta.tasks) == 2
-    assert "Pick the cube" in modified_dataset.meta.tasks.index
-    assert "Place the cube" in modified_dataset.meta.tasks.index
-
-    for ep_idx in range(5):
-        expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Place the cube"
-        assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
-
-
-def test_modify_tasks_replacements_with_episode_overrides(sample_dataset):
-    """Test that explicit episode overrides take precedence over replacements."""
-    modified_dataset = modify_tasks(
-        sample_dataset,
-        task_replacements={
-            "task_0": "Pick the cube",
-            "task_1": "Place the cube",
-        },
-        episode_tasks={1: "Inspect the cube"},
-    )
-
-    assert modified_dataset.meta.episodes[0]["tasks"][0] == "Pick the cube"
-    assert modified_dataset.meta.episodes[1]["tasks"][0] == "Inspect the cube"
-    assert modified_dataset.meta.episodes[3]["tasks"][0] == "Place the cube"
-    assert len(modified_dataset.meta.tasks) == 3
-
-
-def test_modify_tasks_default_task_and_replacements(sample_dataset):
-    """Test that new_task acts as the default for episodes not matched by task_replacements."""
-    modified_dataset = modify_tasks(
-        sample_dataset,
-        new_task="Default task",
-        task_replacements={"task_0": "Pick the cube"},
-    )
-
-    for ep_idx in range(5):
-        expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Default task"
-        assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
-    assert len(modified_dataset.meta.tasks) == 2
-
-
 def test_modify_tasks_no_task_specified(sample_dataset):
     """Test error when no task is specified."""
-    with pytest.raises(
-        ValueError, match="Must specify at least one of new_task, episode_tasks, or task_replacements"
-    ):
+    with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
         modify_tasks(sample_dataset)
 
 
@@ -1189,12 +1137,6 @@ def test_modify_tasks_invalid_episode_indices(sample_dataset):
         modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
 
 
-def test_modify_tasks_invalid_task_replacements(sample_dataset):
-    """Test error when task replacements refer to unknown task strings."""
-    with pytest.raises(ValueError, match="Task replacements reference unknown tasks"):
-        modify_tasks(sample_dataset, task_replacements={"missing_task": "New task"})
-
-
 def test_modify_tasks_updates_info_json(sample_dataset):
     """Test that total_tasks is updated in info.json."""
     episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}

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",
     [

tests/envs/test_robocasa.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from unittest.mock import Mock, call
+
+import pytest
+
+from lerobot.envs import robocasa
+from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig
+
+
+def _instantiate_envs(
+    factories: Sequence[Callable[[], robocasa.RoboCasaEnv]],
+) -> list[robocasa.RoboCasaEnv]:
+    return [factory() for factory in factories]
+
+
+def test_robocasa_config_uses_registered_horizon_by_default() -> None:
+    assert RoboCasaEnvConfig().episode_length is None
+
+
+def test_multi_task_envs_use_registered_horizons(monkeypatch: pytest.MonkeyPatch) -> None:
+    horizons = {"CloseFridge": 900, "SearingMeat": 4350}
+    get_task_horizon = Mock(side_effect=horizons.__getitem__)
+    monkeypatch.setattr(robocasa, "_get_task_horizon", get_task_horizon)
+
+    envs = robocasa.create_robocasa_envs(
+        task="CloseFridge,SearingMeat",
+        n_envs=1,
+        env_cls=_instantiate_envs,
+    )
+
+    assert envs["CloseFridge"][0][0]._max_episode_steps == 900
+    assert envs["SearingMeat"][0][0]._max_episode_steps == 4350
+    assert get_task_horizon.call_args_list == [call("CloseFridge"), call("SearingMeat")]
+
+
+def test_explicit_episode_length_overrides_registered_horizons(monkeypatch: pytest.MonkeyPatch) -> None:
+    get_task_horizon = Mock()
+    monkeypatch.setattr(robocasa, "_get_task_horizon", get_task_horizon)
+
+    envs = robocasa.create_robocasa_envs(
+        task="CloseFridge,SearingMeat",
+        n_envs=1,
+        env_cls=_instantiate_envs,
+        episode_length=1234,
+    )
+
+    assert envs["CloseFridge"][0][0]._max_episode_steps == 1234
+    assert envs["SearingMeat"][0][0]._max_episode_steps == 1234
+    get_task_horizon.assert_not_called()

tests/mocks/mock_robot.py

@@ -19,9 +19,9 @@ from dataclasses import dataclass, field
 from functools import cached_property
 
 from lerobot.cameras import CameraConfig, make_cameras_from_configs
+from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.motors.motors_bus import Motor, MotorNormMode
 from lerobot.robots import Robot, RobotConfig
-from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from tests.mocks.mock_motors_bus import MockMotorsBus