Merge branch 'main' into feat/add-multidataset-training

.github/workflows/stale.yml

@@ -31,11 +31,11 @@ env:
     Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
   WARN_ISSUE_MESSAGE: >
     This issue has been automatically marked as stale because it has not had
-    recent activity (6 months). It will be closed if no further activity occurs.
+    recent activity (1 year). It will be closed if no further activity occurs.
     Thank you for your contributions.
   WARN_PR_MESSAGE: >
     This PR has been automatically marked as stale because it has not had
-    recent activity (6 months). It will be closed if no further activity occurs.
+    recent activity (1 year). It will be closed if no further activity occurs.
     Thank you for your contributions.
 
 jobs:

README.md

@@ -202,7 +202,7 @@ Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/
 You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:
 
 ```bash
-lerobot-dataset-viz \
+python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --episode-index 0
 ```
@@ -210,7 +210,7 @@ lerobot-dataset-viz \
 or from a dataset in a local folder with the `root` option and the `--local-files-only` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
 
 ```bash
-lerobot-dataset-viz \
+python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --root ./my_local_data_dir \
     --local-files-only 1 \
@@ -221,7 +221,7 @@ It will open `rerun.io` and display the camera streams, robot states and actions
 
 https://github-production-user-asset-6210df.s3.amazonaws.com/4681518/328035972-fd46b787-b532-47e2-bb6f-fd536a55a7ed.mov?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240505%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240505T172924Z&X-Amz-Expires=300&X-Amz-Signature=d680b26c532eeaf80740f08af3320d22ad0b8a4e4da1bcc4f33142c15b509eda&X-Amz-SignedHeaders=host&actor_id=24889239&key_id=0&repo_id=748713144
 
-Our script can also visualize datasets stored on a distant server. See `lerobot-dataset-viz --help` for more instructions.
+Our script can also visualize datasets stored on a distant server. See `python -m lerobot.scripts.visualize_dataset --help` for more instructions.
 
 ### The `LeRobotDataset` format
 

benchmarks/video/run_video_benchmark.py

@@ -35,13 +35,12 @@ import torch
 from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
 from tqdm import tqdm
 
-from benchmarks.video.benchmark import TimeBenchmark
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.video_utils import (
     decode_video_frames_torchvision,
     encode_video_frames,
 )
-from lerobot.utils.constants import OBS_IMAGE
+from lerobot.utils.benchmark import TimeBenchmark
 
 BASE_ENCODING = OrderedDict(
     [
@@ -118,7 +117,7 @@ def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
     hf_dataset = dataset.hf_dataset.with_format(None)
 
     # We only save images from the first camera
-    img_keys = [key for key in hf_dataset.features if key.startswith(OBS_IMAGE)]
+    img_keys = [key for key in hf_dataset.features if key.startswith("observation.image")]
     imgs_dataset = hf_dataset.select_columns(img_keys[0])
 
     for i, item in enumerate(

docs/source/async.mdx

@@ -31,7 +31,7 @@ Then, spin up a policy server (in one terminal, or in a separate machine) specif
 You can spin up a policy server running:
 
 ```shell
-python src/lerobot/async_inference/policy_server.py \
+python src/lerobot/scripts/server/policy_server.py \
     --host=127.0.0.1 \
     --port=8080 \
 ```
@@ -39,7 +39,7 @@ python src/lerobot/async_inference/policy_server.py \
 This will start a policy server listening on `127.0.0.1:8080` (`localhost`, port 8080). At this stage, the policy server is empty, as all information related to which policy to run and with which parameters are specified during the first handshake with the client. Spin up a client with:
 
 ```shell
-python src/lerobot/async_inference/robot_client.py \
+python src/lerobot/scripts/server/robot_client.py \
     --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
     --robot.type=so100_follower \ # ROBOT: your robot type
     --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -122,8 +122,8 @@ python -m lerobot.scripts.server.policy_server \
 
 <!-- prettier-ignore-start -->
 ```python
-from lerobot.async_inference.configs import PolicyServerConfig
-from lerobot.async_inference.policy_server import serve
+from lerobot.scripts.server.configs import PolicyServerConfig
+from lerobot.scripts.server.policy_server import serve
 
 config = PolicyServerConfig(
     host="localhost",
@@ -148,7 +148,7 @@ The `RobotClient` streams observations to the `PolicyServer`, and receives actio
 <hfoptions id="start_robot_client">
 <hfoption id="Command">
 ```bash
-python src/lerobot/async_inference/robot_client.py \
+python src/lerobot/scripts/server/robot_client.py \
     --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
     --robot.type=so100_follower \ # ROBOT: your robot type
     --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -171,9 +171,9 @@ python src/lerobot/async_inference/robot_client.py \
 import threading
 from lerobot.robots.so100_follower import SO100FollowerConfig
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.async_inference.configs import RobotClientConfig
-from lerobot.async_inference.robot_client import RobotClient
-from lerobot.async_inference.helpers import visualize_action_queue_size
+from lerobot.scripts.server.configs import RobotClientConfig
+from lerobot.scripts.server.robot_client import RobotClient
+from lerobot.scripts.server.helpers import visualize_action_queue_size
 
 # 1. Create the robot instance
 """Check out the cameras available in your setup by running `python lerobot/find_cameras.py`"""

docs/source/hilserl.mdx

@@ -95,6 +95,7 @@ class HILSerlProcessorConfig:
 class ObservationConfig:
     add_joint_velocity_to_observation: bool = False    # Add joint velocities to state
     add_current_to_observation: bool = False    # Add motor currents to state
+    add_ee_pose_to_observation: bool = False    # Add end-effector pose to state
     display_cameras: bool = False    # Display camera feeds during execution
 
 class ImagePreprocessingConfig:
@@ -104,6 +105,7 @@ class ImagePreprocessingConfig:
 class GripperConfig:
     use_gripper: bool = True    # Enable gripper control
     gripper_penalty: float = 0.0    # Penalty for inappropriate gripper usage
+    gripper_penalty_in_reward: bool = False    # Include gripper penalty in reward
 
 class ResetConfig:
     fixed_reset_joint_positions: Any | None = None    # Joint positions for reset
@@ -286,6 +288,7 @@ You can enable multiple observation processing features simultaneously:
       "observation": {
         "add_joint_velocity_to_observation": true,
         "add_current_to_observation": true,
+        "add_ee_pose_to_observation": false,
         "display_cameras": false
       }
     }
@@ -301,19 +304,19 @@ Before collecting demonstrations, you need to determine the appropriate operatio
 
 This helps simplify the problem of learning on the real robot in two ways: 1) by limiting the robot's operational space to a specific region that solves the task and avoids unnecessary or unsafe exploration, and 2) by allowing training in end-effector space rather than joint space. Empirically, learning in joint space for reinforcement learning in manipulation is often a harder problem - some tasks are nearly impossible to learn in joint space but become learnable when the action space is transformed to end-effector coordinates.
 
-**Using lerobot-find-joint-limits**
+**Using find_joint_limits.py**
 
 This script helps you find the safe operational bounds for your robot's end-effector. Given that you have a follower and leader arm, you can use the script to find the bounds for the follower arm that will be applied during training.
 Bounding the action space will reduce the redundant exploration of the agent and guarantees safety.
 
 ```bash
-lerobot-find-joint-limits \
-  --robot.type=so100_follower \
-  --robot.port=/dev/tty.usbmodem58760431541 \
-  --robot.id=black \
-  --teleop.type=so100_leader \
-  --teleop.port=/dev/tty.usbmodem58760431551 \
-  --teleop.id=blue
+python -m lerobot.scripts.find_joint_limits \
+    --robot.type=so100_follower \
+    --robot.port=/dev/tty.usbmodem58760431541 \
+    --robot.id=black \
+    --teleop.type=so100_leader \
+    --teleop.port=/dev/tty.usbmodem58760431551 \
+    --teleop.id=blue
 ```
 
 **Workflow**

docs/source/il_robots.mdx

@@ -200,7 +200,7 @@ from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderCo
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop
 
 NUM_EPISODES = 5
@@ -237,7 +237,7 @@ dataset = LeRobotDataset.create(
 
 # Initialize the keyboard listener and rerun visualization
 _, events = init_keyboard_listener()
-init_rerun(session_name="recording")
+_init_rerun(session_name="recording")
 
 # Connect the robot and teleoperator
 robot.connect()
@@ -517,7 +517,7 @@ from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerCon
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop
 from lerobot.policies.factory import make_processor
 
@@ -557,7 +557,7 @@ dataset = LeRobotDataset.create(
 
 # Initialize the keyboard listener and rerun visualization
 _, events = init_keyboard_listener()
-init_rerun(session_name="recording")
+_init_rerun(session_name="recording")
 
 # Connect the robot
 robot.connect()

docs/source/koch.mdx

@@ -277,7 +277,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

docs/source/lekiwi.mdx

@@ -323,7 +323,7 @@ To replay an episode run the API example below, make sure to change `remote_ip`,
 python examples/lekiwi/replay.py
 ```
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by the training part of this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by the training part of this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 ## Evaluate your policy
 

docs/source/lerobot-dataset-v3.mdx

@@ -246,7 +246,7 @@ You can also use any `torchvision.transforms.v2` transform by passing it directl
 Use the visualization script to preview how transforms affect your data:
 
 ```bash
-lerobot-imgtransform-viz \
+python -m lerobot.scripts.visualize_image_transforms \
   --repo-id=your-username/your-dataset \
   --output-dir=./transform_examples \
   --n-examples=5

docs/source/libero.mdx

@@ -33,7 +33,7 @@ To Install LIBERO, after following LeRobot official instructions, just do:
 Evaluate a policy on one LIBERO suite:
 
 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
   --policy.path="your-policy-id" \
   --env.type=libero \
   --env.task=libero_object \
@@ -52,7 +52,7 @@ lerobot-eval \
 Benchmark a policy across multiple suites at once:
 
 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
   --policy.path="your-policy-id" \
   --env.type=libero \
   --env.task=libero_object,libero_spatial \
@@ -103,11 +103,10 @@ For reference, here is the **original dataset** published by Physical Intelligen
 ### Example training command
 
 ```bash
-lerobot-train \
+python src/lerobot/scripts/train.py \
   --policy.type=smolvla \
   --policy.repo_id=${HF_USER}/libero-test \
-  --policy.load_vlm_weights=true \
-  --dataset.repo_id=HuggingFaceVLA/libero \
+  --dataset.repo_id=jadechoghari/smol-libero3 \
   --env.type=libero \
   --env.task=libero_10 \
   --output_dir=./outputs/ \

docs/source/phone_teleop.mdx

@@ -136,12 +136,13 @@ Additionally you can customize mapping or safety limits by editing the processor
   ),
   ```
 
-- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` are the step limits for the EE pose and can be modified to change the safety limits.
+- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` and `max_ee_twist_step_rad` are the step limits for the EE pose and can be modified to change the safety limits.
 
   ```examples/phone_to_so100/teleoperate.py
   EEBoundsAndSafety(
       end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
       max_ee_step_m=0.10,
+      max_ee_twist_step_rad=0.50,
   )
   ```
 

docs/source/processors_robots_teleop.mdx

@@ -38,7 +38,7 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotActi
             kinematics=kinematics_solver, end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5}, motor_names=list(robot.bus.motors.keys()),
         ),
         EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20,
+            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20, max_ee_twist_step_rad=0.50,
         ),
         GripperVelocityToJoint(),
     ],

docs/source/smolvla.mdx

@@ -29,7 +29,7 @@ SmolVLA is Hugging Face’s lightweight foundation model for robotics. Designed
 ## Collect a dataset
 
 SmolVLA is a base model, so fine-tuning on your own data is required for optimal performance in your setup.
-We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](./il_robots)
+We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](https://huggingface.co/docs/lerobot/getting_started_real_world_robot#record-a-dataset)
 
 <Tip>
 
@@ -93,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](./getting_started_real_world_robot#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash

docs/source/so100.mdx

@@ -634,7 +634,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

docs/source/so101.mdx

@@ -430,7 +430,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>
 
-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)
 
 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).

examples/backward_compatibility/replay.py

@@ -44,7 +44,6 @@ from lerobot.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import (
     init_logging,
@@ -79,16 +78,16 @@ def replay(cfg: ReplayConfig):
 
     robot = make_robot_from_config(cfg.robot)
     dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode])
-    actions = dataset.hf_dataset.select_columns(ACTION)
+    actions = dataset.hf_dataset.select_columns("action")
     robot.connect()
 
     log_say("Replaying episode", cfg.play_sounds, blocking=True)
     for idx in range(dataset.num_frames):
         start_episode_t = time.perf_counter()
 
-        action_array = actions[idx][ACTION]
+        action_array = actions[idx]["action"]
         action = {}
-        for i, name in enumerate(dataset.features[ACTION]["names"]):
+        for i, name in enumerate(dataset.features["action"]["names"]):
             key = f"{name.removeprefix('main_')}.pos"
             action[key] = action_array[i].item()
 

examples/lekiwi/evaluate.py

@@ -19,12 +19,11 @@ from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
 from lerobot.policies.factory import make_pre_post_processors
 from lerobot.processor import make_default_processors
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.scripts.lerobot_record import record_loop
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 2
 FPS = 30
@@ -42,8 +41,8 @@ robot = LeKiwiClient(robot_config)
 policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
 
 # Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
 dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
@@ -74,7 +73,7 @@ teleop_action_processor, robot_action_processor, robot_observation_processor = m
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_evaluate")
+_init_rerun(session_name="lekiwi_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")

examples/lekiwi/record.py

@@ -17,15 +17,14 @@
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.processor import make_default_processors
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 2
 FPS = 30
@@ -48,8 +47,8 @@ keyboard = KeyboardTeleop(keyboard_config)
 teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
 
 # Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
 dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
@@ -70,7 +69,7 @@ keyboard.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_record")
+_init_rerun(session_name="lekiwi_record")
 
 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/lekiwi/replay.py

@@ -19,7 +19,6 @@ import time
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say
 
@@ -35,7 +34,7 @@ robot = LeKiwiClient(robot_config)
 dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
 episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = episode_frames.select_columns("action")
 
 # Connect to the robot
 robot.connect()
@@ -50,7 +49,7 @@ for idx in range(len(episode_frames)):
 
     # Get recorded action from dataset
     action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
     # Send action to robot

examples/lekiwi/teleoperate.py

@@ -20,7 +20,7 @@ from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
 from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 FPS = 30
 
@@ -41,7 +41,7 @@ leader_arm.connect()
 keyboard.connect()
 
 # Init rerun viewer
-init_rerun(session_name="lekiwi_teleop")
+_init_rerun(session_name="lekiwi_teleop")
 
 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/phone_to_so100/evaluate.py

@@ -34,16 +34,16 @@ from lerobot.processor.converters import (
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 5
 FPS = 30
@@ -137,7 +137,7 @@ robot.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_evaluate")
+_init_rerun(session_name="phone_so100_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")

examples/phone_to_so100/record.py

@@ -26,6 +26,7 @@ from lerobot.processor.converters import (
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     EEBoundsAndSafety,
@@ -35,13 +36,12 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 2
 FPS = 30
@@ -84,6 +84,7 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, Rob
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.20,
+            max_ee_twist_step_rad=0.50,
         ),
         GripperVelocityToJoint(speed_factor=20.0),
     ],
@@ -142,7 +143,7 @@ phone.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_record")
+_init_rerun(session_name="phone_so100_record")
 
 if not robot.is_connected or not phone.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/phone_to_so100/replay.py

@@ -28,7 +28,6 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say
 
@@ -67,7 +66,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
 dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
 episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = episode_frames.select_columns("action")
 
 # Connect to the robot
 robot.connect()
@@ -82,7 +81,7 @@ for idx in range(len(episode_frames)):
 
     # Get recorded action from dataset
     ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
     # Get robot observation

examples/phone_to_so100/teleoperate.py

@@ -33,7 +33,7 @@ from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
 from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 FPS = 30
 
