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/hil_data_collection.mdx

@@ -59,6 +59,7 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
 
 - `bi_openarm_mini` - Bimanual OpenArm Mini
 - `so_leader` - SO100 / SO101 leader arm
+- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
 
 > [!IMPORTANT]
 > The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.

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/so101.mdx

@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
 <hfoption id="Leader">
 
 - Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
-- Attach the handle to motor 5 using 1 M2x6mm screw.
+- Attach the handle to the leader holder using 1 M2x6mm screw.
 - Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
 - Attach the follower trigger with 4 M3x6mm screws.
 

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
 

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     "einops>=0.8.0,<0.9.0",
 
     # Config & Hub
-    "draccus==0.10.0", # TODO: Relax version constraint
+    "draccus>=0.11.6,<0.12.0",
     "huggingface-hub>=1.0.0,<2.0.0",
     "requests>=2.32.0,<3.0.0",
 

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

@@ -163,8 +163,10 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
         return None
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/rewards.py

@@ -103,8 +103,10 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
         pass
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/train.py

@@ -194,7 +194,11 @@ class TrainPipelineConfig(HubMixin):
             )
 
         if Path(config_path).resolve().exists():
-            policy_dir = Path(config_path).parent
+            # `config_path` may point at the checkpoint's train_config.json or at its
+            # pretrained_model/ directory (both documented above) — resolve either to
+            # the pretrained_model/ directory.
+            config_path_obj = Path(config_path)
+            policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
             self.checkpoint_path = policy_dir.parent
         elif self.job.is_remote:
             return

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

@@ -42,6 +42,9 @@ class Evo1Policy(PreTrainedPolicy):
     config_class = Evo1Config
     name = "evo1"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
         super().__init__(config)
         config.validate_features()

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

@@ -68,6 +68,9 @@ class GrootPolicy(PreTrainedPolicy):
     name = "groot"
     config_class = GrootConfig
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(self, config: GrootConfig, **kwargs):
         """Initialize Groot policy wrapper."""
         require_package("transformers", extra="groot")

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

@@ -520,6 +520,9 @@ class MolmoAct2Policy(PreTrainedPolicy):
     config_class = MolmoAct2Config
     name = "molmoact2"
 
