add: note

Signed-off-by: Jade Choghari <chogharijade@gmail.com>

.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 (1 year). It will be closed if no further activity occurs.
+    recent activity (6 months). 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 (1 year). It will be closed if no further activity occurs.
+    recent activity (6 months). 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
-python -m lerobot.scripts.visualize_dataset \
+lerobot-dataset-viz \
     --repo-id lerobot/pusht \
     --episode-index 0
 ```
@@ -210,7 +210,7 @@ python -m lerobot.scripts.visualize_dataset \
 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
-python -m lerobot.scripts.visualize_dataset \
+lerobot-dataset-viz \
     --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 `python -m lerobot.scripts.visualize_dataset --help` for more instructions.
+Our script can also visualize datasets stored on a distant server. See `lerobot-dataset-viz --help` for more instructions.
 
 ### The `LeRobotDataset` format
 

benchmarks/video/run_video_benchmark.py

@@ -35,12 +35,13 @@ 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.benchmark import TimeBenchmark
+from lerobot.utils.constants import OBS_IMAGE
 
 BASE_ENCODING = OrderedDict(
     [
@@ -117,7 +118,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("observation.image")]
+    img_keys = [key for key in hf_dataset.features if key.startswith(OBS_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/scripts/server/policy_server.py \
+python src/lerobot/async_inference/policy_server.py \
     --host=127.0.0.1 \
     --port=8080 \
 ```
@@ -39,7 +39,7 @@ python src/lerobot/scripts/server/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/scripts/server/robot_client.py \
+python src/lerobot/async_inference/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.scripts.server.configs import PolicyServerConfig
-from lerobot.scripts.server.policy_server import serve
+from lerobot.async_inference.configs import PolicyServerConfig
+from lerobot.async_inference.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/scripts/server/robot_client.py \
+python src/lerobot/async_inference/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/scripts/server/robot_client.py \
 import threading
 from lerobot.robots.so100_follower import SO100FollowerConfig
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-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
+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
 
 # 1. Create the robot instance
 """Check out the cameras available in your setup by running `python lerobot/find_cameras.py`"""

docs/source/hilserl.mdx

@@ -304,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 find_joint_limits.py**
+**Using lerobot-find-joint-limits**
 
 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
-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
+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
 ```
 
 **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](./getting_started_real_world_robot)
+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)
 
 > [!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](./getting_started_real_world_robot)
+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)
 
 ## 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
-python -m lerobot.scripts.visualize_image_transforms \
+lerobot-imgtransform-viz \
   --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
-python src/lerobot/scripts/eval.py \
+lerobot-eval \
   --policy.path="your-policy-id" \
   --env.type=libero \
   --env.task=libero_object \
@@ -52,7 +52,7 @@ python src/lerobot/scripts/eval.py \
 Benchmark a policy across multiple suites at once:
 
 ```bash
-python src/lerobot/scripts/eval.py \
+lerobot-eval \
   --policy.path="your-policy-id" \
   --env.type=libero \
   --env.task=libero_object,libero_spatial \
@@ -103,10 +103,11 @@ For reference, here is the **original dataset** published by Physical Intelligen
 ### Example training command
 
 ```bash
-python src/lerobot/scripts/train.py \
+lerobot-train \
   --policy.type=smolvla \
   --policy.repo_id=${HF_USER}/libero-test \
-  --dataset.repo_id=jadechoghari/smol-libero3 \
+  --policy.load_vlm_weights=true \
+  --dataset.repo_id=HuggingFaceVLA/libero \
   --env.type=libero \
   --env.task=libero_10 \
   --output_dir=./outputs/ \

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](https://huggingface.co/docs/lerobot/getting_started_real_world_robot#record-a-dataset)
+We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](./il_robots)
 
 <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](./getting_started_real_world_robot#record-a-dataset).
+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).
 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](./getting_started_real_world_robot)
+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)
 
 > [!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](./getting_started_real_world_robot)
+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)
 
 > [!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,6 +44,7 @@ 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,
@@ -78,16 +79,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,11 +19,12 @@ 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
@@ -41,8 +42,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, "observation")
+action_features = hw_to_dataset_features(robot.action_features, ACTION)
+obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
 dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
@@ -73,7 +74,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,14 +17,15 @@
 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
@@ -47,8 +48,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, "observation")
+action_features = hw_to_dataset_features(robot.action_features, ACTION)
+obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
 dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
@@ -69,7 +70,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,6 +19,7 @@ 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
 
@@ -34,7 +35,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()
@@ -49,7 +50,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,7 +26,6 @@ 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,
@@ -36,12 +35,13 @@ 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
@@ -143,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,6 +28,7 @@ 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
 
@@ -66,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()
@@ -81,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/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
 
@@ -87,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,7 +27,6 @@ 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,11 +34,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.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
@@ -143,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,6 +29,7 @@ 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 +68,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 +83,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
 
@@ -95,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

@@ -1,66 +0,0 @@
-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

@@ -1,16 +0,0 @@
-# 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,16 +162,19 @@ all = [
 ]
 
 [project.scripts]
-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-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-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]
@@ -264,8 +267,83 @@ 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/async_inference/configs.py

@@ -18,7 +18,8 @@ from dataclasses import dataclass, field
 import torch
 
 from lerobot.robots.config import RobotConfig
-from lerobot.scripts.server.constants import (
+
+from .constants import (
     DEFAULT_FPS,
     DEFAULT_INFERENCE_LATENCY,
     DEFAULT_OBS_QUEUE_TIMEOUT,

src/lerobot/async_inference/helpers.py

@@ -22,12 +22,12 @@ from pathlib import Path
 import torch
 
 from lerobot.configs.types import PolicyFeature
-from lerobot.constants import OBS_IMAGES, OBS_STATE
 from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
 
 # NOTE: Configs need to be loaded for the client to be able to instantiate the policy config
 from lerobot.policies import ACTConfig, DiffusionConfig, PI0Config, SmolVLAConfig, VQBeTConfig  # noqa: F401
 from lerobot.robots.robot import Robot
+from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.utils import init_logging
 
 Action = torch.Tensor
@@ -66,7 +66,7 @@ def validate_robot_cameras_for_policy(
 
 
 def map_robot_keys_to_lerobot_features(robot: Robot) -> dict[str, dict]:
-    return hw_to_dataset_features(robot.observation_features, "observation", use_video=False)
+    return hw_to_dataset_features(robot.observation_features, OBS_STR, use_video=False)
 
 
 def is_image_key(k: str) -> bool:
@@ -141,7 +141,7 @@ def make_lerobot_observation(
     lerobot_features: dict[str, dict],
 ) -> LeRobotObservation:
     """Make a lerobot observation from a raw observation."""
-    return build_dataset_frame(lerobot_features, robot_obs, prefix="observation")
+    return build_dataset_frame(lerobot_features, robot_obs, prefix=OBS_STR)
 
 
 def prepare_raw_observation(

src/lerobot/async_inference/policy_server.py

@@ -15,7 +15,7 @@
 """
 Example:
 ```shell