@@ -67,6 +67,7 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.10,
+            max_ee_twist_step_rad=0.50,
         ),
         GripperVelocityToJoint(
             speed_factor=20.0,
@@ -86,7 +87,7 @@ robot.connect()
 teleop_device.connect()
 
 # Init rerun viewer
-init_rerun(session_name="phone_so100_teleop")
+_init_rerun(session_name="phone_so100_teleop")
 
 if not robot.is_connected or not teleop_device.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/so100_to_so100_EE/evaluate.py

@@ -34,16 +34,16 @@ from lerobot.processor.converters import (
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 5
 FPS = 30
@@ -138,7 +138,7 @@ robot.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="so100_so100_evaluate")
+_init_rerun(session_name="so100_so100_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")

examples/so100_to_so100_EE/record.py

@@ -27,6 +27,7 @@ from lerobot.processor.converters import (
     transition_to_observation,
     transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     EEBoundsAndSafety,
@@ -34,12 +35,11 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 
 NUM_EPISODES = 2
 FPS = 30
@@ -101,6 +101,7 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.10,
+            max_ee_twist_step_rad=0.50,
         ),
         InverseKinematicsEEToJoints(
             kinematics=follower_kinematics_solver,
@@ -142,7 +143,7 @@ follower.connect()
 
 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="recording_phone")
+_init_rerun(session_name="recording_phone")
 
 if not leader.is_connected or not follower.is_connected:
     raise ValueError("Robot or teleop is not connected!")

examples/so100_to_so100_EE/replay.py

@@ -29,7 +29,6 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say
 
@@ -68,7 +67,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
 dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
 episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = episode_frames.select_columns("action")
 
 # Connect to the robot
 robot.connect()
@@ -83,7 +82,7 @@ for idx in range(len(episode_frames)):
 
     # Get recorded action from dataset
     ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
     # Get robot observation

examples/so100_to_so100_EE/teleoperate.py

@@ -33,7 +33,7 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 FPS = 30
 
@@ -78,6 +78,7 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
         EEBoundsAndSafety(
             end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
             max_ee_step_m=0.10,
+            max_ee_twist_step_rad=0.50,
         ),
         InverseKinematicsEEToJoints(
             kinematics=follower_kinematics_solver,
@@ -94,7 +95,7 @@ follower.connect()
 leader.connect()
 
 # Init rerun viewer
-init_rerun(session_name="so100_so100_EE_teleop")
+_init_rerun(session_name="so100_so100_EE_teleop")
 
 print("Starting teleop loop...")
 while True:

examples/tester.py

@@ -0,0 +1,66 @@
+from lerobot.datasets.lerobot_dataset import MultiLeRobotDataset
+
+REPO_A = "lerobot/pusht"
+REPO_B = "lerobot/aloha_mobile_cabinet"  # replace with the actual repo id
+
+feature_keys_mapping = {
+    REPO_A: {  # pusht (1 camera, 2-dim)
+        "action": "actions",
+        "observation.state": "obs_state",
+        "observation.image": "obs_image.cam_high",
+    },
+    REPO_B: {  # dual arm (3 cameras, 14-dim)
+        "action": "actions",
+        "observation.state": "obs_state",
+        "observation.images.cam_high": "obs_image.cam_high",
+        "observation.images.cam_left_wrist": "obs_image.cam_left_wrist",
+        "observation.images.cam_right_wrist": "obs_image.cam_right_wrist",
+    },
+}
+
+from torchvision.transforms.v2 import Compose, ToImage, Resize
+image_tf = Compose([
+    ToImage(),          # converts to tensor if needed
+    Resize((224, 224)), # unify sizes across datasets (96x96 vs 480x640)
+])
+
+from torch.utils.data import DataLoader
+
+dataset = MultiLeRobotDataset(
+    repo_ids=[REPO_A, REPO_B],
+    image_transforms=image_tf,              # ensures same HxW
+    feature_keys_mapping=feature_keys_mapping,
+    train_on_all_features=True,             # keep union of cameras; zero-fill missing
+    # optional: override if you want fixed maxima; else inferred:
+    # max_action_dim=14,
+    # max_state_dim=14,
+    max_action_dim=14,
+    max_state_dim=14,
+    max_image_dim=224,
+    ignore_keys=[
+        "next.*", # drop reward/done/success
+        "index",
+        "timestamp",
+        "videos/*",               # drop all video metadata
+        "observation.effort",     # 👈 drop effort everywhere
+    ],
+)
+breakpoint()
+loader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=0, pin_memory=True)
+for _ in range(100):
+    batch = next(iter(loader))
+
+breakpoint()
+# vectors padded to maxima (pusht:2 -> 14; dual-arm:14 -> 14)
+assert batch["actions"].shape[-1] == 14
+assert batch["obs_state"].shape[-1] == 14
+assert batch["actions_padding_mask"].shape[-1] == 14
+assert batch["obs_state_padding_mask"].shape[-1] == 14
+
+# cameras: all canonical keys exist; pusht will have wrists zero-filled
+for cam in ["obs_image.cam_high", "obs_image.cam_left_wrist", "obs_image.cam_right_wrist"]:
+    assert cam in batch
+    assert f"{cam}_is_pad" in batch
+    # images should all be 3x224x224 (or your transform’s size)
+    img = batch[cam]
+    assert img.ndim in (4, 5)  # (B,C,H,W) or (B,T,C,H,W) depending on your loader

examples/tester.sh

@@ -0,0 +1,16 @@
+# storage / caches
+RAID=/raid/jade
+export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
+export HF_HOME=$RAID/.cache/huggingface
+export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
+export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
+export WANDB_CACHE_DIR=$RAID/.cache/wandb
+export TMPDIR=$RAID/.cache/tmp
+mkdir -p $TMPDIR
+export WANDB_MODE=offline
+# export HF_DATASETS_OFFLINE=1
+# export HF_HUB_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+export MUJOCO_GL=egl
+
+python examples/tester.py

examples/training/train_with_streaming.py

@@ -20,13 +20,13 @@ from pathlib import Path
 import torch
 
 from lerobot.configs.types import FeatureType
+from lerobot.constants import ACTION
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.act.modeling_act import ACTPolicy
 from lerobot.policies.factory import make_pre_post_processors
-from lerobot.utils.constants import ACTION
 
 
 def main():

pyproject.toml

@@ -162,19 +162,16 @@ all = [
 ]
 
 [project.scripts]
-lerobot-calibrate="lerobot.scripts.lerobot_calibrate:main"
-lerobot-find-cameras="lerobot.scripts.lerobot_find_cameras:main"
-lerobot-find-port="lerobot.scripts.lerobot_find_port:main"
-lerobot-record="lerobot.scripts.lerobot_record:main"
-lerobot-replay="lerobot.scripts.lerobot_replay:main"
-lerobot-setup-motors="lerobot.scripts.lerobot_setup_motors:main"
-lerobot-teleoperate="lerobot.scripts.lerobot_teleoperate:main"
-lerobot-eval="lerobot.scripts.lerobot_eval:main"
-lerobot-train="lerobot.scripts.lerobot_train:main"
-lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"
+lerobot-calibrate="lerobot.calibrate:main"
+lerobot-find-cameras="lerobot.find_cameras:main"
+lerobot-find-port="lerobot.find_port:main"
+lerobot-record="lerobot.record:main"
+lerobot-replay="lerobot.replay:main"
+lerobot-setup-motors="lerobot.setup_motors:main"
+lerobot-teleoperate="lerobot.teleoperate:main"
+lerobot-eval="lerobot.scripts.eval:main"
+lerobot-train="lerobot.scripts.train:main"
 lerobot-info="lerobot.scripts.lerobot_info:main"
-lerobot-find-joint-limits="lerobot.scripts.lerobot_find_joint_limits:main"
-lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
 
 # ---------------- Tool Configurations ----------------
 [tool.setuptools.packages.find]
@@ -201,7 +198,7 @@ exclude = ["tests/artifacts/**/*.safetensors", "*_pb2.py", "*_pb2_grpc.py"]
 # N: pep8-naming
 # TODO: Uncomment rules when ready to use
 select = [
-    "E", "W", "F", "I", "B", "C4", "T20", "N", "UP", "SIM" #, "A", "S", "D", "RUF"
+    "E", "W", "F", "I", "B", "C4", "T20", "N" # "SIM", "A", "S", "D", "RUF", "UP"
 ]
 ignore = [
     "E501", # Line too long
@@ -267,83 +264,8 @@ default.extend-ignore-identifiers-re = [
 # color = true
 # paths = ["src/lerobot"]
 
-# TODO: Enable mypy gradually module by module across multiple PRs
-# Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations
-
 # [tool.mypy]
 # python_version = "3.10"
 # warn_return_any = true
 # warn_unused_configs = true
 # ignore_missing_imports = false
-# strict = true
-# disallow_untyped_defs = true
-# disallow_incomplete_defs = true
-# check_untyped_defs = true
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.utils.*"
-# # include = "src/lerobot/utils/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.configs.*"
-# # include = "src/lerobot/configs/**/*.py"
-
-# # Data processing modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.processor.*"
-# # include = "src/lerobot/processor/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.datasets.*"
-# # include = "src/lerobot/datasets/**/*.py"
-
-# # Core machine learning modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.optim.*"
-# # include = "src/lerobot/optim/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.model.*"
-# # include = "src/lerobot/model/**/*.py"
-
-# # Hardware interfaces
-# [[tool.mypy.overrides]]
-# module = "lerobot.cameras.*"
-# # include = "src/lerobot/cameras/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.motors.*"
-# # include = "src/lerobot/motors/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.robots.*"
-# # include = "src/lerobot/robots/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.teleoperators.*"
-# # include = "src/lerobot/teleoperators/**/*.py"
-
-# # Complex modules (enable these last)
-# [[tool.mypy.overrides]]
-# module = "lerobot.policies.*"
-# # include = "src/lerobot/policies/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.rl.*"
-# # include = "src/lerobot/rl/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.envs.*"
-# # include = "src/lerobot/envs/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.async_inference.*"
-# # include = "src/lerobot/async_inference/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.transport.*"
-# # include = "src/lerobot/transport/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.scripts.*"
-# # include = "src/lerobot/scripts/**/*.py"

src/lerobot/cameras/opencv/camera_opencv.py

@@ -31,7 +31,7 @@ if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"
 import cv2
 import numpy as np
 
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..utils import get_cv2_backend, get_cv2_rotation

src/lerobot/cameras/reachy2_camera/reachy2_camera.py

@@ -31,7 +31,7 @@ import numpy as np
 from reachy2_sdk.media.camera import CameraView
 from reachy2_sdk.media.camera_manager import CameraManager
 
-from lerobot.utils.errors import DeviceNotConnectedError
+from lerobot.errors import DeviceNotConnectedError
 
 from ..camera import Camera
 from .configuration_reachy2_camera import ColorMode, Reachy2CameraConfig

src/lerobot/cameras/realsense/camera_realsense.py

@@ -29,7 +29,7 @@ try:
 except Exception as e:
     logging.info(f"Could not import realsense: {e}")
 
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..configs import ColorMode

src/lerobot/cameras/utils.py

@@ -15,10 +15,14 @@
 # limitations under the License.
 
 import platform
+from pathlib import Path
+from typing import TypeAlias
 
 from .camera import Camera
 from .configs import CameraConfig, Cv2Rotation
 
+IndexOrPath: TypeAlias = int | Path
+
 
 def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[str, Camera]:
     cameras = {}

src/lerobot/configs/default.py

@@ -16,6 +16,9 @@
 
 from dataclasses import dataclass, field
 
+from lerobot import (
+    policies,  # noqa: F401
+)
 from lerobot.datasets.transforms import ImageTransformsConfig
 from lerobot.datasets.video_utils import get_safe_default_codec
 

src/lerobot/configs/policies.py

@@ -27,9 +27,9 @@ from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import ACTION, OBS_STATE
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
-from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.hub import HubMixin
 from lerobot.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 

src/lerobot/configs/types.py

@@ -15,6 +15,7 @@
 # https://stackoverflow.com/questions/24481852/serialising-an-enum-member-to-json
 from dataclasses import dataclass
 from enum import Enum
+from typing import Any, Protocol
 
 
 class FeatureType(str, Enum):
@@ -37,6 +38,10 @@ class NormalizationMode(str, Enum):
     IDENTITY = "IDENTITY"
 
 
+class DictLike(Protocol):
+    def __getitem__(self, key: Any) -> Any: ...
+
+
 @dataclass
 class PolicyFeature:
     type: FeatureType

src/lerobot/constants.py

@@ -17,22 +17,21 @@ from pathlib import Path
 
 from huggingface_hub.constants import HF_HOME
 
-OBS_STR = "observation"
-OBS_PREFIX = OBS_STR + "."
-OBS_ENV_STATE = OBS_STR + ".environment_state"
-OBS_STATE = OBS_STR + ".state"
-OBS_IMAGE = OBS_STR + ".image"
-OBS_IMAGES = OBS_IMAGE + "s"
-OBS_LANGUAGE = OBS_STR + ".language"
-OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
-OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
-
+OBS_ENV_STATE = "observation.environment_state"
+OBS_STATE = "observation.state"
+OBS_IMAGE = "observation.image"
+OBS_IMAGES = "observation.images"
+OBS_LANGUAGE = "observation.language"
 ACTION = "action"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"
 
+OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
+OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
+
 ROBOTS = "robots"
+ROBOT_TYPE = "robot_type"
 TELEOPERATORS = "teleoperators"
 
 # files & directories

src/lerobot/datasets/aggregate.py

@@ -93,13 +93,14 @@ def update_data_df(df, src_meta, dst_meta):
         pd.DataFrame: Updated DataFrame with adjusted indices.
     """
 
-    df["episode_index"] = df["episode_index"] + dst_meta.info["total_episodes"]
-    df["index"] = df["index"] + dst_meta.info["total_frames"]
+    def _update(row):
+        row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
+        row["index"] = row["index"] + dst_meta.info["total_frames"]
+        task = src_meta.tasks.iloc[row["task_index"]].name
+        row["task_index"] = dst_meta.tasks.loc[task].task_index.item()
+        return row
 
-    src_task_names = src_meta.tasks.index.take(df["task_index"].to_numpy())
-    df["task_index"] = dst_meta.tasks.loc[src_task_names, "task_index"].to_numpy()
-
-    return df
+    return df.apply(_update, axis=1)
 
 
 def update_meta_data(
@@ -125,21 +126,27 @@ def update_meta_data(
         pd.DataFrame: Updated DataFrame with adjusted indices and timestamps.
     """
 
-    df["meta/episodes/chunk_index"] = df["meta/episodes/chunk_index"] + meta_idx["chunk"]
-    df["meta/episodes/file_index"] = df["meta/episodes/file_index"] + meta_idx["file"]
-    df["data/chunk_index"] = df["data/chunk_index"] + data_idx["chunk"]
-    df["data/file_index"] = df["data/file_index"] + data_idx["file"]
-    for key, video_idx in videos_idx.items():
-        df[f"videos/{key}/chunk_index"] = df[f"videos/{key}/chunk_index"] + video_idx["chunk"]
-        df[f"videos/{key}/file_index"] = df[f"videos/{key}/file_index"] + video_idx["file"]
-        df[f"videos/{key}/from_timestamp"] = df[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
-        df[f"videos/{key}/to_timestamp"] = df[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
+    def _update(row):
+        row["meta/episodes/chunk_index"] = row["meta/episodes/chunk_index"] + meta_idx["chunk"]
+        row["meta/episodes/file_index"] = row["meta/episodes/file_index"] + meta_idx["file"]
+        row["data/chunk_index"] = row["data/chunk_index"] + data_idx["chunk"]
+        row["data/file_index"] = row["data/file_index"] + data_idx["file"]
+        for key, video_idx in videos_idx.items():
+            row[f"videos/{key}/chunk_index"] = row[f"videos/{key}/chunk_index"] + video_idx["chunk"]
+            row[f"videos/{key}/file_index"] = row[f"videos/{key}/file_index"] + video_idx["file"]
+            row[f"videos/{key}/from_timestamp"] = (
+                row[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
+            )
+            row[f"videos/{key}/to_timestamp"] = (
+                row[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
+            )
 
-    df["dataset_from_index"] = df["dataset_from_index"] + dst_meta.info["total_frames"]
-    df["dataset_to_index"] = df["dataset_to_index"] + dst_meta.info["total_frames"]
-    df["episode_index"] = df["episode_index"] + dst_meta.info["total_episodes"]
+        row["dataset_from_index"] = row["dataset_from_index"] + dst_meta.info["total_frames"]
+        row["dataset_to_index"] = row["dataset_to_index"] + dst_meta.info["total_frames"]
+        row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
+        return row
 
-    return df
+    return df.apply(_update, axis=1)
 
 
 def aggregate_datasets(

src/lerobot/datasets/compute_stats.py

@@ -174,3 +174,79 @@ def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np
         aggregated_stats[key] = aggregate_feature_stats(stats_with_key)
 
     return aggregated_stats
+
+import numpy as np
+
+def aggregate_stats_multi(
+    stats_list: list[dict[str, dict]],
+    max_action_dim: int | None = None,
+    max_state_dim: int | None = None,
+) -> dict[str, dict[str, np.ndarray]]:
+    """Aggregate stats from multiple compute_stats outputs into a single set of stats.
+    
+    Supports heterogeneous robots by padding action/state stats to the max dim.
+    The final stats will have the union of all data keys from each of the stats dicts.
+
+    - new_min  = elementwise min across datasets
+    - new_max  = elementwise max across datasets
+    - new_mean = weighted mean (by count)
+    - new_std  = recomputed from total variance
+    """
+
+    data_keys = {key for stats in stats_list for key in stats}
+    aggregated_stats = {key: {} for key in data_keys}
+
+    def _pad(arr: np.ndarray, target: int) -> np.ndarray:
+        if arr.ndim == 0:  # scalar
+            return arr
+        if target is None or target <= 0 or arr.shape[-1] == target:
+            return arr
+        pad_width = [(0, 0)] * arr.ndim
+        pad_width[-1] = (0, target - arr.shape[-1])
+        return np.pad(arr, pad_width, mode="constant")
+
+    for key in data_keys:
+        stats_with_key = [stats[key] for stats in stats_list if key in stats]
+
+        # decide if this key should be padded
+        target_dim = None
+        if "action" in key and max_action_dim:
+            target_dim = max_action_dim
+        elif "state" in key and max_state_dim:
+            target_dim = max_state_dim
+
+        padded = []
+        counts = []
+        for s in stats_with_key:
+            mean = _pad(np.array(s["mean"]), target_dim)
+            std = _pad(np.array(s["std"]), target_dim)
+            min_ = _pad(np.array(s["min"]), target_dim)
+            max_ = _pad(np.array(s["max"]), target_dim)
+            count = s.get("count", 1)
+
+            padded.append(dict(mean=mean, std=std, min=min_, max=max_, count=count))
+            counts.append(count)
+
+        counts = np.array(counts, dtype=np.float64)
+        total_count = counts.sum()
+
+        means = np.stack([p["mean"] for p in padded])
+        stds  = np.stack([p["std"]  for p in padded])
+        mins  = np.stack([p["min"]  for p in padded])
+        maxs  = np.stack([p["max"]  for p in padded])
+
+        # weighted mean (broadcast weights properly)
+        new_mean = np.average(means, axis=0, weights=counts)
+        new_var  = np.average(stds**2 + (means - new_mean)**2, axis=0, weights=counts)
+
+        new_std = np.sqrt(new_var)
+
+        aggregated_stats[key] = {
+            "min": mins.min(axis=0),
+            "max": maxs.max(axis=0),
+            "mean": new_mean,
+            "std": new_std,
+            "count": int(total_count),
+        }
+
+    return aggregated_stats

src/lerobot/datasets/factory.py

@@ -27,7 +27,6 @@ from lerobot.datasets.lerobot_dataset import (
 )
 from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.transforms import ImageTransforms
-from lerobot.utils.constants import ACTION, OBS_PREFIX, REWARD
 
 IMAGENET_STATS = {
     "mean": [[[0.485]], [[0.456]], [[0.406]]],  # (c,1,1)
@@ -55,11 +54,11 @@ def resolve_delta_timestamps(
     """
     delta_timestamps = {}
     for key in ds_meta.features:
-        if key == REWARD and cfg.reward_delta_indices is not None:
+        if key == "next.reward" and cfg.reward_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.reward_delta_indices]
-        if key == ACTION and cfg.action_delta_indices is not None:
+        if key == "action" and cfg.action_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices]
-        if key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None:
+        if key.startswith("observation.") and cfg.observation_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.observation_delta_indices]
 
     if len(delta_timestamps) == 0:

src/lerobot/datasets/lerobot_dataset.py

@@ -31,6 +31,8 @@ import torch.utils
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.errors import RevisionNotFoundError
 
+from collections import defaultdict
+from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats
 from lerobot.datasets.image_writer import AsyncImageWriter, write_image
 from lerobot.datasets.utils import (
@@ -78,10 +80,14 @@ from lerobot.datasets.video_utils import (
     get_video_duration_in_s,
     get_video_info,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME
 
 CODEBASE_VERSION = "v3.0"
-
+OBS_IMAGE = "observation.image"
+OBS_IMAGE_2 = "observation.image_2"
+OBS_IMAGE_3 = "observation.image_3"
+OBS_STATE = "observation.state"
+OBS_ENV_STATE = "observation.env_state"
+ACTION = "action"
 
 class LeRobotDatasetMetadata:
     def __init__(
@@ -848,6 +854,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         return item
 
+    def _add_padding_keys(self, item: dict, padding: dict[str, list[bool]]) -> dict:
+        for key, val in padding.items():
+            item[key] = torch.BoolTensor(val)
+        return item
+
     def __len__(self):
         return self.num_frames
 
@@ -1027,7 +1038,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             # Reset episode buffer and clean up temporary images (if not already deleted during video encoding)
             self.clear_episode_buffer(delete_images=len(self.meta.image_keys) > 0)
 
-    def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
+    def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None):
         """
         Batch save videos for multiple episodes.
 
@@ -1153,7 +1164,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         }
         return metadata
 
-    def _save_episode_video(self, video_key: str, episode_index: int) -> dict:
+    def _save_episode_video(self, video_key: str, episode_index: int):
         # Encode episode frames into a temporary video
         ep_path = self._encode_temporary_episode_video(video_key, episode_index)
         ep_size_in_mb = get_video_size_in_mb(ep_path)
@@ -1258,7 +1269,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         if self.image_writer is not None:
             self.image_writer.wait_until_done()
 
-    def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> Path:
+    def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> dict:
         """
         Use ffmpeg to convert frames stored as png into mp4 videos.
         Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
@@ -1317,13 +1328,139 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         return obj
 
+ROBOT_TYPE_KEYS_MAPPING = {
+    "lerobot/stanford_hydra_dataset": "static_single_arm",
+    "lerobot/iamlab_cmu_pickup_insert": "static_single_arm",
+    "lerobot/berkeley_fanuc_manipulation": "static_single_arm",
+    "lerobot/toto": "static_single_arm",
+    "lerobot/roboturk": "static_single_arm",
+    "lerobot/jaco_play": "static_single_arm",
+    "lerobot/taco_play": "static_single_arm_7statedim",
+}
+class MultiLeRobotDatasetMeta:
+    def __init__(
+        self,
+        datasets: list[LeRobotDataset],
+        repo_ids: list[str],
+        keys_to_max_dim: dict[str, int],
+        train_on_all_features: bool = False,
+    ):
+        self.repo_ids = repo_ids
+        self.keys_to_max_dim = keys_to_max_dim
+        self.train_on_all_features = train_on_all_features
+        self.robot_types = [ds.meta.info["robot_type"] for ds in datasets]
 
+        # assign robot_type if missing
+        for ds in datasets:
+            ds.meta.info["robot_type"] = ROBOT_TYPE_KEYS_MAPPING.get(ds.repo_id, ds.meta.info["robot_type"])
+            ds.robot_type = ds.meta.info["robot_type"]
+
+        # step 1: compute disabled features
+        self.disabled_features = set()
+        if not self.train_on_all_features:
+            intersection = set(datasets[0].features)
+            for ds in datasets:
+                intersection.intersection_update(ds.features)
+            if not intersection:
+                raise RuntimeError("No common features across datasets.")
+            for repo_id, ds in zip(repo_ids, datasets, strict=False):
+                extra = set(ds.features) - intersection
+                logging.warning(f"Disabling {extra} for repo {repo_id}")
+                self.disabled_features.update(extra)
+
+        # step 2: build union_features excluding disabled
+        self.union_features = {}
+        for ds in datasets:
+            for k, v in ds.features.items():
+                if k not in self.disabled_features:
+                    self.union_features[k] = v
+
+        # step 3: reshape feature schema
+        self.features = reshape_features_to_max_dim(
+            self.union_features, reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim
+        )
+
+        # step 4: aggregate stats
+        self.stats = aggregate_stats_per_robot_type(datasets)
+        for robot_type_, stats_ in self.stats.items():
+            for feat_key, feat_stats in stats_.items():
+                if feat_key in [ACTION, OBS_ENV_STATE, OBS_STATE]:
+                    for k, v in feat_stats.items():
+                        pad_value = 0 if k in ["min", "mean"] else 1
+                        self.stats[robot_type_][feat_key][k] = pad_tensor(
+                            v,
+                            max_size=self.keys_to_max_dim.get(feat_key, -1),
+                            pad_dim=-1,
+                            pad_value=pad_value,
+                        )
+
+        # step 5: episodes & tasks
+        self.episodes = {repo_id: ds.meta.episodes for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+        self.tasks = {repo_id: ds.meta.tasks for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+        self.info = {repo_id: ds.meta.info for repo_id, ds in zip(repo_ids, datasets, strict=False)}
+
+
+class MultiLeRobotDatasetCleaner:
+    def __init__(
+        self,
+        datasets: list[LeRobotDataset],
+        repo_ids: list[str],
+        sampling_weights: list[float],
+        datasets_repo_ids: list[str],
+        min_fps: int = 1,
+        max_fps: int = 100,
+    ):
+        self.original_datasets = datasets
+        self.original_repo_ids = repo_ids
+        self.original_weights = sampling_weights
+        self.original_datasets_repo_ids = datasets_repo_ids
+
+        # step 1: remove datasets with invalid fps
+
+        # step 2: keep datasets with same features per robot type
+        consistent_datasets, keep_mask = keep_datasets_with_the_same_features_per_robot_type(
+            datasets
+        )
+
+        self.cleaned_datasets = consistent_datasets
+        self.keep_mask = keep_mask
+        self.cleaned_weights = [sampling_weights[i] for i in range(len(datasets)) if keep_mask[i]]
+        self.cleaned_repo_ids = [repo_ids[i] for i in range(len(datasets)) if keep_mask[i]]
+        self.cleaned_datasets_repo_ids = [
+            datasets_repo_ids[i] for i in range(len(datasets)) if keep_mask[i]
+        ]
+
+        self.cumulative_sizes = np.array(
+            [0] + list(torch.cumsum(torch.tensor([len(d) for d in consistent_datasets]), dim=0))
+        )
+        self.cleaned_weights = np.array(self.cleaned_weights, dtype=np.float32)
+
+# --- at the top of the file (same imports as before) ---
+from collections import defaultdict
+from typing import Callable
+import copy
+import numpy as np
+import torch
+import datasets
+from pathlib import Path
+
+# If you already have these in your codebase, reuse them
+try:
+    from lerobot.common.constants import (
+        ACTION, OBS_ENV_STATE, OBS_STATE, OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3
+    )
+except Exception:
+    # Fallbacks if constants are already strings elsewhere
+    ACTION = "action"
+    OBS_ENV_STATE = "observation.env_state"
+    OBS_STATE = "observation.state"
+    OBS_IMAGE = "observation.image"
+    OBS_IMAGE_2 = "observation.image_2"
+    OBS_IMAGE_3 = "observation.image_3"
+
+IGNORED_KEYS = ["observation.effort"]
 class MultiLeRobotDataset(torch.utils.data.Dataset):
-    """A dataset consisting of multiple underlying `LeRobotDataset`s.
-
-    The underlying `LeRobotDataset`s are effectively concatenated, and this class adopts much of the API
-    structure of `LeRobotDataset`.
-    """
+    # ... keep your existing docstring ...
 
     def __init__(
         self,
@@ -1331,108 +1468,261 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         root: str | Path | None = None,
         episodes: dict | None = None,
         image_transforms: Callable | None = None,
-        delta_timestamps: dict[str, list[float]] | None = None,
+        delta_timestamps: dict[list[float]] | None = None,
         tolerances_s: dict | None = None,
         download_videos: bool = True,
         video_backend: str | None = None,
+        # --- NEW: simple add-ons ---
+        sampling_weights: list[float] | None = None,
+        feature_keys_mapping: dict[str, dict[str, str]] | None = None,
+        max_action_dim: int | None = None,
+        max_state_dim: int | None = None,
+        max_num_images: int | None = None,
+        max_image_dim: int | None = None,
+        train_on_all_features: bool = False,
+        min_fps: int = 1,
+        max_fps: int = 100,
+        ignore_keys: list[str] | None = None,  # exact or glob patterns
     ):
         super().__init__()
         self.repo_ids = repo_ids
         self.root = Path(root) if root else HF_LEROBOT_HOME
         self.tolerances_s = tolerances_s if tolerances_s else dict.fromkeys(repo_ids, 0.0001)
-        # Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which
-        # are handled by this class.
-        self._datasets = [
-            LeRobotDataset(
-                repo_id,
-                root=self.root / repo_id,
-                episodes=episodes[repo_id] if episodes else None,
-                image_transforms=image_transforms,
-                delta_timestamps=delta_timestamps,
-                tolerance_s=self.tolerances_s[repo_id],
-                download_videos=download_videos,
-                video_backend=video_backend,
-            )
-            for repo_id in repo_ids
-        ]
 
-        # Disable any data keys that are not common across all of the datasets. Note: we may relax this
-        # restriction in future iterations of this class. For now, this is necessary at least for being able
-        # to use PyTorch's default DataLoader collate function.
-        self.disabled_features = set()
-        intersection_features = set(self._datasets[0].features)
-        for ds in self._datasets:
-            intersection_features.intersection_update(ds.features)
-        if len(intersection_features) == 0:
-            raise RuntimeError(
-                "Multiple datasets were provided but they had no keys common to all of them. "
-                "The multi-dataset functionality currently only keeps common keys."
-            )
-        for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
-            extra_keys = set(ds.features).difference(intersection_features)
-            logging.warning(
-                f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
-                "other datasets."
-            )
-            self.disabled_features.update(extra_keys)
+        # --- NEW: store mapping and simple knobs ---
+        self.feature_keys_mapping: dict[str, dict[str, str]] = feature_keys_mapping or {}
+        self.train_on_all_features = train_on_all_features
+        self.max_action_dim = max_action_dim
+        self.max_state_dim = max_state_dim
+        self.max_image_dim = max_image_dim
+        self.max_num_images = max_num_images  # (optional, we don’t enforce count, we enforce names)
+        self._ignore_patterns = list(ignore_keys or [])
+        # Build underlying single datasets
+        _datasets = []
+        datasets_repo_ids = []
+        self.sampling_weights = []
 
+        sampling_weights = sampling_weights if sampling_weights is not None else [1] * len(repo_ids)
+        assert len(sampling_weights) == len(repo_ids), (
+            "The number of sampling weights must match the number of datasets. "
+            f"Got {len(sampling_weights)} weights for {len(repo_ids)} datasets."
+        )
+        for i, repo_id in enumerate(repo_ids):
+            try:
+                _datasets.append(
+                    LeRobotDataset(
+                        repo_id,
+                        root=self.root / repo_id,
+                        episodes=episodes.get(repo_id, None) if episodes else None,
+                        image_transforms=image_transforms,  # transforms applied inside single ds
+                        delta_timestamps=delta_timestamps.get(repo_id, None) if delta_timestamps else None,
+                        tolerance_s=self.tolerances_s[repo_id],
+                        download_videos=download_videos,
+                        video_backend=video_backend,
+                    )
+                )
+                datasets_repo_ids.append(repo_id)
+                self.sampling_weights.append(float(sampling_weights[i]))
+            except Exception as e:
+                print(f"Failed to load dataset: {repo_id} due to Exception: {e}")
+
+        print(
+            f"Finish loading {len(_datasets)} datasets, with sampling weights: "
+            f"{self.sampling_weights} corresponding to: {datasets_repo_ids}"
+        )
+
+        # Bookkeeping for mapping & canonical image inventory
         self.image_transforms = image_transforms
-        self.delta_timestamps = delta_timestamps
-        # TODO(rcadene, aliberts): We should not perform this aggregation for datasets
-        # with multiple robots of different ranges. Instead we should have one normalization
-        # per robot.
-        self.stats = aggregate_stats([dataset.meta.stats for dataset in self._datasets])
+        self.delta_timestamps = delta_timestamps.get(repo_id, None) if delta_timestamps else None
+        self._datasets = _datasets
+        self.datasets_repo_ids = datasets_repo_ids
+
+        # --- NEW: compute “canonical image keys” (targets across all mappings) ---
+        self._canonical_image_keys: set[str] = set()
+        self._source_keys_per_repo: dict[str, set[str]] = {}
+        self._target_keys_per_repo: dict[str, set[str]] = {}
+        for rid, mapping in self.feature_keys_mapping.items():
+            src_keys = set(mapping.keys())
+            tgt_keys = set(mapping.values())
+            self._source_keys_per_repo[rid] = src_keys
+            self._target_keys_per_repo[rid] = tgt_keys
+            # union of target names (we will ensure these exist at __getitem__)
+            self._canonical_image_keys |= {
+                k for k in tgt_keys if self._is_image_key_like(k)
+            }
+
+        # If user didn’t give any mapping, fall back to native keys (no-ops)
+        if not self._canonical_image_keys and self.train_on_all_features:
+            # discover all image-like keys from raw features
+            for ds in self._datasets:
+                for k, v in ds.hf_features.items():
+                    if isinstance(v, (datasets.Image, VideoFrame)):
+                        self._canonical_image_keys.add(k)
+
+        # Cleaner: keep fps & consistent feature sets per robot type (unchanged)
+        cleaner = MultiLeRobotDatasetCleaner(
+            datasets=self._datasets,
+            repo_ids=repo_ids,
+            sampling_weights=self.sampling_weights,
+            datasets_repo_ids=self.datasets_repo_ids,
+            min_fps=min_fps,
+            max_fps=max_fps,
+        )
+        self._datasets = cleaner.cleaned_datasets
+        self.sampling_weights = cleaner.cleaned_weights
+        self.repo_ids = cleaner.cleaned_repo_ids
+        self.datasets_repo_ids = cleaner.cleaned_datasets_repo_ids
+        self.cumulative_sizes = cleaner.cumulative_sizes
+
+        # Meta (unchanged): we give it dim maxima; it will reshape/pad vectors
+        self.meta = MultiLeRobotDatasetMeta(
+            datasets=self._datasets,
+            repo_ids=self.repo_ids,
+            keys_to_max_dim={
+                ACTION: self.max_action_dim if self.max_action_dim is not None else -1,
+                OBS_ENV_STATE: self.max_state_dim if self.max_state_dim is not None else -1,
+                OBS_STATE: self.max_state_dim if self.max_state_dim is not None else -1,
+                OBS_IMAGE: self.max_image_dim if self.max_image_dim is not None else -1,
+                OBS_IMAGE_2: self.max_image_dim if self.max_image_dim is not None else -1,
+                OBS_IMAGE_3: self.max_image_dim if self.max_image_dim is not None else -1,
+            },
+            train_on_all_features=train_on_all_features,
+        )
+
+        # --- NEW: track dropped (source) keys so collate won’t expect them
+        # Anything that we *rename away* should be considered disabled,
+        # otherwise downstream may expect them to exist.
+        self._dropped_keys = set()
+        for rid, mapping in self.feature_keys_mapping.items():
+            self._dropped_keys |= set(mapping.keys())
+
+        # Merge with meta’s disabled features
+        self.disabled_features = set(self.meta.disabled_features) | self._dropped_keys
+
+        self.stats = self.meta.stats
+
+        # --- NEW: cache an example image shape per canonical key (lazy, filled on first use)
+        self._cached_img_shape: dict[str, torch.Size] = {}
+
+    # ---------------------- NEW small helpers ----------------------
+
+    def _is_image_key_like(self, key: str) -> bool:
+        # A loose heuristic: rely on name OR on features later
+        return ("image" in key) or ("cam_" in key) or ("images." in key)
+    
+    def _should_ignore(self, key: str) -> bool:
+        # exact or glob-style match
+        for pat in self._ignore_patterns:
+            if key == pat or fnmatch.fnmatch(key, pat):
+                return True
+        return False
+    def _apply_feature_mapping(self, item: dict, repo_id: str) -> dict:
+        """
+        Rename features according to feature_keys_mapping[repo_id].
+        - Moves tensor/image under target key.
+        - Drops source key if moved.
+        - Adds *_is_pad=False for image targets we fill/keep.
+        """
+        mapping = self.feature_keys_mapping.get(repo_id, {}) or {}
+        if not mapping:
+            return item
+
+        for src, tgt in mapping.items():
+            if src in item:
+                # Move value
+                item[tgt] = item[src]
+                # Drop the source to avoid duplication
+                del item[src]
+        return item
+
+    def _ensure_union_image_keys(self, item: dict) -> dict:
+        """
+        Ensure that every canonical image key exists.
+        When missing, create a zero tensor matching (B,C,H,W) or (C,H,W) of an available image.
+        Also add boolean mask at f"{key}_is_pad".
+        """
+        if not self.train_on_all_features or not self._canonical_image_keys:
+            return item
+
+        # find any existing image tensor in item to copy shape/dtype
+        exemplar = None
+        for k in list(item.keys()):
+            v = item[k]
+            if torch.is_tensor(v) and v.ndim in (3, 4, 5):  # (C,H,W) or (B,C,H,W) or (B,T,C,H,W)
+                exemplar = v
+                break
+
+        # fallback to a safe 3x224x224 if nothing found
+        def _fallback_image():
+            return torch.zeros(3, 224, 224, dtype=torch.uint8)
+
+        for key in self._canonical_image_keys:
+            if key not in item:
+                img = torch.zeros_like(exemplar) if exemplar is not None else _fallback_image()
+                item[key] = img
+                item[f"{key}_is_pad"] = torch.tensor(True, dtype=torch.bool)
+            else:
+                # Add a mask saying it’s *not* padded
+                if f"{key}_is_pad" not in item:
+                    item[f"{key}_is_pad"] = torch.tensor(False, dtype=torch.bool)
+        return item
+
+    # ---------------------- existing API below (mostly unchanged) ----------------------
 
     @property
     def repo_id_to_index(self):
-        """Return a mapping from dataset repo_id to a dataset index automatically created by this class.
-
-        This index is incorporated as a data key in the dictionary returned by `__getitem__`.
-        """
         return {repo_id: i for i, repo_id in enumerate(self.repo_ids)}
 
     @property
-    def fps(self) -> int:
-        """Frames per second used during data collection.
+    def repo_index_to_id(self):
+        return {v: k for k, v in self.repo_id_to_index}
 
-        NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
-        """
+    @property
+    def fps(self) -> int:
         return self._datasets[0].meta.info["fps"]
 
     @property
     def video(self) -> bool:
-        """Returns True if this dataset loads video frames from mp4 files.
-
-        Returns False if it only loads images from png files.
-
-        NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
-        """
         return self._datasets[0].meta.info.get("video", False)
 
     @property
     def features(self) -> datasets.Features:
-        features = {}
+        """
+        Extend native HF features with any *target* keys introduced by mapping.
+        We copy the source spec for targets that didn’t exist in any raw dataset.
+        """
+        features: dict[str, datasets.features.Feature] = {}
         for dataset in self._datasets:
-            features.update({k: v for k, v in dataset.hf_features.items() if k not in self.disabled_features})
+            for k, v in dataset.hf_features.items():
+                if k not in self.disabled_features:
+                    features[k] = v
+
+        # Add mapped target image specs if not present yet
+        for rid, mapping in self.feature_keys_mapping.items():
+            ds = None
+            # find the dataset object to read feature spec for source
+            for _ds, _rid in zip(self._datasets, self.repo_ids, strict=False):
+                if _rid == rid:
+                    ds = _ds
+                    break
+            if ds is None:
+                continue
+            for src, tgt in mapping.items():
+                if tgt not in features and src in ds.hf_features:
+                    features[tgt] = ds.hf_features[src]
+
         return features
 
     @property
     def camera_keys(self) -> list[str]:
-        """Keys to access image and video stream from cameras."""
         keys = []
         for key, feats in self.features.items():
-            if isinstance(feats, (datasets.Image | VideoFrame)):
+            if isinstance(feats, (datasets.Image, VideoFrame)):
                 keys.append(key)
         return keys
 
     @property
     def video_frame_keys(self) -> list[str]:
-        """Keys to access video frames that requires to be decoded into images.
-
-        Note: It is empty if the dataset contains images only,
-        or equal to `self.cameras` if the dataset contains videos only,
-        or can even be a subset of `self.cameras` in a case of a mixed image/video dataset.
-        """
         video_frame_keys = []
         for key, feats in self.features.items():
             if isinstance(feats, VideoFrame):
@@ -1441,21 +1731,14 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
 
     @property
     def num_frames(self) -> int:
-        """Number of samples/frames."""
         return sum(d.num_frames for d in self._datasets)
 
     @property
     def num_episodes(self) -> int:
-        """Number of episodes."""
         return sum(d.num_episodes for d in self._datasets)
 
     @property
     def tolerance_s(self) -> float:
-        """Tolerance in seconds used to discard loaded frames when their timestamps
-        are not close enough from the requested frames. It is only used when `delta_timestamps`
-        is provided or when loading video frames from mp4 files.
-        """
-        # 1e-4 to account for possible numerical error
         return 1 / self.fps - 1e-4
 
     def __len__(self):
@@ -1464,22 +1747,83 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
     def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
         if idx >= len(self):
             raise IndexError(f"Index {idx} out of bounds.")
-        # Determine which dataset to get an item from based on the index.
-        start_idx = 0
-        dataset_idx = 0
-        for dataset in self._datasets:
-            if idx >= start_idx + dataset.num_frames:
-                start_idx += dataset.num_frames
-                dataset_idx += 1
-                continue
-            break
-        else:
-            raise AssertionError("We expect the loop to break out as long as the index is within bounds.")
-        item = self._datasets[dataset_idx][idx - start_idx]
+        dataset_idx = np.searchsorted(self.cumulative_sizes, idx, side="right").item() - 1
+        local_idx = (idx - self.cumulative_sizes[dataset_idx]).item()
+        item = self._datasets[dataset_idx][local_idx]
+
+        # Identify which repo this sample came from
+        repo_id = self.datasets_repo_ids[dataset_idx]
+
+        # --- NEW: apply mapping and ensure union of image keys ---
+        item = self._apply_feature_mapping(item, repo_id)
+        item = self._ensure_union_image_keys(item)
+
+        # annotate dataset index for downstream
         item["dataset_index"] = torch.tensor(dataset_idx)
+
+        # Pad vector features to max dims using meta (unchanged)
+        item = create_padded_features(item, self.meta.features)
+
+        # Drop any disabled (including original source keys we remapped away)
         for data_key in self.disabled_features:
             if data_key in item:
                 del item[data_key]
+        for k in IGNORED_KEYS:
+            if k in item:
+                item.pop(k)
+        # Convert any datasets.Image still present to tensor
+        if self.image_transforms is not None:
+            for cam in [k for k in item.keys() if self._is_image_key_like(k)]:
+                val = item[cam]
+                if not torch.is_tensor(val):
+                    item[cam] = self.image_transforms(val)
+        # 🔑 Pad actions if too short
+        if "actions" in item and self.max_action_dim is not None:
+            act = item["actions"]
+            if act.shape[-1] < self.max_action_dim:
+                pad_len = self.max_action_dim - act.shape[-1]
+                item["actions"] = torch.cat([act, torch.zeros(pad_len, dtype=act.dtype)], dim=-1)
+                item["actions_padding_mask"] = torch.cat(
+                    [torch.zeros_like(act, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+
+        # pad obs_state if too short
+        if "obs_state" in item and self.max_state_dim is not None:
+            st = item["obs_state"]
+            if st.shape[-1] < self.max_state_dim:
+                pad_len = self.max_state_dim - st.shape[-1]
+                item["obs_state"] = torch.cat([st, torch.zeros(pad_len, dtype=st.dtype)], dim=-1)
+                item["obs_state_padding_mask"] = torch.cat(
+                    [torch.zeros_like(st, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+        # actions
+        if "actions" in item and self.max_action_dim is not None:
+            act = item["actions"]
+            if act.shape[-1] < self.max_action_dim:
+                pad_len = self.max_action_dim - act.shape[-1]
+                item["actions"] = torch.cat([act, torch.zeros(pad_len, dtype=act.dtype)], dim=-1)
+                mask = torch.cat(
+                    [torch.zeros_like(act, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+            else:
+                mask = torch.zeros(self.max_action_dim, dtype=torch.bool)  # 👈 all False if no padding
+            item["actions_padding_mask"] = mask
+        # obs state
+        if "obs_state" in item and self.max_state_dim is not None:
+            st = item["obs_state"]
+            if st.shape[-1] < self.max_state_dim:
+                pad_len = self.max_state_dim - st.shape[-1]
+                item["obs_state"] = torch.cat([st, torch.zeros(pad_len, dtype=st.dtype)], dim=-1)
+                mask = torch.cat(
+                    [torch.zeros_like(st, dtype=torch.bool), torch.ones(pad_len, dtype=torch.bool)],
+                    dim=-1,
+                )
+            else:
+                mask = torch.zeros(self.max_state_dim, dtype=torch.bool)  # 👈 always add mask
+            item["obs_state_padding_mask"] = mask
 
         return item
 
@@ -1496,3 +1840,149 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
             f"  Transformations: {self.image_transforms},\n"
             f")"
         )
+
+def keep_datasets_with_the_same_features_per_robot_type(ls_datasets: list) -> list:
+    """
+    Filters datasets to only keep those with consistent feature shapes per robot type.
+
+    Args:
+        ls_datasets (List): List of datasets, each with a `meta.info['robot_type']`
+            and `meta.episodes_stats` dictionary.
+
+    Returns:
+        List: Filtered list of datasets with consistent feature shapes.
+    """
+    robot_types = {ds.meta.info["robot_type"] for ds in ls_datasets}
+    datasets_to_remove = set()
+
+    for robot_type in robot_types:
+        # Collect all stats dicts for this robot type
+        stats_list = [
+            ep_stats
+            for ds in ls_datasets
+            if ds.meta.info["robot_type"] == robot_type
+            for ep_stats in episode_stats_values(ds.meta)
+        ]
+        if not stats_list:
+            continue
+
+        # Determine the most common shape for each key
+        all_keys = {key for stats in stats_list for key in stats}
+        for ds in ls_datasets:
+            if ds.meta.info["robot_type"] != robot_type:
+                continue
+            for key in all_keys:
+                shape_counter = defaultdict(int)
+
+                for stats in stats_list:
+                    value = stats.get(key)
+                    if (
+                        value and "mean" in value and isinstance(value["mean"], (torch.Tensor, np.ndarray))
+                    ):  # FIXME(mshukor): check all stats; min, mean, max
+                        shape_counter[value["mean"].shape] += 1
+                if not shape_counter:
+                    continue
+
+                # Identify the most frequent shape
+                main_shape = max(shape_counter, key=shape_counter.get)
+                # Flag datasets that don't match the main shape
+                # for ds in ls_datasets:
+                first_ep_stats = next(iter(episode_stats_values(ds.meta)), None)
+                if not first_ep_stats:
+                    continue
+                value = first_ep_stats.get(key)
+                if (
+                    value
+                    and "mean" in value
+                    and isinstance(value["mean"], (torch.Tensor, np.ndarray))
+                    and value["mean"].shape != main_shape
+                ):
+                    datasets_to_remove.add(ds)
+                    break
+
+    # Filter out inconsistent datasets
+    datasets_maks = [ds not in datasets_to_remove for ds in ls_datasets]
+    filtered_datasets = [ds for ds in ls_datasets if ds not in datasets_to_remove]
+    print(
+        f"Keeping {len(filtered_datasets)} datasets. Removed {len(datasets_to_remove)} inconsistent ones. Inconsistent datasets:\n{datasets_to_remove}"
+    )
+    return filtered_datasets, datasets_maks
+
+
+def aggregate_stats_per_robot_type(ls_datasets) -> dict[str, dict[str, torch.Tensor]]:
+    """Aggregate stats of multiple LeRobot datasets into multiple set of stats per robot type.
+
+    The final stats will have the union of all data keys from each of the datasets.
+
+    The final stats will have the union of all data keys from each of the datasets. For instance:
+    - new_max = max(max_dataset_0, max_dataset_1, ...)
+    - new_min = min(min_dataset_0, min_dataset_1, ...)
+    - new_mean = (mean of all data)
+    - new_std = (std of all data)
+    """
+
+    robot_types = {ds.meta.info["robot_type"] for ds in ls_datasets}
+    stats = {robot_type: {} for robot_type in robot_types}
+    for robot_type in robot_types:
+        robot_type_datasets = []
+        for ds in ls_datasets:
+            if ds.meta.info["robot_type"] == robot_type:
+                robot_type_datasets.extend(list(episode_stats_values(ds.meta)))
+        # robot_type_datasets = [list(ds.episodes_stats.values()) for ds in ls_datasets if ds.meta.info["robot_type"] == robot_type]
+        stat = aggregate_stats(robot_type_datasets)
+        stats[robot_type] = stat
+    return stats
+
+def reshape_features_to_max_dim(features: dict, reshape_dim: int = -1, keys_to_max_dim: dict = {}) -> dict:
+    """Reshape features to have a maximum dimension of `max_dim`."""
+    reshaped_features = {}
+    for key in features:
+        if key in keys_to_max_dim and keys_to_max_dim[key] is not None:
+            reshaped_features[key] = features[key]
+            shape = list(features[key]["shape"])
+            if any([k in key for k in [OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3]]):  # Assume square images
+                shape[-3] = keys_to_max_dim[key]
+                shape[-2] = keys_to_max_dim[key]
+            else:
+                shape[reshape_dim] = keys_to_max_dim[key]
+            reshaped_features[key]["shape"] = tuple(shape)
+        else:
+            reshaped_features[key] = features[key]
+    return reshaped_features
+
+def create_padded_features(item: dict, features: dict = {}):
+    for key, ft in features.items():
+        if any([k in key for k in ["cam", "effort", "absolute"]]):  # FIXME(mshukor): temporary hack
+            continue
+        shape = ft["shape"]
+        if len(shape) == 3:  # images to torch format (C, H, W)
+            shape = (shape[2], shape[0], shape[1])
+        if len(shape) == 1 and shape[0] == 1:  # ft with shape are actually tensor(ele)
+            shape = []
+        if key not in item:
+            dtype = str_to_torch_dtype(ft["dtype"])
+            item[key] = torch.zeros(shape, dtype=dtype)
+            item[f"{key}_padding_mask"] = torch.tensor(0, dtype=torch.int64)
+            if "image" in key:  # FIXME(mshukor): support other observations
+                item[f"{key}_is_pad"] = torch.BoolTensor([False])
+        else:
+            item[f"{key}_padding_mask"] = torch.tensor(1, dtype=torch.int64)
+    return item
+
+def str_to_torch_dtype(dtype_str):
+    """Convert a dtype string to a torch dtype."""
+    mapping = {
+        "float32": torch.float32,
+        "int64": torch.int64,
+        "int16": torch.int16,
+        "bool": torch.bool,
+        "video": torch.float32,  # Assuming video is stored as uint8 images
+    }
+    return mapping.get(dtype_str, torch.float32)  # Default to float32
+
+def episode_stats_values(meta):
+    episodes = meta.episodes.to_pandas().to_dict(orient="records")
+    return [
+        {k: v for k, v in ep.items() if isinstance(v, dict) and "mean" in v}
+        for ep in episodes
+    ]

src/lerobot/datasets/pipeline_features.py

@@ -17,9 +17,9 @@ from collections.abc import Sequence
 from typing import Any
 
 from lerobot.configs.types import PipelineFeatureType
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.processor import DataProcessorPipeline
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
 
 
 def create_initial_features(
@@ -92,8 +92,8 @@ def aggregate_pipeline_dataset_features(
 
     # Intermediate storage for categorized and filtered features.
     processed_features: dict[str, dict[str, Any]] = {
-        ACTION: {},
-        OBS_STR: {},
+        "action": {},
+        "observation": {},
     }
     images_token = OBS_IMAGES.split(".")[-1]
 
@@ -125,15 +125,17 @@ def aggregate_pipeline_dataset_features(
             # 3. Add the feature to the appropriate group with a clean name.
             name = strip_prefix(key, PREFIXES_TO_STRIP)
             if is_action:
-                processed_features[ACTION][name] = value
+                processed_features["action"][name] = value
             else:
-                processed_features[OBS_STR][name] = value
+                processed_features["observation"][name] = value
 
     # Convert the processed features into the final dataset format.
     dataset_features = {}
-    if processed_features[ACTION]:
-        dataset_features.update(hw_to_dataset_features(processed_features[ACTION], ACTION, use_videos))
-    if processed_features[OBS_STR]:
-        dataset_features.update(hw_to_dataset_features(processed_features[OBS_STR], OBS_STR, use_videos))
+    if processed_features["action"]:
+        dataset_features.update(hw_to_dataset_features(processed_features["action"], ACTION, use_videos))
+    if processed_features["observation"]:
+        dataset_features.update(
+            hw_to_dataset_features(processed_features["observation"], "observation", use_videos)
+        )
 
     return dataset_features

src/lerobot/datasets/push_dataset_to_hub/utils.py

@@ -13,10 +13,67 @@
 # 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.
+import inspect
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
 
 import datasets
+import numpy
+import PIL
 import torch
 
+from lerobot.datasets.video_utils import encode_video_frames
+
+
+def concatenate_episodes(ep_dicts):
+    data_dict = {}
+
+    keys = ep_dicts[0].keys()
+    for key in keys:
+        if torch.is_tensor(ep_dicts[0][key][0]):
+            data_dict[key] = torch.cat([ep_dict[key] for ep_dict in ep_dicts])
+        else:
+            if key not in data_dict:
+                data_dict[key] = []
+            for ep_dict in ep_dicts:
+                for x in ep_dict[key]:
+                    data_dict[key].append(x)
+
+    total_frames = data_dict["frame_index"].shape[0]
+    data_dict["index"] = torch.arange(0, total_frames, 1)
+    return data_dict
+
+
+def save_images_concurrently(imgs_array: numpy.array, out_dir: Path, max_workers: int = 4):
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    def save_image(img_array, i, out_dir):
+        img = PIL.Image.fromarray(img_array)
+        img.save(str(out_dir / f"frame_{i:06d}.png"), quality=100)
+
+    num_images = len(imgs_array)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        [executor.submit(save_image, imgs_array[i], i, out_dir) for i in range(num_images)]
+
+
+def get_default_encoding() -> dict:
+    """Returns the default ffmpeg encoding parameters used by `encode_video_frames`."""
+    signature = inspect.signature(encode_video_frames)
+    return {
+        k: v.default
+        for k, v in signature.parameters.items()
+        if v.default is not inspect.Parameter.empty and k in ["vcodec", "pix_fmt", "g", "crf"]
+    }
+
+
+def check_repo_id(repo_id: str) -> None:
+    if len(repo_id.split("/")) != 2:
+        raise ValueError(
+            f"""`repo_id` is expected to contain a community or user id `/` the name of the dataset
+            (e.g. 'lerobot/pusht'), but contains '{repo_id}'."""
+        )
+
 
 # TODO(aliberts): remove
 def calculate_episode_data_index(hf_dataset: datasets.Dataset) -> dict[str, torch.Tensor]:

src/lerobot/datasets/streaming_dataset.py

@@ -21,6 +21,7 @@ import numpy as np
 import torch
 from datasets import load_dataset
 
+from lerobot.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
 from lerobot.datasets.utils import (
     Backtrackable,
@@ -37,7 +38,6 @@ from lerobot.datasets.video_utils import (
     VideoDecoderCache,
     decode_video_frames_torchcodec,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 
 
 class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
@@ -298,7 +298,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
 
         return padding_mask
 
-    def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
+    def make_frame(
+        self, dataset_iterator: Backtrackable, previous_dataset_iterator: Backtrackable | None = None
+    ) -> Generator:
         """Makes a frame starting from a dataset iterator"""
         item = next(dataset_iterator)
         item = item_to_torch(item)

src/lerobot/datasets/transforms.py

@@ -120,7 +120,7 @@ class SharpnessJitter(Transform):
         self.sharpness = self._check_input(sharpness)
 
     def _check_input(self, sharpness):
-        if isinstance(sharpness, (int | float)):
+        if isinstance(sharpness, (int, float)):
             if sharpness < 0:
                 raise ValueError("If sharpness is a single number, it must be non negative.")
             sharpness = [1.0 - sharpness, 1.0 + sharpness]

src/lerobot/datasets/utils.py

@@ -21,7 +21,7 @@ from collections import deque
 from collections.abc import Iterable, Iterator
 from pathlib import Path
 from pprint import pformat
-from typing import Any, Generic, TypeVar
+from typing import Any, Deque, Generic, TypeVar
 
 import datasets
 import numpy as np
@@ -43,7 +43,6 @@ from lerobot.datasets.backward_compatibility import (
     BackwardCompatibilityError,
     ForwardCompatibilityError,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STR
 from lerobot.utils.utils import is_valid_numpy_dtype_string
 
 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
@@ -67,6 +66,18 @@ DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/frame-{fram
 LEGACY_EPISODES_PATH = "meta/episodes.jsonl"
 LEGACY_EPISODES_STATS_PATH = "meta/episodes_stats.jsonl"
 LEGACY_TASKS_PATH = "meta/tasks.jsonl"
+LEGACY_DEFAULT_VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
+LEGACY_DEFAULT_PARQUET_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
+
+DATASET_CARD_TEMPLATE = """
+---
+# Metadata will go there
+---
+This dataset was created using [LeRobot](https://github.com/huggingface/lerobot).
+
+## {}
+
+"""
 
 DEFAULT_FEATURES = {
     "timestamp": {"dtype": "float32", "shape": (1,), "names": None},
@@ -207,13 +218,13 @@ def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
     """
     serialized_dict = {}
     for key, value in flatten_dict(stats).items():
-        if isinstance(value, (torch.Tensor | np.ndarray)):
+        if isinstance(value, (torch.Tensor, np.ndarray)):
             serialized_dict[key] = value.tolist()
-        elif isinstance(value, list) and isinstance(value[0], (int | float | list)):
+        elif isinstance(value, list) and isinstance(value[0], (int, float, list)):
             serialized_dict[key] = value
         elif isinstance(value, np.generic):
             serialized_dict[key] = value.item()
-        elif isinstance(value, (int | float)):
+        elif isinstance(value, (int, float)):
             serialized_dict[key] = value
         else:
             raise NotImplementedError(f"The value '{value}' of type '{type(value)}' is not supported.")
@@ -371,6 +382,12 @@ def load_episodes(local_dir: Path) -> datasets.Dataset:
     return episodes
 
 
+def backward_compatible_episodes_stats(
+    stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
+) -> dict[int, dict[str, dict[str, np.ndarray]]]:
+    return dict.fromkeys(episodes, stats)
+
+
 def load_image_as_numpy(
     fpath: str | Path, dtype: np.dtype = np.float32, channel_first: bool = True
 ) -> np.ndarray:
@@ -628,14 +645,14 @@ def hw_to_dataset_features(
     }
     cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)}
 
-    if joint_fts and prefix == ACTION:
+    if joint_fts and prefix == "action":
         features[prefix] = {
             "dtype": "float32",
             "shape": (len(joint_fts),),
             "names": list(joint_fts),
         }
 
-    if joint_fts and prefix == OBS_STR:
+    if joint_fts and prefix == "observation":
         features[f"{prefix}.state"] = {
             "dtype": "float32",
             "shape": (len(joint_fts),),
@@ -711,11 +728,11 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
             # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets.
             if names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
                 shape = (shape[2], shape[0], shape[1])
-        elif key == OBS_ENV_STATE:
+        elif key == "observation.environment_state":
             type = FeatureType.ENV
-        elif key.startswith(OBS_STR):
+        elif key.startswith("observation"):
             type = FeatureType.STATE
-        elif key.startswith(ACTION):
+        elif key.startswith("action"):
             type = FeatureType.ACTION
         else:
             continue
@@ -1179,7 +1196,7 @@ def item_to_torch(item: dict) -> dict:
         dict: Dictionary with all tensor-like items converted to torch.Tensor.
     """
     for key, val in item.items():
-        if isinstance(val, (np.ndarray | list)) and key not in ["task"]:
+        if isinstance(val, (np.ndarray, list)) and key not in ["task"]:
             # Convert numpy arrays and lists to torch tensors
             item[key] = torch.tensor(val)
     return item
@@ -1253,8 +1270,8 @@ class Backtrackable(Generic[T]):
             raise ValueError("lookahead must be > 0")
 
         self._source: Iterator[T] = iter(iterable)
-        self._back_buf: deque[T] = deque(maxlen=history)
-        self._ahead_buf: deque[T] = deque(maxlen=lookahead) if lookahead > 0 else deque()
+        self._back_buf: Deque[T] = deque(maxlen=history)
+        self._ahead_buf: Deque[T] = deque(maxlen=lookahead) if lookahead > 0 else deque()
         self._cursor: int = 0
         self._history = history
         self._lookahead = lookahead
@@ -1328,6 +1345,12 @@ class Backtrackable(Generic[T]):
         # When cursor<0, slice so the order remains chronological
         return list(self._back_buf)[: self._cursor or None]
 
+    def lookahead_buffer(self) -> list[T]:
+        """
+        Return a copy of the current lookahead buffer.
+        """
+        return list(self._ahead_buf)
+
     def can_peek_back(self, steps: int = 1) -> bool:
         """
         Check if we can go back `steps` items without raising an IndexError.
@@ -1353,6 +1376,31 @@ class Backtrackable(Generic[T]):
         except StopIteration:
             return False
 
+    def reset_cursor(self) -> None:
+        """
+        Reset cursor to the most recent position (equivalent to calling next()
+        until you're back to the latest item).
+        """
+        self._cursor = 0
+
+    def clear_ahead_buffer(self) -> None:
+        """
+        Clear the ahead buffer, discarding any pre-fetched items.
+        """
+        self._ahead_buf.clear()
+
+    def switch_source_iterable(self, new_source: Iterable[T]) -> None:
+        """
+        Switch the source of the backtrackable to a new iterable, keeping the history.
+
+        This is useful when iterating over a sequence of datasets. The history from the
+        previous source is kept, but the lookahead buffer is cleared. The cursor is reset
+        to the present.
+        """
+        self._source = iter(new_source)
+        self.clear_ahead_buffer()
+        self.reset_cursor()
+
 
 def safe_shard(dataset: datasets.IterableDataset, index: int, num_shards: int) -> datasets.Dataset:
     """

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -34,7 +34,6 @@ python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
 """
 
 import argparse
-import logging
 import shutil
 from pathlib import Path
 from typing import Any
@@ -47,6 +46,7 @@ from datasets import Dataset, Features, Image
 from huggingface_hub import HfApi, snapshot_download
 from requests import HTTPError
 
+from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
 from lerobot.datasets.utils import (
@@ -71,8 +71,6 @@ from lerobot.datasets.utils import (
     write_tasks,
 )
 from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
-from lerobot.utils.constants import HF_LEROBOT_HOME
-from lerobot.utils.utils import init_logging
 
 V21 = "v2.1"
 
@@ -146,7 +144,6 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
 
 
 def convert_tasks(root, new_root):
-    logging.info(f"Converting tasks from {root} to {new_root}")
     tasks, _ = legacy_load_tasks(root)
     task_indices = tasks.keys()
     task_strings = tasks.values()
@@ -188,10 +185,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
     num_frames = 0
     paths_to_cat = []
     episodes_metadata = []
-
-    logging.info(f"Converting data files from {len(ep_paths)} episodes")
-
-    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
+    for ep_path in ep_paths:
         ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
         ep_num_frames = get_parquet_num_frames(ep_path)
         ep_metadata = {
@@ -215,6 +209,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
 
         # Reset for the next file
         size_in_mb = ep_size_in_mb
+        num_frames = ep_num_frames
         paths_to_cat = [ep_path]
 
         chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
@@ -241,8 +236,6 @@ def get_image_keys(root):
 
 
 def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
-    logging.info(f"Converting videos from {root} to {new_root}")
-
     video_keys = get_video_keys(root)
     if len(video_keys) == 0:
         return None
@@ -261,7 +254,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
     episods_metadata = []
     num_cameras = len(video_keys)
     num_episodes = num_eps_per_cam[0]
-    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
+    for ep_idx in range(num_episodes):
         # Sanity check
         ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
         ep_ids += [ep_idx]
@@ -288,7 +281,6 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
     duration_in_s = 0.0
     paths_to_cat = []
     episodes_metadata = []
-
     for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
         ep_size_in_mb = get_video_size_in_mb(ep_path)
         ep_duration_in_s = get_video_duration_in_s(ep_path)
@@ -382,8 +374,6 @@ def generate_episode_metadata_dict(
 
 
 def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None):
-    logging.info(f"Converting episodes metadata from {root} to {new_root}")
-
     episodes_legacy_metadata = legacy_load_episodes(root)
     episodes_stats = legacy_load_episodes_stats(root)
 
@@ -415,7 +405,6 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
     info["data_path"] = DEFAULT_DATA_PATH
     info["video_path"] = DEFAULT_VIDEO_PATH
     info["fps"] = int(info["fps"])
-    logging.info(f"Converting info from {root} to {new_root}")
     for key in info["features"]:
         if info["features"][key]["dtype"] == "video":
             # already has fps in video_info
@@ -480,7 +469,6 @@ def convert_dataset(
 
 
 if __name__ == "__main__":
-    init_logging()
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--repo-id",

src/lerobot/datasets/video_utils.py

@@ -428,7 +428,7 @@ def concatenate_video_files(
     with tempfile.NamedTemporaryFile(mode="w", suffix=".ffconcat", delete=False) as tmp_concatenate_file:
         tmp_concatenate_file.write("ffconcat version 1.0\n")
         for input_path in input_video_paths:
-            tmp_concatenate_file.write(f"file '{str(input_path.resolve())}'\n")
+            tmp_concatenate_file.write(f"file '{str(input_path)}'\n")
         tmp_concatenate_file.flush()
         tmp_concatenate_path = tmp_concatenate_file.name
 
@@ -437,9 +437,7 @@ def concatenate_video_files(
         tmp_concatenate_path, mode="r", format="concat", options={"safe": "0"}
     )  # safe = 0 allows absolute paths as well as relative paths
 
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_named_file:
-        tmp_output_video_path = tmp_named_file.name
-
+    tmp_output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     output_container = av.open(
         tmp_output_video_path, mode="w", options={"movflags": "faststart"}
     )  # faststart is to move the metadata to the beginning of the file to speed up loading
@@ -587,6 +585,19 @@ def get_video_pixel_channels(pix_fmt: str) -> int:
         raise ValueError("Unknown format")
 
 
+def get_image_pixel_channels(image: Image):
+    if image.mode == "L":
+        return 1  # Grayscale
+    elif image.mode == "LA":
+        return 2  # Grayscale + Alpha
+    elif image.mode == "RGB":
+        return 3  # RGB
+    elif image.mode == "RGBA":
+        return 4  # RGBA
+    else:
+        raise ValueError("Unknown format")
+
+
 def get_video_duration_in_s(video_path: Path | str) -> float:
     """
     Get the duration of a video file in seconds using PyAV.

src/lerobot/envs/configs.py

@@ -19,9 +19,9 @@ from typing import Any
 import draccus
 
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.robots import RobotConfig
 from lerobot.teleoperators.config import TeleoperatorConfig
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
 
 @dataclass
@@ -53,12 +53,12 @@ class AlohaEnv(EnvConfig):
     render_mode: str = "rgb_array"
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "top": f"{OBS_IMAGE}.top",
             "pixels/top": f"{OBS_IMAGES}.top",
@@ -93,13 +93,13 @@ class PushtEnv(EnvConfig):
     visualization_height: int = 384
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
             "agent_pos": PolicyFeature(type=FeatureType.STATE, shape=(2,)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "environment_state": OBS_ENV_STATE,
             "pixels": OBS_IMAGE,
@@ -135,13 +135,13 @@ class XarmEnv(EnvConfig):
     visualization_height: int = 384
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
             "pixels": PolicyFeature(type=FeatureType.VISUAL, shape=(84, 84, 3)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "pixels": OBS_IMAGE,
         }
@@ -193,6 +193,7 @@ class ObservationConfig:
 
     add_joint_velocity_to_observation: bool = False
     add_current_to_observation: bool = False
+    add_ee_pose_to_observation: bool = False
     display_cameras: bool = False
 
 
@@ -202,6 +203,7 @@ class GripperConfig:
 
     use_gripper: bool = True
     gripper_penalty: float = 0.0
+    gripper_penalty_in_reward: bool = False
 
 
 @dataclass
@@ -257,12 +259,12 @@ class LiberoEnv(EnvConfig):
     camera_name_mapping: dict[str, str] | None = (None,)
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
         }
     )
     features_map: dict[str, str] = field(
         default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
             "agent_pos": OBS_STATE,
             "pixels/agentview_image": f"{OBS_IMAGES}.image",
             "pixels/robot0_eye_in_hand_image": f"{OBS_IMAGES}.image2",

src/lerobot/envs/libero.py

@@ -35,7 +35,7 @@ def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
     """Normalize camera_name into a non-empty list of strings."""
     if isinstance(camera_name, str):
         cams = [c.strip() for c in camera_name.split(",") if c.strip()]
-    elif isinstance(camera_name, (list | tuple)):
+    elif isinstance(camera_name, (list, tuple)):
         cams = [str(c).strip() for c in camera_name if str(c).strip()]
     else:
         raise TypeError(f"camera_name must be str or sequence[str], got {type(camera_name).__name__}")

src/lerobot/envs/utils.py

@@ -26,7 +26,6 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.envs.configs import EnvConfig
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.utils.utils import get_channel_first_image_shape
 
 
@@ -42,9 +41,9 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
     return_observations = {}
     if "pixels" in observations:
         if isinstance(observations["pixels"], dict):
-            imgs = {f"{OBS_IMAGES}.{key}": img for key, img in observations["pixels"].items()}
+            imgs = {f"observation.images.{key}": img for key, img in observations["pixels"].items()}
         else:
-            imgs = {OBS_IMAGE: observations["pixels"]}
+            imgs = {"observation.image": observations["pixels"]}
 
         for imgkey, img in imgs.items():
             # TODO(aliberts, rcadene): use transforms.ToTensor()?
@@ -73,13 +72,13 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
         if env_state.dim() == 1:
             env_state = env_state.unsqueeze(0)
 
-        return_observations[OBS_ENV_STATE] = env_state
+        return_observations["observation.environment_state"] = env_state
 
     # TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
     agent_pos = torch.from_numpy(observations["agent_pos"]).float()
     if agent_pos.dim() == 1:
         agent_pos = agent_pos.unsqueeze(0)
-    return_observations[OBS_STATE] = agent_pos
+    return_observations["observation.state"] = agent_pos
 
     return return_observations
 
@@ -183,10 +182,10 @@ def _(env: Mapping) -> None:
 
 @close_envs.register
 def _(envs: Sequence) -> None:
-    if isinstance(envs, (str | bytes)):
+    if isinstance(envs, (str, bytes)):
         return
     for v in envs:
-        if isinstance(v, Mapping) or isinstance(v, Sequence) and not isinstance(v, (str | bytes)):
+        if isinstance(v, Mapping) or isinstance(v, Sequence) and not isinstance(v, (str, bytes)):
             close_envs(v)
         elif hasattr(v, "close"):
             _close_single_env(v)

src/lerobot/errors.py

@@ -30,3 +30,14 @@ class DeviceAlreadyConnectedError(ConnectionError):
     ):
         self.message = message
         super().__init__(self.message)
+
+
+class InvalidActionError(ValueError):
+    """Exception raised when an action is already invalid."""
+
+    def __init__(
+        self,
+        message="The action is invalid. Check the value follows what it is expected from the action space.",
+    ):
+        self.message = message
+        super().__init__(self.message)

src/lerobot/find_cameras.py

@@ -24,7 +24,7 @@ lerobot-find-cameras
 ```
 """
 
-# NOTE(Steven): RealSense can also be identified/opened as OpenCV cameras. If you know the camera is a RealSense, use the `lerobot-find-cameras realsense` flag to avoid confusion.
+# NOTE(Steven): RealSense can also be identified/opened as OpenCV cameras. If you know the camera is a RealSense, use the `lerobot.find_cameras realsense` flag to avoid confusion.
 # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
 
 import argparse

src/lerobot/motors/dynamixel/dynamixel.py

@@ -22,7 +22,7 @@ import logging
 from copy import deepcopy
 from enum import Enum
 
-from lerobot.motors.encoding_utils import decode_twos_complement, encode_twos_complement
+from lerobot.utils.encoding_utils import decode_twos_complement, encode_twos_complement
 
 from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
 from .tables import (

src/lerobot/motors/feetech/feetech.py

@@ -17,7 +17,7 @@ from copy import deepcopy
 from enum import Enum
 from pprint import pformat
 
-from lerobot.motors.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
+from lerobot.utils.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
 
 from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
 from .tables import (

src/lerobot/motors/motors_bus.py

@@ -32,7 +32,7 @@ import serial
 from deepdiff import DeepDiff
 from tqdm import tqdm
 
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.utils.utils import enter_pressed, move_cursor_up
 
 NameOrID: TypeAlias = str | int
@@ -99,6 +99,12 @@ class Motor:
     norm_mode: MotorNormMode
 
 
+class JointOutOfRangeError(Exception):
+    def __init__(self, message="Joint is out of range"):
+        self.message = message
+        super().__init__(self.message)
+
+
 class PortHandler(Protocol):
     def __init__(self, port_name):
         self.is_open: bool
@@ -342,7 +348,7 @@ class MotorsBus(abc.ABC):
             raise TypeError(motors)
 
     def _get_ids_values_dict(self, values: Value | dict[str, Value] | None) -> list[str]:
-        if isinstance(values, (int | float)):
+        if isinstance(values, (int, float)):
             return dict.fromkeys(self.ids, values)
         elif isinstance(values, dict):
             return {self.motors[motor].id: val for motor, val in values.items()}
@@ -669,7 +675,7 @@ class MotorsBus(abc.ABC):
         """
         if motors is None:
             motors = list(self.motors)
-        elif isinstance(motors, (str | int)):
+        elif isinstance(motors, (str, int)):
             motors = [motors]
         elif not isinstance(motors, list):
             raise TypeError(motors)
@@ -697,7 +703,7 @@ class MotorsBus(abc.ABC):
         """
         if motors is None:
             motors = list(self.motors)
-        elif isinstance(motors, (str | int)):
+        elif isinstance(motors, (str, int)):
             motors = [motors]
         elif not isinstance(motors, list):
             raise TypeError(motors)
@@ -733,7 +739,7 @@ class MotorsBus(abc.ABC):
         """
         if motors is None:
             motors = list(self.motors)
-        elif isinstance(motors, (str | int)):
+        elif isinstance(motors, (str, int)):
             motors = [motors]
         elif not isinstance(motors, list):
             raise TypeError(motors)

src/lerobot/optim/optimizers.py

@@ -22,11 +22,11 @@ import draccus
 import torch
 from safetensors.torch import load_file, save_file
 
-from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
-from lerobot.utils.constants import (
+from lerobot.constants import (
     OPTIMIZER_PARAM_GROUPS,
     OPTIMIZER_STATE,
 )
+from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
 from lerobot.utils.io_utils import deserialize_json_into_object
 
 

src/lerobot/optim/schedulers.py

@@ -22,8 +22,8 @@ import draccus
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler
 
+from lerobot.constants import SCHEDULER_STATE
 from lerobot.datasets.utils import write_json
-from lerobot.utils.constants import SCHEDULER_STATE
 from lerobot.utils.io_utils import deserialize_json_into_object
 
 

src/lerobot/policies/act/modeling_act.py

@@ -33,9 +33,9 @@ from torch import Tensor, nn
 from torchvision.models._utils import IntermediateLayerGetter
 from torchvision.ops.misc import FrozenBatchNorm2d
 
+from lerobot.constants import ACTION, OBS_IMAGES
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 
 
 class ACTPolicy(PreTrainedPolicy):
@@ -394,22 +394,25 @@ class ACT(nn.Module):
             latent dimension.
         """
         if self.config.use_vae and self.training:
-            assert ACTION in batch, (
+            assert "action" in batch, (
                 "actions must be provided when using the variational objective in training mode."
             )
 
-        batch_size = batch[OBS_IMAGES][0].shape[0] if OBS_IMAGES in batch else batch[OBS_ENV_STATE].shape[0]
+        if "observation.images" in batch:
+            batch_size = batch["observation.images"][0].shape[0]
+        else:
+            batch_size = batch["observation.environment_state"].shape[0]
 
         # Prepare the latent for input to the transformer encoder.
-        if self.config.use_vae and ACTION in batch and self.training:
+        if self.config.use_vae and "action" in batch and self.training:
             # Prepare the input to the VAE encoder: [cls, *joint_space_configuration, *action_sequence].
             cls_embed = einops.repeat(
                 self.vae_encoder_cls_embed.weight, "1 d -> b 1 d", b=batch_size
             )  # (B, 1, D)
             if self.config.robot_state_feature:
-                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch[OBS_STATE])
+                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch["observation.state"])
                 robot_state_embed = robot_state_embed.unsqueeze(1)  # (B, 1, D)
-            action_embed = self.vae_encoder_action_input_proj(batch[ACTION])  # (B, S, D)
+            action_embed = self.vae_encoder_action_input_proj(batch["action"])  # (B, S, D)
 
             if self.config.robot_state_feature:
                 vae_encoder_input = [cls_embed, robot_state_embed, action_embed]  # (B, S+2, D)
@@ -427,7 +430,7 @@ class ACT(nn.Module):
             cls_joint_is_pad = torch.full(
                 (batch_size, 2 if self.config.robot_state_feature else 1),
                 False,
-                device=batch[OBS_STATE].device,
+                device=batch["observation.state"].device,
             )
             key_padding_mask = torch.cat(
                 [cls_joint_is_pad, batch["action_is_pad"]], axis=1
@@ -451,7 +454,7 @@ class ACT(nn.Module):
             mu = log_sigma_x2 = None
             # TODO(rcadene, alexander-soare): remove call to `.to` to speedup forward ; precompute and use buffer
             latent_sample = torch.zeros([batch_size, self.config.latent_dim], dtype=torch.float32).to(
-                batch[OBS_STATE].device
+                batch["observation.state"].device
             )
 
         # Prepare transformer encoder inputs.
@@ -459,16 +462,18 @@ class ACT(nn.Module):
         encoder_in_pos_embed = list(self.encoder_1d_feature_pos_embed.weight.unsqueeze(1))
         # Robot state token.
         if self.config.robot_state_feature:
-            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch[OBS_STATE]))
+            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch["observation.state"]))
         # Environment state token.
         if self.config.env_state_feature:
-            encoder_in_tokens.append(self.encoder_env_state_input_proj(batch[OBS_ENV_STATE]))
+            encoder_in_tokens.append(
+                self.encoder_env_state_input_proj(batch["observation.environment_state"])
+            )
 
         if self.config.image_features:
             # For a list of images, the H and W may vary but H*W is constant.
             # NOTE: If modifying this section, verify on MPS devices that
             # gradients remain stable (no explosions or NaNs).
-            for img in batch[OBS_IMAGES]:
+            for img in batch["observation.images"]:
                 cam_features = self.backbone(img)["feature_map"]
                 cam_pos_embed = self.encoder_cam_feat_pos_embed(cam_features).to(dtype=cam_features.dtype)
                 cam_features = self.encoder_img_feat_input_proj(cam_features)

src/lerobot/policies/act/processor_act.py

@@ -17,6 +17,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -28,7 +29,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_act_pre_post_processors(

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -33,6 +33,7 @@ from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import (
@@ -41,7 +42,6 @@ from lerobot.policies.utils import (
     get_output_shape,
     populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 
 
 class DiffusionPolicy(PreTrainedPolicy):
@@ -81,13 +81,13 @@ class DiffusionPolicy(PreTrainedPolicy):
     def reset(self):
         """Clear observation and action queues. Should be called on `env.reset()`"""
         self._queues = {
-            OBS_STATE: deque(maxlen=self.config.n_obs_steps),
-            ACTION: deque(maxlen=self.config.n_action_steps),
+            "observation.state": deque(maxlen=self.config.n_obs_steps),
+            "action": deque(maxlen=self.config.n_action_steps),
         }
         if self.config.image_features:
-            self._queues[OBS_IMAGES] = deque(maxlen=self.config.n_obs_steps)
+            self._queues["observation.images"] = deque(maxlen=self.config.n_obs_steps)
         if self.config.env_state_feature:
-            self._queues[OBS_ENV_STATE] = deque(maxlen=self.config.n_obs_steps)
+            self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)
 
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
@@ -234,7 +234,7 @@ class DiffusionModel(nn.Module):
         if self.config.image_features:
             if self.config.use_separate_rgb_encoder_per_camera:
                 # Combine batch and sequence dims while rearranging to make the camera index dimension first.
-                images_per_camera = einops.rearrange(batch[OBS_IMAGES], "b s n ... -> n (b s) ...")
+                images_per_camera = einops.rearrange(batch["observation.images"], "b s n ... -> n (b s) ...")
                 img_features_list = torch.cat(
                     [
                         encoder(images)
@@ -249,7 +249,7 @@ class DiffusionModel(nn.Module):
             else:
                 # Combine batch, sequence, and "which camera" dims before passing to shared encoder.
                 img_features = self.rgb_encoder(
-                    einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ...")
+                    einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
                 )
                 # Separate batch dim and sequence dim back out. The camera index dim gets absorbed into the
                 # feature dim (effectively concatenating the camera features).
@@ -275,7 +275,7 @@ class DiffusionModel(nn.Module):
             "observation.environment_state": (B, n_obs_steps, environment_dim)
         }
         """
-        batch_size, n_obs_steps = batch[OBS_STATE].shape[:2]
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
         assert n_obs_steps == self.config.n_obs_steps
 
         # Encode image features and concatenate them all together along with the state vector.
@@ -306,10 +306,10 @@ class DiffusionModel(nn.Module):
         }
         """
         # Input validation.
-        assert set(batch).issuperset({OBS_STATE, ACTION, "action_is_pad"})
-        assert OBS_IMAGES in batch or OBS_ENV_STATE in batch
-        n_obs_steps = batch[OBS_STATE].shape[1]
-        horizon = batch[ACTION].shape[1]
+        assert set(batch).issuperset({"observation.state", "action", "action_is_pad"})
+        assert "observation.images" in batch or "observation.environment_state" in batch
+        n_obs_steps = batch["observation.state"].shape[1]
+        horizon = batch["action"].shape[1]
         assert horizon == self.config.horizon
         assert n_obs_steps == self.config.n_obs_steps
 
@@ -317,7 +317,7 @@ class DiffusionModel(nn.Module):
         global_cond = self._prepare_global_conditioning(batch)  # (B, global_cond_dim)
 
         # Forward diffusion.
-        trajectory = batch[ACTION]
+        trajectory = batch["action"]
         # Sample noise to add to the trajectory.
         eps = torch.randn(trajectory.shape, device=trajectory.device)
         # Sample a random noising timestep for each item in the batch.
@@ -338,7 +338,7 @@ class DiffusionModel(nn.Module):
         if self.config.prediction_type == "epsilon":
             target = eps
         elif self.config.prediction_type == "sample":
-            target = batch[ACTION]
+            target = batch["action"]
         else:
             raise ValueError(f"Unsupported prediction type {self.config.prediction_type}")
 

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_diffusion_pre_post_processors(

src/lerobot/policies/factory.py

@@ -24,6 +24,7 @@ from typing_extensions import Unpack
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.envs.configs import EnvConfig
@@ -45,7 +46,6 @@ from lerobot.processor.converters import (
     transition_to_batch,
     transition_to_policy_action,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def get_policy_class(name: str) -> type[PreTrainedPolicy]:

src/lerobot/policies/pi0/configuration_pi0.py

@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("pi0")
@@ -114,7 +113,7 @@ class PI0Config(PreTrainedConfig):
         #     raise ValueError("You must provide at least one image or the environment state among the inputs.")
 
         for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
             empty_camera = PolicyFeature(
                 type=FeatureType.VISUAL,
                 shape=(3, 480, 640),

src/lerobot/policies/pi0/conversion_scripts/compare_with_jax.py

@@ -21,7 +21,6 @@ import torch
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.policies.factory import make_policy
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 
 def display(tensor: torch.Tensor):
@@ -61,26 +60,26 @@ def main():
 
     # Override stats
     dataset_meta = LeRobotDatasetMetadata(dataset_repo_id)
-    dataset_meta.stats[OBS_STATE]["mean"] = torch.tensor(
+    dataset_meta.stats["observation.state"]["mean"] = torch.tensor(
         norm_stats["norm_stats"]["state"]["mean"][:num_motors], dtype=torch.float32
     )
-    dataset_meta.stats[OBS_STATE]["std"] = torch.tensor(
+    dataset_meta.stats["observation.state"]["std"] = torch.tensor(
         norm_stats["norm_stats"]["state"]["std"][:num_motors], dtype=torch.float32
     )
 
     # Create LeRobot batch from Jax
     batch = {}
     for cam_key, uint_chw_array in example["images"].items():
-        batch[f"{OBS_IMAGES}.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
-    batch[OBS_STATE] = torch.from_numpy(example["state"])
-    batch[ACTION] = torch.from_numpy(outputs["actions"])
+        batch[f"observation.images.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
+    batch["observation.state"] = torch.from_numpy(example["state"])
+    batch["action"] = torch.from_numpy(outputs["actions"])
     batch["task"] = example["prompt"]
 
     if model_name == "pi0_aloha_towel":
-        del batch[f"{OBS_IMAGES}.cam_low"]
+        del batch["observation.images.cam_low"]
     elif model_name == "pi0_aloha_sim":
-        batch[f"{OBS_IMAGES}.top"] = batch[f"{OBS_IMAGES}.cam_high"]
-        del batch[f"{OBS_IMAGES}.cam_high"]
+        batch["observation.images.top"] = batch["observation.images.cam_high"]
+        del batch["observation.images.cam_high"]
 
     # Batchify
     for key in batch:
@@ -117,7 +116,7 @@ def main():
         actions.append(action)
 
     actions = torch.stack(actions, dim=1)
-    pi_actions = batch[ACTION]
+    pi_actions = batch["action"]
     print("actions")
     display(actions)
     print()

src/lerobot/policies/pi0/modeling_pi0.py

@@ -57,13 +57,13 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0.paligemma_with_expert import (
     PaliGemmaWithExpertConfig,
     PaliGemmaWithExpertModel,
 )
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype
 
 

src/lerobot/policies/pi0/processor_pi0.py

@@ -19,6 +19,7 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -34,7 +35,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 @ProcessorStepRegistry.register(name="pi0_new_line_processor")

src/lerobot/policies/pi0fast/configuration_pi0fast.py

@@ -6,7 +6,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("pi0fast")
@@ -100,7 +99,7 @@ class PI0FASTConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
             empty_camera = PolicyFeature(
                 type=FeatureType.VISUAL,
                 shape=(3, 480, 640),

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -57,9 +57,9 @@ from transformers import AutoProcessor, AutoTokenizer, PaliGemmaForConditionalGe
 from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING
 
+from lerobot.constants import ACTION, OBS_STATE
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_STATE
 
 PRECISION = {
     "float16": torch.float16,

src/lerobot/policies/pi0fast/processor_pi0fast.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_pi0fast_pre_post_processors(

src/lerobot/policies/sac/configuration_sac.py

@@ -19,8 +19,8 @@ from dataclasses import dataclass, field
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
 from lerobot.optim.optimizers import MultiAdamConfig
-from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
 
 
 def is_image_feature(key: str) -> bool:
@@ -139,6 +139,8 @@ class SACConfig(PreTrainedConfig):
     # Training parameter
     # Number of steps for online training
     online_steps: int = 1000000
+    # Seed for the online environment
+    online_env_seed: int = 10000
     # Capacity of the online replay buffer
     online_buffer_capacity: int = 100000
     # Capacity of the offline replay buffer
@@ -223,7 +225,7 @@ class SACConfig(PreTrainedConfig):
                 "You must provide either 'observation.state' or an image observation (key starting with 'observation.image') in the input features"
             )
 
-        if ACTION not in self.output_features:
+        if "action" not in self.output_features:
             raise ValueError("You must provide 'action' in the output features")
 
     @property

src/lerobot/policies/sac/modeling_sac.py

@@ -31,7 +31,6 @@ from torch.distributions import MultivariateNormal, TanhTransform, Transform, Tr
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE
 
 DISCRETE_DIMENSION_INDEX = -1  # Gripper is always the last dimension
 
@@ -51,7 +50,7 @@ class SACPolicy(
         self.config = config
 
         # Determine action dimension and initialize all components
-        continuous_action_dim = config.output_features[ACTION].shape[0]
+        continuous_action_dim = config.output_features["action"].shape[0]
         self._init_encoders()
         self._init_critics(continuous_action_dim)
         self._init_actor(continuous_action_dim)
@@ -158,7 +157,7 @@ class SACPolicy(
             The computed loss tensor
         """
         # Extract common components from batch
-        actions: Tensor = batch[ACTION]
+        actions: Tensor = batch["action"]
         observations: dict[str, Tensor] = batch["state"]
         observation_features: Tensor = batch.get("observation_feature")
 
@@ -514,17 +513,17 @@ class SACObservationEncoder(nn.Module):
             )
 
     def _init_state_layers(self) -> None:
-        self.has_env = OBS_ENV_STATE in self.config.input_features
-        self.has_state = OBS_STATE in self.config.input_features
+        self.has_env = "observation.environment_state" in self.config.input_features
+        self.has_state = "observation.state" in self.config.input_features
         if self.has_env:
-            dim = self.config.input_features[OBS_ENV_STATE].shape[0]
+            dim = self.config.input_features["observation.environment_state"].shape[0]
             self.env_encoder = nn.Sequential(
                 nn.Linear(dim, self.config.latent_dim),
                 nn.LayerNorm(self.config.latent_dim),
                 nn.Tanh(),
             )
         if self.has_state:
-            dim = self.config.input_features[OBS_STATE].shape[0]
+            dim = self.config.input_features["observation.state"].shape[0]
             self.state_encoder = nn.Sequential(
                 nn.Linear(dim, self.config.latent_dim),
                 nn.LayerNorm(self.config.latent_dim),
@@ -550,9 +549,9 @@ class SACObservationEncoder(nn.Module):
                 cache = self.get_cached_image_features(obs)
             parts.append(self._encode_images(cache, detach))
         if self.has_env:
-            parts.append(self.env_encoder(obs[OBS_ENV_STATE]))
+            parts.append(self.env_encoder(obs["observation.environment_state"]))
         if self.has_state:
-            parts.append(self.state_encoder(obs[OBS_STATE]))
+            parts.append(self.state_encoder(obs["observation.state"]))
         if parts:
             return torch.cat(parts, dim=-1)
 
@@ -1061,3 +1060,15 @@ class TanhMultivariateNormalDiag(TransformedDistribution):
             x = transform(x)
 
         return x
+
+
+def _convert_normalization_params_to_tensor(normalization_params: dict) -> dict:
+    converted_params = {}
+    for outer_key, inner_dict in normalization_params.items():
+        converted_params[outer_key] = {}
+        for key, value in inner_dict.items():
+            converted_params[outer_key][key] = torch.tensor(value)
+            if "image" in outer_key:
+                converted_params[outer_key][key] = converted_params[outer_key][key].view(3, 1, 1)
+
+    return converted_params

src/lerobot/policies/sac/processor_sac.py

@@ -19,6 +19,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -30,7 +31,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_sac_pre_post_processors(

src/lerobot/policies/sac/reward_model/configuration_classifier.py

@@ -19,7 +19,6 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 from lerobot.optim.optimizers import AdamWConfig, OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
-from lerobot.utils.constants import OBS_IMAGE
 
 
 @PreTrainedConfig.register_subclass(name="reward_classifier")
@@ -70,7 +69,7 @@ class RewardClassifierConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         """Validate feature configurations."""
-        has_image = any(key.startswith(OBS_IMAGE) for key in self.input_features)
+        has_image = any(key.startswith("observation.image") for key in self.input_features)
         if not has_image:
             raise ValueError(
                 "You must provide an image observation (key starting with 'observation.image') in the input features"

src/lerobot/policies/sac/reward_model/modeling_classifier.py

@@ -19,9 +19,9 @@ import logging
 import torch
 from torch import Tensor, nn
 
+from lerobot.constants import OBS_IMAGE, REWARD
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.utils.constants import OBS_IMAGE, REWARD
 
 
 class ClassifierOutput:

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("smolvla")
@@ -118,7 +117,7 @@ class SmolVLAConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
             empty_camera = PolicyFeature(
                 type=FeatureType.VISUAL,
                 shape=(3, 480, 640),

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -59,13 +59,13 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.policies.utils import (
     populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype
 
 

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -19,6 +19,7 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -33,7 +34,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_smolvla_pre_post_processors(

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -35,10 +35,10 @@ import torch.nn as nn
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_STATE, REWARD
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, OBS_STR, REWARD
 
 
 class TDMPCPolicy(PreTrainedPolicy):
@@ -91,13 +91,13 @@ class TDMPCPolicy(PreTrainedPolicy):
         called on `env.reset()`
         """
         self._queues = {
-            OBS_STATE: deque(maxlen=1),
-            ACTION: deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
+            "observation.state": deque(maxlen=1),
+            "action": deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
         }
         if self.config.image_features:
-            self._queues[OBS_IMAGE] = deque(maxlen=1)
+            self._queues["observation.image"] = deque(maxlen=1)
         if self.config.env_state_feature:
-            self._queues[OBS_ENV_STATE] = deque(maxlen=1)
+            self._queues["observation.environment_state"] = deque(maxlen=1)
         # Previous mean obtained from the cross-entropy method (CEM) used during MPC. It is used to warm start
         # CEM for the next step.
         self._prev_mean: torch.Tensor | None = None
@@ -325,7 +325,7 @@ class TDMPCPolicy(PreTrainedPolicy):
 
         action = batch[ACTION]  # (t, b, action_dim)
         reward = batch[REWARD]  # (t, b)
-        observations = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
+        observations = {k: v for k, v in batch.items() if k.startswith("observation.")}
 
         # Apply random image augmentations.
         if self.config.image_features and self.config.max_random_shift_ratio > 0:
@@ -387,10 +387,10 @@ class TDMPCPolicy(PreTrainedPolicy):
                 temporal_loss_coeffs
                 * F.mse_loss(z_preds[1:], z_targets, reduction="none").mean(dim=-1)
                 # `z_preds` depends on the current observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
                 # `z_targets` depends on the next observation.
-                * ~batch[f"{OBS_STR}.state_is_pad"][1:]
+                * ~batch["observation.state_is_pad"][1:]
             )
             .sum(0)
             .mean()
@@ -403,7 +403,7 @@ class TDMPCPolicy(PreTrainedPolicy):
                 * F.mse_loss(reward_preds, reward, reduction="none")
                 * ~batch["next.reward_is_pad"]
                 # `reward_preds` depends on the current observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
             )
             .sum(0)
@@ -419,11 +419,11 @@ class TDMPCPolicy(PreTrainedPolicy):
                     reduction="none",
                 ).sum(0)  # sum over ensemble
                 # `q_preds_ensemble` depends on the first observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
                 # q_targets depends on the reward and the next observations.
                 * ~batch["next.reward_is_pad"]
-                * ~batch[f"{OBS_STR}.state_is_pad"][1:]
+                * ~batch["observation.state_is_pad"][1:]
             )
             .sum(0)
             .mean()
@@ -441,7 +441,7 @@ class TDMPCPolicy(PreTrainedPolicy):
                 temporal_loss_coeffs
                 * raw_v_value_loss
                 # `v_targets` depends on the first observation and the actions, as does `v_preds`.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                 * ~batch["action_is_pad"]
             )
             .sum(0)
@@ -477,7 +477,7 @@ class TDMPCPolicy(PreTrainedPolicy):
             * mse
             * temporal_loss_coeffs
             # `action_preds` depends on the first observation and the actions.
-            * ~batch[f"{OBS_STR}.state_is_pad"][0]
+            * ~batch["observation.state_is_pad"][0]
             * ~batch["action_is_pad"]
         ).mean()
 

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -18,6 +18,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_tdmpc_pre_post_processors(

src/lerobot/policies/vqbet/configuration_vqbet.py

@@ -82,6 +82,7 @@ class VQBeTConfig(PreTrainedConfig):
         gpt_n_head: Number of headers of GPT
         gpt_hidden_dim: Size of hidden dimensions of GPT
         dropout: Dropout rate for GPT
+        mlp_hidden_dim: Size of hidden dimensions of offset header / bin prediction headers parts of VQ-BeT
         offset_loss_weight:  A constant that is multiplied to the offset loss
         primary_code_loss_weight: A constant that is multiplied to the primary code prediction loss
         secondary_code_loss_weight: A constant that is multiplied to the secondary code prediction loss
@@ -124,6 +125,7 @@ class VQBeTConfig(PreTrainedConfig):
     gpt_n_head: int = 8
     gpt_hidden_dim: int = 512
     dropout: float = 0.1
+    mlp_hidden_dim: int = 1024
     offset_loss_weight: float = 10000.0
     primary_code_loss_weight: float = 5.0
     secondary_code_loss_weight: float = 0.5

src/lerobot/policies/vqbet/modeling_vqbet.py

@@ -27,11 +27,11 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
 
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.policies.vqbet.vqbet_utils import GPT, ResidualVQ
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
 # ruff: noqa: N806
 
@@ -133,7 +133,7 @@ class VQBeTPolicy(PreTrainedPolicy):
             batch.pop(ACTION)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         # NOTE: It's important that this happens after stacking the images into a single key.
-        batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
         if ACTION in batch:
             batch.pop(ACTION)
@@ -340,12 +340,14 @@ class VQBeTModel(nn.Module):
 
     def forward(self, batch: dict[str, Tensor], rollout: bool) -> tuple[dict, dict]:
         # Input validation.
-        assert set(batch).issuperset({OBS_STATE, OBS_IMAGES})
-        batch_size, n_obs_steps = batch[OBS_STATE].shape[:2]
+        assert set(batch).issuperset({"observation.state", "observation.images"})
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
         assert n_obs_steps == self.config.n_obs_steps
 
         # Extract image feature (first combine batch and sequence dims).
-        img_features = self.rgb_encoder(einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ..."))
+        img_features = self.rgb_encoder(
+            einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
+        )
         # Separate batch and sequence dims.
         img_features = einops.rearrange(
             img_features, "(b s n) ... -> b s n ...", b=batch_size, s=n_obs_steps, n=self.num_images
@@ -357,7 +359,9 @@ class VQBeTModel(nn.Module):
             img_features
         )  # (batch, obs_step, number of different cameras, projection dims)
         input_tokens = [rgb_tokens[:, :, i] for i in range(rgb_tokens.size(2))]
-        input_tokens.append(self.state_projector(batch[OBS_STATE]))  # (batch, obs_step, projection dims)
+        input_tokens.append(
+            self.state_projector(batch["observation.state"])
+        )  # (batch, obs_step, projection dims)
         input_tokens.append(einops.repeat(self.action_token, "1 1 d -> b n d", b=batch_size, n=n_obs_steps))
         # Interleave tokens by stacking and rearranging.
         input_tokens = torch.stack(input_tokens, dim=2)

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -19,6 +19,7 @@ from typing import Any
 
 import torch
 
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
@@ -30,7 +31,6 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 def make_vqbet_pre_post_processors(

src/lerobot/policies/vqbet/vqbet_utils.py

@@ -231,6 +231,16 @@ class GPT(nn.Module):
             torch.nn.init.zeros_(module.bias)
             torch.nn.init.ones_(module.weight)
 
+    def crop_block_size(self, gpt_block_size):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert gpt_block_size <= self.config.gpt_block_size
+        self.config.gpt_block_size = gpt_block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:gpt_block_size])
+        for block in self.transformer.h:
+            block.attn.bias = block.attn.bias[:, :, :gpt_block_size, :gpt_block_size]
+
     def configure_parameters(self):
         """
         This long function is unfortunately doing something very simple and is being very defensive:
@@ -260,11 +270,13 @@ class GPT(nn.Module):
         param_dict = dict(self.named_parameters())
         inter_params = decay & no_decay
         union_params = decay | no_decay
-        assert len(inter_params) == 0, (
-            f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
+        assert len(inter_params) == 0, "parameters {} made it into both decay/no_decay sets!".format(
+            str(inter_params)
         )
         assert len(param_dict.keys() - union_params) == 0, (
-            f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"
+            "parameters {} were not separated into either decay/no_decay set!".format(
+                str(param_dict.keys() - union_params),
+            )
         )
 
         decay = [param_dict[pn] for pn in sorted(decay)]

src/lerobot/processor/batch_processor.py

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

src/lerobot/processor/converters.py

@@ -23,8 +23,6 @@ from typing import Any
 import numpy as np
 import torch
 
-from lerobot.utils.constants import ACTION, DONE, OBS_PREFIX, REWARD, TRUNCATED
-
 from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
 
 
@@ -344,20 +342,20 @@ def batch_to_transition(batch: dict[str, Any]) -> EnvTransition:
     if not isinstance(batch, dict):
         raise ValueError(f"EnvTransition must be a dictionary. Got {type(batch).__name__}")
 
-    action = batch.get(ACTION)
+    action = batch.get("action")
     if action is not None and not isinstance(action, PolicyAction):
         raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
 
     # Extract observation and complementary data keys.
-    observation_keys = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
+    observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
     complementary_data = _extract_complementary_data(batch)
 
     return create_transition(
         observation=observation_keys if observation_keys else None,
-        action=batch.get(ACTION),
-        reward=batch.get(REWARD, 0.0),
-        done=batch.get(DONE, False),
-        truncated=batch.get(TRUNCATED, False),
+        action=batch.get("action"),
+        reward=batch.get("next.reward", 0.0),
+        done=batch.get("next.done", False),
+        truncated=batch.get("next.truncated", False),
         info=batch.get("info", {}),
         complementary_data=complementary_data if complementary_data else None,
     )
@@ -379,10 +377,10 @@ def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
         raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
 
     batch = {
-        ACTION: transition.get(TransitionKey.ACTION),
-        REWARD: transition.get(TransitionKey.REWARD, 0.0),
-        DONE: transition.get(TransitionKey.DONE, False),
-        TRUNCATED: transition.get(TransitionKey.TRUNCATED, False),
+        "action": transition.get(TransitionKey.ACTION),
+        "next.reward": transition.get(TransitionKey.REWARD, 0.0),
+        "next.done": transition.get(TransitionKey.DONE, False),
+        "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
         "info": transition.get(TransitionKey.INFO, {}),
     }
 

src/lerobot/processor/delta_action_processor.py

@@ -83,12 +83,14 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
 
     Attributes:
         position_scale: A factor to scale the delta position inputs.
+        rotation_scale: A factor to scale the delta rotation inputs (currently unused).
         noise_threshold: The magnitude below which delta inputs are considered noise
                          and do not trigger an "enabled" state.
     """
 
     # Scale factors for delta movements
     position_scale: float = 1.0
+    rotation_scale: float = 0.0  # No rotation deltas for gamepad/keyboard
     noise_threshold: float = 1e-3  # 1 mm threshold to filter out noise
 
     def action(self, action: RobotAction) -> RobotAction:

src/lerobot/processor/hil_processor.py

@@ -340,7 +340,7 @@ class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
         """
         action = self.transition.get(TransitionKey.ACTION)
 
-        raw_joint_positions = complementary_data.get("raw_joint_positions")
+        raw_joint_positions = complementary_data.get("raw_joint_positions", None)
         if raw_joint_positions is None:
             return complementary_data
 

src/lerobot/processor/joint_observations_processor.py

@@ -20,12 +20,12 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import OBS_STATE
 from lerobot.processor.pipeline import (
     ObservationProcessorStep,
     ProcessorStepRegistry,
 )
 from lerobot.robots import Robot
-from lerobot.utils.constants import OBS_STATE
 
 
 @dataclass

src/lerobot/processor/migrate_policy_normalization.py

@@ -59,7 +59,6 @@ from safetensors.torch import load_file as load_safetensors
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
-from lerobot.utils.constants import ACTION
 
 
 def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
@@ -197,7 +196,7 @@ def detect_features_and_norm_modes(
                 feature_type = FeatureType.VISUAL
             elif "state" in key:
                 feature_type = FeatureType.STATE
-            elif ACTION in key:
+            elif "action" in key:
                 feature_type = FeatureType.ACTION
             else:
                 feature_type = FeatureType.STATE  # Default
@@ -216,7 +215,7 @@ def detect_features_and_norm_modes(
                 feature_type = FeatureType.VISUAL
             elif "state" in key or "joint" in key or "position" in key:
                 feature_type = FeatureType.STATE
-            elif ACTION in key:
+            elif "action" in key:
                 feature_type = FeatureType.ACTION
             else:
                 feature_type = FeatureType.STATE
@@ -322,7 +321,7 @@ def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[
             feature_type = FeatureType.VISUAL
         elif "state" in key:
             feature_type = FeatureType.STATE
-        elif ACTION in key:
+        elif "action" in key:
             feature_type = FeatureType.ACTION
         else:
             feature_type = FeatureType.STATE

src/lerobot/processor/normalize_processor.py

@@ -26,7 +26,6 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.utils.constants import ACTION
 
 from .converters import from_tensor_to_numpy, to_tensor
 from .core import EnvTransition, PolicyAction, TransitionKey
@@ -119,12 +118,13 @@ class _NormalizationMixin:
                     )
                 self.features = reconstructed
 
-        # if keys are strings (JSON), rebuild enum map
-        if self.norm_map and all(isinstance(k, str) for k in self.norm_map):
-            reconstructed = {}
-            for ft_type_str, norm_mode_str in self.norm_map.items():
-                reconstructed[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
-            self.norm_map = reconstructed
+        if self.norm_map:
+            # if keys are strings (JSON), rebuild enum map
+            if all(isinstance(k, str) for k in self.norm_map.keys()):
+                reconstructed = {}
+                for ft_type_str, norm_mode_str in self.norm_map.items():
+                    reconstructed[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
+                self.norm_map = reconstructed
 
         # Convert stats to tensors and move to the target device once during initialization.
         self.stats = self.stats or {}
@@ -272,7 +272,7 @@ class _NormalizationMixin:
         Returns:
             The transformed action tensor.
         """
-        processed_action = self._apply_transform(action, ACTION, FeatureType.ACTION, inverse=inverse)
+        processed_action = self._apply_transform(action, "action", FeatureType.ACTION, inverse=inverse)
         return processed_action
 
     def _apply_transform(

src/lerobot/processor/observation_processor.py

@@ -21,7 +21,7 @@ import torch
 from torch import Tensor
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
+from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
 
@@ -152,7 +152,7 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):
         """
         # Build a new features mapping keyed by the same FeatureType buckets
         # We assume callers already placed features in the correct FeatureType.
-        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = {ft: {} for ft in features}
+        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = {ft: {} for ft in features.keys()}
 
         exact_pairs = {
             "pixels": OBS_IMAGE,
@@ -171,7 +171,7 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):
 
                 # Prefix-based rules (e.g. pixels.cam1 -> OBS_IMAGES.cam1)
                 for old_prefix, new_prefix in prefix_pairs.items():
-                    prefixed_old = f"{OBS_STR}.{old_prefix}"
+                    prefixed_old = f"observation.{old_prefix}"
                     if key.startswith(prefixed_old):
                         suffix = key[len(prefixed_old) :]
                         new_key = f"{new_prefix}{suffix}"
@@ -191,7 +191,7 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):
 
                 # Exact-name rules (pixels, environment_state, agent_pos)
                 for old, new in exact_pairs.items():
-                    if key == old or key == f"{OBS_STR}.{old}":
+                    if key == old or key == f"observation.{old}":
                         new_key = new
                         new_features[src_ft][new_key] = feat
                         handled = True

src/lerobot/processor/pipeline.py

@@ -422,7 +422,7 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
         """
         if save_directory is None:
             # Use default directory in HF_LEROBOT_HOME
-            from lerobot.utils.constants import HF_LEROBOT_HOME
+            from lerobot.constants import HF_LEROBOT_HOME
 
             sanitized_name = re.sub(r"[^a-zA-Z0-9_]", "_", self.name.lower())
             save_directory = HF_LEROBOT_HOME / "processors" / sanitized_name

src/lerobot/processor/policy_robot_bridge.py

@@ -5,7 +5,6 @@ import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.processor import ActionProcessorStep, PolicyAction, ProcessorStepRegistry, RobotAction
-from lerobot.utils.constants import ACTION
 
 
 @dataclass
@@ -24,7 +23,7 @@ class RobotActionToPolicyActionProcessorStep(ActionProcessorStep):
         return asdict(self)
 
     def transform_features(self, features):
-        features[PipelineFeatureType.ACTION][ACTION] = PolicyFeature(
+        features[PipelineFeatureType.ACTION]["action"] = PolicyFeature(
             type=FeatureType.ACTION, shape=(len(self.motor_names),)
         )
         return features

src/lerobot/processor/tokenizer_processor.py

@@ -29,7 +29,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+from lerobot.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
 from lerobot.utils.import_utils import _transformers_available
 
 from .core import EnvTransition, TransitionKey