fix(dataset): bump versions + improvements

The prior test patched _load_metadata as a plain no-op, so __init__'s try
block succeeded and never entered the except branch where the
repo_type != "bucket" guard and get_safe_version live - the assertion passed
vacuously (verified: it still passed with the guard removed).

Use side_effect=[FileNotFoundError, None] so the first _load_metadata raises
(forcing the except path) and the second succeeds after the meta pull, and
stub _pull_from_repo so the path runs without a network call. Now the test
fails if the bucket guard is removed. Thanks @mohitydv09 for the catch.

Signed-off-by: Sundar Raghavan <sdraghav@amazon.com>

AGENT_GUIDE.md

@@ -61,15 +61,19 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
 **4.1 Install**
 
 ```bash
-pip install 'lerobot[feetech]'              # SO-100/SO-101 motor stack
-# pip install 'lerobot[all]'                # everything
-# pip install 'lerobot[aloha,pusht]'        # specific features
-# pip install 'lerobot[smolvla]'            # add SmolVLA deps
-git lfs install && git lfs pull
-hf auth login                               # required to push datasets/policies
-```
+# uv (recommended — see AGENTS.md and CLAUDE.md)
+uv sync --locked --extra feetech          # SO-100/SO-101 motor stack
+# uv sync --locked --extra all            # everything
+# uv sync --locked --extra smolvla        # add SmolVLA deps
 
-Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
+# pip (alternative, e.g. when not working from source)
+# pip install 'lerobot[feetech]'
+# pip install 'lerobot[all]'
+# pip install 'lerobot[smolvla]'
+
+git lfs install && git lfs pull
+hf auth login                             # required to push datasets/policies
+```
 
 **4.2 Find USB ports** — run once per arm, unplug when prompted.
 
@@ -321,11 +325,10 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p
 
 ```bash
 lerobot-train ... --policy.type=smolvla \
-  --policy.fine_tune_vision_encoder=true
+  --policy.freeze_vision_encoder=false \
+  --policy.train_expert_only=false
 ```
 
-This selectively trains the vision encoder and connector while leaving the language model frozen. Their learning rate defaults to `0.1 × optimizer_lr`; adjust it with `--policy.vision_encoder_lr_multiplier` if needed.
-
 ### 7.7 Signals to stop / keep going
 
 - Train loss plateaus → stop, save a Hub checkpoint.

docs/source/env_processor.mdx

@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
 
 The same policy can work with different environment processors, and the same environment processor can work with different policies:
 
-````python
-# Use SmolVLA policy with LIBERO environment
-# Use SmolVLA policy with LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=smolvla_cfg,
-)
-smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
-# Or use ACT policy with the same LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=act_cfg,
-)
-act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
     policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
+```
 
 ### 3. **Easier Experimentation**
 
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
         state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                           gripper_pos, gripper_vel], dim=-1)  # 14D
         return state
-````
+```
 
 ### 4. **Cleaner Environment Code**
 

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

@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 
 ### Record
 
-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
 
 ```bash
 lerobot-record \

docs/source/integrate_hardware.mdx

@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
 
 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
 
 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
 

docs/source/lekiwi.mdx

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```
 
-Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
+Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
 
 # Step-by-Step Assembly Instructions

docs/source/pi0fast.mdx

@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
 
 ## Reproducing π₀Fast results
 
-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
 
 The finetuned model can be found here:
 

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

@@ -70,19 +70,6 @@ cd lerobot && lerobot-train \
   GPU allows it, as long as loading times remain short.
 </Tip>
 
-For tasks that require adapting visual features, such as distinguishing new colors or shapes, selectively
-fine-tune the vision encoder and its connector:
-
-```bash
-lerobot-train ... \
-  --policy.path=lerobot/smolvla_base \
-  --policy.fine_tune_vision_encoder=true
-```
-
-This keeps the language model frozen with the default `train_expert_only=true` setting and trains the vision
-path at `0.1` times the main learning rate by default. Fine-tuning the vision encoder increases memory use and
-can reduce the model's general visual knowledge, so enable it only when the frozen encoder is insufficient.
-
 Fine-tuning is an art. For a complete overview of the options for finetuning, run
 
 ```bash
@@ -106,7 +93,7 @@ lerobot-train --help
 
 ## Evaluate the finetuned model and run it in real-time
 