-python src/lerobot/scripts/server/policy_server.py \
+python src/lerobot/async_inference/policy_server.py \
      --host=127.0.0.1 \
      --port=8080 \
      --fps=30 \
@@ -38,9 +38,15 @@ import grpc
 import torch
 
 from lerobot.policies.factory import get_policy_class
-from lerobot.scripts.server.configs import PolicyServerConfig
-from lerobot.scripts.server.constants import SUPPORTED_POLICIES
-from lerobot.scripts.server.helpers import (
+from lerobot.transport import (
+    services_pb2,  # type: ignore
+    services_pb2_grpc,  # type: ignore
+)
+from lerobot.transport.utils import receive_bytes_in_chunks
+
+from .configs import PolicyServerConfig
+from .constants import SUPPORTED_POLICIES
+from .helpers import (
     FPSTracker,
     Observation,
     RemotePolicyConfig,
@@ -50,11 +56,6 @@ from lerobot.scripts.server.helpers import (
     observations_similar,
     raw_observation_to_observation,
 )
-from lerobot.transport import (
-    services_pb2,  # type: ignore
-    services_pb2_grpc,  # type: ignore
-)
-from lerobot.transport.utils import receive_bytes_in_chunks
 
 
 class PolicyServer(services_pb2_grpc.AsyncInferenceServicer):

src/lerobot/async_inference/robot_client.py

@@ -15,7 +15,7 @@
 """
 Example command:
 ```shell
-python src/lerobot/scripts/server/robot_client.py \
+python src/lerobot/async_inference/robot_client.py \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \
@@ -57,9 +57,15 @@ from lerobot.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.scripts.server.configs import RobotClientConfig
-from lerobot.scripts.server.constants import SUPPORTED_ROBOTS
-from lerobot.scripts.server.helpers import (
+from lerobot.transport import (
+    services_pb2,  # type: ignore
+    services_pb2_grpc,  # type: ignore
+)
+from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks
+
+from .configs import RobotClientConfig
+from .constants import SUPPORTED_ROBOTS
+from .helpers import (
     Action,
     FPSTracker,
     Observation,
@@ -72,11 +78,6 @@ from lerobot.scripts.server.helpers import (
     validate_robot_cameras_for_policy,
     visualize_action_queue_size,
 )
-from lerobot.transport import (
-    services_pb2,  # type: ignore
-    services_pb2_grpc,  # type: ignore
-)
-from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks
 
 
 class RobotClient:

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.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.utils.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.errors import DeviceNotConnectedError
+from lerobot.utils.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.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..configs import ColorMode

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/datasets/compute_stats.py

@@ -174,79 +174,3 @@ 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,6 +27,7 @@ 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)
@@ -54,11 +55,11 @@ def resolve_delta_timestamps(
     """
     delta_timestamps = {}
     for key in ds_meta.features:
-        if key == "next.reward" and cfg.reward_delta_indices is not None:
+        if key == 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("observation.") and cfg.observation_delta_indices is not None:
+        if key.startswith(OBS_PREFIX) 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,8 +31,6 @@ 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 (
@@ -80,14 +78,10 @@ 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__(
@@ -1328,139 +1322,13 @@ 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):
-    # ... keep your existing docstring ...
+    """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`.
+    """
 
     def __init__(
         self,
@@ -1468,253 +1336,99 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         root: str | Path | None = None,
         episodes: dict | None = None,
         image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
+        delta_timestamps: dict[str, 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
+        ]
 
-        # --- 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 = []
+        # 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)
 
-        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.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) ----------------------
+        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])
 
     @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 repo_index_to_id(self):
+        """Return the inverse mapping if repo_id_to_index."""
         return {v: k for k, v in self.repo_id_to_index}
 
     @property
     def fps(self) -> int:
+        """Frames per second used during data collection.
+
+        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["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:
-        """
-        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] = {}
+        features = {}
         for dataset in self._datasets:
-            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]
-
+            features.update({k: v for k, v in dataset.hf_features.items() if k not in self.disabled_features})
         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)):
@@ -1723,6 +1437,12 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
 
     @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):
@@ -1731,14 +1451,21 @@ 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):
@@ -1747,83 +1474,22 @@ 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.")
-        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
+        # 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]
         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
 
@@ -1840,149 +1506,3 @@ 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": {},
-        "observation": {},
+        ACTION: {},
+        OBS_STR: {},
     }
     images_token = OBS_IMAGES.split(".")[-1]
 
@@ -125,17 +125,15 @@ 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["observation"][name] = value
+                processed_features[OBS_STR][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["observation"]:
-        dataset_features.update(
-            hw_to_dataset_features(processed_features["observation"], "observation", use_videos)
-        )
+    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))
 
     return dataset_features

src/lerobot/datasets/streaming_dataset.py

@@ -21,7 +21,6 @@ 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,
@@ -38,6 +37,7 @@ 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):

src/lerobot/datasets/utils.py

@@ -43,6 +43,7 @@ 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
@@ -645,14 +646,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 == "observation":
+    if joint_fts and prefix == OBS_STR:
         features[f"{prefix}.state"] = {
             "dtype": "float32",
             "shape": (len(joint_fts),),
@@ -728,11 +729,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 == "observation.environment_state":
+        elif key == OBS_ENV_STATE:
             type = FeatureType.ENV
-        elif key.startswith("observation"):
+        elif key.startswith(OBS_STR):
             type = FeatureType.STATE
-        elif key.startswith("action"):
+        elif key.startswith(ACTION):
             type = FeatureType.ACTION
         else:
             continue

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -34,6 +34,7 @@ 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
@@ -46,7 +47,6 @@ 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,6 +71,8 @@ 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"
 
@@ -144,6 +146,7 @@ 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()
@@ -185,7 +188,10 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
     num_frames = 0
     paths_to_cat = []
     episodes_metadata = []
-    for ep_path in ep_paths:
+
+    logging.info(f"Converting data files from {len(ep_paths)} episodes")
+
+    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
         ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
         ep_num_frames = get_parquet_num_frames(ep_path)
         ep_metadata = {
@@ -209,7 +215,6 @@ 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)
@@ -236,6 +241,8 @@ 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
@@ -254,7 +261,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 range(num_episodes):
+    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
         # 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]
@@ -281,6 +288,7 @@ 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)
@@ -374,6 +382,8 @@ 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)
 
@@ -405,6 +415,7 @@ 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
@@ -469,6 +480,7 @@ def convert_dataset(
 
 
 if __name__ == "__main__":
+    init_logging()
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--repo-id",

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,
         }
@@ -259,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/utils.py

@@ -26,6 +26,7 @@ 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
 
 
@@ -41,9 +42,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"observation.images.{key}": img for key, img in observations["pixels"].items()}
+            imgs = {f"{OBS_IMAGES}.{key}": img for key, img in observations["pixels"].items()}
         else:
-            imgs = {"observation.image": observations["pixels"]}
+            imgs = {OBS_IMAGE: observations["pixels"]}
 
         for imgkey, img in imgs.items():
             # TODO(aliberts, rcadene): use transforms.ToTensor()?
@@ -72,13 +73,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["observation.environment_state"] = env_state
+        return_observations[OBS_ENV_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["observation.state"] = agent_pos
+    return_observations[OBS_STATE] = agent_pos
 
     return return_observations
 

src/lerobot/motors/dynamixel/dynamixel.py

@@ -22,7 +22,7 @@ import logging
 from copy import deepcopy
 from enum import Enum
 
-from lerobot.utils.encoding_utils import decode_twos_complement, encode_twos_complement
+from lerobot.motors.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.utils.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
+from lerobot.motors.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.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.utils.utils import enter_pressed, move_cursor_up
 
 NameOrID: TypeAlias = str | int

src/lerobot/optim/optimizers.py

@@ -22,11 +22,11 @@ import draccus
 import torch
 from safetensors.torch import load_file, save_file
 
-from lerobot.constants import (
+from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
+from lerobot.utils.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,25 +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."
             )
 
-        if "observation.images" in batch:
-            batch_size = batch["observation.images"][0].shape[0]
+        if OBS_IMAGES in batch:
+            batch_size = batch[OBS_IMAGES][0].shape[0]
         else:
-            batch_size = batch["observation.environment_state"].shape[0]
+            batch_size = batch[OBS_ENV_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["observation.state"])
+                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch[OBS_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)
@@ -430,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["observation.state"].device,
+                device=batch[OBS_STATE].device,
             )
             key_padding_mask = torch.cat(
                 [cls_joint_is_pad, batch["action_is_pad"]], axis=1
@@ -454,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["observation.state"].device
+                batch[OBS_STATE].device
             )
 
         # Prepare transformer encoder inputs.
