fix

- Use normalization processor as default example
- Add section on transform features
- Add section on overrides.

.github/PULL_REQUEST_TEMPLATE.md

@@ -30,7 +30,7 @@ pytest -sx tests/test_stuff.py::test_something
 ```
 
 ```bash
-lerobot-train --some.option=true
+python -m lerobot.scripts.train --some.option=true
 ```
 
 ## SECTION TO REMOVE BEFORE SUBMITTING YOUR PR

.github/workflows/nightly.yml

@@ -29,8 +29,8 @@ on:
 env:
   UV_VERSION: "0.8.0"
   PYTHON_VERSION: "3.10"
-  DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-cpu:latest
-  DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-gpu:latest
+  DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-gpu:latest
+  DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-cpu:latest
 
 # Ensures that only the latest commit is built, canceling older runs.
 concurrency:

Makefile

@@ -44,7 +44,7 @@ test-end-to-end:
 	${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval
 
 test-act-ete-train:
-	lerobot-train \
+	python -m lerobot.scripts.train \
 		--policy.type=act \
 		--policy.dim_model=64 \
 		--policy.n_action_steps=20 \
@@ -68,12 +68,12 @@ test-act-ete-train:
 		--output_dir=tests/outputs/act/
 
 test-act-ete-train-resume:
-	lerobot-train \
+	python -m lerobot.scripts.train \
 		--config_path=tests/outputs/act/checkpoints/000002/pretrained_model/train_config.json \
 		--resume=true
 
 test-act-ete-eval:
-	lerobot-eval \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=aloha \
@@ -82,7 +82,7 @@ test-act-ete-eval:
 		--eval.batch_size=1
 
 test-diffusion-ete-train:
-	lerobot-train \
+	python -m lerobot.scripts.train \
 		--policy.type=diffusion \
 		--policy.down_dims='[64,128,256]' \
 		--policy.diffusion_step_embed_dim=32 \
@@ -106,7 +106,7 @@ test-diffusion-ete-train:
 		--output_dir=tests/outputs/diffusion/
 
 test-diffusion-ete-eval:
-	lerobot-eval \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=pusht \
@@ -115,7 +115,7 @@ test-diffusion-ete-eval:
 		--eval.batch_size=1
 
 test-tdmpc-ete-train:
-	lerobot-train \
+	python -m lerobot.scripts.train \
 		--policy.type=tdmpc \
 		--policy.device=$(DEVICE) \
 		--policy.push_to_hub=false \
@@ -137,7 +137,7 @@ test-tdmpc-ete-train:
 		--output_dir=tests/outputs/tdmpc/
 
 test-tdmpc-ete-eval:
-	lerobot-eval \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=xarm \
@@ -148,7 +148,7 @@ test-tdmpc-ete-eval:
 
 
 test-smolvla-ete-train:
-	lerobot-train \
+	python -m lerobot.scripts.train \
 		--policy.type=smolvla \
 		--policy.n_action_steps=20 \
 		--policy.chunk_size=20 \
@@ -171,7 +171,7 @@ test-smolvla-ete-train:
 		--output_dir=tests/outputs/smolvla/
 
 test-smolvla-ete-eval:
-	lerobot-eval \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=aloha \

README.md

@@ -6,7 +6,7 @@
 
 <div align="center">
 
-[![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml?query=branch%3Amain)
+[![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nighty.yml?query=branch%3Amain)
 [![Python versions](https://img.shields.io/pypi/pyversions/lerobot)](https://www.python.org/downloads/)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/huggingface/lerobot/blob/main/LICENSE)
 [![Status](https://img.shields.io/pypi/status/lerobot)](https://pypi.org/project/lerobot/)
@@ -276,7 +276,7 @@ Check out [example 2](https://github.com/huggingface/lerobot/blob/main/examples/
 We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht):
 
 ```bash
-lerobot-eval \
+python -m lerobot.scripts.eval \
     --policy.path=lerobot/diffusion_pusht \
     --env.type=pusht \
     --eval.batch_size=10 \
@@ -288,10 +288,10 @@ lerobot-eval \
 Note: After training your own policy, you can re-evaluate the checkpoints with:
 
 ```bash
-lerobot-eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
+python -m lerobot.scripts.eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
 ```
 
-See `lerobot-eval --help` for more instructions.
+See `python -m lerobot.scripts.eval --help` for more instructions.
 
 ### Train your own policy
 
@@ -303,7 +303,7 @@ A link to the wandb logs for the run will also show up in yellow in your termina
 
 \<img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/wandb.png" alt="WandB logs example"\>
 
-Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `lerobot-eval --help` for more instructions.
+Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `python -m lerobot.scripts.eval --help` for more instructions.
 
 #### Reproduce state-of-the-art (SOTA)
 
@@ -311,7 +311,7 @@ We provide some pretrained policies on our [hub page](https://huggingface.co/ler
 You can reproduce their training by loading the config from their run. Simply running:
 
 ```bash
-lerobot-train --config_path=lerobot/diffusion_pusht
+python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht
 ```
 
 reproduces SOTA results for Diffusion Policy on the PushT task.

docker/Dockerfile.user

@@ -29,7 +29,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
 
 # Install system dependencies and uv (as root)
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential git curl libglib2.0-0 libegl1-mesa-dev ffmpeg \
+    build-essential git curl libglib2.0-0 libegl1-mesa ffmpeg \
     libusb-1.0-0-dev speech-dispatcher libgeos-dev portaudio19-dev \
     && curl -LsSf https://astral.sh/uv/install.sh | sh \
     && mv /root/.local/bin/uv /usr/local/bin/uv \

docs/source/_toctree.yml

@@ -24,9 +24,16 @@
   - local: smolvla
     title: Finetune SmolVLA
   title: "Policies"
+
+- sections:
+  - local: introduction_processors
+    title: Introduction to Robot Processors
+  - local: implement_your_own_processor
+    title: Implement your own processor
+  - local: processors_robots_teleop
+    title: Processors for Robots and Teleoperators
+  title: "Robot Processors"
 - sections:
-  - local: hope_jr
-    title: Hope Jr
   - local: so101
     title: SO-101
   - local: so100
@@ -35,14 +42,16 @@
     title: Koch v1.1
   - local: lekiwi
     title: LeKiwi
-  - local: reachy2
-    title: Reachy 2
+  - local: hope_jr
+    title: Hope Jr
   title: "Robots"
+- sections:
+  - local: phone_teleop
+    title: Phone
+  title: "Teleoperators"
 - sections:
   - local: notebooks
     title: Notebooks
-  - local: feetech
-    title: Updating Feetech Firmware
   title: "Resources"
 - sections:
   - local: contributing

docs/source/cameras.mdx

@@ -9,7 +9,7 @@ To instantiate a camera, you need a camera identifier. This identifier might cha
 To find the camera indices of the cameras plugged into your system, run the following script:
 
 ```bash
-lerobot-find-cameras opencv # or realsense for Intel Realsense cameras
+python -m lerobot.find_cameras opencv # or realsense for Intel Realsense cameras
 ```
 
 The output will look something like this if you have two cameras connected:

docs/source/feetech.mdx

@@ -1,71 +0,0 @@
-# Feetech Motor Firmware Update
-
-This tutorial guides you through updating the firmware of Feetech motors using the official Feetech software.
-
-## Prerequisites
-
-- Windows computer (Feetech software is only available for Windows)
-- Feetech motor control board
-- USB cable to connect the control board to your computer
-- Feetech motors connected to the control board
-
-## Step 1: Download Feetech Software
-
-1. Visit the official Feetech software download page: [https://www.feetechrc.com/software.html](https://www.feetechrc.com/software.html)
-2. Download the latest version of the Feetech debugging software (FD)
-3. Install the software on your Windows computer
-
-## Step 2: Hardware Setup
-
-1. Connect your Feetech motors to the motor control board
-2. Connect the motor control board to your Windows computer via USB cable
-3. Ensure power is supplied to the motors
-
-## Step 3: Configure Connection
-
-1. Launch the Feetech debugging software
-2. Select the correct COM port from the port dropdown menu
-   - If unsure which port to use, check Windows Device Manager under "Ports (COM & LPT)"
-3. Set the appropriate baud rate (typically 1000000 for most Feetech motors)
-4. Click "Open" to establish communication with the control board
-
-## Step 4: Scan for Motors
-
-1. Once connected, click the "Search" button to detect all connected motors
-2. The software will automatically discover and list all motors on the bus
-3. Each motor will appear with its ID number
-
-## Step 5: Update Firmware
-
-For each motor you want to update:
-
-1. **Select the motor** from the list by clicking on it
-2. **Click on Upgrade tab**:
-3. **Click on Online button**:
-   - If an potential firmware update is found, it will be displayed in the box
-4. **Click on Upgrade button**:
-   - The update progress will be displayed
-
-## Step 6: Verify Update
-
-1. After the update completes, the software should automatically refresh the motor information
-2. Verify that the firmware version has been updated to the expected version
-
-## Important Notes
-
-⚠️ **Warning**: Do not disconnect power or USB during firmware updates, it will potentially brick the motor.
-
-## Bonus: Motor Debugging on Linux/macOS
-
-For debugging purposes only, you can use the open-source Feetech Debug Tool:
-
-- **Repository**: [FT_SCServo_Debug_Qt](https://github.com/CarolinePascal/FT_SCServo_Debug_Qt/tree/fix/port-search-timer)
-
-### Installation Instructions
-
-Follow the instructions in the repository to install the tool, for Ubuntu you can directly install it, for MacOS you need to build it from source.
-
-**Limitations:**
-
-- This tool is for debugging and parameter adjustment only
-- Firmware updates must still be done on Windows with official Feetech software

docs/source/hilserl.mdx

@@ -4,13 +4,7 @@ In this tutorial you will go through the full Human-in-the-Loop Sample-Efficient
 
 HIL-SERL is a sample-efficient reinforcement learning algorithm that combines human demonstrations with online learning and human interventions. The approach starts from a small set of human demonstrations, uses them to train a reward classifier, and then employs an actor-learner architecture where humans can intervene during policy execution to guide exploration and correct unsafe behaviors. In this tutorial, you'll use a gamepad to provide interventions and control the robot during the learning process.
 
-It combines three key ingredients:
-
-1. **Offline demonstrations & reward classifier:** a handful of human-teleop episodes plus a vision-based success detector give the policy a shaped starting point.
-
-2. **On-robot actor / learner loop with human interventions:** a distributed Soft Actor Critic (SAC) learner updates the policy while an actor explores on the physical robot; the human can jump in at any time to correct dangerous or unproductive behaviour.
-
-3. **Safety & efficiency tools:** joint/end-effector (EE) bounds, crop region of interest (ROI) preprocessing and WandB monitoring keep the data useful and the hardware safe.
+It combines three key ingredients: 1. **Offline demonstrations & reward classifier:** a handful of human-teleop episodes plus a vision-based success detector give the policy a shaped starting point. 2. **On-robot actor / learner loop with human interventions:** a distributed Soft Actor Critic (SAC) learner updates the policy while an actor explores on the physical robot; the human can jump in at any time to correct dangerous or unproductive behaviour. 3. **Safety & efficiency tools:** joint/end-effector (EE) bounds, crop region of interest (ROI) preprocessing and WandB monitoring keep the data useful and the hardware safe.
 
 Together these elements let HIL-SERL reach near-perfect task success and faster cycle times than imitation-only baselines.
 
@@ -62,243 +56,30 @@ pip install -e ".[hilserl]"
 
 ### Understanding Configuration
 
-The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/scripts/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
+The training process begins with proper configuration for the HILSerl environment. The configuration class of interest is `HILSerlRobotEnvConfig` in `lerobot/envs/configs.py`. Which is defined as:
 
 <!-- prettier-ignore-start -->
 ```python
-class GymManipulatorConfig:
-    env: HILSerlRobotEnvConfig    # Environment configuration (nested)
-    dataset: DatasetConfig    # Dataset recording/replay configuration (nested)
-    mode: str | None = None    # "record", "replay", or None (for training)
-    device: str = "cpu"    # Compute device
-
 class HILSerlRobotEnvConfig(EnvConfig):
     robot: RobotConfig | None = None    # Main robot agent (defined in `lerobot/robots`)
-    teleop: TeleoperatorConfig | None = None    # Teleoperator agent, e.g., gamepad or leader arm
-    processor: HILSerlProcessorConfig    # Processing pipeline configuration (nested)
-    name: str = "real_robot"    # Environment name
-    task: str | None = None    # Task identifier
+    teleop: TeleoperatorConfig | None = None    # Teleoperator agent, e.g., gamepad or leader arm, (defined in `lerobot/teleoperators`)
+    wrapper: EnvTransformConfig | None = None    # Environment wrapper settings; check `lerobot/scripts/server/gym_manipulator.py`
     fps: int = 10    # Control frequency
-
-# Nested processor configuration
-class HILSerlProcessorConfig:
-    control_mode: str = "gamepad"    # Control mode
-    observation: ObservationConfig | None = None    # Observation processing settings
-    image_preprocessing: ImagePreprocessingConfig | None = None    # Image crop/resize settings
-    gripper: GripperConfig | None = None    # Gripper control and penalty settings
-    reset: ResetConfig | None = None    # Environment reset and timing settings
-    inverse_kinematics: InverseKinematicsConfig | None = None    # IK processing settings
-    reward_classifier: RewardClassifierConfig | None = None    # Reward classifier settings
-    max_gripper_pos: float | None = 100.0    # Maximum gripper position
-
-# Sub-configuration classes
-class ObservationConfig:
-    add_joint_velocity_to_observation: bool = False    # Add joint velocities to state
-    add_current_to_observation: bool = False    # Add motor currents to state
-    add_ee_pose_to_observation: bool = False    # Add end-effector pose to state
-    display_cameras: bool = False    # Display camera feeds during execution
-
-class ImagePreprocessingConfig:
-    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None    # Image cropping parameters
-    resize_size: tuple[int, int] | None = None    # Target image size
-
-class GripperConfig:
-    use_gripper: bool = True    # Enable gripper control
-    gripper_penalty: float = 0.0    # Penalty for inappropriate gripper usage
-    gripper_penalty_in_reward: bool = False    # Include gripper penalty in reward
-
-class ResetConfig:
-    fixed_reset_joint_positions: Any | None = None    # Joint positions for reset
-    reset_time_s: float = 5.0    # Time to wait during reset
-    control_time_s: float = 20.0    # Maximum episode duration
-    terminate_on_success: bool = True    # Whether to terminate episodes on success detection
-
-class InverseKinematicsConfig:
-    urdf_path: str | None = None    # Path to robot URDF file
-    target_frame_name: str | None = None    # End-effector frame name
-    end_effector_bounds: dict[str, list[float]] | None = None    # EE workspace bounds
-    end_effector_step_sizes: dict[str, float] | None = None    # EE step sizes per axis
-
-class RewardClassifierConfig:
-    pretrained_path: str | None = None    # Path to pretrained reward classifier
-    success_threshold: float = 0.5    # Success detection threshold
-    success_reward: float = 1.0    # Reward value for successful episodes
-
-# Dataset configuration
-class DatasetConfig:
-    repo_id: str    # LeRobot dataset repository ID
-    task: str    # Task identifier
-    root: str | None = None    # Local dataset root directory
-    num_episodes_to_record: int = 5    # Number of episodes for recording
-    replay_episode: int | None = None    # Episode index for replay
-    push_to_hub: bool = False    # Whether to push datasets to Hub
+    name: str = "real_robot"    # Environment name
+    mode: str = None    # "record", "replay", or None (for training)
+    repo_id: str | None = None    # LeRobot dataset repository ID
+    dataset_root: str | None = None    # Local dataset root (optional)
+    task: str = ""    # Task identifier
+    num_episodes: int = 10    # Number of episodes for recording
+    episode: int = 0    # episode index for replay
+    device: str = "cuda"    # Compute device
+    push_to_hub: bool = True    # Whether to push the recorded datasets to Hub
+    pretrained_policy_name_or_path: str | None = None    # For policy loading
+    reward_classifier_pretrained_path: str | None = None    # For reward model
+    number_of_steps_after_success: int = 0    # For reward classifier, collect more positive examples after a success to train a classifier
 ```
 <!-- prettier-ignore-end -->
 
-### Processor Pipeline Architecture
-
-HIL-SERL uses a modular processor pipeline architecture that processes robot observations and actions through a series of composable steps. The pipeline is divided into two main components:
-
-#### Environment Processor Pipeline
-
-The environment processor (`env_processor`) handles incoming observations and environment state:
-
-1. **VanillaObservationProcessorStep**: Converts raw robot observations into standardized format
-2. **JointVelocityProcessorStep** (optional): Adds joint velocity information to observations
-3. **MotorCurrentProcessorStep** (optional): Adds motor current readings to observations
-4. **ForwardKinematicsJointsToEE** (optional): Computes end-effector pose from joint positions
-5. **ImageCropResizeProcessorStep** (optional): Crops and resizes camera images
-6. **TimeLimitProcessorStep** (optional): Enforces episode time limits
-7. **GripperPenaltyProcessorStep** (optional): Applies penalties for inappropriate gripper usage
-8. **RewardClassifierProcessorStep** (optional): Automated reward detection using vision models
-9. **AddBatchDimensionProcessorStep**: Converts data to batch format for neural network processing
-10. **DeviceProcessorStep**: Moves data to the specified compute device (CPU/GPU)
-
-#### Action Processor Pipeline
-
-The action processor (`action_processor`) handles outgoing actions and human interventions:
-
-1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging
-2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals
-3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing
-4. **InterventionActionProcessorStep**: Handles human interventions and episode termination
-5. **Inverse Kinematics Pipeline** (when enabled):
-   - **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format
-   - **EEReferenceAndDelta**: Computes end-effector reference and delta movements
-   - **EEBoundsAndSafety**: Enforces workspace safety bounds
-   - **InverseKinematicsEEToJoints**: Converts end-effector actions to joint targets
-   - **GripperVelocityToJoint**: Handles gripper control commands
-
-#### Configuration Examples
-
-**Basic Observation Processing**:
-
-```json
-{
-  "env": {
-    "processor": {
-      "observation": {
-        "add_joint_velocity_to_observation": true,
-        "add_current_to_observation": false,
-        "display_cameras": false
-      }
-    }
-  }
-}
-```
-
-**Image Processing**:
-
-```json
-{
-  "env": {
-    "processor": {
-      "image_preprocessing": {
-        "crop_params_dict": {
-          "observation.images.front": [180, 250, 120, 150],
-          "observation.images.side": [180, 207, 180, 200]
-        },
-        "resize_size": [128, 128]
-      }
-    }
-  }
-}
-```
-
-**Inverse Kinematics Setup**:
-
-```json
-{
-  "env": {
-    "processor": {
-      "inverse_kinematics": {
-        "urdf_path": "path/to/robot.urdf",
-        "target_frame_name": "end_effector",
-        "end_effector_bounds": {
-          "min": [0.16, -0.08, 0.03],
-          "max": [0.24, 0.2, 0.1]
-        },
-        "end_effector_step_sizes": {
-          "x": 0.02,
-          "y": 0.02,
-          "z": 0.02
-        }
-      }
-    }
-  }
-}
-```
-
-### Advanced Observation Processing
-
-The HIL-SERL framework supports additional observation processing features that can improve policy learning:
-
-#### Joint Velocity Processing
-
-Enable joint velocity estimation to provide the policy with motion information:
-
-```json
-{
-  "env": {
-    "processor": {
-      "observation": {
-        "add_joint_velocity_to_observation": true
-      }
-    }
-  }
-}
-```
-
-This processor:
-
-- Estimates joint velocities using finite differences between consecutive joint position readings
-- Adds velocity information to the observation state vector
-- Useful for policies that need motion awareness for dynamic tasks
-
-#### Motor Current Processing
-
-Monitor motor currents to detect contact forces and load conditions:
-
-```json
-{
-  "env": {
-    "processor": {
-      "observation": {
-        "add_current_to_observation": true
-      }
-    }
-  }
-}
-```
-
-This processor:
-
-- Reads motor current values from the robot's control system
-- Adds current measurements to the observation state vector
-- Helps detect contact events, object weights, and mechanical resistance
-- Useful for contact-rich manipulation tasks
-
-#### Combined Observation Processing
-
-You can enable multiple observation processing features simultaneously:
-
-```json
-{
-  "env": {
-    "processor": {
-      "observation": {
-        "add_joint_velocity_to_observation": true,
-        "add_current_to_observation": true,
-        "add_ee_pose_to_observation": false,
-        "display_cameras": false
-      }
-    }
-  }
-}
-```
-
-**Note**: Enabling additional observation features increases the state space dimensionality, which may require adjusting your policy network architecture and potentially collecting more training data.
-
 ### Finding Robot Workspace Bounds
 
 Before collecting demonstrations, you need to determine the appropriate operational bounds for your robot.
@@ -349,56 +130,22 @@ With the bounds defined, you can safely collect demonstrations for training. Tra
 
 Create a configuration file for recording demonstrations (or edit an existing one like [env_config_so100.json](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json)):
 
-1. Set `mode` to `"record"` at the root level
-2. Specify a unique `repo_id` for your dataset in the `dataset` section (e.g., "username/task_name")
-3. Set `num_episodes_to_record` in the `dataset` section to the number of demonstrations you want to collect
-4. Set `env.processor.image_preprocessing.crop_params_dict` to `{}` initially (we'll determine crops later)
-5. Configure `env.robot`, `env.teleop`, and other hardware settings in the `env` section
+1. Set `mode` to `"record"`
+2. Specify a unique `repo_id` for your dataset (e.g., "username/task_name")
+3. Set `num_episodes` to the number of demonstrations you want to collect
+4. Set `crop_params_dict` to `null` initially (we'll determine crops later)
+5. Configure `robot`, `cameras`, and other hardware settings
 
 Example configuration section:
 
 ```json
-{
-  "env": {
-    "type": "gym_manipulator",
-    "name": "real_robot",
-    "fps": 10,
-    "processor": {
-      "control_mode": "gamepad",
-      "observation": {
-        "display_cameras": false
-      },
-      "image_preprocessing": {
-        "crop_params_dict": {},
-        "resize_size": [128, 128]
-      },
-      "gripper": {
-        "use_gripper": true,
-        "gripper_penalty": 0.0
-      },
-      "reset": {
-        "reset_time_s": 5.0,
-        "control_time_s": 20.0
-      }
-    },
-    "robot": {
-      // ... robot configuration ...
-    },
-    "teleop": {
-      // ... teleoperator configuration ...
-    }
-  },
-  "dataset": {
-    "repo_id": "username/pick_lift_cube",
-    "root": null,
-    "task": "pick_and_lift",
-    "num_episodes_to_record": 15,
-    "replay_episode": 0,
-    "push_to_hub": true
-  },
-  "mode": "record",
-  "device": "cpu"
-}
+"mode": "record",
+"repo_id": "username/pick_lift_cube",
+"dataset_root": null,
+"task": "pick_and_lift",
+"num_episodes": 15,
+"episode": 0,
+"push_to_hub": true
 ```
 
 ### Using a Teleoperation Device
@@ -444,20 +191,10 @@ The gamepad provides a very convenient way to control the robot and the episode
 To setup the gamepad, you need to set the `control_mode` to `"gamepad"` and define the `teleop` section in the configuration file.
 
 ```json
-{
-  "env": {
     "teleop": {
-      "type": "gamepad",
-      "use_gripper": true
-    },
-    "processor": {
-      "control_mode": "gamepad",
-      "gripper": {
+        "type": "gamepad",
         "use_gripper": true
-      }
-    }
-  }
-}
+    },
 ```
 
 <p align="center">
@@ -479,21 +216,11 @@ The SO101 leader arm has reduced gears that allows it to move and track the foll
 To setup the SO101 leader, you need to set the `control_mode` to `"leader"` and define the `teleop` section in the configuration file.
 
 ```json
-{
-  "env": {
     "teleop": {
-      "type": "so101_leader",
-      "port": "/dev/tty.usbmodem585A0077921",
-      "use_degrees": true
+        "type": "so101_leader",
+        "port": "/dev/tty.usbmodem585A0077921", # check your port number
+        "use_degrees": true
     },
-    "processor": {
-      "control_mode": "leader",
-      "gripper": {
-        "use_gripper": true
-      }
-    }
-  }
-}
 ```
 
 In order to annotate the success/failure of the episode, **you will need** to use a keyboard to press `s` for success, `esc` for failure.
@@ -524,7 +251,7 @@ python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/e
 
 During recording:
 
-1. The robot will reset to the initial position defined in the configuration file `env.processor.reset.fixed_reset_joint_positions`
+1. The robot will reset to the initial position defined in the configuration file `fixed_reset_joint_positions`
 2. Complete the task successfully
 3. The episode ends with a reward of 1 when you press the "success" button
 4. If the time limit is reached, or the fail button is pressed, the episode ends with a reward of 0
@@ -583,19 +310,11 @@ observation.images.front: [180, 250, 120, 150]
 Add these crop parameters to your training configuration:
 
 ```json
-{
-  "env": {
-    "processor": {
-      "image_preprocessing": {
-        "crop_params_dict": {
-          "observation.images.side": [180, 207, 180, 200],
-          "observation.images.front": [180, 250, 120, 150]
-        },
-        "resize_size": [128, 128]
-      }
-    }
-  }
-}
+"crop_params_dict": {
+    "observation.images.side": [180, 207, 180, 200],
+    "observation.images.front": [180, 250, 120, 150]
+},
+"resize_size": [128, 128]
 ```
 
 **Recommended image resolution**
@@ -624,52 +343,26 @@ python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/r
 
 **Key Parameters for Data Collection**
 
-- **mode**: set it to `"record"` to collect a dataset (at root level)
-- **dataset.repo_id**: `"hf_username/dataset_name"`, name of the dataset and repo on the hub
-- **dataset.num_episodes_to_record**: Number of episodes to record
-- **env.processor.reset.terminate_on_success**: Whether to automatically terminate episodes when success is detected (default: `true`)
-- **env.fps**: Number of frames per second to record
-- **dataset.push_to_hub**: Whether to push the dataset to the hub
+- **mode**: set it to `"record"` to collect a dataset
+- **repo_id**: `"hf_username/dataset_name"`, name of the dataset and repo on the hub
+- **num_episodes**: Number of episodes to record
+- **number_of_steps_after_success**: Number of additional frames to record after a success (reward=1) is detected
+- **fps**: Number of frames per second to record
+- **push_to_hub**: Whether to push the dataset to the hub
 
-The `env.processor.reset.terminate_on_success` parameter allows you to control episode termination behavior. When set to `false`, episodes will continue even after success is detected, allowing you to collect more positive examples with the reward=1 label. This is crucial for training reward classifiers as it provides more success state examples in your dataset. When set to `true` (default), episodes terminate immediately upon success detection.
-
-**Important**: For reward classifier training, set `terminate_on_success: false` to collect sufficient positive examples. For regular HIL-SERL training, keep it as `true` to enable automatic episode termination when the task is completed successfully.
+The `number_of_steps_after_success` parameter is crucial as it allows you to collect more positive examples. When a success is detected, the system will continue recording for the specified number of steps while maintaining the reward=1 label. Otherwise, there won't be enough states in the dataset labeled to 1 to train a good classifier.
 
 Example configuration section for data collection:
 
 ```json
 {
-  "env": {
-    "type": "gym_manipulator",
-    "name": "real_robot",
-    "fps": 10,
-    "processor": {
-      "reset": {
-        "reset_time_s": 5.0,
-        "control_time_s": 20.0,
-        "terminate_on_success": false
-      },
-      "gripper": {
-        "use_gripper": true
-      }
-    },
-    "robot": {
-      // ... robot configuration ...
-    },
-    "teleop": {
-      // ... teleoperator configuration ...
-    }
-  },
-  "dataset": {
-    "repo_id": "hf_username/dataset_name",
-    "dataset_root": "data/your_dataset",
-    "task": "reward_classifier_task",
-    "num_episodes_to_record": 20,
-    "replay_episode": null,
-    "push_to_hub": true
-  },
   "mode": "record",
-  "device": "cpu"
+  "repo_id": "hf_username/dataset_name",
+  "dataset_root": "data/your_dataset",
+  "num_episodes": 20,
+  "push_to_hub": true,
+  "fps": 10,
+  "number_of_steps_after_success": 15
 }
 ```
 
@@ -719,7 +412,7 @@ Example configuration for training the [reward classifier](https://huggingface.c
 To train the classifier, use the `train.py` script with your configuration:
 
 ```bash
-lerobot-train --config_path path/to/reward_classifier_train_config.json
+python -m lerobot.scripts.train --config_path path/to/reward_classifier_train_config.json
 ```
 
 **Deploying and Testing the Model**
@@ -728,17 +421,9 @@ To use your trained reward classifier, configure the `HILSerlRobotEnvConfig` to
 
 <!-- prettier-ignore-start -->
 ```python
-config = GymManipulatorConfig(
-    env=HILSerlRobotEnvConfig(
-        processor=HILSerlProcessorConfig(
-            reward_classifier=RewardClassifierConfig(
-                pretrained_path="path_to_your_pretrained_trained_model"
-            )
-        ),
-        # Other environment parameters
-    ),
-    dataset=DatasetConfig(...),
-    mode=None  # For training
+env_config = HILSerlRobotEnvConfig(
+    reward_classifier_pretrained_path="path_to_your_pretrained_trained_model",
+    # Other environment parameters
 )
 ```
 <!-- prettier-ignore-end -->
@@ -747,18 +432,7 @@ or set the argument in the json config file.
 
 ```json
 {
-  "env": {
-    "processor": {
-      "reward_classifier": {
-        "pretrained_path": "path_to_your_pretrained_model",
-        "success_threshold": 0.7,
-        "success_reward": 1.0
-      },
-      "reset": {
-        "terminate_on_success": true
-      }
-    }
-  }
+  "reward_classifier_pretrained_path": "path_to_your_pretrained_model"
 }
 ```
 
@@ -784,7 +458,7 @@ The reward classifier will automatically provide rewards based on the visual inp
 3. **Train the classifier**:
 
    ```bash
-   lerobot-train --config_path src/lerobot/configs/reward_classifier_train_config.json
+   python -m lerobot.scripts.train --config_path src/lerobot/configs/reward_classifier_train_config.json
    ```
 
 4. **Test the classifier**:

docs/source/hilserl_sim.mdx

@@ -32,12 +32,9 @@ To use `gym_hil` with LeRobot, you need to create a configuration file. An examp
 
 ```json
 {
-  "env": {
-    "type": "gym_manipulator",
-    "name": "gym_hil",
-    "task": "PandaPickCubeGamepad-v0",
-    "fps": 10
-  },
+  "type": "hil",
+  "name": "franka_sim",
+  "task": "PandaPickCubeGamepad-v0",
   "device": "cuda"
 }
 ```
@@ -48,40 +45,28 @@ Available tasks:
 - `PandaPickCubeGamepad-v0`: With gamepad control
 - `PandaPickCubeKeyboard-v0`: With keyboard control
 
-### Processor Configuration
+### Gym Wrappers Configuration
 
 ```json
-{
-  "env": {
-    "processor": {
-      "control_mode": "gamepad",
-      "gripper": {
-        "use_gripper": true,
-        "gripper_penalty": -0.02
-      },
-      "reset": {
-        "control_time_s": 15.0,
-        "fixed_reset_joint_positions": [
-          0.0, 0.195, 0.0, -2.43, 0.0, 2.62, 0.785
-        ]
-      },
-      "inverse_kinematics": {
-        "end_effector_step_sizes": {
-          "x": 0.025,
-          "y": 0.025,
-          "z": 0.025
-        }
-      }
+"wrapper": {
+    "gripper_penalty": -0.02,
+    "control_time_s": 15.0,
+    "use_gripper": true,
+    "fixed_reset_joint_positions": [0.0, 0.195, 0.0, -2.43, 0.0, 2.62, 0.785],
+    "end_effector_step_sizes": {
+        "x": 0.025,
+        "y": 0.025,
+        "z": 0.025
+    },
+    "control_mode": "gamepad"
     }
-  }
-}
 ```
 
 Important parameters:
 
-- `gripper.gripper_penalty`: Penalty for excessive gripper movement
-- `gripper.use_gripper`: Whether to enable gripper control
-- `inverse_kinematics.end_effector_step_sizes`: Size of the steps in the x,y,z axes of the end-effector
+- `gripper_penalty`: Penalty for excessive gripper movement
+- `use_gripper`: Whether to enable gripper control
+- `end_effector_step_sizes`: Size of the steps in the x,y,z axes of the end-effector
 - `control_mode`: Set to `"gamepad"` to use a gamepad controller
 
 ## Running with HIL RL of LeRobot
@@ -90,50 +75,39 @@ Important parameters:
 
 To run the environment, set mode to null:
 
-```bash
+<!-- prettier-ignore-start -->
+```python
 python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
+<!-- prettier-ignore-end -->
 
 ### Recording a Dataset
 
 To collect a dataset, set the mode to `record` whilst defining the repo_id and number of episodes to record:
 
-```json
-{
-  "env": {
-    "type": "gym_manipulator",
-    "name": "gym_hil",
-    "task": "PandaPickCubeGamepad-v0"
-  },
-  "dataset": {
-    "repo_id": "username/sim_dataset",
-    "root": null,
-    "task": "pick_cube",
-    "num_episodes_to_record": 10,
-    "replay_episode": null,
-    "push_to_hub": true
-  },
-  "mode": "record"
-}
-```
-
-```bash
+<!-- prettier-ignore-start -->
+```python
 python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
+<!-- prettier-ignore-end -->
 
 ### Training a Policy
 
 To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_gym_hil_env.json) and run the actor and learner servers:
 
-```bash
+<!-- prettier-ignore-start -->
+```python
 python -m lerobot.scripts.rl.actor --config_path path/to/train_gym_hil_env.json
 ```
+<!-- prettier-ignore-end -->
 
 In a different terminal, run the learner server:
 
-```bash
+<!-- prettier-ignore-start -->
+```python
 python -m lerobot.scripts.rl.learner --config_path path/to/train_gym_hil_env.json
 ```
+<!-- prettier-ignore-end -->
 
 The simulation environment provides a safe and repeatable way to develop and test your Human-In-the-Loop reinforcement learning components before deploying to real robots.
 

docs/source/hope_jr.mdx

@@ -19,7 +19,7 @@ pip install -e ".[hopejr]"
 Before starting calibration and operation, you need to identify the USB ports for each HopeJR component. Run this script to find the USB ports for the arm, hand, glove, and exoskeleton:
 
 ```bash
-lerobot-find-port
+python -m lerobot.find_port
 ```
 
 This will display the available USB ports and their associated devices. Make note of the port paths (e.g., `/dev/tty.usbmodem58760433331`, `/dev/tty.usbmodem11301`) as you'll need to specify them in the `--robot.port` and `--teleop.port` parameters when recording data, replaying episodes, or running teleoperation scripts.
@@ -31,7 +31,7 @@ Before performing teleoperation, HopeJR's limbs need to be calibrated. Calibrati
 ### 1.1 Calibrate Robot Hand
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=blue \
@@ -81,7 +81,7 @@ Once you have set the appropriate boundaries for all joints, click "Save" to sav
 ### 1.2 Calibrate Teleoperator Glove
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=homunculus_glove \
     --teleop.port=/dev/tty.usbmodem11201 \
     --teleop.id=red \
@@ -120,7 +120,7 @@ Once calibration is complete, the system will save the calibration to `/Users/yo
 ### 1.3 Calibrate Robot Arm
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --robot.type=hope_jr_arm \
     --robot.port=/dev/tty.usbserial-1110 \
     --robot.id=white
@@ -146,7 +146,7 @@ Use the calibration interface to set the range boundaries for each joint. Move e
 ### 1.4 Calibrate Teleoperator Exoskeleton
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=homunculus_arm \
     --teleop.port=/dev/tty.usbmodem11201 \
     --teleop.id=black
@@ -178,7 +178,7 @@ Due to global variable conflicts in the Feetech middleware, teleoperation for ar
 ### Hand
 
 ```bash
-lerobot-teleoperate \
+python -m lerobot.teleoperate \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=blue \
@@ -194,7 +194,7 @@ lerobot-teleoperate \
 ### Arm
 
 ```bash
-lerobot-teleoperate \
+python -m lerobot.teleoperate \
     --robot.type=hope_jr_arm \
     --robot.port=/dev/tty.usbserial-1110 \
     --robot.id=white \
@@ -214,7 +214,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
 
 ```bash
-lerobot-record \
+python -m lerobot.record \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=right \
@@ -236,7 +236,7 @@ lerobot-record \
 ### Replay
 
 ```bash
-lerobot-replay \
+python -m lerobot.replay \
     --robot.type=hope_jr_hand \
     --robot.port=/dev/tty.usbmodem58760432281 \
     --robot.id=right \
@@ -248,7 +248,7 @@ lerobot-replay \
 ### Train
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
   --dataset.repo_id=nepyope/hand_record_test_with_video_data \
   --policy.type=act \
   --output_dir=outputs/train/hopejr_hand \
@@ -263,7 +263,7 @@ lerobot-train \
 This training run can be viewed as an example [here](https://wandb.ai/tino/lerobot/runs/rp0k8zvw?nw=nwusertino).
 
 ```bash
-lerobot-record \
+python -m lerobot.record \
   --robot.type=hope_jr_hand \
   --robot.port=/dev/tty.usbmodem58760432281 \
   --robot.id=right \

docs/source/il_robots.mdx

@@ -45,7 +45,7 @@ Note that the `id` associated with a robot is used to store the calibration file
 <hfoptions id="teleoperate_so101">
 <hfoption id="Command">
 ```bash
-lerobot-teleoperate \
+python -m lerobot.teleoperate \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=my_awesome_follower_arm \
@@ -101,7 +101,7 @@ With `rerun`, you can teleoperate again while simultaneously visualizing the cam
 <hfoptions id="teleoperate_koch_camera">
 <hfoption id="Command">
 ```bash
-lerobot-teleoperate \
+python -m lerobot.teleoperate \
     --robot.type=koch_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=my_awesome_follower_arm \
@@ -174,7 +174,7 @@ Now you can record a dataset. To record 5 episodes and upload your dataset to th
 <hfoptions id="record">
 <hfoption id="Command">
 ```bash
-lerobot-record \
+python -m lerobot.record \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem585A0076841 \
     --robot.id=my_awesome_follower_arm \
@@ -376,7 +376,7 @@ You can replay the first episode on your robot with either the command below or
 <hfoptions id="replay">
 <hfoption id="Command">
 ```bash
-lerobot-replay \
+python -m lerobot.replay \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=my_awesome_follower_arm \
@@ -428,10 +428,10 @@ Your robot should replicate movements similar to those you recorded. For example
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
   --dataset.repo_id=${HF_USER}/so101_test \
   --policy.type=act \
   --output_dir=outputs/train/act_so101_test \
@@ -453,7 +453,7 @@ Training should take several hours. You will find checkpoints in `outputs/train/
 To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so101_test` policy:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
   --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \
   --resume=true
 ```
@@ -490,7 +490,7 @@ You can use the `record` script from [`lerobot/record.py`](https://github.com/hu
 <hfoptions id="eval">
 <hfoption id="Command">
 ```bash
-lerobot-record  \
+python -m lerobot.record  \
   --robot.type=so100_follower \
   --robot.port=/dev/ttyACM1 \
   --robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \

docs/source/il_sim.mdx

@@ -24,36 +24,11 @@ pip install -e ".[hilserl]"
 
 To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_gym_hil_il.json).
 
-To teleoperate and collect a dataset, we need to modify this config file. Here's an example configuration for imitation learning data collection:
+To teleoperate and collect a dataset, we need to modify this config file and you should add your `repo_id` here: `"repo_id": "il_gym",` and `"num_episodes": 30,` and make sure you set `mode` to `record`, "mode": "record".
 
-```json
-{
-  "env": {
-    "type": "gym_manipulator",
-    "name": "gym_hil",
-    "task": "PandaPickCubeGamepad-v0",
-    "fps": 10
-  },
-  "dataset": {
-    "repo_id": "your_username/il_gym",
-    "root": null,
-    "task": "pick_cube",
-    "num_episodes_to_record": 30,
-    "replay_episode": null,
-    "push_to_hub": true
-  },
-  "mode": "record",
-  "device": "cuda"
-}
-```
+If you do not have a Nvidia GPU also change `"device": "cuda"` parameter in the config file (for example to `mps` for MacOS).
 
-Key configuration points:
-
-- Set your `repo_id` in the `dataset` section: `"repo_id": "your_username/il_gym"`
-- Set `num_episodes_to_record: 30` to collect 30 demonstration episodes
-- Ensure `mode` is set to `"record"`
-- If you don't have an NVIDIA GPU, change `"device": "cuda"` to `"mps"` for macOS or `"cpu"`
-- To use keyboard instead of gamepad, change `"task"` to `"PandaPickCubeKeyboard-v0"`
+By default the config file assumes you use a controller. To use your keyboard please change the envoirment specified at `"task"` in the config file and set it to `"PandaPickCubeKeyboard-v0"`.
 
 Then we can run this command to start:
 
@@ -121,10 +96,10 @@ If you uploaded your dataset to the hub you can [visualize your dataset online](
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
   --dataset.repo_id=${HF_USER}/il_gym \
   --policy.type=act \
   --output_dir=outputs/train/il_sim_test \
@@ -165,32 +140,9 @@ huggingface-cli upload ${HF_USER}/il_sim_test${CKPT} \
 
 ## Evaluate your policy in Sim
 
-To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/eval_config_gym_hil.json).
+To evaluate your policy we have to use the config file that can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/eval_config_gym_hil.json).
 
-Here's an example evaluation configuration:
-
-```json
-{
-  "env": {
-    "type": "gym_manipulator",
-    "name": "gym_hil",
-    "task": "PandaPickCubeGamepad-v0",
-    "fps": 10
-  },
-  "dataset": {
-    "repo_id": "your_username/il_sim_dataset",
-    "dataset_root": null,
-    "task": "pick_cube"
-  },
-  "pretrained_policy_name_or_path": "your_username/il_sim_model",
-  "device": "cuda"
-}
-```
-
-Make sure to replace:
-
-- `repo_id` with the dataset you trained on (e.g., `your_username/il_sim_dataset`)
-- `pretrained_policy_name_or_path` with your model ID (e.g., `your_username/il_sim_model`)
+Make sure to replace the `repo_id` with the dataset you trained on, for example `pepijn223/il_sim_dataset` and replace the `pretrained_policy_name_or_path` with your model id, for example `pepijn223/il_sim_model`
 
 Then you can run this command to visualize your trained policy
 

docs/source/implement_your_own_processor.mdx

@@ -0,0 +1,323 @@
+# Implement your own Robot Processor
+
+In this tutorial, you'll learn how to implement your own Robot Processor.
+It begins by exploring the need for a custom processor, then uses the Normalization processors as the running example to explain how to implement, configure, and serialize a processor. Finally, it lists all helper processors that ship with LeRobot.
+
+## Why would you need a custom processor?
+
+In most cases, when reading raw data from a sensor like the camera and robot motor encoders,
+you will need to process this data to transform it into a format that is compatible to use with the policies in LeRobot.
+For example, raw images are encoded with `uint8` and the values are in the range `[0, 255]`.
+To use these images with the policies, you will need to cast them to `float32` and normalize them to the range `[0, 1]`.
+
+For example, in LeRobot's `VanillaObservationProcessor`, raw images come from the environment as numpy arrays with `uint8` values in range `[0, 255]` and in channel-last format `(H, W, C)`. The processor transforms them into PyTorch tensors with `float32` values in range `[0, 1]` and channel-first format `(C, H, W)`:
+
+```python
+# Input: numpy array with shape (480, 640, 3) and dtype uint8
+raw_image = env_observation["pixels"]  # Values in [0, 255]
+
+# After processing: torch tensor with shape (1, 3, 480, 640) and dtype float32
+processed_image = processor(transition)["observation"]["observation.image"]  # Values in [0, 1]
+```
+
+On the other hand, when a model returns a certain action to be executed on the robot, it is often that one has to post-process this action to make it compatible to run on the robot.
+For example, the model might return joint positions values that range from `[-1, 1]` and one would need to scale them to the ranges of the minimum and maximum joint angle positions of the robot.
+
+In LeRobot, this normalization workflow is handled by the `NormalizerProcessor` (for inputs) and the `UnnormalizerProcessor` (for outputs). These processors are heavily used by policies (e.g., Pi0, SmolVLA) and integrate tightly with the `RobotProcessor`'s `get_config`, `state_dict`, and `load_state_dict` APIs.
+
+For instance, `UnnormalizerProcessor` converts model outputs in `[-1, 1]` back to actual robot joint ranges:
+
+```python
+# Input: model action with normalized values in [-1, 1]
+normalized_action = torch.tensor([-0.5, 0.8, -1.0, 0.2])  # Model output
+
+# After post-processing: real joint positions in robot's native ranges
+# Example: joints range from [-180.0, 180.0]
+real_action = unnormalizer(transition)["action"]
+# real action after post-processing: [ -90.,  144., -180.,   36.]
+```
+
+The unnormalizer uses the dataset statistics to convert back:
+
+```python
+# For MIN_MAX normalization: action = (normalized + 1) * (max - min) / 2 + min
+real_action = (normalized_action + 1) * (max_val - min_val) / 2 + min_val
+```
+
+All these situations point us towards the need for a mechanism to preprocess the data before being passed to the policies and then post-process the action that are returned to be executed on the robot.
+
+To that end, LeRobot provides a pipeline mechanism to implement a sequence of processing steps for the input data and the output action.
+
+## How to implement your own processor?
+
+We'll use the `NormalizerProcessor` as a concrete running example because it is central to most policies and demonstrates configuration and state serialization cleanly.
+
+Prepare the sequence of processing steps necessary for your problem. A processor step is a class that implements the following methods:
+
+- `__call__`: implements the processing step for the input transition.
+- `get_config`: gets the configuration of the processor step.
+- `state_dict`: gets the state of the processor step.
+- `load_state_dict`: loads the state of the processor step.
+- `reset`: resets the state of the processor step.
+- `feature_contract`: displays the modification to the feature space during the processor step.
+
+### Implement the `__call__` method
+
+The `__call__` method is the core of your processor step. It takes an `EnvTransition` and returns a modified `EnvTransition`. Here's how the `NormalizerProcessor` conceptually works (simplified):
+
+```python
+from dataclasses import dataclass
+import torch
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.processor.pipeline import EnvTransition, TransitionKey
+
+@dataclass
+class NormalizerProcessor:
+    features: dict[str, PolicyFeature]
+    norm_map: dict[FeatureType, NormalizationMode]
+    stats: dict[str, dict[str, torch.Tensor]]
+    eps: float = 1e-8
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        normalized_info = {}
+
+        obs = transition.get(TransitionKey.OBSERVATION)
+        act = transition.get(TransitionKey.ACTION)
+
+        new_obs = self._normalize_observation(obs, normalized_info)
+        new_act = self._normalize_action(act, normalized_info)
+
+        new_transition = transition.copy()
+        new_transition[TransitionKey.OBSERVATION] = new_obs
+        new_transition[TransitionKey.ACTION] = new_act
+
+        # Record what was normalized into complementary_data
+        if normalized_info:
+            comp = new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
+            comp = dict(comp)
+            comp["normalized_keys"] = normalized_info
+            new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp
+
+        return new_transition
+```
+
+See the full implementation in `src/lerobot/processor/normalize_processor.py` for details on mean/std and min/max modes and key selection.
+
+**Key principles:**
+
+- Always check if required data exists before processing
+- Return unchanged transition if no processing is needed
+- Use `transition.copy()` to avoid side effects
+- Only modify the specific keys your processor handles
+
+**Tip**: For observation-only processors, you can inherit from `ObservationProcessor` to avoid writing `__call__` boilerplate. The normalizer is mixed (observations and actions), so it implements `__call__` directly.
+
+### Configuration and State Management
+
+Processors support serialization through three methods that separate configuration from tensor state. This is especially important for normalization processors, which carry dataset statistics (tensors) in their state, and hyperparameters in their config:
+
+```python
+from dataclasses import dataclass, field
+from typing import Any
+import torch
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+
+@dataclass
+class NormalizerProcessor:
+    features: dict[str, PolicyFeature]
+    norm_map: dict[FeatureType, NormalizationMode]
+    eps: float = 1e-8
+    _tensor_stats: dict[str, dict[str, torch.Tensor]] = field(default_factory=dict, init=False, repr=False)
+
+    def get_config(self) -> dict[str, Any]:
+        """JSON-serializable configuration (no tensors)."""
+        return {
+            "eps": self.eps,
+            "features": {k: {"type": v.type.value, "shape": v.shape} for k, v in self.features.items()},
+            "norm_map": {ft.value: nm.value for ft, nm in self.norm_map.items()},
+        }
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Tensor state only (e.g., dataset statistics)."""
+        flat: dict[str, torch.Tensor] = {}
+        for key, sub in self._tensor_stats.items():
+            for stat_name, tensor in sub.items():
+                flat[f"{key}.{stat_name}"] = tensor
+        return flat
+
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Restore tensor state at runtime."""
+        self._tensor_stats.clear()
+        for flat_key, tensor in state.items():
+            key, stat_name = flat_key.rsplit(".", 1)
+            self._tensor_stats.setdefault(key, {})[stat_name] = tensor
+```
+
+**Usage:**
+
+```python
+# Save (e.g., inside a policy)
+config = processor.get_config()
+tensors = processor.state_dict()
+
+# Restore (e.g., loading a pretrained policy)
+new_processor = NormalizerProcessor(**config)
+new_processor.load_state_dict(tensors)
+```
+
+### Transform features
+
+The `transform_features` method defines how your processor transforms feature names and shapes. This is crucial for policy configuration and debugging.
+
+Normalization typically preserves the feature keys and shapes, so `NormalizerProcessor.transform_features` returns the input features unchanged. When your processor renames or reshapes, implement this method to reflect the mapping for downstream components. For example, a simple rename processor:
+
+```python
+def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    # Simple renaming
+    if "pixels" in features:
+        features["observation.image"] = features.pop("pixels")
+
+    # Pattern-based renaming
+    for key in list(features.keys()):
+        if key.startswith("env_state."):
+            suffix = key[len("env_state."):]
+            features[f"observation.{suffix}"] = features.pop(key)
+
+    return features
+```
+
+**Key principles:**
+
+- Use `features.pop(old_key)` to remove and get the old feature
+- Use `features[new_key] = old_feature` to add the renamed feature
+- Always return the modified features dictionary
+- Document transformations clearly in the docstring
+
+### Example of usage from the codebase
+
+`transform_features` is used by `RobotProcessor` to derive the dataset/policy feature contract from an initial feature set by applying each step's transformation. You can see concrete examples in the codebase:
+
+- Phone teleoperation record pipeline (`examples/phone_so100_record.py`): processors like `ForwardKinematicsJointsToEE`, `GripperVelocityToJoint`, and `EEBoundsAndSafety` implement `transform_features` to declare which action/observation keys should be materialized in the dataset.
+- SO100 follower kinematics (`src/lerobot/robots/so100_follower/robot_kinematic_processor.py`): each processor's `transform_features` method adds or refines feature keys such as `observation.state.ee.{x,y,z,wx,wy,wz}` or `action.gripper.pos`.
+- Rename and tokenizer processors (`src/lerobot/processor/rename_processor.py`, `src/lerobot/processor/tokenizer_processor.py`): demonstrate key renaming and adding language token features to the contract.
+
+In practice, you will often aggregate features by running `RobotProcessor.transform_features(...)` with your initial features to compute the final contract before recording or training.
+
+## Helper Classes
+
+LeRobot provides pre-built processor classes for common transformations. Below is a comprehensive list of registered processors in the codebase.
+
+### Core processors (observations, actions, normalization)
+
+- **`VanillaObservationProcessor`** (`observation_processor`): Images and state processing to LeRobot format.
+- **`NormalizerProcessor`** (`normalizer_processor`): Normalize observations/actions (mean/std or min/max to [-1, 1]).
+- **`UnnormalizerProcessor`** (`unnormalizer_processor`): Inverse of the normalizer for model outputs.
+- **`DeviceProcessor`** (`device_processor`): Move tensors to a specific device (CPU/GPU) and optional float dtype.
+- **`ToBatchProcessor`** (`to_batch_processor`): Add batch dimension to observations/actions when missing.
+- **`RenameProcessor`** (`rename_processor`): Rename observation keys using a mapping dictionary.
+- **`TokenizerProcessor`** (`tokenizer_processor`): Tokenize language tasks into `observation.language.*` tensors.
+
+### Teleoperation mapping processors
+
+- **`MapDeltaActionToRobotAction`** (`map_delta_action_to_robot_action`): Map teleop deltas (e.g., gamepad) to `action.target_*` fields.
+- **`MapPhoneActionToRobotAction`** (`map_phone_action_to_robot_action`): Map calibrated phone pose/buttons to `action.target_*` and gripper.
+
+### Robot kinematics processors (SO100 follower example)
+
+- **`EEReferenceAndDelta`** (`ee_reference_and_delta`): Compute desired EE pose from target deltas and current pose.
+- **`EEBoundsAndSafety`** (`ee_bounds_and_safety`): Clip EE pose to bounds and check for jumps.
+- **`InverseKinematicsEEToJoints`** (`inverse_kinematics_ee_to_joints`): Convert EE pose to joint targets via IK.
+- **`GripperVelocityToJoint`** (`gripper_velocity_to_joint`): Convert gripper velocity input to joint position command.
+- **`ForwardKinematicsJointsToEE`** (`forward_kinematics_joints_to_ee`): Compute EE pose features from joint positions via FK.
+- **`AddRobotObservationAsComplimentaryData`** (`add_robot_observation`): Read robot observation and insert `raw_joint_positions` into complementary data.
+
+### Policy-specific utility processors
+
+- **`Pi0NewLineProcessor`** (`pi0_new_line_processor`): Ensure text tasks end with a newline (Pi0 tokenizer compatibility).
+- **`SmolVLANewLineProcessor`** (`smolvla_new_line_processor`): Ensure text tasks end with a newline (SmolVLA tokenizer compatibility).
+
+### Usage Example
+
+```python
+from lerobot.processor import NormalizerProcessor, DeviceProcessor, RobotProcessor, ToBatchProcessor
+
+# Create a processing pipeline (typical policy preprocessor)
+steps = [
+    NormalizerProcessor(features=features, norm_map=norm_map, stats=stats),
+    ToBatchProcessor(),
+    DeviceProcessor(device="cuda"),
+]
+
+# Use in RobotProcessor
+processor = RobotProcessor(steps=steps)
+processed_transition = processor(raw_transition)
+```
+
+### Using overrides
+
+You can override step parameters at load-time using `overrides`. This is handy for non-serializable objects or site-specific settings. It works both in policy factories and with `RobotProcessor.from_pretrained(...)`.
+
+Example: during policy evaluation on the robot, override the device and rename map.
+Use this to run a policy trained on CUDA on a CPU-only robot, or to remap camera keys when the robot uses different names than the dataset.
+
+```437:445:src/lerobot/record.py
+preprocessor, postprocessor = make_processor(
+    policy_cfg=cfg.policy,
+    pretrained_path=cfg.policy.pretrained_path,
+    dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map),
+    preprocessor_overrides={
+        "device_processor": {"device": cfg.policy.device},
+        "rename_processor": {"rename_map": cfg.dataset.rename_map},
+    },
+)
+```
+
+Direct usage with `from_pretrained`:
+
+```python
+from lerobot.processor import RobotProcessor
+
+processor = RobotProcessor.from_pretrained(
+    "username/my-processor",
+    overrides={
+        "device_processor": {"device": "cuda:0"},  # registry name for registered steps
+        "CustomStep": {"param": 42},               # class name for non-registered steps
+    },
+)
+```
+
+## Best Practices
+
+- **Keep processors atomic** - One transformation per processor for reusability and debugging
+- **Use dataclasses** - Clean initialization with `@dataclass`
+- **Always register processors** - Use `@ProcessorStepRegistry.register("name")` for discoverability
+- **Check for None** - Always validate required data exists before processing
+- **Use copy() for safety** - Avoid side effects with `transition.copy()`
+- **Separate config and state** - JSON-serializable config vs tensor state_dict
+- **Use base classes** - Inherit from `ObservationProcessor` for observation-only processing
+
+```python
+@ProcessorStepRegistry.register("my_processor")
+@dataclass
+class MyProcessor(ObservationProcessor):
+    threshold: float = 0.5
+
+    def observation(self, observation):
+        if observation is None:
+            return observation
+        # Your processing logic here
+        return processed_observation
+```
+
+## Conclusion
+
+You now have all the tools to implement custom processors in LeRobot! The key steps are:
+
+1. **Define your processor** as a dataclass with the required methods (`__call__`, `get_config`, `state_dict`, `load_state_dict`, `reset`, `feature_contract`)
+2. **Register it** using `@ProcessorStepRegistry.register("name")` for discoverability
+3. **Integrate it** into a `RobotProcessor` pipeline with other processing steps
+4. **Use base classes** like `ObservationProcessor` when possible to reduce boilerplate
+
+The processor system is designed to be modular and composable, allowing you to build complex data processing pipelines from simple, focused components. Whether you're preprocessing sensor data for training or post-processing model outputs for robot execution, custom processors give you the flexibility to handle any data transformation your robotics application requires. Policies like Pi0 and SmolVLA use the same normalization processors described above, so your understanding here will transfer directly when wiring policy preprocessors and postprocessors.
+
+Start simple, test thoroughly, and leverage the existing helper classes to build robust data processing pipelines for your robot learning workflows.

docs/source/introduction_processors.mdx

@@ -0,0 +1,991 @@
+# Introduction to Processors
+
+In robotics, there's a fundamental mismatch between the data that robots and humans produce and what machine learning models expect. This creates several translation challenges:
+
+**Raw Robot Data → Model Input:**
+
+- Robots output raw sensor data (camera images, joint positions, force readings) that need normalization, batching, and device placement before models can process them
+- Language instructions from humans ("pick up the red cube") must be tokenized into numerical representations
+- Different robots use different coordinate systems and units that need standardization
+
+**Model Output → Robot Commands:**
+
+- Models might output end-effector positions, but robots need joint-space commands
+- Teleoperators (like gamepads) produce relative movements (delta positions), but robots expect absolute commands
+- Model predictions are often normalized and need to be converted back to real-world scales
+
+**Cross-Domain Translation:**
+
+- Training data from one robot setup needs adaptation for deployment on different hardware
+- Models trained with specific camera configurations must work with new camera arrangements
+- Datasets with different naming conventions need harmonization
+
+**That's where processors come in.** They serve as the universal translators that bridge these gaps, ensuring seamless data flow from sensors to models to actuators.
+
+Processors are the data transformation backbone of LeRobot. They handle all the preprocessing and postprocessing steps needed to convert raw environment data into model-ready inputs and vice versa. This guide will walk you through everything you need to know about processors - from basic concepts to advanced usage patterns.
+
+## What are Processors?
+
+In robotics, data comes in many forms - images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
+
+- **Normalized**: Scaled to appropriate ranges for neural network processing
+- **Batched**: Organized with proper dimensions for batch processing
+- **Tokenized**: Text converted to numerical representations
+- **Device-placed**: Moved to the right hardware (CPU/GPU)
+- **Type-converted**: Cast to appropriate data types
+
+Processors handle these transformations through composable, reusable steps that can be chained together into pipelines. Think of them as a modular assembly line where each station performs a specific transformation on your data.
+
+## Core Concepts
+
+### EnvTransition: The Universal Data Container
+
+The `EnvTransition` is the fundamental data structure that flows through all processors. It's a typed dictionary that represents a complete robot-environment interaction:
+
+```python
+from lerobot.processor.pipeline import TransitionKey, EnvTransition
+
+# Example transition from a robot collecting data
+transition: EnvTransition = {
+    TransitionKey.OBSERVATION: {
+        "observation.images.camera0": camera0_image_tensor,  # Shape: (H, W, C)
+        "observation.images.camera1": camera1_image_tensor,  # Shape: (H, W, C)
+        "observation.state": joint_positions_tensor,         # Shape: (7,) for 7-DOF arm
+        "observation.environment_state": env_state_tensor    # Shape: (3,) for object position
+    },
+    TransitionKey.ACTION: action_tensor,                    # Shape: (7,) for joint velocities
+    TransitionKey.REWARD: 0.0,                             # Scalar reward signal
+    TransitionKey.DONE: False,                             # Episode termination flag
+    TransitionKey.TRUNCATED: False,                        # Episode truncation flag
+    TransitionKey.INFO: {"success": False},                # Additional metadata
+    TransitionKey.COMPLEMENTARY_DATA: {
+        "task": "pick up the red cube",                    # Language instruction
+        "task_index": 0,                                   # Task identifier
+        "index": 42                                        # Frame index
+    }
+}
+```
+
+Each key in the transition has a specific purpose:
+
+- **OBSERVATION**: All sensor data (images, states, proprioception)
+- **ACTION**: The action to execute or that was executed
+- **REWARD**: Reinforcement learning signal
+- **DONE/TRUNCATED**: Episode boundary indicators
+- **INFO**: Arbitrary metadata
+- **COMPLEMENTARY_DATA**: Task descriptions, indices, padding flags, inter-step data (e.g., you need to compute the velocities and then use this velocity to clip the action)
+
+### ProcessorStep: The Building Block Interface
+
+A `ProcessorStep` is a single transformation unit that processes transitions. It's a protocol (interface) that any processor step must implement:
+
+```python
+from lerobot.processor.pipeline import ProcessorStep, EnvTransition
+from lerobot.configs.types import PolicyFeature
+from typing import Any
+import torch
+
+class MyProcessorStep:
+    """Example processor step interface - all methods must be implemented."""
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """Transform the transition - this is the main processing logic."""
+        raise NotImplementedError
+
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        """Declare how this step transforms feature shapes/types."""
+        raise NotImplementedError
+
+    def get_config(self) -> dict[str, Any]:
+        """Return JSON-serializable configuration for saving/loading."""
+        raise NotImplementedError
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return any learnable parameters (tensors only)."""
+        raise NotImplementedError
+
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Load learnable parameters from saved state."""
+        raise NotImplementedError
+
+    def reset(self) -> None:
+        """Reset any internal state between episodes."""
+        raise NotImplementedError
+```
+
+### RobotProcessor: The Pipeline Orchestrator
+
+The `RobotProcessor` chains multiple `ProcessorStep` instances together, executing them sequentially. It provides automatic format conversion to handle both batch dictionaries (from datasets) and EnvTransition dictionaries:
+
+```python
+from lerobot.processor.pipeline import RobotProcessor, _default_batch_to_transition, _default_transition_to_batch
+
+# Create a processing pipeline
+processor = RobotProcessor(
+    steps=[
+        step1,  # First transformation
+        step2,  # Second transformation
+        step3   # Third transformation
+    ],
+    name="my_preprocessing_pipeline",
+
+    # Optional: Custom converters for input/output formats
+    to_transition=_default_batch_to_transition,  # How to convert batch dict → EnvTransition
+    to_output=_default_transition_to_batch       # How to convert EnvTransition → output format
+)
+
+# The processor automatically handles different input formats:
+# 1. If input is a batch dict (from dataset), converts to EnvTransition
+# 2. Passes through each step sequentially
+# 3. Converts back to original format (or custom output format)
+
+# Example with batch dict input (common in training)
+batch_dict = {"observation.state": tensor, "action": tensor}
+output = processor(batch_dict)  # Automatically converted to/from EnvTransition
+
+# Example with EnvTransition input (common in inference)
+transition = {TransitionKey.OBSERVATION: {...}, TransitionKey.ACTION: ...}
+output = processor(transition)  # Stays as EnvTransition throughout
+```
+
+The `to_transition` and `to_output` converters enable seamless integration with existing codebases.
+By default, they handle the standard LeRobot batch format, but you can customize them for different data structures.
+
+### Additional Converter Functions
+
+LeRobot provides several specialized converter functions for common robotics scenarios:
+
+```python
+from lerobot.processor.converters import (
+    to_transition_teleop_action,
+    to_transition_robot_observation,
+    to_output_robot_action,
+    to_dataset_frame
+)
+```
+
+**`to_transition_teleop_action`** - Converts teleoperation device actions to EnvTransitions:
+
+```python
+# Use case: Phone, gamepad, or other teleop device control
+phone_action = {"x": 0.1, "y": -0.2, "gripper": 0.8}
+transition = to_transition_teleop_action(phone_action)
+# Creates: {ACTION: {"action.x": 0.1, "action.y": -0.2, "action.gripper": 0.8}, ...}
+```
+
+**`to_transition_robot_observation`** - Converts robot sensor data to EnvTransitions:
+
+```python
+# Use case: Live robot observation during inference
+robot_obs = {
+    "joint_1": 0.5, "joint_2": -0.3,  # joint positions
+    "camera_0": image_array            # camera images
+}
+transition = to_transition_robot_observation(robot_obs)
+# Creates: {OBSERVATION: {"observation.state.joint_1": 0.5, "observation.images.camera_0": image, ...}}
+```
+
+**`to_output_robot_action`** - Extracts robot-executable actions from EnvTransitions:
+
+```python
+# Use case: Converting model outputs back to robot commands
+model_transition = {ACTION: {"action.joint_1": 0.2, "action.joint_2": 0.1}}
+robot_action = to_output_robot_action(model_transition)
+# Returns: {"joint_1": 0.2, "joint_2": 0.1} - ready for robot.send_action()
+```
+
+**`to_dataset_frame`** - Converts transitions to dataset-compatible format:
+
+```python
+# Use case: Saving processed data or creating training batches
+features = {
+    "action": {"names": ["joint_1", "joint_2"]},
+    "observation.state": {"names": ["joint_1", "joint_2"]},
+    "observation.images.camera0": {...}
+}
+batch = to_dataset_frame(transition, features)
+# Returns: {"action": [0.2, 0.1], "observation.state": [0.5, -0.3], ...}
+```
+
+These converters are particularly useful when integrating with real robots, as shown in the examples:
+
+```python
+# Example from phone_so100_teleop.py - Real robot teleoperation
+phone_to_robot_ee_pose = RobotProcessor(
+    steps=[...],
+    to_transition=to_transition_teleop_action,  # Phone → EnvTransition
+    to_output=lambda tr: tr                     # Keep as EnvTransition
+)
+
+# Example from phone_so100_eval.py - Robot action execution
+robot_ee_to_joints = RobotProcessor(
+    steps=[...],
+    to_transition=lambda tr: tr,                # Already EnvTransition
+    to_output=to_output_robot_action           # EnvTransition → Robot action
+)
+
+# Example from phone_so100_record.py - Dataset recording
+robot_joints_to_ee_pose = RobotProcessor(
+    steps=[...],
+    to_transition=to_transition_robot_observation,  # Robot obs → EnvTransition
+    to_output=lambda tr: tr                         # Keep as EnvTransition for dataset
+)
+```
+
+### Data Format Conversion
+
+Different data sources have different formats, but processors need a unified `EnvTransition` structure internally.
+The default converters handle LeRobot datasets, but you can customize them:
+
+```python
+# Default: LeRobot batch format
+lerobot_batch = {
+    "observation.state": torch.tensor(...),
+    "action": torch.tensor(...),
+    "next.reward": torch.tensor(...),
+    "task": ["pick cube", ...]
+}
+# → Converts to EnvTransition → Processes → Converts back
+
+# Custom: Live robot data
+robot_data = {
+    "cameras": {"wrist_cam": np.array(...)},
+    "joint_positions": np.array(...),
+    "gripper_state": 0.5
+}
+
+def robot_to_transition(data: dict) -> EnvTransition:
+    return {
+        TransitionKey.OBSERVATION: {
+            "observation.images.wrist": torch.from_numpy(data["cameras"]["wrist_cam"]),
+            "observation.state": torch.from_numpy(data["joint_positions"])
+        },
+        TransitionKey.ACTION: None,
+        # ... other fields with defaults
+    }
+
+# Use custom converter
+processor = RobotProcessor(
+    steps=[...],
+    to_transition=robot_to_transition,
+    to_output=lambda transition: transition  # Keep as EnvTransition
+)
+```
+
+**When to customize:** Live robot data, Gymnasium environments, legacy datasets, or any non-LeRobot format.
+
+## Common Processor Steps
+
+LeRobot provides a rich set of pre-built processor steps for common transformations.
+Let's explore each in detail:
+
+### Data Normalization
+
+Normalization is crucial for neural network training and inference.
+The `NormalizerProcessor` handles both mean-std normalization and min-max scaling:
+
+```python
+from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor
+from lerobot.configs.types import PolicyFeature, FeatureType, NormalizationMode
+
+# Define what features exist in your data
+features = {
+    "observation.images.camera0": PolicyFeature(
+        type=FeatureType.IMAGE,
+        shape=(224, 224, 3)
+    ),
+    "observation.state": PolicyFeature(
+        type=FeatureType.STATE,
+        shape=(7,)
+    ),
+    "action": PolicyFeature(
+        type=FeatureType.ACTION,
+        shape=(7,)
+    )
+}
+
+# Define normalization strategy per feature type
+norm_map = {
+    FeatureType.IMAGE: NormalizationMode.MEAN_STD,    # Images: (x - mean) / std
+    FeatureType.STATE: NormalizationMode.MIN_MAX,     # States: scale to [-1, 1]
+    FeatureType.ACTION: NormalizationMode.MIN_MAX     # Actions: scale to [-1, 1]
+}
+
+# Create normalizer with dataset statistics
+normalizer = NormalizerProcessor(
+    features=features,
+    norm_map=norm_map,
+    stats=dataset.meta.stats,  # Contains mean, std, min, max per feature
+    normalize_keys={"observation.state", "action"}  # Optional: only normalize specific keys
+)
+
+# For postprocessing: inverse transformation
+unnormalizer = UnnormalizerProcessor(
+    features=features,
+    norm_map=norm_map,
+    stats=dataset.meta.stats
+)
+
+# The normalizer automatically:
+# - Detects which normalization to apply based on feature type
+# - Handles device placement of statistics tensors
+# - Skips keys not in stats or not in normalize_keys
+# - Adds metadata about what was normalized
+```
+
+### Device Management
+
+The `DeviceProcessor` ensures tensors are on the right device with the right dtype:
+
+```python
+from lerobot.processor.device_processor import DeviceProcessor
+
+# Basic GPU placement
+gpu_processor = DeviceProcessor(device="cuda:0")
+
+# Advanced: GPU with half-precision for inference
+efficient_processor = DeviceProcessor(
+    device="cuda:0",
+    float_dtype="float16"  # Convert float32 -> float16 for memory efficiency
+)
+
+# The processor:
+# - Moves all tensors to specified device
+# - Preserves non-tensor data unchanged
+# - Optionally converts float dtypes while preserving int/bool types
+# - Uses non_blocking transfers for CUDA devices
+# - Handles nested structures (observations, complementary_data)
+
+# Supported float dtypes:
+# "float16" / "half": 16-bit floating point
+# "float32" / "float": 32-bit floating point (default)
+# "float64" / "double": 64-bit floating point
+# "bfloat16": Brain floating point (better for training)
+```
+
+### Batch Processing
+
+Models expect batched inputs, but robot interactions often produce unbatched data:
+
+```python
+from lerobot.processor.batch_processor import ToBatchProcessor
+
+batch_processor = ToBatchProcessor()
+
+# Automatically adds batch dimensions where needed:
+# State: (7,) -> (1, 7)
+# Image: (224, 224, 3) -> (1, 224, 224, 3)
+# Action: (4,) -> (1, 4)
+# Task: "pick_cube" -> ["pick_cube"]
+# Already batched: (1, 7) -> (1, 7) [unchanged]
+
+# The processor intelligently:
+# - Detects tensor dimensionality
+# - Adds batch dim to 1D states/actions
+# - Adds batch dim to 3D images
+# - Wraps string tasks in lists
+# - Preserves already-batched data
+
+# Example usage in inference:
+single_observation = robot.get_observation()  # Unbatched
+batched_input = batch_processor({"observation": single_observation})
+model_output = model(batched_input)  # Model expects batch dim
+```
+
+### Text Tokenization
+
+For language-conditioned policies, text instructions must be tokenized:
+
+```python
+from lerobot.processor.tokenizer_processor import TokenizerProcessor
+from transformers import AutoTokenizer
+
+# Option 1: Auto-load tokenizer by name
+tokenizer_proc = TokenizerProcessor(
+    tokenizer_name="google/paligemma-3b-pt-224",
+    max_length=128,
+    task_key="task",           # Where to find text in complementary_data
+    padding="max_length",       # Pad to max_length
+    padding_side="right",
+    truncation=True            # Truncate if longer than max_length
+)
+
+# Option 2: Provide custom tokenizer
+custom_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+custom_proc = TokenizerProcessor(
+    tokenizer=custom_tokenizer,
+    max_length=256,
+    padding_side="left"        # For autoregressive models
+)
+
+# The processor:
+# - Extracts task text from complementary_data
+# - Tokenizes using HuggingFace tokenizer
+# - Adds tokens and attention_mask to observations
+# - Handles both single strings and lists of strings
+# - Preserves original task in complementary_data
+
+# Output structure:
+# observation["observation.language.tokens"] = tensor([101, 2032, ...])
+# observation["observation.language.attention_mask"] = tensor([1, 1, 0, ...])
+```
+
+### Key Renaming
+
+Different datasets and models may use different naming conventions.
+The `RenameProcessor` solves this mismatch:
+
+**Why is this useful?**
+
+- When loading a model trained on a different dataset with different key names
+- When using foundation models that expect specific key naming conventions
+- When standardizing datasets from different sources
+- When adapting legacy code to new naming standards
+
+```python
+from lerobot.processor.rename_processor import RenameProcessor
+
+# Example 1: Dataset uses "top"/"wrist", model expects "camera0"/"camera1"
+rename_proc = RenameProcessor(
+    rename_map={
+        "observation.images.top": "observation.images.camera0",
+        "observation.images.wrist": "observation.images.camera1",
+    }
+)
+
+# Example 2: Foundation model compatibility
+# Your dataset: "observation.state", Foundation model: "proprio"
+foundation_rename = RenameProcessor(
+    rename_map={
+        "observation.state": "proprio",
+        "observation.images.main": "rgb",
+    }
+)
+
+# Example 3: Standardizing multiple datasets
+standardize_rename = RenameProcessor(
+    rename_map={
+        # Different robots might use different names
+        "observation.joint_positions": "observation.state",
+        "observation.gripper_state": "observation.end_effector",
+        "observation.arm_camera": "observation.images.wrist",
+    }
+)
+```
+
+## Building Complete Pipelines
+
+Let's build a real-world preprocessing and postprocessing pipeline for a vision-based
+manipulation policy:
+
+```python
+# Consolidated imports
+from lerobot.processor import (
+    RobotProcessor,
+    NormalizerProcessor,
+    UnnormalizerProcessor,
+    DeviceProcessor,
+    ToBatchProcessor,
+    TokenizerProcessor,
+    RenameProcessor
+)
+
+# Step 1: Define the preprocessing pipeline
+preprocessor = RobotProcessor(
+    steps=[
+        # 1. Standardize naming from dataset
+        RenameProcessor(
+            rename_map={
+                "observation.images.top": "observation.images.camera0",
+                "observation.images.wrist": "observation.images.camera1"
+            }
+        ),
+
+        # 2. Add batch dimensions for model
+        ToBatchProcessor(),
+
+        # 3. Tokenize language instructions if present
+        TokenizerProcessor(
+            tokenizer_name="google/paligemma-3b-pt-224",
+            max_length=64,
+            task_key="task"
+        ),
+
+        # 4. Normalize numerical data
+        NormalizerProcessor(
+            features=policy_features,
+            norm_map={
+                FeatureType.IMAGE: NormalizationMode.MEAN_STD,
+                FeatureType.STATE: NormalizationMode.MIN_MAX,
+                FeatureType.ACTION: NormalizationMode.MIN_MAX
+            },
+            stats=dataset.meta.stats
+        ),
+
+        # 5. Move to GPU and convert to half precision
+        DeviceProcessor(
+            device="cuda:0",
+            float_dtype="float16"
+        )
+    ],
+    name="robot_preprocessor"
+)
+
+# Step 2: Define the postprocessing pipeline
+postprocessor = RobotProcessor(
+    steps=[
+        # 1. Move back to CPU for robot hardware
+        DeviceProcessor(device="cpu"),
+
+        # 2. Denormalize actions to original scale
+        UnnormalizerProcessor(
+            features=policy_features,
+            norm_map={
+                FeatureType.ACTION: NormalizationMode.MIN_MAX
+            },
+            stats=dataset.meta.stats
+        )
+    ],
+    name="robot_postprocessor"
+)
+```
+
+## Using Processors in Practice
+
+### Training Loop Integration
+
+Here's how processors integrate into a training loop using the policy's forward method:
+
+```python
+from torch.utils.data import DataLoader
+
+# Create dataset and dataloader
+dataset = LeRobotDataset(repo_id="your_dataset")
+dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+# Initialize model and processors
+model = YourPolicy.from_pretrained("your_model")
+preprocessor = RobotProcessor.from_pretrained(
+    "your_model",
+    config_filename="robot_preprocessor.json"
+)
+
+# Training loop
+for epoch in range(num_epochs):
+    for batch in dataloader:
+        # Preprocess batch
+        processed_batch = preprocessor(batch)
+
+        # Forward pass - returns loss and optional metrics
+        loss, metrics = model.forward(processed_batch)
+
+        # Backward pass
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Log metrics if available
+        if metrics:
+            wandb.log(metrics)
+```
+
+### Inference Pipeline
+
+For deployment, processors ensure consistent data handling with real robots:
+
+```python
+# Load model and processors
+policy = YourPolicy.from_pretrained("path/to/model")
+preprocessor = RobotProcessor.from_pretrained(
+    "path/to/model",
+    config_filename="robot_preprocessor.json"
+)
+postprocessor = RobotProcessor.from_pretrained(
+    "path/to/model",
+    config_filename="robot_postprocessor.json"
+)
+
+# Connect to robot
+robot = make_robot_from_config(robot_config)
+robot.connect()
+
+# Inference loop
+policy.eval()
+# Reset the policy and processors
+policy.reset()
+preprocessor.reset()
+postprocessor.reset()
+
+with torch.no_grad():
+    while not done:
+        # Get observation from robot
+        observation = robot.get_observation()
+
+        # Build dataset-compatible frame
+        observation_frame = build_dataset_frame(
+            dataset.features,
+            observation,
+            prefix="observation"
+        )
+
+        # Add task instruction to complementary data
+        observation_frame["task"] = "pick up the red cube"
+
+        # Preprocess for model
+        model_input = preprocessor(observation_frame)
+
+        # Run policy
+        raw_action = policy.select_action(model_input)
+
+        # Postprocess action
+        action_transition = {TransitionKey.ACTION: raw_action}
+        processed = postprocessor(action_transition)
+        action = processed[TransitionKey.ACTION]
+
+        # Convert to robot action format
+        robot_action = {
+            key: action[i].item()
+            for i, key in enumerate(robot.action_features)
+        }
+
+        # Execute on robot
+        robot.send_action(robot_action)
+```
+
+## Saving and Loading Processors
+
+Processors can be persisted and shared just like models, making them portable across different
+environments and ensuring reproducibility:
+
+### Local Save/Load
+
+```python
+# Save processor configuration and state
+preprocessor.save_pretrained(
+    "./my_robot_processor",
+    config_filename="preprocessor.json"  # Optional custom name
+)
+
+# The save creates:
+# my_robot_processor/
+# ├── preprocessor.json                    # Configuration
+# ├── preprocessor_step_0_normalizer.safetensors  # Step 0 state (stats)
+# └── preprocessor_step_1_device.safetensors     # Step 1 state (if any)
+
+# Load processor
+loaded = RobotProcessor.from_pretrained(
+    "./my_robot_processor",
+    config_filename="preprocessor.json"
+)
+```
+
+### HuggingFace Hub Integration
+
+The HuggingFace Hub provides a centralized place to share and version your processors.
+This is particularly useful for sharing preprocessing configurations with models,
+ensuring that anyone who downloads your model can reproduce your exact preprocessing pipeline.
+It also enables versioning and collaboration on preprocessing strategies.
+
+```python
+# Save to HuggingFace Hub
+preprocessor.save_pretrained("username/my-robot-policy")
+
+# Load from Hub with automatic download
+hub_processor = RobotProcessor.from_pretrained(
+    "username/my-robot-policy",
+    config_filename="robot_preprocessor.json",
+    revision="main",  # Optional: specific revision
+    cache_dir="./cache"  # Optional: local cache directory
+)
+
+# The Hub integration provides:
+# - Automatic versioning with git
+# - Public or private sharing
+# - Download caching for efficiency
+# - Integration with model repositories
+```
+
+### Loading with Overrides
+
+Sometimes you need to modify loaded processors for new environments or datasets.
+The override mechanism allows you to update specific processor configurations without modifying
+the saved files:
+
+```python
+# Load processor with configuration overrides
+processor = RobotProcessor.from_pretrained(
+    "./saved_processor",
+    overrides={
+        # Change device for different hardware
+        "device_processor": {"device": "cuda:1"},
+
+        # Update statistics for new dataset
+        "normalizer_processor": {"stats": new_dataset.meta.stats},
+
+        # Provide non-serializable objects (like tokenizers)
+        "tokenizer_processor": {"tokenizer": custom_tokenizer}
+    }
+)
+
+# Common override scenarios:
+# 1. Adapting to different hardware (GPU availability)
+# 2. Fine-tuning on new datasets with different statistics
+# 3. Providing runtime dependencies that can't be serialized
+# 4. Testing variations without creating new saved configs
+```
+
+## Creating Custom Processor Steps
+
+Build your own processor steps for specialized transformations.
+The key is implementing the required interface:
+
+### Basic Custom Step with Registration
+
+The registration mechanism allows your custom processors to be saved and loaded by name rather
+than by module path.
+This makes them more portable and easier to share:
+
+```python
+from dataclasses import dataclass
+from lerobot.processor.pipeline import ProcessorStepRegistry, ObservationProcessor
+
+# The @register decorator adds your processor to the global registry
+# Use a unique name, preferably namespaced to avoid conflicts
+@dataclass
+@ProcessorStepRegistry.register("my_company/gaussian_noise")
+class GaussianNoiseProcessor(ObservationProcessor):
+    """Add Gaussian noise to observations for robustness training."""
+
+    noise_std: float = 0.01
+    training_only: bool = True
+    is_training: bool = True
+
+    def observation(self, observation):
+        """Add noise to observation tensors."""
+        if not self.is_training and self.training_only:
+            return observation
+
+        noisy_obs = {}
+        for key, value in observation.items():
+            if isinstance(value, torch.Tensor) and "image" not in key:
+                # Add noise to non-image observations
+                noise = torch.randn_like(value) * self.noise_std
+                noisy_obs[key] = value + noise
+            else:
+                noisy_obs[key] = value
+
+        return noisy_obs
+
+    def get_config(self):
+        return {
+            "noise_std": self.noise_std,
+            "training_only": self.training_only,
+            "is_training": self.is_training
+        }
+
+# Why register?
+# 1. Enables saving by name: config saves "my_company/gaussian_noise" instead of full module path
+# 2. More portable: Others can use your processor without your exact module structure
+# 3. Version-safe: Module refactoring won't break saved configs
+# 4. Cleaner configs: JSON shows readable names instead of long import paths
+```
+
+### Using Base Classes for Common Patterns
+
+LeRobot provides base classes like `ObservationProcessor`, `ActionProcessor`, etc., that handle
+the boilerplate of extracting and reinserting specific components:
+
+```python
+from lerobot.processor import ActionProcessor
+
+@dataclass
+@ProcessorStepRegistry.register("my_company/action_clipper")
+class ActionClipProcessor(ActionProcessor):
+    """Clip actions to safe ranges."""
+
+    min_value: float = -1.0
+    max_value: float = 1.0
+
+    def action(self, action):
+        """Process only the action component."""
+        # No need to handle transition dict - base class does it
+        return torch.clamp(action, self.min_value, self.max_value)
+
+    def get_config(self):
+        return {"min_value": self.min_value, "max_value": self.max_value}
+```
+
+For more advanced processor patterns including stateful processors, see [Implement Your Own Processor](implement_your_own_processor.mdx).
+
+## Advanced Features
+
+### Debugging with Hooks
+
+Processors support hooks for monitoring and debugging without modifying the pipeline code:
+
+```python
+# Define monitoring hooks
+def log_shapes(step_idx: int, transition: EnvTransition):
+    """Log tensor shapes after each step."""
+    obs = transition.get(TransitionKey.OBSERVATION)
+    if obs:
+        print(f"Step {step_idx} shapes:")
+        for key, value in obs.items():
+            if isinstance(value, torch.Tensor):
+                print(f"  {key}: {value.shape}")
+
+def check_nans(step_idx: int, transition: EnvTransition):
+    """Check for NaN values."""
+    obs = transition.get(TransitionKey.OBSERVATION)
+    if obs:
+        for key, value in obs.items():
+            if isinstance(value, torch.Tensor) and torch.isnan(value).any():
+                print(f"Warning: NaN detected in {key} at step {step_idx}")
+
+# Register hooks
+processor.register_after_step_hook(log_shapes)
+processor.register_after_step_hook(check_nans)
+
+# Process data - hooks will be called after each step
+output = processor(input_data)
+
+# Remove hooks when done debugging
+processor.unregister_after_step_hook(log_shapes)
+processor.unregister_after_step_hook(check_nans)
+```
+
+### Step-by-Step Inspection
+
+Use `step_through()` for detailed debugging of the transformation pipeline:
+
+```python
+# Inspect data at each transformation stage
+for i, intermediate in enumerate(processor.step_through(data)):
+    print(f"\n=== After step {i} ===")
+
+    # Check observation shapes
+    obs = intermediate.get(TransitionKey.OBSERVATION)
+    if obs:
+        for key, value in obs.items():
+            if isinstance(value, torch.Tensor):
+                print(f"{key}: shape={value.shape}, "
+                      f"dtype={value.dtype}, "
+                      f"device={value.device}, "
+                      f"range=[{value.min():.3f}, {value.max():.3f}]")
+
+    # Check action if present
+    action = intermediate.get(TransitionKey.ACTION)
+    if action is not None and isinstance(action, torch.Tensor):
+        print(f"action: shape={action.shape}, range=[{action.min():.3f}, {action.max():.3f}]")
+```
+
+### Pipeline Slicing
+
+Extract subsets of a pipeline for testing or creating variations:
+
+```python
+# Get specific steps
+first_three_steps = processor[:3]  # Returns new RobotProcessor
+middle_step = processor[2]         # Returns single ProcessorStep
+
+# Test individual steps
+test_input = {...}
+step_output = processor[0](test_input)  # Test first step only
+
+# Create variations
+variant_processor = RobotProcessor(
+    steps=processor.steps[:-1] + [new_final_step],
+    name="variant"
+)
+```
+
+## Best Practices and Tips
+
+### 1. Order Matters
+
+The sequence of processors is crucial. Follow this general order:
+
+```python
+# Preprocessing: Raw → Model-ready
+1. Rename (standardize keys)
+2. Batch (add dimensions)
+3. Tokenize (text → tokens)
+4. Normalize (scale values)
+5. Device (move to GPU)
+
+# Postprocessing: Model → Robot-ready
+1. Device (move to CPU)
+2. Unnormalize (restore scale)
+3. Unbatch (remove dimensions if needed)
+```
+
+### 2. Registration Best Practices
+
+```python
+# Always register custom steps for better portability
+@ProcessorStepRegistry.register("my_company/special_processor")
+class SpecialProcessor:
+    ...
+
+# Use namespaced names to avoid conflicts
+# Good: "my_company/augmentation"
+# Bad: "augmentation" (too generic)
+
+# Check registered processors
+print(ProcessorStepRegistry.list())  # See all registered processors
+```
+
+### 3. Common Pitfalls and Solutions
+
+**Tensor Device Mismatch:**
+
+```python
+# Problem: RuntimeError: Expected all tensors on same device
+# Solution: Ensure DeviceProcessor is in pipeline
+preprocessor = RobotProcessor(
+    steps=[
+        NormalizerProcessor(...),
+        DeviceProcessor(device="cuda")  # Add this
+    ]
+)
+```
+
+**Missing Statistics:**
+
+```python
+# Problem: NormalizerProcessor has no stats
+# Solution 1: Compute stats from dataset
+from lerobot.datasets.compute_stats import compute_stats
+stats = compute_stats(dataset)
+
+# Solution 2: Load with overrides
+processor = RobotProcessor.from_pretrained(
+    "model_path",
+    overrides={"normalizer_processor": {"stats": dataset.meta.stats}}
+)
+```
+
+## Next Steps
+
+Now that you understand processors, explore these topics:
+
+- [**Implement Your Own Processor**](implement_your_own_processor.mdx) - Deep dive into creating custom processors with advanced features like stateful processing
+- [**Policy Documentation**](policies.mdx) - Learn how different policies use processors
+- [**Dataset Documentation**](datasets.mdx) - Understand the data format that processors transform
+- [**Training Guide**](training.mdx) - See processors in action during model training
+- [**Evaluation Guide**](evaluation.mdx) - Learn about processor usage during policy evaluation
+
+## Summary
+
+Processors are the unsung heroes of robotics pipelines, handling the critical transformations between raw sensor data and model-ready tensors. By understanding and effectively using processors, you can:
+
+- Build robust, reusable data pipelines
+- Share preprocessing configurations across projects
+- Debug data transformations systematically
+- Ensure consistency between training and deployment
+- Create custom transformations for specialized tasks
+
+Remember: good preprocessing is often the difference between a model that works in theory
+and one that works in practice!
+The modular pipeline approach ensures your transformations are testable, reproducible,
+and portable across different robots and environments.

docs/source/koch.mdx

@@ -31,7 +31,7 @@ pip install -e ".[dynamixel]"
 To find the port for each bus servo adapter, run this script:
 
 ```bash
-lerobot-find-port
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -98,7 +98,7 @@ For a visual reference on how to set the motor ids please refer to [this video](
 <hfoption id="Command">
 
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --robot.type=koch_follower \
     --robot.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step
 ```
@@ -174,7 +174,7 @@ Do the same steps for the leader arm but modify the command or script accordingl
 <hfoption id="Command">
 
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --teleop.type=koch_leader \
     --teleop.port=/dev/tty.usbmodem575E0031751 \  # <- paste here the port found at previous step
 ```
@@ -211,7 +211,7 @@ Run the following command or API example to calibrate the follower arm:
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --robot.type=koch_follower \
     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --robot.id=my_awesome_follower_arm  # <- Give the robot a unique name
@@ -249,7 +249,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=koch_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm  # <- Give the robot a unique name

docs/source/lekiwi.mdx

@@ -60,7 +60,7 @@ First, we will assemble the two SO100/SO101 arms. One to attach to the mobile ba
 To find the port for each bus servo adapter, run this script:
 
 ```bash
-lerobot-find-port
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -116,7 +116,7 @@ The instructions for configuring the motors can be found in the SO101 [docs](./s
 You can run this command to setup motors for LeKiwi. It will first setup the motors for arm (id 6..1) and then setup motors for wheels (9,8,7)
 
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --robot.type=lekiwi \
     --robot.port=/dev/tty.usbmodem58760431551 # <- paste here the port found at previous step
 ```
@@ -174,7 +174,7 @@ The calibration process is very important because it allows a neural network tra
 Make sure the arm is connected to the Raspberry Pi and run this script or API example (on the Raspberry Pi via SSH) to launch calibration of the follower arm:
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --robot.type=lekiwi \
     --robot.id=my_awesome_kiwi # <- Give the robot a unique name
 ```
@@ -193,7 +193,7 @@ Then, to calibrate the leader arm (which is attached to the laptop/pc). Run the
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name

docs/source/phone_teleop.mdx

@@ -0,0 +1,195 @@
+# Phone
+
+Use your phone (iOS or Android) to control your robot.
+
+**In this guide you'll learn:**
+
+- How to connect an iOS/Android phone
+- How phone pose is mapped to robot end‑effector (EE) targets
+- How to tweak safety limits, gripper control, and IK settings
+
+To use phone to control your robot, install the relevant dependencies with:
+
+```bash
+pip install lerobot[phone]
+```
+
+## Get started
+
+### Supported platforms
+
+- iOS: Uses the HEBI Mobile I/O app (ARKit pose + buttons). Download the app first, open it and the examples will discover it on your network and stream the phone pose and inputs.
+- Android: Uses the `teleop` package (WebXR). When you start the Python process, it prints a local URL. Open the link on your phone, tap Start, then use Move to stream pose.
+
+Links:
+
+- Android WebXR library: [`teleop` on PyPI](https://pypi.org/project/teleop/)
+- iOS app: [HEBI Mobile I/O](https://docs.hebi.us/tools.html#mobile-io)
+
+### Phone orientation and controls
+
+- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phone’s frame with the robot frame so motion feels natural.
+- Enable/disable:
+  - iOS: Hold `B1` to enable teleoperation, release to stop. The first press captures a reference pose.
+  - Android: Press and hold the `Move` button, release to stop. The first press captures a reference pose.
+- Gripper control:
+  - iOS: Analog input `A3` controls the gripper as velocity input.
+  - Android: Buttons `A` and `B` act like increment/decrement (A opens, B closes). You can tune velocity in the `GripperVelocityToJoint` step.
+
+### Step 1: Choose the platform
+
+Modify the examples to use `PhoneOS.IOS` or `PhoneOS.ANDROID` in `PhoneConfig`. The API is identical across platforms, only the input source differs. All examples are under `examples/` and have `phone_so100_*.py` variants.
+
+Teleoperation example:
+
+```36:43:examples/phone_so100_teleop.py
+from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
+
+teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
+teleop_device = Phone(teleop_config)
+```
+
+### Step 2: Connect and calibrate
+
+When `Phone(teleop_config)` is created and `connect()` is called, calibration is prompted automatically. Hold the phone in the orientation described above, then:
+
+- iOS: press and hold `B1` to capture the reference pose.
+- Android: press `Move` button on the WebXR page to capture the reference pose.
+
+Why calibrate? We capture the current pose so subsequent poses are expressed in a robot aligned frame. When you again press the button to enable control, the position is recaptured to avoid drift when your phone is repositioned while it was disabled.
+
+### Step 3: Run an example
+
+Run on of the examples scripts to teleoperate, record a dataset, replay a dataset or evaluate a policy.
+
+All scripts assume you configured your robot (e.g., SO-100 follower) and set the correct serial port.
+
+- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
+- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
+
+You can customize mapping or safety limits by editing the processor steps shown in the examples.
+
+You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
+
+- Run this example to teleoperate:
+
+  ```bash
+  python examples/phone_so100_teleop.py
+  ```
+
+- Run this example to record a dataset, which saves absolute end effector observations and actions:
+
+  ```bash
+  python examples/phone_so100_record.py
+  ```
+
+- Run this example to replay recorded episodes:
+
+  ```bash
+  python examples/phone_so100_replay.py
+  ```
+
+- Run this example to evaluate a pretrained policy:
+
+  ```bash
+  python examples/phone_so100_eval.py
+  ```
+
+### Important pipeline steps and options
+
+- Kinematics are used in multiple steps. We use [Placo](https://github.com/Rhoban/placo) which is a wrapper around Pinocchio for handling our kinematics. We construct the kinematics object by passing the robot's URDF and target frame. We set `target_frame_name` to the gripper frame.
+
+  ```44:49:examples/phone_so100_teleop.py
+  RobotKinematics(
+      urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
+      target_frame_name="gripper_frame_link",
+      joint_names=list(robot.bus.motors.keys()),
+  )
+  ```
+
+- The `MapPhoneActionToRobotAction` step converts the calibrated phone pose and inputs into target deltas and gripper commands, below is shown what the step outputs.
+
+  ```72:83:src/lerobot/teleoperators/phone/phone_processor.py
+  # Map calibrated phone pose to robot targets (enabled gates the motion)
+  act.update(
+      {
+          "action.enabled": enabled,
+          "action.target_x": -pos[1] if enabled else 0.0,
+          "action.target_y":  pos[0] if enabled else 0.0,
+          "action.target_z":  pos[2] if enabled else 0.0,
+          "action.target_wx": rotvec[1] if enabled else 0.0,
+          "action.target_wy": rotvec[0] if enabled else 0.0,
+          "action.target_wz": -rotvec[2] if enabled else 0.0,
+          "action.gripper":   gripper,
+      }
+  )
+  ```
+
+- The `EEReferenceAndDelta` step converts target deltas to an absolute desired EE pose, storing a reference on enable, the `end_effector_step_sizes` are the step sizes for the EE pose and can be modified to change the motion speed.
+
+  ```56:65:examples/phone_so100_teleop.py
+  EEReferenceAndDelta(
+      kinematics=kinematics_solver,
+      end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
+      motor_names=list(robot.bus.motors.keys()),
+  )
+  ```
+
+- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` and `max_ee_twist_step_rad` are the step limits for the EE pose and can be modified to change the safety limits.
+
+  ```61:66:examples/phone_so100_teleop.py
+  EEBoundsAndSafety(
+      end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
+      max_ee_step_m=0.10,
+      max_ee_twist_step_rad=0.50,
+  )
+  ```
+
+- The `GripperVelocityToJoint` step turns a velocity‑like gripper input into absolute gripper position using the current measured state. The `speed_factor` is the factor by which the velocity is multiplied.
+
+  ```78:81:examples/phone_so100_teleop.py
+  GripperVelocityToJoint(
+      motor_names=list(robot.bus.motors.keys()),
+      speed_factor=20.0,
+  )
+  ```
+
+#### Different IK initial guesses
+
+We use different IK initial guesses in the kinematic steps. As initial guess either the current measured joints or the previous IK solution is used.
+
+- Closed loop (used in record/eval): sets `initial_guess_current_joints=True` so IK starts from the measured joints each frame.
+
+  ```71:76:examples/phone_so100_eval.py
+  InverseKinematicsEEToJoints(
+      kinematics=kinematics_solver,
+      motor_names=list(robot.bus.motors.keys()),
+      initial_guess_current_joints=True,  # closed loop
+  )
+  ```
+
+- Open loop (used in replay): sets `initial_guess_current_joints=False` so IK continues from the previous IK solution rather than the measured state. This preserves action stability when we replay without feedback.
+
+  ```80:86:examples/phone_so100_replay.py
+  InverseKinematicsEEToJoints(
+      kinematics=kinematics_solver,
+      motor_names=list(robot.bus.motors.keys()),
+      initial_guess_current_joints=False,  # open loop
+  )
+  ```
+
+### Pipeline steps explained
+
+- MapPhoneActionToRobotAction: converts calibrated phone pose and inputs into target deltas and a gripper command. Motion is gated by an enable signal (B1 on iOS, Move on Android).
+- AddRobotObservationAsComplimentaryData: reads current robot joints and inserts them under `complementary_data.raw_joint_positions` for FK/IK steps to use.
+- EEReferenceAndDelta: latches a reference EE pose on enable and combines it with target deltas to produce an absolute desired EE pose each frame. When disabled, it keeps sending the last commanded pose.
+- EEBoundsAndSafety: clamps the EE pose to a workspace and rate‑limits jumps for safety. Also declares `action.ee.*` features.
+- InverseKinematicsEEToJoints: turns an EE pose into joint positions with IK. `initial_guess_current_joints=True` is recommended for closed‑loop control; set `False` for open‑loop replay for stability.
+- GripperVelocityToJoint: integrates a velocity‑like gripper input into an absolute gripper position using the current measured state.
+- ForwardKinematicsJointsToEE: computes `observation.state.ee.*` from observed joints for logging and training on EE state.
+
+### Troubleshooting
+
+- iOS not discovered: ensure HEBI Mobile I/O is open and your laptop/phone are on the same network.
+- Android URL not reachable: check local you used `https` instead of `http`, use the exact IP printed by the script and allow your browser to enter and ignore the certificate issue.
+- Motion feels inverted: adjust the sign flips in `MapPhoneActionToRobotAction` or swap axes to match your setup.

docs/source/processors_robots_teleop.mdx

@@ -0,0 +1,148 @@
+# Processors for Robots and Teleoperators
+
+This guide shows how to build and modify processing pipelines that connect teleoperators (e.g., phone) to robots and datasets. Pipelines standardize conversions between different action/observation spaces so you can swap teleops and robots without rewriting glue code.
+
+We use the Phone to SO‑100 follower examples for concreteness, but the same patterns apply to other robots.
+
+**What you'll learn**
+
+- Absolute vs. relative EE control: What each means, trade‑offs, and how to choose for your task.
+- Three-pipeline pattern: How to map teleop actions → dataset actions → robot commands, and robot observations → dataset observations.
+- Adapters (`to_transition` / `to_output`): How these convert raw dicts to `EnvTransition` and back to reduce boilerplate.
+- Dataset feature contracts: How steps declare features via `transform_features(...)`, and how to aggregate/merge them for recording.
+- Choosing a representation: When to store joints, absolute EE poses, or relative EE deltas—and how that affects training.
+- Pipeline customization guidance: How to swap robots/URDFs safely and tune bounds, step sizes, and options like IK initialization.
+
+### Absolute vs relative EE control
+
+The examples in this guide use absolute end effector (EE) poses because they are easy to reason about. In practice, relative EE deltas or joint position are often preferred as learning features.
+
+You can choose what you save and learn from the teleop and robot action spaces, joints, absolute EE, or relative EE by using/implementing the right steps (and `transform_features()`) in your pipelines.
+
+## Three pipelines
+
+We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
+Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
+
+1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
+2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
+3. Pipeline 3: Robot observation space → dataset observation space (joints → EE pose)
+
+Below is an example of the three pipelines that we use in the phone to SO-100 follower examples:
+
+```69:90:examples/phone_so100_record.py
+phone_to_robot_ee_pose = RobotProcessor(  # teleop -> dataset action
+    steps=[MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
+           AddRobotObservationAsComplimentaryData(robot=robot),
+            EEReferenceAndDelta(kinematics=kinematics_solver,
+                               end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
+                               motor_names=list(robot.bus.motors.keys())),
+           EEBoundsAndSafety(end_effector_bounds={"min": [-1, -1, -1], "max": [1, 1, 1]},
+                             max_ee_step_m=0.20, max_ee_twist_step_rad=0.50)],
+    to_transition=to_transition_teleop_action,
+    to_output=lambda tr: tr,
+)
+
+robot_ee_to_joints = RobotProcessor(      # dataset action -> robot
+    steps=[InverseKinematicsEEToJoints(kinematics=kinematics_solver,
+                                       motor_names=list(robot.bus.motors.keys()),
+                                       initial_guess_current_joints=True),
+           GripperVelocityToJoint(motor_names=list(robot.bus.motors.keys()), speed_factor=20.0)],
+    to_transition=lambda tr: tr,
+    to_output=to_output_robot_action,
+)
+
+robot_joints_to_ee_pose = RobotProcessor( # robot obs -> dataset obs
+    steps=[ForwardKinematicsJointsToEE(kinematics=kinematics_solver,
+                                       motor_names=list(robot.bus.motors.keys()))],
+    to_transition=to_transition_robot_observation,
+    to_output=lambda tr: tr,
+)
+```
+
+## Why to_transition / to_output
+
+To convert from robot/teleoperator to pipeline and back, we use the `to_transition` and `to_output` pipeline adapters.
+They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dicts and the pipeline’s `EnvTransition` format.
+In the phone to SO-100 follower examples we use the following adapters:
+
+- `to_transition_teleop_action`: transforms the teleop action dict to a pipeline transition (puts keys under `action.*`, converts scalars/arrays to tensors, keeps objects like `Rotation` intact)
+- `to_output_robot_action`: transforms the pipeline transition to a robot action dict (extracts keys ending with `.pos`/`.vel` and strips `action.` prefix)
+- `to_transition_robot_observation`: transforms the robot observation dict to a pipeline transition (splits state vs images; stores state under `observation.state.*` and images under `observation.images.*`)
+
+See `src/lerobot/processor/converters.py` for more details.
+
+## Dataset feature contracts
+
+Dataset features are the keys saved in the dataset. Each step can declare what its dataset features are via `transform_features(...)`. We can then aggregate features per pipeline with `aggregate_pipeline_dataset_features()` and merge multiple groups with `merge_features(...)`.
+
+Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
+
+```203:211:src/lerobot/robots/so100_follower/robot_kinematic_processor.py
+def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    # Because this is last step we specify the dataset features of this step that we want to be stored in the dataset
+    features["action.ee.x"] = float
+    features["action.ee.y"] = float
+    features["action.ee.z"] = float
+    features["action.ee.wx"] = float
+    features["action.ee.wy"] = float
+    features["action.ee.wz"] = float
+    return features
+```
+
+Tip: declare features at the last step that produces them (e.g., `EEBoundsAndSafety` declares `action.ee.*`, `ForwardKinematicsJointsToEE` declares `observation.state.ee.*`).
+
+Below is an example of how we aggregate and merge features in the phone to SO-100 follower examples:
+
+```121:145:examples/phone_so100_record.py
+action_ee = aggregate_pipeline_dataset_features(
+    pipeline=phone_to_robot_ee_pose,
+    initial_features=phone.action_features,
+    use_videos=True,
+    patterns=["action.ee"],
+)
+
+gripper = aggregate_pipeline_dataset_features(
+    pipeline=robot_ee_to_joints,
+    initial_features={},
+    use_videos=True,
+    patterns=["action.gripper.pos", "observation.state.gripper.pos"],
+)
+
+observation_ee = aggregate_pipeline_dataset_features(
+    pipeline=robot_joints_to_ee_pose,
+    initial_features=robot.observation_features,
+    use_videos=True,
+    patterns=["observation.state.ee"],
+)
+
+dataset_features = merge_features(action_ee, gripper, observation_ee)
+```
+
+How it works:
+
+- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`).
+- `merge_features(...)`: combine multiple feature dicts.
+- Recording uses `to_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
+
+## Guidance when customizing robot pipelines
+
+You can store any of the following features as your action/observation space:
+
+- Joint positions
+- Absolute EE poses
+- Relative EE deltas
+- Other features: joint velocity, etc.
+
+Pick what you want to use for your policy action and observation space and configure/modify the pipelines and steps accordingly.
+
+### Different robots
+
+- Swap `RobotKinematics` URDF and `motor_names`. Ensure `target_frame_name` points to your gripper/wrist.
+
+### Safety first
+
+- When changing pipelines, start with tight bounds, implement safety steps when working with real robots.
+- Its advised to start with simulation first and then move to real robots.
+
+Hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.

docs/source/reachy2.mdx

@@ -1,288 +0,0 @@
-# Reachy 2
-
-Reachy 2 is an open-source humanoid robot made by Pollen Robotics, specifically designed for the development of embodied AI and real-world applications.
-Check out [Pollen Robotics website](https://www.pollen-robotics.com/reachy/), or access [Reachy 2 documentation](https://docs.pollen-robotics.com/) for more information on the platform!
-
-## Teleoperate Reachy 2
-
-Currently, there are two ways to teleoperate Reachy 2:
-
-- Pollen Robotics’ VR teleoperation (not included in LeRobot).
-- Robot-to-robot teleoperation (use one Reachy 2 to control another).
-
-## Reachy 2 Simulation
-
-**(Linux only)** You can run Reachy 2 in simulation (Gazebo or MuJoCo) using the provided [Docker image](https://hub.docker.com/r/pollenrobotics/reachy2_core).
-
-1. Install [Docker Engine](https://docs.docker.com/engine/).
-2. Run (for MuJoCo):
-
-```
-docker run --rm -it \
-  --name reachy \
-  --privileged \
-  --network host \
-  --ipc host \
-  --device-cgroup-rule='c 189:* rwm' \
-  --group-add audio \
-  -e ROS_DOMAIN_ID="$ROS_DOMAIN_ID" \
-  -e DISPLAY="$DISPLAY" \
-  -e RCUTILS_CONSOLE_OUTPUT_FORMAT="[{severity}]: {message}" \
-  -e REACHY2_CORE_SERVICE_FAKE="${REACHY2_CORE_SERVICE_FAKE:-true}" \
-  -v /dev:/dev \
-  -v "$HOME/.reachy_config":/home/reachy/.reachy_config_override \
-  -v "$HOME/.reachy.log":/home/reachy/.ros/log \
-  -v /usr/lib/x86_64-linux-gnu:/opt/host-libs \
-  --entrypoint /package/launch.sh \
-  pollenrobotics/reachy2_core:1.7.5.9_deploy \
-  start_rviz:=true start_sdk_server:=true mujoco:=true
-```
-
-> If MuJoCo runs slowly (low simulation frequency), append `-e LD_LIBRARY_PATH="/opt/host-libs:$LD_LIBRARY_PATH" \` to the previous command to improve performance:
->
-> ```
-> docker run --rm -it \
->   --name reachy \
->   --privileged \
->   --network host \
->   --ipc host \
->   --device-cgroup-rule='c 189:* rwm' \
->   --group-add audio \
->   -e ROS_DOMAIN_ID="$ROS_DOMAIN_ID" \
->   -e DISPLAY="$DISPLAY" \
->   -e RCUTILS_CONSOLE_OUTPUT_FORMAT="[{severity}]: {message}" \
->   -e REACHY2_CORE_SERVICE_FAKE="${REACHY2_CORE_SERVICE_FAKE:-true}" \
->   -e LD_LIBRARY_PATH="/opt/host-libs:$LD_LIBRARY_PATH" \
->   -v /dev:/dev \
->   -v "$HOME/.reachy_config":/home/reachy/.reachy_config_override \
->   -v "$HOME/.reachy.log":/home/reachy/.ros/log \
->   -v /usr/lib/x86_64-linux-gnu:/opt/host-libs \
->   --entrypoint /package/launch.sh \
->   pollenrobotics/reachy2_core:1.7.5.9_deploy \
->   start_rviz:=true start_sdk_server:=true mujoco:=true
-> ```
-
-## Setup
-
-### Prerequisites
-
-- On your robot, check the **service images** meet the minimum versions:
-  - **reachy2-core >= 1.7.5.2**
-  - **webrtc >= 2.0.1.1**
-
-Then, if you want to use VR teleoperation:
-
-- Install the [Reachy 2 teleoperation application](https://docs.pollen-robotics.com/teleoperation/teleoperation-introduction/discover-teleoperation/).
-  Use version **>=v1.2.0**
-
-We recommend using two computers: one for teleoperation (Windows required) and another for recording with LeRobot.
-
-### Install LeRobot
-
-Follow the [installation instructions](https://github.com/huggingface/lerobot#installation) to install LeRobot.
-
-Install LeRobot with Reachy 2 dependencies:
-
-```bash
-pip install -e ".[reachy2]"
-```
-
-### (Optional but recommended) Install pollen_data_acquisition_server
-
-How you manage Reachy 2 recording sessions is up to you, but the **easiest** way is to use this server so you can control sessions directly from the VR teleoperation app.
-
-> **Note:** Currently, only the VR teleoperation application works as a client for this server, so this step primarily targets teleoperation. You’re free to develop custom clients to manage sessions to your needs.
-
-In your LeRobot environment, install the server from source:
-
-```bash
-git clone https://github.com/pollen-robotics/pollen_data_acquisition_server.git
-cd pollen_data_acquisition_server
-pip install -e .
-```
-
-Find the [pollen_data_acquisition_server documentation here](https://github.com/pollen-robotics/pollen_data_acquisition_server).
-
-## Step 1: Recording
-
-### Get Reachy 2 IP address
-
-Before starting teleoperation and data recording, find the [robot's IP address](https://docs.pollen-robotics.com/getting-started/setup-reachy2/connect-reachy2/).
-We strongly recommend connecting all devices (PC and robot) via **Ethernet**.
-
-### Launch recording
-
-There are two ways to manage recording sessions when using the Reachy 2 VR teleoperation application:
-
-- **Using the data acquisition server (recommended for VR teleop)**: The VR app orchestrates sessions (via the server it tells LeRobot when to create datasets, start/stop episodes) while also controlling the robot’s motions.
-- **Using LeRobot’s record script**: LeRobot owns session control and decides when to start/stop episodes. If you also use the VR teleop app, it’s only for motion control.
-
-### Option 1: Using Pollen data acquisition server (recommended for VR teleop)
-
-Make sure you have installed pollen_data_acquisition_server, as explained in the Setup section.
-
-Launch the data acquisition server to be able to manage your session directly from the teleoperation application:
-
-```bash
-python -m pollen_data_acquisition_server.server
-```
-
-Then get into the teleoperation application and choose "Data acquisition session".
-You can finally setup your session by following the screens displayed.
-
-> Even without the VR app, you can use the `pollen_data_acquisition_server` with your own client implementation.
-
-### Option 2: Using lerobot.record
-
-Reachy 2 is fully supported by LeRobot’s recording features.
-If you choose this option but still want to use the VR teleoperation application, select "Standard session" in the app.
-
-**Example: start a recording without the mobile base:**
-First add reachy2 and reachy2_teleoperator to the imports of the record script. Then you can use the following command:
-
-```bash
-python -m lerobot.record \
-    --robot.type=reachy2 \
-    --robot.ip_address=192.168.0.200 \
-    --robot.id=r2-0000 \
-    --robot.use_external_commands=true \
-    --robot.with_mobile_base=false \
-    --teleop.type=reachy2_teleoperator \
-    --teleop.ip_address=192.168.0.200 \
-    --teleop.with_mobile_base=false \
-    --dataset.repo_id=pollen_robotics/record_test \
-    --dataset.single_task="Reachy 2 recording test" \
-    --dataset.num_episodes=1 \
-    --dataset.episode_time_s=5 \
-    --dataset.fps=15 \
-    --dataset.push_to_hub=true \
-    --dataset.private=true \
-    --display_data=true
-```
-
-#### Specific Options
-
-**Extended setup overview (all options included):**
-
-```bash
-python -m lerobot.record \
-    --robot.type=reachy2 \
-    --robot.ip_address=192.168.0.200 \
-    --robot.use_external_commands=true \
-    --robot.with_mobile_base=true \
-    --robot.with_l_arm=true \
-    --robot.with_r_arm=true \
-    --robot.with_neck=true \
-    --robot.with_antennas=true \
-    --robot.with_left_teleop_camera=true \
-    --robot.with_right_teleop_camera=true \
-    --robot.with_torso_camera=false \
-    --robot.disable_torque_on_disconnect=false \
-    --robot.max_relative_target=5.0 \
-    --teleop.type=reachy2_teleoperator \
-    --teleop.ip_address=192.168.0.200 \
-    --teleop.use_present_position=false \
-    --teleop.with_mobile_base=false \
-    --teleop.with_l_arm=true \
-    --teleop.with_r_arm=true \
-    --teleop.with_neck=true \
-    --teleop.with_antennas=true \
-    --dataset.repo_id=pollen_robotics/record_test \
-    --dataset.single_task="Reachy 2 recording test" \
-    --dataset.num_episodes=1 \
-    --dataset.episode_time_s=5 \
-    --dataset.fps=15 \
-    --dataset.push_to_hub=true \
-    --dataset.private=true \
-    --display_data=true
-```
-
-##### `--robot.use_external_commands`
-
-Determine whether LeRobot robot.send_action() sends commands to the robot.
-**Must** be set to false while using the VR teleoperation application, as the app already sends commands.
-
-##### `--teleop.use_present_position`
-
-Determine whether the teleoperator reads the goal or present position of the robot.
-Must be set to true if a compliant Reachy 2 is used to control another one.
-
-##### Use the relevant parts
-
-From our initial tests, recording **all** joints when only some are moving can reduce model quality with certain policies.
-To avoid this, you can exclude specific parts from recording and replay using:
-
-````
---robot.with_<part>=false
-```,
-with `<part>` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`.
-It determine whether the corresponding part is recorded in the observations. True if not set.
-
-By default, **all parts are recorded**.
-
-The same per-part mechanism is available in `reachy2_teleoperator` as well.
-
-````
-
---teleop.with\_<part>
-
-```
-with `<part>` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`.
-Determine whether the corresponding part is recorded in the actions. True if not set.
-
-> **Important:** In a given session, the **enabled parts must match** on both the robot and the teleoperator.
-For example, if the robot runs with `--robot.with_mobile_base=false`, the teleoperator must disable the same part `--teleoperator.with_mobile_base=false`.
-
-##### Use the relevant cameras
-
-You can do the same for **cameras**. By default, only the **teleoperation cameras** are recorded (both `left_teleop_camera` and `right_teleop_camera`). Enable or disable each camera with:
-
-```
-
---robot.with_left_teleop_camera=<true|false>
---robot.with_right_teleop_camera=<true|false>
---robot.with_torso_camera=<true|false>
-
-````
-
-
-## Step 2: Replay
-
-Make sure the robot is configured with the same parts as the dataset:
-
-```bash
-python -m lerobot.replay \
-    --robot.type=reachy2 \
-    --robot.ip_address=192.168.0.200 \
-    --robot.use_external_commands=false \
-    --robot.with_mobile_base=false \
-    --dataset.repo_id=pollen_robotics/record_test \
-    --dataset.episode=0
-    --display_data=true
-````
-
-## Step 3: Train
-
-```bash
-python -m lerobot.scripts.train \
-  --dataset.repo_id=pollen_robotics/record_test \
-  --policy.type=act \
-  --output_dir=outputs/train/reachy2_test \
-  --job_name=reachy2 \
-  --policy.device=mps \
-  --wandb.enable=true \
-  --policy.repo_id=pollen_robotics/record_test_policy
-```
-
-## Step 4: Evaluate
-
-```bash
-python -m lerobot.record \
-  --robot.type=reachy2 \
-  --robot.ip_address=192.168.0.200 \
-  --display_data=false \
-  --dataset.repo_id=pollen_robotics/eval_record_test \
-  --dataset.single_task="Evaluate reachy2 policy" \
-  --dataset.num_episodes=10 \
-  --policy.path=outputs/train/reachy2_test/checkpoints/last/pretrained_model
-```

docs/source/smolvla.mdx

@@ -54,7 +54,7 @@ If you don't have a gpu device, you can train using our notebook on [![Google Co
 Pass your dataset to the training script using `--dataset.repo_id`. If you want to test your installation, run the following command where we use one of the datasets we collected for the [SmolVLA Paper](https://huggingface.co/papers/2506.01844).
 
 ```bash
-cd lerobot && lerobot-train \
+cd lerobot && python -m lerobot.scripts.train \
   --policy.path=lerobot/smolvla_base \
   --dataset.repo_id=${HF_USER}/mydataset \
   --batch_size=64 \
@@ -73,7 +73,7 @@ cd lerobot && lerobot-train \
 Fine-tuning is an art. For a complete overview of the options for finetuning, run
 
 ```bash
-lerobot-train --help
+python -m lerobot.scripts.train --help
 ```
 
 <p align="center">
@@ -97,7 +97,7 @@ Similarly for when recording an episode, it is recommended that you are logged i
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash
-lerobot-record \
+python -m lerobot.record \
   --robot.type=so101_follower \
   --robot.port=/dev/ttyACM0 \ # <- Use your port
   --robot.id=my_blue_follower_arm \ # <- Use your robot id

docs/source/so100.mdx

@@ -26,7 +26,7 @@ Unlike the SO-101, the motor connectors are not easily accessible once the arm i
 To find the port for each bus servo adapter, run this script:
 
 ```bash
-lerobot-find-port
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -93,7 +93,7 @@ For a visual reference on how to set the motor ids please refer to [this video](
 <hfoption id="Command">
 
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
 ```
@@ -168,7 +168,7 @@ Do the same steps for the leader arm.
 <hfoptions id="setup_motors">
 <hfoption id="Command">
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem575E0031751  # <- paste here the port found at previous step
 ```
@@ -568,7 +568,7 @@ Run the following command or API example to calibrate the follower arm:
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -606,7 +606,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name

docs/source/so101.mdx

@@ -162,7 +162,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
 To find the port for each bus servo adapter, connect MotorBus to your computer via USB and power. Run the following script and disconnect the MotorBus when prompted:
 
 ```bash
-lerobot-find-port
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -240,7 +240,7 @@ Connect the usb cable from your computer and the power supply to the follower ar
 <hfoption id="Command">
 
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
 ```
@@ -316,7 +316,7 @@ Do the same steps for the leader arm.
 <hfoption id="Command">
 
 ```bash
-lerobot-setup-motors \
+python -m lerobot.setup_motors \
     --teleop.type=so101_leader \
     --teleop.port=/dev/tty.usbmodem575E0031751  # <- paste here the port found at previous step
 ```
@@ -353,7 +353,7 @@ Run the following command or API example to calibrate the follower arm:
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --robot.type=so101_follower \
     --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -402,7 +402,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
 <hfoption id="Command">
 
 ```bash
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=so101_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
     --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name

examples/4_train_policy_with_script.md

@@ -62,7 +62,7 @@ By default, every field takes its default value specified in the dataclass. If a
 Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --dataset.repo_id=lerobot/pusht \
     --policy.type=diffusion \
     --env.type=pusht
@@ -77,7 +77,7 @@ Let's break this down:
 Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
@@ -90,7 +90,7 @@ We now want to train a different policy for aloha on another task. We'll change
 Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
     --env.type=aloha \
@@ -127,7 +127,7 @@ Now, let's assume that we want to reproduce the run just above. That run has pro
 We can then simply load the config values from this file using:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
     --output_dir=outputs/train/act_aloha_transfer_2
 ```
@@ -137,7 +137,7 @@ lerobot-train \
 Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
     --output_dir=outputs/train/act_aloha_transfer_2
     --policy.n_action_steps=80
@@ -148,7 +148,7 @@ lerobot-train \
 `--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
 
 ```bash
-lerobot-train --config_path=lerobot/diffusion_pusht
+python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht
 ```
 
 will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
@@ -160,7 +160,7 @@ Being able to resume a training run is important in case it crashed or aborted f
 Let's reuse the command from the previous run and add a few more options:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
     --env.type=aloha \
@@ -179,7 +179,7 @@ INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
 Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
     --resume=true
 ```
@@ -190,7 +190,7 @@ Another reason for which you might want to resume a run is simply to extend trai
 You could double the number of steps of the previous run with:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
     --resume=true \
     --steps=200000
@@ -224,7 +224,7 @@ In addition to the features currently in Draccus, we've added a special `.path`
 For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --policy.path=lerobot/act_aloha_sim_transfer_cube_human \
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
@@ -270,7 +270,7 @@ We'll summarize here the main use cases to remember from this tutorial.
 #### Train a policy from scratch – CLI
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --policy.type=act \  # <- select 'act' policy
     --env.type=pusht \  # <- select 'pusht' environment
     --dataset.repo_id=lerobot/pusht  # <- train on this dataset
@@ -279,7 +279,7 @@ lerobot-train \
 #### Train a policy from scratch - config file + CLI
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --config_path=path/to/pretrained_model \  # <- can also be a repo_id
     --policy.n_action_steps=80  # <- you may still override values
 ```
@@ -287,7 +287,7 @@ lerobot-train \
 #### Resume/continue a training run
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --config_path=checkpoint/pretrained_model/ \
     --resume=true \
     --steps=200000  # <- you can change some training parameters
@@ -296,7 +296,7 @@ lerobot-train \
 #### Fine-tuning
 
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
     --policy.path=lerobot/act_aloha_sim_transfer_cube_human \  # <- can also be a local path to a checkpoint
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \

examples/backward_compatibility/replay.py

@@ -18,7 +18,7 @@ Replays the actions of an episode from a dataset on a robot.
 Example:
 
 ```shell
-lerobot-replay \
+python -m lerobot.replay \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=black \

examples/lekiwi/evaluate.py

@@ -1,7 +1,7 @@
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 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.policies.factory import make_processor
 from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.utils.control_utils import init_keyboard_listener
@@ -46,7 +46,7 @@ listener, events = init_keyboard_listener()
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")
 
-preprocessor, postprocessor = make_pre_post_processors(
+preprocessor, postprocessor = make_processor(
     policy_cfg=policy,
     pretrained_path=HF_MODEL_ID,
     dataset_stats=dataset.meta.stats,

examples/phone_so100_eval.py

@@ -16,17 +16,16 @@
 
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
-from lerobot.datasets.utils import combine_feature_dicts
+from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
+from lerobot.datasets.utils import merge_features
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies.act.modeling_act import ACTPolicy
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.processor import RobotProcessorPipeline
+from lerobot.policies.factory import make_processor
 from lerobot.processor.converters import (
-    identity_transition,
-    observation_to_transition,
-    transition_to_action,
+    to_output_robot_action,
+    to_transition_robot_observation,
 )
+from lerobot.processor.pipeline import RobotProcessor
 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 (
@@ -66,7 +65,7 @@ kinematics_solver = RobotKinematics(
 )
 
 # Build pipeline to convert ee pose action to joint action
-robot_ee_to_joints_processor = RobotProcessorPipeline(
+robot_ee_to_joints = RobotProcessor(
     steps=[
         AddRobotObservationAsComplimentaryData(robot=robot),
         InverseKinematicsEEToJoints(
@@ -75,36 +74,36 @@ robot_ee_to_joints_processor = RobotProcessorPipeline(
             initial_guess_current_joints=True,
         ),
     ],
-    to_transition=identity_transition,
-    to_output=transition_to_action,
+    to_transition=lambda tr: tr,
+    to_output=to_output_robot_action,
 )
 
 # Build pipeline to convert joint observation to ee pose observation
-robot_joints_to_ee_pose_processor = RobotProcessorPipeline(
+robot_joints_to_ee_pose = RobotProcessor(
     steps=[
         ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
     ],
-    to_transition=observation_to_transition,
-    to_output=identity_transition,
+    to_transition=to_transition_robot_observation,
+    to_output=lambda tr: tr,
 )
 
 # Build dataset action and gripper features
 action_ee_and_gripper = aggregate_pipeline_dataset_features(
-    pipeline=robot_ee_to_joints_processor,
-    initial_features=create_initial_features(),
+    pipeline=robot_ee_to_joints,
+    initial_features={},
     use_videos=True,
     patterns=["action.ee", "action.gripper.pos", "observation.state.gripper.pos"],
 )  # Get all ee action features + gripper pos action features
 
 # Build dataset observation features
 obs_ee = aggregate_pipeline_dataset_features(
-    pipeline=robot_joints_to_ee_pose_processor,
-    initial_features=create_initial_features(observation=robot.observation_features),
+    pipeline=robot_joints_to_ee_pose,
+    initial_features=robot.observation_features,
     use_videos=True,
     patterns=["observation.state.ee"],
 )  # Get all ee observation features
 
-dataset_features = combine_feature_dicts(obs_ee, action_ee_and_gripper)
+dataset_features = merge_features(obs_ee, action_ee_and_gripper)
 
 print("All dataset features: ", dataset_features)
 
@@ -128,7 +127,7 @@ robot.connect()
 episode_idx = 0
 
 policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
-preprocessor, postprocessor = make_pre_post_processors(
+preprocessor, postprocessor = make_processor(
     policy_cfg=policy,
     pretrained_path=HF_MODEL_ID,
     dataset_stats=dataset.meta.stats,
@@ -148,8 +147,8 @@ for episode_idx in range(NUM_EPISODES):
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,
         display_data=True,
-        robot_action_processor=robot_ee_to_joints_processor,
-        robot_observation_processor=robot_joints_to_ee_pose_processor,
+        robot_action_processor=robot_ee_to_joints,
+        robot_observation_processor=robot_joints_to_ee_pose,
     )
     dataset.save_episode()
 

examples/phone_so100_record.py

@@ -17,16 +17,15 @@
 
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
-from lerobot.datasets.utils import combine_feature_dicts
+from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
+from lerobot.datasets.utils import merge_features
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
-    action_to_transition,
-    identity_transition,
-    observation_to_transition,
-    transition_to_action,
+    to_output_robot_action,
+    to_transition_robot_observation,
+    to_transition_teleop_action,
 )
+from lerobot.processor.pipeline import RobotProcessor
 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 (
@@ -39,8 +38,8 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
+from lerobot.teleoperators.phone.phone import Phone
 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
@@ -74,7 +73,7 @@ kinematics_solver = RobotKinematics(
 )
 
 # Build pipeline to convert phone action to ee pose action
-phone_to_robot_ee_pose_processor = RobotProcessorPipeline(
+phone_to_robot_ee_pose = RobotProcessor(
     steps=[
         MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
         AddRobotObservationAsComplimentaryData(robot=robot),
@@ -89,12 +88,12 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline(
             max_ee_twist_step_rad=0.50,
         ),
     ],
-    to_transition=action_to_transition,
-    to_output=identity_transition,
+    to_transition=to_transition_teleop_action,
+    to_output=lambda tr: tr,
 )
 
 # Build pipeline to convert ee pose action to joint action
-robot_ee_to_joints_processor = RobotProcessorPipeline(
+robot_ee_to_joints = RobotProcessor(
     steps=[
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
@@ -106,31 +105,31 @@ robot_ee_to_joints_processor = RobotProcessorPipeline(
             speed_factor=20.0,
         ),
     ],
-    to_transition=identity_transition,
-    to_output=transition_to_action,
+    to_transition=lambda tr: tr,
+    to_output=to_output_robot_action,
 )
 
 # Build pipeline to convert joint observation to ee pose observation
-robot_joints_to_ee_pose = RobotProcessorPipeline(
+robot_joints_to_ee_pose = RobotProcessor(
     steps=[
         ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
     ],
-    to_transition=observation_to_transition,
-    to_output=identity_transition,
+    to_transition=to_transition_robot_observation,
+    to_output=lambda tr: tr,
 )
 
 # Build dataset ee action features
 action_ee = aggregate_pipeline_dataset_features(
-    pipeline=phone_to_robot_ee_pose_processor,
-    initial_features=create_initial_features(action=phone.action_features),
+    pipeline=phone_to_robot_ee_pose,
+    initial_features=phone.action_features,
     use_videos=True,
     patterns=["action.ee"],
 )
 
 # Get gripper pos action features
 gripper = aggregate_pipeline_dataset_features(
-    pipeline=robot_ee_to_joints_processor,
-    initial_features=create_initial_features(),
+    pipeline=robot_ee_to_joints,
+    initial_features={},
     use_videos=True,
     patterns=["action.gripper.pos", "observation.state.gripper.pos"],
 )
@@ -138,12 +137,12 @@ gripper = aggregate_pipeline_dataset_features(
 # Build dataset ee observation features
 observation_ee = aggregate_pipeline_dataset_features(
     pipeline=robot_joints_to_ee_pose,
-    initial_features=create_initial_features(observation=robot.observation_features),
+    initial_features=robot.observation_features,
     use_videos=True,
     patterns=["observation.state.ee"],
 )
 
-dataset_features = combine_feature_dicts(action_ee, gripper, observation_ee)
+dataset_features = merge_features(action_ee, gripper, observation_ee)
 
 print("All dataset features: ", dataset_features)
 
@@ -178,8 +177,8 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,
         display_data=True,
-        teleop_action_processor=phone_to_robot_ee_pose_processor,
-        robot_action_processor=robot_ee_to_joints_processor,
+        teleop_action_processor=phone_to_robot_ee_pose,
+        robot_action_processor=robot_ee_to_joints,
         robot_observation_processor=robot_joints_to_ee_pose,
     )
 
@@ -194,8 +193,8 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
             control_time_s=RESET_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
-            teleop_action_processor=phone_to_robot_ee_pose_processor,
-            robot_action_processor=robot_ee_to_joints_processor,
+            teleop_action_processor=phone_to_robot_ee_pose,
+            robot_action_processor=robot_ee_to_joints,
             robot_observation_processor=robot_joints_to_ee_pose,
         )
 

examples/phone_so100_replay.py

@@ -19,8 +19,8 @@ import time
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotProcessorPipeline
-from lerobot.processor.converters import action_to_transition, transition_to_action
+from lerobot.processor.converters import to_output_robot_action
+from lerobot.processor.pipeline import RobotProcessor
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     AddRobotObservationAsComplimentaryData,
@@ -49,8 +49,33 @@ kinematics_solver = RobotKinematics(
     joint_names=list(robot.bus.motors.keys()),
 )
 
+
+# This method converts the action from the dataset to a transition for pipeline
+def action_to_transition(action: dict):
+    act = {}
+
+    # EE pose
+    for k in ("ee.x", "ee.y", "ee.z", "ee.wx", "ee.wy", "ee.wz"):
+        if k in action:
+            act[f"action.{k}"] = float(action[k])
+
+    # Gripper: your dataset has absolute position
+    if "gripper.pos" in action:
+        act["action.gripper.pos"] = float(action["gripper.pos"])
+
+    return {
+        "observation": None,
+        "action": act,
+        "reward": None,
+        "done": False,
+        "truncated": False,
+        "info": {},
+        "complementary_data": {},
+    }
+
+
 # Build pipeline to convert ee pose action to joint action
-robot_ee_to_joints_processor = RobotProcessorPipeline(
+robot_ee_to_joints = RobotProcessor(
     steps=[
         AddRobotObservationAsComplimentaryData(robot=robot),
         InverseKinematicsEEToJoints(
@@ -60,10 +85,10 @@ robot_ee_to_joints_processor = RobotProcessorPipeline(
         ),
     ],
     to_transition=action_to_transition,
-    to_output=transition_to_action,
+    to_output=to_output_robot_action,
 )
 
-robot_ee_to_joints_processor.reset()
+robot_ee_to_joints.reset()
 
 log_say(f"Replaying episode {EPISODE_IDX}")
 for idx in range(dataset.num_frames):
@@ -73,7 +98,7 @@ for idx in range(dataset.num_frames):
         name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
 
-    joint_action = robot_ee_to_joints_processor(ee_action)
+    joint_action = robot_ee_to_joints(ee_action)
     action_sent = robot.send_action(joint_action)
 
     busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))

examples/phone_so100_teleop.py

@@ -16,8 +16,8 @@
 import time
 
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotProcessorPipeline
-from lerobot.processor.converters import action_to_transition, transition_to_action
+from lerobot.processor import RobotProcessor
+from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     AddRobotObservationAsComplimentaryData,
@@ -28,8 +28,8 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
+from lerobot.teleoperators.phone.phone import Phone
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.teleoperators.phone.teleop_phone import Phone
 
 # Initialize the robot and teleoperator
 robot_config = SO100FollowerConfig(
@@ -48,8 +48,8 @@ kinematics_solver = RobotKinematics(
     joint_names=list(robot.bus.motors.keys()),
 )
 
-# Build pipeline to convert phone action to ee pose action to joint action
-phone_to_robot_joints_processor = RobotProcessorPipeline(
+# Build pipeline to convert phone action to ee pose action
+phone_to_robot_ee_pose = RobotProcessor(
     steps=[
         MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
         AddRobotObservationAsComplimentaryData(robot=robot),
@@ -63,6 +63,14 @@ phone_to_robot_joints_processor = RobotProcessorPipeline(
             max_ee_step_m=0.10,
             max_ee_twist_step_rad=0.50,
         ),
+    ],
+    to_transition=to_transition_teleop_action,
+    to_output=lambda tr: tr,
+)
+
+# Build pipeline to convert ee pose action to joint action
+robot_ee_to_joints = RobotProcessor(
+    steps=[
         InverseKinematicsEEToJoints(
             kinematics=kinematics_solver,
             motor_names=list(robot.bus.motors.keys()),
@@ -72,8 +80,8 @@ phone_to_robot_joints_processor = RobotProcessorPipeline(
             speed_factor=20.0,
         ),
     ],
-    to_transition=action_to_transition,
-    to_output=transition_to_action,
+    to_transition=lambda tr: tr,
+    to_output=to_output_robot_action,
 )
 
 robot.connect()
@@ -81,11 +89,19 @@ teleop_device.connect()
 
 print("Starting teleop loop. Move your phone to teleoperate the robot.")
 while True:
+    phone_obs = teleop_device.get_action()
+    if not phone_obs:
+        time.sleep(0.01)
+        continue
+
     # Get teleop observation
     phone_obs = teleop_device.get_action()
 
-    # Phone -> EE pose -> Joints transition
-    joint_action = phone_to_robot_joints_processor(phone_obs)
+    # Phone to EE pose transition
+    ee_transition = phone_to_robot_ee_pose(phone_obs)
+
+    # EE pose to Joints transition
+    joint_action = robot_ee_to_joints(ee_transition)
 
     if joint_action:
         robot.send_action(joint_action)

pyproject.toml

@@ -73,6 +73,7 @@ dependencies = [
     "pynput>=1.7.7",
     "pyserial>=3.5",
     "wandb>=0.20.0",
+    "scipy>=1.15.2",
 
     "torch>=2.2.1,<2.8.0", # TODO: Bumb dependency
     "torchcodec>=0.2.1,<0.6.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')", # TODO: Bumb dependency
@@ -106,7 +107,6 @@ dynamixel = ["dynamixel-sdk>=3.7.31"]
 gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0"]
 hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"]
 lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1"]
-reachy2 = ["reachy2_sdk>=1.0.14"]
 kinematics = ["lerobot[placo-dep]"]
 intelrealsense = [
     "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'",
@@ -143,7 +143,6 @@ all = [
     "lerobot[gamepad]",
     "lerobot[hopejr]",
     "lerobot[lekiwi]",
-    "lerobot[reachy2]",
     "lerobot[kinematics]",
     "lerobot[intelrealsense]",
     "lerobot[pi0]",

scripts/convert_videos_to_images.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+
+"""
+Convert video dataset to image dataset for faster training.
+This pre-extracts all frames from MP4 files to PNG images.
+"""
+
+import argparse
+from pathlib import Path
+import logging
+import shutil
+
+def convert_dataset_videos_to_images(repo_id: str, root: str | None = None):
+    """Convert all videos in a LeRobot dataset to individual image files."""
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    from lerobot.datasets.video_utils import decode_video_frames
+    import torch
+    
+    # Load dataset
+    dataset = LeRobotDataset(repo_id, root=root, download_videos=True)
+    
+    total_frames_processed = 0
+    
+    for ep_idx in range(dataset.meta.total_episodes):
+        logging.info(f"Processing episode {ep_idx}/{dataset.meta.total_episodes}")
+        
+        for vid_key in dataset.meta.video_keys:
+            video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, vid_key)
+            
+            if not video_path.exists():
+                logging.warning(f"Video not found: {video_path}")
+                continue
+                
+            # Create image directory  
+            img_dir = dataset.root / f"images/chunk-{dataset.meta.get_episode_chunk(ep_idx)}/{vid_key}"
+            img_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Decode all frames from video
+            # Get episode length to decode all frames
+            ep_length = dataset.meta.episodes[ep_idx]["length"]
+            timestamps = [i / dataset.fps for i in range(ep_length)]
+            
+            try:
+                frames = decode_video_frames(video_path, timestamps, dataset.tolerance_s, dataset.video_backend)
+                
+                # Save each frame as PNG
+                for i, frame in enumerate(frames.squeeze(0)):
+                    img_path = img_dir / f"episode_{ep_idx:06d}_{i:06d}.png"
+                    # Convert tensor to PIL and save
+                    import torchvision.transforms as T
+                    to_pil = T.ToPILImage()
+                    pil_frame = to_pil(frame)
+                    pil_frame.save(img_path)
+                    
+                total_frames_processed += len(frames.squeeze(0))
+                logging.info(f"  Extracted {len(frames.squeeze(0))} frames to {img_dir}")
+                
+            except Exception as e:
+                logging.error(f"Failed to process {video_path}: {e}")
+                continue
+    
+    logging.info(f"Conversion complete! Processed {total_frames_processed} total frames")
+    logging.info(f"You can now use download_videos=False to use the extracted images")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert LeRobot video dataset to images")
+    parser.add_argument("repo_id", help="Dataset repo ID (e.g., 'kenmacken/record-test-2')")
+    parser.add_argument("--root", help="Local root directory", default=None)
+    
+    args = parser.parse_args()
+    
+    logging.basicConfig(level=logging.INFO)
+    
+    convert_dataset_videos_to_images(args.repo_id, args.root)

src/lerobot/calibrate.py

@@ -18,7 +18,7 @@ Helper to recalibrate your device (robot or teleoperator).
 Example:
 
 ```shell
-lerobot-calibrate \
+python -m lerobot.calibrate \
     --teleop.type=so100_leader \
     --teleop.port=/dev/tty.usbmodem58760431551 \
     --teleop.id=blue

src/lerobot/cameras/opencv/camera_opencv.py

@@ -60,7 +60,7 @@ class OpenCVCamera(Camera):
     or port changes, especially on Linux. Use the provided utility script to find
     available camera indices or paths:
     ```bash
-    lerobot-find-cameras opencv
+    python -m lerobot.find_cameras opencv
     ```
 
     The camera's default settings (FPS, resolution, color mode) are used unless
@@ -165,7 +165,8 @@ class OpenCVCamera(Camera):
             self.videocapture.release()
             self.videocapture = None
             raise ConnectionError(
-                f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
+                f"Failed to open {self}."
+                f"Run `python -m lerobot.find_cameras opencv` to find available cameras."
             )
 
         self._configure_capture_settings()

src/lerobot/cameras/reachy2_camera/__init__.py

@@ -1,16 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .configuration_reachy2_camera import Reachy2CameraConfig
-from .reachy2_camera import Reachy2Camera

src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py

@@ -1,78 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-
-from ..configs import CameraConfig, ColorMode
-
-
-@CameraConfig.register_subclass("reachy2_camera")
-@dataclass
-class Reachy2CameraConfig(CameraConfig):
-    """Configuration class for Reachy 2 camera devices.
-
-    This class provides configuration options for Reachy 2 cameras,
-    supporting both the teleop and depth cameras. It includes settings
-    for resolution, frame rate, color mode, and the selection of the cameras.
-
-    Example configurations:
-    ```python
-    # Basic configurations
-    Reachy2CameraConfig(
-        name="teleop",
-        image_type="left",
-        ip_address="192.168.0.200",  # IP address of the robot
-        fps=15,
-        width=640,
-        height=480,
-        color_mode=ColorMode.RGB,
-    )  # Left teleop camera, 640x480 @ 15FPS
-    ```
-
-    Attributes:
-        name: Name of the camera device. Can be "teleop" or "depth".
-        image_type: Type of image stream. For "teleop" camera, can be "left" or "right".
-                    For "depth" camera, can be "rgb" or "depth". (depth is not supported yet)
-        fps: Requested frames per second for the color stream.
-        width: Requested frame width in pixels for the color stream.
-        height: Requested frame height in pixels for the color stream.
-        color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
-        ip_address: IP address of the robot. Defaults to "localhost".
-        port: Port number for the camera server. Defaults to 50065.
-
-    Note:
-        - Only 3-channel color output (RGB/BGR) is currently supported.
-    """
-
-    name: str
-    image_type: str
-    color_mode: ColorMode = ColorMode.RGB
-    ip_address: str | None = "localhost"
-    port: int = 50065
-    # use_depth: bool = False
-
-    def __post_init__(self):
-        if self.name not in ["teleop", "depth"]:
-            raise ValueError(f"`name` is expected to be 'teleop' or 'depth', but {self.name} is provided.")
-        if (self.name == "teleop" and self.image_type not in ["left", "right"]) or (
-            self.name == "depth" and self.image_type not in ["rgb", "depth"]
-        ):
-            raise ValueError(
-                f"`image_type` is expected to be 'left' or 'right' for teleop camera, and 'rgb' or 'depth' for depth camera, but {self.image_type} is provided."
-            )
-
-        if self.color_mode not in ["rgb", "bgr"]:
-            raise ValueError(
-                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
-            )

src/lerobot/cameras/reachy2_camera/reachy2_camera.py

@@ -1,288 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Provides the Reachy2Camera class for capturing frames from Reachy 2 cameras using Reachy 2's CameraManager.
-"""
-
-import logging
-import os
-import platform
-import time
-from threading import Event, Lock, Thread
-from typing import Any
-
-# Fix MSMF hardware transform compatibility for Windows before importing cv2
-if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS" not in os.environ:
-    os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
-import cv2
-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 ..camera import Camera
-from .configuration_reachy2_camera import ColorMode, Reachy2CameraConfig
-
-logger = logging.getLogger(__name__)
-
-
-class Reachy2Camera(Camera):
-    """
-    Manages Reachy 2 camera using Reachy 2 CameraManager.
-
-    This class provides a high-level interface to connect to, configure, and read
-    frames from Reachy 2 cameras. It supports both synchronous and asynchronous
-    frame reading.
-
-    An Reachy2Camera instance requires a camera name (e.g., "teleop") and an image
-    type (e.g., "left") to be specified in the configuration.
-
-    The camera's default settings (FPS, resolution, color mode) are used unless
-    overridden in the configuration.
-    """
-
-    def __init__(self, config: Reachy2CameraConfig):
-        """
-        Initializes the Reachy2Camera instance.
-
-        Args:
-            config: The configuration settings for the camera.
-        """
-        super().__init__(config)
-
-        self.config = config
-
-        self.fps = config.fps
-        self.color_mode = config.color_mode
-
-        self.cam_manager: CameraManager | None = None
-
-        self.thread: Thread | None = None
-        self.stop_event: Event | None = None
-        self.frame_lock: Lock = Lock()
-        self.latest_frame: np.ndarray | None = None
-        self.new_frame_event: Event = Event()
-
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}({self.config.name}, {self.config.image_type})"
-
-    @property
-    def is_connected(self) -> bool:
-        """Checks if the camera is currently connected and opened."""
-        if self.config.name == "teleop":
-            return self.cam_manager._grpc_connected and self.cam_manager.teleop if self.cam_manager else False
-        elif self.config.name == "depth":
-            return self.cam_manager._grpc_connected and self.cam_manager.depth if self.cam_manager else False
-        else:
-            raise ValueError(f"Invalid camera name '{self.config.name}'. Expected 'teleop' or 'depth'.")
-
-    def connect(self, warmup: bool = True):
-        """
-        Connects to the Reachy2 CameraManager as specified in the configuration.
-        """
-        self.cam_manager = CameraManager(host=self.config.ip_address, port=self.config.port)
-        self.cam_manager.initialize_cameras()
-
-        logger.info(f"{self} connected.")
-
-    @staticmethod
-    def find_cameras(ip_address: str = "localhost", port: int = 50065) -> list[dict[str, Any]]:
-        """
-        Detects available Reachy 2 cameras.
-
-        Returns:
-            List[Dict[str, Any]]: A list of dictionaries,
-            where each dictionary contains 'name', 'stereo',
-            and the default profile properties (width, height, fps).
-        """
-        initialized_cameras = []
-        camera_manager = CameraManager(host=ip_address, port=port)
-
-        for camera in [camera_manager.teleop, camera_manager.depth]:
-            if camera is None:
-                continue
-
-            height, width, _, _, _, _, _ = camera.get_parameters()
-
-            camera_info = {
-                "name": camera._cam_info.name,
-                "stereo": camera._cam_info.stereo,
-                "default_profile": {
-                    "width": width,
-                    "height": height,
-                    "fps": 30,
-                },
-            }
-            initialized_cameras.append(camera_info)
-
-        camera_manager.disconnect()
-        return initialized_cameras
-
-    def read(self, color_mode: ColorMode | None = None) -> np.ndarray:
-        """
-        Reads a single frame synchronously from the camera.
-
-        This is a blocking call.
-
-        Args:
-            color_mode (Optional[ColorMode]): If specified, overrides the default
-                color mode (`self.color_mode`) for this read operation (e.g.,
-                request RGB even if default is BGR).
-
-        Returns:
-            np.ndarray: The captured frame as a NumPy array in the format
-                       (height, width, channels), using the specified or default
-                       color mode and applying any configured rotation.
-        """
-        if not self.is_connected:
-            raise DeviceNotConnectedError(f"{self} is not connected.")
-
-        start_time = time.perf_counter()
-
-        frame = None
-
-        if self.cam_manager is None:
-            raise DeviceNotConnectedError(f"{self} is not connected.")
-        else:
-            if self.config.name == "teleop" and hasattr(self.cam_manager, "teleop"):
-                if self.config.image_type == "left":
-                    frame = self.cam_manager.teleop.get_frame(CameraView.LEFT, size=(640, 480))[0]
-                elif self.config.image_type == "right":
-                    frame = self.cam_manager.teleop.get_frame(CameraView.RIGHT, size=(640, 480))[0]
-            elif self.config.name == "depth" and hasattr(self.cam_manager, "depth"):
-                if self.config.image_type == "depth":
-                    frame = self.cam_manager.depth.get_depth_frame()[0]
-                elif self.config.image_type == "rgb":
-                    frame = self.cam_manager.depth.get_frame(size=(640, 480))[0]
-
-            if frame is None:
-                return np.empty((0, 0, 3), dtype=np.uint8)
-
-            if self.config.color_mode == "rgb":
-                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
-        read_duration_ms = (time.perf_counter() - start_time) * 1e3
-        logger.debug(f"{self} read took: {read_duration_ms:.1f}ms")
-
-        return frame
-
-    def _read_loop(self):
-        """
-        Internal loop run by the background thread for asynchronous reading.
-
-        On each iteration:
-        1. Reads a color frame
-        2. Stores result in latest_frame (thread-safe)
-        3. Sets new_frame_event to notify listeners
-
-        Stops on DeviceNotConnectedError, logs other errors and continues.
-        """
-        while not self.stop_event.is_set():
-            try:
-                color_image = self.read()
-
-                with self.frame_lock:
-                    self.latest_frame = color_image
-                self.new_frame_event.set()
-
-            except DeviceNotConnectedError:
-                break
-            except Exception as e:
-                logger.warning(f"Error reading frame in background thread for {self}: {e}")
-
-    def _start_read_thread(self) -> None:
-        """Starts or restarts the background read thread if it's not running."""
-        if self.thread is not None and self.thread.is_alive():
-            self.thread.join(timeout=0.1)
-        if self.stop_event is not None:
-            self.stop_event.set()
-
-        self.stop_event = Event()
-        self.thread = Thread(target=self._read_loop, args=(), name=f"{self}_read_loop")
-        self.thread.daemon = True
-        self.thread.start()
-
-    def _stop_read_thread(self) -> None:
-        """Signals the background read thread to stop and waits for it to join."""
-        if self.stop_event is not None:
-            self.stop_event.set()
-
-        if self.thread is not None and self.thread.is_alive():
-            self.thread.join(timeout=2.0)
-
-        self.thread = None
-        self.stop_event = None
-
-    def async_read(self, timeout_ms: float = 200) -> np.ndarray:
-        """
-        Reads the latest available frame asynchronously.
-
-        This method retrieves the most recent frame captured by the background
-        read thread. It does not block waiting for the camera hardware directly,
-        but may wait up to timeout_ms for the background thread to provide a frame.
-
-        Args:
-            timeout_ms (float): Maximum time in milliseconds to wait for a frame
-                to become available. Defaults to 200ms (0.2 seconds).
-
-        Returns:
-            np.ndarray: The latest captured frame as a NumPy array in the format
-                       (height, width, channels), processed according to configuration.
-
-        Raises:
-            DeviceNotConnectedError: If the camera is not connected.
-            TimeoutError: If no frame becomes available within the specified timeout.
-            RuntimeError: If an unexpected error occurs.
-        """
-        if not self.is_connected:
-            raise DeviceNotConnectedError(f"{self} is not connected.")
-
-        if self.thread is None or not self.thread.is_alive():
-            self._start_read_thread()
-
-        if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0):
-            thread_alive = self.thread is not None and self.thread.is_alive()
-            raise TimeoutError(
-                f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. "
-                f"Read thread alive: {thread_alive}."
-            )
-
-        with self.frame_lock:
-            frame = self.latest_frame
-            self.new_frame_event.clear()
-
-        if frame is None:
-            raise RuntimeError(f"Internal error: Event set but no frame available for {self}.")
-
-        return frame
-
-    def disconnect(self):
-        """
-        Stops the background read thread (if running).
-
-        Raises:
-            DeviceNotConnectedError: If the camera is already disconnected.
-        """
-        if not self.is_connected and self.thread is None:
-            raise DeviceNotConnectedError(f"{self} not connected.")
-
-        if self.thread is not None:
-            self._stop_read_thread()
-
-        if self.cam_manager is not None:
-            self.cam_manager.disconnect()
-
-        logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/camera_realsense.py

@@ -51,7 +51,7 @@ class RealSenseCamera(Camera):
 
     Use the provided utility script to find available camera indices and default profiles:
     ```bash
-    lerobot-find-cameras realsense
+    python -m lerobot.find_cameras realsense
     ```
 
     A `RealSenseCamera` instance requires a configuration object specifying the
@@ -176,7 +176,8 @@ class RealSenseCamera(Camera):
             self.rs_profile = None
             self.rs_pipeline = None
             raise ConnectionError(
-                f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
+                f"Failed to open {self}."
+                "Run `python -m lerobot.find_cameras realsense` to find available cameras."
             ) from e
 
         self._configure_capture_settings()

src/lerobot/cameras/utils.py

@@ -37,14 +37,8 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[s
             from .realsense.camera_realsense import RealSenseCamera
 
             cameras[key] = RealSenseCamera(cfg)
-
-        elif cfg.type == "reachy2_camera":
-            from .reachy2_camera.reachy2_camera import Reachy2Camera
-
-            cameras[key] = Reachy2Camera(cfg)
-
         else:
-            raise ValueError(f"The camera type '{cfg.type}' is not valid.")
+            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
 
     return cameras
 

src/lerobot/configs/default.py

@@ -33,6 +33,8 @@ class DatasetConfig:
     # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | None = None
     episodes: list[int] | None = None
+    # Percentage of dataset to use (0-100). If set, overrides episodes parameter.
+    percentage: float | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True

src/lerobot/configs/types.py

@@ -27,11 +27,6 @@ class FeatureType(str, Enum):
     LANGUAGE = "LANGUAGE"
 
 
-class PipelineFeatureType(str, Enum):
-    ACTION = "ACTION"
-    OBSERVATION = "OBSERVATION"
-
-
 class NormalizationMode(str, Enum):
     MIN_MAX = "MIN_MAX"
     MEAN_STD = "MEAN_STD"

src/lerobot/constants.py

@@ -24,11 +24,6 @@ OBS_IMAGES = "observation.images"
 OBS_LANGUAGE = "observation.language"
 ACTION = "action"
 REWARD = "next.reward"
-TRUNCATED = "next.truncated"
-DONE = "next.done"
-
-OBS_LANGUAGE_TOKENS = "observation.language.tokens"
-OBS_LANGUAGE_ATTENTION_MASK = "observation.language.attention_mask"
 
 ROBOTS = "robots"
 ROBOT_TYPE = "robot_type"
@@ -45,9 +40,6 @@ OPTIMIZER_STATE = "optimizer_state.safetensors"
 OPTIMIZER_PARAM_GROUPS = "optimizer_param_groups.json"
 SCHEDULER_STATE = "scheduler_state.json"
 
-POLICY_PREPROCESSOR_DEFAULT_NAME = "policy_preprocessor"
-POLICY_POSTPROCESSOR_DEFAULT_NAME = "policy_postprocessor"
-
 if "LEROBOT_HOME" in os.environ:
     raise ValueError(
         f"You have a 'LEROBOT_HOME' environment variable set to '{os.getenv('LEROBOT_HOME')}'.\n"

src/lerobot/datasets/factory.py

@@ -13,7 +13,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import logging
 from pprint import pformat
 
 import torch
@@ -87,10 +86,24 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
         )
         delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
+
+        # Handle percentage parameter
+        episodes = cfg.dataset.episodes
+        if cfg.dataset.percentage is not None:
+            # Calculate episodes based on percentage
+            total_episodes = ds_meta.total_episodes
+            num_episodes_to_use = max(1, int(total_episodes * cfg.dataset.percentage / 100))
+            episodes = list(range(num_episodes_to_use))
+            import logging
+
+            logging.info(
+                f"Using {cfg.dataset.percentage}% of dataset: {num_episodes_to_use}/{total_episodes} episodes"
+            )
+
         dataset = LeRobotDataset(
             cfg.dataset.repo_id,
             root=cfg.dataset.root,
-            episodes=cfg.dataset.episodes,
+            episodes=episodes,
             delta_timestamps=delta_timestamps,
             image_transforms=image_transforms,
             revision=cfg.dataset.revision,

src/lerobot/datasets/lerobot_dataset.py

@@ -825,8 +825,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """
         if not episode_data:
             episode_buffer = self.episode_buffer
-        else:
-            episode_buffer = episode_data
 
         validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
 

src/lerobot/datasets/pipeline_features.py

@@ -12,130 +12,83 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
 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
-
-
-def create_initial_features(
-    action: dict[str, Any] | None, observation: dict[str, Any] | None
-) -> dict[PipelineFeatureType, dict[str, Any]]:
-    """
-    Creates the initial features dict for the dataset from action and observation specs.
-
-    Args:
-        action: A dictionary of action feature names to their types/shapes.
-        observation: A dictionary of observation feature names to their types/shapes.
-
-    Returns:
-        The initial features dictionary structured by PipelineFeatureType.
-    """
-    features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
-    if action:
-        features[PipelineFeatureType.ACTION] = action
-    if observation:
-        features[PipelineFeatureType.OBSERVATION] = observation
-    return features
-
-
-# Helper to filter state/action keys based on regex patterns.
-def should_keep(key: str, patterns: tuple[str]) -> bool:
-    if patterns is None:
-        return True
-    return any(re.search(pat, key) for pat in patterns)
-
-
-def strip_prefix(key: str, prefixes_to_strip: tuple[str]) -> str:
-    for prefix in prefixes_to_strip:
-        if key.startswith(prefix):
-            return key[len(prefix) :]
-    return key
-
-
-# Define prefixes to strip from feature keys for clean names.
-# Handles both fully qualified (e.g., "action.state") and short (e.g., "state") forms.
-PREFIXES_TO_STRIP = tuple(
-    f"{token}." for const in (ACTION, OBS_STATE, OBS_IMAGES) for token in (const, const.split(".")[-1])
-)
+from lerobot.processor.pipeline import RobotProcessor
 
 
 def aggregate_pipeline_dataset_features(
-    pipeline: DataProcessorPipeline,
-    initial_features: dict[PipelineFeatureType, dict[str, Any]],
+    pipeline: RobotProcessor,
+    initial_features: dict[str, Any],
     *,
     use_videos: bool = True,
     patterns: Sequence[str] | None = None,
 ) -> dict[str, dict]:
     """
-    Aggregates and filters pipeline features to create a dataset-ready features dictionary.
+    Aggregates the pipeline's features and returns a features dict ready for the dataset,
+    filtered to only those keys matching any of the given patterns (for action/state only).
 
-    This function transforms initial features using the pipeline, categorizes them as action or observations
-    (image or state), filters them based on `use_videos` and `patterns`, and finally
-    formats them for use with a Hugging Face LeRobot Dataset.
-
-    Args:
-        pipeline: The DataProcessorPipeline to apply.
-        initial_features: A dictionary of raw feature specs for actions and observations.
-        use_videos: If False, image features are excluded.
-        patterns: A sequence of regex patterns to filter action and state features.
-                  Image features are not affected by this filter.
-
-    Returns:
-        A dictionary of features formatted for a Hugging Face LeRobot Dataset.
+    - `initial_features`: raw camera specs, e.g. {"front": (h,w,c), ...}
+    - `use_videos`: whether to treat image features as video streams
+    - `patterns`: regexes to filter action & state features; images are included
+                  whenever use_videos=True, regardless of patterns.
     """
+    import re
+
+    # Gather everything the pipeline features specifies, seeded with hardware cams:
     all_features = pipeline.transform_features(initial_features)
 
-    # Intermediate storage for categorized and filtered features.
-    processed_features: dict[str, dict[str, Any]] = {
-        "action": {},
-        "observation": {},
-    }
-    images_token = OBS_IMAGES.split(".")[-1]
+    # Helper to decide which action/state keys survive the `patterns` filter:
+    def keep(key: str) -> bool:
+        if patterns is None:
+            return True
+        return any(re.search(pat, key) for pat in patterns)
 
-    # Iterate through all features transformed by the pipeline.
-    for ptype, feats in all_features.items():
-        if ptype not in [PipelineFeatureType.ACTION, PipelineFeatureType.OBSERVATION]:
+    # Start with hardware dict, injecting initial cameras if videos are ON:
+    hw: dict[str, dict[str, Any]] = {}
+    if use_videos:
+        cams = {
+            name: shape
+            for name, shape in initial_features.items()
+            if isinstance(shape, tuple) and len(shape) == 3
+        }
+        if cams:
+            hw["observation"] = dict(cams)
+
+    # Go over every feature from the pipeline and merge:
+    for full_key, ty in all_features.items():
+        if full_key.startswith("action."):
+            # action.<feat>
+            if not keep(full_key):
+                continue
+            name = full_key[len("action.") :]
+            hw.setdefault("action", {})[name] = ty
+
+        elif full_key.startswith("observation.state."):
+            # observation.state.<feat>
+            if not keep(full_key):
+                continue
+            name = full_key[len("observation.state.") :]
+            hw.setdefault("observation", {})[name] = ty
+
+        elif full_key.startswith("observation.images."):
+            # observation.images.<cam>
+            # images obey ONLY the use_videos flag, not patterns
+            if not use_videos:
+                continue
+            name = full_key[len("observation.images.") :]
+            hw.setdefault("observation", {})[name] = ty
+
+        else:
+            # anything else (e.g. policy-only features) is ignored here
             continue
 
-        for key, value in feats.items():
-            # 1. Categorize the feature.
-            is_action = ptype == PipelineFeatureType.ACTION
-            # Observations are classified as images if their key matches image-related tokens or if the shape of the feature is 3.
-            # All other observations are treated as state.
-            is_image = not is_action and (
-                (isinstance(value, tuple) and len(value) == 3)
-                or (
-                    key.startswith(f"{OBS_IMAGES}.")
-                    or key.startswith(f"{images_token}.")
-                    or f".{images_token}." in key
-                )
-            )
+    out: dict[str, dict] = {}
+    if "action" in hw:
+        out.update(hw_to_dataset_features(hw["action"], "action", use_videos))
+    if "observation" in hw:
+        out.update(hw_to_dataset_features(hw["observation"], "observation", use_videos))
 
-            # 2. Apply filtering rules.
-            if is_image and not use_videos:
-                continue
-            if not is_image and not should_keep(key, patterns):
-                continue
-
-            # 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
-            else:
-                processed_features["observation"][name] = value
-
-    # Convert the processed features into the final dataset format.
-    dataset_features = {}
-    if processed_features["action"]:
-        dataset_features.update(hw_to_dataset_features(processed_features["action"], ACTION, use_videos))
-    if processed_features["observation"]:
-        dataset_features.update(
-            hw_to_dataset_features(processed_features["observation"], "observation", use_videos)
-        )
-
-    return dataset_features
+    return out

src/lerobot/datasets/utils.py

@@ -75,20 +75,13 @@ DEFAULT_FEATURES = {
 
 
 def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
-    """Flatten a nested dictionary by joining keys with a separator.
+    """Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.
 
-    Example:
-        >>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}
-        >>> print(flatten_dict(dct))
-        {'a/b': 1, 'a/c/d': 2, 'e': 3}
-
-    Args:
-        d (dict): The dictionary to flatten.
-        parent_key (str): The base key to prepend to the keys in this level.
-        sep (str): The separator to use between keys.
-
-    Returns:
-        dict: A flattened dictionary.
+    For example:
+    ```
+    >>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}`
+    >>> print(flatten_dict(dct))
+    {"a/b": 1, "a/c/d": 2, "e": 3}
     """
     items = []
     for k, v in d.items():
@@ -101,20 +94,6 @@ def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
 
 
 def unflatten_dict(d: dict, sep: str = "/") -> dict:
-    """Unflatten a dictionary with delimited keys into a nested dictionary.
-
-    Example:
-        >>> flat_dct = {"a/b": 1, "a/c/d": 2, "e": 3}
-        >>> print(unflatten_dict(flat_dct))
-        {'a': {'b': 1, 'c': {'d': 2}}, 'e': 3}
-
-    Args:
-        d (dict): A dictionary with flattened keys.
-        sep (str): The separator used in the keys.
-
-    Returns:
-        dict: A nested dictionary.
-    """
     outdict = {}
     for key, value in d.items():
         parts = key.split(sep)
@@ -128,16 +107,6 @@ def unflatten_dict(d: dict, sep: str = "/") -> dict:
 
 
 def get_nested_item(obj: DictLike, flattened_key: str, sep: str = "/") -> Any:
-    """Access an item in a nested dictionary using a flattened key.
-
-    Args:
-        obj (DictLike): The nested dictionary-like object.
-        flattened_key (str): A key with parts separated by `sep`.
-        sep (str): The separator used in the flattened key.
-
-    Returns:
-        Any: The value from the nested dictionary.
-    """
     split_keys = flattened_key.split(sep)
     getter = obj[split_keys[0]]
     if len(split_keys) == 1:
@@ -150,19 +119,6 @@ def get_nested_item(obj: DictLike, flattened_key: str, sep: str = "/") -> Any:
 
 
 def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
-    """Serialize a dictionary containing tensors or numpy arrays to be JSON-compatible.
-
-    Converts torch.Tensor, np.ndarray, and np.generic types to lists or native Python types.
-
-    Args:
-        stats (dict): A dictionary that may contain non-serializable numeric types.
-
-    Returns:
-        dict: A dictionary with all values converted to JSON-serializable types.
-
-    Raises:
-        NotImplementedError: If a value has an unsupported type.
-    """
     serialized_dict = {}
     for key, value in flatten_dict(stats).items():
         if isinstance(value, (torch.Tensor, np.ndarray)):
@@ -177,17 +133,6 @@ def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
 
 
 def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
-    """Embed image bytes into the dataset table before saving to Parquet.
-
-    This function prepares a Hugging Face dataset for serialization by converting
-    image objects into an embedded format that can be stored in Arrow/Parquet.
-
-    Args:
-        dataset (datasets.Dataset): The input dataset, possibly containing image features.
-
-    Returns:
-        datasets.Dataset: The dataset with images embedded in the table storage.
-    """
     # Embed image bytes into the table before saving to parquet
     format = dataset.format
     dataset = dataset.with_format("arrow")
@@ -197,94 +142,38 @@ def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
 
 
 def load_json(fpath: Path) -> Any:
-    """Load data from a JSON file.
-
-    Args:
-        fpath (Path): Path to the JSON file.
-
-    Returns:
-        Any: The data loaded from the JSON file.
-    """
     with open(fpath) as f:
         return json.load(f)
 
 
 def write_json(data: dict, fpath: Path) -> None:
-    """Write data to a JSON file.
-
-    Creates parent directories if they don't exist.
-
-    Args:
-        data (dict): The dictionary to write.
-        fpath (Path): The path to the output JSON file.
-    """
     fpath.parent.mkdir(exist_ok=True, parents=True)
     with open(fpath, "w") as f:
         json.dump(data, f, indent=4, ensure_ascii=False)
 
 
 def load_jsonlines(fpath: Path) -> list[Any]:
-    """Load data from a JSON Lines file.
-
-    Args:
-        fpath (Path): Path to the JSON Lines file.
-
-    Returns:
-        list[Any]: A list of objects loaded from the file.
-    """
     with jsonlines.open(fpath, "r") as reader:
         return list(reader)
 
 
 def write_jsonlines(data: dict, fpath: Path) -> None:
-    """Write a list of dictionaries to a JSON Lines file.
-
-    Creates parent directories if they don't exist.
-
-    Args:
-        data (dict): The list of dictionaries to write.
-        fpath (Path): The path to the output JSON Lines file.
-    """
     fpath.parent.mkdir(exist_ok=True, parents=True)
     with jsonlines.open(fpath, "w") as writer:
         writer.write_all(data)
 
 
 def append_jsonlines(data: dict, fpath: Path) -> None:
-    """Append a dictionary to a JSON Lines file.
-
-    Creates parent directories if they don't exist.
-
-    Args:
-        data (dict): The dictionary to append.
-        fpath (Path): The path to the JSON Lines file.
-    """
     fpath.parent.mkdir(exist_ok=True, parents=True)
     with jsonlines.open(fpath, "a") as writer:
         writer.write(data)
 
 
 def write_info(info: dict, local_dir: Path):
-    """Write dataset info metadata to its standard file path.
-
-    Args:
-        info (dict): The dataset information dictionary.
-        local_dir (Path): The root directory of the dataset.
-    """
     write_json(info, local_dir / INFO_PATH)
 
 
 def load_info(local_dir: Path) -> dict:
-    """Load dataset info metadata from its standard file path.
-
-    Also converts shape lists to tuples for consistency.
-
-    Args:
-        local_dir (Path): The root directory of the dataset.
-
-    Returns:
-        dict: The dataset information dictionary.
-    """
     info = load_json(local_dir / INFO_PATH)
     for ft in info["features"].values():
         ft["shape"] = tuple(ft["shape"])
@@ -292,40 +181,16 @@ def load_info(local_dir: Path) -> dict:
 
 
 def write_stats(stats: dict, local_dir: Path):
-    """Serialize and write dataset statistics to their standard file path.
-
-    Args:
-        stats (dict): The statistics dictionary (can contain tensors/numpy arrays).
-        local_dir (Path): The root directory of the dataset.
-    """
     serialized_stats = serialize_dict(stats)
     write_json(serialized_stats, local_dir / STATS_PATH)
 
 
 def cast_stats_to_numpy(stats) -> dict[str, dict[str, np.ndarray]]:
-    """Recursively cast numerical values in a stats dictionary to numpy arrays.
-
-    Args:
-        stats (dict): The statistics dictionary.
-
-    Returns:
-        dict: The statistics dictionary with values cast to numpy arrays.
-    """
     stats = {key: np.array(value) for key, value in flatten_dict(stats).items()}
     return unflatten_dict(stats)
 
 
 def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
-    """Load dataset statistics and cast numerical values to numpy arrays.
-
-    Returns None if the stats file doesn't exist.
-
-    Args:
-        local_dir (Path): The root directory of the dataset.
-
-    Returns:
-        A dictionary of statistics or None if the file is not found.
-    """
     if not (local_dir / STATS_PATH).exists():
         return None
     stats = load_json(local_dir / STATS_PATH)
@@ -333,13 +198,6 @@ def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
 
 
 def write_task(task_index: int, task: dict, local_dir: Path):
-    """Write a single task to the tasks metadata file.
-
-    Args:
-        task_index (int): The index of the task.
-        task (dict): The task description dictionary.
-        local_dir (Path): The root directory of the dataset.
-    """
     task_dict = {
         "task_index": task_index,
         "task": task,
@@ -348,16 +206,6 @@ def write_task(task_index: int, task: dict, local_dir: Path):
 
 
 def load_tasks(local_dir: Path) -> tuple[dict, dict]:
-    """Load tasks from the tasks metadata file.
-
-    Args:
-        local_dir (Path): The root directory of the dataset.
-
-    Returns:
-        A tuple containing:
-        - A dictionary mapping task index to task description.
-        - A dictionary mapping task description to task index.
-    """
     tasks = load_jsonlines(local_dir / TASKS_PATH)
     tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
     task_to_task_index = {task: task_index for task_index, task in tasks.items()}
@@ -365,36 +213,15 @@ def load_tasks(local_dir: Path) -> tuple[dict, dict]:
 
 
 def write_episode(episode: dict, local_dir: Path):
-    """Write a single episode's metadata to the episodes metadata file.
-
-    Args:
-        episode (dict): The episode metadata dictionary.
-        local_dir (Path): The root directory of the dataset.
-    """
     append_jsonlines(episode, local_dir / EPISODES_PATH)
 
 
 def load_episodes(local_dir: Path) -> dict:
-    """Load episode metadata from the episodes metadata file.
-
-    Args:
-        local_dir (Path): The root directory of the dataset.
-
-    Returns:
-        dict: A dictionary mapping episode index to episode metadata.
-    """
     episodes = load_jsonlines(local_dir / EPISODES_PATH)
     return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
 
 
 def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
-    """Write statistics for a single episode to the episode stats file.
-
-    Args:
-        episode_index (int): The index of the episode.
-        episode_stats (dict): The statistics for the episode.
-        local_dir (Path): The root directory of the dataset.
-    """
     # We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
     # is a dictionary of stats and not an integer.
     episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
@@ -402,14 +229,6 @@ def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path
 
 
 def load_episodes_stats(local_dir: Path) -> dict:
-    """Load per-episode statistics from the episode stats file.
-
-    Args:
-        local_dir (Path): The root directory of the dataset.
-
-    Returns:
-        dict: A dictionary mapping episode index to its statistics dictionary.
-    """
     episodes_stats = load_jsonlines(local_dir / EPISODES_STATS_PATH)
     return {
         item["episode_index"]: cast_stats_to_numpy(item["stats"])
@@ -420,35 +239,12 @@ def load_episodes_stats(local_dir: Path) -> dict:
 def backward_compatible_episodes_stats(
     stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
 ) -> dict[str, dict[str, np.ndarray]]:
-    """Create a per-episode stats dictionary from a global stats dictionary.
-
-    This is used for backward compatibility with older datasets that only had global stats.
-
-    Args:
-        stats (dict): The global dataset statistics.
-        episodes (list[int]): A list of episode indices.
-
-    Returns:
-        dict: A dictionary mapping each episode index to the global stats.
-    """
     return dict.fromkeys(episodes, stats)
 
 
 def load_image_as_numpy(
     fpath: str | Path, dtype: np.dtype = np.float32, channel_first: bool = True
 ) -> np.ndarray:
-    """Load an image from a file into a numpy array.
-
-    Args:
-        fpath (str | Path): Path to the image file.
-        dtype (np.dtype): The desired data type of the output array. If floating,
-            pixels are scaled to [0, 1].
-        channel_first (bool): If True, converts the image to (C, H, W) format.
-            Otherwise, it remains in (H, W, C) format.
-
-    Returns:
-        np.ndarray: The image as a numpy array.
-    """
     img = PILImage.open(fpath).convert("RGB")
     img_array = np.array(img, dtype=dtype)
     if channel_first:  # (H, W, C) -> (C, H, W)
@@ -459,19 +255,10 @@ def load_image_as_numpy(
 
 
 def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
-    """Convert a batch from a Hugging Face dataset to torch tensors.
-
-    This transform function converts items from Hugging Face dataset format (pyarrow)
-    to torch tensors. Importantly, images are converted from PIL objects (H, W, C, uint8)
-    to a torch image representation (C, H, W, float32) in the range [0, 1]. Other
-    types are converted to torch.tensor.
-
-    Args:
-        items_dict (dict): A dictionary representing a batch of data from a
-            Hugging Face dataset.
-
-    Returns:
-        dict: The batch with items converted to torch tensors.
+    """Get a transform function that convert items from Hugging Face dataset (pyarrow)
+    to torch tensors. Importantly, images are converted from PIL, which corresponds to
+    a channel last representation (h w c) of uint8 type, to a torch image representation
+    with channel first (c h w) of float32 type in range [0,1].
     """
     for key in items_dict:
         first_item = items_dict[key][0]
@@ -486,14 +273,6 @@ def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
 
 
 def is_valid_version(version: str) -> bool:
-    """Check if a string is a valid PEP 440 version.
-
-    Args:
-        version (str): The version string to check.
-
-    Returns:
-        bool: True if the version string is valid, False otherwise.
-    """
     try:
         packaging.version.parse(version)
         return True
@@ -507,18 +286,6 @@ def check_version_compatibility(
     current_version: str | packaging.version.Version,
     enforce_breaking_major: bool = True,
 ) -> None:
-    """Check for version compatibility between a dataset and the current codebase.
-
-    Args:
-        repo_id (str): The repository ID for logging purposes.
-        version_to_check (str | packaging.version.Version): The version of the dataset.
-        current_version (str | packaging.version.Version): The current version of the codebase.
-        enforce_breaking_major (bool): If True, raise an error on major version mismatch.
-
-    Raises:
-        BackwardCompatibilityError: If the dataset version is from a newer, incompatible
-            major version of the codebase.
-    """
     v_check = (
         packaging.version.parse(version_to_check)
         if not isinstance(version_to_check, packaging.version.Version)
@@ -536,14 +303,7 @@ def check_version_compatibility(
 
 
 def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
-    """Return available valid versions (branches and tags) on a given Hub repo.
-
-    Args:
-        repo_id (str): The repository ID on the Hugging Face Hub.
-
-    Returns:
-        list[packaging.version.Version]: A list of valid versions found.
-    """
+    """Returns available valid versions (branches and tags) on given repo."""
     api = HfApi()
     repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
     repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
@@ -556,22 +316,9 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
 
 
 def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
-    """Return the specified version if available on repo, or the latest compatible one.
-
-    If the exact version is not found, it looks for the latest version with the
-    same major version number that is less than or equal to the target minor version.
-
-    Args:
-        repo_id (str): The repository ID on the Hugging Face Hub.
-        version (str | packaging.version.Version): The target version.
-
-    Returns:
-        str: The safe version string (e.g., "v1.2.3") to use as a revision.
-
-    Raises:
-        RevisionNotFoundError: If the repo has no version tags.
-        BackwardCompatibilityError: If only older major versions are available.
-        ForwardCompatibilityError: If only newer major versions are available.
+    """
+    Returns the version if available on repo or the latest compatible one.
+    Otherwise, will throw a `CompatibilityError`.
     """
     target_version = (
         packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
@@ -613,17 +360,6 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
 
 
 def get_hf_features_from_features(features: dict) -> datasets.Features:
-    """Convert a LeRobot features dictionary to a `datasets.Features` object.
-
-    Args:
-        features (dict): A LeRobot-style feature dictionary.
-
-    Returns:
-        datasets.Features: The corresponding Hugging Face `datasets.Features` object.
-
-    Raises:
-        ValueError: If a feature has an unsupported shape.
-    """
     hf_features = {}
     for key, ft in features.items():
         if ft["dtype"] == "video":
@@ -651,14 +387,6 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
 
 
 def _validate_feature_names(features: dict[str, dict]) -> None:
-    """Validate that feature names do not contain invalid characters.
-
-    Args:
-        features (dict): The LeRobot features dictionary.
-
-    Raises:
-        ValueError: If any feature name contains '/'.
-    """
     invalid_features = {name: ft for name, ft in features.items() if "/" in name}
     if invalid_features:
         raise ValueError(f"Feature names should not contain '/'. Found '/' in '{invalid_features}'.")
@@ -667,22 +395,6 @@ def _validate_feature_names(features: dict[str, dict]) -> None:
 def hw_to_dataset_features(
     hw_features: dict[str, type | tuple], prefix: str, use_video: bool = True
 ) -> dict[str, dict]:
-    """Convert hardware-specific features to a LeRobot dataset feature dictionary.
-
-    This function takes a dictionary describing hardware outputs (like joint states
-    or camera image shapes) and formats it into the standard LeRobot feature
-    specification.
-
-    Args:
-        hw_features (dict): Dictionary mapping feature names to their type (float for
-            joints) or shape (tuple for images).
-        prefix (str): The prefix to add to the feature keys (e.g., "observation"
-            or "action").
-        use_video (bool): If True, image features are marked as "video", otherwise "image".
-
-    Returns:
-        dict: A LeRobot features dictionary.
-    """
     features = {}
     joint_fts = {key: ftype for key, ftype in hw_features.items() if ftype is float}
     cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)}
@@ -715,20 +427,6 @@ def hw_to_dataset_features(
 def build_dataset_frame(
     ds_features: dict[str, dict], values: dict[str, Any], prefix: str
 ) -> dict[str, np.ndarray]:
-    """Construct a single data frame from raw values based on dataset features.
-
-    A "frame" is a dictionary containing all the data for a single timestep,
-    formatted as numpy arrays according to the feature specification.
-
-    Args:
-        ds_features (dict): The LeRobot dataset features dictionary.
-        values (dict): A dictionary of raw values from the hardware/environment.
-        prefix (str): The prefix to filter features by (e.g., "observation"
-            or "action").
-
-    Returns:
-        dict: A dictionary representing a single frame of data.
-    """
     frame = {}
     for key, ft in ds_features.items():
         if key in DEFAULT_FEATURES or not key.startswith(prefix):
@@ -742,21 +440,6 @@ def build_dataset_frame(
 
 
 def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFeature]:
-    """Convert dataset features to policy features.
-
-    This function transforms the dataset's feature specification into a format
-    that a policy can use, classifying features by type (e.g., visual, state,
-    action) and ensuring correct shapes (e.g., channel-first for images).
-
-    Args:
-        features (dict): The LeRobot dataset features dictionary.
-
-    Returns:
-        dict: A dictionary mapping feature keys to `PolicyFeature` objects.
-
-    Raises:
-        ValueError: If an image feature does not have a 3D shape.
-    """
     # TODO(aliberts): Implement "type" in dataset features and simplify this
     policy_features = {}
     for key, ft in features.items():
@@ -787,20 +470,12 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
     return policy_features
 
 
-def combine_feature_dicts(*dicts: dict) -> dict:
-    """Merge LeRobot grouped feature dicts.
+def merge_features(*dicts: dict) -> dict:
+    """
+    Merge LeRobot grouped feature dicts.
 
     - For 1D numeric specs (dtype not image/video/string) with "names": we merge the names and recompute the shape.
-    - For others (e.g. `observation.images.*`), the last one wins (if they are identical).
-
-    Args:
-        *dicts: A variable number of LeRobot feature dictionaries to merge.
-
-    Returns:
-        dict: A single merged feature dictionary.
-
-    Raises:
-        ValueError: If there's a dtype mismatch for a feature being merged.
+    - For others (observation.images.*), last one wins (if they are identical).
     """
     out: dict = {}
     for d in dicts:
@@ -846,18 +521,6 @@ def create_empty_dataset_info(
     use_videos: bool,
     robot_type: str | None = None,
 ) -> dict:
-    """Create a template dictionary for a new dataset's `info.json`.
-
-    Args:
-        codebase_version (str): The version of the LeRobot codebase.
-        fps (int): The frames per second of the data.
-        features (dict): The LeRobot features dictionary for the dataset.
-        use_videos (bool): Whether the dataset will store videos.
-        robot_type (str | None): The type of robot used, if any.
-
-    Returns:
-        dict: A dictionary with the initial dataset metadata.
-    """
     return {
         "codebase_version": codebase_version,
         "robot_type": robot_type,
@@ -878,18 +541,6 @@ def create_empty_dataset_info(
 def get_episode_data_index(
     episode_dicts: dict[dict], episodes: list[int] | None = None
 ) -> dict[str, torch.Tensor]:
-    """Calculate the start and end indices for each episode in a flattened dataset.
-
-    Args:
-        episode_dicts (dict): A dictionary mapping episode index to episode metadata,
-            which must contain a "length" key.
-        episodes (list[int] | None): An optional list of episode indices to consider.
-            If None, all episodes are used.
-
-    Returns:
-        dict: A dictionary with "from" and "to" keys, containing torch tensors
-            with the start and end indices for each episode.
-    """
     episode_lengths = {ep_idx: ep_dict["length"] for ep_idx, ep_dict in episode_dicts.items()}
     if episodes is not None:
         episode_lengths = {ep_idx: episode_lengths[ep_idx] for ep_idx in episodes}
@@ -909,19 +560,16 @@ def check_timestamps_sync(
     tolerance_s: float,
     raise_value_error: bool = True,
 ) -> bool:
-    """Check if timestamps are separated by (1/fps) +/- tolerance.
-
-    This check ensures that consecutive timestamps within an episode are spaced
-    correctly, accounting for possible numerical errors. It ignores the boundaries
-    between episodes.
+    """
+    This check is to make sure that each timestamp is separated from the next by (1/fps) +/- tolerance
+    to account for possible numerical error.
 
     Args:
         timestamps (np.ndarray): Array of timestamps in seconds.
         episode_indices (np.ndarray): Array indicating the episode index for each timestamp.
-        episode_data_index (dict): A dictionary that includes 'to',
+        episode_data_index (dict[str, np.ndarray]): A dictionary that includes 'to',
             which identifies indices for the end of each episode.
-        fps (int): Frames per second. Used to check the expected difference between
-            consecutive timestamps.
+        fps (int): Frames per second. Used to check the expected difference between consecutive timestamps.
         tolerance_s (float): Allowed deviation from the expected (1/fps) difference.
         raise_value_error (bool): Whether to raise a ValueError if the check fails.
 
@@ -929,8 +577,7 @@ def check_timestamps_sync(
         bool: True if all checked timestamp differences lie within tolerance, False otherwise.
 
     Raises:
-        ValueError: If `timestamps` and `episode_indices` shapes do not match, or if
-            the check fails and `raise_value_error` is True.
+        ValueError: If the check fails and `raise_value_error` is True.
     """
     if timestamps.shape != episode_indices.shape:
         raise ValueError(
@@ -981,23 +628,9 @@ def check_timestamps_sync(
 def check_delta_timestamps(
     delta_timestamps: dict[str, list[float]], fps: int, tolerance_s: float, raise_value_error: bool = True
 ) -> bool:
-    """Check if delta timestamps are multiples of 1/fps +/- tolerance.
-
-    This ensures that adding these delta timestamps to any existing timestamp in
-    the dataset will result in a value that aligns with the dataset's frame rate.
-
-    Args:
-        delta_timestamps (dict): A dictionary where values are lists of time
-            deltas in seconds.
-        fps (int): The frames per second of the dataset.
-        tolerance_s (float): The allowed tolerance in seconds.
-        raise_value_error (bool): If True, raises an error on failure.
-
-    Returns:
-        bool: True if all deltas are valid, False otherwise.
-
-    Raises:
-        ValueError: If any delta is outside the tolerance and `raise_value_error` is True.
+    """This will check if all the values in delta_timestamps are multiples of 1/fps +/- tolerance.
+    This is to ensure that these delta_timestamps added to any timestamp from a dataset will themselves be
+    actual timestamps from the dataset.
     """
     outside_tolerance = {}
     for key, delta_ts in delta_timestamps.items():
@@ -1023,15 +656,6 @@ def check_delta_timestamps(
 
 
 def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dict[str, list[int]]:
-    """Convert delta timestamps in seconds to delta indices in frames.
-
-    Args:
-        delta_timestamps (dict): A dictionary of time deltas in seconds.
-        fps (int): The frames per second of the dataset.
-
-    Returns:
-        dict: A dictionary of frame delta indices.
-    """
     delta_indices = {}
     for key, delta_ts in delta_timestamps.items():
         delta_indices[key] = [round(d * fps) for d in delta_ts]
@@ -1040,17 +664,9 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
 
 
 def cycle(iterable):
-    """Create a dataloader-safe cyclical iterator.
+    """The equivalent of itertools.cycle, but safe for Pytorch dataloaders.
 
-    This is an equivalent of `itertools.cycle` but is safe for use with
-    PyTorch DataLoaders with multiple workers.
-    See https://github.com/pytorch/pytorch/issues/23900 for details.
-
-    Args:
-        iterable: The iterable to cycle over.
-
-    Yields:
-        Items from the iterable, restarting from the beginning when exhausted.
+    See https://github.com/pytorch/pytorch/issues/23900 for information on why itertools.cycle is not safe.
     """
     iterator = iter(iterable)
     while True:
@@ -1061,14 +677,8 @@ def cycle(iterable):
 
 
 def create_branch(repo_id, *, branch: str, repo_type: str | None = None) -> None:
-    """Create a branch on an existing Hugging Face repo.
-
-    Deletes the branch if it already exists before creating it.
-
-    Args:
-        repo_id (str): The ID of the repository.
-        branch (str): The name of the branch to create.
-        repo_type (str | None): The type of the repository (e.g., "dataset").
+    """Create a branch on a existing Hugging Face repo. Delete the branch if it already
+    exists before creating it.
     """
     api = HfApi()
 
@@ -1086,20 +696,9 @@ def create_lerobot_dataset_card(
     dataset_info: dict | None = None,
     **kwargs,
 ) -> DatasetCard:
-    """Create a `DatasetCard` for a LeRobot dataset.
-
-    Keyword arguments are used to replace values in the card template.
-    Note: If specified, `license` must be a valid license identifier from
-    https://huggingface.co/docs/hub/repositories-licenses.
-
-    Args:
-        tags (list | None): A list of tags to add to the dataset card.
-        dataset_info (dict | None): The dataset's info dictionary, which will
-            be displayed on the card.
-        **kwargs: Additional keyword arguments to populate the card template.
-
-    Returns:
-        DatasetCard: The generated dataset card object.
+    """
+    Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
+    Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
     """
     card_tags = ["LeRobot"]
 
@@ -1131,16 +730,19 @@ def create_lerobot_dataset_card(
 
 
 class IterableNamespace(SimpleNamespace):
-    """A namespace object that supports both dictionary-like iteration and dot notation.
+    """
+    A namespace object that supports both dictionary-like iteration and dot notation access.
+    Automatically converts nested dictionaries into IterableNamespaces.
 
-    This class extends `SimpleNamespace` to provide dictionary-style iteration,
-    access to items via brackets (`obj["key"]`), and dictionary-like methods
-    (`items()`, `keys()`, `values()`). Nested dictionaries are recursively
-    converted to `IterableNamespace` objects.
+    This class extends SimpleNamespace to provide:
+    - Dictionary-style iteration over keys
+    - Access to items via both dot notation (obj.key) and brackets (obj["key"])
+    - Dictionary-like methods: items(), keys(), values()
+    - Recursive conversion of nested dictionaries
 
     Args:
-        dictionary (dict, optional): A dictionary to initialize the namespace with.
-        **kwargs: Additional keyword arguments to initialize the namespace.
+        dictionary: Optional dictionary to initialize the namespace
+        **kwargs: Additional keyword arguments passed to SimpleNamespace
 
     Examples:
         >>> data = {"name": "Alice", "details": {"age": 25}}
@@ -1154,16 +756,10 @@ class IterableNamespace(SimpleNamespace):
         >>> for key, value in ns.items():
         ...     print(f"{key}: {value}")
         name: Alice
-        details: <__main__.IterableNamespace object at ...>
+        details: IterableNamespace(age=25)
     """
 
     def __init__(self, dictionary: dict[str, Any] = None, **kwargs):
-        """Initialize the IterableNamespace.
-
-        Args:
-            dictionary (dict, optional): Dictionary to populate the namespace.
-            **kwargs: Keyword arguments to populate the namespace.
-        """
         super().__init__(**kwargs)
         if dictionary is not None:
             for key, value in dictionary.items():
@@ -1173,46 +769,22 @@ class IterableNamespace(SimpleNamespace):
                     setattr(self, key, value)
 
     def __iter__(self) -> Iterator[str]:
-        """Return an iterator over the keys of the namespace."""
         return iter(vars(self))
 
     def __getitem__(self, key: str) -> Any:
-        """Allow bracket-style access to attributes.
-
-        Args:
-            key (str): The name of the attribute.
-
-        Returns:
-            Any: The value of the attribute.
-        """
         return vars(self)[key]
 
     def items(self):
-        """Return a view of the namespace's (key, value) pairs."""
         return vars(self).items()
 
     def values(self):
-        """Return a view of the namespace's values."""
         return vars(self).values()
 
     def keys(self):
-        """Return a view of the namespace's keys."""
         return vars(self).keys()
 
 
 def validate_frame(frame: dict, features: dict):
-    """Validate a single data frame against the dataset's feature specification.
-
-    Checks for missing/extra features, and validates the dtype and shape of each
-    provided feature.
-
-    Args:
-        frame (dict): The data frame to validate.
-        features (dict): The LeRobot features dictionary for the dataset.
-
-    Raises:
-        ValueError: If the frame does not match the feature specification.
-    """
     expected_features = set(features) - set(DEFAULT_FEATURES)
     actual_features = set(frame)
 
@@ -1227,15 +799,6 @@ def validate_frame(frame: dict, features: dict):
 
 
 def validate_features_presence(actual_features: set[str], expected_features: set[str]):
-    """Check for missing or extra features in a frame.
-
-    Args:
-        actual_features (set[str]): The set of feature names present in the frame.
-        expected_features (set[str]): The set of feature names expected in the frame.
-
-    Returns:
-        str: An error message string if there's a mismatch, otherwise an empty string.
-    """
     error_message = ""
     missing_features = expected_features - actual_features
     extra_features = actual_features - expected_features
@@ -1251,19 +814,6 @@ def validate_features_presence(actual_features: set[str], expected_features: set
 
 
 def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray | PILImage.Image | str):
-    """Validate the dtype and shape of a single feature's value.
-
-    Args:
-        name (str): The name of the feature.
-        feature (dict): The feature specification from the LeRobot features dictionary.
-        value: The value of the feature to validate.
-
-    Returns:
-        str: An error message if validation fails, otherwise an empty string.
-
-    Raises:
-        NotImplementedError: If the feature dtype is not supported for validation.
-    """
     expected_dtype = feature["dtype"]
     expected_shape = feature["shape"]
     if is_valid_numpy_dtype_string(expected_dtype):
@@ -1279,17 +829,6 @@ def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray
 def validate_feature_numpy_array(
     name: str, expected_dtype: str, expected_shape: list[int], value: np.ndarray
 ):
-    """Validate a feature that is expected to be a numpy array.
-
-    Args:
-        name (str): The name of the feature.
-        expected_dtype (str): The expected numpy dtype as a string.
-        expected_shape (list[int]): The expected shape.
-        value (np.ndarray): The numpy array to validate.
-
-    Returns:
-        str: An error message if validation fails, otherwise an empty string.
-    """
     error_message = ""
     if isinstance(value, np.ndarray):
         actual_dtype = value.dtype
@@ -1307,18 +846,6 @@ def validate_feature_numpy_array(
 
 
 def validate_feature_image_or_video(name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image):
-    """Validate a feature that is expected to be an image or video frame.
-
-    Accepts `np.ndarray` (channel-first or channel-last) or `PIL.Image.Image`.
-
-    Args:
-        name (str): The name of the feature.
-        expected_shape (list[str]): The expected shape (C, H, W).
-        value: The image data to validate.
-
-    Returns:
-        str: An error message if validation fails, otherwise an empty string.
-    """
     # Note: The check of pixels range ([0,1] for float and [0,255] for uint8) is done by the image writer threads.
     error_message = ""
     if isinstance(value, np.ndarray):
@@ -1335,35 +862,12 @@ def validate_feature_image_or_video(name: str, expected_shape: list[str], value:
 
 
 def validate_feature_string(name: str, value: str):
-    """Validate a feature that is expected to be a string.
-
-    Args:
-        name (str): The name of the feature.
-        value (str): The value to validate.
-
-    Returns:
-        str: An error message if validation fails, otherwise an empty string.
-    """
     if not isinstance(value, str):
         return f"The feature '{name}' is expected to be of type 'str', but type '{type(value)}' provided instead.\n"
     return ""
 
 
 def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict):
-    """Validate the episode buffer before it's written to disk.
-
-    Ensures the buffer has the required keys, contains at least one frame, and
-    has features consistent with the dataset's specification.
-
-    Args:
-        episode_buffer (dict): The buffer containing data for a single episode.
-        total_episodes (int): The current total number of episodes in the dataset.
-        features (dict): The LeRobot features dictionary for the dataset.
-
-    Raises:
-        ValueError: If the buffer is invalid.
-        NotImplementedError: If the episode index is manually set and doesn't match.
-    """
     if "size" not in episode_buffer:
         raise ValueError("size key not found in episode_buffer")
 

src/lerobot/datasets/v21/convert_dataset_v20_to_v21.py

@@ -13,20 +13,24 @@
 # limitations under the License.
 
 """
-This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.0 to
-2.1. It will:
+This script converts a LeRobot dataset already pushed to the Hub from codebase version 2.0 to 2.1.
+It downloads metadata from a SOURCE dataset repo, computes/validates per-episode stats, updates
+the codebase version in `info.json`, and uploads the result to a DESTINATION dataset repo.
+It will:
 
 - Generate per-episodes stats and writes them in `episodes_stats.jsonl`
 - Check consistency between these new stats and the old ones.
 - Remove the deprecated `stats.json`.
 - Update codebase_version in `info.json`.
-- Push this new version to the hub on the 'main' branch and tags it with "v2.1".
+- Push this new version to the destination repo/branch and tag it with the current codebase version.
 
 Usage:
 
 ```bash
 python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 \
-    --repo-id=aliberts/koch_tutorial
+    --source-repo-id=namespace/source_dataset \
+    --dest-repo-id=namespace/destination_dataset \
+    --branch=main
 ```
 
 """
@@ -54,48 +58,67 @@ class SuppressWarnings:
 
 
 def convert_dataset(
-    repo_id: str,
+    source_repo_id: str,
+    dest_repo_id: str,
     branch: str | None = None,
     num_workers: int = 4,
 ):
+    # Download metadata from the source repo at v2.0
     with SuppressWarnings():
-        dataset = LeRobotDataset(repo_id, revision=V20, force_cache_sync=True)
+        dataset = LeRobotDataset(source_repo_id, revision=V20, force_cache_sync=True)
 
+    # Ensure we recompute fresh episodes stats
     if (dataset.root / EPISODES_STATS_PATH).is_file():
         (dataset.root / EPISODES_STATS_PATH).unlink()
 
+    # Compute and validate stats
     convert_stats(dataset, num_workers=num_workers)
     ref_stats = load_stats(dataset.root)
     check_aggregate_stats(dataset, ref_stats)
 
+    # Update codebase version in info.json
     dataset.meta.info["codebase_version"] = CODEBASE_VERSION
     write_info(dataset.meta.info, dataset.root)
 
-    dataset.push_to_hub(branch=branch, tag_version=False, allow_patterns="meta/")
-
-    # delete old stats.json file
-    if (dataset.root / STATS_PATH).is_file:
+    # Remove deprecated stats.json locally so it won't be uploaded
+    if (dataset.root / STATS_PATH).is_file():
         (dataset.root / STATS_PATH).unlink()
 
+    # Push only meta/ to destination repo
     hub_api = HfApi()
-    if hub_api.file_exists(
-        repo_id=dataset.repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"
-    ):
-        hub_api.delete_file(
-            path_in_repo=STATS_PATH, repo_id=dataset.repo_id, revision=branch, repo_type="dataset"
-        )
+    hub_api.create_repo(repo_id=dest_repo_id, private=False, repo_type="dataset", exist_ok=True)
+    if branch:
+        hub_api.create_branch(repo_id=dest_repo_id, branch=branch, repo_type="dataset", exist_ok=True)
 
-    hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
+    hub_api.upload_folder(
+        repo_id=dest_repo_id,
+        folder_path=str(dataset.root),
+        repo_type="dataset",
+        revision=branch,
+        allow_patterns="meta/",
+    )
+
+    # Ensure old stats.json is deleted on destination
+    if hub_api.file_exists(repo_id=dest_repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"):
+        hub_api.delete_file(path_in_repo=STATS_PATH, repo_id=dest_repo_id, revision=branch, repo_type="dataset")
+
+    # Tag destination with current codebase version
+    hub_api.create_tag(dest_repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--repo-id",
+        "--source-repo-id",
         type=str,
         required=True,
-        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset "
-        "(e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
+        help="Source dataset repo id to download from (must be v2.0).",
+    )
+    parser.add_argument(
+        "--dest-repo-id",
+        type=str,
+        required=True,
+        help="Destination dataset repo id to upload the converted metadata to.",
     )
     parser.add_argument(
         "--branch",

src/lerobot/envs/configs.py

@@ -161,71 +161,33 @@ class XarmEnv(EnvConfig):
 
 
 @dataclass
-class ImagePreprocessingConfig:
-    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
-    resize_size: tuple[int, int] | None = None
+class VideoRecordConfig:
+    """Configuration for video recording in ManiSkill environments."""
+
+    enabled: bool = False
+    record_dir: str = "videos"
+    trajectory_name: str = "trajectory"
 
 
 @dataclass
-class RewardClassifierConfig:
-    """Configuration for reward classification."""
-
-    pretrained_path: str | None = None
-    success_threshold: float = 0.5
-    success_reward: float = 1.0
-
-
-@dataclass
-class InverseKinematicsConfig:
-    """Configuration for inverse kinematics processing."""
-
-    urdf_path: str | None = None
-    target_frame_name: str | None = None
-    end_effector_bounds: dict[str, list[float]] | None = None
-    end_effector_step_sizes: dict[str, float] | None = None
-
-
-@dataclass
-class ObservationConfig:
-    """Configuration for observation processing."""
+class EnvTransformConfig:
+    """Configuration for environment wrappers."""
 
+    # ee_action_space_params: EEActionSpaceConfig = field(default_factory=EEActionSpaceConfig)
+    control_mode: str = "gamepad"
+    display_cameras: bool = False
     add_joint_velocity_to_observation: bool = False
     add_current_to_observation: bool = False
     add_ee_pose_to_observation: bool = False
-    display_cameras: bool = False
-
-
-@dataclass
-class GripperConfig:
-    """Configuration for gripper control and penalties."""
-
-    use_gripper: bool = True
-    gripper_penalty: float = 0.0
-    gripper_penalty_in_reward: bool = False
-
-
-@dataclass
-class ResetConfig:
-    """Configuration for environment reset behavior."""
-
+    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
+    resize_size: tuple[int, int] | None = None
+    control_time_s: float = 20.0
     fixed_reset_joint_positions: Any | None = None
     reset_time_s: float = 5.0
-    control_time_s: float = 20.0
-    terminate_on_success: bool = True
-
-
-@dataclass
-class HILSerlProcessorConfig:
-    """Configuration for environment processing pipeline."""
-
-    control_mode: str = "gamepad"
-    observation: ObservationConfig | None = None
-    image_preprocessing: ImagePreprocessingConfig | None = None
-    gripper: GripperConfig | None = None
-    reset: ResetConfig | None = None
-    inverse_kinematics: InverseKinematicsConfig | None = None
-    reward_classifier: RewardClassifierConfig | None = None
-    max_gripper_pos: float | None = 100.0
+    use_gripper: bool = True
+    gripper_quantization_threshold: float | None = 0.8
+    gripper_penalty: float = 0.0
+    gripper_penalty_in_reward: bool = False
 
 
 @EnvConfig.register_subclass(name="gym_manipulator")
@@ -235,10 +197,77 @@ class HILSerlRobotEnvConfig(EnvConfig):
 
     robot: RobotConfig | None = None
     teleop: TeleoperatorConfig | None = None
-    processor: HILSerlProcessorConfig = field(default_factory=HILSerlProcessorConfig)
-
+    wrapper: EnvTransformConfig | None = None
+    fps: int = 10
     name: str = "real_robot"
+    mode: str | None = None  # Either "record", "replay", None
+    repo_id: str | None = None
+    dataset_root: str | None = None
+    task: str | None = ""
+    num_episodes: int = 10  # only for record mode
+    episode: int = 0
+    device: str = "cuda"
+    push_to_hub: bool = True
+    pretrained_policy_name_or_path: str | None = None
+    reward_classifier_pretrained_path: str | None = None
+    # For the reward classifier, to record more positive examples after a success
+    number_of_steps_after_success: int = 0
 
     @property
     def gym_kwargs(self) -> dict:
         return {}
+
+
+@EnvConfig.register_subclass("hil")
+@dataclass
+class HILEnvConfig(EnvConfig):
+    """Configuration for the HIL environment."""
+
+    name: str = "PandaPickCube"
+    task: str | None = "PandaPickCubeKeyboard-v0"
+    use_viewer: bool = True
+    gripper_penalty: float = 0.0
+    use_gamepad: bool = True
+    state_dim: int = 18
+    action_dim: int = 4
+    fps: int = 100
+    episode_length: int = 100
+    video_record: VideoRecordConfig = field(default_factory=VideoRecordConfig)
+    features: dict[str, PolicyFeature] = field(
+        default_factory=lambda: {
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+            "observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
+            "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(18,)),
+        }
+    )
+    features_map: dict[str, str] = field(
+        default_factory=lambda: {
+            "action": ACTION,
+            "observation.image": OBS_IMAGE,
+            "observation.state": OBS_STATE,
+        }
+    )
+    ################# args from hilserlrobotenv
+    reward_classifier_pretrained_path: str | None = None
+    robot_config: RobotConfig | None = None
+    teleop_config: TeleoperatorConfig | None = None
+    wrapper: EnvTransformConfig | None = None
+    mode: str | None = None  # Either "record", "replay", None
+    repo_id: str | None = None
+    dataset_root: str | None = None
+    num_episodes: int = 10  # only for record mode
+    episode: int = 0
+    device: str = "cuda"
+    push_to_hub: bool = True
+    pretrained_policy_name_or_path: str | None = None
+    # For the reward classifier, to record more positive examples after a success
+    number_of_steps_after_success: int = 0
+    ############################
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {
+            "use_viewer": self.use_viewer,
+            "use_gamepad": self.use_gamepad,
+            "gripper_penalty": self.gripper_penalty,
+        }

src/lerobot/envs/factory.py

@@ -17,7 +17,7 @@ import importlib
 
 import gymnasium as gym
 
-from lerobot.envs.configs import AlohaEnv, EnvConfig, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -27,6 +27,8 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
         return PushtEnv(**kwargs)
     elif env_type == "xarm":
         return XarmEnv(**kwargs)
+    elif env_type == "hil":
+        return HILEnvConfig(**kwargs)
     else:
         raise ValueError(f"Policy type '{env_type}' is not available.")
 

src/lerobot/envs/utils.py

@@ -127,29 +127,9 @@ def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
 def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]:
     """Adds task feature to the observation dict with respect to the first environment attribute."""
     if hasattr(env.envs[0], "task_description"):
-        task_result = env.call("task_description")
-
-        if isinstance(task_result, tuple):
-            task_result = list(task_result)
-
-        if not isinstance(task_result, list):
-            raise TypeError(f"Expected task_description to return a list, got {type(task_result)}")
-        if not all(isinstance(item, str) for item in task_result):
-            raise TypeError("All items in task_description result must be strings")
-
-        observation["task"] = task_result
+        observation["task"] = env.call("task_description")
     elif hasattr(env.envs[0], "task"):
-        task_result = env.call("task")
-
-        if isinstance(task_result, tuple):
-            task_result = list(task_result)
-
-        if not isinstance(task_result, list):
-            raise TypeError(f"Expected task to return a list, got {type(task_result)}")
-        if not all(isinstance(item, str) for item in task_result):
-            raise TypeError("All items in task result must be strings")
-
-        observation["task"] = task_result
+        observation["task"] = env.call("task")
     else:  #  For envs without language instructions, e.g. aloha transfer cube and etc.
         num_envs = observation[list(observation.keys())[0]].shape[0]
         observation["task"] = ["" for _ in range(num_envs)]

src/lerobot/find_cameras.py

@@ -20,7 +20,7 @@ Helper to find the camera devices available in your system.
 Example:
 
 ```shell
-lerobot-find-cameras
+python -m lerobot.find_cameras
 ```
 """
 

src/lerobot/find_port.py

@@ -18,7 +18,7 @@ Helper to find the USB port associated with your MotorsBus.
 Example:
 
 ```shell
-lerobot-find-port
+python -m lerobot.find_port
 ```
 """
 

src/lerobot/motors/dynamixel/tables.py

@@ -107,8 +107,6 @@ X_SERIES_ENCODINGS_TABLE = {
     "Goal_PWM": X_SERIES_CONTROL_TABLE["Goal_PWM"][1],
     "Goal_Current": X_SERIES_CONTROL_TABLE["Goal_Current"][1],
     "Goal_Velocity": X_SERIES_CONTROL_TABLE["Goal_Velocity"][1],
-    "Goal_Position": X_SERIES_CONTROL_TABLE["Goal_Position"][1],
-    "Present_Position": X_SERIES_CONTROL_TABLE["Present_Position"][1],
     "Present_PWM": X_SERIES_CONTROL_TABLE["Present_PWM"][1],
     "Present_Current": X_SERIES_CONTROL_TABLE["Present_Current"][1],
     "Present_Velocity": X_SERIES_CONTROL_TABLE["Present_Velocity"][1],

src/lerobot/motors/motors_bus.py

@@ -222,7 +222,7 @@ class MotorsBus(abc.ABC):
     A MotorsBus subclass instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
     To find the port, you can run our utility script:
     ```bash
-    lerobot-find-port.py
+    python -m lerobot.find_port.py
     >>> Finding all available ports for the MotorsBus.
     >>> ["/dev/tty.usbmodem575E0032081", "/dev/tty.usbmodem575E0031751"]
     >>> Remove the usb cable from your MotorsBus and press Enter when done.
@@ -446,7 +446,7 @@ class MotorsBus(abc.ABC):
         except (FileNotFoundError, OSError, serial.SerialException) as e:
             raise ConnectionError(
                 f"\nCould not connect on port '{self.port}'. Make sure you are using the correct port."
-                "\nTry running `lerobot-find-port`\n"
+                "\nTry running `python -m lerobot.find_port`\n"
             ) from e
 
     @abc.abstractmethod

src/lerobot/policies/__init__.py

@@ -16,6 +16,7 @@ from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
 from .pi0.processor_pi0 import Pi0NewLineProcessor
+from .rlearn.configuration_rlearn import RLearNConfig as RLearNConfig
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
 from .smolvla.processor_smolvla import SmolVLANewLineProcessor
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
@@ -28,4 +29,5 @@ __all__ = [
     "SmolVLAConfig",
     "TDMPCConfig",
     "VQBeTConfig",
+    "RLearNConfig",
 ]

src/lerobot/policies/act/modeling_act.py

@@ -287,7 +287,7 @@ class ACT(nn.Module):
                                 └───────────────────────┘
     """
 
-    def __init__(self, config: ACTConfig):
+    def __init__(self, config: ACTConfig, dataset_stats=None):
         # BERT style VAE encoder with input tokens [cls, robot_state, *action_sequence].
         # The cls token forms parameters of the latent's distribution (like this [*means, *log_variances]).
         super().__init__()

src/lerobot/policies/act/processor_act.py

@@ -15,75 +15,36 @@
 # limitations under the License.
 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,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    UnnormalizerProcessor,
 )
 
 
-def make_act_pre_post_processors(
-    config: ACTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """Creates the pre- and post-processing pipelines for the ACT policy.
-
-    The pre-processing pipeline handles normalization, batching, and device placement for the model inputs.
-    The post-processing pipeline handles unnormalization and moves the model outputs back to the CPU.
-
-    Args:
-        config (ACTConfig): The ACT policy configuration object.
-        dataset_stats (dict[str, dict[str, torch.Tensor]] | None): A dictionary containing dataset
-            statistics (e.g., mean and std) used for normalization. Defaults to None.
-        preprocessor_kwargs (ProcessorKwargs | None): Extra keyword arguments to pass to the
-            preprocessor pipeline's constructor. Defaults to None.
-        postprocessor_kwargs (ProcessorKwargs | None): Extra keyword arguments to pass to the
-            postprocessor pipeline's constructor. Defaults to None.
-
-    Returns:
-        tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: A tuple containing the
-        pre-processor pipeline and the post-processor pipeline.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_act_processor(
+    config: ACTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),
+        NormalizerProcessor(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
-            device=config.device,
         ),
+        ToBatchProcessor(),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -16,81 +16,36 @@
 # limitations under the License.
 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,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    UnnormalizerProcessor,
 )
 
 
-def make_diffusion_pre_post_processors(
-    config: DiffusionConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for a diffusion policy.
-
-    The pre-processing pipeline prepares the input data for the model by:
-    1. Renaming features (if a `rename_map` is provided in `preprocessor_kwargs`).
-    2. Normalizing the input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Moving the data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1. Moving the data to the CPU.
-    2. Unnormalizing the output features to their original scale.
-
-    Args:
-        config: The configuration object for the diffusion policy,
-            containing feature definitions, normalization mappings, and device information.
-        dataset_stats: A dictionary of statistics used for normalization.
-            Defaults to None.
-        preprocessor_kwargs: Additional keyword arguments
-            for the pre-processor pipeline. Defaults to an empty dictionary.
-        postprocessor_kwargs: Additional keyword arguments
-            for the post-processor pipeline. Defaults to an empty dictionary.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_diffusion_processor(
+    config: DiffusionConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),
+        NormalizerProcessor(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
+        ToBatchProcessor(),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

src/lerobot/policies/factory.py

@@ -17,14 +17,14 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, TypedDict
+from typing import Any, TypedDict, cast
 
 import torch
+from torch import nn
 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
@@ -34,31 +34,17 @@ from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.rlearn.configuration_rlearn import RLearNConfig
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.processor import PolicyProcessorPipeline, ProcessorKwargs
+from lerobot.processor.pipeline import RobotProcessor
 
 
 def get_policy_class(name: str) -> type[PreTrainedPolicy]:
-    """
-    Retrieves a policy class by its registered name.
-
-    This function uses dynamic imports to avoid loading all policy classes into memory
-    at once, improving startup time and reducing dependencies.
-
-    Args:
-        name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
-              "vqbet", "pi0", "pi0fast", "sac", "reward_classifier", "smolvla".
-
-    Returns:
-        The policy class corresponding to the given name.
-
-    Raises:
-        NotImplementedError: If the policy name is not recognized.
-    """
+    """Get the policy's class and config class given a name (matching the policy class' `name` attribute)."""
     if name == "tdmpc":
         from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
 
@@ -95,29 +81,15 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
+    elif name == "rlearn":
+        from lerobot.policies.rlearn.modeling_rlearn import RLearNPolicy
+
+        return RLearNPolicy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
 
 def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
-    """
-    Instantiates a policy configuration object based on the policy type.
-
-    This factory function simplifies the creation of policy configuration objects by
-    mapping a string identifier to the corresponding config class.
-
-    Args:
-        policy_type: The type of the policy. Supported types include "tdmpc",
-                     "diffusion", "act", "vqbet", "pi0", "pi0fast", "sac", "smolvla",
-                     "reward_classifier".
-        **kwargs: Keyword arguments to be passed to the configuration class constructor.
-
-    Returns:
-        An instance of a `PreTrainedConfig` subclass.
-
-    Raises:
-        ValueError: If the `policy_type` is not recognized.
-    """
     if policy_type == "tdmpc":
         return TDMPCConfig(**kwargs)
     elif policy_type == "diffusion":
@@ -136,178 +108,130 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SmolVLAConfig(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
+    elif policy_type == "rlearn":
+        return RLearNConfig(**kwargs)
     else:
         raise ValueError(f"Policy type '{policy_type}' is not available.")
 
 
 class ProcessorConfigKwargs(TypedDict, total=False):
-    """
-    A TypedDict defining the keyword arguments for processor configuration.
-
-    This provides type hints for the optional arguments passed to `make_pre_post_processors`,
-    improving code clarity and enabling static analysis.
-
-    Attributes:
-        preprocessor_config_filename: The filename for the preprocessor configuration.
-        postprocessor_config_filename: The filename for the postprocessor configuration.
-        preprocessor_overrides: A dictionary of overrides for the preprocessor configuration.
-        postprocessor_overrides: A dictionary of overrides for the postprocessor configuration.
-        dataset_stats: Dataset statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the `PolicyProcessorPipeline`.
-        postprocessor_kwargs: Additional arguments for the `PolicyProcessorPipeline`.
-    """
+    """Keyword arguments for the processor config."""
 
     preprocessor_config_filename: str | None
     postprocessor_config_filename: str | None
     preprocessor_overrides: dict[str, Any] | None
     postprocessor_overrides: dict[str, Any] | None
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None
-    preprocessor_kwargs: ProcessorKwargs | None
-    postprocessor_kwargs: ProcessorKwargs | None
 
 
-def make_pre_post_processors(
+def make_processor(
     policy_cfg: PreTrainedConfig,
     pretrained_path: str | None = None,
     **kwargs: Unpack[ProcessorConfigKwargs],
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Create or load pre- and post-processor pipelines for a given policy.
+) -> tuple[RobotProcessor, RobotProcessor]:
+    """Make a processor instance for a given policy type.
 
-    This function acts as a factory. It can either load existing processor pipelines
-    from a pretrained path or create new ones from scratch based on the policy
-    configuration. Each policy type has a dedicated factory function for its
-    processors (e.g., `make_tdmpc_pre_post_processors`).
+    This function creates the appropriate processor configuration based on the policy type.
+    Each policy type has its own processor with specific preprocessing steps.
 
     Args:
-        policy_cfg: The configuration of the policy for which to create processors.
-        pretrained_path: An optional path to load pretrained processor pipelines from.
-            If provided, pipelines are loaded from this path.
-        **kwargs: Keyword arguments for processor configuration, as defined in
-            `ProcessorConfigKwargs`.
+        policy_cfg: The config of the policy to create a processor for (e.g., "act", "diffusion", etc.)
+        pretrained_path: Optional path to load a pretrained processor from. If provided, loads
+            the processor from this path instead of creating a new one.
+        **kwargs: Additional keyword arguments passed to the processor creation.
 
     Returns:
-        A tuple containing the input (pre-processor) and output (post-processor) pipelines.
+        Tuple of (input_processor, output_processor) for the policy.
 
     Raises:
-        NotImplementedError: If a processor factory is not implemented for the given
-            policy configuration type.
+        NotImplementedError: If the policy type doesn't have a processor implemented.
     """
     if pretrained_path:
-        # Extract preprocessor and postprocessor kwargs
-        preprocessor_kwargs = kwargs.get("preprocessor_kwargs", {})
-        postprocessor_kwargs = kwargs.get("postprocessor_kwargs", {})
-
+        # Load a pretrained processor
+        # TODO(azouitine): Handle this case.
         return (
-            PolicyProcessorPipeline.from_pretrained(
+            RobotProcessor.from_pretrained(
                 pretrained_model_name_or_path=pretrained_path,
-                config_filename=kwargs.get(
-                    "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
-                ),
+                config_filename=kwargs.get("preprocessor_config_filename", "robot_preprocessor.json"),
                 overrides=kwargs.get("preprocessor_overrides", {}),
-                to_transition=preprocessor_kwargs.get("to_transition"),
-                to_output=preprocessor_kwargs.get("to_output"),
             ),
-            PolicyProcessorPipeline.from_pretrained(
+            RobotProcessor.from_pretrained(
                 pretrained_model_name_or_path=pretrained_path,
-                config_filename=kwargs.get(
-                    "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
-                ),
+                config_filename=kwargs.get("postprocessor_config_filename", "robot_postprocessor.json"),
                 overrides=kwargs.get("postprocessor_overrides", {}),
-                to_transition=postprocessor_kwargs.get("to_transition"),
-                to_output=postprocessor_kwargs.get("to_output"),
             ),
         )
 
     # Create a new processor based on policy type
-    if isinstance(policy_cfg, TDMPCConfig):
-        from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
+    if policy_cfg.type == "tdmpc":
+        from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
+        from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_processor
 
-        processors = make_tdmpc_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_tdmpc_processor(
+            config=cast(TDMPCConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, DiffusionConfig):
-        from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors
+    elif policy_cfg.type == "diffusion":
+        from lerobot.policies.diffusion.processor_diffusion import make_diffusion_processor
 
-        processors = make_diffusion_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_diffusion_processor(
+            cast(DiffusionConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, ACTConfig):
-        from lerobot.policies.act.processor_act import make_act_pre_post_processors
+    elif policy_cfg.type == "act":
+        from lerobot.policies.act.processor_act import make_act_processor
 
-        processors = make_act_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_act_processor(
+            config=cast(ACTConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, VQBeTConfig):
-        from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors
+    elif policy_cfg.type == "vqbet":
+        from lerobot.policies.vqbet.processor_vqbet import make_vqbet_processor
 
-        processors = make_vqbet_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_vqbet_processor(
+            config=cast(VQBeTConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, PI0Config):
-        from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
+    elif policy_cfg.type == "pi0":
+        from lerobot.policies.pi0.processor_pi0 import make_pi0_processor
 
-        processors = make_pi0_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_pi0_processor(
+            config=cast(PI0Config, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, PI0FASTConfig):
-        from lerobot.policies.pi0fast.processor_pi0fast import make_pi0fast_pre_post_processors
+    elif policy_cfg.type == "pi0fast":
+        from lerobot.policies.pi0fast.processor_pi0fast import make_pi0fast_processor
 
-        processors = make_pi0fast_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_pi0fast_processor(
+            cast(PI0Config, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, SACConfig):
-        from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors
+    elif policy_cfg.type == "sac":
+        from lerobot.policies.sac.processor_sac import make_sac_processor
 
-        processors = make_sac_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_sac_processor(
+            cast(SACConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, RewardClassifierConfig):
+    elif policy_cfg.type == "reward_classifier":
         from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
 
         processors = make_classifier_processor(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+            cast(RewardClassifierConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
-    elif isinstance(policy_cfg, SmolVLAConfig):
-        from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors
+    elif policy_cfg.type == "smolvla":
+        from lerobot.policies.smolvla.processor_smolvla import make_smolvla_processor
 
-        processors = make_smolvla_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            preprocessor_kwargs=kwargs.get("preprocessor_kwargs"),
-            postprocessor_kwargs=kwargs.get("postprocessor_kwargs"),
+        processors = make_smolvla_processor(
+            cast(SmolVLAConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
+        )
+
+    elif policy_cfg.type == "rlearn":
+        from lerobot.policies.rlearn.processor_rlearn import make_rlearn_processor
+
+        processors = make_rlearn_processor(
+            cast(RLearNConfig, policy_cfg), dataset_stats=kwargs.get("dataset_stats")
         )
 
     else:
@@ -320,30 +244,27 @@ def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,
     env_cfg: EnvConfig | None = None,
+    episode_data_index: dict | None = None,
 ) -> PreTrainedPolicy:
-    """
-    Instantiate a policy model.
+    """Make an instance of a policy class.
 
-    This factory function handles the logic of creating a policy, which requires
-    determining the input and output feature shapes. These shapes can be derived
-    either from a `LeRobotDatasetMetadata` object or an `EnvConfig` object. The function
-    can either initialize a new policy from scratch or load a pretrained one.
+    This function exists because (for now) we need to parse features from either a dataset or an environment
+    in order to properly dimension and instantiate a policy for that dataset or environment.
 
     Args:
-        cfg: The configuration for the policy to be created. If `cfg.pretrained_path` is
-             set, the policy will be loaded with weights from that path.
-        ds_meta: Dataset metadata used to infer feature shapes and types. Also provides
-                 statistics for normalization layers.
-        env_cfg: Environment configuration used to infer feature shapes and types.
-                 One of `ds_meta` or `env_cfg` must be provided.
-
-    Returns:
-        An instantiated and device-placed policy model.
+        cfg (PreTrainedConfig): The config of the policy to make. If `pretrained_path` is set, the policy will
+            be loaded with the weights from that path.
+        ds_meta (LeRobotDatasetMetadata | None, optional): Dataset metadata to take input/output shapes and
+            statistics to use for (un)normalization of inputs/outputs in the policy. Defaults to None.
+        env_cfg (EnvConfig | None, optional): The config of a gym environment to parse features from. Must be
+            provided if ds_meta is not. Defaults to None.
 
     Raises:
-        ValueError: If both or neither of `ds_meta` and `env_cfg` are provided.
-        NotImplementedError: If attempting to use an unsupported policy-backend
-                             combination (e.g., VQBeT with 'mps').
+        ValueError: Either ds_meta or env and env_cfg must be provided.
+        NotImplementedError: if the policy.type is 'vqbet' and the policy device 'mps' (due to an incompatibility)
+
+    Returns:
+        PreTrainedPolicy: _description_
     """
     if bool(ds_meta) == bool(env_cfg):
         raise ValueError("Either one of a dataset metadata or a sim env must be provided.")
@@ -381,6 +302,10 @@ def make_policy(
     cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
     kwargs["config"] = cfg
 
+    # Pass episode_data_index for RLearN policy to calculate proper progress
+    if cfg.type == "rlearn" and episode_data_index is not None:
+        kwargs["episode_data_index"] = episode_data_index
+
     if cfg.pretrained_path:
         # Load a pretrained policy and override the config if needed (for example, if there are inference-time
         # hyperparameters that we want to vary).
@@ -391,7 +316,7 @@ def make_policy(
         policy = policy_cls(**kwargs)
 
     policy.to(cfg.device)
-    assert isinstance(policy, torch.nn.Module)
+    assert isinstance(policy, nn.Module)
 
     # policy = torch.compile(policy, mode="reduce-overhead")
 

src/lerobot/policies/normalize.py

@@ -0,0 +1,420 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+
+
+def create_stats_buffers(
+    features: dict[str, PolicyFeature],
+    norm_map: dict[str, NormalizationMode],
+    stats: dict[str, dict[str, Tensor]] | None = None,
+) -> dict[str, dict[str, nn.ParameterDict]]:
+    """
+    Create buffers per modality (e.g. "observation.image", "action") containing their mean, std, min, max
+    statistics.
+
+    Args: (see Normalize and Unnormalize)
+
+    Returns:
+        dict: A dictionary where keys are modalities and values are `nn.ParameterDict` containing
+            `nn.Parameters` set to `requires_grad=False`, suitable to not be updated during backpropagation.
+    """
+    stats_buffers = {}
+
+    for key, ft in features.items():
+        norm_mode = norm_map.get(ft.type, NormalizationMode.IDENTITY)
+        if norm_mode is NormalizationMode.IDENTITY:
+            continue
+
+        assert isinstance(norm_mode, NormalizationMode)
+
+        shape = tuple(ft.shape)
+
+        if ft.type is FeatureType.VISUAL:
+            # sanity checks
+            assert len(shape) == 3, f"number of dimensions of {key} != 3 ({shape=}"
+            c, h, w = shape
+            assert c < h and c < w, f"{key} is not channel first ({shape=})"
+            # override image shape to be invariant to height and width
+            shape = (c, 1, 1)
+
+        # Note: we initialize mean, std, min, max to infinity. They should be overwritten
+        # downstream by `stats` or `policy.load_state_dict`, as expected. During forward,
+        # we assert they are not infinity anymore.
+
+        buffer = {}
+        if norm_mode is NormalizationMode.MEAN_STD:
+            mean = torch.ones(shape, dtype=torch.float32) * torch.inf
+            std = torch.ones(shape, dtype=torch.float32) * torch.inf
+            buffer = nn.ParameterDict(
+                {
+                    "mean": nn.Parameter(mean, requires_grad=False),
+                    "std": nn.Parameter(std, requires_grad=False),
+                }
+            )
+        elif norm_mode is NormalizationMode.MIN_MAX:
+            min = torch.ones(shape, dtype=torch.float32) * torch.inf
+            max = torch.ones(shape, dtype=torch.float32) * torch.inf
+            buffer = nn.ParameterDict(
+                {
+                    "min": nn.Parameter(min, requires_grad=False),
+                    "max": nn.Parameter(max, requires_grad=False),
+                }
+            )
+
+        # TODO(aliberts, rcadene): harmonize this to only use one framework (np or torch)
+        if stats:
+            if isinstance(stats[key]["mean"], np.ndarray):
+                if norm_mode is NormalizationMode.MEAN_STD:
+                    buffer["mean"].data = torch.from_numpy(stats[key]["mean"]).to(dtype=torch.float32)
+                    buffer["std"].data = torch.from_numpy(stats[key]["std"]).to(dtype=torch.float32)
+                elif norm_mode is NormalizationMode.MIN_MAX:
+                    buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
+                    buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
+            elif isinstance(stats[key]["mean"], torch.Tensor):
+                # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
+                # tensors anywhere (for example, when we use the same stats for normalization and
+                # unnormalization). See the logic here
+                # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
+                if norm_mode is NormalizationMode.MEAN_STD:
+                    buffer["mean"].data = stats[key]["mean"].clone().to(dtype=torch.float32)
+                    buffer["std"].data = stats[key]["std"].clone().to(dtype=torch.float32)
+                elif norm_mode is NormalizationMode.MIN_MAX:
+                    buffer["min"].data = stats[key]["min"].clone().to(dtype=torch.float32)
+                    buffer["max"].data = stats[key]["max"].clone().to(dtype=torch.float32)
+            else:
+                type_ = type(stats[key]["mean"])
+                raise ValueError(f"np.ndarray or torch.Tensor expected, but type is '{type_}' instead.")
+
+        stats_buffers[key] = buffer
+    return stats_buffers
+
+
+def _no_stats_error_str(name: str) -> str:
+    return (
+        f"`{name}` is infinity. You should either initialize with `stats` as an argument, or use a "
+        "pretrained model."
+    )
+
+
+class Normalize(nn.Module):
+    """Normalizes data (e.g. "observation.image") for more stable and faster convergence during training."""
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """
+        Args:
+            shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
+            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
+            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
+            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
+            modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
+                are their normalization modes among:
+                    - "mean_std": subtract the mean and divide by standard deviation.
+                    - "min_max": map to [-1, 1] range.
+            stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
+                and values are dictionaries of statistic types and their values (e.g.
+                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
+                training the model for the first time, these statistics will overwrite the default buffers. If
+                not provided, as expected for finetuning or evaluation, the default buffers should to be
+                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
+                dataset is not needed to get the stats, since they are already in the policy state_dict.
+        """
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+        self.stats = stats
+        stats_buffers = create_stats_buffers(features, norm_map, stats)
+        for key, buffer in stats_buffers.items():
+            setattr(self, "buffer_" + key.replace(".", "_"), buffer)
+
+    # TODO(rcadene): should we remove torch.no_grad?
+    @torch.no_grad()
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        # TODO: Remove this shallow copy
+        batch = dict(batch)  # shallow copy avoids mutating the input batch
+        for key, ft in self.features.items():
+            if key not in batch:
+                # FIXME(aliberts, rcadene): This might lead to silent fail!
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            buffer = getattr(self, "buffer_" + key.replace(".", "_"))
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = buffer["mean"]
+                std = buffer["std"]
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = (batch[key] - mean) / (std + 1e-8)
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                min = buffer["min"]
+                max = buffer["max"]
+                assert not torch.isinf(min).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max).any(), _no_stats_error_str("max")
+                # normalize to [0,1]
+                batch[key] = (batch[key] - min) / (max - min + 1e-8)
+                # normalize to [-1, 1]
+                batch[key] = batch[key] * 2 - 1
+            else:
+                raise ValueError(norm_mode)
+        return batch
+
+
+class Unnormalize(nn.Module):
+    """
+    Similar to `Normalize` but unnormalizes output data (e.g. `{"action": torch.randn(b,c)}`) in their
+    original range used by the environment.
+    """
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        """
+        Args:
+            shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
+            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
+            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
+            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
+            modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
+                are their normalization modes among:
+                    - "mean_std": subtract the mean and divide by standard deviation.
+                    - "min_max": map to [-1, 1] range.
+            stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
+                and values are dictionaries of statistic types and their values (e.g.
+                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
+                training the model for the first time, these statistics will overwrite the default buffers. If
+                not provided, as expected for finetuning or evaluation, the default buffers should to be
+                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
+                dataset is not needed to get the stats, since they are already in the policy state_dict.
+        """
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+        self.stats = stats
+        # `self.buffer_observation_state["mean"]` contains `torch.tensor(state_dim)`
+        stats_buffers = create_stats_buffers(features, norm_map, stats)
+        for key, buffer in stats_buffers.items():
+            setattr(self, "buffer_" + key.replace(".", "_"), buffer)
+
+    # TODO(rcadene): should we remove torch.no_grad?
+    @torch.no_grad()
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        batch = dict(batch)  # shallow copy avoids mutating the input batch
+        for key, ft in self.features.items():
+            if key not in batch:
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            buffer = getattr(self, "buffer_" + key.replace(".", "_"))
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = buffer["mean"]
+                std = buffer["std"]
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = batch[key] * std + mean
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                min = buffer["min"]
+                max = buffer["max"]
+                assert not torch.isinf(min).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max).any(), _no_stats_error_str("max")
+                batch[key] = (batch[key] + 1) / 2
+                batch[key] = batch[key] * (max - min) + min
+            else:
+                raise ValueError(norm_mode)
+        return batch
+
+
+# TODO (azouitine): We should replace all normalization on the policies with register_buffer normalization
+#       and remove the `Normalize` and `Unnormalize` classes.
+def _initialize_stats_buffers(
+    module: nn.Module,
+    features: dict[str, PolicyFeature],
+    norm_map: dict[str, NormalizationMode],
+    stats: dict[str, dict[str, Tensor]] | None = None,
+) -> None:
+    """Register statistics buffers (mean/std or min/max) on the given *module*.
+
+    The logic matches the previous constructors of `NormalizeBuffer` and `UnnormalizeBuffer`,
+    but is factored out so it can be reused by both classes and stay in sync.
+    """
+    for key, ft in features.items():
+        norm_mode = norm_map.get(ft.type, NormalizationMode.IDENTITY)
+        if norm_mode is NormalizationMode.IDENTITY:
+            continue
+
+        shape: tuple[int, ...] = tuple(ft.shape)
+        if ft.type is FeatureType.VISUAL:
+            # reduce spatial dimensions, keep channel dimension only
+            c, *_ = shape
+            shape = (c, 1, 1)
+
+        prefix = key.replace(".", "_")
+
+        if norm_mode is NormalizationMode.MEAN_STD:
+            mean = torch.full(shape, torch.inf, dtype=torch.float32)
+            std = torch.full(shape, torch.inf, dtype=torch.float32)
+
+            if stats and key in stats and "mean" in stats[key] and "std" in stats[key]:
+                mean_data = stats[key]["mean"]
+                std_data = stats[key]["std"]
+                if isinstance(mean_data, torch.Tensor):
+                    # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
+                    # tensors anywhere (for example, when we use the same stats for normalization and
+                    # unnormalization). See the logic here
+                    # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
+                    mean = mean_data.clone().to(dtype=torch.float32)
+                    std = std_data.clone().to(dtype=torch.float32)
+                else:
+                    raise ValueError(f"Unsupported stats type for key '{key}' (expected ndarray or Tensor).")
+
+            module.register_buffer(f"{prefix}_mean", mean)
+            module.register_buffer(f"{prefix}_std", std)
+            continue
+
+        if norm_mode is NormalizationMode.MIN_MAX:
+            min_val = torch.full(shape, torch.inf, dtype=torch.float32)
+            max_val = torch.full(shape, torch.inf, dtype=torch.float32)
+
+            if stats and key in stats and "min" in stats[key] and "max" in stats[key]:
+                min_data = stats[key]["min"]
+                max_data = stats[key]["max"]
+                if isinstance(min_data, torch.Tensor):
+                    min_val = min_data.clone().to(dtype=torch.float32)
+                    max_val = max_data.clone().to(dtype=torch.float32)
+                else:
+                    raise ValueError(f"Unsupported stats type for key '{key}' (expected ndarray or Tensor).")
+
+            module.register_buffer(f"{prefix}_min", min_val)
+            module.register_buffer(f"{prefix}_max", max_val)
+            continue
+
+        raise ValueError(norm_mode)
+
+
+class NormalizeBuffer(nn.Module):
+    """Same as `Normalize` but statistics are stored as registered buffers rather than parameters."""
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+
+        _initialize_stats_buffers(self, features, norm_map, stats)
+
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        batch = dict(batch)
+        for key, ft in self.features.items():
+            if key not in batch:
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            prefix = key.replace(".", "_")
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = getattr(self, f"{prefix}_mean")
+                std = getattr(self, f"{prefix}_std")
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = (batch[key] - mean) / (std + 1e-8)
+                continue
+
+            if norm_mode is NormalizationMode.MIN_MAX:
+                min_val = getattr(self, f"{prefix}_min")
+                max_val = getattr(self, f"{prefix}_max")
+                assert not torch.isinf(min_val).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max_val).any(), _no_stats_error_str("max")
+                batch[key] = (batch[key] - min_val) / (max_val - min_val + 1e-8)
+                batch[key] = batch[key] * 2 - 1
+                continue
+
+            raise ValueError(norm_mode)
+
+        return batch
+
+
+class UnnormalizeBuffer(nn.Module):
+    """Inverse operation of `NormalizeBuffer`. Uses registered buffers for statistics."""
+
+    def __init__(
+        self,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[str, NormalizationMode],
+        stats: dict[str, dict[str, Tensor]] | None = None,
+    ):
+        super().__init__()
+        self.features = features
+        self.norm_map = norm_map
+
+        _initialize_stats_buffers(self, features, norm_map, stats)
+
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        # batch = dict(batch)
+        for key, ft in self.features.items():
+            if key not in batch:
+                continue
+
+            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
+            if norm_mode is NormalizationMode.IDENTITY:
+                continue
+
+            prefix = key.replace(".", "_")
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                mean = getattr(self, f"{prefix}_mean")
+                std = getattr(self, f"{prefix}_std")
+                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
+                assert not torch.isinf(std).any(), _no_stats_error_str("std")
+                batch[key] = batch[key] * std + mean
+                continue
+
+            if norm_mode is NormalizationMode.MIN_MAX:
+                min_val = getattr(self, f"{prefix}_min")
+                max_val = getattr(self, f"{prefix}_max")
+                assert not torch.isinf(min_val).any(), _no_stats_error_str("min")
+                assert not torch.isinf(max_val).any(), _no_stats_error_str("max")
+                batch[key] = (batch[key] + 1) / 2
+                batch[key] = batch[key] * (max_val - min_val) + min_val
+                continue
+
+            raise ValueError(norm_mode)
+
+        return batch

src/lerobot/policies/pi0/modeling_pi0.py

@@ -30,7 +30,7 @@ pip install -e ".[pi0]"
 
 Example of finetuning the pi0 pretrained model (`pi0_base` in `openpi`):
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
 --policy.path=lerobot/pi0 \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -38,7 +38,7 @@ lerobot-train \
 Example of finetuning the pi0 neural network with PaliGemma and expert Gemma
 pretrained with VLM default parameters before pi0 finetuning:
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
 --policy.type=pi0 \
 --dataset.repo_id=danaaubakirova/koch_test
 ```

src/lerobot/policies/pi0/processor_pi0.py

@@ -14,154 +14,107 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+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.configs.types import PolicyFeature
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    TokenizerProcessor,
+    UnnormalizerProcessor,
+)
+from lerobot.processor.pipeline import (
+    EnvTransition,
     ProcessorStep,
     ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
+    TransitionKey,
 )
+from lerobot.processor.rename_processor import RenameProcessor
 
 
 @ProcessorStepRegistry.register(name="pi0_new_line_processor")
-class Pi0NewLineProcessor(ComplementaryDataProcessorStep):
-    """
-    Ensures that the task description string ends with a newline character.
-
-    This processing step is required for compatibility with the PaliGemma tokenizer,
-    which expects a newline at the end of the text prompt. It handles both single
-    strings and lists of strings for the 'task' key in complementary data.
+class Pi0NewLineProcessor(ProcessorStep):
+    """Add a new line to the end of the task if it doesn't have one.
+    This is required for the PaliGemma tokenizer.
     """
 
-    def complementary_data(self, complementary_data):
-        """
-        Adds a newline to the 'task' field if it doesn't already have one.
-
-        Args:
-            complementary_data: A dictionary that may contain a 'task' key with a
-                                string or list of strings.
-
-        Returns:
-            A new dictionary with the modified 'task' field.
-        """
-        if "task" not in complementary_data:
-            return complementary_data
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        # Check if complementary_data exists
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if complementary_data is None or "task" not in complementary_data:
+            return transition
 
         task = complementary_data["task"]
         if task is None:
-            return complementary_data
-
-        new_complementary_data = dict(complementary_data)
+            return transition
 
         # Handle both string and list of strings
         if isinstance(task, str):
             # Single string: add newline if not present
             if not task.endswith("\n"):
-                new_complementary_data["task"] = f"{task}\n"
+                complementary_data["task"] = f"{task}\n"
         elif isinstance(task, list) and all(isinstance(t, str) for t in task):
             # List of strings: add newline to each if not present
-            new_complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
+            complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
         # If task is neither string nor list of strings, leave unchanged
 
-        return new_complementary_data
+        return transition
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        This step does not alter the feature definitions.
-
-        Args:
-            features: The input feature dictionary.
-
-        Returns:
-            The unchanged feature dictionary.
-        """
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        """Add tokenized task features to the features."""
         return features
 
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return state dictionary (empty for this processor)."""
+        return {}
 
-def make_pi0_pre_post_processors(
-    config: PI0Config,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the PI0 policy.
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Load state dictionary (no-op for this processor)."""
+        pass
 
-    The pre-processing pipeline prepares input data for the model by:
-    1. Renaming features to match pretrained configurations.
-    2. Normalizing input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Appending a newline character to the task description for tokenizer compatibility.
-    5. Tokenizing the text prompt using the PaliGemma tokenizer.
-    6. Moving all data to the specified device.
+    def reset(self) -> None:
+        """Reset processor state (no-op for this processor)."""
+        pass
 
-    The post-processing pipeline handles the model's output by:
-    1. Moving data to the CPU.
-    2. Unnormalizing the output features to their original scale.
+    def get_config(self) -> dict[str, Any]:
+        """Return configuration for serialization."""
+        return {}
 
-    Args:
-        config: The configuration object for the PI0 policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
 
+def make_pi0_processor(
+    config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     # Add remaining processors
     input_steps: list[ProcessorStep] = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
+        RenameProcessor(rename_map={}),  # To mimic the same processor as pretrained one
+        NormalizerProcessor(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        ToBatchProcessor(),
         Pi0NewLineProcessor(),  # Add newlines before tokenization for PaliGemma
-        TokenizerProcessorStep(
+        TokenizerProcessor(
             tokenizer_name="google/paligemma-3b-pt-224",
             max_length=config.tokenizer_max_length,
             padding_side="right",
             padding="max_length",
         ),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        DeviceProcessor(device=config.device),
     ]
 
     output_steps: list[ProcessorStep] = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
 
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -25,14 +25,14 @@ Disclaimer: It is not expected to perform as well as the original implementation
 
 Example of finetuning the pi0+FAST pretrained model (`pi0_fast_base` in `openpi`):
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
 --policy.path=lerobot/pi0fast_base \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
 
 Example of training the pi0+FAST neural network with from scratch:
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
 --policy.type=pi0fast \
 --dataset.repo_id=danaaubakirova/koch_test
 ```

src/lerobot/policies/pi0fast/processor_pi0fast.py

@@ -16,77 +16,36 @@
 
 import torch
 
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
+from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    UnnormalizerProcessor,
 )
 
 
-def make_pi0fast_pre_post_processors(
-    config: PI0FASTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the PI0Fast policy.
-
-    The pre-processing pipeline prepares input data for the model by:
-    1. Renaming features to match pretrained configurations.
-    2. Normalizing input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Moving all data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1. Moving data to the CPU.
-    2. Unnormalizing the output features to their original scale.
-
-    Args:
-        config: The configuration object for the PI0Fast policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_pi0fast_processor(
+    config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),  # To mimic the same processor as pretrained one
+        NormalizerProcessor(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
+        ToBatchProcessor(),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

src/lerobot/policies/rlearn/configuration_rlearn.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import NormalizationMode
+
+
+@PreTrainedConfig.register_subclass("rlearn")
+@dataclass
+class RLearNConfig(PreTrainedConfig):
+    """Configuration for a video-language conditioned reward model (RLearN).
+
+    Inputs:
+      - Visual frames (one or multiple cameras). Optionally a short sequence.
+      - A language instruction/goal string.
+
+    Output:
+      - Per-timestep reward logits or a single-step reward logit.
+
+    Notes:
+      - This follows the ReWiND paper architecture. It uses frozen vision/text encoders
+        (DINOv3 for vision, SigLIP2 for language) and trains a
+        lightweight temporal aggregator + head.
+    """
+
+    # Encoders - Use SigLIP2 for both vision and text (shared checkpoint)
+    vision_model_name: str = "google/siglip2-base-patch16-224"
+    text_model_name: str = "google/siglip2-base-patch16-224"
+    freeze_backbones: bool = True
+
+    # Sequence length, amount of past frames including current one to use in the temporal model
+    max_seq_len: int = 16
+    # Temporal sampling stride
+    temporal_sampling_stride: int = 3 # Open x mostly has fps 10, and rewind has seq len 16, ours is 30fps so 30/10 = 3 stride lenght to have same timeframe!
+
+    # Model dimensions and transformer
+    dim_model: int = 512
+    num_layers: int = 4
+    num_heads: int = 8
+    ff_mult: int = 4  # Feed-forward multiplier, hidden = dim_model * ff_mult
+    dropout: float = 0.05
+
+    # --- reward head options ---
+    use_categorical_rewards: bool = False      # classification over bins
+    num_reward_bins: int = 25
+    reward_min_value: float = 0.0              # for HL-Gauss range
+    reward_max_value: float = 1.0
+    use_hl_gauss_loss: bool = True             # if False -> plain regression
+    hl_gauss_num_bins: int = 25                # histogram resolution
+
+    # Inference-time subsampling and regularization
+    inference_stride: int = 1 # inference_stride is an extra, second downsampling applied in forward after window sampling/rewind. Keep it at 1 to disable extra skipping
+    frame_dropout_p: float = 0.10
+
+    # Training
+    learning_rate: float = 5e-4
+    weight_decay: float = 0.01
+    head_lr_multiplier: float = 5.0
+    logit_eps: float = 1e-4
+    regularizer_warmup_steps: int = 500
+    
+    # Performance optimizations
+    use_amp: bool = False
+    compile_model: bool = True
+
+    # ReWiND augmentation
+    rewind_prob: float = 0.3 #0.8
+    rewind_last3_prob: float = 0.0 #0.3
+    mismatch_prob: float = 0.0 #0.2
+    
+    # Normalization presets
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.MEAN_STD,
+        }
+    )
+
+    # Required path to episodes.jsonl for episode boundaries
+    episodes_jsonl_path: str | None = "meta/episodes.jsonl"
+
+    def validate_features(self) -> None:
+        # Require at least one image feature. Language is recommended but optional (can be blank).
+        if not self.image_features:
+            raise ValueError(
+                "You must provide at least one image feature for RLearN (e.g. 'observation.image')."
+            )
+
+    @property
+    def observation_delta_indices(self) -> list | None:
+        # Request a long enough context so in-window stride sampling can be >1.
+        # We ask for (max_seq_len * temporal_sampling_stride) frames ending at t=0.
+        # Example: max_seq_len=16, temporal_sampling_stride=3 → 48 deltas → ~46 frames available.
+        total_needed = self.max_seq_len * max(1, int(self.temporal_sampling_stride))
+        return list(range(1 - total_needed, 1))
+
+    @property
+    def action_delta_indices(self) -> list | None:
+        # Not an action chunking policy.
+        return None
+
+    @property
+    def reward_delta_indices(self) -> list | None:
+        # ReWiND generates progress labels on-the-fly, doesn't need reward data
+        return None
+
+    def get_optimizer_preset(self):  # type: ignore[override]
+        from lerobot.optim.optimizers import AdamWConfig
+
+        return AdamWConfig(lr=self.learning_rate, weight_decay=self.weight_decay)
+
+    def get_scheduler_preset(self):  # type: ignore[override]
+        # No scheduler by default.
+        return None

src/lerobot/policies/rlearn/eval_script.py

@@ -0,0 +1,392 @@
+#!/usr/bin/env python
+
+"""
+Standalone evaluation script for RLearN models.
+
+This script evaluates RLearN reward models on episodes from a dataset,
+generating comparison plots between ground truth rewards and model predictions.
+
+Usage:
+    python src/lerobot/policies/rlearn/eval_script.py --model MODEL_NAME --dataset DATASET_REPO --episodes N
+
+Example:
+    python src/lerobot/policies/rlearn/eval_script.py --model pepijn223/rlearn_18 --dataset pepijn223/phone_pipeline_pickup1 --episodes 2
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+# Add src to path for imports
+sys.path.append(str(Path(__file__).parent.parent.parent.parent))
+
+import warnings
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from scipy.stats import spearmanr
+from tqdm import tqdm
+
+warnings.filterwarnings("ignore")
+
+# LeRobot imports
+from lerobot.constants import OBS_IMAGE, OBS_IMAGES, OBS_LANGUAGE
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.policies.rlearn.modeling_rlearn import RLearNPolicy
+
+
+def _to_chw_float01(img):
+    """Ensure CHW float in [0,1]."""
+    if isinstance(img, np.ndarray):
+        img = torch.from_numpy(img)
+    # HWC -> CHW if needed
+    if len(img.shape) == 3 and img.shape[-1] in (1, 3, 4):
+        img = img.permute(2, 0, 1)
+    if img.dtype == torch.uint8:
+        img = img.float() / 255.0
+    else:
+        img = img.float()
+    return torch.clamp(img, 0.0, 1.0)
+
+
+def _get_language(frame_data):
+    lang = None
+    if OBS_LANGUAGE in frame_data:
+        lang = frame_data[OBS_LANGUAGE]
+        if isinstance(lang, list) and len(lang) > 0:
+            lang = lang[0]
+    elif "task" in frame_data:
+        lang = frame_data["task"]
+    return lang if isinstance(lang, str) else "No language provided"
+
+
+def _get_ground_truth_reward(frame_data):
+    """Try common keys for ground-truth reward. Return None if unavailable."""
+    for key in ("reward", "rewards", "gt_reward", "progress"):
+        if key in frame_data:
+            r = frame_data[key]
+            # unwrap single-element lists/arrays
+            if isinstance(r, (list, np.ndarray)) and np.array(r).size == 1:
+                r = float(np.array(r).reshape(-1)[0])
+            try:
+                return float(r)
+            except Exception:
+                pass
+    return None
+
+
+def extract_episode_frames_and_gt(dataset, episode_idx):
+    """Load a full episode: frames (T, C, H, W), language (str), gt_rewards (np.ndarray or None)."""
+    ep_start = dataset.episode_data_index["from"][episode_idx].item()
+    ep_end = dataset.episode_data_index["to"][episode_idx].item()
+    T = ep_end - ep_start
+
+    frames = []
+    gt_rewards = []
+    language = None
+
+    for t in range(T):
+        item = dataset[ep_start + t]
+
+        # image(s)
+        if OBS_IMAGES in item:
+            img = item[OBS_IMAGES]
+        elif OBS_IMAGE in item:
+            img = item[OBS_IMAGE]
+        else:
+            # try to find an image-like key
+            img_keys = [k for k in item.keys() if "image" in k.lower()]
+            if not img_keys:
+                continue
+            img = item[img_keys[0]]
+
+        frames.append(_to_chw_float01(img))
+
+        # language once
+        if language is None:
+            language = _get_language(item)
+
+        # ground-truth reward (optional)
+        r = _get_ground_truth_reward(item)
+        gt_rewards.append(r)
+
+    if not frames:
+        return None, None, None
+
+    frames = torch.stack(frames)  # (T, C, H, W)
+
+    # If all GT entries are None, treat as missing
+    if all(r is None for r in gt_rewards):
+        gt_rewards = None
+    else:
+        # Replace None by forward filling
+        arr = np.array([np.nan if r is None else float(r) for r in gt_rewards], dtype=float)
+        # forward/back fill
+        if np.isnan(arr[0]):
+            first_valid = np.flatnonzero(~np.isnan(arr))
+            if len(first_valid) > 0:
+                arr[0] = arr[first_valid[0]]
+            else:
+                arr[0] = 0.0
+        for i in range(1, len(arr)):
+            if np.isnan(arr[i]):
+                arr[i] = arr[i - 1]
+        gt_rewards = arr
+
+    return frames, language or "No language provided", gt_rewards
+
+
+@torch.no_grad()
+def predict_rewards_sliding(model, frames, language, max_seq_len=16, batch_size=64, device="cuda", temporal_stride: int | None = None):
+    """
+    Sliding-window prediction: for each frame i, create a window [max(0, i-L+1) .. i],
+    left-pad by repeating the first frame to length L (<= 16), and take the prediction 
+    corresponding to the current frame's position in the window.
+    Returns np.ndarray of shape (T,).
+    """
+    T = frames.shape[0]
+    cfg = getattr(model, "config", object())
+    L = int(getattr(cfg, "max_seq_len", max_seq_len))
+    L = min(L, max_seq_len)  # hard-cap at 16
+    # Use the same temporal stride as training (skip s-1 frames, take 1)
+    if temporal_stride is None:
+        temporal_stride = int(getattr(cfg, "temporal_sampling_stride", 1))
+    temporal_stride = max(1, int(temporal_stride))
+
+    # Preprocessed tensor on device
+    frames = frames.to(device)
+
+    windows = []
+    frame_positions = []  # Track which temporal position each frame should use
+    left_pad_counts = []  # Number of left-pad (OOB) frames per window
+    
+    for i in range(T):
+        # Build indices with stride s: [..., i-3, i] etc., left-padded by clamping to 0
+        idxs = [i - (L - 1 - j) * temporal_stride for j in range(L)]
+        pad_needed = sum(1 for k in idxs if k < 0)
+        clamped = [0 if k < 0 else (T - 1 if k >= T else k) for k in idxs]
+        window = frames[clamped]  # (L, C, H, W)
+
+        # Use the last temporal position (current frame) for reading model output
+        frame_pos = L - 1
+
+        windows.append(window)
+        frame_positions.append(frame_pos)
+        left_pad_counts.append(pad_needed)
+
+    preds = np.zeros(T, dtype=float)
+
+    for s in range(0, T, batch_size):
+        e = min(s + batch_size, T)
+        batch_windows = torch.stack(windows[s:e])  # (B, L, C, H, W)
+        batch_positions = frame_positions[s:e]
+
+        batch = {OBS_IMAGES: batch_windows, OBS_LANGUAGE: [language] * (e - s)}  # expects (B, L, C, H, W)
+
+        # Model returns (B, L) predictions for each temporal position
+        values = model.predict_rewards(batch)  # torch.Tensor (B, L)
+
+        # Apply eval-time padding rule: predictions for left-padded (OOB) frames are zero
+        if values.dim() == 2 and len(left_pad_counts) >= (e - s):
+            for b_idx in range(e - s):
+                pad_n = left_pad_counts[s + b_idx]
+                if pad_n > 0:
+                    values[b_idx, :pad_n] = 0.0
+
+        # Debug output removed - issue was identified and fixed
+
+        if values.dim() == 2:
+            # Extract the prediction corresponding to each frame's position in its window
+            batch_preds = []
+            for b_idx, pos in enumerate(batch_positions):
+                batch_preds.append(values[b_idx, pos].item())
+            preds[s:e] = np.array(batch_preds)
+        else:
+            # Fallback: if model returns (B,), use as is
+            preds[s:e] = values.detach().float().cpu().numpy()
+
+    return preds
+
+
+def plot_episode_eval(episode_idx, gt, pred, language, save_path=None, show=False, title_prefix="RLearN Eval"):
+    """Plot GT vs Predicted over time. Saves PNG if save_path is provided."""
+    T = len(pred)
+    x = np.arange(T)
+
+    plt.figure(figsize=(14, 8))
+    plt.plot(x, pred, linewidth=2.5, marker="o", markersize=3, label="Predicted Reward", color="blue")
+
+    if gt is not None:
+        plt.plot(x, gt, linestyle="--", linewidth=2.5, label="Ground-Truth Reward", color="orange")
+        # Correlation between GT and Pred
+        corr, p = spearmanr(gt, pred)
+        corr_str = f"ρ(GT, Pred) = {0.0 if np.isnan(corr) else corr:.3f} (p={0.0 if np.isnan(p) else p:.3f})"
+    else:
+        expected = np.linspace(0, 1, T)
+        plt.plot(x, expected, linestyle="--", linewidth=2.5, label="Expected Progress (0→1)", color="orange")
+        corr, p = spearmanr(x, pred)
+        corr_str = f"VOC-S ρ(t, Pred) = {0.0 if np.isnan(corr) else corr:.3f} (p={0.0 if np.isnan(p) else p:.3f})"
+
+    plt.title(f"{title_prefix} — Episode {episode_idx}\n{language}\n{corr_str}", fontsize=14)
+    plt.xlabel("Frame Index", fontsize=12)
+    plt.ylabel("Reward / Progress", fontsize=12)
+    plt.legend(fontsize=11)
+    plt.grid(True, alpha=0.3)
+    plt.tight_layout()
+
+    if save_path is not None:
+        plt.savefig(save_path, dpi=200, bbox_inches="tight")
+        print(f"Saved eval image to: {save_path}")
+
+    if show:
+        plt.show()
+    else:
+        plt.close()
+
+
+def eval_episode_sliding(
+    episode_idx, dataset, model, save_dir=".", device="cuda", max_seq_len=16, batch_size=64, title_prefix="RLearN Eval"
+):
+    """End-to-end: load episode, predict with sliding 16-frame windows, and save PNG."""
+    frames, language, gt = extract_episode_frames_and_gt(dataset, episode_idx)
+    if frames is None:
+        print(f"[Episode {episode_idx}] No frames found.")
+        return None
+
+    model.eval()
+
+    pred = predict_rewards_sliding(
+        model=model, frames=frames, language=language, max_seq_len=max_seq_len, batch_size=batch_size, device=device
+    )
+
+    # Basic stats
+    print(f"Episode {episode_idx}: T={len(pred)}, pred∈[{pred.min():.3f},{pred.max():.3f}]")
+    if gt is not None:
+        print(f"GT available: gt∈[{np.nanmin(gt):.3f},{np.nanmax(gt):.3f}]")
+
+    save_path = f"{save_dir}/episode_{episode_idx:04d}_eval.png"
+    plot_episode_eval(
+        episode_idx=episode_idx, gt=gt, pred=pred, language=language, save_path=save_path, show=False, title_prefix=title_prefix
+    )
+    return save_path
+
+
+def main():
+    """Main evaluation script for RLearN models."""
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Evaluate RLearN model on episodes with GT vs Predicted rewards")
+    parser.add_argument("--model", type=str, required=True, help="Model name/path (e.g., pepijn223/rlearn_mse5)")
+    parser.add_argument("--dataset", type=str, required=True, help="Dataset repo (e.g., pepijn223/phone_pipeline_pickup1)")
+    parser.add_argument("--episodes", type=int, default=5, help="Number of episodes to evaluate")
+    parser.add_argument("--output", type=str, default="./eval_results", help="Output directory for images")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu",
+        help="Device to use",
+    )
+    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for sliding window evaluation")
+
+    args = parser.parse_args()
+
+    # Create output directory
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print("🎯 RLearN Model Evaluation")
+    print("=" * 60)
+    print(f"Model: {args.model}")
+    print(f"Dataset: {args.dataset}")
+    print(f"Episodes: {args.episodes}")
+    print(f"Device: {args.device}")
+    print(f"Output: {output_dir}")
+    print("=" * 60)
+
+    try:
+        # Load dataset
+        print("📁 Loading dataset...")
+
+        dataset = LeRobotDataset(
+            repo_id=args.dataset,
+            episodes=list(range(min(args.episodes, 50))),  # Load enough episodes
+            download_videos=True,
+        )
+
+        print(f"✅ Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
+        print(f"   Features: {list(dataset.features.keys())}")
+        print(f"   FPS: {dataset.fps}")
+
+        # Load model
+        print("\n🤖 Loading model...")
+
+        model = RLearNPolicy.from_pretrained(args.model)
+        model = model.to(args.device)
+        model.eval()
+
+        print(f"✅ Model loaded on {args.device}")
+        print(f"   Parameters: {sum(p.numel() for p in model.parameters()):,}")
+        print(f"   Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
+        print(f"   Max sequence length: {model.config.max_seq_len}")
+
+        # Select episodes to evaluate
+        total_available = min(dataset.num_episodes, args.episodes)
+        episode_indices = list(range(total_available))
+
+        print(f"\n📊 Evaluating {len(episode_indices)} episodes...")
+        print("=" * 60)
+
+        # Run sliding window evaluation on each episode
+        saved_paths = []
+        for i, ep_idx in enumerate(episode_indices):
+            print(f"\n[{i+1}/{len(episode_indices)}] Processing Episode {ep_idx}")
+            print("-" * 40)
+
+            try:
+                save_path = eval_episode_sliding(
+                    episode_idx=ep_idx,
+                    dataset=dataset,
+                    model=model,
+                    save_dir=str(output_dir),
+                    device=args.device,
+                    batch_size=args.batch_size,
+                    title_prefix="RLearN Ground Truth vs Predicted",
+                )
+
+                if save_path:
+                    saved_paths.append(save_path)
+
+            except Exception as e:
+                print(f"❌ Error processing episode {ep_idx}: {e}")
+                import traceback
+
+                traceback.print_exc()
+                continue
+
+        # Summary
+        print("\n" + "=" * 60)
+        print("✅ EVALUATION COMPLETE")
+        print(f"📈 Generated {len(saved_paths)} evaluation plots")
+        print(f"📁 Results saved to: {output_dir}")
+        print("\nGenerated files:")
+        for path in saved_paths:
+            print(f"  • {path}")
+
+        if saved_paths:
+            print(f"\n💡 View the plots to compare ground truth vs predicted rewards!")
+            print(f"   Each plot shows the model's sliding 16-frame window predictions")
+            print(f"   against available ground truth rewards over the episode timeline.")
+
+        return 0
+
+    except Exception as e:
+        print(f"❌ Error during evaluation: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    exit(main())

src/lerobot/policies/rlearn/processor_rlearn.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any
+
+from lerobot.configs.types import PolicyFeature
+from lerobot.constants import OBS_LANGUAGE
+from lerobot.policies.rlearn.configuration_rlearn import RLearNConfig
+from lerobot.processor import (
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    TokenizerProcessor,
+    UnnormalizerProcessor,
+)
+from lerobot.processor.pipeline import (
+    ComplementaryDataProcessor,
+    EnvTransition,
+    ProcessorStepRegistry,
+    TransitionKey,
+)
+
+
+def make_rlearn_processor(
+    config: RLearNConfig, dataset_stats: dict[str, dict[str, Any]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
+    """Build pre/post processors for RLearN.
+
+    Responsibilities moved out of the model:
+      - Normalize inputs (images) using dataset stats
+      - Ensure batching
+      - Map complementary_data.task to observation.language when available
+      - Tokenize language into observation.language.tokens / attention_mask
+      - Move to/from device
+    """
+
+    input_steps = [
+        # No renaming by default, but keep for future extensibility
+        RenameProcessor(rename_map={}),
+        # Move heavy normalization to GPU after transfer for better parallelism
+        ToBatchProcessor(),
+        RLearnLanguageFromTaskProcessor(),
+        # Use SigLIP2 for tokenizer to keep vocab aligned with text tower
+        TokenizerProcessor(
+            tokenizer_name=config.text_model_name,
+            max_length=64,
+            padding="max_length",
+            truncation=True,
+            padding_side="right",
+        ),
+        DeviceProcessor(device=config.device),
+        # Move normalization after GPU transfer to use GPU acceleration
+        NormalizerProcessor(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+
+    output_steps = [
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+    ]
+
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
+    )
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="rlearn_language_from_task")
+class RLearnLanguageFromTaskProcessor(ComplementaryDataProcessor):
+    """Copy complementary_data['task'] into observation['observation.language'] if present.
+
+    This ensures the model can consume a raw language string when tokenization is not used,
+    while TokenizerProcessor can still create tokenized fields.
+    """
+
+    task_key: str = "task"
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:  # type: ignore[override]
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if not complementary_data or self.task_key not in complementary_data:
+            return transition
+
+        task = complementary_data.get(self.task_key)
+        if task is None:
+            return transition
+
+        # Normalize to list[str]
+        if isinstance(task, str):
+            task_list = [task]
+        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
+            task_list = task
+        else:
+            return transition
+
+        observation = transition.get(TransitionKey.OBSERVATION) or {}
+        # Do not overwrite if user already provided observation.language
+        if OBS_LANGUAGE not in observation:
+            observation[OBS_LANGUAGE] = task_list
+            transition[TransitionKey.OBSERVATION] = observation
+        return transition
+
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:  # noqa: D401
+        # Adds nothing to features; only mirrors complementary_data.task into observation
+        return features
+
+    def get_config(self) -> dict[str, Any]:
+        return {"task_key": self.task_key}

src/lerobot/policies/rlearn/rlearn_plan.md

@@ -0,0 +1,101 @@
+## General Value/Reward Learning:
+
+I want to implement a general/universal vision and language value function or reward model for robotics/video tasks. Also called a video language conditioned reward model. Integrated with already existing LeRobot code if convenient, use the LeRobot Dataset for dataset and store the reward for a frame in the lerobot frame itself.
+
+Inspired by these papers:
+
+- ReWiND; https://arxiv.org/pdf/2505.10911 (Most applicable and main paper I want to implement ideas from) and code: https://github.com/lucidrains/rewind-reward-pytorch
+- LIV; https://arxiv.org/pdf/2306.00958 (Most applicable and 2nd main paper I want to implement ideas from) and code https://github.com/penn-pal-lab/LI
+- VLC: Video-Language Critic: Transferable Reward Functions for Language-Conditioned Robotics: https://arxiv.org/pdf/2405.19988 (Most applicable and 3rd paper I want to implement ideas from) and code: https://github.com/minttusofia/video_language_critic
+
+And these papers which are also relevant:
+
+- https://www.dyna.co/dyna-1/research (Main company I want to reproduce the eventual results from)
+- vip; https://arxiv.org/pdf/2210.00030
+- uvd; https://arxiv.org/pdf/2310.08581
+- vlm in context; https://arxiv.org/pdf/2411.04549
+- https://www.youtube.com/watch?v=JfZYtpEisoM
+
+Little less relevant but still similar papers:
+
+- Learning Generalizable Robotic Reward Functions from “In-The-Wild” Human Videos,
+- XIRL: Cross-embodiment Inverse Reinforcement Learning,
+- Video-Language Critic: Transferable Reward https://arxiv.org/pdf/2405.19988
+- Functions for Language-Conditioned Robotics,
+- LORel, Language-Driven Representation Learning for Robotics https://sites.google.com/view/robotlorel
+- RoboCLIP: One Demonstration is Enough to Learn Robot Policies https://arxiv.org/pdf/2310.07899
+- Points2Rewards: learn first key points and then uses the keypoints to learn general value function/policy https://semrob.github.io/docs/2025_rss_semrob.github.io_paper20.pdf
+- Language-Driven Representation Learning for Robotics: https://arxiv.org/pdf/2302.12766v1
+- R3M: A Universal Visual Representation for Robot Manipulation: https://arxiv.org/pdf/2203.12601v3
+
+Input should be the current image or whole video and the task goal specified in text/language. Output is current reward.
+Archiutecture:
+_ inputs: video o1:T (or current o1:t), language z;
+_ DINO v3 ViT-B/16 (86M params): https://huggingface.co/facebook/dinov3-vitb16-pretrain-lvd1689m for vision encoding
+\_ sentence-transformers/all-MiniLM-L12-v2: https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2 for text encoding \* Temporal module: small causal transformer ("cross-modal sequential aggregator"), with first-frame positional embedding (to avoid position cheating), frame-dropout, and stride sampling; outputs per-timestep logits.
+
+Loss: See this chatgpt thread: https://chatgpt.com/s/t_68999a50a0b081919abc365cdd205e01
+
+Past images: (for example a reward method go to 3rd floor, has to know what floor it was on and what pas actions it did, can we attend or encorperate images of decision from history in one way?) Maybe via this paper: Learning Long-Context Diffusion Policies via Past-Token Prediction
+
+Amount of frames needed for test/generalization: 1M frames? or ~20% of IPEC-COMMUNITY/bc_z_lerobot
+
+Eval:
+Implement something like voc score , or ROC rank order correlation between reward leanredna and ev reward from sim, or use something else to do additional evaluation
+
+Ideas:
+
+- Incorporate training on multiple horizons: as in label same dataset for longer horizons: make a sandwich (long), put cheese on bread (medium) and even smaller horizons: go down or close gripper (small)
+- Incorporate navigation goals “walk towards the kitchen”, make sure we fix CLIP contrastive learning issue of positional text misunderstanding where model doesnnt learn difference between "horse right of cow" and "horse left of cow" “Move right” potentially train with more other data or even actionable world models such as Genie 3 (https://deepmind.google/discover/blog/genie-3-a-new-frontier-for-world-models/)
+
+How to use a general reward model (use cases): - Train rl policy on it - Success detection - Do exploraion - Do task via planning and search to optimize reward - Filter out bad episodes in large datasets from imitation learning
+
+Potential Datasets: (start with dataset that is most clean for this and works best with chosen way of doing evals)
+_ Epic-Kitchens-100
+_ Something-Something v. 2 Dataset https://www.qualcomm.com/developer/software/something-something-v-2-dataset
+_ Ego4D (3000 hours)
+_ Open X-Embodiment (OXE)
+\_ Agi bot world: https://huggingface.co/datasets/agibot-world/AgiBotWorld-Alpha
+
+- GalexiAI dataset: https://opengalaxea.github.io/G0/
+  _ GTEA+ Gaze: https://cbs.ic.gatech.edu/fpv/
+  _ YouCook2 dataset
+  \_ HOWTO100M: https://www.di.ens.fr/willow/research/howto100m/
+- Genie generated dataset?
+
+### TODOs:
+
+- Implement first architecture [x]
+- Implement processors [x]
+- Choose right loss metric(s) [x]
+- Make dataset with script that generated the dataset (IPEC-COMMUNITY/bc_z_lerobot) ready in lerobot format (and be able to visualize in dataset visualizer)
+  - Annotate with ReWiND-style 0→1 progress rewards [x]
+  - Visualize to check [x]
+- Implement eval score or metric that is robust and can deal with generalization/is a good metric to try different architectures. And use it in an eval jupyter notebook with visalization of the live reward next to the video for part of the dataset: VOC score and score with correct and incorrect language captions [x]
+- Do first training [x]
+- Implement on-the-fly progress label generation (no need for pre-annotated rewards) [x]
+- Try different losses
+  - Only rewind loss [x]
+  - Exactly similar to: https://github.com/lucidrains/rewind-reward-pytorch/blob/main/rewind_reward_pytorch/rewind_reward.py#L11 [x]
+  - Try DINO v2 as encoder Base 86 M: with https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2 [x]
+  - Test rewind (evaluate) [x]
+- benchmark siglip 2 vs this implementation forward pass, debug speed [x]
+- use siglip 2 [x]
+- Fix evaluation bug !!! []
+- Fix sample episode padding bug !!! []
+- Overfit on one episode []
+- Cleanup code? [] + enable language loss
+- Convert python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id=IPEC-COMMUNITY/bc_z_lerobot and train on 1 percent
+- Then on 10 percent []
+- Ablation 16 sucessive frame vs 16 frame samples with stride 2 or 4 []
+- Add more artificial text to dataset generated by vlm (google gemini) []
+  - See google gemini vlm caption [] https://gemini.google.com/app/7e332ffaf32580f2
+  - Multiple captions per video, creat method to generate as much data as possible etc [] https://arxiv.org/abs/2508.13446, https://arxiv.org/pdf/2412.04453
+- Add other datasets from OXE metioned in rewind []
+- Extend evaluation []
+- Ablation for size vision encoder, language encoder, temporal head []
+- Ablation one mlp head per frame or single mlp head []
+- Add other datasets metnioned here []
+- How can we improve spatial aware learning? solve issue of Contrastive learning and position []
+
+

src/lerobot/policies/sac/processor_sac.py

@@ -17,77 +17,36 @@
 
 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,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    UnnormalizerProcessor,
 )
 
 
-def make_sac_pre_post_processors(
-    config: SACConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the SAC policy.
-
-    The pre-processing pipeline prepares input data for the model by:
-    1. Renaming features to match pretrained configurations.
-    2. Normalizing input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Moving all data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1. Moving data to the CPU.
-    2. Unnormalizing the output features to their original scale.
-
-    Args:
-        config: The configuration object for the SAC policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_sac_processor(
+    config: SACConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),
+        NormalizerProcessor(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
+        ToBatchProcessor(),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

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

@@ -17,65 +17,26 @@ import torch
 
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.processor import (
-    DeviceProcessorStep,
-    IdentityProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
+    DeviceProcessor,
+    IdentityProcessor,
+    NormalizerProcessor,
+    RobotProcessor,
 )
 
 
 def make_classifier_processor(
-    config: RewardClassifierConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the reward classifier.
-
-    The pre-processing pipeline prepares input data for the classifier by:
-    1. Normalizing both input and output features based on dataset statistics.
-    2. Moving the data to the specified device.
-
-    The post-processing pipeline handles the classifier's output by:
-    1. Moving the data to the CPU.
-    2. Applying an identity step, as no unnormalization is needed for the output logits.
-
-    Args:
-        config: The configuration object for the RewardClassifier.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+    config: RewardClassifierConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        NormalizerProcessorStep(
+        NormalizerProcessor(
             features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        NormalizerProcessorStep(
+        NormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
-        DeviceProcessorStep(device=config.device),
+        DeviceProcessor(device=config.device),
     ]
-    output_steps = [DeviceProcessorStep(device="cpu"), IdentityProcessorStep()]
-
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name="classifier_preprocessor",
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name="classifier_postprocessor",
-            **postprocessor_kwargs,
-        ),
+    output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()]
+    return RobotProcessor(steps=input_steps, name="classifier_preprocessor"), RobotProcessor(
+        steps=output_steps, name="classifier_postprocessor"
     )

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -28,7 +28,7 @@ pip install -e ".[smolvla]"
 
 Example of finetuning the smolvla pretrained model (`smolvla_base`):
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
 --policy.path=lerobot/smolvla_base \
 --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
 --batch_size=64 \
@@ -38,7 +38,7 @@ lerobot-train \
 Example of finetuning a smolVLA. SmolVLA is composed of a pretrained VLM,
 and an action expert.
 ```bash
-lerobot-train \
+python -m lerobot.scripts.train \
 --policy.type=smolvla \
 --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
 --batch_size=64 \

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -13,131 +13,97 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+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.configs.types import PolicyFeature
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    TokenizerProcessor,
+    UnnormalizerProcessor,
 )
+from lerobot.processor.pipeline import EnvTransition, ProcessorStep, ProcessorStepRegistry, TransitionKey
 
 
-def make_smolvla_pre_post_processors(
-    config: SmolVLAConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the SmolVLA policy.
-
-    The pre-processing pipeline prepares input data for the model by:
-    1.  Renaming features to match pretrained configurations.
-    2.  Normalizing input and output features based on dataset statistics.
-    3.  Adding a batch dimension.
-    4.  Ensuring the language task description ends with a newline character.
-    5.  Tokenizing the language task description.
-    6.  Moving all data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1.  Moving data to the CPU.
-    2.  Unnormalizing the output actions to their original scale.
-
-    Args:
-        config: The configuration object for the SmolVLA policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_smolvla_processor(
+    config: SmolVLAConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
+        RenameProcessor(rename_map={}),  # To mimic the same processor as pretrained one
+        NormalizerProcessor(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        ToBatchProcessor(),
         SmolVLANewLineProcessor(),
-        TokenizerProcessorStep(
+        TokenizerProcessor(
             tokenizer_name=config.vlm_model_name,
             padding=config.pad_language_to,
             padding_side="right",
             max_length=config.tokenizer_max_length,
         ),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )
 
 
 @ProcessorStepRegistry.register(name="smolvla_new_line_processor")
-class SmolVLANewLineProcessor(ComplementaryDataProcessorStep):
-    """
-    A processor step that ensures the 'task' description ends with a newline character.
+class SmolVLANewLineProcessor(ProcessorStep):
+    """Add a new line to the end of the task if it doesn't have one."""
 
-    This step is necessary for certain tokenizers (e.g., PaliGemma) that expect a
-    newline at the end of the prompt. It handles both single string tasks and lists
-    of string tasks.
-    """
-
-    def complementary_data(self, complementary_data):
-        if "task" not in complementary_data:
-            return complementary_data
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        # Check if complementary_data exists
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if complementary_data is None or "task" not in complementary_data:
+            return transition
 
         task = complementary_data["task"]
         if task is None:
-            return complementary_data
-
-        new_complementary_data = dict(complementary_data)
+            return transition
 
         # Handle both string and list of strings
         if isinstance(task, str):
             # Single string: add newline if not present
             if not task.endswith("\n"):
-                new_complementary_data["task"] = f"{task}\n"
+                complementary_data["task"] = f"{task}\n"
         elif isinstance(task, list) and all(isinstance(t, str) for t in task):
             # List of strings: add newline to each if not present
-            new_complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
+            complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
         # If task is neither string nor list of strings, leave unchanged
 
-        return new_complementary_data
+        return transition
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        """Adds nothing to the features."""
         return features
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return state dictionary (empty for this processor)."""
+        return {}
+
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Load state dictionary (no-op for this processor)."""
+        pass
+
+    def reset(self) -> None:
+        """Reset processor state (no-op for this processor)."""
+        pass
+
+    def get_config(self) -> dict[str, Any]:
+        """Return configuration for serialization."""
+        return {}

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -16,77 +16,36 @@
 # limitations under the License.
 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,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    UnnormalizerProcessor,
 )
 
 
-def make_tdmpc_pre_post_processors(
-    config: TDMPCConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the TDMPC policy.
-
-    The pre-processing pipeline prepares input data for the model by:
-    1. Renaming features to match pretrained configurations.
-    2. Normalizing input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Moving all data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1. Moving data to the CPU.
-    2. Unnormalizing the output features to their original scale.
-
-    Args:
-        config: The configuration object for the TDMPC policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_tdmpc_processor(
+    config: TDMPCConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),
+        NormalizerProcessor(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
+        ToBatchProcessor(),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -17,77 +17,36 @@
 # limitations under the License.
 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,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    DeviceProcessor,
+    NormalizerProcessor,
+    RenameProcessor,
+    RobotProcessor,
+    ToBatchProcessor,
+    UnnormalizerProcessor,
 )
 
 
-def make_vqbet_pre_post_processors(
-    config: VQBeTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    preprocessor_kwargs: ProcessorKwargs | None = None,
-    postprocessor_kwargs: ProcessorKwargs | None = None,
-) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
-    """
-    Constructs pre-processor and post-processor pipelines for the VQ-BeT policy.
-
-    The pre-processing pipeline prepares input data for the model by:
-    1. Renaming features, allowing customization to match pretrained configurations.
-    2. Normalizing input and output features based on dataset statistics.
-    3. Adding a batch dimension.
-    4. Moving all data to the specified device.
-
-    The post-processing pipeline handles the model's output by:
-    1. Moving data to the CPU.
-    2. Unnormalizing the output features to their original scale.
-
-    Args:
-        config: The configuration object for the VQ-BeT policy.
-        dataset_stats: A dictionary of statistics for normalization.
-        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
-        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-    if preprocessor_kwargs is None:
-        preprocessor_kwargs = {}
-    if postprocessor_kwargs is None:
-        postprocessor_kwargs = {}
-
+def make_vqbet_processor(
+    config: VQBeTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+) -> tuple[RobotProcessor, RobotProcessor]:
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # Let the possibility to the user to rename the keys
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),  # Let the possibility to the user to rename the keys
+        NormalizerProcessor(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
+        ToBatchProcessor(),
+        DeviceProcessor(device=config.device),
     ]
     output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(
             features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
         ),
     ]
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-            **preprocessor_kwargs,
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            **postprocessor_kwargs,
-        ),
+    return RobotProcessor(steps=input_steps, name="robot_preprocessor"), RobotProcessor(
+        steps=output_steps, name="robot_postprocessor"
     )

src/lerobot/processor/__init__.py

@@ -14,90 +14,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .batch_processor import AddBatchDimensionProcessorStep
-from .converters import (
-    batch_to_transition,
-    create_transition,
-    merge_transitions,
-    transition_to_batch,
-    transition_to_dataset_frame,
-)
-from .core import EnvTransition, TransitionKey
-from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
-from .device_processor import DeviceProcessorStep
-from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep
-from .hil_processor import (
-    AddTeleopActionAsComplimentaryDataStep,
-    AddTeleopEventsAsInfoStep,
-    GripperPenaltyProcessorStep,
-    ImageCropResizeProcessorStep,
-    InterventionActionProcessorStep,
-    RewardClassifierProcessorStep,
-    TimeLimitProcessorStep,
-)
-from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep
-from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats
-from .observation_processor import VanillaObservationProcessorStep
+from .batch_processor import ToBatchProcessor
+from .device_processor import DeviceProcessor
+from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats
+from .observation_processor import VanillaObservationProcessor
 from .pipeline import (
-    ActionProcessorStep,
-    ComplementaryDataProcessorStep,
-    DataProcessorPipeline,
-    DoneProcessorStep,
-    IdentityProcessorStep,
-    InfoProcessorStep,
-    ObservationProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
+    ActionProcessor,
+    DoneProcessor,
+    EnvTransition,
+    IdentityProcessor,
+    InfoProcessor,
+    ObservationProcessor,
     ProcessorStep,
     ProcessorStepRegistry,
-    RewardProcessorStep,
-    RobotProcessorPipeline,
-    TruncatedProcessorStep,
+    RewardProcessor,
+    RobotProcessor,
+    TransitionKey,
+    TruncatedProcessor,
 )
-from .rename_processor import RenameObservationsProcessorStep
-from .tokenizer_processor import TokenizerProcessorStep
+from .rename_processor import RenameProcessor
+from .tokenizer_processor import TokenizerProcessor
 
 __all__ = [
-    "ActionProcessorStep",
-    "AddTeleopActionAsComplimentaryDataStep",
-    "AddTeleopEventsAsInfoStep",
-    "ComplementaryDataProcessorStep",
-    "batch_to_transition",
-    "create_transition",
-    "DeviceProcessorStep",
-    "DoneProcessorStep",
+    "ActionProcessor",
+    "DeviceProcessor",
+    "DoneProcessor",
     "EnvTransition",
-    "GripperPenaltyProcessorStep",
+    "IdentityProcessor",
+    "InfoProcessor",
+    "NormalizerProcessor",
+    "UnnormalizerProcessor",
     "hotswap_stats",
-    "IdentityProcessorStep",
-    "ImageCropResizeProcessorStep",
-    "InfoProcessorStep",
-    "InterventionActionProcessorStep",
-    "JointVelocityProcessorStep",
-    "MapDeltaActionToRobotActionStep",
-    "MapTensorToDeltaActionDictStep",
-    "merge_transitions",
-    "MotorCurrentProcessorStep",
-    "NormalizerProcessorStep",
-    "Numpy2TorchActionProcessorStep",
-    "ObservationProcessorStep",
-    "PolicyProcessorPipeline",
-    "ProcessorKwargs",
+    "ObservationProcessor",
     "ProcessorStep",
     "ProcessorStepRegistry",
-    "RenameObservationsProcessorStep",
-    "RewardClassifierProcessorStep",
-    "RewardProcessorStep",
-    "DataProcessorPipeline",
-    "TimeLimitProcessorStep",
-    "AddBatchDimensionProcessorStep",
-    "RobotProcessorPipeline",
-    "TokenizerProcessorStep",
-    "Torch2NumpyActionProcessorStep",
-    "transition_to_batch",
-    "transition_to_dataset_frame",
+    "RenameProcessor",
+    "RewardProcessor",
+    "RobotProcessor",
+    "ToBatchProcessor",
+    "TokenizerProcessor",
     "TransitionKey",
-    "TruncatedProcessorStep",
-    "UnnormalizerProcessorStep",
-    "VanillaObservationProcessorStep",
+    "TruncatedProcessor",
+    "VanillaObservationProcessor",
 ]

src/lerobot/processor/batch_processor.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 # Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,92 +11,66 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from dataclasses import dataclass
+from typing import Any
 
-"""
-This script defines processor steps for adding a batch dimension to various components of an environment transition.
-
-These steps are designed to process actions, observations, and complementary data, making them suitable for batch processing by adding a leading dimension. This is a common requirement before feeding data into a neural network model.
-"""
-
-from dataclasses import dataclass, field
-
+import torch
 from torch import Tensor
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs.types import PolicyFeature
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-
-from .core import EnvTransition
-from .pipeline import (
-    ActionProcessorStep,
-    ComplementaryDataProcessorStep,
-    ObservationProcessorStep,
-    ProcessorStep,
-    ProcessorStepRegistry,
-)
+from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
 
 
 @dataclass
-@ProcessorStepRegistry.register(name="to_batch_processor_action")
-class AddBatchDimensionActionStep(ActionProcessorStep):
-    """
-    Processor step to add a batch dimension to a 1D tensor action.
+@ProcessorStepRegistry.register(name="to_batch_processor")
+class ToBatchProcessor:
+    """Processor that adds batch dimensions to observations and actions when needed.
 
-    This is useful for creating a batch of size 1 from a single action sample.
+    This processor ensures that observations and actions have proper batch dimensions for model processing:
+
+    - For state observations (observation.state, observation.environment_state):
+      Adds batch dimension (unsqueeze at dim=0) if tensor is 1-dimensional
+
+    - For image observations (observation.image, observation.images.*):
+      Adds batch dimension (unsqueeze at dim=0) if tensor is 3-dimensional (H, W, C)
+
+    - For actions:
+      Adds batch dimension (unsqueeze at dim=0) if tensor is 1-dimensional
+
+    - For task field in complementary data:
+      Wraps string task in a list to add batch dimension
+      (task must be a string or list of strings)
+
+    This is useful when processing single transitions that need to be batched for
+    model inference or when converting from unbatched environment outputs to
+    batched model inputs.
+
+    The processor only modifies tensors that need batching and leaves already
+    batched tensors unchanged.
+
+    Example:
+        ```python
+        # State: (7,) -> (1, 7)
+        # Image: (224, 224, 3) -> (1, 224, 224, 3)
+        # Action: (4,) -> (1, 4)
+        # Task: "pick_cube" -> ["pick_cube"]
+        # Already batched: (1, 7) -> (1, 7) [unchanged]
+        ```
     """
 
-    def action(self, action: Tensor) -> Tensor:
-        """
-        Adds a batch dimension to the action if it's a 1D tensor.
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        self._process_observation(transition)
+        self._process_action(transition)
+        self._process_complementary_data(transition)
+        return transition
 
-        Args:
-            action: The action tensor.
+    def _process_observation(self, transition: EnvTransition) -> None:
+        """Process observation component in-place, adding batch dimensions where needed."""
+        observation = transition.get(TransitionKey.OBSERVATION)
+        if observation is None:
+            return
 
-        Returns:
-            The action tensor with an added batch dimension.
-        """
-        if not isinstance(action, Tensor) or action.dim() != 1:
-            return action
-        return action.unsqueeze(0)
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Returns the input features unchanged.
-
-        Adding a batch dimension does not alter the feature definition.
-
-        Args:
-            features: A dictionary of policy features.
-
-        Returns:
-            The original dictionary of policy features.
-        """
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="to_batch_processor_observation")
-class AddBatchDimensionObservationStep(ObservationProcessorStep):
-    """
-    Processor step to add a batch dimension to observations.
-
-    It handles different types of observations:
-    - State vectors (1D tensors).
-    - Single images (3D tensors).
-    - Dictionaries of multiple images (3D tensors).
-    """
-
-    def observation(self, observation: dict[str, Tensor]) -> dict[str, Tensor]:
-        """
-        Adds a batch dimension to tensor-based observations in the observation dictionary.
-
-        Args:
-            observation: The observation dictionary.
-
-        Returns:
-            The observation dictionary with batch dimensions added to tensors.
-        """
         # Process state observations - add batch dim if 1D
         for state_key in [OBS_STATE, OBS_ENV_STATE]:
             if state_key in observation:
@@ -116,46 +88,19 @@ class AddBatchDimensionObservationStep(ObservationProcessorStep):
         for key, value in observation.items():
             if key.startswith(f"{OBS_IMAGES}.") and isinstance(value, Tensor) and value.dim() == 3:
                 observation[key] = value.unsqueeze(0)
-        return observation
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Returns the input features unchanged.
+    def _process_action(self, transition: EnvTransition) -> None:
+        """Process action component in-place, adding batch dimension if needed."""
+        action = transition.get(TransitionKey.ACTION)
+        if action is not None and isinstance(action, Tensor) and action.dim() == 1:
+            transition[TransitionKey.ACTION] = action.unsqueeze(0)
 
-        Adding a batch dimension does not alter the feature definition.
+    def _process_complementary_data(self, transition: EnvTransition) -> None:
+        """Process complementary data in-place, handling task field batching."""
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if complementary_data is None:
+            return
 
-        Args:
-            features: A dictionary of policy features.
-
-        Returns:
-            The original dictionary of policy features.
-        """
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="to_batch_processor_complementary_data")
-class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
-    """
-    Processor step to add a batch dimension to complementary data fields.
-
-    Handles specific keys like 'task', 'index', and 'task_index' to make them batched.
-    - 'task' (str) is wrapped in a list.
-    - 'index' and 'task_index' (0D tensors) get a batch dimension.
-    """
-
-    def complementary_data(self, complementary_data: dict) -> dict:
-        """
-        Adds a batch dimension to specific fields in the complementary data dictionary.
-
-        Args:
-            complementary_data: The complementary data dictionary.
-
-        Returns:
-            The complementary data dictionary with batch dimensions added.
-        """
         # Process task field - wrap string in list to add batch dimension
         if "task" in complementary_data:
             task_value = complementary_data["task"]
@@ -173,78 +118,22 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
             task_index_value = complementary_data["task_index"]
             if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
                 complementary_data["task_index"] = task_index_value.unsqueeze(0)
-        return complementary_data
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Returns the input features unchanged.
+    def get_config(self) -> dict[str, Any]:
+        """Return configuration for serialization."""
+        return {}
 
-        Adding a batch dimension does not alter the feature definition.
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return state dictionary (empty for this processor)."""
+        return {}
 
-        Args:
-            features: A dictionary of policy features.
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Load state dictionary (no-op for this processor)."""
+        pass
 
-        Returns:
-            The original dictionary of policy features.
-        """
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="to_batch_processor")
-class AddBatchDimensionProcessorStep(ProcessorStep):
-    """
-    A composite processor step that adds a batch dimension to the entire environment transition.
-
-    This step combines individual processors for actions, observations, and complementary data
-    to create a batched transition (batch size 1) from a single-instance transition.
-
-    Attributes:
-        to_batch_action_processor: Processor for the action component.
-        to_batch_observation_processor: Processor for the observation component.
-        to_batch_complementary_data_processor: Processor for the complementary data component.
-    """
-
-    to_batch_action_processor: AddBatchDimensionActionStep = field(
-        default_factory=AddBatchDimensionActionStep
-    )
-    to_batch_observation_processor: AddBatchDimensionObservationStep = field(
-        default_factory=AddBatchDimensionObservationStep
-    )
-    to_batch_complementary_data_processor: AddBatchDimensionComplementaryDataStep = field(
-        default_factory=AddBatchDimensionComplementaryDataStep
-    )
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """
-        Applies the batching process to all relevant parts of an environment transition.
-
-        Args:
-            transition: The environment transition to process.
-
-        Returns:
-            The environment transition with a batch dimension added.
-        """
-        transition = self.to_batch_action_processor(transition)
-        transition = self.to_batch_observation_processor(transition)
-        transition = self.to_batch_complementary_data_processor(transition)
-        return transition
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Returns the input features unchanged.
-
-        Adding a batch dimension does not alter the feature definition.
-
-        Args:
-            features: A dictionary of policy features.
-
-        Returns:
-            The original dictionary of policy features.
-        """
-        # NOTE: We ignore the batch dimension when transforming features
+    def reset(self) -> None:
+        """Reset processor state (no-op for this processor)."""
+        pass
+
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features

src/lerobot/processor/converters.py

@@ -16,508 +16,210 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from copy import deepcopy
-from functools import singledispatch
 from typing import Any
 
 import numpy as np
 import torch
+from scipy.spatial.transform import Rotation
 
-from lerobot.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD, TRUNCATED
-
-from .core import EnvTransition, TransitionKey
+from .pipeline import EnvTransition, TransitionKey
 
 
-@singledispatch
-def to_tensor(
-    value: Any,
-    *,
-    dtype: torch.dtype | None = torch.float32,
-    device: torch.device | str | None = None,
-) -> torch.Tensor:
-    """
-    Convert various data types to PyTorch tensors with configurable options.
-
-    This is a unified tensor conversion function using single dispatch to handle
-    different input types appropriately.
-
-    Args:
-        value: Input value to convert (tensor, array, scalar, sequence, etc.).
-        dtype: Target tensor dtype. If None, preserves original dtype.
-        device: Target device for the tensor.
-
-    Returns:
-        A PyTorch tensor.
-
-    Raises:
-        TypeError: If the input type is not supported.
-    """
-    raise TypeError(f"Unsupported type for tensor conversion: {type(value)}")
+def _to_tensor(x: torch.Tensor | np.ndarray | Sequence[int | float]):
+    if isinstance(x, torch.Tensor):
+        return x
+    if isinstance(x, np.ndarray):
+        # Keep images (uint8 HWC) and python objects as-is
+        if x.dtype == np.uint8 or x.dtype == np.object_:
+            return x
+        # Scalars/arrays to float32 tensor
+        return torch.as_tensor(x, dtype=torch.float32)
+    # Anything else to float32 tensor
+    return torch.as_tensor(x, dtype=torch.float32)
 
 
-@to_tensor.register(torch.Tensor)
-def _(value: torch.Tensor, *, dtype=torch.float32, device=None, **kwargs) -> torch.Tensor:
-    """Handle conversion for existing PyTorch tensors."""
-    if dtype is not None:
-        value = value.to(dtype=dtype)
-    if device is not None:
-        value = value.to(device=device)
-    return value
-
-
-@to_tensor.register(np.ndarray)
-def _(
-    value: np.ndarray,
-    *,
-    dtype=torch.float32,
-    device=None,
-    **kwargs,
-) -> torch.Tensor:
-    """Handle conversion for numpy arrays."""
-    # Check for numpy scalars (0-dimensional arrays) and treat them as scalars.
-    if value.ndim == 0:
-        # Numpy scalars should be converted to 0-dimensional tensors.
-        scalar_value = value.item()
-        return torch.tensor(scalar_value, dtype=dtype, device=device)
-
-    # Create tensor from numpy array.
-    tensor = torch.from_numpy(value)
-
-    # Apply dtype and device conversion if specified.
-    if dtype is not None:
-        tensor = tensor.to(dtype=dtype)
-    if device is not None:
-        tensor = tensor.to(device=device)
-
-    return tensor
-
-
-@to_tensor.register(int)
-@to_tensor.register(float)
-@to_tensor.register(np.integer)
-@to_tensor.register(np.floating)
-def _(value, *, dtype=torch.float32, device=None, **kwargs) -> torch.Tensor:
-    """Handle conversion for scalar values including numpy scalars."""
-    return torch.tensor(value, dtype=dtype, device=device)
-
-
-@to_tensor.register(list)
-@to_tensor.register(tuple)
-def _(value: Sequence, *, dtype=torch.float32, device=None, **kwargs) -> torch.Tensor:
-    """Handle conversion for sequences (lists, tuples)."""
-    return torch.tensor(value, dtype=dtype, device=device)
-
-
-@to_tensor.register(dict)
-def _(value: dict, *, device=None, **kwargs) -> dict:
-    """Handle conversion for dictionaries by recursively converting their values to tensors."""
-    if not value:
-        return {}
-
-    result = {}
-    for key, sub_value in value.items():
-        if sub_value is None:
-            continue
-
-        if isinstance(sub_value, dict):
-            # Recursively process nested dictionaries.
-            result[key] = to_tensor(
-                sub_value,
-                device=device,
-                **kwargs,
-            )
-            continue
-
-        # Convert individual values to tensors.
-        result[key] = to_tensor(
-            sub_value,
-            device=device,
-            **kwargs,
-        )
-    return result
-
-
-def from_tensor_to_numpy(x: torch.Tensor | Any) -> np.ndarray | float | int | Any:
-    """
-    Convert a PyTorch tensor to a numpy array or scalar if applicable.
-
-    If the input is not a tensor, it is returned unchanged.
-
-    Args:
-        x: The input, which can be a tensor or any other type.
-
-    Returns:
-        A numpy array, a scalar, or the original input.
-    """
+def _from_tensor(x: Any):
     if isinstance(x, torch.Tensor):
         return x.item() if x.numel() == 1 else x.detach().cpu().numpy()
     return x
 
 
 def _is_image(arr: Any) -> bool:
-    """
-    Check if a given array is likely an image (uint8, 3D).
-
-    Args:
-        arr: The array to check.
-
-    Returns:
-        True if the array matches the image criteria, False otherwise.
-    """
     return isinstance(arr, np.ndarray) and arr.dtype == np.uint8 and arr.ndim == 3
 
 
 def _split_obs_to_state_and_images(obs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
-    """
-    Separate an observation dictionary into state and image components.
-
-    Args:
-        obs: The observation dictionary.
-
-    Returns:
-        A tuple containing two dictionaries: one for state and one for images.
-    """
     state, images = {}, {}
     for k, v in obs.items():
-        if "image" in k.lower() or _is_image(v):
+        if _is_image(v):
             images[k] = v
         else:
             state[k] = v
     return state, images
 
 
-# Private Helper Functions (Common Logic)
-
-
-def _extract_complementary_data(batch: dict[str, Any]) -> dict[str, Any]:
-    """
-    Extract complementary data from a batch dictionary.
-
-    This includes padding flags, task description, and indices.
-
-    Args:
-        batch: The batch dictionary.
-
-    Returns:
-        A dictionary with the extracted complementary data.
-    """
-    pad_keys = {k: v for k, v in batch.items() if "_is_pad" in k}
-    task_key = {"task": batch["task"]} if "task" in batch else {}
-    index_key = {"index": batch["index"]} if "index" in batch else {}
-    task_index_key = {"task_index": batch["task_index"]} if "task_index" in batch else {}
-
-    return {**pad_keys, **task_key, **index_key, **task_index_key}
-
-
-def _merge_transitions(base: EnvTransition, other: EnvTransition) -> EnvTransition:
-    """
-    Merge two transitions, with the second one taking precedence in case of conflicts.
-
-    Args:
-        base: The base transition.
-        other: The transition to merge, which will overwrite base values.
-
-    Returns:
-        The merged transition dictionary.
-    """
-    out = deepcopy(base)
-
-    for key in (
-        TransitionKey.OBSERVATION,
-        TransitionKey.ACTION,
-        TransitionKey.INFO,
-        TransitionKey.COMPLEMENTARY_DATA,
-    ):
-        if other.get(key):
-            out.setdefault(key, {}).update(deepcopy(other[key]))
-
-    for k in (TransitionKey.REWARD, TransitionKey.DONE, TransitionKey.TRUNCATED):
-        if k in other:
-            out[k] = other[k]
-    return out
-
-
-# Core Conversion Functions
-
-
-def create_transition(
-    observation: dict[str, Any] | None = None,
-    action: dict[str, Any] | None = None,
-    reward: float = 0.0,
-    done: bool = False,
-    truncated: bool = False,
-    info: dict[str, Any] | None = None,
-    complementary_data: dict[str, Any] | None = None,
+def make_obs_act_transition(
+    *, obs: dict[str, Any] | None = None, act: dict[str, Any] | None = None
 ) -> EnvTransition:
-    """
-    Create an `EnvTransition` dictionary with sensible defaults.
-
-    Args:
-        observation: Observation dictionary.
-        action: Action dictionary.
-        reward: Scalar reward value.
-        done: Episode termination flag.
-        truncated: Episode truncation flag.
-        info: Additional info dictionary.
-        complementary_data: Complementary data dictionary.
-
-    Returns:
-        A complete `EnvTransition` dictionary.
-    """
     return {
-        TransitionKey.OBSERVATION: observation,
-        TransitionKey.ACTION: action,
-        TransitionKey.REWARD: reward,
-        TransitionKey.DONE: done,
-        TransitionKey.TRUNCATED: truncated,
-        TransitionKey.INFO: info if info is not None else {},
-        TransitionKey.COMPLEMENTARY_DATA: complementary_data if complementary_data is not None else {},
+        TransitionKey.OBSERVATION: {} if obs is None else obs,
+        TransitionKey.ACTION: {} if act is None else act,
+        TransitionKey.INFO: {},
+        TransitionKey.COMPLEMENTARY_DATA: {},
+        TransitionKey.REWARD: None,
+        TransitionKey.DONE: None,
+        TransitionKey.TRUNCATED: None,
     }
 
 
-def action_to_transition(action: dict[str, Any]) -> EnvTransition:
+def to_transition_teleop_action(action: dict[str, Any]) -> EnvTransition:
     """
-    Convert a raw action dictionary into a standardized `EnvTransition`.
-
-    The keys in the action dictionary are prefixed with "action." and stored under
-    the `ACTION` key in the transition. Values are converted to tensors, except for
-    special types like `Rotation`.
-
-    Args:
-        action: The raw action dictionary from a teleoperation device or controller.
-
-    Returns:
-        An `EnvTransition` containing the formatted action.
+    Convert a raw teleop action dict into an EnvTransition under the ACTION TransitionKey.
     """
+    act_dict: dict[str, Any] = {}
+    for k, v in action.items():
+        # Check if the value is a type that should not be converted to a tensor.
+        if isinstance(v, (Rotation, dict)):
+            act_dict[f"action.{k}"] = v
+            continue
 
-    return create_transition(observation={}, action=action)
+        arr = np.array(v) if np.isscalar(v) else v
+        act_dict[f"action.{k}"] = _to_tensor(arr)
+
+    return make_obs_act_transition(act=act_dict)
 
 
-def observation_to_transition(observation: dict[str, Any]) -> EnvTransition:
+# TODO(Adil, Pepijn): Overtime we can maybe add these converters to pipeline.py itself
+def to_transition_robot_observation(observation: dict[str, Any]) -> EnvTransition:
     """
-    Convert a raw robot observation dictionary into a standardized `EnvTransition`.
-
-    The observation is split into state and image components. State keys are prefixed
-    with "observation.state." and image keys with "observation.images.". The result is
-    stored under the `OBSERVATION` key in the transition.
-
-    Args:
-        observation: The raw observation dictionary from the environment.
-
-    Returns:
-        An `EnvTransition` containing the formatted observation.
+    Convert a raw robot observation dict into an EnvTransition under the OBSERVATION TransitionKey.
     """
     state, images = _split_obs_to_state_and_images(observation)
 
-    image_observations = {f"{OBS_IMAGES}.{cam}": img for cam, img in images.items()}
+    obs_dict: dict[str, Any] = {}
+    for k, v in state.items():
+        arr = np.array(v) if np.isscalar(v) else v
+        obs_dict[f"observation.state.{k}"] = _to_tensor(arr)
 
-    return create_transition(observation={**state, **image_observations}, action={})
+    for cam, img in images.items():
+        obs_dict[f"observation.images.{cam}"] = img
+
+    return make_obs_act_transition(obs=obs_dict)
 
 
-def transition_to_action(transition: EnvTransition) -> dict[str, Any]:
+def to_output_robot_action(transition: EnvTransition) -> dict[str, Any]:
     """
-    Extract a raw action dictionary for a robot from an `EnvTransition`.
+    Converts a EnvTransition under the ACTION TransitionKey to a dict with keys ending in '.pos' for raw robot actions.
+    """
+    out: dict[str, Any] = {}
+    action_dict = transition.get(TransitionKey.ACTION) or {}
 
-    This function searches for keys in the format "action.*.pos" or "action.*.vel"
-    and converts them into a flat dictionary suitable for sending to a robot controller.
+    for k, v in action_dict.items():
+        if isinstance(k, str) and k.startswith("action.") and k.endswith((".pos", ".vel")):
+            out_key = k[len("action.") :]  # Strip the 'action.' prefix.
+            out[out_key] = float(v)
+
+    return out
+
+
+def to_dataset_frame(
+    transitions_or_transition: EnvTransition | Iterable[EnvTransition], features: dict[str, dict]
+) -> dict[str, any]:
+    """
+    Converts a single EnvTransition or an iterable of them into a flat,
+    dataset-friendly dictionary for training or evaluation, according to
+    the provided `features` spec.
 
     Args:
-        transition: The `EnvTransition` containing the action.
+        transitions_or_transition: Either a single EnvTransition dict
+            or an iterable of them (which will be merged).
+        features (dict[str, dict]):
+            A feature specification dictionary:
+              - 'action': dict with 'names': list of action feature names
+              - 'observation.state': dict with 'names': list of state feature names
+              - keys starting with 'observation.images.' are passed through
 
     Returns:
-        A dictionary representing the raw robot action.
+        batch (dict[str, any]): Flat dictionary containing:
+          - numpy arrays for "observation.state" and "action"
+          - any image tensors defined in features
+          - next.{reward,done,truncated}
+          - info dict
+          - *_is_pad flags and task from complementary_data
     """
-    return transition.get(TransitionKey.ACTION)
+    action_names = features.get("action", {}).get("names", [])
+    obs_state_names = features.get("observation.state", {}).get("names", [])
+    image_keys = [k for k in features if k.startswith("observation.images.")]
 
+    def _merge(base: EnvTransition, other: EnvTransition) -> EnvTransition:
+        out = deepcopy(base)
+        for key in (
+            TransitionKey.OBSERVATION,
+            TransitionKey.ACTION,
+            TransitionKey.INFO,
+            TransitionKey.COMPLEMENTARY_DATA,
+        ):
+            if other.get(key):
+                out.setdefault(key, {}).update(deepcopy(other[key]))
+        for k in (TransitionKey.REWARD, TransitionKey.DONE, TransitionKey.TRUNCATED):
+            if k in other:
+                out[k] = other[k]
+        return out
 
-def merge_transitions(transitions: Sequence[EnvTransition] | EnvTransition) -> EnvTransition:
-    """
-    Merge a sequence of transitions into a single one.
+    def _ensure_transition(obj) -> EnvTransition:
+        # single transition
+        if isinstance(obj, dict) and any(isinstance(k, TransitionKey) for k in obj):
+            return obj
+        # iterable of transitions
+        if isinstance(obj, Iterable):
+            items = list(obj)
+            if not items:
+                return {}
+            acc = items[0]
+            for t in items[1:]:
+                acc = _merge(acc, t)
+            return acc
+        raise TypeError("Expected EnvTransition or iterable of them")
 
-    If a single transition is provided, it is returned as is. For a sequence,
-    transitions are merged sequentially, with later transitions in the sequence
-    overwriting earlier ones.
-
-    Args:
-        transitions: A single transition or a sequence of them.
-
-    Returns:
-        A single merged `EnvTransition`.
-
-    Raises:
-        ValueError: If an empty sequence of transitions is provided.
-    """
-
-    if not isinstance(transitions, Sequence):  # Single transition
-        return transitions
-
-    items = list(transitions)
-    if not items:
-        raise ValueError("merge_transitions() requires a non-empty sequence of transitions")
-
-    result = items[0]
-    for t in items[1:]:
-        result = _merge_transitions(result, t)
-    return result
-
-
-def transition_to_dataset_frame(
-    transitions_or_transition: EnvTransition | Sequence[EnvTransition], features: dict[str, dict]
-) -> dict[str, Any]:
-    """
-    Convert one or more transitions into a flat dictionary suitable for a dataset frame.
-
-    This function processes `EnvTransition` objects according to a feature
-    specification, producing a format ready for training or evaluation.
-
-    Args:
-        transitions_or_transition: A single `EnvTransition` or a sequence to be merged.
-        features: A feature specification dictionary.
-
-    Returns:
-        A flat dictionary representing a single frame of data for a dataset.
-    """
-    action_names = features.get(ACTION, {}).get("names", [])
-    obs_state_names = features.get(OBS_STATE, {}).get("names", [])
-    image_keys = [k for k in features if k.startswith(OBS_IMAGES)]
-
-    tr = merge_transitions(transitions_or_transition)
+    tr = _ensure_transition(transitions_or_transition)
     obs = tr.get(TransitionKey.OBSERVATION, {}) or {}
     act = tr.get(TransitionKey.ACTION, {}) or {}
-    batch: dict[str, Any] = {}
+    batch: dict[str, any] = {}
 
-    # Passthrough for images.
+    # Images passthrough
     for k in image_keys:
         if k in obs:
             batch[k] = obs[k]
 
-    # Create observation.state vector.
+    # Observation.state vector
     if obs_state_names:
-        vals = [from_tensor_to_numpy(obs.get(f"{OBS_STATE}.{n}", 0.0)) for n in obs_state_names]
-        batch[OBS_STATE] = np.asarray(vals, dtype=np.float32)
+        vals = [_from_tensor(obs.get(f"observation.state.{n}", 0.0)) for n in obs_state_names]
+        batch["observation.state"] = np.asarray(vals, dtype=np.float32)
 
-    # Create action vector.
+    # Action vector
     if action_names:
-        vals = [from_tensor_to_numpy(act.get(f"{ACTION}.{n}", 0.0)) for n in action_names]
-        batch[ACTION] = np.asarray(vals, dtype=np.float32)
+        vals = [_from_tensor(act.get(f"action.{n}", 0.0)) for n in action_names]
+        batch["action"] = np.asarray(vals, dtype=np.float32)
 
-    # Add transition metadata.
+    # Next.* fields
     if tr.get(TransitionKey.REWARD) is not None:
-        reward_val = from_tensor_to_numpy(tr[TransitionKey.REWARD])
-        # Check if features expect array format, otherwise keep as scalar.
-        if REWARD in features and features[REWARD].get("shape") == (1,):
-            batch[REWARD] = np.array([reward_val], dtype=np.float32)
-        else:
-            batch[REWARD] = reward_val
-
+        batch["next.reward"] = _from_tensor(tr[TransitionKey.REWARD])
     if tr.get(TransitionKey.DONE) is not None:
-        done_val = from_tensor_to_numpy(tr[TransitionKey.DONE])
-        if DONE in features and features[DONE].get("shape") == (1,):
-            batch[DONE] = np.array([done_val], dtype=bool)
-        else:
-            batch[DONE] = done_val
-
+        batch["next.done"] = _from_tensor(tr[TransitionKey.DONE])
     if tr.get(TransitionKey.TRUNCATED) is not None:
-        truncated_val = from_tensor_to_numpy(tr[TransitionKey.TRUNCATED])
-        if TRUNCATED in features and features[TRUNCATED].get("shape") == (1,):
-            batch[TRUNCATED] = np.array([truncated_val], dtype=bool)
-        else:
-            batch[TRUNCATED] = truncated_val
+        batch["next.truncated"] = _from_tensor(tr[TransitionKey.TRUNCATED])
 
-    # Add complementary data flags and task.
+    # Complementary data flags and task
     comp = tr.get(TransitionKey.COMPLEMENTARY_DATA) or {}
     if comp:
-        # Padding flags.
+        # pad flags
         for k, v in comp.items():
             if k.endswith("_is_pad"):
                 batch[k] = v
-        # Task label.
+        # task label
         if comp.get("task") is not None:
             batch["task"] = comp["task"]
 
     return batch
-
-
-def batch_to_transition(batch: dict[str, Any]) -> EnvTransition:
-    """
-    Convert a batch dictionary from a dataset/dataloader into an `EnvTransition`.
-
-    This function maps recognized keys from a batch to the `EnvTransition` structure,
-    filling in missing keys with sensible defaults.
-
-    Args:
-        batch: A batch dictionary.
-
-    Returns:
-        An `EnvTransition` dictionary.
-
-    Raises:
-        ValueError: If the input is not a dictionary.
-    """
-
-    # Validate input type.
-    if not isinstance(batch, dict):
-        raise ValueError(f"EnvTransition must be a dictionary. Got {type(batch).__name__}")
-
-    # Extract observation and complementary data keys.
-    observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
-    complementary_data = _extract_complementary_data(batch)
-
-    return create_transition(
-        observation=observation_keys if observation_keys else None,
-        action=batch.get("action"),
-        reward=batch.get("next.reward", 0.0),
-        done=batch.get("next.done", False),
-        truncated=batch.get("next.truncated", False),
-        info=batch.get("info", {}),
-        complementary_data=complementary_data if complementary_data else None,
-    )
-
-
-def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
-    """
-    Convert an `EnvTransition` back to the canonical batch format used in LeRobot.
-
-    This is the inverse of `batch_to_transition`.
-
-    Args:
-        transition: The `EnvTransition` to convert.
-
-    Returns:
-        A batch dictionary with canonical LeRobot field names.
-    """
-    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),
-        "info": transition.get(TransitionKey.INFO, {}),
-    }
-
-    # Add complementary data.
-    comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
-    if comp_data:
-        batch.update(comp_data)
-
-    # Flatten observation dictionary.
-    observation = transition.get(TransitionKey.OBSERVATION)
-    if isinstance(observation, dict):
-        batch.update(observation)
-
-    return batch
-
-
-def identity_transition(tr: EnvTransition) -> EnvTransition:
-    """
-    An identity function for transitions, returning the input unchanged.
-
-    Useful as a default or placeholder in processing pipelines.
-
-    Args:
-        tr: An `EnvTransition`.
-
-    Returns:
-        The same `EnvTransition`.
-    """
-    return tr

src/lerobot/processor/core.py

@@ -1,49 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from enum import Enum
-from typing import Any, TypedDict
-
-import torch
-
-
-class TransitionKey(str, Enum):
-    """Keys for accessing EnvTransition dictionary components."""
-
-    # TODO(Steven): Use consts
-    OBSERVATION = "observation"
-    ACTION = "action"
-    REWARD = "reward"
-    DONE = "done"
-    TRUNCATED = "truncated"
-    INFO = "info"
-    COMPLEMENTARY_DATA = "complementary_data"
-
-
-EnvTransition = TypedDict(
-    "EnvTransition",
-    {
-        TransitionKey.OBSERVATION.value: dict[str, Any] | None,
-        TransitionKey.ACTION.value: Any | torch.Tensor | None,
-        TransitionKey.REWARD.value: float | torch.Tensor | None,
-        TransitionKey.DONE.value: bool | torch.Tensor | None,
-        TransitionKey.TRUNCATED.value: bool | torch.Tensor | None,
-        TransitionKey.INFO.value: dict[str, Any] | None,
-        TransitionKey.COMPLEMENTARY_DATA.value: dict[str, Any] | None,
-    },
-)

src/lerobot/processor/delta_action_processor.py

@@ -1,147 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-
-from torch import Tensor
-
-from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-
-from .pipeline import ActionProcessorStep, ProcessorStepRegistry
-
-
-@ProcessorStepRegistry.register("map_tensor_to_delta_action_dict")
-@dataclass
-class MapTensorToDeltaActionDictStep(ActionProcessorStep):
-    """
-    Maps a flat action tensor from a policy to a structured delta action dictionary.
-
-    This step is typically used after a policy outputs a continuous action vector.
-    It decomposes the vector into named components for delta movements of the
-    end-effector (x, y, z) and optionally the gripper.
-
-    Attributes:
-        use_gripper: If True, assumes the 4th element of the tensor is the
-                     gripper action.
-    """
-
-    use_gripper: bool = True
-
-    def action(self, action: Tensor) -> dict:
-        if action.dim() > 1:
-            action = action.squeeze(0)
-
-        # TODO (maractingi): add rotation
-        delta_action = {
-            "delta_x": action[0],
-            "delta_y": action[1],
-            "delta_z": action[2],
-        }
-        if self.use_gripper:
-            delta_action["gripper"] = action[3]
-        return delta_action
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        features[PipelineFeatureType.ACTION]["delta_x"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["delta_y"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["delta_z"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        if self.use_gripper:
-            features[PipelineFeatureType.ACTION]["gripper"] = PolicyFeature(
-                type=FeatureType.ACTION, shape=(1,)
-            )
-        return features
-
-
-@ProcessorStepRegistry.register("map_delta_action_to_robot_action")
-@dataclass
-class MapDeltaActionToRobotActionStep(ActionProcessorStep):
-    """
-    Maps delta actions from teleoperators to robot target actions for inverse kinematics.
-
-    This step converts a dictionary of delta movements (e.g., from a gamepad)
-    into a target action format that includes an "enabled" flag and target
-    end-effector positions. It also handles scaling and noise filtering.
-
-    Attributes:
-        position_scale: A factor to scale the delta position inputs.
-        rotation_scale: A factor to scale the delta rotation inputs (currently unused).
-        noise_threshold: The magnitude below which delta inputs are considered noise
-                         and do not trigger an "enabled" state.
-    """
-
-    # Scale factors for delta movements
-    position_scale: float = 1.0
-    rotation_scale: float = 0.0  # No rotation deltas for gamepad/keyboard
-    noise_threshold: float = 1e-3  # 1 mm threshold to filter out noise
-
-    def action(self, action: dict) -> dict:
-        # NOTE (maractingi): Action can be a dict from the teleop_devices or a tensor from the policy
-        # TODO (maractingi): changing this target_xyz naming convention from the teleop_devices
-        delta_x = action.pop("delta_x", 0.0)
-        delta_y = action.pop("delta_y", 0.0)
-        delta_z = action.pop("delta_z", 0.0)
-        gripper = action.pop("gripper", 1.0)  # Default to "stay" (1.0)
-
-        # Determine if the teleoperator is actively providing input
-        # Consider enabled if any significant movement delta is detected
-        position_magnitude = (delta_x**2 + delta_y**2 + delta_z**2) ** 0.5  # Use Euclidean norm for position
-        enabled = position_magnitude > self.noise_threshold  # Small threshold to avoid noise
-
-        # Scale the deltas appropriately
-        scaled_delta_x = delta_x * self.position_scale
-        scaled_delta_y = delta_y * self.position_scale
-        scaled_delta_z = delta_z * self.position_scale
-
-        # For gamepad/keyboard, we don't have rotation input, so set to 0
-        # These could be extended in the future for more sophisticated teleoperators
-        target_wx = 0.0
-        target_wy = 0.0
-        target_wz = 0.0
-
-        # Update action with robot target format
-        action = {
-            "enabled": enabled,
-            "target_x": scaled_delta_x,
-            "target_y": scaled_delta_y,
-            "target_z": scaled_delta_z,
-            "target_wx": target_wx,
-            "target_wy": target_wy,
-            "target_wz": target_wz,
-            "gripper": float(gripper),
-        }
-
-        return action
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """Transform features to match output format."""
-        features[PipelineFeatureType.ACTION].pop("delta_x", None)
-        features[PipelineFeatureType.ACTION].pop("delta_y", None)
-        features[PipelineFeatureType.ACTION].pop("delta_z", None)
-        features[PipelineFeatureType.ACTION].pop("gripper", None)
-
-        features[PipelineFeatureType.ACTION]["enabled"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["target_x"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["target_y"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["target_z"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["target_wx"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["target_wy"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["target_wz"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        features[PipelineFeatureType.ACTION]["gripper"] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-        return features

src/lerobot/processor/device_processor.py

@@ -13,101 +13,62 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-"""
-This script defines a processor step for moving environment transition data to a specific torch device and casting
-its floating-point precision.
-"""
-
 from dataclasses import dataclass
 from typing import Any
 
 import torch
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs.types import PolicyFeature
+from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
 from lerobot.utils.utils import get_safe_torch_device
 
-from .core import EnvTransition, TransitionKey
-from .pipeline import ProcessorStep, ProcessorStepRegistry
-
 
 @ProcessorStepRegistry.register("device_processor")
 @dataclass
-class DeviceProcessorStep(ProcessorStep):
-    """
-    Processor step to move all tensors within an `EnvTransition` to a specified device and optionally cast their
-    floating-point data type.
+class DeviceProcessor:
+    """Processes transitions by moving tensors to the specified device and optionally converting float dtypes.
 
-    This is crucial for preparing data for model training or inference on hardware like GPUs.
-
-    Attributes:
-        device: The target device for tensors (e.g., "cpu", "cuda", "cuda:0").
-        float_dtype: The target floating-point dtype as a string (e.g., "float32", "float16", "bfloat16").
-                     If None, the dtype is not changed.
+    This processor ensures that all tensors in the transition are moved to the
+    specified device (CPU or GPU) before they are returned. It can also convert
+    floating-point tensors to a specified dtype while preserving non-float types
+    (int, long, bool, etc.).
     """
 
     device: str = "cpu"
     float_dtype: str | None = None
-
-    DTYPE_MAPPING = {
-        "float16": torch.float16,
-        "float32": torch.float32,
-        "float64": torch.float64,
-        "bfloat16": torch.bfloat16,
-        "half": torch.float16,
-        "float": torch.float32,
-        "double": torch.float64,
-    }
+    _device: torch.device | None = None
 
     def __post_init__(self):
-        """
-        Initializes the processor by converting string configurations to torch objects.
-
-        This method sets up the `torch.device`, determines if transfers can be non-blocking, and validates the
-        `float_dtype` string, converting it to a `torch.dtype` object.
-        """
-        self.tensor_device: torch.device = get_safe_torch_device(self.device)
-        # Update device string in case a specific GPU was selected (e.g., "cuda" -> "cuda:0")
-        self.device = self.tensor_device.type
+        self._device = get_safe_torch_device(self.device)
+        self.device = self._device.type
         self.non_blocking = "cuda" in str(self.device)
 
         # Validate and convert float_dtype string to torch dtype
         if self.float_dtype is not None:
-            if self.float_dtype not in self.DTYPE_MAPPING:
+            dtype_mapping = {
+                "float16": torch.float16,
+                "float32": torch.float32,
+                "float64": torch.float64,
+                "bfloat16": torch.bfloat16,
+                "half": torch.float16,
+                "float": torch.float32,
+                "double": torch.float64,
+            }
+
+            if self.float_dtype not in dtype_mapping:
+                available_dtypes = list(dtype_mapping.keys())
                 raise ValueError(
-                    f"Invalid float_dtype '{self.float_dtype}'. Available options: {list(self.DTYPE_MAPPING.keys())}"
+                    f"Invalid float_dtype '{self.float_dtype}'. Available options: {available_dtypes}"
                 )
-            self._target_float_dtype = self.DTYPE_MAPPING[self.float_dtype]
+
+            self._target_float_dtype = dtype_mapping[self.float_dtype]
         else:
             self._target_float_dtype = None
 
     def _process_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
-        """
-        Moves a single tensor to the target device and casts its dtype.
-
-        Handles multi-GPU scenarios by not moving a tensor if it's already on a different CUDA device than
-        the target, which is useful when using frameworks like Accelerate.
-
-        Args:
-            tensor: The input torch.Tensor.
-
-        Returns:
-            The processed tensor on the correct device and with the correct dtype.
-        """
-        # Determine target device
-        if tensor.is_cuda and self.tensor_device.type == "cuda":
-            # Both tensor and target are on GPU - preserve tensor's GPU placement.
-            # This handles multi-GPU scenarios where Accelerate has already placed
-            # tensors on the correct GPU for each process.
-            target_device = tensor.device
-        else:
-            # Either tensor is on CPU, or we're configured for CPU.
-            # In both cases, use the configured device.
-            target_device = self.tensor_device
-
-        # Only move if necessary
-        if tensor.device != target_device:
-            tensor = tensor.to(target_device, non_blocking=self.non_blocking)
+        """Process a tensor by moving to device and optionally converting float dtype."""
+        # Move to device first
+        tensor = tensor.to(self.device, non_blocking=self.non_blocking)
 
         # Convert float dtype if specified and tensor is floating point
         if self._target_float_dtype is not None and tensor.is_floating_point():
@@ -116,71 +77,69 @@ class DeviceProcessorStep(ProcessorStep):
         return tensor
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """
-        Applies device and dtype conversion to all tensors in an environment transition.
-
-        It iterates through the transition, finds all `torch.Tensor` objects (including those nested in
-        dictionaries like `observation`), and processes them.
-
-        Args:
-            transition: The input `EnvTransition` object.
-
-        Returns:
-            A new `EnvTransition` object with all tensors moved to the target device and dtype.
-        """
+        # Create a copy of the transition
         new_transition = transition.copy()
 
-        simple_tensor_keys = [
-            TransitionKey.ACTION,
-            TransitionKey.REWARD,
-            TransitionKey.DONE,
-            TransitionKey.TRUNCATED,
-        ]
+        # Process observation tensors
+        observation = transition.get(TransitionKey.OBSERVATION)
+        if observation is not None:
+            new_observation = {
+                k: self._process_tensor(v) if isinstance(v, torch.Tensor) else v
+                for k, v in observation.items()
+            }
+            new_transition[TransitionKey.OBSERVATION] = new_observation
 
-        dict_tensor_keys = [
-            TransitionKey.OBSERVATION,
-            TransitionKey.COMPLEMENTARY_DATA,
-        ]
+        # Process action tensor
+        action = transition.get(TransitionKey.ACTION)
+        if action is not None and isinstance(action, torch.Tensor):
+            new_transition[TransitionKey.ACTION] = self._process_tensor(action)
 
-        # Process simple, top-level tensors
-        for key in simple_tensor_keys:
-            value = transition.get(key)
-            if isinstance(value, torch.Tensor):
-                new_transition[key] = self._process_tensor(value)
+        # Process reward tensor
+        reward = transition.get(TransitionKey.REWARD)
+        if reward is not None and isinstance(reward, torch.Tensor):
+            new_transition[TransitionKey.REWARD] = self._process_tensor(reward)
 
-        # Process tensors nested within dictionaries
-        for key in dict_tensor_keys:
-            data_dict = transition.get(key)
-            if data_dict is not None:
-                new_data_dict = {
-                    k: self._process_tensor(v) if isinstance(v, torch.Tensor) else v
-                    for k, v in data_dict.items()
-                }
-                new_transition[key] = new_data_dict
+        # Process done tensor
+        done = transition.get(TransitionKey.DONE)
+        if done is not None and isinstance(done, torch.Tensor):
+            new_transition[TransitionKey.DONE] = self._process_tensor(done)
+
+        # Process truncated tensor
+        truncated = transition.get(TransitionKey.TRUNCATED)
+        if truncated is not None and isinstance(truncated, torch.Tensor):
+            new_transition[TransitionKey.TRUNCATED] = self._process_tensor(truncated)
+
+        # Process complementary data tensors
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if complementary_data is not None:
+            new_complementary_data = {}
+
+            # Process all items in complementary_data
+            for key, value in complementary_data.items():
+                if isinstance(value, torch.Tensor):
+                    new_complementary_data[key] = self._process_tensor(value)
+                else:
+                    new_complementary_data[key] = value
+
+            new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
 
         return new_transition
 
     def get_config(self) -> dict[str, Any]:
-        """
-        Returns the serializable configuration of the processor.
-
-        Returns:
-            A dictionary containing the device and float_dtype settings.
-        """
+        """Return configuration for serialization."""
         return {"device": self.device, "float_dtype": self.float_dtype}
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Returns the input features unchanged.
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return state dictionary (empty for this processor)."""
+        return {}
 
-        Device and dtype transformations do not alter the fundamental definition of the features (e.g., shape).
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Load state dictionary (no-op for this processor)."""
+        pass
 
-        Args:
-            features: A dictionary of policy features.
+    def reset(self) -> None:
+        """Reset processor state (no-op for this processor)."""
+        pass
 
-        Returns:
-            The original dictionary of policy features.
-        """
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features

src/lerobot/processor/gym_action_processor.py

@@ -1,95 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-
-from .converters import to_tensor
-from .pipeline import ActionProcessorStep, ProcessorStepRegistry
-
-
-@ProcessorStepRegistry.register("torch2numpy_action_processor")
-@dataclass
-class Torch2NumpyActionProcessorStep(ActionProcessorStep):
-    """
-    Converts a PyTorch tensor action to a NumPy array.
-
-    This step is useful when the output of a policy (typically a torch.Tensor)
-    needs to be passed to an environment or component that expects a NumPy array.
-
-    Attributes:
-        squeeze_batch_dim: If True, removes the first dimension of the array
-                           if it is of size 1. This is useful for converting a
-                           batched action of size (1, D) to a single action of size (D,).
-    """
-
-    squeeze_batch_dim: bool = True
-
-    def action(self, action: torch.Tensor) -> np.ndarray:
-        if not isinstance(action, torch.Tensor):
-            raise TypeError(
-                f"Expected torch.Tensor or None, got {type(action).__name__}. "
-                "Use appropriate processor for non-tensor actions."
-            )
-
-        numpy_action = action.detach().cpu().numpy()
-
-        # Remove batch dimensions but preserve action dimensions.
-        # Only squeeze if there's a batch dimension (first dim == 1).
-        if (
-            self.squeeze_batch_dim
-            and numpy_action.shape
-            and len(numpy_action.shape) > 1
-            and numpy_action.shape[0] == 1
-        ):
-            numpy_action = numpy_action.squeeze(0)
-
-        return numpy_action
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@ProcessorStepRegistry.register("numpy2torch_action_processor")
-@dataclass
-class Numpy2TorchActionProcessorStep(ActionProcessorStep):
-    """
-    Converts a NumPy array action to a PyTorch tensor.
-
-    This step is useful for converting actions from environments or hardware,
-    which are often NumPy arrays, into PyTorch tensors that can be processed
-    by a policy or model.
-    """
-
-    def action(self, action: np.ndarray) -> torch.Tensor:
-        if not isinstance(action, np.ndarray):
-            raise TypeError(
-                f"Expected np.ndarray or None, got {type(action).__name__}. "
-                "Use appropriate processor for non-tensor actions."
-            )
-        torch_action = to_tensor(action, dtype=None)  # Preserve original dtype
-        return torch_action
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features

src/lerobot/processor/hil_processor.py

@@ -1,592 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-import time
-from dataclasses import dataclass
-from typing import Any, Protocol, TypeVar, runtime_checkable
-
-import numpy as np
-import torch
-import torchvision.transforms.functional as F  # noqa: N812
-
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.constants import ACTION
-from lerobot.teleoperators.teleoperator import Teleoperator
-from lerobot.teleoperators.utils import TeleopEvents
-
-from .core import EnvTransition, TransitionKey
-from .pipeline import (
-    ComplementaryDataProcessorStep,
-    InfoProcessorStep,
-    ObservationProcessorStep,
-    ProcessorStep,
-    ProcessorStepRegistry,
-    TruncatedProcessorStep,
-)
-
-GRIPPER_KEY = "gripper"
-DISCRETE_PENALTY_KEY = "discrete_penalty"
-TELEOP_ACTION_KEY = "teleop_action"
-
-
-@runtime_checkable
-class HasTeleopEvents(Protocol):
-    """
-    Minimal protocol for objects that provide teleoperation events.
-
-    This protocol defines the `get_teleop_events()` method, allowing processor
-    steps to interact with teleoperators that support event-based controls
-    (like episode termination or success flagging) without needing to know the
-    teleoperator's specific class.
-    """
-
-    def get_teleop_events(self) -> dict[str, Any]:
-        """
-        Get extra control events from the teleoperator.
-
-        Returns:
-            A dictionary containing control events such as:
-            - `is_intervention`: bool - Whether the human is currently intervening.
-            - `terminate_episode`: bool - Whether to terminate the current episode.
-            - `success`: bool - Whether the episode was successful.
-            - `rerecord_episode`: bool - Whether to rerecord the episode.
-        """
-        ...
-
-
-# Type variable constrained to Teleoperator subclasses that also implement events
-TeleopWithEvents = TypeVar("TeleopWithEvents", bound=Teleoperator)
-
-
-def _check_teleop_with_events(teleop: Teleoperator) -> None:
-    """
-    Runtime check that a teleoperator implements the `HasTeleopEvents` protocol.
-
-    Args:
-        teleop: The teleoperator instance to check.
-
-    Raises:
-        TypeError: If the teleoperator does not have a `get_teleop_events` method.
-    """
-    if not isinstance(teleop, HasTeleopEvents):
-        raise TypeError(
-            f"Teleoperator {type(teleop).__name__} must implement get_teleop_events() method. "
-            f"Compatible teleoperators: GamepadTeleop, KeyboardEndEffectorTeleop"
-        )
-
-
-@ProcessorStepRegistry.register("add_teleop_action_as_complementary_data")
-@dataclass
-class AddTeleopActionAsComplimentaryDataStep(ComplementaryDataProcessorStep):
-    """
-    Adds the raw action from a teleoperator to the transition's complementary data.
-
-    This is useful for human-in-the-loop scenarios where the human's input needs to
-    be available to downstream processors, for example, to override a policy's action
-    during an intervention.
-
-    Attributes:
-        teleop_device: The teleoperator instance to get the action from.
-    """
-
-    teleop_device: Teleoperator
-
-    def complementary_data(self, complementary_data: dict) -> dict:
-        """
-        Retrieves the teleoperator's action and adds it to the complementary data.
-
-        Args:
-            complementary_data: The incoming complementary data dictionary.
-
-        Returns:
-            A new dictionary with the teleoperator action added under the
-            `teleop_action` key.
-        """
-        new_complementary_data = dict(complementary_data)
-        new_complementary_data[TELEOP_ACTION_KEY] = self.teleop_device.get_action()
-        return new_complementary_data
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@ProcessorStepRegistry.register("add_teleop_action_as_info")
-@dataclass
-class AddTeleopEventsAsInfoStep(InfoProcessorStep):
-    """
-    Adds teleoperator control events (e.g., terminate, success) to the transition's info.
-
-    This step extracts control events from teleoperators that support event-based
-    interaction, making these signals available to other parts of the system.
-
-    Attributes:
-        teleop_device: An instance of a teleoperator that implements the
-                       `HasTeleopEvents` protocol.
-    """
-
-    teleop_device: TeleopWithEvents
-
-    def __post_init__(self):
-        """Validates that the provided teleoperator supports events after initialization."""
-        _check_teleop_with_events(self.teleop_device)
-
-    def info(self, info: dict) -> dict:
-        """
-        Retrieves teleoperator events and updates the info dictionary.
-
-        Args:
-            info: The incoming info dictionary.
-
-        Returns:
-            A new dictionary including the teleoperator events.
-        """
-        new_info = dict(info)
-
-        teleop_events = self.teleop_device.get_teleop_events()
-        new_info.update(teleop_events)
-        return new_info
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@ProcessorStepRegistry.register("image_crop_resize_processor")
-@dataclass
-class ImageCropResizeProcessorStep(ObservationProcessorStep):
-    """
-    Crops and/or resizes image observations.
-
-    This step iterates through all image keys in an observation dictionary and applies
-    the specified transformations. It handles device placement, moving tensors to the
-    CPU if necessary for operations not supported on certain accelerators like MPS.
-
-    Attributes:
-        crop_params_dict: A dictionary mapping image keys to cropping parameters
-                          (top, left, height, width).
-        resize_size: A tuple (height, width) to resize all images to.
-    """
-
-    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
-    resize_size: tuple[int, int] | None = None
-
-    def observation(self, observation: dict) -> dict:
-        """
-        Applies cropping and resizing to all images in the observation dictionary.
-
-        Args:
-            observation: The observation dictionary, potentially containing image tensors.
-
-        Returns:
-            A new observation dictionary with transformed images.
-        """
-        if self.resize_size is None and not self.crop_params_dict:
-            return observation
-
-        new_observation = dict(observation)
-
-        # Process all image keys in the observation
-        for key in observation:
-            if "image" not in key:
-                continue
-
-            image = observation[key]
-            device = image.device
-            # NOTE (maractingi): No mps kernel for crop and resize, so we need to move to cpu
-            if device.type == "mps":
-                image = image.cpu()
-            # Crop if crop params are provided for this key
-            if self.crop_params_dict is not None and key in self.crop_params_dict:
-                crop_params = self.crop_params_dict[key]
-                image = F.crop(image, *crop_params)
-            if self.resize_size is not None:
-                image = F.resize(image, self.resize_size)
-                image = image.clamp(0.0, 1.0)
-            new_observation[key] = image.to(device)
-
-        return new_observation
-
-    def get_config(self) -> dict[str, Any]:
-        """
-        Returns the configuration of the step for serialization.
-
-        Returns:
-            A dictionary with the crop parameters and resize dimensions.
-        """
-        return {
-            "crop_params_dict": self.crop_params_dict,
-            "resize_size": self.resize_size,
-        }
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Updates the image feature shapes in the policy features dictionary if resizing is applied.
-
-        Args:
-            features: The policy features dictionary.
-
-        Returns:
-            The updated policy features dictionary with new image shapes.
-        """
-        if self.resize_size is None:
-            return features
-        for key in features[PipelineFeatureType.OBSERVATION]:
-            if "image" in key:
-                nb_channel = features[PipelineFeatureType.OBSERVATION][key].shape[0]
-                features[PipelineFeatureType.OBSERVATION][key] = PolicyFeature(
-                    type=features[PipelineFeatureType.OBSERVATION][key].type,
-                    shape=(nb_channel, *self.resize_size),
-                )
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register("time_limit_processor")
-class TimeLimitProcessorStep(TruncatedProcessorStep):
-    """
-    Tracks episode steps and enforces a time limit by truncating the episode.
-
-    Attributes:
-        max_episode_steps: The maximum number of steps allowed per episode.
-        current_step: The current step count for the active episode.
-    """
-
-    max_episode_steps: int
-    current_step: int = 0
-
-    def truncated(self, truncated: bool) -> bool:
-        """
-        Increments the step counter and sets the truncated flag if the time limit is reached.
-
-        Args:
-            truncated: The incoming truncated flag.
-
-        Returns:
-            True if the episode step limit is reached, otherwise the incoming value.
-        """
-        self.current_step += 1
-        if self.current_step >= self.max_episode_steps:
-            truncated = True
-        # TODO (steven): missing an else truncated = False?
-        return truncated
-
-    def get_config(self) -> dict[str, Any]:
-        """
-        Returns the configuration of the step for serialization.
-
-        Returns:
-            A dictionary containing the `max_episode_steps`.
-        """
-        return {
-            "max_episode_steps": self.max_episode_steps,
-        }
-
-    def reset(self) -> None:
-        """Resets the step counter, typically called at the start of a new episode."""
-        self.current_step = 0
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register("gripper_penalty_processor")
-class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
-    """
-    Applies a penalty for inefficient gripper usage.
-
-    This step penalizes actions that attempt to close an already closed gripper or
-    open an already open one, based on position thresholds.
-
-    Attributes:
-        penalty: The negative reward value to apply.
-        max_gripper_pos: The maximum position value for the gripper, used for normalization.
-    """
-
-    penalty: float = -0.01
-    max_gripper_pos: float = 30.0
-
-    def complementary_data(self, complementary_data: dict) -> dict:
-        """
-        Calculates the gripper penalty and adds it to the complementary data.
-
-        Args:
-            complementary_data: The incoming complementary data, which should contain
-                                raw joint positions.
-
-        Returns:
-            A new complementary data dictionary with the `discrete_penalty` key added.
-        """
-        action = self.transition.get(TransitionKey.ACTION)
-
-        current_gripper_pos = complementary_data.get("raw_joint_positions", None).get(GRIPPER_KEY, None)
-        if current_gripper_pos is None:
-            return complementary_data
-
-        gripper_action = action[f"{ACTION}.{GRIPPER_KEY}.pos"]
-        gripper_action_normalized = gripper_action / self.max_gripper_pos
-
-        # Normalize gripper state and action
-        gripper_state_normalized = current_gripper_pos / self.max_gripper_pos
-
-        # Calculate penalty boolean as in original
-        gripper_penalty_bool = (gripper_state_normalized < 0.5 and gripper_action_normalized > 0.5) or (
-            gripper_state_normalized > 0.75 and gripper_action_normalized < 0.5
-        )
-
-        gripper_penalty = self.penalty * int(gripper_penalty_bool)
-
-        # Create new complementary data with penalty info
-        new_complementary_data = dict(complementary_data)
-        new_complementary_data[DISCRETE_PENALTY_KEY] = gripper_penalty
-
-        return new_complementary_data
-
-    def get_config(self) -> dict[str, Any]:
-        """
-        Returns the configuration of the step for serialization.
-
-        Returns:
-            A dictionary containing the penalty value and max gripper position.
-        """
-        return {
-            "penalty": self.penalty,
-            "max_gripper_pos": self.max_gripper_pos,
-        }
-
-    def reset(self) -> None:
-        """Resets the processor's internal state."""
-        pass
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register("intervention_action_processor")
-class InterventionActionProcessorStep(ProcessorStep):
-    """
-    Handles human intervention, overriding policy actions and managing episode termination.
-
-    When an intervention is detected (via teleoperator events in the `info` dict),
-    this step replaces the policy's action with the human's teleoperated action.
-    It also processes signals to terminate the episode or flag success.
-
-    Attributes:
-        use_gripper: Whether to include the gripper in the teleoperated action.
-        terminate_on_success: If True, automatically sets the `done` flag when a
-                              `success` event is received.
-    """
-
-    use_gripper: bool = False
-    terminate_on_success: bool = True
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """
-        Processes the transition to handle interventions.
-
-        Args:
-            transition: The incoming environment transition.
-
-        Returns:
-            The modified transition, potentially with an overridden action, updated
-            reward, and termination status.
-        """
-        action = transition.get(TransitionKey.ACTION)
-        if action is None:
-            return transition
-
-        # Get intervention signals from complementary data
-        info = transition.get(TransitionKey.INFO, {})
-        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
-        teleop_action = complementary_data.get(TELEOP_ACTION_KEY, {})
-        is_intervention = info.get(TeleopEvents.IS_INTERVENTION, False)
-        terminate_episode = info.get(TeleopEvents.TERMINATE_EPISODE, False)
-        success = info.get(TeleopEvents.SUCCESS, False)
-        rerecord_episode = info.get(TeleopEvents.RERECORD_EPISODE, False)
-
-        new_transition = transition.copy()
-
-        # Override action if intervention is active
-        if is_intervention and teleop_action is not None:
-            if isinstance(teleop_action, dict):
-                # Convert teleop_action dict to tensor format
-                action_list = [
-                    teleop_action.get(f"{ACTION}.delta_x", 0.0),
-                    teleop_action.get(f"{ACTION}.delta_y", 0.0),
-                    teleop_action.get(f"{ACTION}.delta_z", 0.0),
-                ]
-                if self.use_gripper:
-                    action_list.append(teleop_action.get(GRIPPER_KEY, 1.0))
-            elif isinstance(teleop_action, np.ndarray):
-                action_list = teleop_action.tolist()
-            else:
-                action_list = teleop_action
-
-            teleop_action_tensor = torch.tensor(action_list, dtype=action.dtype, device=action.device)
-            new_transition[TransitionKey.ACTION] = teleop_action_tensor
-
-        # Handle episode termination
-        new_transition[TransitionKey.DONE] = bool(terminate_episode) or (
-            self.terminate_on_success and success
-        )
-        new_transition[TransitionKey.REWARD] = float(success)
-
-        # Update info with intervention metadata
-        info = new_transition.get(TransitionKey.INFO, {})
-        info[TeleopEvents.IS_INTERVENTION] = is_intervention
-        info[TeleopEvents.RERECORD_EPISODE] = rerecord_episode
-        info[TeleopEvents.SUCCESS] = success
-        new_transition[TransitionKey.INFO] = info
-
-        # Update complementary data with teleop action
-        complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
-        complementary_data[TELEOP_ACTION_KEY] = new_transition.get(TransitionKey.ACTION)
-        new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
-
-        return new_transition
-
-    def get_config(self) -> dict[str, Any]:
-        """
-        Returns the configuration of the step for serialization.
-
-        Returns:
-            A dictionary containing the step's configuration attributes.
-        """
-        return {
-            "use_gripper": self.use_gripper,
-            "terminate_on_success": self.terminate_on_success,
-        }
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register("reward_classifier_processor")
-class RewardClassifierProcessorStep(ProcessorStep):
-    """
-    Applies a pretrained reward classifier to image observations to predict success.
-
-    This step uses a model to determine if the current state is successful, updating
-    the reward and potentially terminating the episode.
-
-    Attributes:
-        pretrained_path: Path to the pretrained reward classifier model.
-        device: The device to run the classifier on.
-        success_threshold: The probability threshold to consider a prediction as successful.
-        success_reward: The reward value to assign on success.
-        terminate_on_success: If True, terminates the episode upon successful classification.
-        reward_classifier: The loaded classifier model instance.
-    """
-
-    pretrained_path: str | None = None
-    device: str = "cpu"
-    success_threshold: float = 0.5
-    success_reward: float = 1.0
-    terminate_on_success: bool = True
-
-    reward_classifier: Any = None
-
-    def __post_init__(self):
-        """Initializes the reward classifier model after the dataclass is created."""
-        if self.pretrained_path is not None:
-            from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
-
-            self.reward_classifier = Classifier.from_pretrained(self.pretrained_path)
-            self.reward_classifier.to(self.device)
-            self.reward_classifier.eval()
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """
-        Processes a transition, applying the reward classifier to its image observations.
-
-        Args:
-            transition: The incoming environment transition.
-
-        Returns:
-            The modified transition with an updated reward and done flag based on the
-            classifier's prediction.
-        """
-        new_transition = transition.copy()
-        observation = new_transition.get(TransitionKey.OBSERVATION)
-        if observation is None or self.reward_classifier is None:
-            return new_transition
-
-        # Extract images from observation
-        images = {key: value for key, value in observation.items() if "image" in key}
-
-        if not images:
-            return new_transition
-
-        # Run reward classifier
-        start_time = time.perf_counter()
-        with torch.inference_mode():
-            success = self.reward_classifier.predict_reward(images, threshold=self.success_threshold)
-
-        classifier_frequency = 1 / (time.perf_counter() - start_time)
-
-        # Calculate reward and termination
-        reward = new_transition.get(TransitionKey.REWARD, 0.0)
-        terminated = new_transition.get(TransitionKey.DONE, False)
-
-        if math.isclose(success, 1, abs_tol=1e-2):
-            reward = self.success_reward
-            if self.terminate_on_success:
-                terminated = True
-
-        # Update transition
-        new_transition[TransitionKey.REWARD] = reward
-        new_transition[TransitionKey.DONE] = terminated
-
-        # Update info with classifier frequency
-        info = new_transition.get(TransitionKey.INFO, {})
-        info["reward_classifier_frequency"] = classifier_frequency
-        new_transition[TransitionKey.INFO] = info
-
-        return new_transition
-
-    def get_config(self) -> dict[str, Any]:
-        """
-        Returns the configuration of the step for serialization.
-
-        Returns:
-            A dictionary containing the step's configuration attributes.
-        """
-        return {
-            "device": self.device,
-            "success_threshold": self.success_threshold,
-            "success_reward": self.success_reward,
-            "terminate_on_success": self.terminate_on_success,
-        }
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features

src/lerobot/processor/joint_observations_processor.py

@@ -1,211 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-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
-
-
-@dataclass
-@ProcessorStepRegistry.register("joint_velocity_processor")
-class JointVelocityProcessorStep(ObservationProcessorStep):
-    """
-    Calculates and appends joint velocity information to the observation state.
-
-    This step computes the velocity of each joint by calculating the finite
-    difference between the current and the last observed joint positions. The
-    resulting velocity vector is then concatenated to the original state vector.
-
-    Attributes:
-        dt: The time step (delta time) in seconds between observations, used for
-            calculating velocity.
-        last_joint_positions: Stores the joint positions from the previous step
-                              to enable velocity calculation.
-    """
-
-    dt: float = 0.1
-
-    last_joint_positions: torch.Tensor | None = None
-
-    def observation(self, observation: dict) -> dict:
-        """
-        Computes joint velocities and adds them to the observation state.
-
-        Args:
-            observation: The input observation dictionary, expected to contain
-                         an `observation.state` key with joint positions.
-
-        Returns:
-            A new observation dictionary with the `observation.state` tensor
-            extended to include joint velocities.
-
-        Raises:
-            ValueError: If `observation.state` is not found in the observation.
-        """
-        # Get current joint positions (assuming they're in observation.state)
-        current_positions = observation.get(OBS_STATE)
-        if current_positions is None:
-            raise ValueError(f"{OBS_STATE} is not in observation")
-
-        # Initialize last joint positions if not already set
-        if self.last_joint_positions is None:
-            self.last_joint_positions = current_positions.clone()
-            joint_velocities = torch.zeros_like(current_positions)
-        else:
-            # Compute velocities
-            joint_velocities = (current_positions - self.last_joint_positions) / self.dt
-
-        self.last_joint_positions = current_positions.clone()
-
-        # Extend observation with velocities
-        extended_state = torch.cat([current_positions, joint_velocities], dim=-1)
-
-        # Create new observation dict
-        new_observation = dict(observation)
-        new_observation[OBS_STATE] = extended_state
-
-        return new_observation
-
-    def get_config(self) -> dict[str, Any]:
-        """
-        Returns the configuration of the step for serialization.
-
-        Returns:
-            A dictionary containing the time step `dt`.
-        """
-        return {
-            "dt": self.dt,
-        }
-
-    def reset(self) -> None:
-        """Resets the internal state, clearing the last known joint positions."""
-        self.last_joint_positions = None
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Updates the `observation.state` feature to reflect the added velocities.
-
-        This method doubles the size of the first dimension of the `observation.state`
-        shape to account for the concatenation of position and velocity vectors.
-
-        Args:
-            features: The policy features dictionary.
-
-        Returns:
-            The updated policy features dictionary.
-        """
-        if OBS_STATE in features[PipelineFeatureType.OBSERVATION]:
-            original_feature = features[PipelineFeatureType.OBSERVATION][OBS_STATE]
-            # Double the shape to account for positions + velocities
-            new_shape = (original_feature.shape[0] * 2,) + original_feature.shape[1:]
-
-            features[PipelineFeatureType.OBSERVATION][OBS_STATE] = PolicyFeature(
-                type=original_feature.type, shape=new_shape
-            )
-        return features
-
-
-@dataclass
-@ProcessorStepRegistry.register("current_processor")
-class MotorCurrentProcessorStep(ObservationProcessorStep):
-    """
-    Reads motor currents from a robot and appends them to the observation state.
-
-    This step queries the robot's hardware interface to get the present current
-    for each motor and concatenates this information to the existing state vector.
-
-    Attributes:
-        robot: An instance of a `lerobot` Robot class that provides access to
-               the hardware bus.
-    """
-
-    robot: Robot | None = None
-
-    def observation(self, observation: dict) -> dict:
-        """
-        Fetches motor currents and adds them to the observation state.
-
-        Args:
-            observation: The input observation dictionary.
-
-        Returns:
-            A new observation dictionary with the `observation.state` tensor
-            extended to include motor currents.
-
-        Raises:
-            ValueError: If the `robot` attribute has not been set.
-        """
-        # Get current values from robot state
-        if self.robot is None:
-            raise ValueError("Robot is not set")
-
-        present_current_dict = self.robot.bus.sync_read("Present_Current")  # type: ignore[attr-defined]
-        motor_currents = torch.tensor(
-            [present_current_dict[name] for name in self.robot.bus.motors],  # type: ignore[attr-defined]
-            dtype=torch.float32,
-        ).unsqueeze(0)
-
-        current_state = observation.get(OBS_STATE)
-        if current_state is None:
-            return observation
-
-        extended_state = torch.cat([current_state, motor_currents], dim=-1)
-
-        # Create new observation dict
-        new_observation = dict(observation)
-        new_observation[OBS_STATE] = extended_state
-
-        return new_observation
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Updates the `observation.state` feature to reflect the added motor currents.
-
-        This method increases the size of the first dimension of the `observation.state`
-        shape by the number of motors in the robot.
-
-        Args:
-            features: The policy features dictionary.
-
-        Returns:
-            The updated policy features dictionary.
-        """
-        if OBS_STATE in features[PipelineFeatureType.OBSERVATION] and self.robot is not None:
-            original_feature = features[PipelineFeatureType.OBSERVATION][OBS_STATE]
-            # Add motor current dimensions to the original state shape
-            num_motors = 0
-            if hasattr(self.robot, "bus") and hasattr(self.robot.bus, "motors"):  # type: ignore[attr-defined]
-                num_motors = len(self.robot.bus.motors)  # type: ignore[attr-defined]
-
-            if num_motors > 0:
-                new_shape = (original_feature.shape[0] + num_motors,) + original_feature.shape[1:]
-                features[PipelineFeatureType.OBSERVATION][OBS_STATE] = PolicyFeature(
-                    type=original_feature.type, shape=new_shape
-                )
-        return features

src/lerobot/processor/migrate_policy_normalization.py

@@ -15,22 +15,16 @@
 # limitations under the License.
 
 """
-A generic script to migrate LeRobot policies with built-in normalization layers to the new
-pipeline-based processor system.
+Generic script to migrate any policy model with normalization layers to the new pipeline-based system.
 
-This script performs the following steps:
-1.  Loads a pretrained policy model and its configuration from a local path or the
-    Hugging Face Hub.
-2.  Scans the model's state dictionary to extract normalization statistics (e.g., mean,
-    std, min, max) for all features.
-3.  Creates two new processor pipelines:
-    - A preprocessor that normalizes inputs (observations) and outputs (actions).
-    - A postprocessor that unnormalizes outputs (actions) for inference.
-4.  Removes the original normalization layers from the model's state dictionary,
-    creating a "clean" model.
-5.  Saves the new clean model, the preprocessor, the postprocessor, and a generated
-    model card to a new directory.
-6.  Optionally pushes all the new artifacts to the Hugging Face Hub.
+This script:
+1. Loads an existing pretrained policy model
+2. Extracts normalization statistics from the model
+3. Creates both preprocessor and postprocessor:
+   - Preprocessor: normalizes both inputs (observations) and outputs (actions) for training
+   - Postprocessor: unnormalizes outputs (actions) for inference
+4. Removes normalization layers from the model state_dict
+5. Saves the new model and both processors
 
 Usage:
     python src/lerobot/processor/migrate_policy_normalization.py \
@@ -52,12 +46,11 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file as load_safetensors
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-
-from .batch_processor import AddBatchDimensionProcessorStep
-from .device_processor import DeviceProcessorStep
-from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
-from .pipeline import PolicyProcessorPipeline
-from .rename_processor import RenameObservationsProcessorStep
+from lerobot.processor.batch_processor import ToBatchProcessor
+from lerobot.processor.device_processor import DeviceProcessor
+from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor
+from lerobot.processor.pipeline import RobotProcessor
+from lerobot.processor.rename_processor import RenameProcessor
 
 # Policy type to class mapping
 POLICY_CLASSES = {
@@ -74,21 +67,7 @@ POLICY_CLASSES = {
 
 
 def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
-    """
-    Scans a model's state_dict to find and extract normalization statistics.
-
-    This function identifies keys corresponding to normalization layers (e.g., those
-    for mean, std, min, max) based on a set of predefined patterns and organizes
-    them into a nested dictionary.
-
-    Args:
-        state_dict: The state dictionary of a pretrained policy model.
-
-    Returns:
-        A nested dictionary where outer keys are feature names (e.g.,
-        'observation.state') and inner keys are statistic types ('mean', 'std'),
-        mapping to their corresponding tensor values.
-    """
+    """Extract normalization statistics from model state_dict."""
     stats = {}
 
     # Define patterns to match and their prefixes to remove
@@ -132,25 +111,7 @@ def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str
 def detect_features_and_norm_modes(
     config: dict[str, Any], stats: dict[str, dict[str, torch.Tensor]]
 ) -> tuple[dict[str, PolicyFeature], dict[FeatureType, NormalizationMode]]:
-    """
-    Infers policy features and normalization modes from the model config and stats.
-
-    This function first attempts to find feature definitions and normalization
-    mappings directly from the policy's configuration file. If this information is
-    not present, it infers it from the extracted normalization statistics, using
-    tensor shapes to determine feature shapes and the presence of specific stat
-    keys (e.g., 'mean'/'std' vs 'min'/'max') to determine the normalization mode.
-    It applies sensible defaults if inference is not possible.
-
-    Args:
-        config: The policy's configuration dictionary from `config.json`.
-        stats: The normalization statistics extracted from the model's state_dict.
-
-    Returns:
-        A tuple containing:
-        - A dictionary mapping feature names to `PolicyFeature` objects.
-        - A dictionary mapping `FeatureType` enums to `NormalizationMode` enums.
-    """
+    """Detect features and normalization modes from config and stats."""
     features = {}
     norm_modes = {}
 
@@ -242,19 +203,7 @@ def detect_features_and_norm_modes(
 
 
 def remove_normalization_layers(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-    """
-    Creates a new state_dict with all normalization-related layers removed.
-
-    This function filters the original state dictionary, excluding any keys that
-    match a set of predefined patterns associated with normalization modules.
-
-    Args:
-        state_dict: The original model state dictionary.
-
-    Returns:
-        A new state dictionary containing only the core model weights, without
-        any normalization parameters.
-    """
+    """Remove normalization layers from state_dict."""
     new_state_dict = {}
 
     # Patterns to remove
@@ -278,16 +227,7 @@ def remove_normalization_layers(state_dict: dict[str, torch.Tensor]) -> dict[str
 
 
 def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[str, PolicyFeature]:
-    """
-    Converts a feature dictionary from the old config format to the new `PolicyFeature` format.
-
-    Args:
-        features_dict: The feature dictionary in the old format, where values are
-                       simple dictionaries (e.g., `{"shape": [7]}`).
-
-    Returns:
-        A dictionary mapping feature names to `PolicyFeature` dataclass objects.
-    """
+    """Convert features from old format to PolicyFeature objects."""
     converted_features = {}
 
     for key, feature_dict in features_dict.items():
@@ -313,18 +253,8 @@ def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[
 def load_model_from_hub(
     repo_id: str, revision: str = None
 ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any]]:
-    """
-    Downloads and loads a model's state_dict and configs from the Hugging Face Hub.
-
-    Args:
-        repo_id: The repository ID on the Hub (e.g., 'lerobot/aloha').
-        revision: The specific git revision (branch, tag, or commit hash) to use.
-
-    Returns:
-        A tuple containing the model's state dictionary, the policy configuration,
-        and the training configuration.
-    """
-    # Download files.
+    """Load model state_dict and config from hub."""
+    # Download files
     safetensors_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", revision=revision)
 
     config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
@@ -473,8 +403,8 @@ def main():
     # Now create preprocessor and postprocessor with cleaned_config available
     print("Creating preprocessor and postprocessor...")
     # The pattern from existing processor factories:
-    # - Preprocessor has two NormalizerProcessorSteps: one for input_features, one for output_features
-    # - Postprocessor has one UnnormalizerProcessorStep for output_features only
+    # - Preprocessor has two NormalizerProcessors: one for input_features, one for output_features
+    # - Postprocessor has one UnnormalizerProcessor for output_features only
 
     # Get features from cleaned_config (now they're PolicyFeature objects)
     input_features = cleaned_config.get("input_features", {})
@@ -482,23 +412,23 @@ def main():
 
     # Create preprocessor with two normalizers (following the pattern from processor factories)
     preprocessor_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        NormalizerProcessorStep(
+        RenameProcessor(rename_map={}),
+        NormalizerProcessor(
             features={**input_features, **output_features},
             norm_map=norm_map,
             stats=stats,
         ),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=policy_config.device),
+        ToBatchProcessor(),
+        DeviceProcessor(device=policy_config.device),
     ]
-    preprocessor = PolicyProcessorPipeline(steps=preprocessor_steps, name="robot_preprocessor")
+    preprocessor = RobotProcessor(steps=preprocessor_steps, name="robot_preprocessor")
 
     # Create postprocessor with unnormalizer for outputs only
     postprocessor_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(features=output_features, norm_map=norm_map, stats=stats),
+        DeviceProcessor(device="cpu"),
+        UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats),
     ]
-    postprocessor = PolicyProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor")
+    postprocessor = RobotProcessor(steps=postprocessor_steps, name="robot_postprocessor")
 
     # Determine hub repo ID if pushing to hub
     if args.push_to_hub:

src/lerobot/processor/normalize_processor.py

@@ -1,182 +1,234 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 from __future__ import annotations
 
+from collections.abc import Mapping
 from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any
 
+import numpy as np
 import torch
 from torch import Tensor
 
-from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, RobotProcessor, TransitionKey
 
-from .converters import from_tensor_to_numpy, to_tensor
-from .core import EnvTransition, TransitionKey
-from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry
+
+def _convert_stats_to_tensors(stats: dict[str, dict[str, Any]]) -> dict[str, dict[str, Tensor]]:
+    """Convert numpy arrays and other types to torch tensors."""
+    tensor_stats: dict[str, dict[str, Tensor]] = {}
+    for key, sub in stats.items():
+        tensor_stats[key] = {}
+        for stat_name, value in sub.items():
+            if isinstance(value, np.ndarray):
+                tensor_val = torch.from_numpy(value.astype(np.float32))
+            elif isinstance(value, torch.Tensor):
+                tensor_val = value.to(dtype=torch.float32)
+            elif isinstance(value, (int, float, list, tuple)):
+                tensor_val = torch.tensor(value, dtype=torch.float32)
+            else:
+                raise TypeError(f"Unsupported type for stats['{key}']['{stat_name}']: {type(value)}")
+            tensor_stats[key][stat_name] = tensor_val
+    return tensor_stats
 
 
 @dataclass
-class _NormalizationMixin:
-    """
-    A mixin class providing core functionality for normalization and unnormalization.
+@ProcessorStepRegistry.register(name="normalizer_processor")
+class NormalizerProcessor:
+    """Normalizes observations and actions in a single processor step.
 
-    This class manages normalization statistics (`stats`), converts them to tensors for
-    efficient computation, handles device placement, and implements the logic for
-    applying normalization transformations (mean/std and min/max). It is designed to
-    be inherited by concrete `ProcessorStep` implementations and should not be used
-    directly.
+    This processor handles normalization of both observation and action tensors
+    using either mean/std normalization or min/max scaling to a [-1, 1] range.
 
-    Attributes:
-        features: A dictionary mapping feature names to `PolicyFeature` objects, defining
-            the data structure to be processed.
-        norm_map: A dictionary mapping `FeatureType` to `NormalizationMode`, specifying
-            which normalization method to use for each type of feature.
-        stats: A dictionary containing the normalization statistics (e.g., mean, std,
-            min, max) for each feature.
-        device: The PyTorch device on which to store and perform tensor operations.
-        eps: A small epsilon value to prevent division by zero in normalization
-            calculations.
-        normalize_observation_keys: An optional set of keys to selectively apply
-            normalization to specific observation features.
-        _tensor_stats: An internal dictionary holding the normalization statistics as
-            PyTorch tensors.
+    For each tensor key in the stats dictionary, the processor will:
+    - Use mean/std normalization if those statistics are provided: (x - mean) / std
+    - Use min/max scaling if those statistics are provided: 2 * (x - min) / (max - min) - 1
+
+    The processor can be configured to normalize only specific keys by setting
+    the normalize_keys parameter.
     """
 
+    # Features and normalisation map are mandatory to match the design of normalize.py
     features: dict[str, PolicyFeature]
     norm_map: dict[FeatureType, NormalizationMode]
+
+    # Pre-computed statistics coming from dataset.meta.stats for instance.
     stats: dict[str, dict[str, Any]] | None = None
-    device: torch.device | str | None = None
-    dtype: torch.dtype | None = None
+
+    # Explicit subset of keys to normalise. If ``None`` every key (except
+    # "action") found in ``stats`` will be normalised. Using a ``set`` makes
+    # membership checks O(1).
+    normalize_keys: set[str] | None = None
+
     eps: float = 1e-8
-    normalize_observation_keys: set[str] | None = None
 
     _tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
 
+    @classmethod
+    def from_lerobot_dataset(
+        cls,
+        dataset: LeRobotDataset,
+        features: dict[str, PolicyFeature],
+        norm_map: dict[FeatureType, NormalizationMode],
+        *,
+        normalize_keys: set[str] | None = None,
+        eps: float = 1e-8,
+    ) -> NormalizerProcessor:
+        """Factory helper that pulls statistics from a :class:`LeRobotDataset`.
+
+        The features and norm_map parameters are mandatory to match the design
+        pattern used in normalize.py.
+        """
+
+        return cls(
+            features=features,
+            norm_map=norm_map,
+            stats=dataset.meta.stats,
+            normalize_keys=normalize_keys,
+            eps=eps,
+        )
+
     def __post_init__(self):
-        """
-        Initializes the mixin after dataclass construction.
+        # Handle deserialization from JSON config
+        if self.features and isinstance(list(self.features.values())[0], dict):
+            # Features came from JSON - need to reconstruct PolicyFeature objects
+            reconstructed_features = {}
+            for key, ft_dict in self.features.items():
+                reconstructed_features[key] = PolicyFeature(
+                    type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
+                )
+            self.features = reconstructed_features
 
-        This method handles the robust deserialization of `features` and `norm_map`
-        from JSON-compatible formats (where enums become strings and tuples become
-        lists) and converts the provided `stats` dictionary into a dictionary of
-        tensors (`_tensor_stats`) on the specified device.
-        """
-        # Robust JSON deserialization handling (guard empty maps).
-        if self.features:
-            first_val = next(iter(self.features.values()))
-            if isinstance(first_val, dict):
-                reconstructed = {}
-                for key, ft_dict in self.features.items():
-                    reconstructed[key] = PolicyFeature(
-                        type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
-                    )
-                self.features = reconstructed
+        if self.norm_map and isinstance(list(self.norm_map.keys())[0], str):
+            # norm_map came from JSON - need to reconstruct enum keys and values
+            reconstructed_norm_map = {}
+            for ft_type_str, norm_mode_str in self.norm_map.items():
+                reconstructed_norm_map[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
+            self.norm_map = reconstructed_norm_map
 
-        if self.norm_map:
-            # if keys are strings (JSON), rebuild enum map
-            if all(isinstance(k, str) for k in self.norm_map.keys()):
-                reconstructed = {}
-                for ft_type_str, norm_mode_str in self.norm_map.items():
-                    reconstructed[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
-                self.norm_map = reconstructed
-
-        # Convert stats to tensors and move to the target device once during initialization.
+        # Convert statistics once so we avoid repeated numpy→Tensor conversions
+        # during runtime.
         self.stats = self.stats or {}
-        if self.dtype is None:
-            self.dtype = torch.float32
-        self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype)
+        self._tensor_stats = _convert_stats_to_tensors(self.stats)
 
-    def to(
-        self, device: torch.device | str | None = None, dtype: torch.dtype | None = None
-    ) -> _NormalizationMixin:
-        """
-        Moves the processor's normalization stats to the specified device.
+        # Ensure *normalize_keys* is a set for fast look-ups and compare by
+        # value later when returning the configuration.
+        if self.normalize_keys is not None and not isinstance(self.normalize_keys, set):
+            self.normalize_keys = set(self.normalize_keys)
 
-        Args:
-            device: The target PyTorch device.
+    def _normalize_obs(self, observation, normalized_info):
+        if observation is None:
+            return None
 
-        Returns:
-            The instance of the class, allowing for method chaining.
-        """
-        if device is not None:
-            self.device = device
-        if dtype is not None:
-            self.dtype = dtype
-        self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype)
-        return self
+        # Decide which keys should be normalised for this call.
+        if self.normalize_keys is not None:
+            keys_to_norm = self.normalize_keys
+        else:
+            # Use feature map to skip action keys.
+            keys_to_norm = {k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION}
 
-    def state_dict(self) -> dict[str, Tensor]:
-        """
-        Returns the normalization statistics as a flat state dictionary.
+        processed = dict(observation)
+        for key in keys_to_norm:
+            if key not in processed or key not in self.features:
+                continue
 
-        All tensors are moved to the CPU before being returned, which is standard practice
-        for saving state dictionaries.
+            # Check the normalization mode for this feature type
+            feature = self.features[key]
+            norm_mode = self.norm_map.get(feature.type, NormalizationMode.IDENTITY)
 
-        Returns:
-            A flat dictionary mapping from `'feature_name.stat_name'` to the
-            corresponding statistics tensor on the CPU.
-        """
-        flat: dict[str, Tensor] = {}
-        for key, sub in self._tensor_stats.items():
-            for stat_name, tensor in sub.items():
-                flat[f"{key}.{stat_name}"] = tensor.cpu()  # Always save to CPU
-        return flat
+            # Skip normalization if mode is IDENTITY
+            if norm_mode is NormalizationMode.IDENTITY:
+                normalized_info[key] = "IDENTITY"
+                continue
 
-    def load_state_dict(self, state: dict[str, Tensor]) -> None:
-        """
-        Loads normalization statistics from a state dictionary.
+            # Skip if no stats available for this key
+            if key not in self._tensor_stats:
+                continue
 
-        The loaded tensors are moved to the processor's configured device.
-
-        Args:
-            state: A flat state dictionary with keys in the format
-                   `'feature_name.stat_name'`.
-        """
-        self._tensor_stats.clear()
-        for flat_key, tensor in state.items():
-            key, stat_name = flat_key.rsplit(".", 1)
-            # Load to the processor's configured device.
-            self._tensor_stats.setdefault(key, {})[stat_name] = tensor.to(
-                dtype=torch.float32, device=self.device
+            orig_val = processed[key]
+            tensor = (
+                orig_val.to(dtype=torch.float32)
+                if isinstance(orig_val, torch.Tensor)
+                else torch.as_tensor(orig_val, dtype=torch.float32)
             )
+            stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
 
-        # Reconstruct the original stats dict from tensor stats for compatibility with to() method
-        # and other functions that rely on self.stats
+            if norm_mode is NormalizationMode.MEAN_STD:
+                if "mean" in stats and "std" in stats:
+                    mean, std = stats["mean"], stats["std"]
+                    processed[key] = (tensor - mean) / (std + self.eps)
+                    normalized_info[key] = "MEAN_STD"
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                if "min" in stats and "max" in stats:
+                    min_val, max_val = stats["min"], stats["max"]
+                    processed[key] = 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
+                    normalized_info[key] = "MIN_MAX"
+            else:
+                raise ValueError(f"Unsupported normalization mode: {norm_mode}")
 
-        self.stats = {}
-        for key, tensor_dict in self._tensor_stats.items():
-            self.stats[key] = {}
-            for stat_name, tensor in tensor_dict.items():
-                # Convert tensor back to python/numpy format
-                self.stats[key][stat_name] = from_tensor_to_numpy(tensor)
+        return processed
+
+    def _normalize_action(self, action, normalized_info):
+        if action is None:
+            return action
+
+        # Check the normalization mode for actions
+        norm_mode = self.norm_map.get(FeatureType.ACTION, NormalizationMode.IDENTITY)
+
+        # Skip normalization if mode is IDENTITY
+        if norm_mode is NormalizationMode.IDENTITY:
+            normalized_info["action"] = "IDENTITY"
+            return action
+
+        # Skip if no stats available for actions
+        if "action" not in self._tensor_stats:
+            return action
+
+        tensor = (
+            action.to(dtype=torch.float32)
+            if isinstance(action, torch.Tensor)
+            else torch.as_tensor(action, dtype=torch.float32)
+        )
+        stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
+
+        if norm_mode is NormalizationMode.MEAN_STD:
+            if "mean" in stats and "std" in stats:
+                mean, std = stats["mean"], stats["std"]
+                normalized_info["action"] = "MEAN_STD"
+                return (tensor - mean) / (std + self.eps)
+        elif norm_mode is NormalizationMode.MIN_MAX:
+            if "min" in stats and "max" in stats:
+                min_val, max_val = stats["min"], stats["max"]
+                normalized_info["action"] = "MIN_MAX"
+                return 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
+        else:
+            raise ValueError(f"Unsupported normalization mode: {norm_mode}")
+
+        # If we reach here, the required stats for the normalization mode are not available
+        raise ValueError(f"Action stats must contain appropriate values for {norm_mode} normalization")
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        # Track what was normalized
+        normalized_info = {}
+
+        observation = self._normalize_obs(transition.get(TransitionKey.OBSERVATION), normalized_info)
+        action = self._normalize_action(transition.get(TransitionKey.ACTION), normalized_info)
+
+        # Create a new transition with normalized values
+        new_transition = transition.copy()
+        new_transition[TransitionKey.OBSERVATION] = observation
+        new_transition[TransitionKey.ACTION] = action
+
+        # Add normalization info to complementary data
+        if normalized_info:
+            comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+            comp_data = {} if comp_data is None else dict(comp_data)
+            comp_data["normalized_keys"] = normalized_info
+            new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
+
+        return new_transition
 
     def get_config(self) -> dict[str, Any]:
-        """
-        Returns a serializable dictionary of the processor's configuration.
-
-        This method is used when saving the processor to disk, ensuring that its
-        configuration can be reconstructed later.
-
-        Returns:
-            A JSON-serializable dictionary containing the configuration.
-        """
         config = {
             "eps": self.eps,
             "features": {
@@ -184,257 +236,240 @@ class _NormalizationMixin:
             },
             "norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
         }
-        if self.normalize_observation_keys is not None:
-            config["normalize_observation_keys"] = sorted(self.normalize_observation_keys)
+        if self.normalize_keys is not None:
+            # Serialise as a list for YAML / JSON friendliness
+            config["normalize_keys"] = sorted(self.normalize_keys)
         return config
 
-    def _normalize_observation(self, observation: dict[str, Any], inverse: bool) -> dict[str, Tensor]:
-        """
-        Applies (un)normalization to all relevant features in an observation dictionary.
+    def state_dict(self) -> dict[str, Tensor]:
+        flat = {}
+        for key, sub in self._tensor_stats.items():
+            for stat_name, tensor in sub.items():
+                flat[f"{key}.{stat_name}"] = tensor
+        return flat
 
-        Args:
-            observation: The observation dictionary to process.
-            inverse: If `True`, applies unnormalization; otherwise, applies normalization.
+    def load_state_dict(self, state: Mapping[str, Tensor]) -> None:
+        self._tensor_stats.clear()
+        for flat_key, tensor in state.items():
+            key, stat_name = flat_key.rsplit(".", 1)
+            self._tensor_stats.setdefault(key, {})[stat_name] = tensor
 
-        Returns:
-            A new observation dictionary with the transformed tensor values.
-        """
-        new_observation = dict(observation)
-        for key, feature in self.features.items():
-            if self.normalize_observation_keys is not None and key not in self.normalize_observation_keys:
-                continue
-            if feature.type != FeatureType.ACTION and key in new_observation:
-                # Convert to tensor but preserve original dtype for adaptation logic
-                tensor = torch.as_tensor(new_observation[key])
-                new_observation[key] = self._apply_transform(tensor, key, feature.type, inverse=inverse)
-        return new_observation
+    def reset(self):
+        pass
 
-    def _normalize_action(self, action: Any, inverse: bool) -> Tensor:
-        # Convert to tensor but preserve original dtype for adaptation logic
-        """
-        Applies (un)normalization to an action tensor.
-
-        Args:
-            action: The action tensor to process.
-            inverse: If `True`, applies unnormalization; otherwise, applies normalization.
-
-        Returns:
-            The transformed action tensor.
-        """
-        tensor = torch.as_tensor(action)
-        processed_action = self._apply_transform(tensor, "action", FeatureType.ACTION, inverse=inverse)
-        return processed_action
-
-    def _apply_transform(
-        self, tensor: Tensor, key: str, feature_type: FeatureType, *, inverse: bool = False
-    ) -> Tensor:
-        """
-        Core logic to apply a normalization or unnormalization transformation to a tensor.
-
-        This method selects the appropriate normalization mode (e.g., mean/std, min/max)
-        based on the feature type and applies the corresponding mathematical operation.
-
-        Args:
-            tensor: The input tensor to transform.
-            key: The feature key corresponding to the tensor.
-            feature_type: The `FeatureType` of the tensor.
-            inverse: If `True`, applies the inverse transformation (unnormalization).
-
-        Returns:
-            The transformed tensor.
-
-        Raises:
-            ValueError: If an unsupported normalization mode is encountered.
-        """
-        norm_mode = self.norm_map.get(feature_type, NormalizationMode.IDENTITY)
-        if norm_mode == NormalizationMode.IDENTITY or key not in self._tensor_stats:
-            return tensor
-
-        if norm_mode not in (NormalizationMode.MEAN_STD, NormalizationMode.MIN_MAX):
-            raise ValueError(f"Unsupported normalization mode: {norm_mode}")
-
-        # For Accelerate compatibility: Ensure stats are on the same device and dtype as the input tensor
-        if self._tensor_stats and key in self._tensor_stats:
-            first_stat = next(iter(self._tensor_stats[key].values()))
-            if first_stat.device != tensor.device or first_stat.dtype != tensor.dtype:
-                self.to(device=tensor.device, dtype=tensor.dtype)
-
-        stats = self._tensor_stats[key]
-
-        if norm_mode == NormalizationMode.MEAN_STD and "mean" in stats and "std" in stats:
-            mean, std = stats["mean"], stats["std"]
-            # Avoid division by zero by adding a small epsilon.
-            denom = std + self.eps
-            if inverse:
-                return tensor * std + mean
-            return (tensor - mean) / denom
-
-        if norm_mode == NormalizationMode.MIN_MAX and "min" in stats and "max" in stats:
-            min_val, max_val = stats["min"], stats["max"]
-            denom = max_val - min_val
-            # When min_val == max_val, substitute the denominator with a small epsilon
-            # to prevent division by zero. This consistently maps an input equal to
-            # min_val to -1, ensuring a stable transformation.
-            denom = torch.where(
-                denom == 0, torch.tensor(self.eps, device=tensor.device, dtype=tensor.dtype), denom
-            )
-            if inverse:
-                # Map from [-1, 1] back to [min, max]
-                return (tensor + 1) / 2 * denom + min_val
-            # Map from [min, max] to [-1, 1]
-            return 2 * (tensor - min_val) / denom - 1
-
-        # If necessary stats are missing, return input unchanged.
-        return tensor
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="normalizer_processor")
-class NormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
-    """
-    A processor step that applies normalization to observations and actions in a transition.
-
-    This class uses the logic from `_NormalizationMixin` to perform forward normalization
-    (e.g., scaling data to have zero mean and unit variance, or to the range [-1, 1]).
-    It is typically used in the pre-processing pipeline before feeding data to a policy.
-    """
-
-    @classmethod
-    def from_lerobot_dataset(
-        cls,
-        dataset: LeRobotDataset,
-        features: dict[str, PolicyFeature],
-        norm_map: dict[FeatureType, NormalizationMode],
-        *,
-        normalize_observation_keys: set[str] | None = None,
-        eps: float = 1e-8,
-        device: torch.device | str | None = None,
-    ) -> NormalizerProcessorStep:
-        """
-        Creates a `NormalizerProcessorStep` instance using statistics from a `LeRobotDataset`.
-
-        Args:
-            dataset: The dataset from which to extract normalization statistics.
-            features: The feature definition for the processor.
-            norm_map: The mapping from feature types to normalization modes.
-            normalize_observation_keys: An optional set of observation keys to normalize.
-            eps: A small epsilon value for numerical stability.
-            device: The target device for the processor.
-
-        Returns:
-            A new instance of `NormalizerProcessorStep`.
-        """
-        return cls(
-            features=features,
-            norm_map=norm_map,
-            stats=dataset.meta.stats,
-            normalize_observation_keys=normalize_observation_keys,
-            eps=eps,
-            device=device,
-        )
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        new_transition = transition.copy()
-
-        # Handle observation normalization.
-        observation = new_transition.get(TransitionKey.OBSERVATION)
-        if observation is not None:
-            new_transition[TransitionKey.OBSERVATION] = self._normalize_observation(
-                observation, inverse=False
-            )
-
-        # Handle action normalization.
-        action = new_transition.get(TransitionKey.ACTION)
-        if action is not None:
-            new_transition[TransitionKey.ACTION] = self._normalize_action(action, inverse=False)
-
-        return new_transition
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
 @dataclass
 @ProcessorStepRegistry.register(name="unnormalizer_processor")
-class UnnormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
-    """
-    A processor step that applies unnormalization to observations and actions.
+class UnnormalizerProcessor:
+    """Inverse normalisation for observations and actions.
 
-    This class inverts the normalization process, scaling data back to its original
-    range. It is typically used in the post-processing pipeline to convert a policy's
-    normalized action output into a format that can be executed by a robot or
-    environment.
+    Exactly mirrors :class:`NormalizerProcessor` but applies the inverse
+    transform.
     """
 
+    features: dict[str, PolicyFeature]
+    norm_map: dict[FeatureType, NormalizationMode]
+    stats: dict[str, dict[str, Any]] | None = None
+
+    _tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
+
     @classmethod
     def from_lerobot_dataset(
         cls,
         dataset: LeRobotDataset,
         features: dict[str, PolicyFeature],
         norm_map: dict[FeatureType, NormalizationMode],
-        *,
-        device: torch.device | str | None = None,
-    ) -> UnnormalizerProcessorStep:
-        """
-        Creates an `UnnormalizerProcessorStep` using statistics from a `LeRobotDataset`.
+    ) -> UnnormalizerProcessor:
+        return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats)
 
-        Args:
-            dataset: The dataset from which to extract normalization statistics.
-            features: The feature definition for the processor.
-            norm_map: The mapping from feature types to normalization modes.
-            device: The target device for the processor.
+    def __post_init__(self):
+        # Handle deserialization from JSON config
+        if self.features and isinstance(list(self.features.values())[0], dict):
+            # Features came from JSON - need to reconstruct PolicyFeature objects
+            reconstructed_features = {}
+            for key, ft_dict in self.features.items():
+                reconstructed_features[key] = PolicyFeature(
+                    type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
+                )
+            self.features = reconstructed_features
 
-        Returns:
-            A new instance of `UnnormalizerProcessorStep`.
-        """
-        return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, device=device)
+        if self.norm_map and isinstance(list(self.norm_map.keys())[0], str):
+            # norm_map came from JSON - need to reconstruct enum keys and values
+            reconstructed_norm_map = {}
+            for ft_type_str, norm_mode_str in self.norm_map.items():
+                reconstructed_norm_map[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
+            self.norm_map = reconstructed_norm_map
+
+        self.stats = self.stats or {}
+        self._tensor_stats = _convert_stats_to_tensors(self.stats)
+
+    def _unnormalize_obs(self, observation, unnormalized_info):
+        if observation is None:
+            return None
+        keys = [k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION]
+        processed = dict(observation)
+        for key in keys:
+            if key not in processed or key not in self.features:
+                continue
+
+            # Check the normalization mode for this feature type
+            feature = self.features[key]
+            norm_mode = self.norm_map.get(feature.type, NormalizationMode.IDENTITY)
+
+            # Skip unnormalization if mode is IDENTITY
+            if norm_mode is NormalizationMode.IDENTITY:
+                unnormalized_info[key] = "IDENTITY"
+                continue
+
+            # Skip if no stats available for this key
+            if key not in self._tensor_stats:
+                continue
+
+            orig_val = processed[key]
+            tensor = (
+                orig_val.to(dtype=torch.float32)
+                if isinstance(orig_val, torch.Tensor)
+                else torch.as_tensor(orig_val, dtype=torch.float32)
+            )
+            stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
+
+            if norm_mode is NormalizationMode.MEAN_STD:
+                if "mean" in stats and "std" in stats:
+                    mean, std = stats["mean"], stats["std"]
+                    processed[key] = tensor * std + mean
+                    unnormalized_info[key] = "MEAN_STD"
+            elif norm_mode is NormalizationMode.MIN_MAX:
+                if "min" in stats and "max" in stats:
+                    min_val, max_val = stats["min"], stats["max"]
+                    processed[key] = (tensor + 1) / 2 * (max_val - min_val) + min_val
+                    unnormalized_info[key] = "MIN_MAX"
+            else:
+                raise ValueError(f"Unsupported normalization mode: {norm_mode}")
+
+        return processed
+
+    def _unnormalize_action(self, action, unnormalized_info):
+        if action is None:
+            return action
+
+        # Check the normalization mode for actions
+        norm_mode = self.norm_map.get(FeatureType.ACTION, NormalizationMode.IDENTITY)
+
+        # Skip unnormalization if mode is IDENTITY
+        if norm_mode is NormalizationMode.IDENTITY:
+            unnormalized_info["action"] = "IDENTITY"
+            return action
+
+        # Skip if no stats available for actions
+        if "action" not in self._tensor_stats:
+            return action
+
+        tensor = (
+            action.to(dtype=torch.float32)
+            if isinstance(action, torch.Tensor)
+            else torch.as_tensor(action, dtype=torch.float32)
+        )
+        stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
+
+        if norm_mode is NormalizationMode.MEAN_STD:
+            if "mean" in stats and "std" in stats:
+                mean, std = stats["mean"], stats["std"]
+                unnormalized_info["action"] = "MEAN_STD"
+                return tensor * std + mean
+        elif norm_mode is NormalizationMode.MIN_MAX:
+            if "min" in stats and "max" in stats:
+                min_val, max_val = stats["min"], stats["max"]
+                unnormalized_info["action"] = "MIN_MAX"
+                return (tensor + 1) / 2 * (max_val - min_val) + min_val
+        else:
+            raise ValueError(f"Unsupported normalization mode: {norm_mode}")
+
+        # If we reach here, the required stats for the normalization mode are not available
+        raise ValueError(f"Action stats must contain appropriate values for {norm_mode} normalization")
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
+        # Track what was unnormalized
+        unnormalized_info = {}
+
+        observation = self._unnormalize_obs(transition.get(TransitionKey.OBSERVATION), unnormalized_info)
+        action = self._unnormalize_action(transition.get(TransitionKey.ACTION), unnormalized_info)
+
+        # Create a new transition with unnormalized values
         new_transition = transition.copy()
+        new_transition[TransitionKey.OBSERVATION] = observation
+        new_transition[TransitionKey.ACTION] = action
 
-        # Handle observation unnormalization.
-        observation = new_transition.get(TransitionKey.OBSERVATION)
-        if observation is not None:
-            new_transition[TransitionKey.OBSERVATION] = self._normalize_observation(observation, inverse=True)
-
-        # Handle action unnormalization.
-        action = new_transition.get(TransitionKey.ACTION)
-        if action is not None:
-            new_transition[TransitionKey.ACTION] = self._normalize_action(action, inverse=True)
+        # Add unnormalization info to complementary data
+        if unnormalized_info:
+            comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+            comp_data = {} if comp_data is None else dict(comp_data)
+            comp_data["unnormalized_keys"] = unnormalized_info
+            new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
 
         return new_transition
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "features": {
+                key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
+            },
+            "norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
+        }
+
+    def state_dict(self) -> dict[str, Tensor]:
+        flat = {}
+        for key, sub in self._tensor_stats.items():
+            for stat_name, tensor in sub.items():
+                flat[f"{key}.{stat_name}"] = tensor
+        return flat
+
+    def load_state_dict(self, state: Mapping[str, Tensor]) -> None:
+        self._tensor_stats.clear()
+        for flat_key, tensor in state.items():
+            key, stat_name = flat_key.rsplit(".", 1)
+            self._tensor_stats.setdefault(key, {})[stat_name] = tensor
+
+    def reset(self):
+        pass
+
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
-def hotswap_stats(
-    policy_processor: PolicyProcessorPipeline, stats: dict[str, dict[str, Any]]
-) -> PolicyProcessorPipeline:
-    """
-    Replaces normalization statistics in an existing `PolicyProcessorPipeline` instance.
+def hotswap_stats(robot_processor: RobotProcessor, stats: dict[str, dict[str, Any]]) -> RobotProcessor:
+    robot_processor = deepcopy(robot_processor)
+    for step in robot_processor.steps:
+        if isinstance(step, NormalizerProcessor) or isinstance(step, UnnormalizerProcessor):
+            step: NormalizerProcessor | UnnormalizerProcessor
+            step.stats = stats
+            step._tensor_stats = _convert_stats_to_tensors(stats)
+    return robot_processor
 
-    This function creates a deep copy of the provided pipeline and updates the
-    statistics of any `NormalizerProcessorStep` or `UnnormalizerProcessorStep` it
-    contains. This is useful for adapting a trained policy to a new environment or
-    dataset with different data distributions without having to reconstruct the entire
-    pipeline.
+
+def rename_stats(stats: dict[str, dict[str, Any]], rename_map: dict[str, str]) -> dict[str, dict[str, Any]]:
+    """Rename keys in the stats dictionary according to the provided mapping.
 
     Args:
-        policy_processor: The policy processor pipeline to modify.
-        stats: The new dictionary of normalization statistics to apply.
+        stats: The statistics dictionary with structure {feature_key: {stat_name: value}}
+        rename_map: Dictionary mapping old key names to new key names
 
     Returns:
-        A new `PolicyProcessorPipeline` instance with the updated statistics.
+        A new stats dictionary with renamed keys
+
+    Example:
+        >>> stats = {"observation.state": {"mean": 0.0, "std": 1.0}, "action": {"mean": 0.5, "std": 0.5}}
+        >>> rename_map = {"observation.state": "observation.robot_state"}
+        >>> new_stats = rename_stats(stats, rename_map)
+        >>> # new_stats will have "observation.robot_state" instead of "observation.state"
     """
-    rp = deepcopy(policy_processor)
-    for step in rp.steps:
-        if isinstance(step, _NormalizationMixin):
-            step.stats = stats
-            # Re-initialize tensor_stats on the correct device.
-            step._tensor_stats = to_tensor(stats, device=step.device, dtype=step.dtype)
-    return rp
+    renamed_stats = {}
+
+    for old_key, sub_stats in stats.items():
+        # Use the new key if it exists in the rename map, otherwise keep the old key
+        new_key = rename_map.get(old_key, old_key)
+        renamed_stats[new_key] = deepcopy(sub_stats)
+
+    return renamed_stats

src/lerobot/processor/observation_processor.py

@@ -20,54 +20,32 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs.types import PolicyFeature
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-
-from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
+from lerobot.processor.pipeline import ObservationProcessor, ProcessorStepRegistry
 
 
 @dataclass
 @ProcessorStepRegistry.register(name="observation_processor")
-class VanillaObservationProcessorStep(ObservationProcessorStep):
+class VanillaObservationProcessor(ObservationProcessor):
     """
-    Processes standard Gymnasium observations into the LeRobot format.
+    Processes environment observations into the LeRobot format by handling both images and states.
 
-    This step handles both image and state data from a typical observation dictionary,
-    preparing it for use in a LeRobot policy.
+    Image processing:
+        - Converts channel-last (H, W, C) images to channel-first (C, H, W)
+        - Normalizes uint8 images ([0, 255]) to float32 ([0, 1])
+        - Adds a batch dimension if missing
+        - Supports single images and image dictionaries
 
-    **Image Processing:**
-    -   Converts channel-last (H, W, C), `uint8` images to channel-first (C, H, W),
-        `float32` tensors.
-    -   Normalizes pixel values from the [0, 255] range to [0, 1].
-    -   Adds a batch dimension if one is not already present.
-    -   Recognizes a single image under the key `"pixels"` and maps it to
-        `"observation.image"`.
-    -   Recognizes a dictionary of images under the key `"pixels"` and maps them
-        to `"observation.images.{camera_name}"`.
-
-    **State Processing:**
-    -   Maps the `"environment_state"` key to `"observation.environment_state"`.
-    -   Maps the `"agent_pos"` key to `"observation.state"`.
-    -   Converts NumPy arrays to PyTorch tensors.
-    -   Adds a batch dimension if one is not already present.
+    State processing:
+        - Maps 'environment_state' to observation.environment_state
+        - Maps 'agent_pos' to observation.state
+        - Converts numpy arrays to tensors
+        - Adds a batch dimension if missing
     """
 
     def _process_single_image(self, img: np.ndarray) -> Tensor:
-        """
-        Processes a single NumPy image array into a channel-first, normalized tensor.
-
-        Args:
-            img: A NumPy array representing the image, expected to be in channel-last
-                 (H, W, C) format with a `uint8` dtype.
-
-        Returns:
-            A `float32` PyTorch tensor in channel-first (B, C, H, W) format, with
-            pixel values normalized to the [0, 1] range.
-
-        Raises:
-            ValueError: If the input image does not appear to be in channel-last
-                        format or is not of `uint8` dtype.
-        """
+        """Process a single image array."""
         # Convert to tensor
         img_tensor = torch.from_numpy(img)
 
@@ -128,32 +106,18 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):
     def observation(self, observation):
         return self._process_observation(observation)
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        """Transforms feature keys to a standardized contract.
+        This method handles several renaming patterns:
+        - Exact matches (e.g., 'pixels' -> 'OBS_IMAGE').
+        - Prefixed exact matches (e.g., 'observation.pixels' -> 'OBS_IMAGE').
+        - Prefix matches (e.g., 'pixels.cam1' -> 'OBS_IMAGES.cam1').
+        - Prefixed prefix matches (e.g., 'observation.pixels.cam1' -> 'OBS_IMAGES.cam1').
+        - environment_state -> OBS_ENV_STATE,
+        - agent_pos -> OBS_STATE,
+        - observation.environment_state -> OBS_ENV_STATE,
+        - observation.agent_pos -> OBS_STATE
         """
-        Transforms feature keys from the Gym standard to the LeRobot standard.
-
-        This method standardizes the feature dictionary by renaming keys according
-        to LeRobot's conventions, ensuring that policies can be constructed correctly.
-        It handles various raw key formats, including those with an "observation." prefix.
-
-        **Renaming Rules:**
-        - `pixels` or `observation.pixels` -> `observation.image`
-        - `pixels.{cam}` or `observation.pixels.{cam}` -> `observation.images.{cam}`
-        - `environment_state` or `observation.environment_state` -> `observation.environment_state`
-        - `agent_pos` or `observation.agent_pos` -> `observation.state`
-
-        Args:
-            features: The policy features dictionary with Gym-style keys.
-
-        Returns:
-            The policy features dictionary with standardized LeRobot keys.
-        """
-        # Build a new features mapping keyed by the same FeatureType buckets
-        # We assume callers already placed features in the correct FeatureType.
-        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = {ft: {} for ft in features.keys()}
-
         exact_pairs = {
             "pixels": OBS_IMAGE,
             "environment_state": OBS_ENV_STATE,
@@ -164,43 +128,29 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):
             "pixels.": f"{OBS_IMAGES}.",
         }
 
-        # Iterate over all incoming feature buckets and normalize/move each entry
-        for src_ft, bucket in features.items():
-            for key, feat in list(bucket.items()):
-                handled = False
+        for key in list(features.keys()):
+            matched_prefix = False
+            for old_prefix, new_prefix in prefix_pairs.items():
+                prefixed_old = f"observation.{old_prefix}"
+                if key.startswith(prefixed_old):
+                    suffix = key[len(prefixed_old) :]
+                    features[f"{new_prefix}{suffix}"] = features.pop(key)
+                    matched_prefix = True
+                    break
 
-                # 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}"
-                    if key.startswith(prefixed_old):
-                        suffix = key[len(prefixed_old) :]
-                        new_key = f"{new_prefix}{suffix}"
-                        new_features[src_ft][new_key] = feat
-                        handled = True
+                if key.startswith(old_prefix):
+                    suffix = key[len(old_prefix) :]
+                    features[f"{new_prefix}{suffix}"] = features.pop(key)
+                    matched_prefix = True
+                    break
+
+            if matched_prefix:
+                continue
+
+            for old, new in exact_pairs.items():
+                if key == old or key == f"observation.{old}":
+                    if key in features:
+                        features[new] = features.pop(key)
                         break
 
-                    if key.startswith(old_prefix):
-                        suffix = key[len(old_prefix) :]
-                        new_key = f"{new_prefix}{suffix}"
-                        new_features[src_ft][new_key] = feat
-                        handled = True
-                        break
-
-                if handled:
-                    continue
-
-                # Exact-name rules (pixels, environment_state, agent_pos)
-                for old, new in exact_pairs.items():
-                    if key == old or key == f"observation.{old}":
-                        new_key = new
-                        new_features[src_ft][new_key] = feat
-                        handled = True
-                        break
-
-                if handled:
-                    continue
-
-                # Default: keep key in the same source FeatureType bucket
-                new_features[src_ft][key] = feat
-
-        return new_features
+        return features

src/lerobot/processor/rename_processor.py

@@ -13,30 +13,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-
-from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
+from lerobot.configs.types import PolicyFeature
+from lerobot.processor.pipeline import (
+    ObservationProcessor,
+    ProcessorStepRegistry,
+)
 
 
 @dataclass
-@ProcessorStepRegistry.register(name="rename_observations_processor")
-class RenameObservationsProcessorStep(ObservationProcessorStep):
-    """
-    A processor step that renames keys in an observation dictionary.
-
-    This step is useful for creating a standardized data interface by mapping keys
-    from an environment's format to the format expected by a LeRobot policy or
-    other downstream components.
-
-    Attributes:
-        rename_map: A dictionary mapping from old key names to new key names.
-                    Keys present in an observation that are not in this map will
-                    be kept with their original names.
-    """
+@ProcessorStepRegistry.register(name="rename_processor")
+class RenameProcessor(ObservationProcessor):
+    """Rename processor that renames keys in the observation."""
 
     rename_map: dict[str, str] = field(default_factory=dict)
 
@@ -53,41 +43,9 @@ class RenameObservationsProcessorStep(ObservationProcessorStep):
     def get_config(self) -> dict[str, Any]:
         return {"rename_map": self.rename_map}
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         """Transforms:
         - Each key in the observation that appears in `rename_map` is renamed to its value.
         - Keys not in `rename_map` remain unchanged.
         """
-        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = features.copy()
-        new_features[PipelineFeatureType.OBSERVATION] = {
-            self.rename_map.get(k, k): v for k, v in features[PipelineFeatureType.OBSERVATION].items()
-        }
-        return new_features
-
-
-def rename_stats(stats: dict[str, dict[str, Any]], rename_map: dict[str, str]) -> dict[str, dict[str, Any]]:
-    """
-    Renames the top-level keys in a statistics dictionary using a provided mapping.
-
-    This is a helper function typically used to keep normalization statistics
-    consistent with renamed observation or action features. It performs a defensive
-    deep copy to avoid modifying the original `stats` dictionary.
-
-    Args:
-        stats: A nested dictionary of statistics, where top-level keys are
-               feature names (e.g., `{"observation.state": {"mean": 0.5}}`).
-        rename_map: A dictionary mapping old feature names to new feature names.
-
-    Returns:
-        A new statistics dictionary with its top-level keys renamed. Returns an
-        empty dictionary if the input `stats` is empty.
-    """
-    if not stats:
-        return {}
-    renamed: dict[str, dict[str, Any]] = {}
-    for old_key, sub_stats in stats.items():
-        new_key = rename_map.get(old_key, old_key)
-        renamed[new_key] = deepcopy(sub_stats) if sub_stats is not None else {}
-    return renamed
+        return {self.rename_map.get(k, k): v for k, v in features.items()}

src/lerobot/processor/tokenizer_processor.py

@@ -1,24 +1,5 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 """
-This script defines a processor for tokenizing natural language instructions from an environment transition.
-
-It uses a tokenizer from the Hugging Face `transformers` library to convert task descriptions (text) into
-token IDs and attention masks, which are then added to the observation dictionary.
+Tokenizer processor for handling text tokenization in robot transitions.
 """
 
 from __future__ import annotations
@@ -28,14 +9,11 @@ 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.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import OBS_LANGUAGE
+from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
 from lerobot.utils.import_utils import _transformers_available
 
-from .core import EnvTransition, TransitionKey
-from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
-
-# Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
     from transformers import AutoTokenizer
 else:
@@ -44,62 +22,68 @@ else:
 
 @dataclass
 @ProcessorStepRegistry.register(name="tokenizer_processor")
-class TokenizerProcessorStep(ObservationProcessorStep):
-    """
-    Processor step to tokenize a natural language task description.
+class TokenizerProcessor:
+    """Tokenizes text tasks in complementary data using a huggingface tokenizer.
 
-    This step extracts a task string from the `complementary_data` of an `EnvTransition`,
-    tokenizes it using a Hugging Face `transformers` tokenizer, and adds the resulting
-    token IDs and attention mask to the `observation` dictionary.
+    This processor handles tokenization of task strings found in the complementary_data
+    using a specified pretrained tokenizer from Hugging Face. It adds tokenized versions
+    to the observation data for model processing while preserving the original task string.
 
-    Requires the `transformers` library to be installed.
+    The processor supports both single strings and lists of strings as task inputs.
 
-    Attributes:
-        tokenizer_name: The name of a pretrained tokenizer from the Hugging Face Hub (e.g., "bert-base-uncased").
-        tokenizer: A pre-initialized tokenizer object. If provided, `tokenizer_name` is ignored.
-        max_length: The maximum length to pad or truncate sequences to.
-        task_key: The key in `complementary_data` where the task string is stored.
-        padding_side: The side to pad on ('left' or 'right').
-        padding: The padding strategy ('max_length', 'longest', etc.).
-        truncation: Whether to truncate sequences longer than `max_length`.
-        input_tokenizer: The internal tokenizer instance, loaded during initialization.
+    Args:
+        tokenizer_name: Name of the pretrained tokenizer to load from Hugging Face Hub
+            (e.g., "bert-base-uncased", "microsoft/DialoGPT-medium"). This will be used
+            with AutoTokenizer.from_pretrained(). If tokenizer is provided, this is ignored.
+        tokenizer: A tokenizer object (e.g., from transformers library) that implements
+            the __call__ method. If provided, tokenizer_name is ignored. This parameter
+            is not serialized and must be provided via overrides when loading.
+        max_length: Maximum sequence length for tokenization. Defaults to 512.
+        task_key: Key in complementary_data containing the task text. Defaults to "task".
+        padding: Padding strategy for tokenization. Defaults to "max_length".
+        truncation: Whether to truncate sequences longer than max_length. Defaults to True.
+
+    Examples:
+        Using tokenizer name (auto-loaded):
+        ```python
+        processor = TokenizerProcessor(tokenizer_name="bert-base-uncased", max_length=128)
+        ```
+
+        Using custom tokenizer object:
+        ```python
+        from transformers import AutoTokenizer
+
+        custom_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        processor = TokenizerProcessor(tokenizer=custom_tokenizer, max_length=128)
+        ```
     """
 
     tokenizer_name: str | None = None
-    tokenizer: Any | None = None  # Use `Any` for compatibility without a hard dependency
+    tokenizer: Any | None = None  # Otherwise transformers is not available in the core dependencies
     max_length: int = 512
     task_key: str = "task"
     padding_side: str = "right"
     padding: str = "max_length"
     truncation: bool = True
 
-    # Internal tokenizer instance (not part of the config)
-    input_tokenizer: Any = field(default=None, init=False, repr=False)
+    # Internal tokenizer instance (not serialized)
+    _tokenizer: Any = field(default=None, init=False, repr=False)
 
     def __post_init__(self):
-        """
-        Initializes the tokenizer after the dataclass is created.
-
-        It checks for the availability of the `transformers` library and loads the tokenizer
-        either from a provided object or by name from the Hugging Face Hub.
-
-        Raises:
-            ImportError: If the `transformers` library is not installed.
-            ValueError: If neither `tokenizer` nor `tokenizer_name` is provided.
-        """
+        """Initialize the tokenizer from the provided tokenizer or tokenizer name."""
         if not _transformers_available:
             raise ImportError(
                 "The 'transformers' library is not installed. "
-                "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessorStep."
+                "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessor."
             )
 
         if self.tokenizer is not None:
             # Use provided tokenizer object directly
-            self.input_tokenizer = self.tokenizer
+            self._tokenizer = self.tokenizer
         elif self.tokenizer_name is not None:
             if AutoTokenizer is None:
                 raise ImportError("AutoTokenizer is not available")
-            self.input_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
+            self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
         else:
             raise ValueError(
                 "Either 'tokenizer' or 'tokenizer_name' must be provided. "
@@ -107,14 +91,13 @@ class TokenizerProcessorStep(ObservationProcessorStep):
             )
 
     def get_task(self, transition: EnvTransition) -> list[str] | None:
-        """
-        Extracts the task description(s) from the transition's complementary data.
+        """Extract and normalize task from complementary data.
 
         Args:
-            transition: The environment transition.
+            transition: Input transition containing complementary_data.
 
         Returns:
-            A list of task strings, or None if the task key is not found or the value is None.
+            List of task strings if task is present, None otherwise.
         """
         complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
         if complementary_data is None:
@@ -127,7 +110,7 @@ class TokenizerProcessorStep(ObservationProcessorStep):
         if task is None:
             return None
 
-        # Standardize to a list of strings for the tokenizer
+        # Convert to list of strings
         if isinstance(task, str):
             return [task]
         elif isinstance(task, list) and all(isinstance(t, str) for t in task):
@@ -135,82 +118,54 @@ class TokenizerProcessorStep(ObservationProcessorStep):
 
         return None
 
-    def observation(self, observation: dict[str, Any]) -> dict[str, Any]:
-        """
-        Tokenizes the task description and adds it to the observation dictionary.
-
-        This method retrieves the task, tokenizes it, moves the resulting tensors to the
-        same device as other data in the transition, and updates the observation.
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """Process the transition by tokenizing the task text.
 
         Args:
-            observation: The original observation dictionary.
+            transition: Input transition containing complementary_data with task text.
 
         Returns:
-            The updated observation dictionary including token IDs and an attention mask.
-        """
-        task = self.get_task(self.transition)
-        if task is None:
-            return observation
+            Modified transition with tokenized task added to observation.
 
-        # Tokenize the task (this will create CPU tensors)
+        Raises:
+            ValueError: If tokenizer initialization failed.
+        """
+        task = self.get_task(transition)
+        if task is None:
+            return transition
+
+        # Tokenize the task
         tokenized_prompt = self._tokenize_text(task)
 
-        # Detect the device from existing tensors in the transition to ensure consistency
-        target_device = self._detect_device(self.transition)
-
-        # Move new tokenized tensors to the detected device
-        if target_device is not None:
-            tokenized_prompt = {
-                k: v.to(target_device) if isinstance(v, torch.Tensor) else v
-                for k, v in tokenized_prompt.items()
-            }
-
-        # Create a new observation dict to avoid modifying the original in place
-        new_observation = dict(observation)
-
-        # Add tokenized data to the observation
-        new_observation[OBS_LANGUAGE_TOKENS] = tokenized_prompt["input_ids"]
-        new_observation[OBS_LANGUAGE_ATTENTION_MASK] = tokenized_prompt["attention_mask"].to(dtype=torch.bool)
-
-        return new_observation
-
-    def _detect_device(self, transition: EnvTransition) -> torch.device | None:
-        """
-        Detects the torch.device from existing tensors in the transition.
-
-        It checks tensors in the observation dictionary first, then the action tensor.
-
-        Args:
-            transition: The environment transition.
-
-        Returns:
-            The detected `torch.device`, or None if no tensors are found.
-        """
-        # Check observation tensors first (most likely place to find tensors)
+        # Get or create observation dict
         observation = transition.get(TransitionKey.OBSERVATION)
-        if observation:
-            for value in observation.values():
-                if isinstance(value, torch.Tensor):
-                    return value.device
+        if observation is None:
+            observation = {}
+        else:
+            observation = dict(observation)  # Make a copy
 
-        # Fallback to checking the action tensor
-        action = transition.get(TransitionKey.ACTION)
-        if isinstance(action, torch.Tensor):
-            return action.device
+        # Add tokenized data to observation
+        input_ids = tokenized_prompt["input_ids"]
+        attention_mask = tokenized_prompt.get("attention_mask")
+        if attention_mask is None:
+            # Some tokenizers (e.g., SigLIP text) may not return attention_mask; default to ones
+            attention_mask = torch.ones_like(input_ids)
+        observation[f"{OBS_LANGUAGE}.tokens"] = input_ids
+        observation[f"{OBS_LANGUAGE}.attention_mask"] = attention_mask.to(dtype=torch.bool)
 
-        return None  # No tensors found, default will be CPU
+        transition[TransitionKey.OBSERVATION.value] = observation  # type: ignore[misc]
+        return transition
 
     def _tokenize_text(self, text: str | list[str]) -> dict[str, torch.Tensor]:
-        """
-        A wrapper around the tokenizer call.
+        """Tokenize text using the configured tokenizer.
 
         Args:
-            text: A string or list of strings to tokenize.
+            text: Text string or list of strings to tokenize.
 
         Returns:
-            A dictionary containing tokenized 'input_ids' and 'attention_mask' as PyTorch tensors.
+            Dictionary containing tokenized output with keys like 'input_ids', 'attention_mask'.
         """
-        return self.input_tokenizer(
+        return self._tokenizer(
             text,
             max_length=self.max_length,
             truncation=self.truncation,
@@ -220,14 +175,10 @@ class TokenizerProcessorStep(ObservationProcessorStep):
         )
 
     def get_config(self) -> dict[str, Any]:
-        """
-        Returns the serializable configuration of the processor.
+        """Return configuration for serialization.
 
-        Note: The tokenizer object itself is not serialized. If the processor was initialized
-        with a tokenizer name, that name will be included in the config.
-
-        Returns:
-            A dictionary with the processor's configuration parameters.
+        Note: Only tokenizer_name is saved, not the tokenizer object itself.
+        When loading, provide the tokenizer via overrides if needed.
         """
         config = {
             "max_length": self.max_length,
@@ -237,37 +188,42 @@ class TokenizerProcessorStep(ObservationProcessorStep):
             "truncation": self.truncation,
         }
 
-        # Only save tokenizer_name if it was used to create the tokenizer
-        if self.tokenizer_name is not None and self.tokenizer is None:
+        # Only include tokenizer_name if it was used (not when tokenizer object was provided)
+        if self.tokenizer_name is not None:
             config["tokenizer_name"] = self.tokenizer_name
 
         return config
 
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Adds feature definitions for the language tokens and attention mask.
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return state dictionary (empty for this processor)."""
+        return {}
 
-        This updates the policy features dictionary to include the new data added to the
-        observation, ensuring downstream components are aware of their shape and type.
+    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Load state dictionary (no-op for this processor)."""
+        pass
+
+    def reset(self) -> None:
+        """Reset processor state (no-op for this processor)."""
+        pass
+
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        """Add tokenized task features to the feature contract.
 
         Args:
-            features: The dictionary of existing policy features.
+            features: Input feature dictionary.
 
         Returns:
-            The updated dictionary of policy features.
+            Updated feature dictionary with tokenized task features added.
         """
-        # Add a feature for the token IDs if it doesn't already exist
-        if OBS_LANGUAGE_TOKENS not in features[PipelineFeatureType.OBSERVATION]:
-            features[PipelineFeatureType.OBSERVATION][OBS_LANGUAGE_TOKENS] = PolicyFeature(
-                type=FeatureType.LANGUAGE, shape=(self.max_length,)
-            )
+        # Add features for tokenized output if they don't exist
+        # Standard tokenizer output includes tokens and attention_mask
+        tokens_key = f"{OBS_LANGUAGE}.tokens"
+        attention_mask_key = f"{OBS_LANGUAGE}.attention_mask"
 
-        # Add a feature for the attention mask if it doesn't already exist
-        if OBS_LANGUAGE_ATTENTION_MASK not in features[PipelineFeatureType.OBSERVATION]:
-            features[PipelineFeatureType.OBSERVATION][OBS_LANGUAGE_ATTENTION_MASK] = PolicyFeature(
-                type=FeatureType.LANGUAGE, shape=(self.max_length,)
-            )
+        if tokens_key not in features:
+            features[tokens_key] = PolicyFeature(type=FeatureType.LANGUAGE, shape=(self.max_length,))
+
+        if attention_mask_key not in features:
+            features[attention_mask_key] = PolicyFeature(type=FeatureType.LANGUAGE, shape=(self.max_length,))
 
         return features

src/lerobot/record.py

@@ -18,7 +18,7 @@ Records a dataset. Actions for the robot can be either generated by teleoperatio
 Example:
 
 ```shell
-lerobot-record \
+python -m lerobot.record \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \
@@ -36,7 +36,7 @@ lerobot-record \
 
 Example recording with bimanual so100:
 ```shell
-lerobot-record \
+python -m lerobot.record \
   --robot.type=bi_so100_follower \
   --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \
   --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \
@@ -62,7 +62,6 @@ import time
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from pprint import pformat
-from typing import Any
 
 from lerobot.cameras import (  # noqa: F401
     CameraConfig,  # noqa: F401
@@ -75,23 +74,17 @@ from lerobot.datasets.image_writer import safe_stop_image_writer
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.datasets.video_utils import VideoEncodingManager
-from lerobot.policies.factory import make_policy, make_pre_post_processors
+from lerobot.policies.factory import make_policy, make_processor
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.processor import (
-    EnvTransition,
-    IdentityProcessorStep,
-    PolicyProcessorPipeline,
-    RobotProcessorPipeline,
-    TransitionKey,
-)
+from lerobot.processor import RobotProcessor
 from lerobot.processor.converters import (
-    action_to_transition,
-    identity_transition,
-    observation_to_transition,
-    transition_to_action,
-    transition_to_dataset_frame,
+    to_dataset_frame,
+    to_output_robot_action,
+    to_transition_robot_observation,
+    to_transition_teleop_action,
 )
-from lerobot.processor.rename_processor import rename_stats
+from lerobot.processor.normalize_processor import rename_stats
+from lerobot.processor.pipeline import IdentityProcessor, TransitionKey
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
@@ -243,36 +236,23 @@ def record_loop(
     dataset: LeRobotDataset | None = None,
     teleop: Teleoperator | list[Teleoperator] | None = None,
     policy: PreTrainedPolicy | None = None,
-    preprocessor: PolicyProcessorPipeline | None = None,
-    postprocessor: PolicyProcessorPipeline | None = None,
+    preprocessor: RobotProcessor | None = None,
+    postprocessor: RobotProcessor | None = None,
     control_time_s: int | None = None,
-    teleop_action_processor: RobotProcessorPipeline[EnvTransition] | None = None,  # runs after teleop
-    robot_action_processor: RobotProcessorPipeline[dict[str, Any]] | None = None,  # runs before robot
-    robot_observation_processor: RobotProcessorPipeline[EnvTransition] | None = None,  # runs after robot
+    teleop_action_processor: RobotProcessor | None = None,  # runs after teleop
+    robot_action_processor: RobotProcessor | None = None,  # runs before robot
+    robot_observation_processor: RobotProcessor | None = None,  # runs after robot
     single_task: str | None = None,
     display_data: bool = False,
 ):
-    teleop_action_processor: RobotProcessorPipeline[EnvTransition] = (
-        teleop_action_processor
-        or RobotProcessorPipeline(
-            steps=[IdentityProcessorStep()], to_transition=action_to_transition, to_output=identity_transition
-        )
+    teleop_action_processor = teleop_action_processor or RobotProcessor(
+        steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr
     )
-    robot_action_processor: RobotProcessorPipeline[dict[str, Any]] = (
-        robot_action_processor
-        or RobotProcessorPipeline(
-            steps=[IdentityProcessorStep()],
-            to_transition=identity_transition,
-            to_output=transition_to_action,
-        )
+    robot_action_processor = robot_action_processor or RobotProcessor(
+        steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action
     )
-    robot_observation_processor: RobotProcessorPipeline[EnvTransition] = (
-        robot_observation_processor
-        or RobotProcessorPipeline(
-            steps=[IdentityProcessorStep()],
-            to_transition=observation_to_transition,
-            to_output=identity_transition,
-        )
+    robot_observation_processor = robot_observation_processor or RobotProcessor(
+        steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr
     )
 
     if dataset is not None and dataset.fps != fps:
@@ -285,14 +265,7 @@ def record_loop(
             (
                 t
                 for t in teleop
-                if isinstance(
-                    t,
-                    (
-                        so100_leader.SO100Leader,
-                        so101_leader.SO101Leader,
-                        koch_leader.KochLeader,
-                    ),
-                )
+                if isinstance(t, (so100_leader.SO100Leader, so101_leader.SO101Leader, koch_leader.KochLeader))
             ),
             None,
         )
@@ -335,7 +308,7 @@ def record_loop(
         # Get action from either policy or teleop
         if policy is not None and preprocessor is not None and postprocessor is not None:
             if dataset is not None:
-                observation_frame = transition_to_dataset_frame(
+                observation_frame = to_dataset_frame(
                     obs_transition, dataset.features
                 )  # Convert the observation to the dataset format
 
@@ -361,7 +334,6 @@ def record_loop(
             act = teleop.get_action()
 
             # Applies a pipeline to the raw teleop action, default is IdentityProcessor
-            # TODO(Steven): This assumes that the processor passed by the user should have identity_transition as to_output.
             teleop_transition = teleop_action_processor(act)
 
         elif isinstance(teleop, list):
@@ -374,14 +346,13 @@ def record_loop(
         else:
             logging.info(
                 "No policy or teleoperator provided, skipping action generation. "
-                "This is likely to happen when resetting the environment without a teleop device."
-                "The robot won't be at its rest position at the start of the next episode."
+                "This is likely to happen during environment reset."
             )
-            continue
+            # Still continue to next loop to respect timing
 
         # Applies a pipeline to the action, default is IdentityProcessor
         # IMPORTANT: action_pipeline.to_output must return a dict suitable for robot.send_action()
-        if policy is not None and policy_transition is not None:
+        if policy_transition is not None:
             robot_action_to_send = robot_action_processor(policy_transition)
         else:
             robot_action_to_send = robot_action_processor(teleop_transition)
@@ -394,7 +365,7 @@ def record_loop(
 
         # Write to dataset
         if dataset is not None:
-            # If transition_to_dataset_frame is provided, use it to merge the transitions.
+            # If to_dataset_frame is provided, use it to merge the transitions.
             merged = []
             if obs_transition is not None:  # The observation from the robot
                 merged.append(obs_transition)
@@ -402,15 +373,13 @@ def record_loop(
                 merged.append(teleop_transition)
             if policy_transition is not None:  # The action from policy
                 merged.append(policy_transition)
-            frame = transition_to_dataset_frame(
+            frame = to_dataset_frame(
                 merged if len(merged) > 1 else merged[0], dataset.features
             )  # Convert the observation to the dataset format
             dataset.add_frame(frame, task=single_task)
 
         if display_data:
-            log_rerun_data(
-                observation=obs_transition.get(TransitionKey.OBSERVATION), action=robot_action_to_send
-            )
+            log_rerun_data([obs_transition, teleop_transition or policy_transition])
 
         dt_s = time.perf_counter() - start_loop_t
         busy_wait(1 / fps - dt_s)
@@ -430,15 +399,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
 
     action_features = hw_to_dataset_features(robot.action_features, "action", cfg.dataset.video)
     obs_features = hw_to_dataset_features(robot.observation_features, "observation", cfg.dataset.video)
-
-    # Add next.* features that are generated during recording
-    transition_features = {
-        "next.reward": {"dtype": "float32", "shape": (1,), "names": None},
-        "next.done": {"dtype": "bool", "shape": (1,), "names": None},
-        "next.truncated": {"dtype": "bool", "shape": (1,), "names": None},
-    }
-
-    dataset_features = {**action_features, **obs_features, **transition_features}
+    dataset_features = {**action_features, **obs_features}
 
     if cfg.resume:
         dataset = LeRobotDataset(
@@ -473,7 +434,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
     preprocessor = None
     postprocessor = None
     if cfg.policy is not None:
-        preprocessor, postprocessor = make_pre_post_processors(
+        preprocessor, postprocessor = make_processor(
             policy_cfg=cfg.policy,
             pretrained_path=cfg.policy.pretrained_path,
             dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map),
@@ -549,9 +510,5 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
     return dataset
 
 
-def main():
-    record()
-
-
 if __name__ == "__main__":
-    main()
+    record()

src/lerobot/replay.py

@@ -18,7 +18,7 @@ Replays the actions of an episode from a dataset on a robot.
 Examples:
 
 ```shell
-lerobot-replay \
+python -m lerobot.replay \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=black \
@@ -28,7 +28,7 @@ lerobot-replay \
 
 Example replay with bimanual so100:
 ```shell
-lerobot-replay \
+python -m lerobot.replay \
   --robot.type=bi_so100_follower \
   --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \
   --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \
@@ -45,10 +45,9 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 from pprint import pformat
 
-from lerobot.configs import parser
+import draccus
+
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.processor import IdentityProcessorStep, RobotProcessorPipeline
-from lerobot.processor.converters import action_to_transition, transition_to_action
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
@@ -56,7 +55,6 @@ from lerobot.robots import (  # noqa: F401
     hope_jr,
     koch_follower,
     make_robot_from_config,
-    reachy2,
     so100_follower,
     so101_follower,
 )
@@ -85,25 +83,13 @@ class ReplayConfig:
     dataset: DatasetReplayConfig
     # Use vocal synthesis to read events.
     play_sounds: bool = True
-    # Optional processor for actions before sending to robot
-    robot_action_processor: RobotProcessorPipeline | None = None
 
 
-@parser.wrap()
+@draccus.wrap()
 def replay(cfg: ReplayConfig):
     init_logging()
     logging.info(pformat(asdict(cfg)))
 
-    # Initialize robot action processor with default if not provided
-    robot_action_processor = cfg.robot_action_processor or RobotProcessorPipeline(
-        steps=[IdentityProcessorStep()],
-        to_transition=action_to_transition,
-        to_output=transition_to_action,  # type: ignore[arg-type]
-    )
-
-    # Reset processor
-    robot_action_processor.reset()
-
     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")
@@ -118,10 +104,7 @@ def replay(cfg: ReplayConfig):
         for i, name in enumerate(dataset.features["action"]["names"]):
             action[name] = action_array[i]
 
-        # Process action through robot action processor
-        processed_action = robot_action_processor(action)
-
-        robot.send_action(processed_action)  # type: ignore[arg-type]
+        robot.send_action(action)
 
         dt_s = time.perf_counter() - start_episode_t
         busy_wait(1 / dataset.fps - dt_s)

src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py

@@ -29,10 +29,10 @@ class BiSO100FollowerConfig(RobotConfig):
 
     # Optional
     left_arm_disable_torque_on_disconnect: bool = True
-    left_arm_max_relative_target: float | dict[str, float] | None = None
+    left_arm_max_relative_target: int | None = None
     left_arm_use_degrees: bool = False
     right_arm_disable_torque_on_disconnect: bool = True
-    right_arm_max_relative_target: float | dict[str, float] | None = None
+    right_arm_max_relative_target: int | None = None
     right_arm_use_degrees: bool = False
 
     # cameras (shared between both arms)

src/lerobot/robots/hope_jr/config_hope_jr.py

@@ -44,8 +44,8 @@ class HopeJrArmConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
-    # names to the max_relative_target value for that motor.
-    max_relative_target: float | dict[str, float] | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
 
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/koch_follower/config_koch_follower.py

@@ -28,9 +28,9 @@ class KochFollowerConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
-    # names to the max_relative_target value for that motor.
-    max_relative_target: float | dict[str, float] | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
 
     # cameras
     cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/koch_follower/koch_follower.py

@@ -110,7 +110,6 @@ class KochFollower(Robot):
         return self.bus.is_calibrated
 
     def calibrate(self) -> None:
-        self.bus.disable_torque()
         if self.calibration:
             # Calibration file exists, ask user whether to use it or run new calibration
             user_input = input(
@@ -121,6 +120,7 @@ class KochFollower(Robot):
                 self.bus.write_calibration(self.calibration)
                 return
         logger.info(f"\nRunning calibration of {self}")
+        self.bus.disable_torque()
         for motor in self.bus.motors:
             self.bus.write("Operating_Mode", motor, OperatingMode.EXTENDED_POSITION.value)
 

src/lerobot/robots/lekiwi/config_lekiwi.py

@@ -39,9 +39,9 @@ class LeKiwiConfig(RobotConfig):
     disable_torque_on_disconnect: bool = True
 
     # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors, or a dictionary that maps motor
-    # names to the max_relative_target value for that motor.
-    max_relative_target: float | dict[str, float] | None = None
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
 
     cameras: dict[str, CameraConfig] = field(default_factory=lekiwi_cameras_config)
 

src/lerobot/robots/reachy2/__init__.py

@@ -1,25 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .configuration_reachy2 import Reachy2RobotConfig
-from .robot_reachy2 import (
-    REACHY2_ANTENNAS_JOINTS,
-    REACHY2_L_ARM_JOINTS,
-    REACHY2_NECK_JOINTS,
-    REACHY2_R_ARM_JOINTS,
-    REACHY2_VEL,
-    Reachy2Robot,
-)

src/lerobot/robots/reachy2/configuration_reachy2.py

@@ -1,107 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass, field
-
-from lerobot.cameras import CameraConfig
-from lerobot.cameras.configs import ColorMode
-from lerobot.cameras.reachy2_camera import Reachy2CameraConfig
-
-from ..config import RobotConfig
-
-
-@RobotConfig.register_subclass("reachy2")
-@dataclass
-class Reachy2RobotConfig(RobotConfig):
-    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
-    # Set this to a positive scalar to have the same value for all motors.
-    max_relative_target: float | None = None
-
-    # IP address of the Reachy 2 robot
-    ip_address: str | None = "localhost"
-
-    # If True, turn_off_smoothly() will be sent to the robot before disconnecting.
-    disable_torque_on_disconnect: bool = False
-
-    # Tag for external commands control
-    # Set to True if you use an external commands system to control the robot,
-    # such as the official teleoperation application: https://github.com/pollen-robotics/Reachy2Teleoperation
-    # If True, robot.send_action() will not send commands to the robot.
-    use_external_commands: bool = False
-
-    # Robot parts
-    # Set to False to not add the corresponding joints part to the robot list of joints.
-    # By default, all parts are set to True.
-    with_mobile_base: bool = True
-    with_l_arm: bool = True
-    with_r_arm: bool = True
-    with_neck: bool = True
-    with_antennas: bool = True
-
-    # Robot cameras
-    # Set to True if you want to use the corresponding cameras in the observations.
-    # By default, only the teleop cameras are used.
-    with_left_teleop_camera: bool = True
-    with_right_teleop_camera: bool = True
-    with_torso_camera: bool = False
-
-    cameras: dict[str, CameraConfig] = field(default_factory=dict)
-
-    def __post_init__(self) -> None:
-        # Add cameras with same ip_address as the robot
-        if self.with_left_teleop_camera:
-            self.cameras["teleop_left"] = Reachy2CameraConfig(
-                name="teleop",
-                image_type="left",
-                ip_address=self.ip_address,
-                fps=15,
-                width=640,
-                height=480,
-                color_mode=ColorMode.RGB,
-            )
-        if self.with_right_teleop_camera:
-            self.cameras["teleop_right"] = Reachy2CameraConfig(
-                name="teleop",
-                image_type="right",
-                ip_address=self.ip_address,
-                fps=15,
-                width=640,
-                height=480,
-                color_mode=ColorMode.RGB,
-            )
-        if self.with_torso_camera:
-            self.cameras["torso_rgb"] = Reachy2CameraConfig(
-                name="depth",
-                image_type="rgb",
-                ip_address=self.ip_address,
-                fps=15,
-                width=640,
-                height=480,
-                color_mode=ColorMode.RGB,
-            )
-
-        super().__post_init__()
-
-        if not (
-            self.with_mobile_base
-            or self.with_l_arm
-            or self.with_r_arm
-            or self.with_neck
-            or self.with_antennas
-        ):
-            raise ValueError(
-                "No Reachy2Robot part used.\n"
-                "At least one part of the robot must be set to True "
-                "(with_mobile_base, with_l_arm, with_r_arm, with_neck, with_antennas)"
-            )

src/lerobot/robots/reachy2/robot_reachy2.py

@@ -1,230 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import time
-from typing import Any
-
-import numpy as np
-from reachy2_sdk import ReachySDK
-
-from lerobot.cameras.utils import make_cameras_from_configs
-
-from ..robot import Robot
-from ..utils import ensure_safe_goal_position
-from .configuration_reachy2 import Reachy2RobotConfig
-
-# {lerobot_keys: reachy2_sdk_keys}
-REACHY2_NECK_JOINTS = {
-    "neck_yaw.pos": "head.neck.yaw",
-    "neck_pitch.pos": "head.neck.pitch",
-    "neck_roll.pos": "head.neck.roll",
-}
-
-REACHY2_ANTENNAS_JOINTS = {
-    "l_antenna.pos": "head.l_antenna",
-    "r_antenna.pos": "head.r_antenna",
-}
-
-REACHY2_R_ARM_JOINTS = {
-    "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
-    "r_shoulder_roll.pos": "r_arm.shoulder.roll",
-    "r_elbow_yaw.pos": "r_arm.elbow.yaw",
-    "r_elbow_pitch.pos": "r_arm.elbow.pitch",
-    "r_wrist_roll.pos": "r_arm.wrist.roll",
-    "r_wrist_pitch.pos": "r_arm.wrist.pitch",
-    "r_wrist_yaw.pos": "r_arm.wrist.yaw",
-    "r_gripper.pos": "r_arm.gripper",
-}
-
-REACHY2_L_ARM_JOINTS = {
-    "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
-    "l_shoulder_roll.pos": "l_arm.shoulder.roll",
-    "l_elbow_yaw.pos": "l_arm.elbow.yaw",
-    "l_elbow_pitch.pos": "l_arm.elbow.pitch",
-    "l_wrist_roll.pos": "l_arm.wrist.roll",
-    "l_wrist_pitch.pos": "l_arm.wrist.pitch",
-    "l_wrist_yaw.pos": "l_arm.wrist.yaw",
-    "l_gripper.pos": "l_arm.gripper",
-}
-
-REACHY2_VEL = {
-    "mobile_base.vx": "vx",
-    "mobile_base.vy": "vy",
-    "mobile_base.vtheta": "vtheta",
-}
-
-
-class Reachy2Robot(Robot):
-    """
-    [Reachy 2](https://www.pollen-robotics.com/reachy/), by Pollen Robotics.
-    """
-
-    config_class = Reachy2RobotConfig
-    name = "reachy2"
-
-    def __init__(self, config: Reachy2RobotConfig):
-        super().__init__(config)
-
-        self.config = config
-        self.robot_type = self.config.type
-        self.use_external_commands = self.config.use_external_commands
-
-        self.reachy: None | ReachySDK = None
-        self.cameras = make_cameras_from_configs(config.cameras)
-
-        self.logs: dict[str, float] = {}
-
-        self.joints_dict: dict[str, str] = self._generate_joints_dict()
-
-    @property
-    def observation_features(self) -> dict[str, Any]:
-        return {**self.motors_features, **self.camera_features}
-
-    @property
-    def action_features(self) -> dict[str, type]:
-        return self.motors_features
-
-    @property
-    def camera_features(self) -> dict[str, tuple[int | None, int | None, int]]:
-        return {cam: (self.cameras[cam].height, self.cameras[cam].width, 3) for cam in self.cameras}
-
-    @property
-    def motors_features(self) -> dict[str, type]:
-        if self.config.with_mobile_base:
-            return {
-                **dict.fromkeys(
-                    self.joints_dict.keys(),
-                    float,
-                ),
-                **dict.fromkeys(
-                    REACHY2_VEL.keys(),
-                    float,
-                ),
-            }
-        else:
-            return dict.fromkeys(self.joints_dict.keys(), float)
-
-    @property
-    def is_connected(self) -> bool:
-        return self.reachy.is_connected() if self.reachy is not None else False
-
-    def connect(self, calibrate: bool = False) -> None:
-        self.reachy = ReachySDK(self.config.ip_address)
-        if not self.is_connected:
-            raise ConnectionError()
-
-        for cam in self.cameras.values():
-            cam.connect()
-
-        self.configure()
-
-    def configure(self) -> None:
-        if self.reachy is not None:
-            self.reachy.turn_on()
-            self.reachy.reset_default_limits()
-
-    @property
-    def is_calibrated(self) -> bool:
-        return True
-
-    def calibrate(self) -> None:
-        pass
-
-    def _generate_joints_dict(self) -> dict[str, str]:
-        joints = {}
-        if self.config.with_neck:
-            joints.update(REACHY2_NECK_JOINTS)
-        if self.config.with_l_arm:
-            joints.update(REACHY2_L_ARM_JOINTS)
-        if self.config.with_r_arm:
-            joints.update(REACHY2_R_ARM_JOINTS)
-        if self.config.with_antennas:
-            joints.update(REACHY2_ANTENNAS_JOINTS)
-        return joints
-
-    def _get_state(self) -> dict[str, float]:
-        if self.reachy is not None:
-            pos_dict = {k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()}
-            if not self.config.with_mobile_base:
-                return pos_dict
-            vel_dict = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()}
-            return {**pos_dict, **vel_dict}
-        else:
-            return {}
-
-    def get_observation(self) -> dict[str, np.ndarray]:
-        obs_dict: dict[str, Any] = {}
-
-        # Read Reachy 2 state
-        before_read_t = time.perf_counter()
-        obs_dict.update(self._get_state())
-        self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t
-
-        # Capture images from cameras
-        for cam_key, cam in self.cameras.items():
-            obs_dict[cam_key] = cam.async_read()
-
-        return obs_dict
-
-    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
-        if self.reachy is not None:
-            if not self.is_connected:
-                raise ConnectionError()
-
-            before_write_t = time.perf_counter()
-
-            vel = {}
-            goal_pos = {}
-            for key, val in action.items():
-                if key not in self.joints_dict:
-                    if key not in REACHY2_VEL:
-                        raise KeyError(f"Key '{key}' is not a valid motor key in Reachy 2.")
-                    else:
-                        vel[REACHY2_VEL[key]] = float(val)
-                else:
-                    if not self.use_external_commands and self.config.max_relative_target is not None:
-                        goal_pos[key] = float(val)
-                        goal_present_pos = {
-                            key: (
-                                goal_pos[key],
-                                self.reachy.joints[self.joints_dict[key]].present_position,
-                            )
-                        }
-                        safe_goal_pos = ensure_safe_goal_position(
-                            goal_present_pos, float(self.config.max_relative_target)
-                        )
-                        val = safe_goal_pos[key]
-                    self.reachy.joints[self.joints_dict[key]].goal_position = float(val)
-
-            if self.config.with_mobile_base:
-                self.reachy.mobile_base.set_goal_speed(vel["vx"], vel["vy"], vel["vtheta"])
-
-            # We don't send the goal positions if we control Reachy 2 externally
-            if not self.use_external_commands:
-                self.reachy.send_goal_positions()
-                if self.config.with_mobile_base:
-                    self.reachy.mobile_base.send_speed_command()
-
-            self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t
-        return action
-
-    def disconnect(self) -> None:
-        if self.reachy is not None:
-            for cam in self.cameras.values():
-                cam.disconnect()
-            if self.config.disable_torque_on_disconnect:
-                self.reachy.turn_off_smoothly()
-            self.reachy.disconnect()