-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```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,8 +67,8 @@ dependencies = [
     "einops>=0.8.0,<0.9.0",
 
     # Config & Hub
-    "draccus==0.10.0", # TODO: Relax version constraint
-    "huggingface-hub>=1.0.0,<2.0.0",
+    "draccus>=0.11.6,<0.12.0",
+    "huggingface-hub>=1.6.0,<2.0.0",
     "requests>=2.32.0,<3.0.0",
 
     # Environments
@@ -95,7 +95,7 @@ dependencies = [
 
 # ── Feature-scoped extras ──────────────────────────────────
 dataset = [
-    "datasets>=4.7.0,<5.0.0",
+    "datasets>=4.8.0,<5.0.0",
     "pandas>=2.0.0,<3.0.0", # NOTE: Transitive dependency of datasets
     "pyarrow>=21.0.0,<30.0.0", # NOTE: Transitive dependency of datasets
     "lerobot[av-dep]",

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

@@ -18,13 +18,15 @@ import logging
 from collections.abc import Callable, Iterable
 from copy import deepcopy
 from pathlib import Path
+from typing import Literal
 
 import numpy as np
 import packaging.version
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download, sync_bucket
+from huggingface_hub.utils import WeakFileLock
 
 from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig
 from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
@@ -74,6 +76,7 @@ class LeRobotDatasetMetadata:
         force_cache_sync: bool = False,
         metadata_buffer_size: int = 10,
         *,
+        repo_type: Literal["dataset", "bucket"] = "dataset",
         token: str | bool | None = None,
     ):
         """Load or download metadata for an existing LeRobot dataset.
@@ -96,36 +99,53 @@ class LeRobotDatasetMetadata:
                 even when local files exist.
             metadata_buffer_size: Number of episode metadata records to buffer
                 in memory before flushing to parquet.
+            repo_type: Repository type: "dataset" (default) or "bucket" for an
+                HF Storage Bucket streamed over hf://buckets/.
             token: Authentication token used for Hub requests. Pass a string
                 token, ``True`` to require the locally stored token, ``False``
                 to disable authentication, or ``None`` to use the Hugging Face
                 Hub default.
         """
+        if repo_type not in ("dataset", "bucket"):
+            raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {repo_type!r}")
+
         self.repo_id = repo_id
+        self.repo_type = repo_type
         self.revision = revision if revision else CODEBASE_VERSION
         self._requested_root = Path(root) if root is not None else None
-        self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id
+        if self._requested_root is not None:
+            self.root = self._requested_root
+        elif self.repo_type == "bucket":
+            self.root = HF_LEROBOT_HUB_CACHE / ("buckets--" + self.repo_id.replace("/", "--"))
+        else:
+            self.root = HF_LEROBOT_HOME / repo_id
         self._pq_writer = None
         self.latest_episode = None
         self._metadata_buffer: list[dict] = []
         self._metadata_buffer_size = metadata_buffer_size
         self._finalized = False
 
-        try:
-            if force_cache_sync or (
-                self._requested_root is None and has_legacy_hub_download_metadata(self.root)
-            ):
-                raise FileNotFoundError
-            self._load_metadata()
-        except (FileNotFoundError, NotADirectoryError):
-            if is_valid_version(self.revision):
-                if token is None:
-                    self.revision = get_safe_version(self.repo_id, self.revision)
-                else:
-                    self.revision = get_safe_version(self.repo_id, self.revision, token=token)
+        metadata_lock = contextlib.nullcontext()
+        if self.repo_type == "bucket":
+            self.root.parent.mkdir(parents=True, exist_ok=True)
+            metadata_lock = WeakFileLock(self.root.parent / f".{self.root.name}.lock")
 
-            self._pull_from_repo(allow_patterns="meta/", token=token)
-            self._load_metadata()
+        with metadata_lock:
+            try:
+                if force_cache_sync or (
+                    self._requested_root is None and has_legacy_hub_download_metadata(self.root)
+                ):
+                    raise FileNotFoundError
+                self._load_metadata()
+            except (FileNotFoundError, NotADirectoryError):
+                if self.repo_type != "bucket" and is_valid_version(self.revision):
+                    if token is None:
+                        self.revision = get_safe_version(self.repo_id, self.revision)
+                    else:
+                        self.revision = get_safe_version(self.repo_id, self.revision, token=token)
+
+                self._pull_from_repo(allow_patterns="meta/", token=token)
+                self._load_metadata()
 
     def _flush_metadata_buffer(self) -> None:
         """Write all buffered episode metadata to parquet file."""
@@ -232,6 +252,16 @@ class LeRobotDatasetMetadata:
         *,
         token: str | bool | None = None,
     ) -> None:
+        if self.repo_type == "bucket":
+            self.root.mkdir(parents=True, exist_ok=True)
+            sync_bucket(
+                f"hf://buckets/{self.repo_id}/meta",
+                str(self.root / "meta"),
+                delete=True,
+                quiet=True,
+                token=token,
+            )
+            return
         token_kwargs = {} if token is None else {"token": token}
         if self._requested_root is None:
             self.root = Path(
@@ -262,6 +292,8 @@ class LeRobotDatasetMetadata:
     @property
     def url_root(self) -> str:
         """Hugging Face Hub URL root for this dataset."""
+        if self.repo_type == "bucket":
+            return f"hf://buckets/{self.repo_id}"
         return f"hf://datasets/{self.repo_id}"
 
     @property

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

@@ -16,6 +16,7 @@
 from collections import deque
 from collections.abc import Callable, Generator, Iterable, Iterator
 from pathlib import Path
+from typing import Literal
 
 import datasets
 import numpy as np
@@ -257,15 +258,16 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         return_uint8: bool = False,
         depth_output_unit: str = DEFAULT_DEPTH_UNIT,
         *,
+        repo_type: Literal["dataset", "bucket"] = "dataset",
         token: str | bool | None = None,
     ):
         """Initialize a StreamingLeRobotDataset.
 
         Args:
             repo_id (str): This is the repo id that will be used to fetch the dataset.
-            root (Path | None, optional): Local directory to use for local datasets. When omitted, Hub
-                metadata is resolved through a revision-safe snapshot cache under
-                ``$HF_LEROBOT_HOME/hub``.
+            root (Path | None, optional): Local directory to use for local datasets. In bucket mode,
+                this is an optional local metadata-cache directory; parquet and video data remain remote.
+                When omitted, Hub metadata is resolved through the cache under ``$HF_LEROBOT_HOME/hub``.
             episodes (list[int] | None, optional): If specified, this will only load episodes specified by
                 their episode_index in this list.
             image_transforms (Callable | None, optional): Transform to apply to image data.
@@ -280,6 +282,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
             depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
                 Defaults to "mm".
+            repo_type: "dataset" (default) or "bucket" to stream from an HF Storage Bucket
+                over ``hf://buckets/``.
             token: Authentication token used while streaming this dataset from
                 the Hub. Pass a string token, ``True`` to require the locally
                 stored token, ``False`` to disable authentication, or ``None``
@@ -287,10 +291,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                 on the dataset instance after initialization.
         """
         super().__init__()
+        if repo_type not in ("dataset", "bucket"):
+            raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {repo_type!r}")
+
         self.repo_id = repo_id
-        self._requested_root = Path(root) if root else None
+        self.repo_type = repo_type
+        self._requested_root = Path(root) if root is not None else None
         self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id
-        self.streaming_from_local = root is not None
+        self.streaming_from_local = root is not None and self.repo_type == "dataset"
 
         self.image_transforms = image_transforms
         self.episodes = episodes
@@ -317,6 +325,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             self._requested_root,
             self.revision,
             force_cache_sync=force_cache_sync,
+            repo_type=self.repo_type,
             token=token,
         )
         self.root = self.meta.root
@@ -345,15 +354,26 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             self.delta_timestamps = delta_timestamps
             self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
 
-        token_kwargs = {} if token is None or self.streaming_from_local else {"token": token}
-        self.hf_dataset: datasets.IterableDataset = load_dataset(
-            self.repo_id if not self.streaming_from_local else str(self.root),
-            split="train",
-            streaming=self.streaming,
-            data_files="data/*/*.parquet",
-            revision=self.revision,
-            **token_kwargs,
-        )
+        token_kwargs = {} if token is None else {"token": token}
+        if self.repo_type == "bucket":
+            self.hf_dataset: datasets.IterableDataset = load_dataset(
+                "parquet",
+                data_files=f"hf://buckets/{self.repo_id}/data/*/*.parquet",
+                split="train",
+                streaming=self.streaming,
+                **token_kwargs,
+            )
+        else:
+            if self.streaming_from_local:
+                token_kwargs = {}
+            self.hf_dataset: datasets.IterableDataset = load_dataset(
+                self.repo_id if not self.streaming_from_local else str(self.root),
+                split="train",
+                streaming=self.streaming,
+                data_files="data/*/*.parquet",
+                revision=self.revision,
+                **token_kwargs,
+            )
 
         self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
 

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

@@ -67,8 +67,6 @@ class SmolVLAConfig(PreTrainedConfig):
 
     # Finetuning settings
     freeze_vision_encoder: bool = True
-    fine_tune_vision_encoder: bool = False  # Fine-tune vision + connector; takes priority over freezing.
-    vision_encoder_lr_multiplier: float = 0.1
     train_expert_only: bool = True
     train_state_proj: bool = True
 
@@ -112,12 +110,6 @@ class SmolVLAConfig(PreTrainedConfig):
         super().__post_init__()
 
         """Input validation (not exhaustive)."""
-        if self.fine_tune_vision_encoder:
-            self.freeze_vision_encoder = False
-        if self.vision_encoder_lr_multiplier <= 0:
-            raise ValueError(
-                f"`vision_encoder_lr_multiplier` must be positive, got {self.vision_encoder_lr_multiplier}."
-            )
         if self.n_action_steps > self.chunk_size:
             raise ValueError(
                 f"The chunk size is the upper bound for the number of action steps per model invocation. Got "

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,
@@ -186,27 +189,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
             if model_value is not None:
                 model_value.rtc_processor = self.rtc_processor
 
-    def get_optim_params(self):
-        if not self.config.fine_tune_vision_encoder:
-            return self.parameters()
-
-        vision_params = []
-        other_params = []
-        for name, param in self.named_parameters():
-            if not param.requires_grad:
-                continue
-            if ".vision_model." in name or ".connector." in name:
-                vision_params.append(param)
-            else:
-                other_params.append(param)
-
-        return [
-            {"params": other_params},
-            {
-                "params": vision_params,
-                "lr": self.config.optimizer_lr * self.config.vision_encoder_lr_multiplier,
-            },
-        ]
+    def get_optim_params(self) -> dict:
+        return self.parameters()
 
     def _get_action_chunk(
         self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
@@ -512,7 +496,6 @@ class VLAFlowMatching(nn.Module):
         self.vlm_with_expert = SmolVLMWithExpertModel(
             model_id=self.config.vlm_model_name,
             freeze_vision_encoder=self.config.freeze_vision_encoder,
-            fine_tune_vision_encoder=self.config.fine_tune_vision_encoder,
             train_expert_only=self.config.train_expert_only,
             load_vlm_weights=self.config.load_vlm_weights,
             attention_mode=self.config.attention_mode,

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -78,7 +78,6 @@ class SmolVLMWithExpertModel(nn.Module):
         load_vlm_weights: bool = True,
         train_expert_only: bool = True,
         freeze_vision_encoder: bool = False,
-        fine_tune_vision_encoder: bool = False,
         attention_mode: str = "self_attn",
         num_expert_layers: int = -1,
         num_vlm_layers: int = -1,
@@ -142,7 +141,6 @@ class SmolVLMWithExpertModel(nn.Module):
         self.num_key_value_heads = self.config.text_config.num_key_value_heads
 
         self.freeze_vision_encoder = freeze_vision_encoder
-        self.fine_tune_vision_encoder = fine_tune_vision_encoder
         self.train_expert_only = train_expert_only
         self.attention_mode = attention_mode
         self.expert_hidden_size = lm_expert_config.hidden_size
@@ -152,6 +150,10 @@ class SmolVLMWithExpertModel(nn.Module):
         return self.vlm.model
 
     def set_requires_grad(self):
+        if self.freeze_vision_encoder:
+            self.get_vlm_model().vision_model.eval()
+            for params in self.get_vlm_model().vision_model.parameters():
+                params.requires_grad = False
         if self.train_expert_only:
             self.vlm.eval()
             for params in self.vlm.parameters():
@@ -166,26 +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
-
-        if self.freeze_vision_encoder:
-            self.get_vlm_model().vision_model.eval()
-            for params in self.get_vlm_model().vision_model.parameters():
-                params.requires_grad = False
-
-        if self.fine_tune_vision_encoder:
-            for params in self.get_vlm_model().vision_model.parameters():
-                params.requires_grad = True
-            for params in self.get_vlm_model().connector.parameters():
-                params.requires_grad = True
-
+                    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:
@@ -194,15 +193,11 @@ class SmolVLMWithExpertModel(nn.Module):
     def train(self, mode: bool = True):
         super().train(mode)
 
-        if self.train_expert_only:
-            self.vlm.eval()
-
         if self.freeze_vision_encoder:
             self.get_vlm_model().vision_model.eval()
 
-        if self.fine_tune_vision_encoder:
-            self.get_vlm_model().vision_model.train(mode)
-            self.get_vlm_model().connector.train(mode)
+        if self.train_expert_only:
+            self.vlm.eval()
 
     def embed_image(self, image: torch.Tensor):
         patch_attention_mask = None

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

@@ -18,7 +18,7 @@ import functools
 import threading
 from collections.abc import Callable, Sequence
 from contextlib import suppress
-from typing import TypedDict
+from typing import NotRequired, TypedDict
 
 import torch
 import torch.nn.functional as F  # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
     next_state: dict[str, torch.Tensor]
     done: torch.Tensor
     truncated: torch.Tensor
-    complementary_info: dict[str, torch.Tensor | float | int] | None = None
+    complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
 
 
 def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:

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
 

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