@@ -462,18 +462,16 @@ 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["observation.state"]))
+            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch[OBS_STATE]))
         # Environment state token.
         if self.config.env_state_feature:
-            encoder_in_tokens.append(
-                self.encoder_env_state_input_proj(batch["observation.environment_state"])
-            )
+            encoder_in_tokens.append(self.encoder_env_state_input_proj(batch[OBS_ENV_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["observation.images"]:
+            for img in batch[OBS_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,7 +17,6 @@ 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,
@@ -29,6 +28,7 @@ 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,7 +33,6 @@ 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 (
@@ -42,6 +41,7 @@ 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 = {
-            "observation.state": deque(maxlen=self.config.n_obs_steps),
-            "action": deque(maxlen=self.config.n_action_steps),
+            OBS_STATE: deque(maxlen=self.config.n_obs_steps),
+            ACTION: deque(maxlen=self.config.n_action_steps),
         }
         if self.config.image_features:
-            self._queues["observation.images"] = deque(maxlen=self.config.n_obs_steps)
+            self._queues[OBS_IMAGES] = deque(maxlen=self.config.n_obs_steps)
         if self.config.env_state_feature:
-            self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)
+            self._queues[OBS_ENV_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["observation.images"], "b s n ... -> n (b s) ...")
+                images_per_camera = einops.rearrange(batch[OBS_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["observation.images"], "b s n ... -> (b s n) ...")
+                    einops.rearrange(batch[OBS_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["observation.state"].shape[:2]
+        batch_size, n_obs_steps = batch[OBS_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({"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 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 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,7 +18,6 @@ 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,
@@ -30,6 +29,7 @@ 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,7 +24,6 @@ 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
@@ -46,6 +45,7 @@ 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,6 +20,7 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
+from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("pi0")
@@ -113,7 +114,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"observation.images.empty_camera_{i}"
+            key = f"{OBS_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,6 +21,7 @@ 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):
@@ -60,26 +61,26 @@ def main():
 
     # Override stats
     dataset_meta = LeRobotDatasetMetadata(dataset_repo_id)
-    dataset_meta.stats["observation.state"]["mean"] = torch.tensor(
+    dataset_meta.stats[OBS_STATE]["mean"] = torch.tensor(
         norm_stats["norm_stats"]["state"]["mean"][:num_motors], dtype=torch.float32
     )
-    dataset_meta.stats["observation.state"]["std"] = torch.tensor(
+    dataset_meta.stats[OBS_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"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[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["task"] = example["prompt"]
 
     if model_name == "pi0_aloha_towel":
-        del batch["observation.images.cam_low"]
+        del batch[f"{OBS_IMAGES}.cam_low"]
     elif model_name == "pi0_aloha_sim":
-        batch["observation.images.top"] = batch["observation.images.cam_high"]
-        del batch["observation.images.cam_high"]
+        batch[f"{OBS_IMAGES}.top"] = batch[f"{OBS_IMAGES}.cam_high"]
+        del batch[f"{OBS_IMAGES}.cam_high"]
 
     # Batchify
     for key in batch:
@@ -116,7 +117,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,7 +19,6 @@ 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,
@@ -35,6 +34,7 @@ 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,6 +6,7 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
+from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("pi0fast")
@@ -99,7 +100,7 @@ class PI0FASTConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
-            key = f"observation.images.empty_camera_{i}"
+            key = f"{OBS_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,7 +18,6 @@ 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,
@@ -30,6 +29,7 @@ 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:
@@ -225,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,6 +31,7 @@ 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
 
@@ -50,7 +51,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)
@@ -157,7 +158,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")
 
@@ -513,17 +514,17 @@ class SACObservationEncoder(nn.Module):
             )
 
     def _init_state_layers(self) -> None:
-        self.has_env = "observation.environment_state" in self.config.input_features
-        self.has_state = "observation.state" in self.config.input_features
+        self.has_env = OBS_ENV_STATE in self.config.input_features
+        self.has_state = OBS_STATE in self.config.input_features
         if self.has_env:
-            dim = self.config.input_features["observation.environment_state"].shape[0]
+            dim = self.config.input_features[OBS_ENV_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["observation.state"].shape[0]
+            dim = self.config.input_features[OBS_STATE].shape[0]
             self.state_encoder = nn.Sequential(
                 nn.Linear(dim, self.config.latent_dim),
                 nn.LayerNorm(self.config.latent_dim),
@@ -549,9 +550,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["observation.environment_state"]))
+            parts.append(self.env_encoder(obs[OBS_ENV_STATE]))
         if self.has_state:
-            parts.append(self.state_encoder(obs["observation.state"]))
+            parts.append(self.state_encoder(obs[OBS_STATE]))
         if parts:
             return torch.cat(parts, dim=-1)
 

src/lerobot/policies/sac/processor_sac.py

@@ -19,7 +19,6 @@ 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,
@@ -31,6 +30,7 @@ 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,6 +19,7 @@ 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")
@@ -69,7 +70,7 @@ class RewardClassifierConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         """Validate feature configurations."""
-        has_image = any(key.startswith("observation.image") for key in self.input_features)
+        has_image = any(key.startswith(OBS_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:
@@ -59,7 +59,9 @@ class SpatialLearnedEmbeddings(nn.Module):
         super().__init__()
         self.height = height
         self.width = width
-        self.channel = channel
+        self.channel = (
+            channel  # TODO(fracapuano): this gives issues with non-square images bc is hardcoded to 4
+        )
         self.num_features = num_features
 
         self.kernel = nn.Parameter(torch.empty(channel, height, width, num_features))

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -20,6 +20,7 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
+from lerobot.utils.constants import OBS_IMAGES
 
 
 @PreTrainedConfig.register_subclass("smolvla")
@@ -117,7 +118,7 @@ class SmolVLAConfig(PreTrainedConfig):
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
-            key = f"observation.images.empty_camera_{i}"
+            key = f"{OBS_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,7 +19,6 @@ 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,
@@ -34,6 +33,7 @@ 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 = {
-            "observation.state": deque(maxlen=1),
-            "action": deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
+            OBS_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["observation.image"] = deque(maxlen=1)
+            self._queues[OBS_IMAGE] = deque(maxlen=1)
         if self.config.env_state_feature:
-            self._queues["observation.environment_state"] = deque(maxlen=1)
+            self._queues[OBS_ENV_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("observation.")}
+        observations = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
 
         # 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["observation.state_is_pad"][0]
+                * ~batch[f"{OBS_STR}.state_is_pad"][0]
                 * ~batch["action_is_pad"]
                 # `z_targets` depends on the next observation.
-                * ~batch["observation.state_is_pad"][1:]
+                * ~batch[f"{OBS_STR}.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["observation.state_is_pad"][0]
+                * ~batch[f"{OBS_STR}.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["observation.state_is_pad"][0]
+                * ~batch[f"{OBS_STR}.state_is_pad"][0]
                 * ~batch["action_is_pad"]
                 # q_targets depends on the reward and the next observations.
                 * ~batch["next.reward_is_pad"]
-                * ~batch["observation.state_is_pad"][1:]
+                * ~batch[f"{OBS_STR}.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["observation.state_is_pad"][0]
+                * ~batch[f"{OBS_STR}.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["observation.state_is_pad"][0]
+            * ~batch[f"{OBS_STR}.state_is_pad"][0]
             * ~batch["action_is_pad"]
         ).mean()
 

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -18,7 +18,6 @@ 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,
@@ -30,6 +29,7 @@ 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/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["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch[OBS_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,14 +340,12 @@ class VQBeTModel(nn.Module):
 
     def forward(self, batch: dict[str, Tensor], rollout: bool) -> tuple[dict, dict]:
         # Input validation.
-        assert set(batch).issuperset({"observation.state", "observation.images"})
-        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
+        assert set(batch).issuperset({OBS_STATE, OBS_IMAGES})
+        batch_size, n_obs_steps = batch[OBS_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["observation.images"], "b s n ... -> (b s n) ...")
-        )
+        img_features = self.rgb_encoder(einops.rearrange(batch[OBS_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
@@ -359,9 +357,7 @@ 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["observation.state"])
-        )  # (batch, obs_step, projection dims)
+        input_tokens.append(self.state_projector(batch[OBS_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,7 +19,6 @@ 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,
@@ -31,6 +30,7 @@ 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/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.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from lerobot.utils.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,6 +23,8 @@ 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
 
 
@@ -342,20 +344,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("observation.")}
+    observation_keys = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
     complementary_data = _extract_complementary_data(batch)
 
     return create_transition(
         observation=observation_keys if observation_keys else None,
-        action=batch.get("action"),
-        reward=batch.get("next.reward", 0.0),
-        done=batch.get("next.done", False),
-        truncated=batch.get("next.truncated", False),
+        action=batch.get(ACTION),
+        reward=batch.get(REWARD, 0.0),
+        done=batch.get(DONE, False),
+        truncated=batch.get(TRUNCATED, False),
         info=batch.get("info", {}),
         complementary_data=complementary_data if complementary_data else None,
     )
@@ -377,10 +379,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),
-        "next.reward": transition.get(TransitionKey.REWARD, 0.0),
-        "next.done": transition.get(TransitionKey.DONE, False),
-        "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
+        ACTION: transition.get(TransitionKey.ACTION),
+        REWARD: transition.get(TransitionKey.REWARD, 0.0),
+        DONE: transition.get(TransitionKey.DONE, False),
+        TRUNCATED: transition.get(TransitionKey.TRUNCATED, False),
         "info": transition.get(TransitionKey.INFO, {}),
     }
 

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,6 +59,7 @@ 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]]:
@@ -196,7 +197,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
@@ -215,7 +216,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
@@ -321,7 +322,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,6 +26,7 @@ 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
@@ -272,7 +273,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.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
 
 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
 
@@ -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"observation.{old_prefix}"
+                    prefixed_old = f"{OBS_STR}.{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"observation.{old}":
+                    if key == old or key == f"{OBS_STR}.{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.constants import HF_LEROBOT_HOME
+            from lerobot.utils.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,6 +5,7 @@ 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
@@ -23,7 +24,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.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
 from lerobot.utils.import_utils import _transformers_available
 
 from .core import EnvTransition, TransitionKey

src/lerobot/rl/actor.py

@@ -35,7 +35,7 @@ gamepad to take control of the robot during training. Initially intervene freque
 reduce interventions as the policy improves.
 
 **WORKFLOW**:
-1. Determine robot workspace bounds using `find_joint_limits.py`
+1. Determine robot workspace bounds using `lerobot-find-joint-limits`
 2. Record demonstrations with `gym_manipulator.py` in record mode
 3. Process the dataset and determine camera crops with `crop_dataset_roi.py`
 4. Start the learner server with the training configuration
@@ -63,6 +63,8 @@ from lerobot.configs.train import TrainRLServerPipelineConfig
 from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
 from lerobot.processor import TransitionKey
+from lerobot.rl.process import ProcessSignalHandler
+from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -75,8 +77,6 @@ from lerobot.transport.utils import (
     send_bytes_in_chunks,
     transitions_to_bytes,
 )
-from lerobot.utils.process import ProcessSignalHandler
-from lerobot.utils.queue import get_last_item_from_queue
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.transition import (

src/lerobot/rl/buffer.py

@@ -24,6 +24,7 @@ import torch.nn.functional as F  # noqa: N812
 from tqdm import tqdm
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, REWARD
 from lerobot.utils.transition import Transition
 
 
@@ -240,7 +241,7 @@ class ReplayBuffer:
         idx = torch.randint(low=0, high=high, size=(batch_size,), device=self.storage_device)
 
         # Identify image keys that need augmentation
-        image_keys = [k for k in self.states if k.startswith("observation.image")] if self.use_drq else []
+        image_keys = [k for k in self.states if k.startswith(OBS_IMAGE)] if self.use_drq else []
 
         # Create batched state and next_state
         batch_state = {}
@@ -466,7 +467,7 @@ class ReplayBuffer:
         if list_transition:
             first_transition = list_transition[0]
             first_state = {k: v.to(device) for k, v in first_transition["state"].items()}
-            first_action = first_transition["action"].to(device)
+            first_action = first_transition[ACTION].to(device)
 
             # Get complementary info if available
             first_complementary_info = None
@@ -491,7 +492,7 @@ class ReplayBuffer:
                 elif isinstance(v, torch.Tensor):
                     data[k] = v.to(storage_device)
 
-            action = data["action"]
+            action = data[ACTION]
 
             replay_buffer.add(
                 state=data["state"],
@@ -529,12 +530,12 @@ class ReplayBuffer:
 
         # Add "action"
         sample_action = self.actions[0]
-        act_info = guess_feature_info(t=sample_action, name="action")
-        features["action"] = act_info
+        act_info = guess_feature_info(t=sample_action, name=ACTION)
+        features[ACTION] = act_info
 
         # Add "reward" and "done"
-        features["next.reward"] = {"dtype": "float32", "shape": (1,)}
-        features["next.done"] = {"dtype": "bool", "shape": (1,)}
+        features[REWARD] = {"dtype": "float32", "shape": (1,)}
+        features[DONE] = {"dtype": "bool", "shape": (1,)}
 
         # Add state keys
         for key in self.states:
@@ -576,9 +577,9 @@ class ReplayBuffer:
                 frame_dict[key] = self.states[key][actual_idx].cpu()
 
             # Fill action, reward, done
-            frame_dict["action"] = self.actions[actual_idx].cpu()
-            frame_dict["next.reward"] = torch.tensor([self.rewards[actual_idx]], dtype=torch.float32).cpu()
-            frame_dict["next.done"] = torch.tensor([self.dones[actual_idx]], dtype=torch.bool).cpu()
+            frame_dict[ACTION] = self.actions[actual_idx].cpu()
+            frame_dict[REWARD] = torch.tensor([self.rewards[actual_idx]], dtype=torch.float32).cpu()
+            frame_dict[DONE] = torch.tensor([self.dones[actual_idx]], dtype=torch.bool).cpu()
             frame_dict["task"] = task_name
 
             # Add complementary_info if available
@@ -647,7 +648,7 @@ class ReplayBuffer:
 
         # Check if the dataset has "next.done" key
         sample = dataset[0]
-        has_done_key = "next.done" in sample
+        has_done_key = DONE in sample
 
         # Check for complementary_info keys
         complementary_info_keys = [key for key in sample if key.startswith("complementary_info.")]
@@ -667,14 +668,14 @@ class ReplayBuffer:
                 current_state[key] = val.unsqueeze(0)  # Add batch dimension
 
             # ----- 2) Action -----
-            action = current_sample["action"].unsqueeze(0)  # Add batch dimension
+            action = current_sample[ACTION].unsqueeze(0)  # Add batch dimension
 
             # ----- 3) Reward and done -----
-            reward = float(current_sample["next.reward"].item())  # ensure float
+            reward = float(current_sample[REWARD].item())  # ensure float
 
             # Determine done flag - use next.done if available, otherwise infer from episode boundaries
             if has_done_key:
-                done = bool(current_sample["next.done"].item())  # ensure bool
+                done = bool(current_sample[DONE].item())  # ensure bool
             else:
                 # If this is the last frame or if next frame is in a different episode, mark as done
                 done = False
@@ -787,8 +788,8 @@ def concatenate_batch_transitions(
     }
 
     # Concatenate basic fields
-    left_batch_transitions["action"] = torch.cat(
-        [left_batch_transitions["action"], right_batch_transition["action"]], dim=0
+    left_batch_transitions[ACTION] = torch.cat(
+        [left_batch_transitions[ACTION], right_batch_transition[ACTION]], dim=0
     )
     left_batch_transitions["reward"] = torch.cat(
         [left_batch_transitions["reward"], right_batch_transition["reward"]], dim=0

src/lerobot/rl/crop_dataset_roi.py

@@ -25,6 +25,7 @@ import torchvision.transforms.functional as F  # type: ignore  # noqa: N812
 from tqdm import tqdm  # type: ignore
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.constants import DONE, REWARD
 
 
 def select_rect_roi(img):
@@ -212,7 +213,7 @@ def convert_lerobot_dataset_to_cropper_lerobot_dataset(
         for key, value in frame.items():
             if key in ("task_index", "timestamp", "episode_index", "frame_index", "index", "task"):
                 continue
-            if key in ("next.done", "next.reward"):
+            if key in (DONE, REWARD):
                 # if not isinstance(value, str) and len(value.shape) == 0:
                 value = value.unsqueeze(0)
 

src/lerobot/rl/gym_manipulator.py

@@ -73,6 +73,7 @@ from lerobot.teleoperators import (
 )
 from lerobot.teleoperators.teleoperator import Teleoperator
 from lerobot.teleoperators.utils import TeleopEvents
+from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say
 
@@ -180,7 +181,7 @@ class RobotEnv(gym.Env):
 
         # Define observation spaces for images and other states.
         if current_observation is not None and "pixels" in current_observation:
-            prefix = "observation.images"
+            prefix = OBS_IMAGES
             observation_spaces = {
                 f"{prefix}.{key}": gym.spaces.Box(
                     low=0, high=255, shape=current_observation["pixels"][key].shape, dtype=np.uint8
@@ -190,7 +191,7 @@ class RobotEnv(gym.Env):
 
         if current_observation is not None:
             agent_pos = current_observation["agent_pos"]
-            observation_spaces["observation.state"] = gym.spaces.Box(
+            observation_spaces[OBS_STATE] = gym.spaces.Box(
                 low=0,
                 high=10,
                 shape=agent_pos.shape,
@@ -600,9 +601,9 @@ def control_loop(
     if cfg.mode == "record":
         action_features = teleop_device.action_features
         features = {
-            "action": action_features,
-            "next.reward": {"dtype": "float32", "shape": (1,), "names": None},
-            "next.done": {"dtype": "bool", "shape": (1,), "names": None},
+            ACTION: action_features,
+            REWARD: {"dtype": "float32", "shape": (1,), "names": None},
+            DONE: {"dtype": "bool", "shape": (1,), "names": None},
         }
         if use_gripper:
             features["complementary_info.discrete_penalty"] = {
@@ -612,7 +613,7 @@ def control_loop(
             }
 
         for key, value in transition[TransitionKey.OBSERVATION].items():
-            if key == "observation.state":
+            if key == OBS_STATE:
                 features[key] = {
                     "dtype": "float32",
                     "shape": value.squeeze(0).shape,
@@ -671,9 +672,9 @@ def control_loop(
             )
             frame = {
                 **observations,
-                "action": action_to_record.cpu(),
-                "next.reward": np.array([transition[TransitionKey.REWARD]], dtype=np.float32),
-                "next.done": np.array([terminated or truncated], dtype=bool),
+                ACTION: action_to_record.cpu(),
+                REWARD: np.array([transition[TransitionKey.REWARD]], dtype=np.float32),
+                DONE: np.array([terminated or truncated], dtype=bool),
             }
             if use_gripper:
                 discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)
@@ -732,7 +733,7 @@ def replay_trajectory(
         download_videos=False,
     )
     episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == cfg.dataset.replay_episode)
-    actions = episode_frames.select_columns("action")
+    actions = episode_frames.select_columns(ACTION)
 
     _, info = env.reset()
 
@@ -740,7 +741,7 @@ def replay_trajectory(
         start_time = time.perf_counter()
         transition = create_transition(
             observation=env.get_raw_joint_positions() if hasattr(env, "get_raw_joint_positions") else {},
-            action=action_data["action"],
+            action=action_data[ACTION],
         )
         transition = action_processor(transition)
         env.step(transition[TransitionKey.ACTION])

src/lerobot/rl/learner.py

@@ -62,16 +62,13 @@ from torch.optim.optimizer import Optimizer
 from lerobot.cameras import opencv  # noqa: F401
 from lerobot.configs import parser
 from lerobot.configs.train import TrainRLServerPipelineConfig
-from lerobot.constants import (
-    CHECKPOINTS_DIR,
-    LAST_CHECKPOINT_LINK,
-    PRETRAINED_MODEL_DIR,
-    TRAINING_STATE_DIR,
-)
 from lerobot.datasets.factory import make_dataset
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
+from lerobot.rl.buffer import ReplayBuffer, concatenate_batch_transitions
+from lerobot.rl.process import ProcessSignalHandler
+from lerobot.rl.wandb_utils import WandBLogger
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -82,8 +79,13 @@ from lerobot.transport.utils import (
     bytes_to_transitions,
     state_to_bytes,
 )
-from lerobot.utils.buffer import ReplayBuffer, concatenate_batch_transitions
-from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.constants import (
+    ACTION,
+    CHECKPOINTS_DIR,
+    LAST_CHECKPOINT_LINK,
+    PRETRAINED_MODEL_DIR,
+    TRAINING_STATE_DIR,
+)
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.train_utils import (
     get_step_checkpoint_dir,
@@ -97,7 +99,6 @@ from lerobot.utils.utils import (
     get_safe_torch_device,
     init_logging,
 )
-from lerobot.utils.wandb_utils import WandBLogger
 
 from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
 
@@ -153,7 +154,7 @@ def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None):
 
     # Setup WandB logging if enabled
     if cfg.wandb.enable and cfg.wandb.project:
-        from lerobot.utils.wandb_utils import WandBLogger
+        from lerobot.rl.wandb_utils import WandBLogger
 
         wandb_logger = WandBLogger(cfg)
     else:
@@ -402,7 +403,7 @@ def add_actor_information_and_train(
                     left_batch_transitions=batch, right_batch_transition=batch_offline
                 )
 
-            actions = batch["action"]
+            actions = batch[ACTION]
             rewards = batch["reward"]
             observations = batch["state"]
             next_observations = batch["next_state"]
@@ -415,7 +416,7 @@ def add_actor_information_and_train(
 
             # Create a batch dictionary with all required elements for the forward method
             forward_batch = {
-                "action": actions,
+                ACTION: actions,
                 "reward": rewards,
                 "state": observations,
                 "next_state": next_observations,
@@ -460,7 +461,7 @@ def add_actor_information_and_train(
                 left_batch_transitions=batch, right_batch_transition=batch_offline
             )
 
-        actions = batch["action"]
+        actions = batch[ACTION]
         rewards = batch["reward"]
         observations = batch["state"]
         next_observations = batch["next_state"]
@@ -474,7 +475,7 @@ def add_actor_information_and_train(
 
         # Create a batch dictionary with all required elements for the forward method
         forward_batch = {
-            "action": actions,
+            ACTION: actions,
             "reward": rewards,
             "state": observations,
             "next_state": next_observations,
@@ -1155,7 +1156,7 @@ def process_transitions(
             # Skip transitions with NaN values
             if check_nan_in_transition(
                 observations=transition["state"],
-                actions=transition["action"],
+                actions=transition[ACTION],
                 next_state=transition["next_state"],
             ):
                 logging.warning("[LEARNER] NaN detected in transition, skipping")

src/lerobot/rl/learner_service.py

@@ -19,9 +19,9 @@ import logging
 import time
 from multiprocessing import Event, Queue
 
+from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.transport import services_pb2, services_pb2_grpc
 from lerobot.transport.utils import receive_bytes_in_chunks, send_bytes_in_chunks
-from lerobot.utils.queue import get_last_item_from_queue
 
 MAX_WORKERS = 3  # Stream parameters, send transitions and interactions
 SHUTDOWN_TIMEOUT = 10

src/lerobot/rl/wandb_utils.py

@@ -23,7 +23,7 @@ from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from termcolor import colored
 
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.constants import PRETRAINED_MODEL_DIR
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR
 
 
 def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[str] | str:

src/lerobot/robots/hope_jr/hope_jr_arm.py

@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any
 
 from lerobot.cameras.utils import make_cameras_from_configs
-from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorNormMode
 from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
 )
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..robot import Robot
 from ..utils import ensure_safe_goal_position

src/lerobot/robots/hope_jr/hope_jr_hand.py

@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any
 
 from lerobot.cameras.utils import make_cameras_from_configs
-from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorNormMode
 from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
 )
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..robot import Robot
 from .config_hope_jr import HopeJrHandConfig

src/lerobot/robots/koch_follower/koch_follower.py

@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any
 
 from lerobot.cameras.utils import make_cameras_from_configs
-from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.dynamixel import (
     DynamixelMotorsBus,
     OperatingMode,
 )
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..robot import Robot
 from ..utils import ensure_safe_goal_position