+    def supports_rtc(self) -> bool:
+        return self.config.inference_action_mode == "continuous"
+
     def __init__(
         self,
         config: MolmoAct2Config,

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/pi0/modeling_pi0.py

@@ -749,6 +749,9 @@ class PI0Policy(PreTrainedPolicy):
     config_class = PI0Config
     name = "pi0"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(
         self,
         config: PI0Config,

src/lerobot/policies/pi05/modeling_pi05.py

@@ -714,6 +714,9 @@ class PI05Policy(PreTrainedPolicy):
     config_class = PI05Config
     name = "pi05"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(
         self,
         config: PI05Config,

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

@@ -249,6 +249,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         """
         raise NotImplementedError
 
+    def supports_rtc(self) -> bool:
+        """Whether this policy implements Real-Time Chunking inference semantics."""
+        return False
+
     # TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
     @abc.abstractmethod
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -145,6 +145,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
     config_class = SmolVLAConfig
     name = "smolvla"
 
+    def supports_rtc(self) -> bool:
+        return True
+
     def __init__(
         self,
         config: SmolVLAConfig,

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -168,14 +168,23 @@ class SmolVLMWithExpertModel(nn.Module):
                 last_layers.append(self.num_vlm_layers - 2)
             frozen_layers = [
                 "lm_head",
-                "text_model.model.norm.weight",
+                "text_model.norm.weight",
             ]
             for layer in last_layers:
-                frozen_layers.append(f"text_model.model.layers.{layer}.")
+                frozen_layers.append(f"text_model.layers.{layer}.")
 
+            unmatched_patterns = set(frozen_layers)
             for name, params in self.vlm.named_parameters():
-                if any(k in name for k in frozen_layers):
+                matched_patterns = [k for k in frozen_layers if k in name]
+                if matched_patterns:
                     params.requires_grad = False
+                    unmatched_patterns.difference_update(matched_patterns)
+            if unmatched_patterns:
+                raise RuntimeError(
+                    "Some frozen layer patterns matched no VLM parameters, so the corresponding layers "
+                    "would silently remain trainable (parameter naming may have changed in transformers): "
+                    f"{sorted(unmatched_patterns)}"
+                )
         # To avoid unused params issue with distributed training
         for name, params in self.lm_expert.named_parameters():
             if "lm_head" in name:

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

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import abc
 import builtins
+import json
 import logging
 import os
 from dataclasses import dataclass, field
@@ -78,8 +79,10 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
 
     def _save_pretrained(self, save_directory: Path) -> None:
         """Serialize this config as ``config.json`` inside ``save_directory``."""
-        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
-            draccus.dump(self, f, indent=4)
+        # Encode against the base class so draccus includes the choice "type" key,
+        # which `from_pretrained` needs to resolve the concrete subclass.
+        with open(save_directory / CONFIG_NAME, "w") as f:
+            json.dump(draccus.encode(self, RLAlgorithmConfig), f, indent=4)
 
     @classmethod
     def from_pretrained(

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

@@ -68,10 +68,6 @@ class UnitreeG1Config(RobotConfig):
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Locomotion controller class name, e.g. "GrootLocomotionController",
-    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
-    # Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
-    # latent-token action/observation interface (``motion_token.{i}.pos`` action and a
-    # ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
-    # policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
+    # Lower-body controller class name, e.g. "GrootLocomotionController" or
+    # "HolosomaLocomotionController". None disables it.
     controller: str | None = None

src/lerobot/robots/unitree_g1/controllers/__init__.py

@@ -1,27 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
-
-from .gr00t_locomotion import GrootLocomotionController
-from .holosoma_locomotion import HolosomaLocomotionController
-from .sonic_whole_body import SonicWholeBodyController
-
-__all__ = [
-    "GrootLocomotionController",
-    "HolosomaLocomotionController",
-    "SonicWholeBodyController",
-]

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -1,360 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
-
-Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
-The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
-supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
-maps ``token + recent proprioception history`` to a residual action that is scaled and
-added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
-for the robot's PD controller.
-
-Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
-**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
-convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-
-import numpy as np
-import onnx
-import onnxruntime as ort
-from huggingface_hub import hf_hub_download
-
-from ..g1_utils import (
-    ISAACLAB_TO_MUJOCO,
-    MUJOCO_TO_ISAACLAB,
-    G1_29_JointIndex,
-    get_gravity_orientation,
-)
-
-logger = logging.getLogger(__name__)
-
-# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
-CONTROL_DT = 0.02  # 50 Hz control period (s)
-TOKEN_DIM = 64  # decoder latent size
-
-# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
-# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
-# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
-# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
-# motor-physics math happens at deploy time.
-DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
-DECODER_FILENAME = "model_decoder.onnx"
-DECODER_INPUT_DIM = 994  # token(64) + 10-frame proprio history + gravity
-
-
-def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
-    """Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
-
-    Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
-    gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
-    (64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
-    recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
-    """
-    decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
-    so = ort.SessionOptions()
-    so.log_severity_level = 3  # quiet ORT logs
-    session = ort.InferenceSession(decoder_path, sess_options=so)
-    dec_dim = int(session.get_inputs()[0].shape[1])
-    if dec_dim != DECODER_INPUT_DIM:
-        raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
-
-    meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
-    required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
-    missing = [k for k in required if k not in meta]
-    if missing:
-        raise ValueError(
-            f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
-            "re-run upload_sonic_decoder.py to (re)provision the checkpoint."
-        )
-    arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
-    logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
-    return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
-
-
-# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
-TOKEN_ACTION_PREFIX = "motion_token"  # nosec B105 - feature-key prefix, not a secret
-# Proprio-state prefix for the token interface: the robot echoes the last commanded token
-# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
-TOKEN_STATE_PREFIX = "motion_token_state"  # nosec B105 - feature-key prefix, not a secret
-
-
-def token_action_key(i: int) -> str:
-    """Action-dict key for the i-th component of the 64-D SONIC latent token.
-
-    The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which
-    only routes ``.pos`` scalar features onto the policy action vector.
-    """
-    return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
-
-
-def token_state_key(i: int) -> str:
-    """Observation key for the i-th component of the 64-D SONIC latent token state."""
-    return f"{TOKEN_STATE_PREFIX}.{i}.pos"
-
-
-# Startup blend duration: over the first control ticks, linearly interpolate every joint
-# from the robot's initial measured pose into the policy's commanded target, so control
-# eases in without a snap on the first command.
-INIT_RAMP_S = 3.0
-
-
-def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
-    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
-
-    The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
-    token-output VLA's action), which the decoder consumes with the encoder bypassed.
-    Requires the full dense token; a partial one is ignored (returns None).
-    """
-    if not action:
-        return None
-    keys = [token_action_key(i) for i in range(TOKEN_DIM)]
-    if any(key not in action for key in keys):
-        return None
-    return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
-
-
-class SonicDecoder:
-    """Runs the SONIC decoder ONNX model and owns the proprioception history.
-
-    Each tick it appends the latest robot state to 10-frame history buffers, then maps the
-    supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
-    The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
-    ``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
-    """
-
-    def __init__(self, decoder, default_angles, action_scale):
-        self.decoder = decoder
-        self.decoder_input = decoder.get_inputs()[0].name
-        self.default_angles = np.asarray(default_angles, np.float32)
-        self.action_scale = np.asarray(action_scale, np.float32)
-        self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-
-    def reset(self):
-        """Clear the token and 10-frame proprioception history.
-
-        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
-        are not contaminated by the previous episode's state.
-        """
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-
-    def update_history(self, q, dq, ang, quat):
-        """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
-        quat = quat / (np.linalg.norm(quat) + 1e-8)
-        # Reorder IsaacLab-order state into the MuJoCo order the decoder consumes. This
-        # permutation direction is validated against the deployed SONIC ONNX; don't flip it.
-        q_mj = q[MUJOCO_TO_ISAACLAB]
-        dq_mj = dq[MUJOCO_TO_ISAACLAB]
-        self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
-        self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
-        self.h_ang = [ang.copy()] + self.h_ang[:-1]
-        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
-        self.h_quat = [quat.copy()] + self.h_quat[:-1]
-
-    def build_decoder_obs(self):
-        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
-        obs = np.zeros(994, np.float32)
-        off = 0
-        obs[off : off + 64] = self.token
-        off += 64
-        for h, sz in [
-            (list(reversed(self.h_ang)), 3),
-            (list(reversed(self.h_q_mj)), 29),
-            (list(reversed(self.h_dq_mj)), 29),
-            (list(reversed(self.h_act_mj)), 29),
-        ]:
-            for f in range(10):
-                obs[off : off + sz] = h[f]
-                off += sz
-        for q in reversed(self.h_quat):
-            obs[off : off + 3] = get_gravity_orientation(q)
-            off += 3
-        assert off == 994, f"Decoder obs mismatch: {off}"
-        return obs
-
-    def step(self, robot_obs, token, debug=False):
-        """One control tick: read robot obs, decode the supplied token -> joint targets.
-
-        Args:
-            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
-            token: 64-D latent supplied by the policy (encoder bypassed).
-            debug: log action/delta norms.
-
-        Returns:
-            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
-        """
-        self.token = np.asarray(token, np.float32)
-        jnames = [m.name for m in G1_29_JointIndex]
-        q = np.array(
-            [
-                robot_obs.get(f"{n}.q", self.default_angles[m.value])
-                for m, n in zip(G1_29_JointIndex, jnames, strict=False)
-            ],
-            np.float32,
-        )
-        dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
-        quat = np.array(
-            [
-                robot_obs.get("imu.quat.w", 1),
-                robot_obs.get("imu.quat.x", 0),
-                robot_obs.get("imu.quat.y", 0),
-                robot_obs.get("imu.quat.z", 0),
-            ],
-            np.float32,
-        )
-        ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
-        self.update_history(q, dq, ang, quat)
-        action_mj = (
-            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
-            .squeeze()
-            .astype(np.float32)
-        )
-        self.last_action_mj = action_mj.copy()
-        target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
-        if debug:
-            delta = target - q
-            logger.debug(
-                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
-                np.linalg.norm(self.token),
-                np.linalg.norm(action_mj),
-                np.max(np.abs(delta)),
-                np.sqrt(np.mean(delta**2)),
-            )
-        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
-
-
-class SonicRuntime:
-    """Loads the SONIC decoder ONNX model and owns the decode controller.
-
-    Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
-    latent token supplied directly by the policy.
-    """
-
-    def __init__(self):
-        decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
-        self.default_angles = default_angles
-        self.neutral_token = neutral_token
-        self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
-
-    @property
-    def pipeline(self):
-        return self.controller
-
-    def reset(self):
-        self.controller.reset()
-
-    def shutdown(self):
-        pass
-
-
-class SonicWholeBodyController:
-    """Full-body SONIC controller for UnitreeG1's background controller thread."""
-
-    control_dt = CONTROL_DT
-    full_body = True
-
-    def __init__(self):
-        logger.info("Loading SONIC whole-body controller...")
-        self._runtime = SonicRuntime()
-        self.kp = self._runtime.kp
-        self.kd = self._runtime.kd
-        self.controller = self._runtime.controller
-        self._default_angles = self._runtime.default_angles
-        self._neutral_token = self._runtime.neutral_token
-
-        # Startup blend: ease from the robot's initial pose into the first commanded policy
-        # targets over INIT_RAMP_S (captured on the first control tick).
-        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
-        self._init_step = 0
-        self._start_pose: dict[str, float] = {}
-
-        # Token-interface state. The controller holds a stable *neutral* token until the first
-        # real token arrives, and afterwards holds the *last* token received between ticks (the
-        # async controller runs ~50 Hz while a token VLA streams ~30 Hz).
-        self._last_token: np.ndarray | None = None
-
-        logger.info("SONIC ready (decoder, 64-D token command path)")
-
-    def _startup_blend(self, obs: dict, out: dict) -> dict:
-        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
-        interpolate between the robot's pose captured on the first tick and the policy's
-        live commanded target, so the handoff has no snap.
-
-        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
-        climbs 0->1 over the ramp, after which the raw policy target passes through.
-        """
-        if self._init_step >= self._init_ramp_steps or not out:
-            return out
-        if self._init_step == 0:
-            # Capture the robot's actual pose as the interpolation start point.
-            self._start_pose = {
-                f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
-                for m in G1_29_JointIndex
-            }
-        self._init_step += 1
-        ratio = min(1.0, self._init_step / self._init_ramp_steps)
-        blended = {
-            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
-            for k, tgt in out.items()
-        }
-        if self._init_step >= self._init_ramp_steps:
-            logger.info("SONIC startup blend complete -> full policy control")
-        return blended
-
-    def run_step(self, action: dict, obs: dict) -> dict:
-        if not obs:
-            return {}
-
-        # Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
-        # is decoded directly, encoder bypassed.
-        token = _extract_token_from_action(action)
-        if token is not None:
-            self._last_token = token
-        elif self._last_token is None:
-            # No token has arrived yet: hold the checkpoint's neutral token, which the decoder
-            # maps to a stable, natural standing pose.
-            self._last_token = self._neutral_token.copy()
-        # Either a fresh token this tick or the last one received (held between the ~30 Hz
-        # token stream and the ~50 Hz control loop).
-        return self._startup_blend(obs, self.controller.step(obs, self._last_token))
-
-    def reset(self):
-        self._runtime.reset()
-        self._init_step = 0  # re-run the startup blend after a reset
-        self._start_pose = {}
-        # Drop the held token so the neutral token is re-seeded after a reset.
-        self._last_token = None
-
-    def shutdown(self):
-        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,47 +23,6 @@ import numpy as np
 
 NUM_MOTORS = 29
 
-# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
-# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
-# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
-ISAACLAB_TO_MUJOCO = np.array(
-    [
-        0,
-        3,
-        6,
-        9,
-        13,
-        17,
-        1,
-        4,
-        7,
-        10,
-        14,
-        18,
-        2,
-        5,
-        8,
-        11,
-        15,
-        19,
-        21,
-        23,
-        25,
-        27,
-        12,
-        16,
-        20,
-        22,
-        24,
-        26,
-        28,
-    ],
-    dtype=np.int32,
-)
-# The two orderings are inverses of each other, so derive one from the other (argsort) to
-# guarantee they can never drift out of sync.
-MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
-
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -109,9 +68,8 @@ def make_locomotion_controller(name: str | None):
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
-        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/gr00t_locomotion.py

@@ -21,7 +21,7 @@ import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/holosoma_locomotion.py

@@ -22,7 +22,7 @@ import onnx
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,

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
@@ -34,6 +34,7 @@ from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
     REMOTE_AXES,
+    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
@@ -147,54 +148,22 @@ class UnitreeG1(Robot):
 
         self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
-        # Lower-body / whole-body controller loaded dynamically
+        # Lower-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
+
         # Controller thread state
         self._controller_thread = None
-        # When set, the controller loop stops publishing low commands so reset() can
-        # drive the joints directly without two publishers fighting (single-publisher).
-        self._controller_paused = threading.Event()
         self._controller_action_lock = threading.Lock()
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
-        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
-        # echoed back as ``observation.state`` so a token-output VLA closes the loop
-        # on its own previous token. Implicit whenever the SONIC whole-body controller
-        # is active. Seeded to zeros; the controller's startup blend eases joints in.
-        self._last_token: np.ndarray | None = None
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import TOKEN_DIM
-
-            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
-
-    @property
-    def _sonic_token(self) -> bool:
-        """Whether the SONIC whole-body decoder is active.
-
-        A SONIC controller consumes a 64-D latent motion token as its action and echoes
-        the last commanded token as ``observation.state``. Keyed purely off the selected
-        controller so the token interface is implicit -- no separate config flag.
-        """
-        return self.config.controller == "SonicWholeBodyController"
-
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
-                try:
-                    self.sim_env.step()
-                except ValueError as e:
-                    # Startup race: the sim thread can step once before reset() has
-                    # written a valid base pose, giving a zero-norm pelvis quaternion
-                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
-                    # the thread survives instead of dying and freezing the sim.
-                    if "zero norm" not in str(e).lower():
-                        raise
-                    time.sleep(self.control_dt)
-                    continue
+                self.sim_env.step()
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -262,38 +231,15 @@ class UnitreeG1(Robot):
                 features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
         return features
 
-    @property
-    def _token_state_ft(self) -> dict[str, type]:
-        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
-
-        Exposed only when a SONIC whole-body controller is active; aggregated by the
-        rollout into a 64-D ``observation.state`` (the last token the policy commanded).
-        """
-        if not self._sonic_token:
-            return {}
-        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
-
-        return {token_state_key(i): float for i in range(TOKEN_DIM)}
-
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
+        return {**self._motors_ft, **self._cameras_ft}
 
     @cached_property
     def action_features(self) -> dict[str, type]:
-        # No controller configured at all: raw 29-DoF joint teleop.
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
-        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
-        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
-        # straight onto the decoder, bypassing the encoder.
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
-
-            return {token_action_key(i): float for i in range(TOKEN_DIM)}
-
-        # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
         arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
         remote_features = dict.fromkeys(REMOTE_AXES, float)
         return {**arm_features, **remote_features}
@@ -309,11 +255,6 @@ class UnitreeG1(Robot):
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
-            # Paused during reset() so the reset routine is the sole low-cmd publisher.
-            if self._controller_paused.is_set():
-                time.sleep(control_dt)
-                continue
-
             with self._lowstate_lock:
                 lowstate = self._lowstate
 
@@ -330,12 +271,8 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
-                # Full-body controllers (SONIC) consume the full observation dict; others
-                # take the raw lowstate. get_observation() is the single lowstate -> obs builder.
-                controller_state = (
-                    self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
-                )
-                controller_action = self.controller.run_step(controller_input, controller_state)
+                # Run controller step
+                controller_action = self.controller.run_step(controller_input, lowstate)
 
                 # Write controller output snapshot
                 with self._controller_action_lock:
@@ -406,9 +343,6 @@ class UnitreeG1(Robot):
 
         self.kp = np.array(self.config.kp, dtype=np.float32)
         self.kd = np.array(self.config.kd, dtype=np.float32)
-        if self.controller is not None and hasattr(self.controller, "kp"):
-            self.kp = np.array(self.controller.kp, dtype=np.float32)
-            self.kd = np.array(self.controller.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -457,10 +391,6 @@ class UnitreeG1(Robot):
             if self._controller_thread.is_alive():
                 logger.warning("Controller thread did not stop cleanly")
 
-        # Release controller resources (e.g. SONIC decoder sessions).
-        if self.controller is not None and hasattr(self.controller, "shutdown"):
-            self.controller.shutdown()
-
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -531,15 +461,6 @@ class UnitreeG1(Robot):
         if lowstate.wireless_remote:
             obs["wireless_remote"] = lowstate.wireless_remote
 
-        # Token mode: echo the last commanded latent token as observation.state so a
-        # token-output VLA closes the loop on its own previous token.
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import token_state_key
-
-            token = self._last_token if self._last_token is not None else []
-            for i, v in enumerate(token):
-                obs[token_state_key(i)] = float(v)
-
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
             if getattr(cam, "use_rgb", True):
@@ -552,22 +473,9 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
-            # SONIC decoder: pull the 64-D latent token out of the action and remember it
-            # for the observation.state echo. The controller thread reads it back from
-            # controller_input (populated below) and decodes it into a 29-DoF command.
-            if self._sonic_token:
-                from .controllers.sonic_whole_body import _extract_token_from_action
-
-                token = _extract_token_from_action(action)
-                if token is not None:
-                    self._last_token = token
-            self._update_controller_action(action)
-            # Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
-            # publish here (the controller thread is the sole publisher).
-            if getattr(self.controller, "full_body", False):
-                return action
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
+            self._update_controller_action(action)
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -595,17 +503,11 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from an incoming teleop action.
-
-        Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
-        SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
-        and each controller extracts only the keys it understands. The robot deliberately
-        does not enumerate any controller's key schema here.
-        """
+        """Update controller input state from incoming teleop action."""
         with self._controller_action_lock:
-            for key, value in action.items():
-                if isinstance(key, str) and value is not None:
-                    self.controller_input[key] = value
+            for key in REMOTE_KEYS:
+                if key in action:
+                    self.controller_input[key] = action[key]
 
     @property
     def is_calibrated(self) -> bool:
@@ -635,64 +537,43 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        # Full-body controllers (SONIC) own the whole 29-DoF command and ignore
-        # ``<joint>.q`` in send_action(), so reset() must publish the default pose
-        # directly. Pause the background controller first so the two aren't both writing
-        # low commands while the robot moves to the default pose.
-        full_body = getattr(self.controller, "full_body", False)
-        paused = False
-        if full_body and self._controller_thread is not None:
-            self._controller_paused.set()
-            paused = True
-            time.sleep(control_dt)  # let any in-flight controller tick settle
+        if self.config.is_simulation and self.sim_env is not None:
+            self.sim_env.reset()
+            self.publish_lowcmd(
+                {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
+            )
+        else:
+            total_time = 3.0
+            num_steps = int(total_time / control_dt)
 
-        try:
-            if self.config.is_simulation and self.sim_env is not None:
-                self.sim_env.reset()
-                self.publish_lowcmd(
-                    {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
-                )
-            else:
-                total_time = 3.0
-                num_steps = int(total_time / control_dt)
+            # get current state
+            obs = self.get_observation()
 
-                # get current state
-                obs = self.get_observation()
+            # record current positions
+            init_dof_pos = np.zeros(29, dtype=np.float32)
+            for motor in G1_29_JointIndex:
+                init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                # record current positions
-                init_dof_pos = np.zeros(29, dtype=np.float32)
+            # Interpolate to default position
+            for step in range(num_steps):
+                start_time = time.time()
+
+                alpha = step / num_steps
+                action_dict = {}
                 for motor in G1_29_JointIndex:
-                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
+                    target_pos = default_positions[motor.value]
+                    interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
+                    action_dict[f"{motor.name}.q"] = float(interp_pos)
 
-                # Interpolate to default position
-                for step in range(num_steps):
-                    start_time = time.time()
+                self.send_action(action_dict)
 
-                    alpha = step / num_steps
-                    action_dict = {}
-                    for motor in G1_29_JointIndex:
-                        target_pos = default_positions[motor.value]
-                        interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
-                        action_dict[f"{motor.name}.q"] = float(interp_pos)
+                # Maintain constant control rate
+                elapsed = time.time() - start_time
+                sleep_time = max(0, control_dt - elapsed)
+                time.sleep(sleep_time)
 
-                    # Full-body controllers no-op in send_action(); publish the pose
-                    # directly (arm-only controllers keep the send_action() path).
-                    if full_body:
-                        self.publish_lowcmd(action_dict)
-                    else:
-                        self.send_action(action_dict)
-
-                    # Maintain constant control rate
-                    elapsed = time.time() - start_time
-                    sleep_time = max(0, control_dt - elapsed)
-                    time.sleep(sleep_time)
-
-            # Reset controller internal state (gait phase, obs history, etc.) before
-            # resuming so its buffers reflect the post-reset pose.
-            if self.controller is not None and hasattr(self.controller, "reset"):
-                self.controller.reset()
-        finally:
-            if paused:
-                self._controller_paused.clear()
+        # Reset controller internal state (gait phase, obs history, etc.)
+        if self.controller is not None and hasattr(self.controller, "reset"):
+            self.controller.reset()
 
         logger.info("Reached default position")

src/lerobot/rollout/context.py

@@ -57,6 +57,7 @@ from .inference import (
     SyncInferenceConfig,
     create_inference_engine,
 )
+from .inference.rtc import supports_rtc_inference
 from .robot_wrapper import ThreadSafeRobot
 
 if TYPE_CHECKING or _peft_available:
@@ -226,6 +227,12 @@ def build_rollout_context(
     policy = _load_pretrained_policy(policy_config)
 
     if is_rtc:
+        if not supports_rtc_inference(policy):
+            raise ValueError(
+                f"RTC inference is not supported by policy type '{policy_config.type}': "
+                "the policy must implement RTC semantics and predict_action_chunk must accept "
+                "inference_delay and prev_chunk_left_over. Use '--inference.type=sync' instead."
+            )
         policy.config.rtc_config = cfg.inference.rtc
         if hasattr(policy, "init_rtc_processor"):
             policy.init_rtc_processor()