feat(dataset-tools): add dataset utilities and example script

- Introduced dataset tools for LeRobotDataset, including functions for deleting episodes, splitting datasets, adding/removing features, and merging datasets.
- Added an example script demonstrating the usage of these utilities.
- Implemented comprehensive tests for all new functionalities to ensure reliability and correctness.

.gitignore

@@ -173,7 +173,3 @@ outputs/
 
 # Dev folders
 .cache/*
-*.stl
-*.urdf
-*.xml
-*.part

docker/Dockerfile.internal

@@ -39,7 +39,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     software-properties-common build-essential git curl \
     libglib2.0-0 libgl1-mesa-glx libegl1-mesa ffmpeg \
     libusb-1.0-0-dev speech-dispatcher libgeos-dev portaudio19-dev \
-    cmake pkg-config ninja-build \
     && add-apt-repository -y ppa:deadsnakes/ppa \
     && apt-get update \
     && apt-get install -y --no-install-recommends \

docker/Dockerfile.user

@@ -29,9 +29,8 @@ 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 \
-    cmake pkg-config ninja-build \
     && curl -LsSf https://astral.sh/uv/install.sh | sh \
     && mv /root/.local/bin/uv /usr/local/bin/uv \
     && useradd --create-home --shell /bin/bash user_lerobot \

docs/source/_toctree.yml

@@ -19,31 +19,16 @@
     title: Train RL in Simulation
   - local: async
     title: Use Async Inference
-  title: "Tutorials"
-- sections:
-  - local: lerobot-dataset-v3
-    title: Using LeRobotDataset
-  - local: libero
-    title: Using Libero
   - local: porting_datasets_v3
     title: Porting Large Datasets
-  title: "Datasets"
+  title: "Tutorials"
 - sections:
   - local: smolvla
     title: Finetune SmolVLA
   title: "Policies"
-
-- sections:
-  - local: introduction_processors
-    title: Introduction to Robot Processors
-  - local: debug_processor_pipeline
-    title: Debug your processor pipeline
-  - 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
@@ -52,20 +37,10 @@
     title: Koch v1.1
   - local: lekiwi
     title: LeKiwi
-  - local: hope_jr
-    title: Hope Jr
-  - local: reachy2
-    title: Reachy 2
   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/backwardcomp.mdx

@@ -1,61 +1,5 @@
 # Backward compatibility
 
-## Policy Normalization Migration (PR #1452)
-
-**Breaking Change**: LeRobot policies no longer have built-in normalization layers embedded in their weights. Normalization is now handled by external `PolicyProcessorPipeline` components.
-
-### What changed?
-
-|                            | Before PR #1452                                  | After PR #1452                                               |
-| -------------------------- | ------------------------------------------------ | ------------------------------------------------------------ |
-| **Normalization Location** | Embedded in model weights (`normalize_inputs.*`) | External `PolicyProcessorPipeline` components                |
-| **Model State Dict**       | Contains normalization statistics                | **Clean weights only** - no normalization parameters         |
-| **Usage**                  | `policy(batch)` handles everything               | `preprocessor(batch)` → `policy(...)` → `postprocessor(...)` |
-
-### Impact on existing models
-
-- Models trained **before** PR #1452 have normalization embedded in their weights
-- These models need migration to work with the new `PolicyProcessorPipeline` system
-- The migration extracts normalization statistics and creates separate processor pipelines
-
-### Migrating old models
-
-Use the migration script to convert models with embedded normalization:
-
-```shell
-python src/lerobot/processor/migrate_policy_normalization.py \
-    --pretrained-path lerobot/act_aloha_sim_transfer_cube_human \
-    --push-to-hub \
-    --branch migrated
-```
-
-The script:
-
-1. **Extracts** normalization statistics from model weights
-2. **Creates** external preprocessor and postprocessor pipelines
-3. **Removes** normalization layers from model weights
-4. **Saves** clean model + processor pipelines
-5. **Pushes** to Hub with automatic PR creation
-
-### Using migrated models
-
-```python
-# New usage pattern (after migration)
-from lerobot.policies.factory import make_policy, make_pre_post_processors
-
-# Load model and processors separately
-policy = make_policy(config, ds_meta=dataset.meta)
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=config,
-    dataset_stats=dataset.meta.stats
-)
-
-# Process data through pipeline
-processed_batch = preprocessor(raw_batch)
-action = policy.select_action(processed_batch)
-final_action = postprocessor(action)
-```
-
 ## Hardware API redesign
 
 PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.

docs/source/debug_processor_pipeline.mdx

@@ -1,299 +0,0 @@
-# Debug Your Processor Pipeline
-
-Processor pipelines can be complex, especially when chaining multiple transformation steps.
-Unlike simple function calls, pipelines lack natural observability, you can't easily see what happens
-between each step or where things go wrong.
-This guide provides debugging tools and techniques specifically designed to address these challenges
-and help you understand data flow through your pipelines.
-
-We'll explore three complementary debugging approaches: **hooks** for runtime monitoring, **step-through debugging** for detailed inspection, and **feature validation** for catching structural mismatches. Each serves a different purpose and together they provide complete visibility into your pipeline's behavior.
-
-## Understanding Hooks
-
-Hooks are functions that get called at specific points during pipeline execution.
-They provide a way to inspect, monitor, or modify data without changing your pipeline code.
-Think of them as "event listeners" for your pipeline.
-
-### What is a Hook?
-
-A hook is a callback function that gets automatically invoked at specific moments during pipeline execution.
-The concept comes from event-driven programming, imagine you could "hook into" the pipeline's execution flow to observe or react to what's happening.
-
-Think of hooks like inserting checkpoints into your pipeline. Every time the pipeline reaches one of these checkpoints, it pauses briefly to call your hook function, giving you a chance to inspect the current state, log information, and validate data.
-
-A hook is simply a function that accepts two parameters:
-
-- `step_idx: int` - The index of the current processing step (0, 1, 2, etc.)
-- `transition: EnvTransition` - The data transition at that point in the pipeline
-
-The beauty of hooks is their non-invasive nature: you can add monitoring, validation, or debugging logic without changing a single line of your pipeline code. The pipeline remains clean and focused on its core logic, while hooks handle the cross-cutting concerns like logging, monitoring, and debugging.
-
-### Before vs After Hooks
-
-The pipeline supports two types of hooks:
-
-- **Before hooks** (`register_before_step_hook`) - Called before each step executes
-- **After hooks** (`register_after_step_hook`) - Called after each step completes
-
-```python
-def before_hook(step_idx: int, transition: EnvTransition):
-    """Called before step processes the transition."""
-    print(f"About to execute step {step_idx}")
-    # Useful for: logging, validation, setup
-
-def after_hook(step_idx: int, transition: EnvTransition):
-    """Called after step has processed the transition."""
-    print(f"Completed step {step_idx}")
-    # Useful for: monitoring results, cleanup, debugging
-
-processor.register_before_step_hook(before_hook)
-processor.register_after_step_hook(after_hook)
-```
-
-### Implementing a NaN Detection Hook
-
-Here's a practical example of a hook that detects NaN values:
-
-```python
-def check_nans(step_idx: int, transition: EnvTransition):
-    """Check for NaN values in observations."""
-    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"NaN detected in {key} at step {step_idx}")
-
-# Register the hook to run after each step
-processor.register_after_step_hook(check_nans)
-
-# Process your data - the hook will be called automatically
-output = processor(input_data)
-
-# Remove the hook when done debugging
-processor.unregister_after_step_hook(check_nans)
-```
-
-### How Hooks Work Internally
-
-Understanding the internal mechanism helps you use hooks more effectively. The pipeline maintains two separate lists: one for before-step hooks and another for after-step hooks. When you register a hook, it's simply appended to the appropriate list.
-
-During execution, the pipeline follows a strict sequence: for each processing step, it first calls all before-hooks in registration order, then executes the actual step transformation, and finally calls all after-hooks in registration order. This creates a predictable, sandwich-like structure around each step.
-
-The key insight is that hooks don't change the core pipeline logic—they're purely additive. The pipeline's `_forward` method orchestrates this dance between hooks and processing steps, ensuring that your debugging or monitoring code runs at exactly the right moments without interfering with the main data flow.
-
-Here's a simplified view of how the pipeline executes hooks:
-
-```python
-class DataProcessorPipeline:
-    def __init__(self):
-        self.steps = [...]
-        self.before_step_hooks = []  # List of before hooks
-        self.after_step_hooks = []   # List of after hooks
-
-    def _forward(self, transition):
-        """Internal method that processes the transition through all steps."""
-        for step_idx, processor_step in enumerate(self.steps):
-            # 1. Call all BEFORE hooks
-            for hook in self.before_step_hooks:
-                hook(step_idx, transition)
-
-            # 2. Execute the actual processing step
-            transition = processor_step(transition)
-
-            # 3. Call all AFTER hooks
-            for hook in self.after_step_hooks:
-                hook(step_idx, transition)
-
-        return transition
-
-    def register_before_step_hook(self, hook_fn):
-        self.before_step_hooks.append(hook_fn)
-
-    def register_after_step_hook(self, hook_fn):
-        self.after_step_hooks.append(hook_fn)
-```
-
-### Execution Flow
-
-The execution flow looks like this:
-
-```
-Input → Before Hook → Step 0 → After Hook → Before Hook → Step 1 → After Hook → ... → Output
-```
-
-For example, with 3 steps and both hook types:
-
-```python
-def timing_before(step_idx, transition):
-    print(f"⏱️  Starting step {step_idx}")
-
-def validation_after(step_idx, transition):
-    print(f"✅ Completed step {step_idx}")
-
-processor.register_before_step_hook(timing_before)
-processor.register_after_step_hook(validation_after)
-
-# This will output:
-# ⏱️  Starting step 0
-# ✅ Completed step 0
-# ⏱️  Starting step 1
-# ✅ Completed step 1
-# ⏱️  Starting step 2
-# ✅ Completed step 2
-```
-
-### Multiple Hooks
-
-You can register multiple hooks of the same type - they execute in the order registered:
-
-```python
-def log_shapes(step_idx: int, transition: EnvTransition):
-    obs = transition.get(TransitionKey.OBSERVATION)
-    if obs:
-        print(f"Step {step_idx} observation shapes:")
-        for key, value in obs.items():
-            if isinstance(value, torch.Tensor):
-                print(f"  {key}: {value.shape}")
-
-processor.register_after_step_hook(check_nans)      # Executes first
-processor.register_after_step_hook(log_shapes)     # Executes second
-
-# Both hooks will be called after each step in registration order
-output = processor(input_data)
-```
-
-While hooks are excellent for monitoring specific issues (like NaN detection) or gathering metrics during normal pipeline execution, sometimes you need to dive deeper. When you want to understand exactly what happens at each step or debug complex transformation logic, step-through debugging provides the detailed inspection you need.
-
-## Step-Through Debugging
-
-Step-through debugging is like having a slow-motion replay for your pipeline. Instead of watching your data get transformed in one quick blur from input to output, you can pause and examine what happens after each individual step.
-
-This approach is particularly valuable when you're trying to understand a complex pipeline, debug unexpected behavior, or verify that each transformation is working as expected. Unlike hooks, which are great for automated monitoring, step-through debugging gives you manual, interactive control over the inspection process.
-
-The `step_through()` method is a generator that yields the transition state after each processing step, allowing you to inspect intermediate results. Think of it as creating a series of snapshots of your data as it flows through the pipeline—each snapshot shows you exactly what your data looks like after one more transformation has been applied.
-
-### How Step-Through Works
-
-The `step_through()` method fundamentally changes how the pipeline executes. Instead of running all steps in sequence and only returning the final result, it transforms the pipeline into an iterator that yields intermediate results.
-
-Here's what happens internally: the method starts by converting your input data into the pipeline's internal transition format, then yields this initial state. Next, it applies the first processing step and yields the result. Then it applies the second step to that result and yields again, and so on. Each `yield` gives you a complete snapshot of the transition at that point.
-
-This generator pattern is powerful because it's lazy—the pipeline only computes the next step when you ask for it. This means you can stop at any point, inspect the current state thoroughly, and decide whether to continue. You're not forced to run the entire pipeline just to debug one problematic step.
-
-Instead of running the entire pipeline and only seeing the final result, `step_through()` pauses after each step and gives you the intermediate transition:
-
-```python
-# This creates a generator that yields intermediate states
-for i, intermediate_result in enumerate(processor.step_through(input_data)):
-    print(f"=== After step {i} ===")
-
-    # Inspect the observation at this stage
-    obs = intermediate_result.get(TransitionKey.OBSERVATION)
-    if obs:
-        for key, value in obs.items():
-            if isinstance(value, torch.Tensor):
-                print(f"{key}: shape={value.shape}, dtype={value.dtype}")
-```
-
-### Interactive Debugging with Breakpoints
-
-You can add breakpoints in the step-through loop to interactively debug:
-
-```python
-# Step through the pipeline with debugging
-for i, intermediate in enumerate(processor.step_through(data)):
-    print(f"Step {i}: {processor.steps[i].__class__.__name__}")
-
-    # Set a breakpoint to inspect the current state
-    breakpoint()  # Debugger will pause here
-
-    # You can now inspect 'intermediate' in the debugger:
-    # - Check tensor shapes and values
-    # - Verify expected transformations
-    # - Look for unexpected changes
-```
-
-During the debugger session, you can:
-
-- Examine `intermediate[TransitionKey.OBSERVATION]` to see observation data
-- Check `intermediate[TransitionKey.ACTION]` for action transformations
-- Inspect any part of the transition to understand what each step does
-
-Step-through debugging is perfect for understanding the _data_ transformations, but what about the _structure_ of that data? While hooks and step-through help you debug runtime behavior, you also need to ensure your pipeline produces data in the format expected by downstream components. This is where feature contract validation comes in.
-
-## Validating Feature Contracts
-
-Feature contracts define what data structure your pipeline expects as input and produces as output.
-Validating these contracts helps catch mismatches early.
-
-### Understanding Feature Contracts
-
-Each processor step has a `transform_features()` method that describes how it changes the data structure:
-
-```python
-# Get the expected output features from your pipeline
-initial_features = {
-    PipelineFeatureType.OBSERVATION: {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(7,)),
-        "observation.image": PolicyFeature(type=FeatureType.IMAGE, shape=(3, 224, 224))
-    },
-    PipelineFeatureType.ACTION: {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,))
-    }
-}
-
-# Check what your pipeline will output
-output_features = processor.transform_features(initial_features)
-
-print("Input features:")
-for feature_type, features in initial_features.items():
-    print(f"  {feature_type}:")
-    for key, feature in features.items():
-        print(f"    {key}: {feature.type.value}, shape={feature.shape}")
-
-print("\nOutput features:")
-for feature_type, features in output_features.items():
-    print(f"  {feature_type}:")
-    for key, feature in features.items():
-        print(f"    {key}: {feature.type.value}, shape={feature.shape}")
-```
-
-### Verifying Expected Features
-
-Check that your pipeline produces the features you expect:
-
-```python
-# Define what features you expect the pipeline to produce
-expected_keys = ["observation.state", "observation.image", "action"]
-
-print("Validating feature contract...")
-for expected_key in expected_keys:
-    found = False
-    for feature_type, features in output_features.items():
-        if expected_key in features:
-            feature = features[expected_key]
-            print(f"✅ {expected_key}: {feature.type.value}, shape={feature.shape}")
-            found = True
-            break
-
-    if not found:
-        print(f"❌ Missing expected feature: {expected_key}")
-```
-
-This validation helps ensure your pipeline will work correctly with downstream components that expect specific data structures.
-
-## Summary
-
-Now that you understand the three debugging approaches, you can tackle any pipeline issue systematically:
-
-1. **Hooks** - For runtime monitoring and validation without modifying pipeline code
-2. **Step-through** - For inspecting intermediate states and understanding transformations
-3. **Feature validation** - For ensuring data structure contracts are met
-
-**When to use each approach:**
-
-- Start with **step-through debugging** when you need to understand what your pipeline does or when something unexpected happens
-- Add **hooks** for continuous monitoring during development and production to catch issues automatically
-- Use **feature validation** before deployment to ensure your pipeline works with downstream components
-
-These three tools work together to give you the complete observability that complex pipelines naturally lack. With hooks watching for issues, step-through helping you understand behavior, and feature validation ensuring compatibility, you'll be able to debug any pipeline confidently and efficiently.

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
 }
 ```
 
@@ -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"
 }
 ```
 

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

@@ -519,14 +519,11 @@ from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop
-from lerobot.policies.factory import make_processor
 
 NUM_EPISODES = 5
 FPS = 30
 EPISODE_TIME_SEC = 60
 TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
 
 # Create the robot configuration
 camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
@@ -538,7 +535,7 @@ robot_config = SO100FollowerConfig(
 robot = SO100Follower(robot_config)
 
 # Initialize the policy
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+policy = ACTPolicy.from_pretrained("<hf_username>/<my_policy_repo_id>")
 
 # Configure the dataset features
 action_features = hw_to_dataset_features(robot.action_features, "action")
@@ -547,7 +544,7 @@ dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
 dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
+    repo_id="<hf_username>/eval_<dataset_repo_id>",
     fps=FPS,
     features=dataset_features,
     robot_type=robot.name,
@@ -562,12 +559,6 @@ _init_rerun(session_name="recording")
 # Connect the robot
 robot.connect()
 
-preprocessor, postprocessor = make_processor(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
 for episode_idx in range(NUM_EPISODES):
     log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
 
@@ -577,8 +568,6 @@ for episode_idx in range(NUM_EPISODES):
         events=events,
         fps=FPS,
         policy=policy,
-        preprocessor=preprocessor,
-        postprocessor=postprocessor,
         dataset=dataset,
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,

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:
 
@@ -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

@@ -1,273 +0,0 @@
-# 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 `NormalizerProcessorStep` 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 sensors or when models output actions, you need to process this data to make it compatible with your target system. For example, a common need is normalizing data ranges to make them suitable for neural networks.
-
-LeRobot's `NormalizerProcessorStep` handles this crucial task:
-
-```python
-# Input: raw joint positions in [0, 180] degrees
-raw_action = torch.tensor([90.0, 45.0, 135.0])
-
-# After processing: normalized to [-1, 1] range for model training
-normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=dataset_stats)
-normalized_result = normalizer(transition)
-# ...
-```
-
-Other common processing needs include:
-
-- **Device placement**: Moving tensors between CPU/GPU and converting data types
-- **Format conversion**: Transforming between different data structures
-- **Batching**: Adding/removing batch dimensions for model compatibility
-- **Safety constraints**: Applying limits to robot commands
-
-```python
-# Example pipeline combining multiple processors
-pipeline = PolicyProcessorPipeline([
-    RenameObservationsProcessorStep(rename_map={}),
-    AddBatchDimensionProcessorStep(),
-    NormalizerProcessorStep(features=features, stats=stats),
-    DeviceProcessorStep(device="cuda"),
-    # ...
-])
-```
-
-LeRobot provides a pipeline mechanism to implement sequences of processing steps for both input data and output actions, making it easy to compose these transformations in the right order for optimal performance.
-
-## How to implement your own processor?
-
-We'll use the `NormalizerProcessorStep` as our main example because it demonstrates essential processor patterns including state management, configuration serialization, and tensor handling that you'll commonly need.
-
-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 `NormalizerProcessorStep` works:
-
-```python
-@dataclass
-@ProcessorStepRegistry.register("normalizer_processor")
-class NormalizerProcessorStep(ProcessorStep):
-    """Normalize observations/actions using dataset statistics."""
-
-    features: dict[str, PolicyFeature]
-    norm_map: dict[FeatureType, NormalizationMode]
-    stats: dict[str, dict[str, Any]] | None = None
-    eps: float = 1e-8
-    _tensor_stats: dict = field(default_factory=dict, init=False, repr=False)
-
-    def __post_init__(self):
-        """Convert stats to tensors for efficient computation."""
-        self.stats = self.stats or {}
-        self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=torch.float32)
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        new_transition = transition.copy()
-        # Normalize observations
-        # ...
-        # Normalize action
-        # ...
-        return new_transition
-
-```
-
-See the full implementation in `src/lerobot/processor/normalize_processor.py` for complete details.
-
-**Key principles:**
-
-- **Always use `transition.copy()`** to avoid side effects
-- **Handle both observations and actions** consistently
-- **Separate config from state**: `get_config()` returns JSON-serializable params, `state_dict()` returns tensors
-- **Convert stats to tensors** in `__post_init__()` for efficient computation
-
-### Configuration and State Management
-
-Processors support serialization through three methods that separate configuration from tensor state. The `NormalizerProcessorStep` demonstrates this perfectly - it carries dataset statistics (tensors) in its state, and hyperparameters in its config:
-
-```python
-# Continuing the NormalizerProcessorStep example...
-
-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.cpu()  # Always save to CPU
-    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)
-        # Load to processor's configured device
-        self._tensor_stats.setdefault(key, {})[stat_name] = tensor.to(
-            dtype=torch.float32, device=self.device
-        )
-        # ...
-```
-
-**Usage:**
-
-```python
-# Save (e.g., inside a policy)
-config = normalizer.get_config()
-tensors = normalizer.state_dict()
-
-# Restore (e.g., loading a pretrained policy)
-new_normalizer = NormalizerProcessorStep(**config)
-new_normalizer.load_state_dict(tensors)
-# Now new_normalizer has the same stats and configuration
-```
-
-### Transform features
-
-The `transform_features` method defines how your processor transforms feature names and shapes. This is crucial for policy configuration and debugging.
-
-For `NormalizerProcessorStep`, features are typically preserved unchanged since normalization doesn't alter keys or shapes:
-
-```python
-def transform_features(self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-    """Normalization preserves all feature definitions."""
-    return features  # No changes to feature structure
-    # ...
-```
-
-When your processor renames or reshapes data, 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
-
-### 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 `DataProcessorPipeline.from_pretrained(...)`.
-
-**Foundational model adaptation**: This is particularly useful when working with foundational pretrained policies where you rarely have access to the original training statistics. You can inject your own dataset statistics to adapt the normalizer to your specific robot or environment data.
-
-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.
-
-Direct usage with `from_pretrained`:
-
-```python
-from lerobot.processor import RobotProcessorPipeline
-
-# Load a foundational policy trained on diverse robot data
-# but adapt normalization to your specific robot/environment
-new_stats = LeRobotDataset(repo_id="username/my-dataset").meta.stats
-processor = RobotProcessorPipeline.from_pretrained(
-    "huggingface/foundational-robot-policy",  # Pretrained foundation model
-    overrides={
-        "normalizer_processor": {"stats": new_stats},     # Inject your robot's statistics
-        "device_processor": {"device": "cuda:0"},         # registry name for registered steps
-        "rename_processor": {"rename_map": robot_key_map}, # Map your robot's observation keys
-        # ...
-    },
-)
-```
-
-## Best Practices
-
-Based on analysis of all LeRobot processor implementations, here are the key patterns and practices:
-
-### 1. **Safe Data Handling**
-
-Always create copies of input data to avoid unintended side effects. Use `transition.copy()` and `observation.copy()` rather than modifying data in-place. This prevents your processor from accidentally affecting other components in the pipeline.
-
-Check for required data before processing and handle missing data gracefully. If your processor expects certain keys (like `"pixels"` for image processing), validate their presence first. For optional data, use safe access patterns like `transition.get()` and handle `None` values appropriately.
-
-When data validation fails, provide clear, actionable error messages that help users understand what went wrong and how to fix it.
-
-### 2. **Choose Appropriate Base Classes**
-
-LeRobot provides specialized base classes that reduce boilerplate code and ensure consistency. Use `ObservationProcessorStep` when you only need to modify observations, `ActionProcessorStep` for action-only processing, and `RobotActionProcessorStep` specifically for dictionary-based robot actions.
-
-Only inherit directly from `ProcessorStep` when you need full control over the entire transition or when processing multiple transition components simultaneously. The specialized base classes handle the transition management for you and provide type safety.
-
-### 3. **Registration and Naming**
-
-Register your processors with descriptive, namespaced names using `@ProcessorStepRegistry.register()`. Use organization prefixes like `"robotics_lab/safety_clipper"` or `"acme_corp/vision_enhancer"` to avoid naming conflicts. Avoid generic names like `"processor"` or `"step"` that could clash with other implementations.
-
-Good registration makes your processors discoverable and enables clean serialization/deserialization when saving and loading pipelines.
-
-### 4. **State Management Patterns**
-
-Distinguish between configuration parameters (JSON-serializable values) and internal state (tensors, buffers). Use dataclass fields with `init=False, repr=False` for internal state that shouldn't appear in the constructor or string representation.
-
-Implement the `reset()` method to clear internal state between episodes. This is crucial for stateful processors that accumulate data over time, like moving averages or temporal filters.
-
-Remember that `get_config()` should only return JSON-serializable configuration, while `state_dict()` handles tensor state separately.
-
-### 5. **Input Validation and Error Handling**
-
-Validate input types and shapes before processing. Check tensor properties like `dtype` and dimensions to ensure compatibility with your algorithms. For robot actions, verify that required pose components or joint values are present and within expected ranges.
-
-Use early returns for edge cases where no processing is needed. Provide clear, descriptive error messages that include the expected vs. actual data types or shapes. This makes debugging much easier for users.
-
-### 6. **Device and Dtype Awareness**
-
-Design your processors to automatically adapt to the device and dtype of input tensors. Internal tensors (like normalization statistics) should match the input tensor's device and dtype to ensure compatibility with multi-GPU training, mixed precision, and distributed setups.
-
-Implement a `to()` method that moves your processor's internal state to the specified device. Check device/dtype compatibility at runtime and automatically migrate internal state when needed. This pattern enables seamless operation across different hardware configurations without manual intervention.
-
-## 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`, `transform_features`)
-2. **Register it** using `@ProcessorStepRegistry.register("name")` for discoverability
-3. **Integrate it** into a `DataProcessorPipeline` with other processing steps
-4. **Use base classes** like `ObservationProcessorStep` when possible to reduce boilerplate
-5. **Implement device/dtype awareness** to support multi-GPU and mixed precision setups
-
-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.
-
-Key principles for robust processors:
-
-- **Device/dtype adaptation**: Internal tensors should match input tensors
-- **Clear error messages**: Help users understand what went wrong
-- **Base class usage**: Leverage specialized base classes to reduce boilerplate
-- **Feature contracts**: Declare data structure changes with `transform_features()`
-
-Start simple, test thoroughly, and ensure your processors work seamlessly across different hardware configurations!

docs/source/introduction_processors.mdx

@@ -1,314 +0,0 @@
-# Introduction to Processors
-
-In robotics, there's a fundamental mismatch between the data that robots and humans produce and what machine learning models expect.
-Robots output raw sensor data like camera images and joint positions that need normalization, batching, and device placement before models can process them.
-Language instructions from humans must be tokenized into numerical representations, and different robots use different coordinate systems that need standardization.
-
-The challenge extends to model outputs as well.
-Models might output end-effector positions while robots need joint-space commands, or teleoperators produce relative movements while robots expect absolute commands.
-Model predictions are often normalized and need conversion back to real-world scales.
-
-Cross-domain translation adds another layer of complexity.
-Training data from one robot setup needs adaptation for deployment on different hardware, models trained with specific camera configurations must work with new arrangements, and datasets with different naming conventions need harmonization.
-
-**That's where processors come in.** They serve as universal translators that bridge these gaps, ensuring seamless data flow from sensors to models to actuators.
-Processors handle all the preprocessing and postprocessing steps needed to convert raw environment data into model-ready inputs and vice versa.
-
-Now your favorite policy can be used like this:
-
-```python
-import torch
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.policies.factory import make_pre_post_processors
-from lerobot.policies.your_policy import YourPolicy
-from lerobot.processor.pipeline import RobotProcessorPipeline, PolicyProcessorPipeline
-dataset = LeRobotDataset("hf_user/dataset", episodes=[0])
-sample = dataset[10]
-
-model = YourPolicy.from_pretrained(
-    "hf_user/model",
-)
-model.eval()
-model.to("cuda")
-preprocessor, postprocessor = make_pre_post_processors(model.config, pretrained_path="hf_user/model", dataset_stats=dataset.meta.stats)
-
-preprocessed_sample = preprocessor(sample)
-action = model.select_action(preprocessed_sample)
-postprocessed_action = postprocessor(action)
-```
-
-## 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:
-
-- **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
-
-### ProcessorStep: The Building Block
-
-A `ProcessorStep` is a single transformation unit that processes transitions. It's an abstract base class with two required methods:
-
-```python
-from lerobot.processor import ProcessorStep, EnvTransition
-
-class MyProcessorStep(ProcessorStep):
-    """Example processor step - inherit and implement abstract methods."""
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """Transform the transition - REQUIRED abstract method."""
-        # Your processing logic here
-        return transition
-
-    def transform_features(self, features):
-        """Declare how this step transforms feature shapes/types - REQUIRED abstract method."""
-        return features  # Most processors return features unchanged
-```
-
-`__call__` is the core of your processor step. It takes an `EnvTransition` and returns a modified `EnvTransition`.
-
-`transform_features` is used to declare how this step transforms feature shapes/types.
-
-### DataProcessorPipeline: The Generic Orchestrator
-
-The `DataProcessorPipeline[TInput, TOutput]` chains multiple `ProcessorStep` instances together:
-
-```python
-from lerobot.processor import RobotProcessorPipeline, PolicyProcessorPipeline
-
-# For robot hardware (unbatched data)
-robot_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[step1, step2, step3],
-    name="robot_pipeline"
-)
-
-# For model training/inference (batched data)
-policy_processor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-    steps=[step1, step2, step3],
-    name="policy_pipeline"
-)
-```
-
-## RobotProcessorPipeline vs PolicyProcessorPipeline
-
-The key distinction is in the data structures they handle:
-
-| Aspect          | RobotProcessorPipeline                       | PolicyProcessorPipeline                  |
-| --------------- | -------------------------------------------- | ---------------------------------------- |
-| **Input**       | `dict[str, Any]` - Individual robot values   | `dict[str, Any]` - Batched tensors       |
-| **Output**      | `dict[str, Any]` - Individual robot commands | `torch.Tensor` - Policy predictions      |
-| **Use Case**    | Real-time robot control                      | Model training/inference                 |
-| **Data Format** | Unbatched, heterogeneous                     | Batched, homogeneous                     |
-| **Examples**    | `{"joint_1": 0.5}`                           | `{"observation.state": tensor([[0.5]])}` |
-
-**Use `RobotProcessorPipeline`** for robot hardware interfaces:
-
-```python
-# Robot data structures: dict[str, Any] for observations and actions
-robot_obs: dict[str, Any] = {
-    "joint_1": 0.5,           # Individual joint values
-    "joint_2": -0.3,
-    "camera_0": image_array   # Raw camera data
-}
-
-robot_action: dict[str, Any] = {
-    "joint_1": 0.2,          # Target joint positions
-    "joint_2": 0.1,
-    "gripper": 0.8
-}
-```
-
-**Use `PolicyProcessorPipeline`** for model training and batch processing:
-
-```python
-# Policy data structures: batch dicts and tensors
-policy_batch: dict[str, Any] = {
-    "observation.state": torch.tensor([[0.5, -0.3]]),      # Batched states
-    "observation.images.camera0": torch.tensor(...),        # Batched images
-    "action": torch.tensor([[0.2, 0.1, 0.8]])              # Batched actions
-}
-
-policy_action: torch.Tensor = torch.tensor([[0.2, 0.1, 0.8]])  # Model output tensor
-```
-
-## Converter Functions
-
-LeRobot provides converter functions to bridge different data formats in `lerobot.processor.converters`. These functions handle the crucial translations between robot hardware data structures, policy model formats, and the internal `EnvTransition` representation that flows through processor pipelines.
-
-| Category                       | Function                      | Description                     |
-| ------------------------------ | ----------------------------- | ------------------------------- |
-| **Robot Hardware Converters**  | `robot_action_to_transition`  | Robot dict → EnvTransition      |
-|                                | `observation_to_transition`   | Robot obs → EnvTransition       |
-|                                | `transition_to_robot_action`  | EnvTransition → Robot dict      |
-| **Policy/Training Converters** | `batch_to_transition`         | Batch dict → EnvTransition      |
-|                                | `transition_to_batch`         | EnvTransition → Batch dict      |
-|                                | `policy_action_to_transition` | Policy tensor → EnvTransition   |
-|                                | `transition_to_policy_action` | EnvTransition → Policy tensor   |
-| **Utilities**                  | `create_transition`           | Build transitions with defaults |
-|                                | `identity_transition`         | Pass-through converter          |
-
-The key insight is that **robot hardware converters** work with individual values and dictionaries, while **policy/training converters** work with batched tensors and model outputs. The converter functions automatically handle the structural differences, so your processor steps can focus on the core transformations without worrying about data format compatibility.
-
-## Processor Examples
-
-The following examples demonstrate real-world processor configurations for policy training and inference.
-
-Here is an example processor for policy training and inference:
-
-```python
-# Training data preprocessing (optimized order for GPU performance)
-training_preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-    steps=[
-        RenameObservationsProcessorStep(rename_map={}),     # Standardize keys
-        AddBatchDimensionProcessorStep(),                   # Add batch dims
-        TokenizerProcessorStep(tokenizer_name="...", ...),  # Tokenize language
-        DeviceProcessorStep(device="cuda"),                 # Move to GPU first ⚡
-        NormalizerProcessorStep(features=..., stats=...),   # Normalize on GPU ⚡
-    ]
-)
-
-# Model output postprocessing
-training_postprocessor = PolicyProcessorPipeline[torch.Tensor, torch.Tensor](
-    steps=[
-        DeviceProcessorStep(device="cpu"),                  # Move to CPU
-        UnnormalizerProcessorStep(features=..., stats=...), # Denormalize
-    ]
-    to_transition=policy_action_to_transition,
-    to_output=transition_to_policy_action,
-)
-```
-
-### An interaction between a robot and a policy with processors
-
-The most common real-world scenario combines both pipeline types robot hardware generates observations that need policy processing, and policy outputs need robot-compatible postprocessing:
-
-```python
-# Real deployment: Robot sensors → Model → Robot commands
-with torch.no_grad():
-    while not done:
-        raw_obs = robot.get_observation()  # dict[str, Any]
-
-        # Add your robot observation to policy observation processor
-
-        policy_input = policy_preprocessor(raw_obs)  # Batched dict
-
-        policy_output = policy.select_action(policy_input)  # Policy tensor
-
-        policy_action = policy_postprocessor(policy_output)
-
-        # Add your robot action to policy action processor
-
-        robot.send_action(policy_action)
-```
-
-## Feature Contracts: Shape and Type Transformation
-
-Processors don't just transform data - they can also **change the data structure itself**. The `transform_features()` method declares these changes, which is crucial for dataset recording and policy creation.
-
-### Why Feature Contracts Matter
-
-When building datasets or policies, LeRobot needs to know:
-
-- **What data fields will exist** after processing
-- **What shapes and types** each field will have
-- **How to configure models** for the expected data structure
-
-```python
-# Example: A processor that adds velocity to observations
-class VelocityProcessor(ObservationProcessorStep):
-    def observation(self, obs):
-        new_obs = obs.copy()
-        if "observation.state" in obs:
-            # concatenate computed velocity field to the state
-            new_obs["observation.state"] = self._compute_velocity(obs["observation.state"])
-        return new_obs
-
-    def transform_features(self, features):
-        """Declare the new velocity field we're adding."""
-        state_feature = features[PipelineFeatureType.OBSERVATION].get("observation.state")
-        if state_feature:
-            double_shape = (state_feature.shape[0] * 2,) if state_feature.shape else (2,)
-            features[PipelineFeatureType.OBSERVATION]["observation.state"] = PolicyFeature(
-                type=FeatureType.STATE, shape=double_shape
-            )
-        return features
-```
-
-### Feature Specification Functions
-
-`create_initial_features()` and `aggregate_pipeline_dataset_features()` solve a critical dataset creation problem: determining the exact final data structure before any data is processed.
-Since processor pipelines can add new features (like velocity fields), change tensor shapes (like cropping images), or rename keys, datasets need to know the complete output specification upfront to allocate proper storage and define schemas.
-These functions work together by starting with robot hardware specifications (`create_initial_features()`) then simulating the entire pipeline transformation (`aggregate_pipeline_dataset_features()`) to compute the final feature dictionary that gets passed to `LeRobotDataset.create()`, ensuring perfect alignment between what processors output and what datasets expect to store.
-
-```python
-from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
-
-# Start with robot's raw features
-initial_features = create_initial_features(
-    observation=robot.observation_features,  # {"joint_1.pos": float, "camera_0": (480,640,3)}
-    action=robot.action_features            # {"joint_1.pos": float, "gripper.pos": float}
-)
-
-# Apply processor pipeline to compute final features
-final_features = aggregate_pipeline_dataset_features(
-    pipeline=my_processor_pipeline,
-    initial_features=initial_features,
-    use_videos=True
-)
-
-# Use for dataset creation
-dataset = LeRobotDataset.create(
-    repo_id="my_dataset",
-    features=final_features,  # Knows exactly what data to expect
-    ...
-)
-```
-
-## Common Processor Steps
-
-LeRobot provides many registered processor steps. Here are the most commonly used core processors:
-
-### Essential Processors
-
-- **`normalizer_processor`**: Normalize observations/actions using dataset statistics (mean/std or min/max)
-- **`device_processor`**: Move tensors to CPU/GPU with optional dtype conversion
-- **`to_batch_processor`**: Add batch dimensions to transitions for model compatibility
-- **`rename_observations_processor`**: Rename observation keys using mapping dictionaries
-- **`tokenizer_processor`**: Tokenize natural language task descriptions into tokens and attention masks
-
-### Next Steps
-
-- **[Implement Your Own Processor](implement_your_own_processor.mdx)** - Create custom processor steps
-- **[Debug Your Pipeline](debug_processor_pipeline.mdx)** - Troubleshoot and optimize pipelines
-- **[Processors for Robots and Teleoperators](processors_robots_teleop.mdx)** - Real-world integration patterns
-
-## Summary
-
-Processors solve the data translation problem in robotics by providing:
-
-- **Modular transformations**: Composable, reusable processing steps
-- **Type safety**: Generic pipelines with compile-time checking
-- **Performance optimization**: GPU-accelerated operations
-- **Robot/Policy distinction**: Separate pipelines for different data structures
-- **Comprehensive ecosystem**: 30+ registered processors for common tasks
-
-The key insight: `RobotProcessorPipeline` handles unbatched robot hardware data, while `PolicyProcessorPipeline` handles batched model data. Choose the right tool for your data structure!

docs/source/lerobot-dataset-v3.mdx

@@ -1,169 +0,0 @@
-# LeRobotDataset v3.0
-
-`LeRobotDataset v3.0` is a standardized format for robot learning data. It provides unified access to multi-modal time-series data, sensorimotor signals and multi‑camera video, as well as rich metadata for indexing, search, and visualization on the Hugging Face Hub.
-
-This docs will guide you to:
-
-- Understand the v3.0 design and directory layout
-- Record a dataset and push it to the Hub
-- Load datasets for training with `LeRobotDataset`
-- Stream datasets without downloading using `StreamingLeRobotDataset`
-- Migrate existing `v2.1` datasets to `v3.0`
-
-## What’s new in `v3`
-
-- **File-based storage**: Many episodes per Parquet/MP4 file (v2 used one file per episode).
-- **Relational metadata**: Episode boundaries and lookups are resolved through metadata, not filenames.
-- **Hub-native streaming**: Consume datasets directly from the Hub with `StreamingLeRobotDataset`.
-- **Lower file-system pressure**: Fewer, larger files ⇒ faster initialization and fewer issues at scale.
-- **Unified organization**: Clean directory layout with consistent path templates across data and videos.
-
-## Installation
-
-`LeRobotDataset v3.0` will be included in `lerobot >= 0.4.0`.
-
-Until that stable release, you can use the main branch by following the [build from source instructions](./installation#from-source).
-
-## Record a dataset
-
-Run the command below to record a dataset with the SO-101 and push to the Hub:
-
-```bash
-lerobot-record \
-  --robot.type=so101_follower \
-  --robot.port=/dev/tty.usbmodem585A0076841 \
-  --robot.id=my_awesome_follower_arm \
-  --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \
-  --teleop.type=so101_leader \
-  --teleop.port=/dev/tty.usbmodem58760431551 \
-  --teleop.id=my_awesome_leader_arm \
-  --display_data=true \
-  --dataset.repo_id=${HF_USER}/record-test \
-  --dataset.num_episodes=5 \
-  --dataset.single_task="Grab the black cube"
-```
-
-See the [recording guide](./il_robots#record-a-dataset) for more details.
-
-## Format design
-
-A core v3 principle is **decoupling storage from the user API**: data is stored efficiently (few large files), while the public API exposes intuitive episode-level access.
-
-`v3` has three pillars:
-
-1. **Tabular data**: Low‑dimensional, high‑frequency signals (states, actions, timestamps) stored in **Apache Parquet**. Access is memory‑mapped or streamed via the `datasets` stack.
-2. **Visual data**: Camera frames concatenated and encoded into **MP4**. Frames from the same episode are grouped; videos are sharded per camera for practical sizes.
-3. **Metadata**: JSON/Parquet records describing schema (feature names, dtypes, shapes), frame rates, normalization stats, and **episode segmentation** (start/end offsets into shared Parquet/MP4 files).
-
-> To scale to millions of episodes, tabular rows and video frames from multiple episodes are **concatenated** into larger files. Episode‑specific views are reconstructed **via metadata**, not file boundaries.
-
-<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
-  <figure style="margin:0; text-align:center;">
-    <img
-      src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobotdataset-v3/asset1datasetv3.png"
-      alt="LeRobotDataset v3 diagram"
-      width="220"
-    />
-    <figcaption style="font-size:0.9em; color:#666;">
-      From episode‑based to file‑based datasets
-    </figcaption>
-  </figure>
-</div>
-
-### Directory layout (simplified)
-
-- **`meta/info.json`**: canonical schema (features, shapes/dtypes), FPS, codebase version, and **path templates** to locate data/video shards.
-- **`meta/stats.json`**: global feature statistics (mean/std/min/max) used for normalization; exposed as `dataset.meta.stats`.
-- **`meta/tasks.jsonl`**: natural‑language task descriptions mapped to integer IDs for task‑conditioned policies.
-- **`meta/episodes/`**: per‑episode records (lengths, tasks, offsets) stored as **chunked Parquet** for scalability.
-- **`data/`**: frame‑by‑frame **Parquet** shards; each file typically contains **many episodes**.
-- **`videos/`**: **MP4** shards per camera; each file typically contains **many episodes**.
-
-## Load a dataset for training
-
-`LeRobotDataset` returns Python dictionaries of PyTorch tensors and integrates with `torch.utils.data.DataLoader`. Here is a code example showing its use:
-
-```python
-import torch
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-repo_id = "yaak-ai/L2D-v3"
-
-# 1) Load from the Hub (cached locally)
-dataset = LeRobotDataset(repo_id)
-
-# 2) Random access by index
-sample = dataset[100]
-print(sample)
-# {
-#   'observation.state': tensor([...]),
-#   'action': tensor([...]),
-#   'observation.images.front_left': tensor([C, H, W]),
-#   'timestamp': tensor(1.234),
-#   ...
-# }
-
-# 3) Temporal windows via delta_timestamps (seconds relative to t)
-delta_timestamps = {
-    "observation.images.front_left": [-0.2, -0.1, 0.0]  # 0.2s and 0.1s before current frame
-}
-
-dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
-
-# Accessing an index now returns a stack for the specified key(s)
-sample = dataset[100]
-print(sample["observation.images.front_left"].shape)  # [T, C, H, W], where T=3
-
-# 4) Wrap with a DataLoader for training
-batch_size = 16
-data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-for batch in data_loader:
-    observations = batch["observation.state"].to(device)
-    actions = batch["action"].to(device)
-    images = batch["observation.images.front_left"].to(device)
-    # model.forward(batch)
-```
-
-## Stream a dataset (no downloads)
-
-Use `StreamingLeRobotDataset` to iterate directly from the Hub without local copies. This allows to stream large datasets without the need to downloading them onto disk or loading them onto memory, and is a key feature of the new dataset format.
-
-```python
-from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
-
-repo_id = "yaak-ai/L2D-v3"
-dataset = StreamingLeRobotDataset(repo_id)  # streams directly from the Hub
-```
-
-<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
-  <figure style="margin:0; text-align:center;">
-    <img
-      src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobotdataset-v3/streaming-lerobot.png"
-      alt="StreamingLeRobotDataset"
-      width="520"
-    />
-    <figcaption style="font-size:0.9em; color:#666;">
-      Stream directly from the Hub for on‑the‑fly training.
-    </figcaption>
-  </figure>
-</div>
-
-## Migrate `v2.1` → `v3.0`
-
-A converter aggregates per‑episode files into larger shards and writes episode offsets/metadata. Convert your dataset using the instructions below.
-
-```bash
-# Pre-release build with v3 support:
-pip install "https://github.com/huggingface/lerobot/archive/33cad37054c2b594ceba57463e8f11ee374fa93c.zip"
-
-# Convert an existing v2.1 dataset hosted on the Hub:
-python -m lerobot.datasets.v30.convert_dataset_v21_to_v30 --repo-id=<HF_USER/DATASET_ID>
-```
-
-**What it does**
-
-- Aggregates parquet files: `episode-0000.parquet`, `episode-0001.parquet`, … → **`file-0000.parquet`**, …
-- Aggregates mp4 files: `episode-0000.mp4`, `episode-0001.mp4`, … → **`file-0000.mp4`**, …
-- Updates `meta/episodes/*` (chunked Parquet) with per‑episode lengths, tasks, and byte/frame offsets.

docs/source/libero.mdx

@@ -1,126 +0,0 @@
-# LIBERO
-
-**LIBERO** is a benchmark designed to study **lifelong robot learning**. The idea is that robots won’t just be pretrained once in a factory, they’ll need to keep learning and adapting with their human users over time. This ongoing adaptation is called **lifelong learning in decision making (LLDM)**, and it’s a key step toward building robots that become truly personalized helpers.
-
-- 📄 [LIBERO paper](https://arxiv.org/abs/2306.03310)
-- 💻 [Original LIBERO repo](https://github.com/Lifelong-Robot-Learning/LIBERO)
-
-To make progress on this challenge, LIBERO provides a set of standardized tasks that focus on **knowledge transfer**: how well a robot can apply what it has already learned to new situations. By evaluating on LIBERO, different algorithms can be compared fairly and researchers can build on each other’s work.
-
-LIBERO includes **five task suites**:
-
-- **LIBERO-Spatial (`libero_spatial`)** – tasks that require reasoning about spatial relations.
-- **LIBERO-Object (`libero_object`)** – tasks centered on manipulating different objects.
-- **LIBERO-Goal (`libero_goal`)** – goal-conditioned tasks where the robot must adapt to changing targets.
-- **LIBERO-90 (`libero_90`)** – 90 short-horizon tasks from the LIBERO-100 collection.
-- **LIBERO-Long (`libero_10`)** – 10 long-horizon tasks from the LIBERO-100 collection.
-
-Together, these suites cover **130 tasks**, ranging from simple object manipulations to complex multi-step scenarios. LIBERO is meant to grow over time, and to serve as a shared benchmark where the community can test and improve lifelong learning algorithms.
-
-![An overview of the LIBERO benchmark](https://libero-project.github.io/assets/img/libero/fig1.png)
-
-## Evaluating with LIBERO
-
-At **LeRobot**, we ported [LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO) into our framework and used it mainly to **evaluate [SmolVLA](https://huggingface.co/docs/lerobot/en/smolvla)**, our lightweight Vision-Language-Action model.
-
-LIBERO is now part of our **multi-eval supported simulation**, meaning you can benchmark your policies either on a **single suite of tasks** or across **multiple suites at once** with just a flag.
-
-To Install LIBERO, after following LeRobot official instructions, just do:
-`pip install -e ".[libero]"`
-
-### Single-suite evaluation
-
-Evaluate a policy on one LIBERO suite:
-
-```bash
-python src/lerobot/scripts/eval.py \
-  --policy.path="your-policy-id" \
-  --env.type=libero \
-  --env.task=libero_object \
-  --eval.batch_size=2 \
-  --eval.n_episodes=3
-```
-
-- `--env.task` picks the suite (`libero_object`, `libero_spatial`, etc.).
-- `--eval.batch_size` controls how many environments run in parallel.
-- `--eval.n_episodes` sets how many episodes to run in total.
-
----
-
-### Multi-suite evaluation
-
-Benchmark a policy across multiple suites at once:
-
-```bash
-python src/lerobot/scripts/eval.py \
-  --policy.path="your-policy-id" \
-  --env.type=libero \
-  --env.task=libero_object,libero_spatial \
-  --eval.batch_size=1 \
-  --eval.n_episodes=2
-```
-
-- Pass a comma-separated list to `--env.task` for multi-suite evaluation.
-
-### Policy inputs and outputs
-
-When using LIBERO through LeRobot, policies interact with the environment via **observations** and **actions**:
-
-- **Observations**
-  - `observation.state` – proprioceptive features (agent state).
-  - `observation.images.image` – main camera view (`agentview_image`).
-  - `observation.images.image2` – wrist camera view (`robot0_eye_in_hand_image`).
-
-  ⚠️ **Note:** LeRobot enforces the `.images.*` prefix for any multi-modal visual features. Always ensure that your policy config `input_features` use the same naming keys, and that your dataset metadata keys follow this convention during evaluation.
-  If your data contains different keys, you must rename the observations to match what the policy expects, since naming keys are encoded inside the normalization statistics layer.
-  This will be fixed with the upcoming Pipeline PR.
-
-- **Actions**
-  - Continuous control values in a `Box(-1, 1, shape=(7,))` space.
-
-We also provide a notebook for quick testing:
-Training with LIBERO
-
-## Training with LIBERO
-
-When training on LIBERO tasks, make sure your dataset parquet and metadata keys follow the LeRobot convention.
-
-The environment expects:
-
-- `observation.state` → 8-dim agent state
-- `observation.images.image` → main camera (`agentview_image`)
-- `observation.images.image2` → wrist camera (`robot0_eye_in_hand_image`)
-
-⚠️ Cleaning the dataset upfront is **cleaner and more efficient** than remapping keys inside the code.
-To avoid potential mismatches and key errors, we provide a **preprocessed LIBERO dataset** that is fully compatible with the current LeRobot codebase and requires no additional manipulation:
-👉 [HuggingFaceVLA/libero](https://huggingface.co/datasets/HuggingFaceVLA/libero)
-
-For reference, here is the **original dataset** published by Physical Intelligence:
-👉 [physical-intelligence/libero](https://huggingface.co/datasets/physical-intelligence/libero)
-
----
-
-### Example training command
-
-```bash
-python src/lerobot/scripts/train.py \
-  --policy.type=smolvla \
-  --policy.repo_id=${HF_USER}/libero-test \
-  --dataset.repo_id=jadechoghari/smol-libero3 \
-  --env.type=libero \
-  --env.task=libero_10 \
-  --output_dir=./outputs/ \
-  --steps=100000 \
-  --batch_size=4 \
-  --eval.batch_size=1 \
-  --eval.n_episodes=1 \
-  --eval_freq=1000 \
-```
-
----
-
-### Note on rendering
-
-LeRobot uses MuJoCo for simulation. You need to set the rendering backend before training or evaluation:
-
-- `export MUJOCO_GL=egl` → for headless servers (e.g. HPC, cloud)

docs/source/phone_teleop.mdx

@@ -1,195 +0,0 @@
-# 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/porting_datasets_v3.mdx

@@ -150,7 +150,7 @@ gsutil -m cp -r gs://gresearch/robotics/droid_100 /your/data/
 ### Step 3: Port the Dataset
 
 ```bash
-python examples/port_datasets/port_droid.py \
+python examples/port_datasets/droid_rlds/port_droid.py \
     --raw-dir /your/data/droid/1.0.1 \
     --repo-id your_id/droid_1.0.1 \
     --push-to-hub
@@ -161,7 +161,7 @@ python examples/port_datasets/port_droid.py \
 For development, you can port a single shard:
 
 ```bash
-python examples/port_datasets/port_droid.py \
+python examples/port_datasets/droid_rlds/port_droid.py \
     --raw-dir /your/data/droid/1.0.1 \
     --repo-id your_id/droid_1.0.1_test \
     --num-shards 2048 \
@@ -194,7 +194,7 @@ Choose a **CPU partition** - no GPU needed for dataset porting.
 ### Step 3: Launch Parallel Porting Jobs
 
 ```bash
-python examples/port_datasets/slurm_port_shards.py \
+python examples/port_datasets/droid_rlds/slurm_port_shards.py \
     --raw-dir /your/data/droid/1.0.1 \
     --repo-id your_id/droid_1.0.1 \
     --logs-dir /your/logs \
@@ -245,7 +245,7 @@ failed_logs /your/logs/port_droid
 Once all porting jobs complete:
 
 ```bash
-python examples/port_datasets/slurm_aggregate_shards.py \
+python examples/port_datasets/droid_rlds/slurm_aggregate_shards.py \
     --repo-id your_id/droid_1.0.1 \
     --logs-dir /your/logs \
     --job-name aggr_droid \
@@ -258,7 +258,7 @@ python examples/port_datasets/slurm_aggregate_shards.py \
 ### Step 6: Upload to Hub
 
 ```bash
-python examples/port_datasets/slurm_upload.py \
+python examples/port_datasets/droid_rlds/slurm_upload.py \
     --repo-id your_id/droid_1.0.1 \
     --logs-dir /your/logs \
     --job-name upload_droid \

docs/source/processors_robots_teleop.mdx

@@ -1,148 +0,0 @@
-# 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
-```

examples/5_train_with_streaming.py

@@ -1,116 +0,0 @@
-# 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 demonstrates how to train a Diffusion Policy on the PushT environment,
-using a dataset processed in streaming mode.
-
-Once you have trained a model with this script, you can try to evaluate it on
-examples/2_evaluate_pretrained_policy.py
-"""
-
-from pathlib import Path
-
-import torch
-
-from lerobot.configs.types import FeatureType
-from lerobot.constants import ACTION
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
-from lerobot.datasets.utils import dataset_to_policy_features
-from lerobot.policies.act.configuration_act import ACTConfig
-from lerobot.policies.act.modeling_act import ACTPolicy
-
-
-def main():
-    # Create a directory to store the training checkpoint.
-    output_directory = Path("outputs/train/example_streaming_dataset")
-    output_directory.mkdir(parents=True, exist_ok=True)
-
-    # Selects the "best" device available
-    device = (
-        torch.device("cuda")
-        if torch.cuda.is_available()
-        else torch.device("mps")
-        if torch.backends.mps.is_available()
-        else torch.device("cpu")
-    )
-    print(f"Using device: {device}")
-
-    training_steps = 10
-    log_freq = 1
-
-    dataset_id = (
-        "aractingi/droid_1.0.1"  # 26M frames! Would require 4TB of disk space if installed locally (:
-    )
-    dataset_metadata = LeRobotDatasetMetadata(dataset_id)
-    features = dataset_to_policy_features(dataset_metadata.features)
-    output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
-    input_features = {key: ft for key, ft in features.items() if key not in output_features}
-
-    # We can now instantiate our policy with this config and the dataset stats.
-    cfg = ACTConfig(input_features=input_features, output_features=output_features)
-    policy = ACTPolicy(cfg, dataset_stats=dataset_metadata.stats)
-    policy.train()
-    policy.to(device)
-
-    # Delta timestamps are used to (1) augment frames used during training and (2) supervise the policy.
-    # Here, we use delta-timestamps to only provide ground truth actions for supervision
-    delta_timestamps = {
-        ACTION: [t / dataset_metadata.fps for t in range(cfg.n_action_steps)],
-    }
-
-    # Instantiating the training dataset in streaming mode allows to not consume up memory as the data is fetched
-    # iteratively rather than being load into memory all at once. Retrieved frames are shuffled across epochs
-    dataset = StreamingLeRobotDataset(dataset_id, delta_timestamps=delta_timestamps, tolerance_s=1e-3)
-
-    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        num_workers=4,
-        batch_size=16,
-        pin_memory=device.type != "cpu",
-        drop_last=True,
-        prefetch_factor=2,  # loads batches with multiprocessing while policy trains
-    )
-
-    # Run training loop.
-    step = 0
-    done = False
-    while not done:
-        for batch in dataloader:
-            batch = {
-                k: (v.type(torch.float32) if isinstance(v, torch.Tensor) and v.dtype != torch.bool else v)
-                for k, v in batch.items()
-            }
-            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
-
-            # batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
-            loss, _ = policy.forward(batch)
-            loss.backward()
-            optimizer.step()
-            optimizer.zero_grad()
-
-            if step % log_freq == 0:
-                print(f"step: {step} loss: {loss.item():.3f}")
-            step += 1
-            if step >= training_steps:
-                done = True
-                break
-
-    # Save a policy checkpoint.
-    policy.save_pretrained(output_directory)
-
-
-if __name__ == "__main__":
-    main()

examples/lekiwi/evaluate.py

@@ -1,24 +1,6 @@
-# !/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 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.processor import make_default_processors
 from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.utils.control_utils import init_keyboard_listener
@@ -29,16 +11,12 @@ NUM_EPISODES = 2
 FPS = 30
 EPISODE_TIME_SEC = 60
 TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
 
-# Create the robot configuration & robot
+# Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
-
 robot = LeKiwiClient(robot_config)
 
-# Create policy
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
+policy = ACTPolicy.from_pretrained("<hf_username>/<policy_repo_id>")
 
 # Configure the dataset features
 action_features = hw_to_dataset_features(robot.action_features, "action")
@@ -47,7 +25,7 @@ dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
 dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
+    repo_id="<hf_username>/<eval_dataset_repo_id>",
     fps=FPS,
     features=dataset_features,
     robot_type=robot.name,
@@ -55,50 +33,33 @@ dataset = LeRobotDataset.create(
     image_writer_threads=4,
 )
 
-# Build Policy Processors
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
-# Connect the robot
 # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
 
-# TODO(Steven): Update this example to use pipelines
-teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+_init_rerun(session_name="recording")
 
-# Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-_init_rerun(session_name="lekiwi_evaluate")
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")
 
-print("Starting evaluate loop...")
 recorded_episodes = 0
 while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
     log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
 
-    # Main record loop
+    # Run the policy inference loop
     record_loop(
         robot=robot,
         events=events,
         fps=FPS,
         policy=policy,
-        preprocessor=preprocessor,  # Pass the pre and post policy processors
-        postprocessor=postprocessor,
         dataset=dataset,
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,
         display_data=True,
-        teleop_action_processor=teleop_action_processor,
-        robot_action_processor=robot_action_processor,
-        robot_observation_processor=robot_observation_processor,
     )
 
-    # Reset the environment if not stopping or re-recording
+    # Logic for reset env
     if not events["stop_recording"] and (
         (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
     ):
@@ -110,9 +71,6 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
             control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
-            teleop_action_processor=teleop_action_processor,
-            robot_action_processor=robot_action_processor,
-            robot_observation_processor=robot_observation_processor,
         )
 
     if events["rerecord_episode"]:
@@ -122,12 +80,11 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
         dataset.clear_episode_buffer()
         continue
 
-    # Save episode
     dataset.save_episode()
     recorded_episodes += 1
 
-# Clean up
-log_say("Stop recording")
+# Upload to hub and clean up
+dataset.push_to_hub()
+
 robot.disconnect()
 listener.stop()
-dataset.push_to_hub()

examples/lekiwi/record.py

@@ -1,22 +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.
-
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.processor import make_default_processors
 from lerobot.record import record_loop
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
@@ -31,21 +14,16 @@ FPS = 30
 EPISODE_TIME_SEC = 30
 RESET_TIME_SEC = 10
 TASK_DESCRIPTION = "My task description"
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
 
 # Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
 keyboard_config = KeyboardTeleopConfig()
 
-# Initialize the robot and teleoperator
 robot = LeKiwiClient(robot_config)
 leader_arm = SO100Leader(leader_arm_config)
 keyboard = KeyboardTeleop(keyboard_config)
 
-# TODO(Steven): Update this example to use pipelines
-teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
-
 # Configure the dataset features
 action_features = hw_to_dataset_features(robot.action_features, "action")
 obs_features = hw_to_dataset_features(robot.observation_features, "observation")
@@ -53,7 +31,7 @@ dataset_features = {**action_features, **obs_features}
 
 # Create the dataset
 dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
+    repo_id="<hf_username>/<dataset_repo_id>",
     fps=FPS,
     features=dataset_features,
     robot_type=robot.name,
@@ -61,25 +39,23 @@ dataset = LeRobotDataset.create(
     image_writer_threads=4,
 )
 
-# Connect the robot and teleoperator
 # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
 leader_arm.connect()
 keyboard.connect()
 
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
 _init_rerun(session_name="lekiwi_record")
 
-if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
+listener, events = init_keyboard_listener()
+
+if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
+    raise ValueError("Robot, leader arm of keyboard is not connected!")
 
-print("Starting record loop...")
 recorded_episodes = 0
 while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
     log_say(f"Recording episode {recorded_episodes}")
 
-    # Main record loop
+    # Run the record loop
     record_loop(
         robot=robot,
         events=events,
@@ -89,12 +65,9 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
         control_time_s=EPISODE_TIME_SEC,
         single_task=TASK_DESCRIPTION,
         display_data=True,
-        teleop_action_processor=teleop_action_processor,
-        robot_action_processor=robot_action_processor,
-        robot_observation_processor=robot_observation_processor,
     )
 
-    # Reset the environment if not stopping or re-recording
+    # Logic for reset env
     if not events["stop_recording"] and (
         (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
     ):
@@ -107,9 +80,6 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
             control_time_s=RESET_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
-            teleop_action_processor=teleop_action_processor,
-            robot_action_processor=robot_action_processor,
-            robot_observation_processor=robot_observation_processor,
         )
 
     if events["rerecord_episode"]:
@@ -119,14 +89,13 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
         dataset.clear_episode_buffer()
         continue
 
-    # Save episode
     dataset.save_episode()
     recorded_episodes += 1
 
-# Clean up
-log_say("Stop recording")
+# Upload to hub and clean up
+dataset.push_to_hub()
+
 robot.disconnect()
 leader_arm.disconnect()
 keyboard.disconnect()
 listener.stop()
-dataset.push_to_hub()

examples/lekiwi/replay.py

@@ -1,19 +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");
-# 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 lerobot.datasets.lerobot_dataset import LeRobotDataset
@@ -24,34 +8,25 @@ from lerobot.utils.utils import log_say
 
 EPISODE_IDX = 0
 
-# Initialize the robot config
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
-
-# Initialize the robot
 robot = LeKiwiClient(robot_config)
 
-# Fetch the dataset to replay
 dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
 actions = dataset.hf_dataset.select_columns("action")
 
-# Connect to the robot
 robot.connect()
 
 if not robot.is_connected:
     raise ValueError("Robot is not connected!")
 
-print("Starting replay loop...")
 log_say(f"Replaying episode {EPISODE_IDX}")
 for idx in range(dataset.num_frames):
     t0 = time.perf_counter()
 
-    # Get recorded action from dataset
     action = {
         name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
     }
-
-    # Send action to robot
-    _ = robot.send_action(action)
+    robot.send_action(action)
 
     busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
 

examples/lekiwi/teleoperate.py

@@ -1,19 +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");
-# 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 lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
@@ -29,44 +13,35 @@ robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
 teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
 keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
 
-# Initialize the robot and teleoperator
 robot = LeKiwiClient(robot_config)
 leader_arm = SO100Leader(teleop_arm_config)
 keyboard = KeyboardTeleop(keyboard_config)
 
-# Connect to the robot and teleoperator
 # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
 leader_arm.connect()
 keyboard.connect()
 
-# Init rerun viewer
 _init_rerun(session_name="lekiwi_teleop")
 
 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
+    raise ValueError("Robot, leader arm of keyboard is not connected!")
 
-print("Starting teleop loop...")
 while True:
     t0 = time.perf_counter()
 
-    # Get robot observation
     observation = robot.get_observation()
 
-    # Get teleop action
-    # Arm
     arm_action = leader_arm.get_action()
     arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
-    # Keyboard
+
     keyboard_keys = keyboard.get_action()
     base_action = robot._from_keyboard_to_base_action(keyboard_keys)
 
+    log_rerun_data(observation, {**arm_action, **base_action})
+
     action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
 
-    # Send action to robot
-    _ = robot.send_action(action)
-
-    # Visualize
-    log_rerun_data(observation=observation, action=action)
+    robot.send_action(action)
 
     busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))

examples/phone_to_so100/evaluate.py

@@ -1,197 +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 lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.configs.types import FeatureType, PolicyFeature
-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.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 (
-    RobotAction,
-    RobotObservation,
-    RobotProcessorPipeline,
-    make_default_teleop_action_processor,
-)
-from lerobot.processor.converters import (
-    observation_to_transition,
-    robot_action_to_transition,
-    transition_to_observation,
-    transition_to_robot_action,
-)
-from lerobot.record import record_loop
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    ForwardKinematicsJointsToEE,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import _init_rerun
-
-NUM_EPISODES = 5
-FPS = 30
-EPISODE_TIME_SEC = 60
-TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
-
-# Create the robot configuration & robot
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
-
-robot = SO100Follower(robot_config)
-
-# Create policy
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        AddRobotObservationAsComplimentaryData(robot=robot),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert joints observation to EE observation
-robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose_processor,
-            initial_features=create_initial_features(observation=robot.observation_features),
-            use_videos=True,
-        ),
-        # User for now should be explicit on the feature keys that were used for record
-        # Alternatively, the user can pass the processor step that has the right features
-        aggregate_pipeline_dataset_features(
-            pipeline=make_default_teleop_action_processor(),
-            initial_features=create_initial_features(
-                action={
-                    f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-                    for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
-                }
-            ),
-            use_videos=True,
-        ),
-    ),
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Build Policy Processors
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
-# Connect the robot
-robot.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-_init_rerun(session_name="phone_so100_evaluate")
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting evaluate loop...")
-episode_idx = 0
-for episode_idx in range(NUM_EPISODES):
-    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        policy=policy,
-        preprocessor=preprocessor,  # Pass the pre and post policy processors
-        postprocessor=postprocessor,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=make_default_teleop_action_processor(),
-        robot_action_processor=robot_ee_to_joints_processor,
-        robot_observation_processor=robot_joints_to_ee_pose_processor,
-    )
-
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
-        log_say("Reset the environment")
-        record_loop(
-            robot=robot,
-            events=events,
-            fps=FPS,
-            control_time_s=EPISODE_TIME_SEC,
-            single_task=TASK_DESCRIPTION,
-            display_data=True,
-            teleop_action_processor=make_default_teleop_action_processor(),
-            robot_action_processor=robot_ee_to_joints_processor,
-            robot_observation_processor=robot_joints_to_ee_pose_processor,
-        )
-
-    if events["rerecord_episode"]:
-        log_say("Re-record episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
-
-    # Save episode
-    dataset.save_episode()
-    episode_idx += 1
-
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-listener.stop()
-dataset.push_to_hub()

examples/phone_to_so100/record.py

@@ -1,205 +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 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.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
-from lerobot.processor.converters import (
-    observation_to_transition,
-    robot_action_to_transition,
-    transition_to_observation,
-    transition_to_robot_action,
-)
-from lerobot.record import record_loop
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    EEBoundsAndSafety,
-    EEReferenceAndDelta,
-    ForwardKinematicsJointsToEE,
-    GripperVelocityToJoint,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
-from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.teleoperators.phone.teleop_phone import Phone
-from lerobot.utils.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import _init_rerun
-
-NUM_EPISODES = 10
-FPS = 30
-EPISODE_TIME_SEC = 60
-RESET_TIME_SEC = 30
-TASK_DESCRIPTION = "My task description"
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
-
-# Create the robot and teleoperator configurations
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
-
-# Initialize the robot and teleoperator
-robot = SO100Follower(robot_config)
-phone = Phone(teleop_config)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert phone action to EE action
-phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    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.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.20,
-            max_ee_twist_step_rad=0.50,
-        ),
-        GripperVelocityToJoint(),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert joint observation to EE observation
-robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        # Run the feature contract of the pipelines
-        # This tells you how the features would look like after the pipeline steps
-        aggregate_pipeline_dataset_features(
-            pipeline=phone_to_robot_ee_pose_processor,
-            initial_features=create_initial_features(action=phone.action_features),
-            use_videos=True,
-        ),
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose,
-            initial_features=create_initial_features(observation=robot.observation_features),
-            use_videos=True,
-        ),
-    ),
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Connect the robot and teleoperator
-robot.connect()
-phone.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-_init_rerun(session_name="phone_so100_record")
-
-if not robot.is_connected or not phone.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
-
-
-print("Starting record loop. Move your phone to teleoperate the robot...")
-episode_idx = 0
-while episode_idx < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        teleop=phone,
-        dataset=dataset,
-        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,
-        robot_observation_processor=robot_joints_to_ee_pose,
-    )
-
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        log_say("Reset the environment")
-        record_loop(
-            robot=robot,
-            events=events,
-            fps=FPS,
-            teleop=phone,
-            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,
-            robot_observation_processor=robot_joints_to_ee_pose,
-        )
-
-    if events["rerecord_episode"]:
-        log_say("Re-recording episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
-
-    # Save episode
-    dataset.save_episode()
-    episode_idx += 1
-
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-phone.disconnect()
-listener.stop()
-dataset.push_to_hub()

examples/phone_to_so100/replay.py

@@ -1,93 +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.
-
-import time
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotProcessorPipeline
-from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.utils import log_say
-
-EPISODE_IDX = 0
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
-
-# Initialize the robot config
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
-)
-
-# Initialize the robot
-robot = SO100Follower(robot_config)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        AddRobotObservationAsComplimentaryData(robot=robot),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=False,  # Because replay is open loop
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Fetch the dataset to replay
-dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
-actions = dataset.hf_dataset.select_columns("action")
-
-# Connect to the robot
-robot.connect()
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting replay loop...")
-log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(dataset.num_frames):
-    t0 = time.perf_counter()
-
-    # Get recorded action from dataset
-    ee_action = {
-        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
-    }
-
-    # Dataset EE -> robot joints
-    joint_action = robot_ee_to_joints_processor(ee_action)
-
-    # Send action to robot
-    _ = robot.send_action(joint_action)
-
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
-
-# Clean up
-robot.disconnect()

examples/phone_to_so100/teleoperate.py

@@ -1,108 +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 specif
-
-import time
-
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotProcessorPipeline
-from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    EEBoundsAndSafety,
-    EEReferenceAndDelta,
-    GripperVelocityToJoint,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
-from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
-from lerobot.teleoperators.phone.teleop_phone import Phone
-from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
-
-FPS = 30
-
-# Initialize the robot and teleoperator
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
-)
-teleop_config = PhoneConfig(phone_os=PhoneOS.IOS)  # or PhoneOS.ANDROID
-
-# Initialize the robot and teleoperator
-robot = SO100Follower(robot_config)
-teleop_device = Phone(teleop_config)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    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[RobotAction, RobotAction](
-    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.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.10,
-            max_ee_twist_step_rad=0.50,
-        ),
-        GripperVelocityToJoint(
-            speed_factor=20.0,
-        ),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Connect to the robot and teleoperator
-robot.connect()
-teleop_device.connect()
-
-# Init rerun viewer
-_init_rerun(session_name="phone_so100_teleop")
-
-if not robot.is_connected or not teleop_device.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
-
-print("Starting teleop loop. Move your phone to teleoperate the robot...")
-while True:
-    t0 = time.perf_counter()
-
-    # Get teleop action
-    phone_obs = teleop_device.get_action()
-
-    # Phone -> EE pose -> Joints transition
-    joint_action = phone_to_robot_joints_processor(phone_obs)
-
-    # Send action to robot
-    _ = robot.send_action(joint_action)
-
-    # Visualize
-    log_rerun_data(observation=phone_obs, action=joint_action)
-
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))

examples/port_datasets/agibot_hdf5/port_agibot.py

@@ -0,0 +1,503 @@
+import json
+import logging
+import shutil
+import time
+from pathlib import Path
+
+import h5py
+import numpy as np
+import pandas as pd
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    EPISODES_DIR,
+    get_video_duration_in_s,
+    get_video_size_in_mb,
+    update_chunk_file_indices,
+    write_info,
+)
+from lerobot.datasets.video_utils import concat_video_files
+from lerobot.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
+
+AGIBOT_FPS = 30
+AGIBOT_ROBOT_TYPE = "AgiBot_A2D"
+AGIBOT_FEATURES = {
+    # gripper open range in mm (0 for pull open, 1 for full close)
+    "observation.state.effector.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["left_gripper", "right_gripper"],
+        },
+    },
+    # flange xyz in meters
+    "observation.state.end.position": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "right_x", "right_y", "right_z"],
+        },
+    },
+    # flange quaternion with xyzw
+    "observation.state.end.orientation": {
+        "dtype": "float32",
+        "shape": (8,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "left_w", "right_x", "right_y", "right_z", "right_w"],
+        },
+    },
+    # in radians
+    "observation.state.head.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["yaw", "pitch"],
+        },
+    },
+    # in motor steps
+    "observation.state.joint.current_value": {
+        "dtype": "float32",
+        "shape": (14,),
+        "names": {
+            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
+        },
+    },
+    # same as current_value but in radians
+    "observation.state.joint.position": {
+        "dtype": "float32",
+        "shape": (14,),
+        "names": {
+            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
+        },
+    },
+    # pitch in radians, lift in meters
+    "observation.state.waist.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["pitch", "lift"],
+        },
+    },
+    # concatenation of head.position, joint.position, effector.position, waist.position
+    "observation.state": {
+        "dtype": "float32",
+        "shape": (20,),
+        "names": {
+            "axes": ["head_yaw", "head_pitch"]
+            + [f"left_joint_{i}" for i in range(7)]
+            + ["left_gripper"]
+            + [f"right_joint_{i}" for i in range(7)]
+            + ["right_gripper"]
+            + ["waist_pitch", "waist_lift"],
+        },
+    },
+    # gripper open range in mm (0 for pull open, 1 for full close)
+    "action.effector.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["left_gripper", "right_gripper"],
+        },
+    },
+    # flange xyz in meters
+    "action.end.position": {
+        "dtype": "float32",
+        "shape": (6,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "right_x", "right_y", "right_z"],
+        },
+    },
+    # flange quaternion with xyzw
+    "action.end.orientation": {
+        "dtype": "float32",
+        "shape": (8,),
+        "names": {
+            "axes": ["left_x", "left_y", "left_z", "left_w", "right_x", "right_y", "right_z", "right_w"],
+        },
+    },
+    # in radians
+    "action.head.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["yaw", "pitch"],
+        },
+    },
+    # goal joint position in radians
+    "action.joint.position": {
+        "dtype": "float32",
+        "shape": (14,),
+        "names": {
+            "axes": [f"left_joint_{i}" for i in range(7)] + [f"right_joint_{i}" for i in range(7)],
+        },
+    },
+    "action.robot.velocity": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["velocity_x", "yaw_rate"],
+        },
+    },
+    # pitch in radians, lift in meters
+    "action.waist.position": {
+        "dtype": "float32",
+        "shape": (2,),
+        "names": {
+            "axes": ["pitch", "lift"],
+        },
+    },
+    # concatenation of head.position, joint.position, effector.position, waist.position, robot.velocity
+    "action": {
+        "dtype": "float32",
+        "shape": (22,),
+        "names": {
+            "axes": ["head_yaw", "head_pitch"]
+            + [f"left_joint_{i}" for i in range(7)]
+            + ["left_gripper"]
+            + [f"right_joint_{i}" for i in range(7)]
+            + ["right_gripper"]
+            + ["waist_pitch", "waist_lift"]
+            + ["velocity_x", "yaw_rate"],
+        },
+    },
+    # episode level annotation
+    "init_scene_text": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    # frame level annotation
+    "action_text": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+    # frame level annotation
+    "skill": {
+        "dtype": "string",
+        "shape": (1,),
+        "names": None,
+    },
+}
+
+AGIBOT_IMAGES_FEATURES = {
+    "observation.images.top_head": {
+        "dtype": "video",
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.hand_left": {
+        "dtype": "video",
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.hand_right": {
+        "dtype": "video",
+        "shape": (480, 640, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.head_center_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.head_left_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.head_right_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.back_left_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+    "observation.images.back_right_fisheye": {
+        "dtype": "video",
+        "shape": (748, 960, 3),
+        "names": ["height", "width", "channel"],
+    },
+}
+
+
+def load_info_per_task(raw_dir):
+    info_per_task = {}
+    task_info_dir = raw_dir / "task_info"
+    for path in task_info_dir.glob("task_*.json"):
+        task_index = int(path.name.replace("task_", "").replace(".json", ""))
+        with open(path) as f:
+            task_info = json.load(f)
+
+        task_info = {ep["episode_id"]: ep for ep in task_info}
+        info_per_task[task_index] = task_info
+
+    return info_per_task
+
+
+def create_frame_idx_to_frames_label_idx(ep_info):
+    frame_idx_to_frames_label_idx = {}
+    for label_idx, frames_label in enumerate(ep_info["label_info"]["action_config"]):
+        for frame_idx in range(frames_label["start_frame"], frames_label["end_frame"]):
+            frame_idx_to_frames_label_idx[frame_idx] = label_idx
+    return frame_idx_to_frames_label_idx
+
+
+def generate_lerobot_frames(raw_dir: Path, task_index: int, episode_index: int):
+    r"""/!\ The frames dont contain observation.cameras.*"""
+    info_per_task = load_info_per_task(raw_dir)
+    ep_info = info_per_task[task_index][episode_index]
+    frame_idx_to_frames_label_idx = create_frame_idx_to_frames_label_idx(ep_info)
+
+    # Empty features are commented out.
+    keys_mapping = {
+        # STATE
+        # "observation.state.effector.force": "state/effector/force",
+        "observation.state.effector.position": "state/effector/position",
+        # "observation.state.end.angular": "state/end/angular",
+        "observation.state.end.position": "state/end/position",
+        "observation.state.end.orientation": "state/end/orientation",
+        # "observation.state.end.velocity": "state/end/velocity",
+        # "observation.state.end.wrench": "state/end/wrench",
+        # "observation.state.head.effort": "state/head/effort",
+        "observation.state.head.position": "state/head/position",
+        # "observation.state.head.velocity": "state/head/velocity",
+        "observation.state.joint.current_value": "state/joint/current_value",
+        # "observation.state.joint.effort": "state/joint/effort",
+        "observation.state.joint.position": "state/joint/position",
+        # "observation.state.joint.velocity": "state/joint/velocity",
+        # "observation.state.robot.orientation": "state/robot/orientation",
+        # "observation.state.robot.orientation_drift": "state/robot/orientation_drift",
+        # "observation.state.robot.position": "state/robot/position",
+        # "observation.state.robot.position_drift": "state/robot/position_drift",
+        # "observation.state.waist.effort": "state/waist/effort",
+        "observation.state.waist.position": "state/waist/position",
+        # "observation.state.waist.velocity": "state/waist/velocity",
+        # ----- ACTION (index are also commented out) -----
+        # "action.effector.index": "action/effector/index",
+        "action.effector.position": "action/effector/position",
+        # "action.effector.force": "action/effector/force",
+        # "action.end.index": "action/end/index",
+        "action.end.position": "action/end/position",
+        "action.end.orientation": "action/end/orientation",
+        # "action.head.index": "action/head/index",
+        "action.head.position": "action/head/position",
+        # "action.joint.index": "action/joint/index",
+        "action.joint.position": "action/joint/position",
+        # "action.joint.effort": "action/joint/effort",
+        # "action.joint.velocity": "action/joint/velocity",
+        # "action.robot.index": "action/robot/index",
+        # "action.robot.position": "action/robot/position",
+        # "action.robot.orientation": "action/robot/orientation",
+        # "action.robot.angular": "action/robot/angular",
+        "action.robot.velocity": "action/robot/velocity",
+        # "action.waist.index": "action/waist/index",
+        "action.waist.position": "action/waist/position",
+    }
+
+    h5_path = raw_dir / f"proprio_stats/{task_index}/{episode_index}/proprio_stats.h5"
+    with h5py.File(h5_path) as h5:
+        num_frames = len(h5["state/joint/position"])
+
+        for h5_key in keys_mapping.values():
+            col_num_frames = h5[h5_key].shape[0]
+            if col_num_frames != num_frames:
+                raise ValueError(
+                    f"HDF5 column '{h5_key}' is expected to have {num_frames} but has {col_num_frames}' frames instead."
+                )
+
+        for i in range(num_frames):
+            # Create frame
+            f = {new_key: h5[h5_key][i] for new_key, h5_key in keys_mapping.items()}
+
+            for key in f:
+                f[key] = np.array(f[key]).astype(np.float32)
+
+            f["observation.state.end.position"] = f["observation.state.end.position"].reshape(6)
+            f["observation.state.end.orientation"] = f["observation.state.end.orientation"].reshape(8)
+            f["observation.state"] = np.concatenate(
+                [
+                    f["observation.state.head.position"],
+                    f["observation.state.joint.position"][:7],  # left
+                    f["observation.state.effector.position"][[0]],  # left
+                    f["observation.state.joint.position"][7:],  # right
+                    f["observation.state.effector.position"][[1]],  # right
+                    f["observation.state.waist.position"],
+                ]
+            )
+
+            f["action.end.position"] = f["action.end.position"].reshape(6)
+            f["action.end.orientation"] = f["action.end.orientation"].reshape(8)
+            f["action"] = np.concatenate(
+                [
+                    f["action.head.position"],
+                    f["action.joint.position"][:7],  # left
+                    f["action.effector.position"][[0]],  # left
+                    f["action.joint.position"][7:],  # right
+                    f["action.effector.position"][[1]],  # right
+                    f["action.waist.position"],
+                    f["action.robot.velocity"],
+                ]
+            )
+
+            # episode level annotation
+            f["task"] = ep_info["task_name"]
+            f["init_scene_text"] = ep_info["init_scene_text"]
+
+            # frame level annotation
+            if i in frame_idx_to_frames_label_idx:
+                frames_label_idx = frame_idx_to_frames_label_idx[i]
+                frames_label = ep_info["label_info"]["action_config"][frames_label_idx]
+                f["action_text"] = frames_label["action_text"]
+                f["skill"] = frames_label["skill"]
+            else:
+                f["action_text"] = ""
+                f["skill"] = ""
+
+            yield f
+
+
+def update_meta_data(
+    df,
+    ep_to_meta,
+):
+    def _update(row):
+        ep_idx = row["episode_index"]
+        for key, meta in ep_to_meta[ep_idx].items():
+            row[f"videos/{key}/chunk_index"] = meta["chunk_index"]
+            row[f"videos/{key}/file_index"] = meta["file_index"]
+            row[f"videos/{key}/from_timestamp"] = meta["from_timestamp"]
+            row[f"videos/{key}/to_timestamp"] = meta["to_timestamp"]
+        return row
+
+    return df.apply(_update, axis=1)
+
+
+def move_videos_to_lerobot_directory(lerobot_dataset, raw_dir, task_index, episode_names):
+    keys_mapping = {
+        "observation.images.top_head": "head_color",
+        "observation.images.hand_left": "hand_left_color",
+        "observation.images.hand_right": "hand_right_color",
+        "observation.images.head_center_fisheye": "head_center_fisheye_color",
+        "observation.images.head_left_fisheye": "head_left_fisheye_color",
+        "observation.images.head_right_fisheye": "head_right_fisheye_color",
+        "observation.images.back_left_fisheye": "back_left_fisheye_color",
+        "observation.images.back_right_fisheye": "back_right_fisheye_color",
+    }
+
+    # sanity check
+    for key in keys_mapping:
+        if key not in lerobot_dataset.meta.info["features"]:
+            raise ValueError(f"Key '{key}' not found in features.")
+
+    video_keys = keys_mapping.keys()
+    chunk_idx = dict.fromkeys(video_keys, 0)
+    file_idx = dict.fromkeys(video_keys, 0)
+    latest_duration_in_s = dict.fromkeys(video_keys, 0)
+    ep_to_meta = {}
+    for ep_idx, ep_name in enumerate(episode_names):
+        for key in video_keys:
+            raw_videos_dir = raw_dir / f"observations/{task_index}/{ep_name}/videos"
+            old_key = keys_mapping[key]
+            ep_path = raw_videos_dir / f"{old_key}.mp4"
+            ep_duration_in_s = get_video_duration_in_s(ep_path)
+
+            aggr_path = lerobot_dataset.root / DEFAULT_VIDEO_PATH.format(
+                video_key=key,
+                chunk_index=chunk_idx[key],
+                file_index=file_idx[key],
+            )
+            if not aggr_path.exists():
+                # First video
+                aggr_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy(str(ep_path), str(aggr_path))
+            else:
+                size_in_mb = get_video_size_in_mb(ep_path)
+                aggr_size_in_mb = get_video_size_in_mb(aggr_path)
+
+                if aggr_size_in_mb + size_in_mb >= DEFAULT_VIDEO_FILE_SIZE_IN_MB:
+                    # Size limit is reached, prepare new parquet file
+                    chunk_idx[key], file_idx[key] = update_chunk_file_indices(
+                        chunk_idx[key], file_idx[key], DEFAULT_CHUNK_SIZE
+                    )
+                    aggr_path = lerobot_dataset.root / DEFAULT_VIDEO_PATH.format(
+                        video_key=key,
+                        chunk_index=chunk_idx[key],
+                        file_index=file_idx[key],
+                    )
+                    aggr_path.parent.mkdir(parents=True, exist_ok=True)
+                    shutil.copy(str(ep_path), str(aggr_path))
+                    latest_duration_in_s[key] = 0
+                else:
+                    # Update the existing parquet file with new rows
+                    concat_video_files(
+                        [aggr_path, ep_path],
+                        lerobot_dataset.root,
+                        key,
+                        chunk_idx[key],
+                        file_idx[key],
+                    )
+
+            if ep_idx not in ep_to_meta:
+                ep_to_meta[ep_idx] = {}
+            ep_to_meta[ep_idx][key] = {
+                "chunk_index": chunk_idx[key],
+                "file_index": file_idx[key],
+                "from_timestamp": latest_duration_in_s[key],
+                "to_timestamp": latest_duration_in_s[key] + ep_duration_in_s,
+            }
+            latest_duration_in_s[key] += ep_duration_in_s
+
+    # Update episodes meta data
+    for meta_path in (lerobot_dataset.root / EPISODES_DIR).glob("chunk-*/file-*.parquet"):
+        df = pd.read_parquet(meta_path)
+        df = update_meta_data(df, ep_to_meta)
+        df.to_parquet(meta_path)
+
+
+def port_agibot(
+    raw_dir: Path, repo_id: str, task_index: int, episode_indices: list[int], push_to_hub: bool = False
+):
+    lerobot_dataset = LeRobotDataset.create(
+        repo_id=repo_id,
+        robot_type=AGIBOT_ROBOT_TYPE,
+        fps=AGIBOT_FPS,
+        features=AGIBOT_FEATURES,
+    )
+
+    start_time = time.time()
+    num_episodes = len(episode_indices)
+    logging.info(f"Number of episodes {num_episodes}")
+
+    for i, episode_index in enumerate(episode_indices):
+        elapsed_time = time.time() - start_time
+        d, h, m, s = get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time)
+
+        logging.info(
+            f"{i} / {num_episodes} episodes processed (after {d} days, {h} hours, {m} minutes, {s:.3f} seconds)"
+        )
+
+        for frame in generate_lerobot_frames(raw_dir, task_index, episode_index):
+            lerobot_dataset.add_frame(frame)
+
+        lerobot_dataset.save_episode()
+        logging.info("Save_episode")
+
+    # Videos have already been encoded with the proper format, so we rely on hacks
+    # HACK: Add extra images features
+    lerobot_dataset.meta.info["features"].update(AGIBOT_IMAGES_FEATURES)
+    write_info(lerobot_dataset.meta.info, lerobot_dataset.meta.root)
+    move_videos_to_lerobot_directory(lerobot_dataset, raw_dir, task_index, episode_indices)
+
+    if push_to_hub:
+        lerobot_dataset.push_to_hub(
+            # Add agibot tag, since it belongs to the agibot collection of datasets
+            tags=["agibot"],
+            private=False,
+        )

examples/port_datasets/agibot_hdf5/slurm_port_shards.py

@@ -0,0 +1,198 @@
+#!/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.
+
+import argparse
+import logging
+import tarfile
+from pathlib import Path
+
+from datatrove.executor import LocalPipelineExecutor
+from datatrove.executor.slurm import SlurmPipelineExecutor
+from datatrove.pipeline.base import PipelineStep
+from port_datasets.agibot_hdf5.download import (
+    RAW_REPO_ID,
+    download_meta_data,
+    get_observations_files,
+)
+
+
+class PortAgiBotShards(PipelineStep):
+    def __init__(
+        self,
+        raw_dir: Path | str,
+        repo_id: str = None,
+    ):
+        super().__init__()
+        self.raw_dir = Path(raw_dir)
+        self.repo_id = repo_id
+
+    def run(self, data=None, rank: int = 0, world_size: int = 1):
+        import shutil
+
+        from datasets.utils.tqdm import disable_progress_bars
+        from port_datasets.agibot_hdf5.download import (
+            RAW_REPO_ID,
+            download,
+            get_observations_files,
+            no_depth,
+        )
+        from port_datasets.agibot_hdf5.port_agibot import port_agibot
+        from port_datasets.droid_rlds.port_droid import validate_dataset
+
+        from lerobot.constants import HF_LEROBOT_HOME
+        from lerobot.utils.utils import init_logging
+
+        init_logging()
+        disable_progress_bars()
+
+        shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
+
+        dataset_dir = HF_LEROBOT_HOME / shard_repo_id
+        if dataset_dir.exists():
+            shutil.rmtree(dataset_dir)
+
+        obs_files, _ = get_observations_files(self.raw_dir, RAW_REPO_ID)
+        obs_file = obs_files[rank]
+
+        # Download subset
+        download(self.raw_dir, allow_patterns=obs_file)
+
+        tar_path = self.raw_dir / obs_file
+        with tarfile.open(tar_path, "r") as tar:
+            extracted_files = tar.getnames()
+
+        task_index = int(tar_path.parent.name)
+        episode_names = [int(p) for p in extracted_files if "/" not in p]
+
+        # Untar if needed
+        if not all((tar_path.parent / f"{ep_name}").exists() for ep_name in episode_names):
+            logging.info(f"Untar-ing {tar_path}...")
+            with tarfile.open(tar_path, "r") as tar:
+                tar.extractall(path=tar_path.parent, filter=no_depth)  # nosec B202
+
+        port_agibot(self.raw_dir, shard_repo_id, task_index, episode_names, push_to_hub=False)
+
+        for ep_name in episode_names:
+            shutil.rmtree(str(tar_path.parent / f"{ep_name}"))
+
+        tar_path.unlink()
+
+        validate_dataset(shard_repo_id)
+
+
+def make_port_executor(
+    raw_dir, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
+):
+    download_meta_data(raw_dir)
+    obs_files, _ = get_observations_files(raw_dir, RAW_REPO_ID)
+    num_shards = len(obs_files)
+
+    kwargs = {
+        "pipeline": [
+            PortAgiBotShards(raw_dir, repo_id),
+        ],
+        "logging_dir": str(logs_dir / job_name),
+    }
+
+    if slurm:
+        kwargs.update(
+            {
+                "job_name": job_name,
+                "tasks": num_shards,
+                "workers": workers,
+                "time": "08:00:00",
+                "partition": partition,
+                "cpus_per_task": cpus_per_task,
+                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
+            }
+        )
+        executor = SlurmPipelineExecutor(**kwargs)
+    else:
+        kwargs.update(
+            {
+                "tasks": num_shards,
+                "workers": 1,
+            }
+        )
+        executor = LocalPipelineExecutor(**kwargs)
+
+    return executor
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--raw-dir",
+        type=Path,
+        required=True,
+        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=Path,
+        help="Path to logs directory for `datatrove`.",
+    )
+    parser.add_argument(
+        "--job-name",
+        type=str,
+        default="port_droid",
+        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
+    )
+    parser.add_argument(
+        "--slurm",
+        type=int,
+        default=1,
+        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=2048,
+        help="Number of slurm workers. It should be less than the maximum number of shards.",
+    )
+    parser.add_argument(
+        "--partition",
+        type=str,
+        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
+    )
+    parser.add_argument(
+        "--cpus-per-task",
+        type=int,
+        default=8,
+        help="Number of cpus that each slurm worker will use.",
+    )
+    parser.add_argument(
+        "--mem-per-cpu",
+        type=str,
+        default="1950M",
+        help="Memory per cpu that each worker will use.",
+    )
+
+    args = parser.parse_args()
+    kwargs = vars(args)
+    kwargs["slurm"] = kwargs.pop("slurm") == 1
+    port_executor = make_port_executor(**kwargs)
+    port_executor.run()
+
+
+if __name__ == "__main__":
+    main()

examples/so100_to_so100_EE/evaluate.py

@@ -1,198 +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 lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.configs.types import FeatureType, PolicyFeature
-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.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 (
-    RobotAction,
-    RobotObservation,
-    RobotProcessorPipeline,
-    make_default_teleop_action_processor,
-)
-from lerobot.processor.converters import (
-    observation_to_transition,
-    robot_action_to_transition,
-    transition_to_observation,
-    transition_to_robot_action,
-)
-from lerobot.record import record_loop
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    ForwardKinematicsJointsToEE,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import _init_rerun
-
-NUM_EPISODES = 5
-FPS = 30
-EPISODE_TIME_SEC = 60
-TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
-
-# Create the robot configuration & robot
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411",
-    id="my_awesome_follower_arm",
-    cameras=camera_config,
-    use_degrees=True,
-)
-
-robot = SO100Follower(robot_config)
-
-# Create policy
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        AddRobotObservationAsComplimentaryData(robot=robot),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=True,
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert joints observation to EE observation
-robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        aggregate_pipeline_dataset_features(
-            pipeline=robot_joints_to_ee_pose_processor,
-            initial_features=create_initial_features(observation=robot.observation_features),
-            use_videos=True,
-        ),
-        # User for now should be explicit on the feature keys that were used for record
-        # Alternatively, the user can pass the processor step that has the right features
-        aggregate_pipeline_dataset_features(
-            pipeline=make_default_teleop_action_processor(),
-            initial_features=create_initial_features(
-                action={
-                    f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
-                    for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
-                }
-            ),
-            use_videos=True,
-        ),
-    ),
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Build Policy Processors
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
-# Connect the robot and teleoperator
-robot.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-_init_rerun(session_name="so100_so100_evaluate")
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting evaluate loop...")
-episode_idx = 0
-for episode_idx in range(NUM_EPISODES):
-    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        policy=policy,
-        preprocessor=preprocessor,  # Pass the pre and post policy processors
-        postprocessor=postprocessor,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=make_default_teleop_action_processor(),
-        robot_action_processor=robot_ee_to_joints_processor,
-        robot_observation_processor=robot_joints_to_ee_pose_processor,
-    )
-
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and ((episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]):
-        log_say("Reset the environment")
-        record_loop(
-            robot=robot,
-            events=events,
-            fps=FPS,
-            control_time_s=EPISODE_TIME_SEC,
-            single_task=TASK_DESCRIPTION,
-            display_data=True,
-            teleop_action_processor=make_default_teleop_action_processor(),
-            robot_action_processor=robot_ee_to_joints_processor,
-            robot_observation_processor=robot_joints_to_ee_pose_processor,
-        )
-
-    if events["rerecord_episode"]:
-        log_say("Re-record episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
-
-    # Save episode
-    dataset.save_episode()
-    episode_idx += 1
-
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-listener.stop()
-dataset.push_to_hub()

examples/so100_to_so100_EE/record.py

@@ -1,204 +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 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.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
-from lerobot.processor.converters import (
-    observation_to_transition,
-    robot_action_to_transition,
-    transition_to_observation,
-    transition_to_robot_action,
-)
-from lerobot.record import record_loop
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    EEBoundsAndSafety,
-    ForwardKinematicsJointsToEE,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
-from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
-from lerobot.utils.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import _init_rerun
-
-NUM_EPISODES = 10
-FPS = 30
-EPISODE_TIME_SEC = 60
-RESET_TIME_SEC = 30
-TASK_DESCRIPTION = "My task description"
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
-
-# Create the robot and teleoperator configurations
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-follower_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", cameras=camera_config, use_degrees=True
-)
-leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
-
-# Initialize the robot and teleoperator
-follower = SO100Follower(follower_config)
-leader = SO100Leader(leader_config)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-follower_kinematics_solver = RobotKinematics(
-    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(follower.bus.motors.keys()),
-)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-leader_kinematics_solver = RobotKinematics(
-    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(leader.bus.motors.keys()),
-)
-
-# Build pipeline to convert follower joints to EE observation
-follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation](
-    steps=[
-        ForwardKinematicsJointsToEE(
-            kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys())
-        ),
-    ],
-    to_transition=observation_to_transition,
-    to_output=transition_to_observation,
-)
-
-# Build pipeline to convert leader joints to EE action
-leader_joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        ForwardKinematicsJointsToEE(
-            kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Build pipeline to convert EE action to follower joints
-ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
-    [
-        AddRobotObservationAsComplimentaryData(robot=follower),
-        EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.10,
-            max_ee_twist_step_rad=0.50,
-        ),
-        InverseKinematicsEEToJoints(
-            kinematics=follower_kinematics_solver,
-            motor_names=list(follower.bus.motors.keys()),
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_REPO_ID,
-    fps=FPS,
-    features=combine_feature_dicts(
-        # Run the feature contract of the pipelines
-        # This tells you how the features would look like after the pipeline steps
-        aggregate_pipeline_dataset_features(
-            pipeline=leader_joints_to_ee,
-            initial_features=create_initial_features(action=leader.action_features),
-            use_videos=True,
-        ),
-        aggregate_pipeline_dataset_features(
-            pipeline=follower_joints_to_ee,
-            initial_features=create_initial_features(observation=follower.observation_features),
-            use_videos=True,
-        ),
-    ),
-    robot_type=follower.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-
-# Connect the robot and teleoperator
-leader.connect()
-follower.connect()
-
-# Initialize the keyboard listener and rerun visualization
-listener, events = init_keyboard_listener()
-_init_rerun(session_name="recording_phone")
-
-if not leader.is_connected or not follower.is_connected:
-    raise ValueError("Robot or teleop is not connected!")
-
-print("Starting record loop...")
-episode_idx = 0
-while episode_idx < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Main record loop
-    record_loop(
-        robot=follower,
-        events=events,
-        fps=FPS,
-        teleop=leader,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-        teleop_action_processor=leader_joints_to_ee,
-        robot_action_processor=ee_to_follower_joints,
-        robot_observation_processor=follower_joints_to_ee,
-    )
-
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        log_say("Reset the environment")
-        record_loop(
-            robot=follower,
-            events=events,
-            fps=FPS,
-            teleop=leader,
-            control_time_s=RESET_TIME_SEC,
-            single_task=TASK_DESCRIPTION,
-            display_data=True,
-            teleop_action_processor=leader_joints_to_ee,
-            robot_action_processor=ee_to_follower_joints,
-            robot_observation_processor=follower_joints_to_ee,
-        )
-
-    if events["rerecord_episode"]:
-        log_say("Re-recording episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
-
-    # Save episode
-    dataset.save_episode()
-    episode_idx += 1
-
-# Clean up
-log_say("Stop recording")
-leader.disconnect()
-follower.disconnect()
-listener.stop()
-dataset.push_to_hub()

examples/so100_to_so100_EE/replay.py

@@ -1,94 +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.
-
-
-import time
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotProcessorPipeline
-from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.utils import log_say
-
-EPISODE_IDX = 0
-HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
-
-# Initialize the robot config
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
-
-# Initialize the robot
-robot = SO100Follower(robot_config)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-kinematics_solver = RobotKinematics(
-    urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(robot.bus.motors.keys()),
-)
-
-# Build pipeline to convert EE action to joints action
-robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        AddRobotObservationAsComplimentaryData(robot=robot),
-        InverseKinematicsEEToJoints(
-            kinematics=kinematics_solver,
-            motor_names=list(robot.bus.motors.keys()),
-            initial_guess_current_joints=False,  # Because replay is open loop
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Fetch the dataset to replay
-dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
-actions = dataset.hf_dataset.select_columns("action")
-
-# Connect to the robot
-robot.connect()
-
-if not robot.is_connected:
-    raise ValueError("Robot is not connected!")
-
-print("Starting replay loop...")
-log_say(f"Replaying episode {EPISODE_IDX}")
-for idx in range(dataset.num_frames):
-    t0 = time.perf_counter()
-
-    # Get recorded action from dataset
-    ee_action = {
-        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
-    }
-
-    # Dataset EE -> robot joints
-    joint_action = robot_ee_to_joints_processor(ee_action)
-
-    # Send action to robot
-    _ = robot.send_action(joint_action)
-
-    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
-
-# Clean up
-robot.disconnect()

examples/so100_to_so100_EE/teleoperate.py

@@ -1,119 +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.
-
-import time
-
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotProcessorPipeline
-from lerobot.processor.converters import (
-    robot_action_to_transition,
-    transition_to_robot_action,
-)
-from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
-from lerobot.robots.so100_follower.robot_kinematic_processor import (
-    AddRobotObservationAsComplimentaryData,
-    EEBoundsAndSafety,
-    ForwardKinematicsJointsToEE,
-    InverseKinematicsEEToJoints,
-)
-from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
-from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
-from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
-
-FPS = 30
-
-# Initialize the robot and teleoperator config
-follower_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
-)
-leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
-
-# Initialize the robot and teleoperator
-follower = SO100Follower(follower_config)
-leader = SO100Leader(leader_config)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-follower_kinematics_solver = RobotKinematics(
-    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(follower.bus.motors.keys()),
-)
-
-# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
-leader_kinematics_solver = RobotKinematics(
-    urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
-    target_frame_name="gripper_frame_link",
-    joint_names=list(leader.bus.motors.keys()),
-)
-
-# Build pipeline to convert teleop joints to EE action
-leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
-    steps=[
-        ForwardKinematicsJointsToEE(
-            kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# build pipeline to convert EE action to robot joints
-ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
-    [
-        AddRobotObservationAsComplimentaryData(robot=follower),
-        EEBoundsAndSafety(
-            end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
-            max_ee_step_m=0.10,
-            max_ee_twist_step_rad=0.50,
-        ),
-        InverseKinematicsEEToJoints(
-            kinematics=follower_kinematics_solver,
-            motor_names=list(follower.bus.motors.keys()),
-        ),
-    ],
-    to_transition=robot_action_to_transition,
-    to_output=transition_to_robot_action,
-)
-
-# Connect to the robot and teleoperator
-follower.connect()
-leader.connect()
-
-# Init rerun viewer
-_init_rerun(session_name="so100_so100_EE_teleop")
-
-print("Starting teleop loop...")
-while True:
-    t0 = time.perf_counter()
-
-    # Get teleop observation
-    leader_joints_obs = leader.get_action()
-
-    # teleop joints -> teleop EE action
-    leader_ee_act = leader_to_ee(leader_joints_obs)
-
-    # teleop EE -> robot joints
-    follower_joints_act = ee_to_follower_joints(leader_ee_act)
-
-    # Send action to robot
-    _ = follower.send_action(follower_joints_act)
-
-    # Visualize
-    log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
-
-    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))

examples/test.sh

@@ -1,58 +0,0 @@
-#!/bin/bash
-
-# storage / caches
-RAID=/raid/jade
-export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
-export HF_HOME=$RAID/.cache/huggingface
-export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
-export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
-export WANDB_CACHE_DIR=$RAID/.cache/wandb
-export TMPDIR=$RAID/.cache/tmp
-mkdir -p $TMPDIR
-export WANDB_MODE=offline
-export HF_DATASETS_OFFLINE=1
-export HF_HUB_OFFLINE=1
-export TOKENIZERS_PARALLELISM=false
-export MUJOCO_GL=egl
-export CUDA_VISIBLE_DEVICES=2
-
-# CONFIGURATION
-POLICY_PATH="/raid/jade/logs/lerobot/lerobot_2_HuggingFaceVLA_libero_smolvla_lr1e-4bs32steps100000/checkpoints/100000/pretrained_model"
-POLICY_PATH="/raid/jade/models/smolvla_pipe"
-TASK=libero_spatial
-ENV_TYPE="libero"
-BATCH_SIZE=1
-N_EPISODES=1
-# storage / caches
-RAID=/raid/jade
-N_ACTION_STEPS=1
-export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
-export HF_HOME=$RAID/.cache/huggingface
-export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
-export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
-export WANDB_CACHE_DIR=$RAID/.cache/wandb
-export TMPDIR=$RAID/.cache/tmp
-mkdir -p $TMPDIR
-export WANDB_MODE=offline
-# export HF_DATASETS_OFFLINE=1
-# export HF_HUB_OFFLINE=1
-export TOKENIZERS_PARALLELISM=false
-export MUJOCO_GL=egl
-export MUJOCO_GL=egl
-unset HF_HUB_OFFLINE
-# RUN EVALUATION
-python src/lerobot/scripts/eval.py \
-    --policy.path="$POLICY_PATH" \
-    --env.type="$ENV_TYPE" \
-    --eval.batch_size="$BATCH_SIZE" \
-    --eval.n_episodes="$N_EPISODES" \
-    --env.task=$TASK \
-    --env.max_parallel_tasks=10 \
-# python examples/evaluate_libero.py \
-#     --policy_path "$POLICY_PATH" \
-#     --task_suite_name "$TASK" \
-#     --num_steps_wait 10 \
-#     --num_trials_per_task 10 \
-#     --video_out_path "data/libero/videos" \
-#     --device "cuda" \
-#     --seed 7

examples/use_dataset_tools.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+"""
+Example script demonstrating dataset tools utilities.
+
+This script shows how to:
+1. Delete episodes from a dataset
+2. Split a dataset into train/val sets
+3. Add/remove features
+4. Merge datasets
+
+Usage:
+    python examples/use_dataset_tools.py
+"""
+
+import numpy as np
+
+from lerobot.datasets.dataset_tools import (
+    add_feature,
+    delete_episodes,
+    merge_datasets,
+    remove_feature,
+    split_dataset,
+)
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+
+def main():
+    # Load an existing dataset (replace with your dataset)
+    dataset = LeRobotDataset("lerobot/pusht")
+
+    print(f"Original dataset: {dataset.meta.total_episodes} episodes, {dataset.meta.total_frames} frames")
+    print(f"Features: {list(dataset.meta.features.keys())}")
+
+    # Example 1: Delete episodes
+    print("\n1. Deleting episodes 0 and 2...")
+    filtered_dataset = delete_episodes(dataset, episode_indices=[0, 2], repo_id="pusht_filtered")
+    print(f"Filtered dataset: {filtered_dataset.meta.total_episodes} episodes")
+
+    # Example 2: Split dataset
+    print("\n2. Splitting dataset into train/val...")
+    splits = split_dataset(
+        dataset,
+        splits={"train": 0.8, "val": 0.2},
+    )
+    print(f"Train split: {splits['train'].meta.total_episodes} episodes")
+    print(f"Val split: {splits['val'].meta.total_episodes} episodes")
+
+    # Example 3: Add a feature
+    print("\n3. Adding a reward feature...")
+
+    # Method 1: Pre-computed values
+    reward_values = np.random.randn(dataset.meta.total_frames).astype(np.float32)
+    dataset_with_reward = add_feature(
+        dataset,
+        feature_name="reward",
+        feature_values=reward_values,
+        feature_info={
+            "dtype": "float32",
+            "shape": (1,),
+            "names": None,
+        },
+        repo_id="pusht_with_reward",
+    )
+
+    # Method 2: Using a callable
+    def compute_success(frame_dict, episode_idx, frame_idx):
+        # Example: mark last 10 frames of each episode as successful
+        episode_length = 10  # You'd get this from episode metadata
+        return float(frame_idx >= episode_length - 10)
+
+    dataset_with_success = add_feature(
+        dataset_with_reward,
+        feature_name="success",
+        feature_values=compute_success,
+        feature_info={
+            "dtype": "float32",
+            "shape": (1,),
+            "names": None,
+        },
+        repo_id="pusht_with_reward_and_success",
+    )
+
+    print(f"New features: {list(dataset_with_success.meta.features.keys())}")
+
+    # Example 4: Remove features
+    print("\n4. Removing the success feature...")
+    dataset_cleaned = remove_feature(dataset_with_success, feature_names="success", repo_id="pusht_cleaned")
+    print(f"Features after removal: {list(dataset_cleaned.meta.features.keys())}")
+
+    # Example 5: Merge datasets
+    print("\n5. Merging train and val splits back together...")
+    merged = merge_datasets([splits["train"], splits["val"]], output_repo_id="pusht_merged")
+    print(f"Merged dataset: {merged.meta.total_episodes} episodes")
+
+    # Example 6: Complex workflow
+    print("\n6. Complex workflow example...")
+
+    # Remove a camera if dataset has multiple
+    if len(dataset.meta.camera_keys) > 1:
+        camera_to_remove = dataset.meta.camera_keys[0]
+        print(f"Removing camera: {camera_to_remove}")
+        dataset_no_cam = remove_feature(
+            dataset, feature_names=camera_to_remove, repo_id="pusht_no_first_camera"
+        )
+        print(f"Remaining cameras: {dataset_no_cam.meta.camera_keys}")
+
+    print("\nDone! Check ~/.cache/huggingface/lerobot/ for the created datasets.")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -84,6 +84,7 @@ dependencies = [
 
     # Support dependencies
     "deepdiff>=7.0.1,<9.0.0",
+    "flask>=3.0.3,<4.0.0",
     "imageio[ffmpeg]>=2.34.0,<3.0.0",
     "termcolor>=2.4.0,<4.0.0",
 ]
@@ -94,7 +95,7 @@ dependencies = [
 # Common
 pygame-dep = ["pygame>=2.5.1"]
 placo-dep = ["placo>=0.9.6"]
-transformers-dep = ["transformers>=4.52.0"]
+transformers-dep = ["transformers>=4.50.3,<4.52.0"] # TODO: Bumb dependency
 grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"]
 
 # Motors
@@ -105,13 +106,11 @@ 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'",
     "pyrealsense2-macosx>=2.54 ; sys_platform == 'darwin'",
 ]
-phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
 # stretch = [
 #     "hello-robot-stretch-body>=0.7.27 ; sys_platform == 'linux'",
 #     "pyrender @ git+https://github.com/mmatl/pyrender.git ; sys_platform == 'linux'",
@@ -135,28 +134,13 @@ video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
 aloha = ["gym-aloha>=0.1.1"]
 pusht = ["gym-pusht>=0.1.5", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
 xarm = ["gym-xarm>=0.1.1"]
-libero = [
-    "hydra-core>=1.2,<1.4",
-    "easydict>=1.9",
-    "lerobot[transformers-dep]",
-    "robomimic==0.2.0",
-    "thop>=0.1.0.post2206102148",
-    "robosuite==1.4.0",
-    "bddl==1.0.1",
-    "matplotlib>=3.5.3",
-    "cloudpickle>=2.0.0",
-    "gym>=0.25,<0.27",
-    "future>=0.18.3",
-    "egl_probe @ git+https://github.com/jadechoghari/egl_probe.git#egg=egl_probe",
-    "libero @ git+https://github.com/jadechoghari/LIBERO.git@main#egg=libero",
-]
+
 # All
 all = [
     "lerobot[dynamixel]",
     "lerobot[gamepad]",
     "lerobot[hopejr]",
     "lerobot[lekiwi]",
-    "lerobot[reachy2]",
     "lerobot[kinematics]",
     "lerobot[intelrealsense]",
     "lerobot[pi0]",
@@ -168,9 +152,7 @@ all = [
     "lerobot[video_benchmark]",
     "lerobot[aloha]",
     "lerobot[pusht]",
-    "lerobot[xarm]",
-    "lerobot[phone]",
-    "lerobot[libero]"
+    "lerobot[xarm]"
 ]
 
 [project.scripts]

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/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

@@ -37,7 +37,6 @@ class DatasetConfig:
     revision: str | None = None
     use_imagenet_stats: bool = True
     video_backend: str = field(default_factory=get_safe_default_codec)
-    streaming: bool = False
 
 
 @dataclass

src/lerobot/configs/policies.py

@@ -26,7 +26,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.constants import ACTION, OBS_STATE
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
@@ -53,6 +53,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     """
 
     n_obs_steps: int = 1
+    normalization_mapping: dict[str, NormalizationMode] = field(default_factory=dict)
 
     input_features: dict[str, PolicyFeature] = field(default_factory=dict)
     output_features: dict[str, PolicyFeature] = field(default_factory=dict)

src/lerobot/configs/types.py

@@ -24,12 +24,6 @@ class FeatureType(str, Enum):
     ENV = "ENV"
     ACTION = "ACTION"
     REWARD = "REWARD"
-    LANGUAGE = "LANGUAGE"
-
-
-class PipelineFeatureType(str, Enum):
-    ACTION = "ACTION"
-    OBSERVATION = "OBSERVATION"
 
 
 class NormalizationMode(str, Enum):

src/lerobot/constants.py

@@ -21,14 +21,8 @@ OBS_ENV_STATE = "observation.environment_state"
 OBS_STATE = "observation.state"
 OBS_IMAGE = "observation.image"
 OBS_IMAGES = "observation.images"
-OBS_LANGUAGE = "observation.language"
 ACTION = "action"
 REWARD = "next.reward"
-TRUNCATED = "next.truncated"
-DONE = "next.done"
-
-OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
-OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
 
 ROBOTS = "robots"
 ROBOT_TYPE = "robot_type"
@@ -45,9 +39,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"
@@ -61,8 +52,3 @@ HF_LEROBOT_HOME = Path(os.getenv("HF_LEROBOT_HOME", default_cache_path)).expandu
 # calibration dir
 default_calibration_path = HF_LEROBOT_HOME / "calibration"
 HF_LEROBOT_CALIBRATION = Path(os.getenv("HF_LEROBOT_CALIBRATION", default_calibration_path)).expanduser()
-
-
-# streaming datasets
-LOOKBACK_BACKTRACKTABLE = 100
-LOOKAHEAD_BACKTRACKTABLE = 100

src/lerobot/datasets/aggregate.py

@@ -39,7 +39,7 @@ from lerobot.datasets.utils import (
     write_stats,
     write_tasks,
 )
-from lerobot.datasets.video_utils import concatenate_video_files
+from lerobot.datasets.video_utils import concat_video_files
 
 
 def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
@@ -152,11 +152,11 @@ def update_meta_data(
 def aggregate_datasets(
     repo_ids: list[str],
     aggr_repo_id: str,
-    roots: list[Path] | None = None,
-    aggr_root: Path | None = None,
-    data_files_size_in_mb: float | None = None,
-    video_files_size_in_mb: float | None = None,
-    chunk_size: int | None = None,
+    roots: list[Path] = None,
+    aggr_root: Path = None,
+    data_files_size_in_mb: float = None,
+    video_files_size_in_mb: float = None,
+    chunk_size: int = None,
 ):
     """Aggregates multiple LeRobot datasets into a single unified dataset.
 
@@ -298,9 +298,12 @@ def aggregate_videos(src_meta, dst_meta, videos_idx, video_files_size_in_mb, chu
                 timestamps_shift_s = dst_meta.info["total_frames"] / dst_meta.info["fps"]
 
                 # Append to existing video file
-                concatenate_video_files(
+                concat_video_files(
                     [dst_path, src_path],
-                    dst_path,
+                    dst_meta.root,
+                    key,
+                    chunk_idx,
+                    file_idx,
                 )
                 # Update the latest_duration when appending (shifts timestamps!)
                 update_latest_duration = not update_latest_duration

src/lerobot/datasets/backward_compatibility.py

@@ -14,11 +14,43 @@
 
 import packaging.version
 
-V30_MESSAGE = """
+V2_MESSAGE = """
 The dataset you requested ({repo_id}) is in {version} format.
 
-We introduced a new format since v3.0 which is not backward compatible with v2.1.
-Please, update your dataset to the new format using this command:
+We introduced a new format since v2.0 which is not backward compatible with v1.x.
+Please, use our conversion script. Modify the following command with your own task description:
+```
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \\
+    --repo-id {repo_id} \\
+    --single-task "TASK DESCRIPTION."  # <---- /!\\ Replace TASK DESCRIPTION /!\\
+```
+
+A few examples to replace TASK DESCRIPTION: "Pick up the blue cube and place it into the bin.", "Insert the
+peg into the socket.", "Slide open the ziploc bag.", "Take the elevator to the 1st floor.", "Open the top
+cabinet, store the pot inside it then close the cabinet.", "Push the T-shaped block onto the T-shaped
+target.", "Grab the spray paint on the shelf and place it in the bin on top of the robot dog.", "Fold the
+sweatshirt.", ...
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+"""
+
+V21_MESSAGE = """
+The dataset you requested ({repo_id}) is in {version} format.
+While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
+stats instead of per-episode stats. Update your dataset stats to the new format using this command:
+```
+python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id={repo_id}
+```
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+"""
+
+V30_MESSAGE = """
+The dataset you requested ({repo_id}) is in {version} format.
+While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
+stats instead of per-episode stats. Update your dataset stats to the new format using this command:
 ```
 python -m lerobot.datasets.v30.convert_dataset_v21_to_v30 --repo-id={repo_id}
 ```
@@ -38,8 +70,10 @@ class CompatibilityError(Exception): ...
 
 class BackwardCompatibilityError(CompatibilityError):
     def __init__(self, repo_id: str, version: packaging.version.Version):
-        if version.major == 2 and version.minor == 1:
+        if version.major == 3:
             message = V30_MESSAGE.format(repo_id=repo_id, version=version)
+        elif version.major == 2:
+            message = V2_MESSAGE.format(repo_id=repo_id, version=version)
         else:
             raise NotImplementedError(
                 "Contact the maintainer on [Discord](https://discord.com/invite/s3KuuzsPFb)."

src/lerobot/datasets/dataset_tools.py

@@ -0,0 +1,761 @@
+#!/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.
+"""Dataset tools utilities for LeRobotDataset.
+
+This module provides utilities for:
+- Deleting episodes from datasets
+- Splitting datasets into multiple smaller datasets
+- Adding/removing features from datasets
+- Merging datasets (wrapper around aggregate functionality)
+"""
+
+import logging
+import shutil
+from collections.abc import Callable
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from tqdm import tqdm
+
+from lerobot.constants import HF_LEROBOT_HOME
+from lerobot.datasets.aggregate import aggregate_datasets
+from lerobot.datasets.compute_stats import aggregate_stats
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_DATA_PATH,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    get_parquet_file_size_in_mb,
+    get_video_size_in_mb,
+    to_parquet_with_hf_images,
+    update_chunk_file_indices,
+    write_info,
+    write_stats,
+    write_tasks,
+)
+
+
+def delete_episodes(
+    dataset: LeRobotDataset,
+    episode_indices: list[int],
+    output_dir: str | Path | None = None,
+    repo_id: str | None = None,
+) -> LeRobotDataset:
+    """Delete episodes from a LeRobotDataset and create a new dataset.
+
+    Args:
+        dataset: The source LeRobotDataset.
+        episode_indices: List of episode indices to delete.
+        output_dir: Directory to save the new dataset. If None, uses default location.
+        repo_id: Repository ID for the new dataset. If None, appends "_filtered" to original.
+
+    Returns:
+        LeRobotDataset: New dataset with episodes removed.
+    """
+    if not episode_indices:
+        raise ValueError("No episodes to delete")
+
+    # Validate episode indices
+    valid_indices = set(range(dataset.meta.total_episodes))
+    invalid = set(episode_indices) - valid_indices
+    if invalid:
+        raise ValueError(f"Invalid episode indices: {invalid}")
+
+    logging.info(f"Deleting {len(episode_indices)} episodes from dataset")
+
+    # Create new dataset metadata
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_filtered"
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Get episodes to keep
+    episodes_to_keep = [i for i in range(dataset.meta.total_episodes) if i not in episode_indices]
+    if not episodes_to_keep:
+        raise ValueError("Cannot delete all episodes from dataset")
+
+    # Create new dataset
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=dataset.meta.features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(dataset.meta.video_keys) > 0,
+    )
+
+    # Process episodes
+    episode_mapping = {}  # old_idx -> new_idx
+    new_episode_idx = 0
+
+    for old_idx in tqdm(episodes_to_keep, desc="Processing episodes"):
+        episode_mapping[old_idx] = new_episode_idx
+        new_episode_idx += 1
+
+    # Copy data files and update indices
+    _copy_and_reindex_data(dataset, new_meta, episode_mapping)
+
+    # Copy video files if present
+    if dataset.meta.video_keys:
+        _copy_and_reindex_videos(dataset, new_meta, episode_mapping)
+
+    # Create new dataset instance
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    logging.info(f"Created new dataset with {len(episodes_to_keep)} episodes")
+    return new_dataset
+
+
+def split_dataset(
+    dataset: LeRobotDataset,
+    splits: dict[str, list[int]] | dict[str, float],
+    output_dir: str | Path | None = None,
+) -> dict[str, LeRobotDataset]:
+    """Split a LeRobotDataset into multiple smaller datasets.
+
+    Args:
+        dataset: The source LeRobotDataset to split.
+        splits: Either a dict mapping split names to episode indices, or a dict mapping
+                split names to fractions (must sum to <= 1.0).
+        output_dir: Base directory for output datasets. If None, uses default location.
+
+    Returns:
+        dict[str, LeRobotDataset]: Dictionary mapping split names to new datasets.
+
+    Examples:
+        # Split by specific episodes
+        splits = {"train": [0, 1, 2], "val": [3, 4]}
+        datasets = split_dataset(dataset, splits)
+
+        # Split by fractions
+        splits = {"train": 0.8, "val": 0.2}
+        datasets = split_dataset(dataset, splits)
+    """
+    if not splits:
+        raise ValueError("No splits provided")
+
+    # Convert fractions to episode indices if needed
+    if all(isinstance(v, float) for v in splits.values()):
+        splits = _fractions_to_episode_indices(dataset.meta.total_episodes, splits)
+
+    # Validate episodes
+    all_episodes = set()
+    for split_name, episodes in splits.items():
+        if not episodes:
+            raise ValueError(f"Split '{split_name}' has no episodes")
+        episode_set = set(episodes)
+        if episode_set & all_episodes:
+            raise ValueError("Episodes cannot appear in multiple splits")
+        all_episodes.update(episode_set)
+
+    # Validate all episodes are valid
+    valid_indices = set(range(dataset.meta.total_episodes))
+    invalid = all_episodes - valid_indices
+    if invalid:
+        raise ValueError(f"Invalid episode indices: {invalid}")
+
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME
+    else:
+        output_dir = Path(output_dir)
+
+    result_datasets = {}
+
+    for split_name, episodes in splits.items():
+        logging.info(f"Creating split '{split_name}' with {len(episodes)} episodes")
+
+        # Create repo_id for split
+        split_repo_id = f"{dataset.repo_id}_{split_name}"
+        split_output_dir = output_dir / split_repo_id
+
+        # Create episode mapping
+        episode_mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(sorted(episodes))}
+
+        # Create new dataset metadata
+        new_meta = LeRobotDatasetMetadata.create(
+            repo_id=split_repo_id,
+            fps=dataset.meta.fps,
+            features=dataset.meta.features,
+            robot_type=dataset.meta.robot_type,
+            root=split_output_dir,
+            use_videos=len(dataset.meta.video_keys) > 0,
+        )
+
+        # Copy data and videos
+        _copy_and_reindex_data(dataset, new_meta, episode_mapping)
+        if dataset.meta.video_keys:
+            _copy_and_reindex_videos(dataset, new_meta, episode_mapping)
+
+        # Create new dataset instance
+        new_dataset = LeRobotDataset(
+            repo_id=split_repo_id,
+            root=split_output_dir,
+            image_transforms=dataset.image_transforms,
+            delta_timestamps=dataset.delta_timestamps,
+            tolerance_s=dataset.tolerance_s,
+        )
+
+        result_datasets[split_name] = new_dataset
+
+    return result_datasets
+
+
+def merge_datasets(
+    datasets: list[LeRobotDataset],
+    output_repo_id: str,
+    output_dir: str | Path | None = None,
+) -> LeRobotDataset:
+    """Merge multiple LeRobotDatasets into a single dataset.
+
+    This is a wrapper around the aggregate_datasets functionality with a cleaner API.
+
+    Args:
+        datasets: List of LeRobotDatasets to merge.
+        output_repo_id: Repository ID for the merged dataset.
+        output_dir: Directory to save the merged dataset. If None, uses default location.
+
+    Returns:
+        LeRobotDataset: The merged dataset.
+    """
+    if not datasets:
+        raise ValueError("No datasets to merge")
+
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / output_repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Extract repo_ids and roots
+    repo_ids = [ds.repo_id for ds in datasets]
+    roots = [ds.root for ds in datasets]
+
+    # Call aggregate_datasets
+    aggregate_datasets(
+        repo_ids=repo_ids,
+        aggr_repo_id=output_repo_id,
+        roots=roots,
+        aggr_root=output_dir,
+    )
+
+    # Create and return the merged dataset
+    merged_dataset = LeRobotDataset(
+        repo_id=output_repo_id,
+        root=output_dir,
+        image_transforms=datasets[0].image_transforms,
+        delta_timestamps=datasets[0].delta_timestamps,
+        tolerance_s=datasets[0].tolerance_s,
+    )
+
+    return merged_dataset
+
+
+def add_feature(
+    dataset: LeRobotDataset,
+    feature_name: str,
+    feature_values: np.ndarray | torch.Tensor | Callable,
+    feature_info: dict,
+    output_dir: str | Path | None = None,
+    repo_id: str | None = None,
+) -> LeRobotDataset:
+    """Add a new feature to a LeRobotDataset.
+
+    Args:
+        dataset: The source LeRobotDataset.
+        feature_name: Name of the new feature.
+        feature_values: Either:
+            - Array/tensor of shape (num_frames, ...) with values for each frame
+            - Callable that takes (frame_dict, episode_index, frame_index) and returns feature value
+        feature_info: Dictionary with feature metadata (dtype, shape, names).
+        output_dir: Directory to save the new dataset. If None, uses default location.
+        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
+
+    Returns:
+        LeRobotDataset: New dataset with the added feature.
+    """
+    if feature_name in dataset.meta.features:
+        raise ValueError(f"Feature '{feature_name}' already exists in dataset")
+
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_modified"
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Validate feature_info
+    required_keys = {"dtype", "shape"}
+    if not required_keys.issubset(feature_info.keys()):
+        raise ValueError(f"feature_info must contain keys: {required_keys}")
+
+    # Create new features dict
+    new_features = dataset.meta.features.copy()
+    new_features[feature_name] = feature_info
+
+    # Create new dataset metadata
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=new_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(dataset.meta.video_keys) > 0,
+    )
+
+    # Process data with new feature
+    _copy_data_with_feature_changes(
+        dataset=dataset,
+        new_meta=new_meta,
+        add_features={feature_name: (feature_values, feature_info)},
+    )
+
+    # Copy videos if present
+    if dataset.meta.video_keys:
+        _copy_videos(dataset, new_meta)
+
+    # Create new dataset instance
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    return new_dataset
+
+
+def remove_feature(
+    dataset: LeRobotDataset,
+    feature_names: str | list[str],
+    output_dir: str | Path | None = None,
+    repo_id: str | None = None,
+) -> LeRobotDataset:
+    """Remove features from a LeRobotDataset.
+
+    Args:
+        dataset: The source LeRobotDataset.
+        feature_names: Name(s) of features to remove. Can be a single string or list.
+        output_dir: Directory to save the new dataset. If None, uses default location.
+        repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
+
+    Returns:
+        LeRobotDataset: New dataset with features removed.
+    """
+    if isinstance(feature_names, str):
+        feature_names = [feature_names]
+
+    # Validate features exist
+    for name in feature_names:
+        if name not in dataset.meta.features:
+            raise ValueError(f"Feature '{name}' not found in dataset")
+
+    # Check if trying to remove required features
+    required_features = {"timestamp", "frame_index", "episode_index", "index", "task_index"}
+    if any(name in required_features for name in feature_names):
+        raise ValueError(f"Cannot remove required features: {required_features}")
+
+    if repo_id is None:
+        repo_id = f"{dataset.repo_id}_modified"
+    if output_dir is None:
+        output_dir = HF_LEROBOT_HOME / repo_id
+    else:
+        output_dir = Path(output_dir)
+
+    # Create new features dict
+    new_features = {k: v for k, v in dataset.meta.features.items() if k not in feature_names}
+
+    # Check if removing video features
+    video_keys_to_remove = [name for name in feature_names if name in dataset.meta.video_keys]
+
+    # Check if videos will remain after removal
+    remaining_video_keys = [k for k in dataset.meta.video_keys if k not in video_keys_to_remove]
+
+    # Create new dataset metadata
+    new_meta = LeRobotDatasetMetadata.create(
+        repo_id=repo_id,
+        fps=dataset.meta.fps,
+        features=new_features,
+        robot_type=dataset.meta.robot_type,
+        root=output_dir,
+        use_videos=len(remaining_video_keys) > 0,
+    )
+
+    # Process data with removed features
+    _copy_data_with_feature_changes(
+        dataset=dataset,
+        new_meta=new_meta,
+        remove_features=feature_names,
+    )
+
+    # Copy videos (excluding removed ones)
+    if new_meta.video_keys:
+        _copy_videos(dataset, new_meta, exclude_keys=video_keys_to_remove)
+
+    # Create new dataset instance
+    new_dataset = LeRobotDataset(
+        repo_id=repo_id,
+        root=output_dir,
+        image_transforms=dataset.image_transforms,
+        delta_timestamps=dataset.delta_timestamps,
+        tolerance_s=dataset.tolerance_s,
+    )
+
+    return new_dataset
+
+
+# Helper functions
+
+
+def _fractions_to_episode_indices(
+    total_episodes: int,
+    splits: dict[str, float],
+) -> dict[str, list[int]]:
+    """Convert split fractions to episode indices."""
+    if sum(splits.values()) > 1.0:
+        raise ValueError("Split fractions must sum to <= 1.0")
+
+    indices = list(range(total_episodes))
+    result = {}
+    start_idx = 0
+
+    for split_name, fraction in splits.items():
+        num_episodes = int(total_episodes * fraction)
+        end_idx = start_idx + num_episodes
+        if split_name == list(splits.keys())[-1]:  # Last split gets remaining episodes
+            end_idx = total_episodes
+        result[split_name] = indices[start_idx:end_idx]
+        start_idx = end_idx
+
+    return result
+
+
+def _copy_and_reindex_data(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_mapping: dict[int, int],
+) -> None:
+    """Copy data files and reindex episodes."""
+    # Get unique data files from episodes to keep
+    file_paths = set()
+    for old_idx in episode_mapping:
+        file_paths.add(src_dataset.meta.get_data_file_path(old_idx))
+
+    # Track global index
+    global_index = 0
+    chunk_idx, file_idx = 0, 0
+
+    # Process each data file
+    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
+        df = pd.read_parquet(src_dataset.root / src_path)
+
+        # Filter to keep only mapped episodes
+        mask = df["episode_index"].isin(episode_mapping.keys())
+        df = df[mask].copy()
+
+        if len(df) == 0:
+            continue
+
+        # Update episode indices
+        df["episode_index"] = df["episode_index"].map(episode_mapping)
+
+        # Update global index to be continuous
+        df["index"] = range(global_index, global_index + len(df))
+        global_index += len(df)
+
+        # Update task indices if needed
+        if dst_meta.tasks is None:
+            # Get unique tasks from filtered data
+            task_indices = df["task_index"].unique()
+            tasks = [src_dataset.meta.tasks.iloc[idx].name for idx in task_indices]
+            dst_meta.save_episode_tasks(list(set(tasks)))
+
+        # Remap task indices
+        task_mapping = {}
+        for old_task_idx in df["task_index"].unique():
+            task_name = src_dataset.meta.tasks.iloc[old_task_idx].name
+            new_task_idx = dst_meta.get_task_index(task_name)
+            task_mapping[old_task_idx] = new_task_idx
+        df["task_index"] = df["task_index"].map(task_mapping)
+
+        # Save processed data
+        chunk_idx, file_idx = _save_data_chunk(df, dst_meta, chunk_idx, file_idx)
+
+    # Process episodes metadata
+    _copy_and_reindex_episodes_metadata(src_dataset, dst_meta, episode_mapping)
+
+
+def _copy_and_reindex_videos(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_mapping: dict[int, int],
+) -> None:
+    """Copy video files and update metadata."""
+    for video_key in src_dataset.meta.video_keys:
+        video_files = set()
+        for old_idx in episode_mapping:
+            video_files.add(src_dataset.meta.get_video_file_path(old_idx, video_key))
+
+        chunk_idx, file_idx = 0, 0
+
+        for src_path in tqdm(sorted(video_files), desc=f"Processing {video_key} videos"):
+            dst_path = dst_meta.root / DEFAULT_VIDEO_PATH.format(
+                video_key=video_key,
+                chunk_index=chunk_idx,
+                file_index=file_idx,
+            )
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
+
+            # For simplicity, copy entire video files
+            # In production, you might want to extract only relevant segments
+            shutil.copy(src_dataset.root / src_path, dst_path)
+
+            # Update indices for next file
+            file_size = get_video_size_in_mb(dst_path)
+            if file_size >= DEFAULT_VIDEO_FILE_SIZE_IN_MB * 0.9:  # 90% threshold
+                chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+
+
+def _copy_and_reindex_episodes_metadata(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    episode_mapping: dict[int, int],
+) -> None:
+    """Copy and reindex episodes metadata."""
+    all_stats = []
+    frame_offset = 0
+
+    for old_idx, new_idx in tqdm(
+        sorted(episode_mapping.items(), key=lambda x: x[1]), desc="Processing episodes metadata"
+    ):
+        # Get episode from source
+        src_episode = src_dataset.meta.episodes[old_idx]
+
+        # Create episode dict
+        episode_dict = {
+            "episode_index": new_idx,
+            "tasks": src_episode["tasks"],  # Already a list of task names
+            "length": src_episode["length"],
+        }
+
+        # Copy other metadata
+        episode_metadata = {
+            "data/chunk_index": 0,  # Will be recalculated when saving
+            "data/file_index": 0,  # Will be recalculated when saving
+            "dataset_from_index": frame_offset,
+            "dataset_to_index": frame_offset + src_episode["length"],
+        }
+
+        # Update frame offset for next episode
+        frame_offset += src_episode["length"]
+
+        # Copy stats metadata
+        for key in src_episode.keys():
+            if key.startswith("stats/"):
+                episode_dict[key] = src_episode[key]
+
+        # Add episode metadata
+        stats_dict = {
+            key.replace("stats/", ""): value
+            for key, value in episode_dict.items()
+            if key.startswith("stats/")
+        }
+        all_stats.append(stats_dict)
+
+        # Calculate stats from dict
+        episode_stats = {}
+        for key in dst_meta.features:
+            if key in stats_dict:
+                episode_stats[key] = stats_dict[key]
+
+        dst_meta.save_episode(
+            new_idx, episode_dict["length"], episode_dict["tasks"], episode_stats, episode_metadata
+        )
+
+    # Aggregate all stats
+    if all_stats:
+        aggregated_stats = aggregate_stats(all_stats)
+        write_stats(aggregated_stats, dst_meta.root)
+
+
+def _save_data_chunk(
+    df: pd.DataFrame,
+    meta: LeRobotDatasetMetadata,
+    chunk_idx: int = 0,
+    file_idx: int = 0,
+) -> tuple[int, int]:
+    """Save a data chunk and return updated indices."""
+    path = meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if len(meta.image_keys) > 0:
+        to_parquet_with_hf_images(df, path)
+    else:
+        df.to_parquet(path)
+
+    # Check if we need to rotate files
+    file_size = get_parquet_file_size_in_mb(path)
+    if file_size >= DEFAULT_DATA_FILE_SIZE_IN_MB * 0.9:  # 90% threshold
+        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+
+    return chunk_idx, file_idx
+
+
+def _copy_data_with_feature_changes(
+    dataset: LeRobotDataset,
+    new_meta: LeRobotDatasetMetadata,
+    add_features: dict[str, tuple] | None = None,
+    remove_features: list[str] | None = None,
+) -> None:
+    """Copy data while adding or removing features."""
+    # Get all unique data files
+    file_paths = set()
+    for ep_idx in range(dataset.meta.total_episodes):
+        file_paths.add(dataset.meta.get_data_file_path(ep_idx))
+
+    frame_idx = 0
+
+    # Process each data file
+    for src_path in tqdm(sorted(file_paths), desc="Processing data files"):
+        df = pd.read_parquet(dataset.root / src_path)
+
+        # Remove features
+        if remove_features:
+            df = df.drop(columns=remove_features, errors="ignore")
+
+        # Add features
+        if add_features:
+            for feature_name, (values, _) in add_features.items():
+                if callable(values):
+                    # Compute values for each frame
+                    feature_values = []
+                    for _, row in df.iterrows():
+                        ep_idx = row["episode_index"]
+                        frame_in_ep = row["frame_index"]
+                        value = values(row.to_dict(), ep_idx, frame_in_ep)
+                        # Convert numpy arrays to scalars for single-element arrays
+                        if isinstance(value, np.ndarray) and value.size == 1:
+                            value = value.item()
+                        feature_values.append(value)
+                    df[feature_name] = feature_values
+                else:
+                    # Use provided values
+                    end_idx = frame_idx + len(df)
+                    # Convert to list to ensure proper shape handling
+                    feature_slice = values[frame_idx:end_idx]
+                    if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
+                        # Flatten single-element arrays to scalars for pandas
+                        df[feature_name] = feature_slice.flatten()
+                    else:
+                        df[feature_name] = feature_slice
+                    frame_idx = end_idx
+
+        # Save chunk
+        _save_data_chunk(df, new_meta)
+
+    # Copy episodes metadata and update stats
+    _copy_episodes_metadata_and_stats(dataset, new_meta)
+
+
+def _copy_videos(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+    exclude_keys: list[str] | None = None,
+) -> None:
+    """Copy video files, optionally excluding certain keys."""
+    if exclude_keys is None:
+        exclude_keys = []
+
+    for video_key in src_dataset.meta.video_keys:
+        if video_key in exclude_keys:
+            continue
+
+        # Get all video files for this key
+        video_files = set()
+        for ep_idx in range(src_dataset.meta.total_episodes):
+            video_files.add(src_dataset.meta.get_video_file_path(ep_idx, video_key))
+
+        # Copy video files
+        for src_path in tqdm(sorted(video_files), desc=f"Copying {video_key} videos"):
+            # Maintain same structure
+            rel_path = src_path.relative_to(src_dataset.root)
+            dst_path = dst_meta.root / rel_path
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy(src_dataset.root / src_path, dst_path)
+
+
+def _copy_episodes_metadata_and_stats(
+    src_dataset: LeRobotDataset,
+    dst_meta: LeRobotDatasetMetadata,
+) -> None:
+    """Copy episodes metadata and recalculate stats."""
+    # Copy tasks
+    if src_dataset.meta.tasks is not None:
+        write_tasks(src_dataset.meta.tasks, dst_meta.root)
+        dst_meta.tasks = src_dataset.meta.tasks.copy()
+
+    # Copy episodes metadata files
+    episodes_dir = src_dataset.root / "meta/episodes"
+    dst_episodes_dir = dst_meta.root / "meta/episodes"
+    if episodes_dir.exists():
+        shutil.copytree(episodes_dir, dst_episodes_dir, dirs_exist_ok=True)
+
+    # Update info
+    dst_meta.info.update(
+        {
+            "total_episodes": src_dataset.meta.total_episodes,
+            "total_frames": src_dataset.meta.total_frames,
+            "total_tasks": src_dataset.meta.total_tasks,
+            "splits": src_dataset.meta.info.get("splits", {"train": f"0:{src_dataset.meta.total_episodes}"}),
+        }
+    )
+
+    # Update video info if needed
+    if dst_meta.video_keys and src_dataset.meta.video_keys:
+        for key in dst_meta.video_keys:
+            if key in src_dataset.meta.features:
+                dst_meta.info["features"][key]["info"] = src_dataset.meta.info["features"][key].get(
+                    "info", {}
+                )
+
+    write_info(dst_meta.info, dst_meta.root)
+
+    # Recalculate stats if features changed
+    if set(dst_meta.features.keys()) != set(src_dataset.meta.features.keys()):
+        # Need to recalculate stats
+        logging.info("Recalculating dataset statistics...")
+        # This is a simplified version - in production you'd want to properly recalculate
+        if src_dataset.meta.stats:
+            new_stats = {}
+            for key in dst_meta.features:
+                if key in src_dataset.meta.stats:
+                    new_stats[key] = src_dataset.meta.stats[key]
+            write_stats(new_stats, dst_meta.root)
+    else:
+        # Copy existing stats
+        if src_dataset.meta.stats:
+            write_stats(src_dataset.meta.stats, dst_meta.root)

src/lerobot/datasets/factory.py

@@ -25,7 +25,6 @@ from lerobot.datasets.lerobot_dataset import (
     LeRobotDatasetMetadata,
     MultiLeRobotDataset,
 )
-from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.transforms import ImageTransforms
 
 IMAGENET_STATS = {
@@ -88,26 +87,15 @@ 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)
-        if not cfg.dataset.streaming:
-            dataset = LeRobotDataset(
-                cfg.dataset.repo_id,
-                root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
-                delta_timestamps=delta_timestamps,
-                image_transforms=image_transforms,
-                revision=cfg.dataset.revision,
-                video_backend=cfg.dataset.video_backend,
-            )
-        else:
-            dataset = StreamingLeRobotDataset(
-                cfg.dataset.repo_id,
-                root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
-                delta_timestamps=delta_timestamps,
-                image_transforms=image_transforms,
-                revision=cfg.dataset.revision,
-                max_num_shards=cfg.num_workers,
-            )
+        dataset = LeRobotDataset(
+            cfg.dataset.repo_id,
+            root=cfg.dataset.root,
+            episodes=cfg.dataset.episodes,
+            delta_timestamps=delta_timestamps,
+            image_transforms=image_transforms,
+            revision=cfg.dataset.revision,
+            video_backend=cfg.dataset.video_backend,
+        )
     else:
         raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
         dataset = MultiLeRobotDataset(

src/lerobot/datasets/lerobot_dataset.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
-import gc
 import logging
 import shutil
 import tempfile
@@ -28,7 +27,9 @@ import pandas as pd
 import PIL.Image
 import torch
 import torch.utils
+from datasets import Dataset
 from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.constants import REPOCARD_NAME
 from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.constants import HF_LEROBOT_HOME
@@ -47,12 +48,12 @@ from lerobot.datasets.utils import (
     embed_images,
     flatten_dict,
     get_delta_indices,
-    get_hf_dataset_cache_dir,
     get_hf_dataset_size_in_mb,
     get_hf_features_from_features,
     get_parquet_file_size_in_mb,
     get_parquet_num_frames,
     get_safe_version,
+    get_video_duration_in_s,
     get_video_size_in_mb,
     hf_transform_to_torch,
     is_valid_version,
@@ -72,11 +73,10 @@ from lerobot.datasets.utils import (
 )
 from lerobot.datasets.video_utils import (
     VideoFrame,
-    concatenate_video_files,
+    concat_video_files,
     decode_video_frames,
     encode_video_frames,
     get_safe_default_codec,
-    get_video_duration_in_s,
     get_video_info,
 )
 
@@ -104,6 +104,9 @@ class LeRobotDatasetMetadata:
                 self.revision = get_safe_version(self.repo_id, self.revision)
 
             (self.root / "meta").mkdir(exist_ok=True, parents=True)
+            # TODO(rcadene): instead of downloading all episodes metadata files,
+            # download only the ones associated to the requested episodes. This would
+            # require adding `episodes: list[int]` as argument.
             self.pull_from_repo(allow_patterns="meta/")
             self.load_metadata()
 
@@ -128,10 +131,6 @@ class LeRobotDatasetMetadata:
             ignore_patterns=ignore_patterns,
         )
 
-    @property
-    def url_root(self) -> str:
-        return f"hf://datasets/{self.repo_id}"
-
     @property
     def _version(self) -> packaging.version.Version:
         """Codebase version used to create this dataset."""
@@ -274,7 +273,7 @@ class LeRobotDatasetMetadata:
         """
         # Convert buffer into HF Dataset
         episode_dict = {key: [value] for key, value in episode_dict.items()}
-        ep_dataset = datasets.Dataset.from_dict(episode_dict)
+        ep_dataset = Dataset.from_dict(episode_dict)
         ep_size_in_mb = get_hf_dataset_size_in_mb(ep_dataset)
         df = pd.DataFrame(ep_dataset)
         num_frames = episode_dict["length"][0]
@@ -310,21 +309,14 @@ class LeRobotDatasetMetadata:
                 latest_df = pd.read_parquet(latest_path)
                 df = pd.concat([latest_df, df], ignore_index=True)
 
-                # Memort optimization
-                del latest_df
-                gc.collect()
-
         # Write the resulting dataframe from RAM to disk
         path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
         path.parent.mkdir(parents=True, exist_ok=True)
         df.to_parquet(path, index=False)
 
-        if self.episodes is not None:
-            # Remove the episodes cache directory, necessary to avoid cache bloat
-            cached_dir = get_hf_dataset_cache_dir(self.episodes)
-            if cached_dir is not None:
-                shutil.rmtree(cached_dir)
-
+        # Update the Hugging Face dataset by reloading it.
+        # This process should be fast because only the latest Parquet file has been modified.
+        # Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
         self.episodes = load_episodes(self.root)
 
     def save_episode(
@@ -349,75 +341,23 @@ class LeRobotDatasetMetadata:
         self.info["total_frames"] += episode_length
         self.info["total_tasks"] = len(self.tasks)
         self.info["splits"] = {"train": f"0:{self.info['total_episodes']}"}
-
+        if len(self.video_keys) > 0:
+            self.update_video_info()
         write_info(self.info, self.root)
 
         self.stats = aggregate_stats([self.stats, episode_stats]) if self.stats is not None else episode_stats
         write_stats(self.stats, self.root)
 
-    def update_video_info(self, video_key: str | None = None) -> None:
+    def update_video_info(self) -> None:
         """
         Warning: this function writes info from first episode videos, implicitly assuming that all videos have
         been encoded the same way. Also, this means it assumes the first episode exists.
         """
-        if video_key is not None and video_key not in self.video_keys:
-            raise ValueError(f"Video key {video_key} not found in dataset")
-
-        video_keys = [video_key] if video_key is not None else self.video_keys
-        for key in video_keys:
+        for key in self.video_keys:
             if not self.features[key].get("info", None):
-                video_path = self.root / self.video_path.format(
-                    video_key=video_key, chunk_index=0, file_index=0
-                )
+                video_path = self.root / self.get_video_file_path(ep_index=0, vid_key=key)
                 self.info["features"][key]["info"] = get_video_info(video_path)
 
-    def update_chunk_settings(
-        self,
-        chunks_size: int | None = None,
-        data_files_size_in_mb: int | None = None,
-        video_files_size_in_mb: int | None = None,
-    ) -> None:
-        """Update chunk and file size settings after dataset creation.
-
-        This allows users to customize storage organization without modifying the constructor.
-        These settings control how episodes are chunked and how large files can grow before
-        creating new ones.
-
-        Args:
-            chunks_size: Maximum number of files per chunk directory. If None, keeps current value.
-            data_files_size_in_mb: Maximum size for data parquet files in MB. If None, keeps current value.
-            video_files_size_in_mb: Maximum size for video files in MB. If None, keeps current value.
-        """
-        if chunks_size is not None:
-            if chunks_size <= 0:
-                raise ValueError(f"chunks_size must be positive, got {chunks_size}")
-            self.info["chunks_size"] = chunks_size
-
-        if data_files_size_in_mb is not None:
-            if data_files_size_in_mb <= 0:
-                raise ValueError(f"data_files_size_in_mb must be positive, got {data_files_size_in_mb}")
-            self.info["data_files_size_in_mb"] = data_files_size_in_mb
-
-        if video_files_size_in_mb is not None:
-            if video_files_size_in_mb <= 0:
-                raise ValueError(f"video_files_size_in_mb must be positive, got {video_files_size_in_mb}")
-            self.info["video_files_size_in_mb"] = video_files_size_in_mb
-
-        # Update the info file on disk
-        write_info(self.info, self.root)
-
-    def get_chunk_settings(self) -> dict[str, int]:
-        """Get current chunk and file size settings.
-
-        Returns:
-            Dict containing chunks_size, data_files_size_in_mb, and video_files_size_in_mb.
-        """
-        return {
-            "chunks_size": self.chunks_size,
-            "data_files_size_in_mb": self.data_files_size_in_mb,
-            "video_files_size_in_mb": self.video_files_size_in_mb,
-        }
-
     def __repr__(self):
         feature_keys = list(self.features)
         return (
@@ -467,13 +407,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
         root: str | Path | None = None,
         episodes: list[int] | None = None,
         image_transforms: Callable | None = None,
-        delta_timestamps: dict[str, list[float]] | None = None,
+        delta_timestamps: dict[list[float]] | None = None,
         tolerance_s: float = 1e-4,
         revision: str | None = None,
         force_cache_sync: bool = False,
         download_videos: bool = True,
         video_backend: str | None = None,
-        batch_encoding_size: int = 1,
     ):
         """
         2 modes are available for instantiating this class, depending on 2 different use cases:
@@ -584,8 +523,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 True.
             video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to 'pyav'.
                 You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
-            batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
-                Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
         """
         super().__init__()
         self.repo_id = repo_id
@@ -597,8 +534,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.revision = revision if revision else CODEBASE_VERSION
         self.video_backend = video_backend if video_backend else get_safe_default_codec()
         self.delta_indices = None
-        self.batch_encoding_size = batch_encoding_size
-        self.episodes_since_last_encoding = 0
 
         # Unused attributes
         self.image_writer = None
@@ -674,10 +609,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         else:
             hub_api.upload_folder(**upload_kwargs)
 
-        card = create_lerobot_dataset_card(
-            tags=tags, dataset_info=self.meta.info, license=license, **card_kwargs
-        )
-        card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch)
+        if not hub_api.file_exists(self.repo_id, REPOCARD_NAME, repo_type="dataset", revision=branch):
+            card = create_lerobot_dataset_card(
+                tags=tags, dataset_info=self.meta.info, license=license, **card_kwargs
+            )
+            card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch)
 
         if tag_version:
             with contextlib.suppress(RevisionNotFoundError):
@@ -739,10 +675,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
             return False
 
         # Get available episode indices from cached dataset
-        available_episodes = {
-            ep_idx.item() if isinstance(ep_idx, torch.Tensor) else ep_idx
-            for ep_idx in self.hf_dataset["episode_index"]
-        }
+        available_episodes = set(self.hf_dataset["episode_index"])
 
         # Determine requested episodes
         if self.episodes is None:
@@ -969,10 +902,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """
         This will save to disk the current episode in self.episode_buffer.
 
-        Video encoding is handled automatically based on batch_encoding_size:
-        - If batch_encoding_size == 1: Videos are encoded immediately after each episode
-        - If batch_encoding_size > 1: Videos are encoded in batches.
-
         Args:
             episode_data (dict | None, optional): Dict containing the episode data to save. If None, this will
                 save the current episode in self.episode_buffer, which is filled with 'add_frame'. Defaults to
@@ -1009,81 +938,15 @@ class LeRobotDataset(torch.utils.data.Dataset):
         ep_stats = compute_episode_stats(episode_buffer, self.features)
 
         ep_metadata = self._save_episode_data(episode_buffer)
-        has_video_keys = len(self.meta.video_keys) > 0
-        use_batched_encoding = self.batch_encoding_size > 1
-
-        if has_video_keys and not use_batched_encoding:
-            for video_key in self.meta.video_keys:
-                ep_metadata.update(self._save_episode_video(video_key, episode_index))
+        for video_key in self.meta.video_keys:
+            ep_metadata.update(self._save_episode_video(video_key, episode_index))
 
         # `meta.save_episode` need to be executed after encoding the videos
         self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats, ep_metadata)
 
-        if has_video_keys and use_batched_encoding:
-            # Check if we should trigger batch encoding
-            self.episodes_since_last_encoding += 1
-            if self.episodes_since_last_encoding == self.batch_encoding_size:
-                start_ep = self.num_episodes - self.batch_encoding_size
-                end_ep = self.num_episodes
-                self._batch_save_episode_video(start_ep, end_ep)
-                self.episodes_since_last_encoding = 0
-
         if not episode_data:
-            # Reset episode buffer and clean up temporary images (if not already deleted during video encoding)
-            self.clear_episode_buffer(delete_images=len(self.meta.image_keys) > 0)
-
-    def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None):
-        """
-        Batch save videos for multiple episodes.
-
-        Args:
-            start_episode: Starting episode index (inclusive)
-            end_episode: Ending episode index (exclusive). If None, encodes all episodes from start_episode to the current episode.
-        """
-        if end_episode is None:
-            end_episode = self.num_episodes
-
-        logging.info(
-            f"Batch encoding {self.batch_encoding_size} videos for episodes {start_episode} to {end_episode - 1}"
-        )
-
-        chunk_idx = self.meta.episodes[start_episode]["data/chunk_index"]
-        file_idx = self.meta.episodes[start_episode]["data/file_index"]
-        episode_df_path = self.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-        episode_df = pd.read_parquet(episode_df_path)
-
-        for ep_idx in range(start_episode, end_episode):
-            logging.info(f"Encoding videos for episode {ep_idx}")
-
-            if (
-                self.meta.episodes[ep_idx]["data/chunk_index"] != chunk_idx
-                or self.meta.episodes[ep_idx]["data/file_index"] != file_idx
-            ):
-                # The current episode is in a new chunk or file.
-                # Save previous episode dataframe and update the Hugging Face dataset by reloading it.
-                episode_df.to_parquet(episode_df_path)
-                self.meta.episodes = load_episodes(self.root)
-
-                # Load new episode dataframe
-                chunk_idx = self.meta.episodes[ep_idx]["data/chunk_index"]
-                file_idx = self.meta.episodes[ep_idx]["data/file_index"]
-                episode_df_path = self.root / DEFAULT_EPISODES_PATH.format(
-                    chunk_index=chunk_idx, file_index=file_idx
-                )
-                episode_df = pd.read_parquet(episode_df_path)
-
-            # Save the current episode's video metadata to the dataframe
-            video_ep_metadata = {}
-            for video_key in self.meta.video_keys:
-                video_ep_metadata.update(self._save_episode_video(video_key, ep_idx))
-            video_ep_metadata.pop("episode_index")
-            video_ep_df = pd.DataFrame(video_ep_metadata, index=[ep_idx]).convert_dtypes(
-                dtype_backend="pyarrow"
-            )  # allows NaN values along with integers
-
-            episode_df = episode_df.combine_first(video_ep_df)
-            episode_df.to_parquet(episode_df_path)
-            self.meta.episodes = load_episodes(self.root)
+            # Reset episode buffer and clean up temporary images
+            self.clear_episode_buffer()
 
     def _save_episode_data(self, episode_buffer: dict) -> dict:
         """Save episode data to a parquet file and update the Hugging Face dataset of frames data.
@@ -1130,10 +993,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 latest_df = pd.read_parquet(latest_path)
                 df = pd.concat([latest_df, df], ignore_index=True)
 
-                # Memort optimization
-                del latest_df
-                gc.collect()
-
         # Write the resulting dataframe from RAM to disk
         path = self.root / self.meta.data_path.format(chunk_index=chunk_idx, file_index=file_idx)
         path.parent.mkdir(parents=True, exist_ok=True)
@@ -1142,12 +1001,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
         else:
             df.to_parquet(path)
 
-        if self.hf_dataset is not None:
-            # Remove hf dataset cache directory, necessary to avoid cache bloat
-            cached_dir = get_hf_dataset_cache_dir(self.hf_dataset)
-            if cached_dir is not None:
-                shutil.rmtree(cached_dir)
-
+        # Update the Hugging Face dataset by reloading it.
+        # This process should be fast because only the latest Parquet file has been modified.
+        # Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
         self.hf_dataset = self.load_hf_dataset()
 
         metadata = {
@@ -1164,21 +1020,18 @@ class LeRobotDataset(torch.utils.data.Dataset):
         ep_size_in_mb = get_video_size_in_mb(ep_path)
         ep_duration_in_s = get_video_duration_in_s(ep_path)
 
-        if self.meta.episodes is None or (
-            f"videos/{video_key}/chunk_index" not in self.meta.episodes.column_names
-            or f"videos/{video_key}/file_index" not in self.meta.episodes.column_names
-        ):
+        if self.meta.episodes is None:
             # Initialize indices for a new dataset made of the first episode data
             chunk_idx, file_idx = 0, 0
-            latest_duration_in_s = 0.0
+            latest_duration_in_s = 0
             new_path = self.root / self.meta.video_path.format(
                 video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
             )
             new_path.parent.mkdir(parents=True, exist_ok=True)
             shutil.move(str(ep_path), str(new_path))
         else:
-            # Retrieve information from the latest updated video file (possibly several episodes ago)
-            latest_ep = self.meta.episodes[episode_index - 1]
+            # Retrieve information from the latest video file
+            latest_ep = self.meta.episodes[-1]
             chunk_idx = latest_ep[f"videos/{video_key}/chunk_index"]
             file_idx = latest_ep[f"videos/{video_key}/file_index"]
 
@@ -1196,22 +1049,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
                 )
                 new_path.parent.mkdir(parents=True, exist_ok=True)
                 shutil.move(str(ep_path), str(new_path))
-                latest_duration_in_s = 0.0
             else:
                 # Update latest video file
-                concatenate_video_files(
-                    [latest_path, ep_path],
-                    latest_path,
-                )
+                concat_video_files([latest_path, ep_path], self.root, video_key, chunk_idx, file_idx)
 
         # Remove temporary directory
         shutil.rmtree(str(ep_path.parent))
 
-        # Update video info (only needed when first episode is encoded since it reads from episode 0)
-        if episode_index == 0:
-            self.meta.update_video_info(video_key)
-            write_info(self.meta.info, self.meta.root)  # ensure video info always written properly
-
         metadata = {
             "episode_index": episode_index,
             f"videos/{video_key}/chunk_index": chunk_idx,
@@ -1221,17 +1065,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
         }
         return metadata
 
-    def clear_episode_buffer(self, delete_images: bool = True) -> None:
-        # Clean up image files for the current episode buffer
-        if delete_images:
-            # Wait for the async image writer to finish
-            if self.image_writer is not None:
-                self._wait_image_writer()
-            episode_index = self.episode_buffer["episode_index"]
-            if isinstance(episode_index, np.ndarray):
-                episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
+    def clear_episode_buffer(self) -> None:
+        if self.image_writer is not None:
             for cam_key in self.meta.camera_keys:
-                img_dir = self._get_image_file_dir(episode_index, cam_key)
+                img_dir = self.root / "images" / cam_key
                 if img_dir.is_dir():
                     shutil.rmtree(img_dir)
 
@@ -1272,7 +1109,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         temp_path = Path(tempfile.mkdtemp(dir=self.root)) / f"{video_key}_{episode_index:03d}.mp4"
         img_dir = self._get_image_file_dir(episode_index, video_key)
         encode_video_frames(img_dir, temp_path, self.fps, overwrite=True)
-        shutil.rmtree(img_dir)
         return temp_path
 
     @classmethod
@@ -1288,7 +1124,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         image_writer_processes: int = 0,
         image_writer_threads: int = 0,
         video_backend: str | None = None,
-        batch_encoding_size: int = 1,
     ) -> "LeRobotDataset":
         """Create a LeRobot Dataset from scratch in order to record data."""
         obj = cls.__new__(cls)
@@ -1305,8 +1140,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.revision = None
         obj.tolerance_s = tolerance_s
         obj.image_writer = None
-        obj.batch_encoding_size = batch_encoding_size
-        obj.episodes_since_last_encoding = 0
 
         if image_writer_processes or image_writer_threads:
             obj.start_image_writer(image_writer_processes, image_writer_threads)
@@ -1336,7 +1169,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         root: str | Path | None = None,
         episodes: dict | None = None,
         image_transforms: Callable | None = None,
-        delta_timestamps: dict[str, list[float]] | None = None,
+        delta_timestamps: dict[list[float]] | None = None,
         tolerances_s: dict | None = None,
         download_videos: bool = True,
         video_backend: str | None = None,

src/lerobot/datasets/pipeline_features.py

@@ -1,141 +0,0 @@
-# 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.
-
-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 = None, observation: dict[str, Any] | None = 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])
-)
-
-
-def aggregate_pipeline_dataset_features(
-    pipeline: DataProcessorPipeline,
-    initial_features: dict[PipelineFeatureType, 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.
-
-    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.
-    """
-    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]
-
-    # Iterate through all features transformed by the pipeline.
-    for ptype, feats in all_features.items():
-        if ptype not in [PipelineFeatureType.ACTION, PipelineFeatureType.OBSERVATION]:
-            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
-                )
-            )
-
-            # 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

src/lerobot/datasets/streaming_dataset.py

@@ -1,535 +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 collections.abc import Callable, Generator, Iterator
-from pathlib import Path
-
-import datasets
-import numpy as np
-import torch
-from datasets import load_dataset
-
-from lerobot.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
-from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
-from lerobot.datasets.utils import (
-    Backtrackable,
-    LookAheadError,
-    LookBackError,
-    check_version_compatibility,
-    find_float_index,
-    get_delta_indices,
-    is_float_in_list,
-    item_to_torch,
-    safe_shard,
-)
-from lerobot.datasets.video_utils import (
-    VideoDecoderCache,
-    decode_video_frames_torchcodec,
-)
-
-
-class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
-    """LeRobotDataset with streaming capabilities.
-
-    This class extends LeRobotDataset to add streaming functionality, allowing data to be streamed
-    rather than loaded entirely into memory. This is especially useful for large datasets that may
-    not fit in memory or when you want to quickly explore a dataset without downloading it completely.
-
-    The key innovation is using a Backtrackable iterator that maintains a bounded buffer of recent
-    items, allowing us to access previous frames for delta timestamps without loading the entire
-    dataset into memory.
-
-    Example:
-        Basic usage:
-        ```python
-        from lerobot.common.datasets.streaming_dataset import StreamingLeRobotDataset
-
-        # Create a streaming dataset with delta timestamps
-        delta_timestamps = {
-            "observation.image": [-1.0, -0.5, 0.0],  # 1 sec ago, 0.5 sec ago, current
-            "action": [0.0, 0.1, 0.2],  # current, 0.1 sec future, 0.2 sec future
-        }
-
-        dataset = StreamingLeRobotDataset(
-            repo_id="your-dataset-repo-id",
-            delta_timestamps=delta_timestamps,
-            streaming=True,
-            buffer_size=1000,
-        )
-
-        # Iterate over the dataset
-        for i, item in enumerate(dataset):
-            print(f"Sample {i}: Episode {item['episode_index']} Frame {item['frame_index']}")
-            # item will contain stacked frames according to delta_timestamps
-            if i >= 10:
-                break
-        ```
-    """
-
-    def __init__(
-        self,
-        repo_id: str,
-        root: str | Path | None = None,
-        episodes: list[int] | None = None,
-        image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
-        tolerance_s: float = 1e-4,
-        revision: str | None = None,
-        force_cache_sync: bool = False,
-        streaming: bool = True,
-        buffer_size: int = 1000,
-        max_num_shards: int = 16,
-        seed: int = 42,
-        rng: np.random.Generator | None = None,
-        shuffle: bool = True,
-    ):
-        """Initialize a StreamingLeRobotDataset.
-
-        Args:
-            repo_id (str): This is the repo id that will be used to fetch the dataset.
-            root (Path | None, optional): Local directory to use for downloading/writing files.
-            episodes (list[int] | None, optional): If specified, this will only load episodes specified by
-                their episode_index in this list.
-            image_transforms (Callable | None, optional): Transform to apply to image data.
-            tolerance_s (float, optional): Tolerance in seconds for timestamp matching.
-            revision (str, optional): Git revision id (branch name, tag, or commit hash).
-            force_cache_sync (bool, optional): Flag to sync and refresh local files first.
-            streaming (bool, optional): Whether to stream the dataset or load it all. Defaults to True.
-            buffer_size (int, optional): Buffer size for shuffling when streaming. Defaults to 1000.
-            max_num_shards (int, optional): Number of shards to re-shard the input dataset into. Defaults to 16.
-            seed (int, optional): Reproducibility random seed.
-            rng (np.random.Generator | None, optional): Random number generator.
-            shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
-        """
-        super().__init__()
-        self.repo_id = repo_id
-        self.root = Path(root) if root else HF_LEROBOT_HOME / repo_id
-        self.streaming_from_local = root is not None
-
-        self.image_transforms = image_transforms
-        self.episodes = episodes
-        self.tolerance_s = tolerance_s
-        self.revision = revision if revision else CODEBASE_VERSION
-        self.seed = seed
-        self.rng = rng if rng is not None else np.random.default_rng(seed)
-        self.shuffle = shuffle
-
-        self.streaming = streaming
-        self.buffer_size = buffer_size
-
-        # We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown)
-        self.video_decoder_cache = None
-
-        self.root.mkdir(exist_ok=True, parents=True)
-
-        # Load metadata
-        self.meta = LeRobotDatasetMetadata(
-            self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync
-        )
-        # Check version
-        check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
-
-        self.delta_timestamps = None
-        self.delta_indices = None
-
-        if delta_timestamps is not None:
-            self._validate_delta_timestamp_keys(delta_timestamps)  # raises ValueError if invalid
-            self.delta_timestamps = delta_timestamps
-            self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
-
-        self.hf_dataset: datasets.IterableDataset = load_dataset(
-            self.repo_id if not self.streaming_from_local else str(self.root),
-            split="train",
-            streaming=self.streaming,
-            data_files="data/*/*.parquet",
-            revision=self.revision,
-        )
-
-        self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
-
-    @property
-    def num_frames(self):
-        return self.meta.total_frames
-
-    @property
-    def num_episodes(self):
-        return self.meta.total_episodes
-
-    @property
-    def fps(self):
-        return self.meta.fps
-
-    @staticmethod
-    def _iter_random_indices(
-        rng: np.random.Generator, buffer_size: int, random_batch_size=100
-    ) -> Iterator[int]:
-        while True:
-            yield from (int(i) for i in rng.integers(0, buffer_size, size=random_batch_size))
-
-    @staticmethod
-    def _infinite_generator_over_elements(rng: np.random.Generator, elements: list[int]) -> Iterator[int]:
-        while True:
-            yield rng.choice(elements)
-
-    # TODO(fracapuano): Implement multi-threaded prefetching to accelerate data loading.
-    # The current sequential iteration is a bottleneck. A producer-consumer pattern
-    # could be used with a ThreadPoolExecutor to run `make_frame` (especially video decoding)
-    # in parallel, feeding a queue from which this iterator will yield processed items.
-    def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
-        if self.video_decoder_cache is None:
-            self.video_decoder_cache = VideoDecoderCache()
-
-        # keep the same seed across exhaustions if shuffle is False, otherwise shuffle data across exhaustions
-        rng = np.random.default_rng(self.seed) if not self.shuffle else self.rng
-
-        buffer_indices_generator = self._iter_random_indices(rng, self.buffer_size)
-
-        idx_to_backtrack_dataset = {
-            idx: self._make_backtrackable_dataset(safe_shard(self.hf_dataset, idx, self.num_shards))
-            for idx in range(self.num_shards)
-        }
-
-        # This buffer is populated while iterating on the dataset's shards
-        # the logic is to add 2 levels of randomness:
-        # (1) sample one shard at random from the ones available, and
-        # (2) sample one frame from the shard sampled at (1)
-        frames_buffer = []
-        while available_shards := list(idx_to_backtrack_dataset.keys()):
-            shard_key = next(self._infinite_generator_over_elements(rng, available_shards))
-            backtrack_dataset = idx_to_backtrack_dataset[shard_key]  # selects which shard to iterate on
-
-            try:
-                for frame in self.make_frame(backtrack_dataset):
-                    if len(frames_buffer) == self.buffer_size:
-                        i = next(buffer_indices_generator)  # samples a element from the buffer
-                        yield frames_buffer[i]
-                        frames_buffer[i] = frame
-                    else:
-                        frames_buffer.append(frame)
-                    break  # random shard sampled, switch shard
-            except (
-                RuntimeError,
-                StopIteration,
-            ):  # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7
-                del idx_to_backtrack_dataset[shard_key]  # Remove exhausted shard, onto another shard
-
-        # Once shards are all exhausted, shuffle the buffer and yield the remaining frames
-        rng.shuffle(frames_buffer)
-        yield from frames_buffer
-
-    def _get_window_steps(
-        self, delta_timestamps: dict[str, list[float]] | None = None, dynamic_bounds: bool = False
-    ) -> tuple[int, int]:
-        if delta_timestamps is None:
-            return 1, 1
-
-        if not dynamic_bounds:
-            # Fix the windows
-            lookback = LOOKBACK_BACKTRACKTABLE
-            lookahead = LOOKAHEAD_BACKTRACKTABLE
-        else:
-            # Dynamically adjust the windows based on the given delta_timesteps
-            all_timestamps = sum(delta_timestamps.values(), [])
-            lookback = min(all_timestamps) * self.fps
-            lookahead = max(all_timestamps) * self.fps
-
-            # When lookback is >=0 it means no negative timesteps have been provided
-            lookback = 0 if lookback >= 0 else (lookback * -1)
-
-        return lookback, lookahead
-
-    def _make_backtrackable_dataset(self, dataset: datasets.IterableDataset) -> Backtrackable:
-        lookback, lookahead = self._get_window_steps(self.delta_timestamps)
-        return Backtrackable(dataset, history=lookback, lookahead=lookahead)
-
-    def _make_timestamps_from_indices(
-        self, start_ts: float, indices: dict[str, list[int]] | None = None
-    ) -> dict[str, list[float]]:
-        if indices is not None:
-            return {
-                key: (
-                    start_ts + torch.tensor(indices[key]) / self.fps
-                ).tolist()  # NOTE: why not delta_timestamps directly?
-                for key in self.delta_timestamps
-            }
-        else:
-            return dict.fromkeys(self.meta.video_keys, [start_ts])
-
-    def _make_padding_camera_frame(self, camera_key: str):
-        """Variable-shape padding frame for given camera keys, given in (H, W, C)"""
-        return torch.zeros(self.meta.info["features"][camera_key]["shape"]).permute(-1, 0, 1)
-
-    def _get_video_frame_padding_mask(
-        self,
-        video_frames: dict[str, torch.Tensor],
-        query_timestamps: dict[str, list[float]],
-        original_timestamps: dict[str, list[float]],
-    ) -> dict[str, torch.BoolTensor]:
-        padding_mask = {}
-
-        for video_key, timestamps in original_timestamps.items():
-            if video_key not in video_frames:
-                continue  # only padding on video keys that are available
-            frames = []
-            mask = []
-            padding_frame = self._make_padding_camera_frame(video_key)
-            for ts in timestamps:
-                if is_float_in_list(ts, query_timestamps[video_key]):
-                    idx = find_float_index(ts, query_timestamps[video_key])
-                    frames.append(video_frames[video_key][idx, :])
-                    mask.append(False)
-                else:
-                    frames.append(padding_frame)
-                    mask.append(True)
-
-            padding_mask[f"{video_key}_is_pad"] = torch.BoolTensor(mask)
-
-        return padding_mask
-
-    def make_frame(
-        self, dataset_iterator: Backtrackable, previous_dataset_iterator: Backtrackable | None = None
-    ) -> Generator:
-        """Makes a frame starting from a dataset iterator"""
-        item = next(dataset_iterator)
-        item = item_to_torch(item)
-
-        updates = []  # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features)
-
-        # Get episode index from the item
-        ep_idx = item["episode_index"]
-
-        # "timestamp" restarts from 0 for each episode, whereas we need a global timestep within the single .mp4 file (given by index/fps)
-        current_ts = item["index"] / self.fps
-
-        episode_boundaries_ts = {
-            key: (
-                self.meta.episodes[ep_idx][f"videos/{key}/from_timestamp"],
-                self.meta.episodes[ep_idx][f"videos/{key}/to_timestamp"],
-            )
-            for key in self.meta.video_keys
-        }
-
-        # Apply delta querying logic if necessary
-        if self.delta_indices is not None:
-            query_result, padding = self._get_delta_frames(dataset_iterator, item)
-            updates.append(query_result)
-            updates.append(padding)
-
-        # Load video frames, when needed
-        if len(self.meta.video_keys) > 0:
-            original_timestamps = self._make_timestamps_from_indices(current_ts, self.delta_indices)
-
-            # Some timestamps might not result available considering the episode's boundaries
-            query_timestamps = self._get_query_timestamps(
-                current_ts, self.delta_indices, episode_boundaries_ts
-            )
-            video_frames = self._query_videos(query_timestamps, ep_idx)
-
-            if self.image_transforms is not None:
-                image_keys = self.meta.camera_keys
-                for cam in image_keys:
-                    video_frames[cam] = self.image_transforms(video_frames[cam])
-
-            updates.append(video_frames)
-
-            if self.delta_indices is not None:
-                # We always return the same number of frames. Unavailable frames are padded.
-                padding_mask = self._get_video_frame_padding_mask(
-                    video_frames, query_timestamps, original_timestamps
-                )
-                updates.append(padding_mask)
-
-        result = item.copy()
-        for update in updates:
-            result.update(update)
-
-        result["task"] = self.meta.tasks.iloc[item["task_index"]].name
-
-        yield result
-
-    def _get_query_timestamps(
-        self,
-        current_ts: float,
-        query_indices: dict[str, list[int]] | None = None,
-        episode_boundaries_ts: dict[str, tuple[float, float]] | None = None,
-    ) -> dict[str, list[float]]:
-        query_timestamps = {}
-        keys_to_timestamps = self._make_timestamps_from_indices(current_ts, query_indices)
-        for key in self.meta.video_keys:
-            if query_indices is not None and key in query_indices:
-                timestamps = keys_to_timestamps[key]
-                # Clamp out timesteps outside of episode boundaries
-                query_timestamps[key] = torch.clamp(
-                    torch.tensor(timestamps), *episode_boundaries_ts[key]
-                ).tolist()
-
-            else:
-                query_timestamps[key] = [current_ts]
-
-        return query_timestamps
-
-    def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict:
-        """Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
-        in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
-        Segmentation Fault. This probably happens because a memory reference to the video loader is created in
-        the main process and a subprocess fails to access it.
-        """
-
-        item = {}
-        for video_key, query_ts in query_timestamps.items():
-            root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
-            video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}"
-            frames = decode_video_frames_torchcodec(
-                video_path, query_ts, self.tolerance_s, decoder_cache=self.video_decoder_cache
-            )
-
-            item[video_key] = frames.squeeze(0) if len(query_ts) == 1 else frames
-
-        return item
-
-    def _get_delta_frames(self, dataset_iterator: Backtrackable, current_item: dict):
-        # TODO(fracapuano): Modularize this function, refactor the code
-        """Get frames with delta offsets using the backtrackable iterator.
-
-        Args:
-            current_item (dict): Current item from the iterator.
-            ep_idx (int): Episode index.
-
-        Returns:
-            tuple: (query_result, padding) - frames at delta offsets and padding info.
-        """
-        current_episode_idx = current_item["episode_index"]
-
-        # Prepare results
-        query_result = {}
-        padding = {}
-
-        for key, delta_indices in self.delta_indices.items():
-            if key in self.meta.video_keys:
-                continue  # visual frames are decoded separately
-
-            target_frames = []
-            is_pad = []
-
-            # Create a results dictionary to store frames in processing order, then reconstruct original order for stacking
-            delta_results = {}
-
-            # Separate and sort deltas by difficulty (easier operations first)
-            negative_deltas = sorted([d for d in delta_indices if d < 0], reverse=True)  # [-1, -2, -3, ...]
-            positive_deltas = sorted([d for d in delta_indices if d > 0])  # [1, 2, 3, ...]
-            zero_deltas = [d for d in delta_indices if d == 0]
-
-            # Process zero deltas (current frame)
-            for delta in zero_deltas:
-                delta_results[delta] = (
-                    current_item[key],
-                    False,
-                )
-
-            # Process negative deltas in order of increasing difficulty
-            lookback_failed = False
-
-            last_successful_frame = current_item[key]
-
-            for delta in negative_deltas:
-                if lookback_failed:
-                    delta_results[delta] = (last_successful_frame, True)
-                    continue
-
-                try:
-                    steps_back = abs(delta)
-                    if dataset_iterator.can_peek_back(steps_back):
-                        past_item = dataset_iterator.peek_back(steps_back)
-                        past_item = item_to_torch(past_item)
-
-                        if past_item["episode_index"] == current_episode_idx:
-                            delta_results[delta] = (past_item[key], False)
-                            last_successful_frame = past_item[key]
-
-                        else:
-                            raise LookBackError("Retrieved frame is from different episode!")
-                    else:
-                        raise LookBackError("Cannot go back further than the history buffer!")
-
-                except LookBackError:
-                    delta_results[delta] = (last_successful_frame, True)
-                    lookback_failed = True  # All subsequent negative deltas will also fail
-
-            # Process positive deltas in order of increasing difficulty
-            lookahead_failed = False
-            last_successful_frame = current_item[key]
-
-            for delta in positive_deltas:
-                if lookahead_failed:
-                    delta_results[delta] = (last_successful_frame, True)
-                    continue
-
-                try:
-                    if dataset_iterator.can_peek_ahead(delta):
-                        future_item = dataset_iterator.peek_ahead(delta)
-                        future_item = item_to_torch(future_item)
-
-                        if future_item["episode_index"] == current_episode_idx:
-                            delta_results[delta] = (future_item[key], False)
-                            last_successful_frame = future_item[key]
-
-                        else:
-                            raise LookAheadError("Retrieved frame is from different episode!")
-                    else:
-                        raise LookAheadError("Cannot go ahead further than the lookahead buffer!")
-
-                except LookAheadError:
-                    delta_results[delta] = (last_successful_frame, True)
-                    lookahead_failed = True  # All subsequent positive deltas will also fail
-
-            # Reconstruct original order for stacking
-            for delta in delta_indices:
-                frame, is_padded = delta_results[delta]
-
-                # add batch dimension for stacking
-                target_frames.append(frame)  # frame.unsqueeze(0))
-                is_pad.append(is_padded)
-
-            # Stack frames and add to results
-            if target_frames:
-                query_result[key] = torch.stack(target_frames)
-                padding[f"{key}_is_pad"] = torch.BoolTensor(is_pad)
-
-        return query_result, padding
-
-    def _validate_delta_timestamp_keys(self, delta_timestamps: dict[list[float]]) -> None:
-        """
-        Validate that all keys in delta_timestamps correspond to actual features in the dataset.
-
-        Raises:
-            ValueError: If any delta timestamp key doesn't correspond to a dataset feature.
-        """
-        if delta_timestamps is None:
-            return
-
-        # Get all available feature keys from the dataset metadata
-        available_features = set(self.meta.features.keys())
-
-        # Get all keys from delta_timestamps
-        delta_keys = set(delta_timestamps.keys())
-
-        # Find any keys that don't correspond to features
-        invalid_keys = delta_keys - available_features
-
-        if invalid_keys:
-            raise ValueError(
-                f"The following delta_timestamp keys do not correspond to dataset features: {invalid_keys}. "
-                f"Available features are: {sorted(available_features)}"
-            )

src/lerobot/datasets/v2/batch_convert_dataset_v1_to_v2.py

@@ -0,0 +1,884 @@
+#!/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.
+
+"""
+This script is for internal use to convert all datasets under the 'lerobot' hub user account to v2.
+
+Note: Since the original Aloha datasets don't use shadow motors, you need to comment those out in
+lerobot/configs/robot/aloha.yaml before running this script.
+"""
+
+import traceback
+from pathlib import Path
+from textwrap import dedent
+
+from lerobot import available_datasets
+from lerobot.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset
+from lerobot.robots.aloha.configuration_aloha import AlohaRobotConfig
+
+LOCAL_DIR = Path("data/")
+
+# spellchecker:off
+ALOHA_MOBILE_INFO = {
+    "robot_config": AlohaRobotConfig(),
+    "license": "mit",
+    "url": "https://mobile-aloha.github.io/",
+    "paper": "https://huggingface.co/papers/2401.02117",
+    "citation_bibtex": dedent(r"""
+        @inproceedings{fu2024mobile,
+            author    = {Fu, Zipeng and Zhao, Tony Z. and Finn, Chelsea},
+            title     = {Mobile ALOHA: Learning Bimanual Mobile Manipulation with Low-Cost Whole-Body Teleoperation},
+            booktitle = {arXiv},
+            year      = {2024},
+        }""").lstrip(),
+}
+ALOHA_STATIC_INFO = {
+    "robot_config": AlohaRobotConfig(),
+    "license": "mit",
+    "url": "https://tonyzhaozh.github.io/aloha/",
+    "paper": "https://huggingface.co/papers/2304.13705",
+    "citation_bibtex": dedent(r"""
+        @article{Zhao2023LearningFB,
+            title={Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware},
+            author={Tony Zhao and Vikash Kumar and Sergey Levine and Chelsea Finn},
+            journal={RSS},
+            year={2023},
+            volume={abs/2304.13705},
+            url={https://huggingface.co/papers/2304.13705}
+        }""").lstrip(),
+}
+PUSHT_INFO = {
+    "license": "mit",
+    "url": "https://diffusion-policy.cs.columbia.edu/",
+    "paper": "https://huggingface.co/papers/2303.04137",
+    "citation_bibtex": dedent(r"""
+        @article{chi2024diffusionpolicy,
+            author = {Cheng Chi and Zhenjia Xu and Siyuan Feng and Eric Cousineau and Yilun Du and Benjamin Burchfiel and Russ Tedrake and Shuran Song},
+            title ={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
+            journal = {The International Journal of Robotics Research},
+            year = {2024},
+        }""").lstrip(),
+}
+XARM_INFO = {
+    "license": "mit",
+    "url": "https://www.nicklashansen.com/td-mpc/",
+    "paper": "https://huggingface.co/papers/2203.04955",
+    "citation_bibtex": dedent(r"""
+        @inproceedings{Hansen2022tdmpc,
+            title={Temporal Difference Learning for Model Predictive Control},
+            author={Nicklas Hansen and Xiaolong Wang and Hao Su},
+            booktitle={ICML},
+            year={2022}
+        }
+    """),
+}
+UNITREEH_INFO = {
+    "license": "apache-2.0",
+}
+
+DATASETS = {
+    "aloha_mobile_cabinet": {
+        "single_task": "Open the top cabinet, store the pot inside it then close the cabinet.",
+        **ALOHA_MOBILE_INFO,
+    },
+    "aloha_mobile_chair": {
+        "single_task": "Push the chairs in front of the desk to place them against it.",
+        **ALOHA_MOBILE_INFO,
+    },
+    "aloha_mobile_elevator": {
+        "single_task": "Take the elevator to the 1st floor.",
+        **ALOHA_MOBILE_INFO,
+    },
+    "aloha_mobile_shrimp": {
+        "single_task": "Sauté the raw shrimp on both sides, then serve it in the bowl.",
+        **ALOHA_MOBILE_INFO,
+    },
+    "aloha_mobile_wash_pan": {
+        "single_task": "Pick up the pan, rinse it in the sink and then place it in the drying rack.",
+        **ALOHA_MOBILE_INFO,
+    },
+    "aloha_mobile_wipe_wine": {
+        "single_task": "Pick up the wet cloth on the faucet and use it to clean the spilled wine on the table and underneath the glass.",
+        **ALOHA_MOBILE_INFO,
+    },
+    "aloha_static_battery": {
+        "single_task": "Place the battery into the slot of the remote controller.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_candy": {"single_task": "Pick up the candy and unwrap it.", **ALOHA_STATIC_INFO},
+    "aloha_static_coffee": {
+        "single_task": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray, then push the 'Hot Water' and 'Travel Mug' buttons.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_coffee_new": {
+        "single_task": "Place the coffee capsule inside the capsule container, then place the cup onto the center of the cup tray.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_cups_open": {
+        "single_task": "Pick up the plastic cup and open its lid.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_fork_pick_up": {
+        "single_task": "Pick up the fork and place it on the plate.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_pingpong_test": {
+        "single_task": "Transfer one of the two balls in the right glass into the left glass, then transfer it back to the right glass.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_pro_pencil": {
+        "single_task": "Pick up the pencil with the right arm, hand it over to the left arm then place it back onto the table.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_screw_driver": {
+        "single_task": "Pick up the screwdriver with the right arm, hand it over to the left arm then place it into the cup.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_tape": {
+        "single_task": "Cut a small piece of tape from the tape dispenser then place it on the cardboard box's edge.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_thread_velcro": {
+        "single_task": "Pick up the velcro cable tie with the left arm, then insert the end of the velcro tie into the other end's loop with the right arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_towel": {
+        "single_task": "Pick up a piece of paper towel and place it on the spilled liquid.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_vinh_cup": {
+        "single_task": "Pick up the plastic cup with the right arm, then pop its lid open with the left arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_vinh_cup_left": {
+        "single_task": "Pick up the plastic cup with the left arm, then pop its lid open with the right arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_static_ziploc_slide": {"single_task": "Slide open the ziploc bag.", **ALOHA_STATIC_INFO},
+    "aloha_sim_insertion_scripted": {"single_task": "Insert the peg into the socket.", **ALOHA_STATIC_INFO},
+    "aloha_sim_insertion_scripted_image": {
+        "single_task": "Insert the peg into the socket.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_sim_insertion_human": {"single_task": "Insert the peg into the socket.", **ALOHA_STATIC_INFO},
+    "aloha_sim_insertion_human_image": {
+        "single_task": "Insert the peg into the socket.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_sim_transfer_cube_scripted": {
+        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_sim_transfer_cube_scripted_image": {
+        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_sim_transfer_cube_human": {
+        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "aloha_sim_transfer_cube_human_image": {
+        "single_task": "Pick up the cube with the right arm and transfer it to the left arm.",
+        **ALOHA_STATIC_INFO,
+    },
+    "pusht": {"single_task": "Push the T-shaped block onto the T-shaped target.", **PUSHT_INFO},
+    "pusht_image": {"single_task": "Push the T-shaped block onto the T-shaped target.", **PUSHT_INFO},
+    "unitreeh1_fold_clothes": {"single_task": "Fold the sweatshirt.", **UNITREEH_INFO},
+    "unitreeh1_rearrange_objects": {"single_task": "Put the object into the bin.", **UNITREEH_INFO},
+    "unitreeh1_two_robot_greeting": {
+        "single_task": "Greet the other robot with a high five.",
+        **UNITREEH_INFO,
+    },
+    "unitreeh1_warehouse": {
+        "single_task": "Grab the spray paint on the shelf and place it in the bin on top of the robot dog.",
+        **UNITREEH_INFO,
+    },
+    "xarm_lift_medium": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
+    "xarm_lift_medium_image": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
+    "xarm_lift_medium_replay": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
+    "xarm_lift_medium_replay_image": {"single_task": "Pick up the cube and lift it.", **XARM_INFO},
+    "xarm_push_medium": {"single_task": "Push the cube onto the target.", **XARM_INFO},
+    "xarm_push_medium_image": {"single_task": "Push the cube onto the target.", **XARM_INFO},
+    "xarm_push_medium_replay": {"single_task": "Push the cube onto the target.", **XARM_INFO},
+    "xarm_push_medium_replay_image": {"single_task": "Push the cube onto the target.", **XARM_INFO},
+    "umi_cup_in_the_wild": {
+        "single_task": "Put the cup on the plate.",
+        "license": "apache-2.0",
+    },
+    "asu_table_top": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://link.springer.com/article/10.1007/s10514-023-10129-1",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{zhou2023modularity,
+                title={Modularity through Attention: Efficient Training and Transfer of Language-Conditioned Policies for Robot Manipulation},
+                author={Zhou, Yifan and Sonawani, Shubham and Phielipp, Mariano and Stepputtis, Simon and Amor, Heni},
+                booktitle={Conference on Robot Learning},
+                pages={1684--1695},
+                year={2023},
+                organization={PMLR}
+            }
+            @article{zhou2023learning,
+                title={Learning modular language-conditioned robot policies through attention},
+                author={Zhou, Yifan and Sonawani, Shubham and Phielipp, Mariano and Ben Amor, Heni and Stepputtis, Simon},
+                journal={Autonomous Robots},
+                pages={1--21},
+                year={2023},
+                publisher={Springer}
+            }""").lstrip(),
+    },
+    "austin_buds_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://ut-austin-rpl.github.io/BUDS-website/",
+        "paper": "https://huggingface.co/papers/2109.13841",
+        "citation_bibtex": dedent(r"""
+            @article{zhu2022bottom,
+                title={Bottom-Up Skill Discovery From Unsegmented Demonstrations for Long-Horizon Robot Manipulation},
+                author={Zhu, Yifeng and Stone, Peter and Zhu, Yuke},
+                journal={IEEE Robotics and Automation Letters},
+                volume={7},
+                number={2},
+                pages={4126--4133},
+                year={2022},
+                publisher={IEEE}
+            }""").lstrip(),
+    },
+    "austin_sailor_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://ut-austin-rpl.github.io/sailor/",
+        "paper": "https://huggingface.co/papers/2210.11435",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{nasiriany2022sailor,
+                title={Learning and Retrieval from Prior Data for Skill-based Imitation Learning},
+                author={Soroush Nasiriany and Tian Gao and Ajay Mandlekar and Yuke Zhu},
+                booktitle={Conference on Robot Learning (CoRL)},
+                year={2022}
+            }""").lstrip(),
+    },
+    "austin_sirius_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://ut-austin-rpl.github.io/sirius/",
+        "paper": "https://huggingface.co/papers/2211.08416",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{liu2022robot,
+                title = {Robot Learning on the Job: Human-in-the-Loop Autonomy and Learning During Deployment},
+                author = {Huihan Liu and Soroush Nasiriany and Lance Zhang and Zhiyao Bao and Yuke Zhu},
+                booktitle = {Robotics: Science and Systems (RSS)},
+                year = {2023}
+            }""").lstrip(),
+    },
+    "berkeley_autolab_ur5": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "url": "https://sites.google.com/view/berkeley-ur5/home",
+        "citation_bibtex": dedent(r"""
+            @misc{BerkeleyUR5Website,
+                title = {Berkeley {UR5} Demonstration Dataset},
+                author = {Lawrence Yunliang Chen and Simeon Adebola and Ken Goldberg},
+                howpublished = {https://sites.google.com/view/berkeley-ur5/home},
+            }""").lstrip(),
+    },
+    "berkeley_cable_routing": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "url": "https://sites.google.com/view/cablerouting/home",
+        "paper": "https://huggingface.co/papers/2307.08927",
+        "citation_bibtex": dedent(r"""
+            @article{luo2023multistage,
+                author    = {Jianlan Luo and Charles Xu and Xinyang Geng and Gilbert Feng and Kuan Fang and Liam Tan and Stefan Schaal and Sergey Levine},
+                title     = {Multi-Stage Cable Routing through Hierarchical Imitation Learning},
+                journal   = {arXiv pre-print},
+                year      = {2023},
+                url       = {https://huggingface.co/papers/2307.08927},
+            }""").lstrip(),
+    },
+    "berkeley_fanuc_manipulation": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://sites.google.com/berkeley.edu/fanuc-manipulation",
+        "citation_bibtex": dedent(r"""
+            @article{fanuc_manipulation2023,
+                title={Fanuc Manipulation: A Dataset for Learning-based Manipulation with FANUC Mate 200iD Robot},
+                author={Zhu, Xinghao and Tian, Ran and Xu, Chenfeng and Ding, Mingyu and Zhan, Wei and Tomizuka, Masayoshi},
+                year={2023},
+            }""").lstrip(),
+    },
+    "berkeley_gnm_cory_hall": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://huggingface.co/papers/1709.10489",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{kahn2018self,
+                title={Self-supervised deep reinforcement learning with generalized computation graphs for robot navigation},
+                author={Kahn, Gregory and Villaflor, Adam and Ding, Bosen and Abbeel, Pieter and Levine, Sergey},
+                booktitle={2018 IEEE international conference on robotics and automation (ICRA)},
+                pages={5129--5136},
+                year={2018},
+                organization={IEEE}
+            }""").lstrip(),
+    },
+    "berkeley_gnm_recon": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://sites.google.com/view/recon-robot",
+        "paper": "https://huggingface.co/papers/2104.05859",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{shah2021rapid,
+                title={Rapid Exploration for Open-World Navigation with Latent Goal Models},
+                author={Dhruv Shah and Benjamin Eysenbach and Nicholas Rhinehart and Sergey Levine},
+                booktitle={5th Annual Conference on Robot Learning },
+                year={2021},
+                url={https://openreview.net/forum?id=d_SWJhyKfVw}
+            }""").lstrip(),
+    },
+    "berkeley_gnm_sac_son": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://sites.google.com/view/SACSoN-review",
+        "paper": "https://huggingface.co/papers/2306.01874",
+        "citation_bibtex": dedent(r"""
+            @article{hirose2023sacson,
+                title={SACSoN: Scalable Autonomous Data Collection for Social Navigation},
+                author={Hirose, Noriaki and Shah, Dhruv and Sridhar, Ajay and Levine, Sergey},
+                journal={arXiv preprint arXiv:2306.01874},
+                year={2023}
+            }""").lstrip(),
+    },
+    "berkeley_mvp": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://huggingface.co/papers/2203.06173",
+        "citation_bibtex": dedent(r"""
+            @InProceedings{Radosavovic2022,
+                title = {Real-World Robot Learning with Masked Visual Pre-training},
+                author = {Ilija Radosavovic and Tete Xiao and Stephen James and Pieter Abbeel and Jitendra Malik and Trevor Darrell},
+                booktitle = {CoRL},
+                year = {2022}
+            }""").lstrip(),
+    },
+    "berkeley_rpt": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://huggingface.co/papers/2306.10007",
+        "citation_bibtex": dedent(r"""
+            @article{Radosavovic2023,
+                title={Robot Learning with Sensorimotor Pre-training},
+                author={Ilija Radosavovic and Baifeng Shi and Letian Fu and Ken Goldberg and Trevor Darrell and Jitendra Malik},
+                year={2023},
+                journal={arXiv:2306.10007}
+            }""").lstrip(),
+    },
+    "cmu_franka_exploration_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://human-world-model.github.io/",
+        "paper": "https://huggingface.co/papers/2308.10901",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{mendonca2023structured,
+                title={Structured World Models from Human Videos},
+                author={Mendonca, Russell  and Bahl, Shikhar and Pathak, Deepak},
+                journal={RSS},
+                year={2023}
+            }""").lstrip(),
+    },
+    "cmu_play_fusion": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://play-fusion.github.io/",
+        "paper": "https://huggingface.co/papers/2312.04549",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{chen2023playfusion,
+                title={PlayFusion: Skill Acquisition via Diffusion from Language-Annotated Play},
+                author={Chen, Lili and Bahl, Shikhar and Pathak, Deepak},
+                booktitle={CoRL},
+                year={2023}
+            }""").lstrip(),
+    },
+    "cmu_stretch": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://robo-affordances.github.io/",
+        "paper": "https://huggingface.co/papers/2304.08488",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{bahl2023affordances,
+                title={Affordances from Human Videos as a Versatile Representation for Robotics},
+                author={Bahl, Shikhar and Mendonca, Russell and Chen, Lili and Jain, Unnat and Pathak, Deepak},
+                booktitle={CVPR},
+                year={2023}
+            }
+                @article{mendonca2023structured,
+                title={Structured World Models from Human Videos},
+                author={Mendonca, Russell and Bahl, Shikhar and Pathak, Deepak},
+                journal={CoRL},
+                year={2023}
+            }""").lstrip(),
+    },
+    "columbia_cairlab_pusht_real": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://diffusion-policy.cs.columbia.edu/",
+        "paper": "https://huggingface.co/papers/2303.04137",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{chi2023diffusionpolicy,
+                title={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
+                author={Chi, Cheng and Feng, Siyuan and Du, Yilun and Xu, Zhenjia and Cousineau, Eric and Burchfiel, Benjamin and Song, Shuran},
+                booktitle={Proceedings of Robotics: Science and Systems (RSS)},
+                year={2023}
+            }""").lstrip(),
+    },
+    "conq_hose_manipulation": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://sites.google.com/view/conq-hose-manipulation-dataset/home",
+        "citation_bibtex": dedent(r"""
+            @misc{ConqHoseManipData,
+                author={Peter Mitrano and Dmitry Berenson},
+                title={Conq Hose Manipulation Dataset, v1.15.0},
+                year={2024},
+                howpublished={https://sites.google.com/view/conq-hose-manipulation-dataset}
+            }""").lstrip(),
+    },
+    "dlr_edan_shared_control": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://ieeexplore.ieee.org/document/9341156",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{vogel_edan_2020,
+                title = {EDAN - an EMG-Controlled Daily Assistant to Help People with Physical Disabilities},
+                language = {en},
+                booktitle = {2020 {IEEE}/{RSJ} {International} {Conference} on {Intelligent} {Robots} and {Systems} ({IROS})},
+                author = {Vogel, Jörn and Hagengruber, Annette and Iskandar, Maged and Quere, Gabriel and Leipscher, Ulrike and Bustamante, Samuel and Dietrich, Alexander and Hoeppner, Hannes and Leidner, Daniel and Albu-Schäffer, Alin},
+                year = {2020}
+            }
+            @inproceedings{quere_shared_2020,
+                address = {Paris, France},
+                title = {Shared {Control} {Templates} for {Assistive} {Robotics}},
+                language = {en},
+                booktitle = {2020 {IEEE} {International} {Conference} on {Robotics} and {Automation} ({ICRA})},
+                author = {Quere, Gabriel and Hagengruber, Annette and Iskandar, Maged and Bustamante, Samuel and Leidner, Daniel and Stulp, Freek and Vogel, Joern},
+                year = {2020},
+                pages = {7},
+            }""").lstrip(),
+    },
+    "dlr_sara_grid_clamp": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://www.researchsquare.com/article/rs-3289569/v1",
+        "citation_bibtex": dedent(r"""
+            @article{padalkar2023guided,
+                title={A guided reinforcement learning approach using shared control templates for learning manipulation skills in the real world},
+                author={Padalkar, Abhishek and Quere, Gabriel and Raffin, Antonin and Silv{\'e}rio, Jo{\~a}o and Stulp, Freek},
+                journal={Research square preprint rs-3289569/v1},
+                year={2023}
+            }""").lstrip(),
+    },
+    "dlr_sara_pour": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "paper": "https://elib.dlr.de/193739/1/padalkar2023rlsct.pdf",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{padalkar2023guiding,
+                title={Guiding Reinforcement Learning with Shared Control Templates},
+                author={Padalkar, Abhishek and Quere, Gabriel and Steinmetz, Franz and Raffin, Antonin and Nieuwenhuisen, Matthias and Silv{\'e}rio, Jo{\~a}o and Stulp, Freek},
+                booktitle={40th IEEE International Conference on Robotics and Automation, ICRA 2023},
+                year={2023},
+                organization={IEEE}
+            }""").lstrip(),
+    },
+    "droid_100": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://droid-dataset.github.io/",
+        "paper": "https://huggingface.co/papers/2403.12945",
+        "citation_bibtex": dedent(r"""
+            @article{khazatsky2024droid,
+                title   = {DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset},
+                author  = {Alexander Khazatsky and Karl Pertsch and Suraj Nair and Ashwin Balakrishna and Sudeep Dasari and Siddharth Karamcheti and Soroush Nasiriany and Mohan Kumar Srirama and Lawrence Yunliang Chen and Kirsty Ellis and Peter David Fagan and Joey Hejna and Masha Itkina and Marion Lepert and Yecheng Jason Ma and Patrick Tree Miller and Jimmy Wu and Suneel Belkhale and Shivin Dass and Huy Ha and Arhan Jain and Abraham Lee and Youngwoon Lee and Marius Memmel and Sungjae Park and Ilija Radosavovic and Kaiyuan Wang and Albert Zhan and Kevin Black and Cheng Chi and Kyle Beltran Hatch and Shan Lin and Jingpei Lu and Jean Mercat and Abdul Rehman and Pannag R Sanketi and Archit Sharma and Cody Simpson and Quan Vuong and Homer Rich Walke and Blake Wulfe and Ted Xiao and Jonathan Heewon Yang and Arefeh Yavary and Tony Z. Zhao and Christopher Agia and Rohan Baijal and Mateo Guaman Castro and Daphne Chen and Qiuyu Chen and Trinity Chung and Jaimyn Drake and Ethan Paul Foster and Jensen Gao and David Antonio Herrera and Minho Heo and Kyle Hsu and Jiaheng Hu and Donovon Jackson and Charlotte Le and Yunshuang Li and Kevin Lin and Roy Lin and Zehan Ma and Abhiram Maddukuri and Suvir Mirchandani and Daniel Morton and Tony Nguyen and Abigail O'Neill and Rosario Scalise and Derick Seale and Victor Son and Stephen Tian and Emi Tran and Andrew E. Wang and Yilin Wu and Annie Xie and Jingyun Yang and Patrick Yin and Yunchu Zhang and Osbert Bastani and Glen Berseth and Jeannette Bohg and Ken Goldberg and Abhinav Gupta and Abhishek Gupta and Dinesh Jayaraman and Joseph J Lim and Jitendra Malik and Roberto Martín-Martín and Subramanian Ramamoorthy and Dorsa Sadigh and Shuran Song and Jiajun Wu and Michael C. Yip and Yuke Zhu and Thomas Kollar and Sergey Levine and Chelsea Finn},
+                year    = {2024},
+            }""").lstrip(),
+    },
+    "fmb": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "url": "https://functional-manipulation-benchmark.github.io/",
+        "paper": "https://huggingface.co/papers/2401.08553",
+        "citation_bibtex": dedent(r"""
+            @article{luo2024fmb,
+                title={FMB: a Functional Manipulation Benchmark for Generalizable Robotic Learning},
+                author={Luo, Jianlan and Xu, Charles and Liu, Fangchen and Tan, Liam and Lin, Zipeng and Wu, Jeffrey and Abbeel, Pieter and Levine, Sergey},
+                journal={arXiv preprint arXiv:2401.08553},
+                year={2024}
+            }""").lstrip(),
+    },
+    "iamlab_cmu_pickup_insert": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://openreview.net/forum?id=WuBv9-IGDUA",
+        "paper": "https://huggingface.co/papers/2401.14502",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{saxena2023multiresolution,
+                title={Multi-Resolution Sensing for Real-Time Control with Vision-Language Models},
+                author={Saumya Saxena and Mohit Sharma and Oliver Kroemer},
+                booktitle={7th Annual Conference on Robot Learning},
+                year={2023},
+                url={https://openreview.net/forum?id=WuBv9-IGDUA}
+            }""").lstrip(),
+    },
+    "imperialcollege_sawyer_wrist_cam": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+    },
+    "jaco_play": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "url": "https://github.com/clvrai/clvr_jaco_play_dataset",
+        "citation_bibtex": dedent(r"""
+            @software{dass2023jacoplay,
+                author = {Dass, Shivin and Yapeter, Jullian and Zhang, Jesse and Zhang, Jiahui
+                            and Pertsch, Karl and Nikolaidis, Stefanos and Lim, Joseph J.},
+                title = {CLVR Jaco Play Dataset},
+                url = {https://github.com/clvrai/clvr_jaco_play_dataset},
+                version = {1.0.0},
+                year = {2023}
+            }""").lstrip(),
+    },
+    "kaist_nonprehensile": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "url": "https://github.com/JaeHyung-Kim/rlds_dataset_builder",
+        "citation_bibtex": dedent(r"""
+            @article{kimpre,
+                title={Pre-and post-contact policy decomposition for non-prehensile manipulation with zero-shot sim-to-real transfer},
+                author={Kim, Minchan and Han, Junhyek and Kim, Jaehyung and Kim, Beomjoon},
+                booktitle={2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
+                year={2023},
+                organization={IEEE}
+            }""").lstrip(),
+    },
+    "nyu_door_opening_surprising_effectiveness": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://jyopari.github.io/VINN/",
+        "paper": "https://huggingface.co/papers/2112.01511",
+        "citation_bibtex": dedent(r"""
+            @misc{pari2021surprising,
+                title={The Surprising Effectiveness of Representation Learning for Visual Imitation},
+                author={Jyothish Pari and Nur Muhammad Shafiullah and Sridhar Pandian Arunachalam and Lerrel Pinto},
+                year={2021},
+                eprint={2112.01511},
+                archivePrefix={arXiv},
+                primaryClass={cs.RO}
+            }""").lstrip(),
+    },
+    "nyu_franka_play_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://play-to-policy.github.io/",
+        "paper": "https://huggingface.co/papers/2210.10047",
+        "citation_bibtex": dedent(r"""
+            @article{cui2022play,
+                title   = {From Play to Policy: Conditional Behavior Generation from Uncurated Robot Data},
+                author  = {Cui, Zichen Jeff and Wang, Yibin and Shafiullah, Nur Muhammad Mahi and Pinto, Lerrel},
+                journal = {arXiv preprint arXiv:2210.10047},
+                year    = {2022}
+            }""").lstrip(),
+    },
+    "nyu_rot_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://rot-robot.github.io/",
+        "paper": "https://huggingface.co/papers/2206.15469",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{haldar2023watch,
+                title={Watch and match: Supercharging imitation with regularized optimal transport},
+                author={Haldar, Siddhant and Mathur, Vaibhav and Yarats, Denis and Pinto, Lerrel},
+                booktitle={Conference on Robot Learning},
+                pages={32--43},
+                year={2023},
+                organization={PMLR}
+            }""").lstrip(),
+    },
+    "roboturk": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://roboturk.stanford.edu/dataset_real.html",
+        "paper": "PAPER",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{mandlekar2019scaling,
+                title={Scaling robot supervision to hundreds of hours with roboturk: Robotic manipulation dataset through human reasoning and dexterity},
+                author={Mandlekar, Ajay and Booher, Jonathan and Spero, Max and Tung, Albert and Gupta, Anchit and Zhu, Yuke and Garg, Animesh and Savarese, Silvio and Fei-Fei, Li},
+                booktitle={2019 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
+                pages={1048--1055},
+                year={2019},
+                organization={IEEE}
+            }""").lstrip(),
+    },
+    "stanford_hydra_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://sites.google.com/view/hydra-il-2023",
+        "paper": "https://huggingface.co/papers/2306.17237",
+        "citation_bibtex": dedent(r"""
+            @article{belkhale2023hydra,
+                title={HYDRA: Hybrid Robot Actions for Imitation Learning},
+                author={Belkhale, Suneel and Cui, Yuchen and Sadigh, Dorsa},
+                journal={arxiv},
+                year={2023}
+            }""").lstrip(),
+    },
+    "stanford_kuka_multimodal_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://sites.google.com/view/visionandtouch",
+        "paper": "https://huggingface.co/papers/1810.10191",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{lee2019icra,
+                title={Making sense of vision and touch: Self-supervised learning of multimodal representations for contact-rich tasks},
+                author={Lee, Michelle A and Zhu, Yuke and Srinivasan, Krishnan and Shah, Parth and Savarese, Silvio and Fei-Fei, Li and  Garg, Animesh and Bohg, Jeannette},
+                booktitle={2019 IEEE International Conference on Robotics and Automation (ICRA)},
+                year={2019},
+                url={https://huggingface.co/papers/1810.10191}
+            }""").lstrip(),
+    },
+    "stanford_robocook": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://hshi74.github.io/robocook/",
+        "paper": "https://huggingface.co/papers/2306.14447",
+        "citation_bibtex": dedent(r"""
+            @article{shi2023robocook,
+                title={RoboCook: Long-Horizon Elasto-Plastic Object Manipulation with Diverse Tools},
+                author={Shi, Haochen and Xu, Huazhe and Clarke, Samuel and Li, Yunzhu and Wu, Jiajun},
+                journal={arXiv preprint arXiv:2306.14447},
+                year={2023}
+            }""").lstrip(),
+    },
+    "taco_play": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "url": "https://www.kaggle.com/datasets/oiermees/taco-robot",
+        "paper": "https://huggingface.co/papers/2209.08959, https://huggingface.co/papers/2210.01911",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{rosete2022tacorl,
+                author = {Erick Rosete-Beas and Oier Mees and Gabriel Kalweit and Joschka Boedecker and Wolfram Burgard},
+                title = {Latent Plans for Task Agnostic Offline Reinforcement Learning},
+                journal = {Proceedings of the 6th Conference on Robot Learning (CoRL)},
+                year = {2022}
+            }
+            @inproceedings{mees23hulc2,
+                title={Grounding  Language  with  Visual  Affordances  over  Unstructured  Data},
+                author={Oier Mees and Jessica Borja-Diaz and Wolfram Burgard},
+                booktitle = {Proceedings of the IEEE International Conference on Robotics and Automation (ICRA)},
+                year={2023},
+                address = {London, UK}
+            }""").lstrip(),
+    },
+    "tokyo_u_lsmo": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "URL",
+        "paper": "https://huggingface.co/papers/2107.05842",
+        "citation_bibtex": dedent(r"""
+            @Article{Osa22,
+                author  = {Takayuki Osa},
+                journal = {The International Journal of Robotics Research},
+                title   = {Motion Planning by Learning the Solution Manifold in Trajectory Optimization},
+                year    = {2022},
+                number  = {3},
+                pages   = {291--311},
+                volume  = {41},
+            }""").lstrip(),
+    },
+    "toto": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://toto-benchmark.org/",
+        "paper": "https://huggingface.co/papers/2306.00942",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{zhou2023train,
+                author={Zhou, Gaoyue and Dean, Victoria and Srirama, Mohan Kumar and Rajeswaran, Aravind and Pari, Jyothish and Hatch, Kyle and Jain, Aryan and Yu, Tianhe and Abbeel, Pieter and Pinto, Lerrel and Finn, Chelsea and Gupta, Abhinav},
+                booktitle={2023 IEEE International Conference on Robotics and Automation (ICRA)},
+                title={Train Offline, Test Online: A Real Robot Learning Benchmark},
+                year={2023},
+            }""").lstrip(),
+    },
+    "ucsd_kitchen_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "citation_bibtex": dedent(r"""
+            @ARTICLE{ucsd_kitchens,
+                author = {Ge Yan, Kris Wu, and Xiaolong Wang},
+                title = {{ucsd kitchens Dataset}},
+                year = {2023},
+                month = {August}
+            }""").lstrip(),
+    },
+    "ucsd_pick_and_place_dataset": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://owmcorl.github.io/#",
+        "paper": "https://huggingface.co/papers/2310.16029",
+        "citation_bibtex": dedent(r"""
+            @preprint{Feng2023Finetuning,
+                title={Finetuning Offline World Models in the Real World},
+                author={Yunhai Feng, Nicklas Hansen, Ziyan Xiong, Chandramouli Rajagopalan, Xiaolong Wang},
+                year={2023}
+            }""").lstrip(),
+    },
+    "uiuc_d3field": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://robopil.github.io/d3fields/",
+        "paper": "https://huggingface.co/papers/2309.16118",
+        "citation_bibtex": dedent(r"""
+            @article{wang2023d3field,
+                title={D^3Field: Dynamic 3D Descriptor Fields for Generalizable Robotic Manipulation},
+                author={Wang, Yixuan and Li, Zhuoran and Zhang, Mingtong and Driggs-Campbell, Katherine and Wu, Jiajun and Fei-Fei, Li and Li, Yunzhu},
+                journal={arXiv preprint arXiv:},
+                year={2023},
+            }""").lstrip(),
+    },
+    "usc_cloth_sim": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://uscresl.github.io/dmfd/",
+        "paper": "https://huggingface.co/papers/2207.10148",
+        "citation_bibtex": dedent(r"""
+            @article{salhotra2022dmfd,
+                author={Salhotra, Gautam and Liu, I-Chun Arthur and Dominguez-Kuhne, Marcus and Sukhatme, Gaurav S.},
+                journal={IEEE Robotics and Automation Letters},
+                title={Learning Deformable Object Manipulation From Expert Demonstrations},
+                year={2022},
+                volume={7},
+                number={4},
+                pages={8775-8782},
+                doi={10.1109/LRA.2022.3187843}
+            }""").lstrip(),
+    },
+    "utaustin_mutex": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://ut-austin-rpl.github.io/MUTEX/",
+        "paper": "https://huggingface.co/papers/2309.14320",
+        "citation_bibtex": dedent(r"""
+            @inproceedings{shah2023mutex,
+                title={{MUTEX}: Learning Unified Policies from Multimodal Task Specifications},
+                author={Rutav Shah and Roberto Mart{\'\i}n-Mart{\'\i}n and Yuke Zhu},
+                booktitle={7th Annual Conference on Robot Learning},
+                year={2023},
+                url={https://openreview.net/forum?id=PwqiqaaEzJ}
+            }""").lstrip(),
+    },
+    "utokyo_pr2_opening_fridge": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "citation_bibtex": dedent(r"""
+            @misc{oh2023pr2utokyodatasets,
+                author={Jihoon Oh and Naoaki Kanazawa and Kento Kawaharazuka},
+                title={X-Embodiment U-Tokyo PR2 Datasets},
+                year={2023},
+                url={https://github.com/ojh6404/rlds_dataset_builder},
+            }""").lstrip(),
+    },
+    "utokyo_pr2_tabletop_manipulation": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "citation_bibtex": dedent(r"""
+            @misc{oh2023pr2utokyodatasets,
+                author={Jihoon Oh and Naoaki Kanazawa and Kento Kawaharazuka},
+                title={X-Embodiment U-Tokyo PR2 Datasets},
+                year={2023},
+                url={https://github.com/ojh6404/rlds_dataset_builder},
+            }""").lstrip(),
+    },
+    "utokyo_saytap": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://saytap.github.io/",
+        "paper": "https://huggingface.co/papers/2306.07580",
+        "citation_bibtex": dedent(r"""
+            @article{saytap2023,
+                author = {Yujin Tang and Wenhao Yu and Jie Tan and Heiga Zen and Aleksandra Faust and
+                Tatsuya Harada},
+                title  = {SayTap: Language to Quadrupedal Locomotion},
+                eprint = {arXiv:2306.07580},
+                url    = {https://saytap.github.io},
+                note   = {https://saytap.github.io},
+                year   = {2023}
+            }""").lstrip(),
+    },
+    "utokyo_xarm_bimanual": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "citation_bibtex": dedent(r"""
+            @misc{matsushima2023weblab,
+                title={Weblab xArm Dataset},
+                author={Tatsuya Matsushima and Hiroki Furuta and Yusuke Iwasawa and Yutaka Matsuo},
+                year={2023},
+            }""").lstrip(),
+    },
+    "utokyo_xarm_pick_and_place": {
+        "tasks_col": "language_instruction",
+        "license": "cc-by-4.0",
+        "citation_bibtex": dedent(r"""
+            @misc{matsushima2023weblab,
+                title={Weblab xArm Dataset},
+                author={Tatsuya Matsushima and Hiroki Furuta and Yusuke Iwasawa and Yutaka Matsuo},
+                year={2023},
+            }""").lstrip(),
+    },
+    "viola": {
+        "tasks_col": "language_instruction",
+        "license": "mit",
+        "url": "https://ut-austin-rpl.github.io/VIOLA/",
+        "paper": "https://huggingface.co/papers/2210.11339",
+        "citation_bibtex": dedent(r"""
+            @article{zhu2022viola,
+                title={VIOLA: Imitation Learning for Vision-Based Manipulation with Object Proposal Priors},
+                author={Zhu, Yifeng and Joshi, Abhishek and Stone, Peter and Zhu, Yuke},
+                journal={6th Annual Conference on Robot Learning (CoRL)},
+                year={2022}
+            }""").lstrip(),
+    },
+}
+# spellchecker:on
+
+
+def batch_convert():
+    status = {}
+    logfile = LOCAL_DIR / "conversion_log.txt"
+    assert set(DATASETS) == {id_.split("/")[1] for id_ in available_datasets}
+    for num, (name, kwargs) in enumerate(DATASETS.items()):
+        repo_id = f"lerobot/{name}"
+        print(f"\nConverting {repo_id} ({num}/{len(DATASETS)})")
+        print("---------------------------------------------------------")
+        try:
+            convert_dataset(repo_id, LOCAL_DIR, **kwargs)
+            status = f"{repo_id}: success."
+            with open(logfile, "a") as file:
+                file.write(status + "\n")
+        except Exception:
+            status = f"{repo_id}: failed\n    {traceback.format_exc()}"
+            with open(logfile, "a") as file:
+                file.write(status + "\n")
+            continue
+
+
+if __name__ == "__main__":
+    batch_convert()

src/lerobot/datasets/v2/convert_dataset_v1_to_v2.py

@@ -0,0 +1,687 @@
+#!/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.
+
+"""
+This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 1.6 to
+2.0. You will be required to provide the 'tasks', which is a short but accurate description in plain English
+for each of the task performed in the dataset. This will allow to easily train models with task-conditioning.
+
+We support 3 different scenarios for these tasks (see instructions below):
+    1. Single task dataset: all episodes of your dataset have the same single task.
+    2. Single task episodes: the episodes of your dataset each contain a single task but they can differ from
+      one episode to the next.
+    3. Multi task episodes: episodes of your dataset may each contain several different tasks.
+
+
+Can you can also provide a robot config .yaml file (not mandatory) to this script via the option
+'--robot-config' so that it writes information about the robot (robot type, motors names) this dataset was
+recorded with. For now, only Aloha/Koch type robots are supported with this option.
+
+
+# 1. Single task dataset
+If your dataset contains a single task, you can simply provide it directly via the CLI with the
+'--single-task' option.
+
+Examples:
+
+```bash
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
+    --repo-id lerobot/aloha_sim_insertion_human_image \
+    --single-task "Insert the peg into the socket." \
+    --robot-config lerobot/configs/robot/aloha.yaml \
+    --local-dir data
+```
+
+```bash
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
+    --repo-id aliberts/koch_tutorial \
+    --single-task "Pick the Lego block and drop it in the box on the right." \
+    --robot-config lerobot/configs/robot/koch.yaml \
+    --local-dir data
+```
+
+
+# 2. Single task episodes
+If your dataset is a multi-task dataset, you have two options to provide the tasks to this script:
+
+- If your dataset already contains a language instruction column in its parquet file, you can simply provide
+  this column's name with the '--tasks-col' arg.
+
+    Example:
+
+    ```bash
+    python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
+        --repo-id lerobot/stanford_kuka_multimodal_dataset \
+        --tasks-col "language_instruction" \
+        --local-dir data
+    ```
+
+- If your dataset doesn't contain a language instruction, you should provide the path to a .json file with the
+  '--tasks-path' arg. This file should have the following structure where keys correspond to each
+  episode_index in the dataset, and values are the language instruction for that episode.
+
+    Example:
+
+    ```json
+    {
+        "0": "Do something",
+        "1": "Do something else",
+        "2": "Do something",
+        "3": "Go there",
+        ...
+    }
+    ```
+
+# 3. Multi task episodes
+If you have multiple tasks per episodes, your dataset should contain a language instruction column in its
+parquet file, and you must provide this column's name with the '--tasks-col' arg.
+
+Example:
+
+```bash
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
+    --repo-id lerobot/stanford_kuka_multimodal_dataset \
+    --tasks-col "language_instruction" \
+    --local-dir data
+```
+"""
+
+import argparse
+import contextlib
+import filecmp
+import json
+import logging
+import math
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+import datasets
+import pyarrow.compute as pc
+import pyarrow.parquet as pq
+import torch
+from datasets import Dataset
+from huggingface_hub import HfApi
+from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
+from safetensors.torch import load_file
+
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    INFO_PATH,
+    LEGACY_DEFAULT_PARQUET_PATH,
+    LEGACY_DEFAULT_VIDEO_PATH,
+    LEGACY_EPISODES_PATH,
+    LEGACY_TASKS_PATH,
+    STATS_PATH,
+    create_branch,
+    create_lerobot_dataset_card,
+    flatten_dict,
+    get_safe_version,
+    load_json,
+    unflatten_dict,
+    write_json,
+    write_jsonlines,
+)
+from lerobot.datasets.video_utils import (
+    VideoFrame,  # noqa: F401
+    get_image_pixel_channels,
+    get_video_info,
+)
+from lerobot.robots import RobotConfig
+
+V16 = "v1.6"
+V20 = "v2.0"
+
+GITATTRIBUTES_REF = "aliberts/gitattributes_reference"
+V1_VIDEO_FILE = "{video_key}_episode_{episode_index:06d}.mp4"
+V1_INFO_PATH = "meta_data/info.json"
+V1_STATS_PATH = "meta_data/stats.safetensors"
+
+
+def parse_robot_config(robot_cfg: RobotConfig) -> tuple[str, dict]:
+    if robot_cfg.type in ["aloha", "koch"]:
+        state_names = [
+            f"{arm}_{motor}" if len(robot_cfg.follower_arms) > 1 else motor
+            for arm in robot_cfg.follower_arms
+            for motor in robot_cfg.follower_arms[arm].motors
+        ]
+        action_names = [
+            # f"{arm}_{motor}" for arm in ["left", "right"] for motor in robot_cfg["leader_arms"][arm]["motors"]
+            f"{arm}_{motor}" if len(robot_cfg.leader_arms) > 1 else motor
+            for arm in robot_cfg.leader_arms
+            for motor in robot_cfg.leader_arms[arm].motors
+        ]
+    # elif robot_cfg["robot_type"] == "stretch3": TODO
+    else:
+        raise NotImplementedError(
+            "Please provide robot_config={'robot_type': ..., 'names': ...} directly to convert_dataset()."
+        )
+
+    return {
+        "robot_type": robot_cfg.type,
+        "names": {
+            "observation.state": state_names,
+            "observation.effort": state_names,
+            "action": action_names,
+        },
+    }
+
+
+def convert_stats_to_json(v1_dir: Path, v2_dir: Path) -> None:
+    safetensor_path = v1_dir / V1_STATS_PATH
+    stats = load_file(safetensor_path)
+    serialized_stats = {key: value.tolist() for key, value in stats.items()}
+    serialized_stats = unflatten_dict(serialized_stats)
+
+    json_path = v2_dir / STATS_PATH
+    json_path.parent.mkdir(exist_ok=True, parents=True)
+    with open(json_path, "w") as f:
+        json.dump(serialized_stats, f, indent=4)
+
+    # Sanity check
+    with open(json_path) as f:
+        stats_json = json.load(f)
+
+    stats_json = flatten_dict(stats_json)
+    stats_json = {key: torch.tensor(value) for key, value in stats_json.items()}
+    for key in stats:
+        torch.testing.assert_close(stats_json[key], stats[key])
+
+
+def get_features_from_hf_dataset(
+    dataset: Dataset, robot_config: RobotConfig | None = None
+) -> dict[str, list]:
+    robot_config = parse_robot_config(robot_config)
+    features = {}
+    for key, ft in dataset.features.items():
+        if isinstance(ft, datasets.Value):
+            dtype = ft.dtype
+            shape = (1,)
+            names = None
+        if isinstance(ft, datasets.Sequence):
+            assert isinstance(ft.feature, datasets.Value)
+            dtype = ft.feature.dtype
+            shape = (ft.length,)
+            motor_names = (
+                robot_config["names"][key] if robot_config else [f"motor_{i}" for i in range(ft.length)]
+            )
+            assert len(motor_names) == shape[0]
+            names = {"motors": motor_names}
+        elif isinstance(ft, datasets.Image):
+            dtype = "image"
+            image = dataset[0][key]  # Assuming first row
+            channels = get_image_pixel_channels(image)
+            shape = (image.height, image.width, channels)
+            names = ["height", "width", "channels"]
+        elif ft._type == "VideoFrame":
+            dtype = "video"
+            shape = None  # Add shape later
+            names = ["height", "width", "channels"]
+
+        features[key] = {
+            "dtype": dtype,
+            "shape": shape,
+            "names": names,
+        }
+
+    return features
+
+
+def add_task_index_by_episodes(dataset: Dataset, tasks_by_episodes: dict) -> tuple[Dataset, list[str]]:
+    df = dataset.to_pandas()
+    tasks = list(set(tasks_by_episodes.values()))
+    tasks_to_task_index = {task: task_idx for task_idx, task in enumerate(tasks)}
+    episodes_to_task_index = {ep_idx: tasks_to_task_index[task] for ep_idx, task in tasks_by_episodes.items()}
+    df["task_index"] = df["episode_index"].map(episodes_to_task_index).astype(int)
+
+    features = dataset.features
+    features["task_index"] = datasets.Value(dtype="int64")
+    dataset = Dataset.from_pandas(df, features=features, split="train")
+    return dataset, tasks
+
+
+def add_task_index_from_tasks_col(
+    dataset: Dataset, tasks_col: str
+) -> tuple[Dataset, dict[str, list[str]], list[str]]:
+    df = dataset.to_pandas()
+
+    # HACK: This is to clean some of the instructions in our version of Open X datasets
+    prefix_to_clean = "tf.Tensor(b'"
+    suffix_to_clean = "', shape=(), dtype=string)"
+    df[tasks_col] = df[tasks_col].str.removeprefix(prefix_to_clean).str.removesuffix(suffix_to_clean)
+
+    # Create task_index col
+    tasks_by_episode = df.groupby("episode_index")[tasks_col].unique().apply(lambda x: x.tolist()).to_dict()
+    tasks = df[tasks_col].unique().tolist()
+    tasks_to_task_index = {task: idx for idx, task in enumerate(tasks)}
+    df["task_index"] = df[tasks_col].map(tasks_to_task_index).astype(int)
+
+    # Build the dataset back from df
+    features = dataset.features
+    features["task_index"] = datasets.Value(dtype="int64")
+    dataset = Dataset.from_pandas(df, features=features, split="train")
+    dataset = dataset.remove_columns(tasks_col)
+
+    return dataset, tasks, tasks_by_episode
+
+
+def split_parquet_by_episodes(
+    dataset: Dataset,
+    total_episodes: int,
+    total_chunks: int,
+    output_dir: Path,
+) -> list:
+    table = dataset.data.table
+    episode_lengths = []
+    for ep_chunk in range(total_chunks):
+        ep_chunk_start = DEFAULT_CHUNK_SIZE * ep_chunk
+        ep_chunk_end = min(DEFAULT_CHUNK_SIZE * (ep_chunk + 1), total_episodes)
+        chunk_dir = "/".join(LEGACY_DEFAULT_PARQUET_PATH.split("/")[:-1]).format(episode_chunk=ep_chunk)
+        (output_dir / chunk_dir).mkdir(parents=True, exist_ok=True)
+        for ep_idx in range(ep_chunk_start, ep_chunk_end):
+            ep_table = table.filter(pc.equal(table["episode_index"], ep_idx))
+            episode_lengths.insert(ep_idx, len(ep_table))
+            output_file = output_dir / LEGACY_DEFAULT_PARQUET_PATH.format(
+                episode_chunk=ep_chunk, episode_index=ep_idx
+            )
+            pq.write_table(ep_table, output_file)
+
+    return episode_lengths
+
+
+def move_videos(
+    repo_id: str,
+    video_keys: list[str],
+    total_episodes: int,
+    total_chunks: int,
+    work_dir: Path,
+    clean_gittatributes: Path,
+    branch: str = "main",
+) -> None:
+    """
+    HACK: Since HfApi() doesn't provide a way to move files directly in a repo, this function will run git
+    commands to fetch git lfs video files references to move them into subdirectories without having to
+    actually download them.
+    """
+    _lfs_clone(repo_id, work_dir, branch)
+
+    videos_moved = False
+    video_files = [str(f.relative_to(work_dir)) for f in work_dir.glob("videos*/*.mp4")]
+    if len(video_files) == 0:
+        video_files = [str(f.relative_to(work_dir)) for f in work_dir.glob("videos*/*/*/*.mp4")]
+        videos_moved = True  # Videos have already been moved
+
+    assert len(video_files) == total_episodes * len(video_keys)
+
+    lfs_untracked_videos = _get_lfs_untracked_videos(work_dir, video_files)
+
+    current_gittatributes = work_dir / ".gitattributes"
+    if not filecmp.cmp(current_gittatributes, clean_gittatributes, shallow=False):
+        fix_gitattributes(work_dir, current_gittatributes, clean_gittatributes)
+
+    if lfs_untracked_videos:
+        fix_lfs_video_files_tracking(work_dir, video_files)
+
+    if videos_moved:
+        return
+
+    video_dirs = sorted(work_dir.glob("videos*/"))
+    for ep_chunk in range(total_chunks):
+        ep_chunk_start = DEFAULT_CHUNK_SIZE * ep_chunk
+        ep_chunk_end = min(DEFAULT_CHUNK_SIZE * (ep_chunk + 1), total_episodes)
+        for vid_key in video_keys:
+            chunk_dir = "/".join(LEGACY_DEFAULT_VIDEO_PATH.split("/")[:-1]).format(
+                episode_chunk=ep_chunk, video_key=vid_key
+            )
+            (work_dir / chunk_dir).mkdir(parents=True, exist_ok=True)
+
+            for ep_idx in range(ep_chunk_start, ep_chunk_end):
+                target_path = LEGACY_DEFAULT_VIDEO_PATH.format(
+                    episode_chunk=ep_chunk, video_key=vid_key, episode_index=ep_idx
+                )
+                video_file = V1_VIDEO_FILE.format(video_key=vid_key, episode_index=ep_idx)
+                if len(video_dirs) == 1:
+                    video_path = video_dirs[0] / video_file
+                else:
+                    for dir in video_dirs:
+                        if (dir / video_file).is_file():
+                            video_path = dir / video_file
+                            break
+
+                video_path.rename(work_dir / target_path)
+
+    commit_message = "Move video files into chunk subdirectories"
+    subprocess.run(["git", "add", "."], cwd=work_dir, check=True)
+    subprocess.run(["git", "commit", "-m", commit_message], cwd=work_dir, check=True)
+    subprocess.run(["git", "push"], cwd=work_dir, check=True)
+
+
+def fix_lfs_video_files_tracking(work_dir: Path, lfs_untracked_videos: list[str]) -> None:
+    """
+    HACK: This function fixes the tracking by git lfs which was not properly set on some repos. In that case,
+    there's no other option than to download the actual files and reupload them with lfs tracking.
+    """
+    for i in range(0, len(lfs_untracked_videos), 100):
+        files = lfs_untracked_videos[i : i + 100]
+        try:
+            subprocess.run(["git", "rm", "--cached", *files], cwd=work_dir, capture_output=True, check=True)
+        except subprocess.CalledProcessError as e:
+            print("git rm --cached ERROR:")
+            print(e.stderr)
+        subprocess.run(["git", "add", *files], cwd=work_dir, check=True)
+
+    commit_message = "Track video files with git lfs"
+    subprocess.run(["git", "commit", "-m", commit_message], cwd=work_dir, check=True)
+    subprocess.run(["git", "push"], cwd=work_dir, check=True)
+
+
+def fix_gitattributes(work_dir: Path, current_gittatributes: Path, clean_gittatributes: Path) -> None:
+    shutil.copyfile(clean_gittatributes, current_gittatributes)
+    subprocess.run(["git", "add", ".gitattributes"], cwd=work_dir, check=True)
+    subprocess.run(["git", "commit", "-m", "Fix .gitattributes"], cwd=work_dir, check=True)
+    subprocess.run(["git", "push"], cwd=work_dir, check=True)
+
+
+def _lfs_clone(repo_id: str, work_dir: Path, branch: str) -> None:
+    subprocess.run(["git", "lfs", "install"], cwd=work_dir, check=True)
+    repo_url = f"https://huggingface.co/datasets/{repo_id}"
+    env = {"GIT_LFS_SKIP_SMUDGE": "1"}  # Prevent downloading LFS files
+    subprocess.run(
+        ["git", "clone", "--branch", branch, "--single-branch", "--depth", "1", repo_url, str(work_dir)],
+        check=True,
+        env=env,
+    )
+
+
+def _get_lfs_untracked_videos(work_dir: Path, video_files: list[str]) -> list[str]:
+    lfs_tracked_files = subprocess.run(
+        ["git", "lfs", "ls-files", "-n"], cwd=work_dir, capture_output=True, text=True, check=True
+    )
+    lfs_tracked_files = set(lfs_tracked_files.stdout.splitlines())
+    return [f for f in video_files if f not in lfs_tracked_files]
+
+
+def get_videos_info(repo_id: str, local_dir: Path, video_keys: list[str], branch: str) -> dict:
+    # Assumes first episode
+    video_files = [
+        LEGACY_DEFAULT_VIDEO_PATH.format(episode_chunk=0, video_key=vid_key, episode_index=0)
+        for vid_key in video_keys
+    ]
+    hub_api = HfApi()
+    hub_api.snapshot_download(
+        repo_id=repo_id, repo_type="dataset", local_dir=local_dir, revision=branch, allow_patterns=video_files
+    )
+    videos_info_dict = {}
+    for vid_key, vid_path in zip(video_keys, video_files, strict=True):
+        videos_info_dict[vid_key] = get_video_info(local_dir / vid_path)
+
+    return videos_info_dict
+
+
+def convert_dataset(
+    repo_id: str,
+    local_dir: Path,
+    single_task: str | None = None,
+    tasks_path: Path | None = None,
+    tasks_col: Path | None = None,
+    robot_config: RobotConfig | None = None,
+    test_branch: str | None = None,
+    **card_kwargs,
+):
+    v1 = get_safe_version(repo_id, V16)
+    v1x_dir = local_dir / V16 / repo_id
+    v20_dir = local_dir / V20 / repo_id
+    v1x_dir.mkdir(parents=True, exist_ok=True)
+    v20_dir.mkdir(parents=True, exist_ok=True)
+
+    hub_api = HfApi()
+    hub_api.snapshot_download(
+        repo_id=repo_id, repo_type="dataset", revision=v1, local_dir=v1x_dir, ignore_patterns="videos*/"
+    )
+    branch = "main"
+    if test_branch:
+        branch = test_branch
+        create_branch(repo_id=repo_id, branch=test_branch, repo_type="dataset")
+
+    metadata_v1 = load_json(v1x_dir / V1_INFO_PATH)
+    dataset = datasets.load_dataset("parquet", data_dir=v1x_dir / "data", split="train")
+    features = get_features_from_hf_dataset(dataset, robot_config)
+    video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
+
+    if single_task and "language_instruction" in dataset.column_names:
+        logging.warning(
+            "'single_task' provided but 'language_instruction' tasks_col found. Using 'language_instruction'.",
+        )
+        single_task = None
+        tasks_col = "language_instruction"
+
+    # Episodes & chunks
+    episode_indices = sorted(dataset.unique("episode_index"))
+    total_episodes = len(episode_indices)
+    assert episode_indices == list(range(total_episodes))
+    total_videos = total_episodes * len(video_keys)
+    total_chunks = total_episodes // DEFAULT_CHUNK_SIZE
+    if total_episodes % DEFAULT_CHUNK_SIZE != 0:
+        total_chunks += 1
+
+    # Tasks
+    if single_task:
+        tasks_by_episodes = dict.fromkeys(episode_indices, single_task)
+        dataset, tasks = add_task_index_by_episodes(dataset, tasks_by_episodes)
+        tasks_by_episodes = {ep_idx: [task] for ep_idx, task in tasks_by_episodes.items()}
+    elif tasks_path:
+        tasks_by_episodes = load_json(tasks_path)
+        tasks_by_episodes = {int(ep_idx): task for ep_idx, task in tasks_by_episodes.items()}
+        dataset, tasks = add_task_index_by_episodes(dataset, tasks_by_episodes)
+        tasks_by_episodes = {ep_idx: [task] for ep_idx, task in tasks_by_episodes.items()}
+    elif tasks_col:
+        dataset, tasks, tasks_by_episodes = add_task_index_from_tasks_col(dataset, tasks_col)
+    else:
+        raise ValueError
+
+    assert set(tasks) == {task for ep_tasks in tasks_by_episodes.values() for task in ep_tasks}
+    tasks = [{"task_index": task_idx, "task": task} for task_idx, task in enumerate(tasks)]
+    write_jsonlines(tasks, v20_dir / LEGACY_TASKS_PATH)
+    features["task_index"] = {
+        "dtype": "int64",
+        "shape": (1,),
+        "names": None,
+    }
+
+    # Videos
+    if video_keys:
+        assert metadata_v1.get("video", False)
+        dataset = dataset.remove_columns(video_keys)
+        clean_gitattr = Path(
+            hub_api.hf_hub_download(
+                repo_id=GITATTRIBUTES_REF, repo_type="dataset", local_dir=local_dir, filename=".gitattributes"
+            )
+        ).absolute()
+        with tempfile.TemporaryDirectory() as tmp_video_dir:
+            move_videos(
+                repo_id, video_keys, total_episodes, total_chunks, Path(tmp_video_dir), clean_gitattr, branch
+            )
+        videos_info = get_videos_info(repo_id, v1x_dir, video_keys=video_keys, branch=branch)
+        for key in video_keys:
+            features[key]["shape"] = (
+                videos_info[key].pop("video.height"),
+                videos_info[key].pop("video.width"),
+                videos_info[key].pop("video.channels"),
+            )
+            features[key]["video_info"] = videos_info[key]
+            assert math.isclose(videos_info[key]["video.fps"], metadata_v1["fps"], rel_tol=1e-3)
+            if "encoding" in metadata_v1:
+                assert videos_info[key]["video.pix_fmt"] == metadata_v1["encoding"]["pix_fmt"]
+    else:
+        assert metadata_v1.get("video", 0) == 0
+        videos_info = None
+
+    # Split data into 1 parquet file by episode
+    episode_lengths = split_parquet_by_episodes(dataset, total_episodes, total_chunks, v20_dir)
+
+    if robot_config is not None:
+        robot_type = robot_config.type
+        repo_tags = [robot_type]
+    else:
+        robot_type = "unknown"
+        repo_tags = None
+
+    # Episodes
+    episodes = [
+        {"episode_index": ep_idx, "tasks": tasks_by_episodes[ep_idx], "length": episode_lengths[ep_idx]}
+        for ep_idx in episode_indices
+    ]
+    write_jsonlines(episodes, v20_dir / LEGACY_EPISODES_PATH)
+
+    # Assemble metadata v2.0
+    metadata_v2_0 = {
+        "codebase_version": V20,
+        "robot_type": robot_type,
+        "total_episodes": total_episodes,
+        "total_frames": len(dataset),
+        "total_tasks": len(tasks),
+        "total_videos": total_videos,
+        "total_chunks": total_chunks,
+        "chunks_size": DEFAULT_CHUNK_SIZE,
+        "fps": metadata_v1["fps"],
+        "splits": {"train": f"0:{total_episodes}"},
+        "data_path": LEGACY_DEFAULT_PARQUET_PATH,
+        "video_path": LEGACY_DEFAULT_VIDEO_PATH if video_keys else None,
+        "features": features,
+    }
+    write_json(metadata_v2_0, v20_dir / INFO_PATH)
+    convert_stats_to_json(v1x_dir, v20_dir)
+    card = create_lerobot_dataset_card(tags=repo_tags, dataset_info=metadata_v2_0, **card_kwargs)
+
+    with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
+        hub_api.delete_folder(repo_id=repo_id, path_in_repo="data", repo_type="dataset", revision=branch)
+
+    with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
+        hub_api.delete_folder(repo_id=repo_id, path_in_repo="meta_data", repo_type="dataset", revision=branch)
+
+    with contextlib.suppress(EntryNotFoundError, HfHubHTTPError):
+        hub_api.delete_folder(repo_id=repo_id, path_in_repo="meta", repo_type="dataset", revision=branch)
+
+    hub_api.upload_folder(
+        repo_id=repo_id,
+        path_in_repo="data",
+        folder_path=v20_dir / "data",
+        repo_type="dataset",
+        revision=branch,
+    )
+    hub_api.upload_folder(
+        repo_id=repo_id,
+        path_in_repo="meta",
+        folder_path=v20_dir / "meta",
+        repo_type="dataset",
+        revision=branch,
+    )
+
+    card.push_to_hub(repo_id=repo_id, repo_type="dataset", revision=branch)
+
+    if not test_branch:
+        create_branch(repo_id=repo_id, branch=V20, repo_type="dataset")
+
+
+def make_robot_config(robot_type: str, **kwargs) -> RobotConfig:
+    if robot_type == "aloha":
+        raise NotImplementedError  # TODO
+
+    elif robot_type == "koch_follower":
+        from lerobot.robots.koch_follower import KochFollowerConfig
+
+        return KochFollowerConfig(**kwargs)
+    elif robot_type == "so100_follower":
+        from lerobot.robots.so100_follower import SO100FollowerConfig
+
+        return SO100FollowerConfig(**kwargs)
+    elif robot_type == "stretch":
+        from lerobot.robots.stretch3 import Stretch3RobotConfig
+
+        return Stretch3RobotConfig(**kwargs)
+    elif robot_type == "lekiwi":
+        from lerobot.robots.lekiwi import LeKiwiConfig
+
+        return LeKiwiConfig(**kwargs)
+    else:
+        raise ValueError(f"Robot type '{robot_type}' is not available.")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    task_args = parser.add_mutually_exclusive_group(required=True)
+
+    parser.add_argument(
+        "--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`).",
+    )
+    task_args.add_argument(
+        "--single-task",
+        type=str,
+        help="A short but accurate description of the single task performed in the dataset.",
+    )
+    task_args.add_argument(
+        "--tasks-col",
+        type=str,
+        help="The name of the column containing language instructions",
+    )
+    task_args.add_argument(
+        "--tasks-path",
+        type=Path,
+        help="The path to a .json file containing one language instruction for each episode_index",
+    )
+    parser.add_argument(
+        "--robot",
+        type=str,
+        default=None,
+        help="Robot config used for the dataset during conversion (e.g. 'koch', 'aloha', 'so100', etc.)",
+    )
+    parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=None,
+        help="Local directory to store the dataset during conversion. Defaults to /tmp/lerobot_dataset_v2",
+    )
+    parser.add_argument(
+        "--license",
+        type=str,
+        default="apache-2.0",
+        help="Repo license. Must be one of https://huggingface.co/docs/hub/repositories-licenses. Defaults to mit.",
+    )
+    parser.add_argument(
+        "--test-branch",
+        type=str,
+        default=None,
+        help="Repo branch to test your conversion first (e.g. 'v2.0.test')",
+    )
+
+    args = parser.parse_args()
+    if not args.local_dir:
+        args.local_dir = Path("/tmp/lerobot_dataset_v2")
+
+    if args.robot is not None:
+        robot_config = make_robot_config(args.robot)
+
+    del args.robot
+
+    convert_dataset(**vars(args), robot_config=robot_config)
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/datasets/v21/_remove_language_instruction.py

@@ -0,0 +1,87 @@
+# 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 logging
+import traceback
+from pathlib import Path
+
+from datasets import get_dataset_config_info
+from huggingface_hub import HfApi
+
+from lerobot import available_datasets
+from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from lerobot.datasets.utils import INFO_PATH, write_info
+from lerobot.datasets.v21.convert_dataset_v20_to_v21 import V20, SuppressWarnings
+
+LOCAL_DIR = Path("data/")
+
+hub_api = HfApi()
+
+
+def fix_dataset(repo_id: str) -> str:
+    if not hub_api.revision_exists(repo_id, V20, repo_type="dataset"):
+        return f"{repo_id}: skipped (not in {V20})."
+
+    dataset_info = get_dataset_config_info(repo_id, "default")
+    with SuppressWarnings():
+        lerobot_metadata = LeRobotDatasetMetadata(repo_id, revision=V20, force_cache_sync=True)
+
+    meta_features = {key for key, ft in lerobot_metadata.features.items() if ft["dtype"] != "video"}
+    parquet_features = set(dataset_info.features)
+
+    diff_parquet_meta = parquet_features - meta_features
+    diff_meta_parquet = meta_features - parquet_features
+
+    if diff_parquet_meta:
+        raise ValueError(f"In parquet not in info.json: {parquet_features - meta_features}")
+
+    if not diff_meta_parquet:
+        return f"{repo_id}: skipped (no diff)"
+
+    if diff_meta_parquet:
+        logging.warning(f"In info.json not in parquet: {meta_features - parquet_features}")
+        assert diff_meta_parquet == {"language_instruction"}
+        lerobot_metadata.features.pop("language_instruction")
+        write_info(lerobot_metadata.info, lerobot_metadata.root)
+        commit_info = hub_api.upload_file(
+            path_or_fileobj=lerobot_metadata.root / INFO_PATH,
+            path_in_repo=INFO_PATH,
+            repo_id=repo_id,
+            repo_type="dataset",
+            revision=V20,
+            commit_message="Remove 'language_instruction'",
+            create_pr=True,
+        )
+        return f"{repo_id}: success - PR: {commit_info.pr_url}"
+
+
+def batch_fix():
+    status = {}
+    LOCAL_DIR.mkdir(parents=True, exist_ok=True)
+    logfile = LOCAL_DIR / "fix_features_v20.txt"
+    for num, repo_id in enumerate(available_datasets):
+        print(f"\nConverting {repo_id} ({num}/{len(available_datasets)})")
+        print("---------------------------------------------------------")
+        try:
+            status = fix_dataset(repo_id)
+        except Exception:
+            status = f"{repo_id}: failed\n    {traceback.format_exc()}"
+
+        logging.info(status)
+        with open(logfile, "a") as file:
+            file.write(status + "\n")
+
+
+if __name__ == "__main__":
+    batch_fix()

src/lerobot/datasets/v21/batch_convert_dataset_v20_to_v21.py

@@ -0,0 +1,54 @@
+#!/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.
+
+"""
+This script is for internal use to convert all datasets under the 'lerobot' hub user account to v2.1.
+"""
+
+import traceback
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+from lerobot import available_datasets
+from lerobot.datasets.v21.convert_dataset_v20_to_v21 import V21, convert_dataset
+
+LOCAL_DIR = Path("data/")
+
+
+def batch_convert():
+    status = {}
+    LOCAL_DIR.mkdir(parents=True, exist_ok=True)
+    logfile = LOCAL_DIR / "conversion_log_v21.txt"
+    hub_api = HfApi()
+    for num, repo_id in enumerate(available_datasets):
+        print(f"\nConverting {repo_id} ({num}/{len(available_datasets)})")
+        print("---------------------------------------------------------")
+        try:
+            if hub_api.revision_exists(repo_id, V21, repo_type="dataset"):
+                status = f"{repo_id}: success (already in {V21})."
+            else:
+                convert_dataset(repo_id)
+                status = f"{repo_id}: success."
+        except Exception:
+            status = f"{repo_id}: failed\n    {traceback.format_exc()}"
+
+        with open(logfile, "a") as file:
+            file.write(status + "\n")
+
+
+if __name__ == "__main__":
+    batch_convert()

src/lerobot/datasets/v21/convert_dataset_v20_to_v21.py

@@ -0,0 +1,111 @@
+# 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.
+
+"""
+This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.0 to
+2.1. 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".
+
+Usage:
+
+```bash
+python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 \
+    --repo-id=aliberts/koch_tutorial
+```
+
+"""
+
+import argparse
+import logging
+
+from huggingface_hub import HfApi
+
+from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+from lerobot.datasets.utils import STATS_PATH, load_stats, write_info
+from lerobot.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
+
+V20 = "v2.0"
+V21 = "v2.1"
+
+
+class SuppressWarnings:
+    def __enter__(self):
+        self.previous_level = logging.getLogger().getEffectiveLevel()
+        logging.getLogger().setLevel(logging.ERROR)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        logging.getLogger().setLevel(self.previous_level)
+
+
+def convert_dataset(
+    repo_id: str,
+    branch: str | None = None,
+    num_workers: int = 4,
+):
+    with SuppressWarnings():
+        dataset = LeRobotDataset(repo_id, revision=V20, force_cache_sync=True)
+
+    convert_stats(dataset, num_workers=num_workers)
+    ref_stats = load_stats(dataset.root)
+    check_aggregate_stats(dataset, ref_stats)
+
+    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:
+        (dataset.root / STATS_PATH).unlink()
+
+    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_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--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`).",
+    )
+    parser.add_argument(
+        "--branch",
+        type=str,
+        default=None,
+        help="Repo branch to push your dataset. Defaults to the main branch.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of workers for parallelizing stats compute. Defaults to 4.",
+    )
+
+    args = parser.parse_args()
+    convert_dataset(**vars(args))

src/lerobot/datasets/v21/convert_stats.py

@@ -0,0 +1,114 @@
+# 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 concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import jsonlines
+import numpy as np
+from tqdm import tqdm
+
+from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import LEGACY_EPISODES_STATS_PATH, serialize_dict
+
+
+def append_jsonlines(data: dict, fpath: Path) -> None:
+    fpath.parent.mkdir(exist_ok=True, parents=True)
+    with jsonlines.open(fpath, "a") as writer:
+        writer.write(data)
+
+
+def legacy_write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
+    # 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)}
+    append_jsonlines(episode_stats, local_dir / LEGACY_EPISODES_STATS_PATH)
+
+
+def sample_episode_video_frames(dataset: LeRobotDataset, episode_index: int, ft_key: str) -> np.ndarray:
+    ep_len = dataset.meta.episodes[episode_index]["length"]
+    sampled_indices = sample_indices(ep_len)
+    query_timestamps = dataset._get_query_timestamps(0.0, {ft_key: sampled_indices})
+    video_frames = dataset._query_videos(query_timestamps, episode_index)
+    return video_frames[ft_key].numpy()
+
+
+def convert_episode_stats(dataset: LeRobotDataset, ep_idx: int):
+    ep_start_idx = dataset.episode_data_index["from"][ep_idx]
+    ep_end_idx = dataset.episode_data_index["to"][ep_idx]
+    ep_data = dataset.hf_dataset.select(range(ep_start_idx, ep_end_idx))
+
+    ep_stats = {}
+    for key, ft in dataset.features.items():
+        if ft["dtype"] == "video":
+            # We sample only for videos
+            ep_ft_data = sample_episode_video_frames(dataset, ep_idx, key)
+        else:
+            ep_ft_data = np.array(ep_data[key])
+
+        axes_to_reduce = (0, 2, 3) if ft["dtype"] in ["image", "video"] else 0
+        keepdims = True if ft["dtype"] in ["image", "video"] else ep_ft_data.ndim == 1
+        ep_stats[key] = get_feature_stats(ep_ft_data, axis=axes_to_reduce, keepdims=keepdims)
+
+        if ft["dtype"] in ["image", "video"]:  # remove batch dim
+            ep_stats[key] = {
+                k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
+            }
+
+    dataset.meta.episodes_stats[ep_idx] = ep_stats
+
+
+def convert_stats(dataset: LeRobotDataset, num_workers: int = 0):
+    assert dataset.episodes is None
+    print("Computing episodes stats")
+    total_episodes = dataset.meta.total_episodes
+    if num_workers > 0:
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            futures = {
+                executor.submit(convert_episode_stats, dataset, ep_idx): ep_idx
+                for ep_idx in range(total_episodes)
+            }
+            for future in tqdm(as_completed(futures), total=total_episodes):
+                future.result()
+    else:
+        for ep_idx in tqdm(range(total_episodes)):
+            convert_episode_stats(dataset, ep_idx)
+
+    for ep_idx in tqdm(range(total_episodes)):
+        legacy_write_episode_stats(ep_idx, dataset.meta.episodes_stats[ep_idx], dataset.root)
+
+
+def check_aggregate_stats(
+    dataset: LeRobotDataset,
+    reference_stats: dict[str, dict[str, np.ndarray]],
+    video_rtol_atol: tuple[float] = (1e-2, 1e-2),
+    default_rtol_atol: tuple[float] = (5e-6, 6e-5),
+):
+    """Verifies that the aggregated stats from episodes_stats are close to reference stats."""
+    agg_stats = aggregate_stats(list(dataset.meta.episodes_stats.values()))
+    for key, ft in dataset.features.items():
+        # These values might need some fine-tuning
+        if ft["dtype"] == "video":
+            # to account for image sub-sampling
+            rtol, atol = video_rtol_atol
+        else:
+            rtol, atol = default_rtol_atol
+
+        for stat, val in agg_stats[key].items():
+            if key in reference_stats and stat in reference_stats[key]:
+                err_msg = f"feature='{key}' stats='{stat}'"
+                np.testing.assert_allclose(
+                    val, reference_stats[key][stat], rtol=rtol, atol=atol, err_msg=err_msg
+                )

src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py

@@ -62,6 +62,7 @@ from lerobot.datasets.utils import (
     flatten_dict,
     get_parquet_file_size_in_mb,
     get_parquet_num_frames,
+    get_video_duration_in_s,
     get_video_size_in_mb,
     load_info,
     update_chunk_file_indices,
@@ -70,7 +71,7 @@ from lerobot.datasets.utils import (
     write_stats,
     write_tasks,
 )
-from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
+from lerobot.datasets.video_utils import concat_video_files
 
 V21 = "v2.1"
 
@@ -172,7 +173,7 @@ def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
     concatenated_df.to_parquet(path, index=False, schema=schema)
 
 
-def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
+def convert_data(root, new_root):
     data_dir = root / "data"
     ep_paths = sorted(data_dir.glob("*/*.parquet"))
 
@@ -200,12 +201,11 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
         episodes_metadata.append(ep_metadata)
         ep_idx += 1
 
-        if size_in_mb < data_file_size_in_mb:
+        if size_in_mb < DEFAULT_DATA_FILE_SIZE_IN_MB:
             paths_to_cat.append(ep_path)
             continue
 
-        if paths_to_cat:
-            concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
+        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
 
         # Reset for the next file
         size_in_mb = ep_size_in_mb
@@ -235,7 +235,7 @@ def get_image_keys(root):
     return image_keys
 
 
-def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
+def convert_videos(root: Path, new_root: Path):
     video_keys = get_video_keys(root)
     if len(video_keys) == 0:
         return None
@@ -244,7 +244,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
 
     eps_metadata_per_cam = []
     for camera in video_keys:
-        eps_metadata = convert_videos_of_camera(root, new_root, camera, video_file_size_in_mb)
+        eps_metadata = convert_videos_of_camera(root, new_root, camera)
         eps_metadata_per_cam.append(eps_metadata)
 
     num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam]
@@ -269,7 +269,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
     return episods_metadata
 
 
-def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int):
+def convert_videos_of_camera(root: Path, new_root: Path, video_key):
     # Access old paths to mp4
     videos_dir = root / "videos"
     ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4"))
@@ -286,13 +286,9 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
         ep_duration_in_s = get_video_duration_in_s(ep_path)
 
         # Check if adding this episode would exceed the limit
-        if size_in_mb + ep_size_in_mb >= video_file_size_in_mb and len(paths_to_cat) > 0:
+        if size_in_mb + ep_size_in_mb >= DEFAULT_VIDEO_FILE_SIZE_IN_MB and len(paths_to_cat) > 0:
             # Size limit would be exceeded, save current accumulation WITHOUT this episode
-            concatenate_video_files(
-                paths_to_cat,
-                new_root
-                / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
-            )
+            concat_video_files(paths_to_cat, new_root, video_key, chunk_idx, file_idx)
 
             # Update episodes metadata for the file we just saved
             for i, _ in enumerate(paths_to_cat):
@@ -324,11 +320,7 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
 
     # Write remaining videos if any
     if paths_to_cat:
-        concatenate_video_files(
-            paths_to_cat,
-            new_root
-            / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
-        )
+        concat_video_files(paths_to_cat, new_root, video_key, chunk_idx, file_idx)
 
         # Update episodes metadata for the final file
         for i, _ in enumerate(paths_to_cat):
@@ -395,13 +387,13 @@ def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_
     write_stats(stats, new_root)
 
 
-def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
+def convert_info(root, new_root):
     info = load_info(root)
     info["codebase_version"] = "v3.0"
     del info["total_chunks"]
     del info["total_videos"]
-    info["data_files_size_in_mb"] = data_file_size_in_mb
-    info["video_files_size_in_mb"] = video_file_size_in_mb
+    info["data_files_size_in_mb"] = DEFAULT_DATA_FILE_SIZE_IN_MB
+    info["video_files_size_in_mb"] = DEFAULT_VIDEO_FILE_SIZE_IN_MB
     info["data_path"] = DEFAULT_DATA_PATH
     info["video_path"] = DEFAULT_VIDEO_PATH
     info["fps"] = float(info["fps"])
@@ -416,18 +408,12 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
 def convert_dataset(
     repo_id: str,
     branch: str | None = None,
-    data_file_size_in_mb: int | None = None,
-    video_file_size_in_mb: int | None = None,
+    num_workers: int = 4,
 ):
     root = HF_LEROBOT_HOME / repo_id
     old_root = HF_LEROBOT_HOME / f"{repo_id}_old"
     new_root = HF_LEROBOT_HOME / f"{repo_id}_v30"
 
-    if data_file_size_in_mb is None:
-        data_file_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
-    if video_file_size_in_mb is None:
-        video_file_size_in_mb = DEFAULT_VIDEO_FILE_SIZE_IN_MB
-
     if old_root.is_dir() and root.is_dir():
         shutil.rmtree(str(root))
         shutil.move(str(old_root), str(root))
@@ -442,10 +428,10 @@ def convert_dataset(
         local_dir=root,
     )
 
-    convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb)
+    convert_info(root, new_root)
     convert_tasks(root, new_root)
-    episodes_metadata = convert_data(root, new_root, data_file_size_in_mb)
-    episodes_videos_metadata = convert_videos(root, new_root, video_file_size_in_mb)
+    episodes_metadata = convert_data(root, new_root)
+    episodes_videos_metadata = convert_videos(root, new_root)
     convert_episodes_metadata(root, new_root, episodes_metadata, episodes_videos_metadata)
 
     shutil.move(str(root), str(old_root))
@@ -484,16 +470,10 @@ if __name__ == "__main__":
         help="Repo branch to push your dataset. Defaults to the main branch.",
     )
     parser.add_argument(
-        "--data-file-size-in-mb",
+        "--num-workers",
         type=int,
-        default=None,
-        help="File size in MB. Defaults to 100 for data and 500 for videos.",
-    )
-    parser.add_argument(
-        "--video-file-size-in-mb",
-        type=int,
-        default=None,
-        help="File size in MB. Defaults to 100 for data and 500 for videos.",
+        default=4,
+        help="Number of workers for parallelizing stats compute. Defaults to 4.",
     )
 
     args = parser.parse_args()

src/lerobot/datasets/video_utils.py

@@ -13,25 +13,26 @@
 # 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 glob
 import importlib
+import json
 import logging
 import shutil
+import subprocess
 import tempfile
 import warnings
+from collections import OrderedDict
 from dataclasses import dataclass, field
 from pathlib import Path
-from threading import Lock
 from typing import Any, ClassVar
 
-import av
-import fsspec
 import pyarrow as pa
 import torch
 import torchvision
 from datasets.features.features import register_feature
 from PIL import Image
 
+from lerobot.datasets.utils import DEFAULT_VIDEO_PATH
+
 
 def get_safe_default_codec():
     if importlib.util.find_spec("torchcodec"):
@@ -105,7 +106,7 @@ def decode_video_frames_torchvision(
     keyframes_only = False
     torchvision.set_video_backend(backend)
     if backend == "pyav":
-        keyframes_only = True  # pyav doesn't support accurate seek
+        keyframes_only = True  # pyav doesnt support accuracte seek
 
     # set a video stream reader
     # TODO(rcadene): also load audio stream at the same time
@@ -158,6 +159,7 @@ def decode_video_frames_torchvision(
     )
 
     # get closest frames to the query timestamps
+    # TODO(rcadene): remove torch.stack
     closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
     closest_ts = loaded_ts[argmin_]
 
@@ -171,68 +173,15 @@ def decode_video_frames_torchvision(
     return closest_frames
 
 
-class VideoDecoderCache:
-    """Thread-safe cache for video decoders to avoid expensive re-initialization."""
-
-    def __init__(self):
-        self._cache: dict[str, tuple[Any, Any]] = {}
-        self._lock = Lock()
-
-    def get_decoder(self, video_path: str):
-        """Get a cached decoder or create a new one."""
-        if importlib.util.find_spec("torchcodec"):
-            from torchcodec.decoders import VideoDecoder
-        else:
-            raise ImportError("torchcodec is required but not available.")
-
-        video_path = str(video_path)
-
-        with self._lock:
-            if video_path not in self._cache:
-                file_handle = fsspec.open(video_path).__enter__()
-                decoder = VideoDecoder(file_handle, seek_mode="approximate")
-                self._cache[video_path] = (decoder, file_handle)
-
-            return self._cache[video_path][0]
-
-    def clear(self):
-        """Clear the cache and close file handles."""
-        with self._lock:
-            for _, file_handle in self._cache.values():
-                file_handle.close()
-            self._cache.clear()
-
-    def size(self) -> int:
-        """Return the number of cached decoders."""
-        with self._lock:
-            return len(self._cache)
-
-
-class FrameTimestampError(ValueError):
-    """Helper error to indicate the retrieved timestamps exceed the queried ones"""
-
-    pass
-
-
-_default_decoder_cache = VideoDecoderCache()
-
-
 def decode_video_frames_torchcodec(
     video_path: Path | str,
     timestamps: list[float],
     tolerance_s: float,
+    device: str = "cpu",
     log_loaded_timestamps: bool = False,
-    decoder_cache: VideoDecoderCache | None = None,
 ) -> torch.Tensor:
     """Loads frames associated with the requested timestamps of a video using torchcodec.
 
-    Args:
-        video_path: Path to the video file.
-        timestamps: List of timestamps to extract frames.
-        tolerance_s: Allowed deviation in seconds for frame retrieval.
-        log_loaded_timestamps: Whether to log loaded timestamps.
-        decoder_cache: Optional decoder cache instance. Uses default if None.
-
     Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
 
     Note: Video benefits from inter-frame compression. Instead of storing every frame individually,
@@ -241,24 +190,27 @@ def decode_video_frames_torchcodec(
     and all subsequent frames until reaching the requested frame. The number of key frames in a video
     can be adjusted during encoding to take into account decoding time and video size in bytes.
     """
-    if decoder_cache is None:
-        decoder_cache = _default_decoder_cache
 
-    # Use cached decoder instead of creating new one each time
-    decoder = decoder_cache.get_decoder(str(video_path))
+    if importlib.util.find_spec("torchcodec"):
+        from torchcodec.decoders import VideoDecoder
+    else:
+        raise ImportError("torchcodec is required but not available.")
 
-    loaded_ts = []
+    # initialize video decoder
+    decoder = VideoDecoder(video_path, device=device, seek_mode="approximate")
     loaded_frames = []
-
+    loaded_ts = []
     # get metadata for frame information
     metadata = decoder.metadata
     average_fps = metadata.average_fps
+
     # convert timestamps to frame indices
     frame_indices = [round(ts * average_fps) for ts in timestamps]
+
     # retrieve frames based on indices
     frames_batch = decoder.get_frames_at(indices=frame_indices)
 
-    for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=True):
+    for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=False):
         loaded_frames.append(frame)
         loaded_ts.append(pts.item())
         if log_loaded_timestamps:
@@ -289,14 +241,10 @@ def decode_video_frames_torchcodec(
     if log_loaded_timestamps:
         logging.info(f"{closest_ts=}")
 
-    # convert to float32 in [0,1] range
-    closest_frames = (closest_frames / 255.0).type(torch.float32)
-
-    if not len(timestamps) == len(closest_frames):
-        raise FrameTimestampError(
-            f"Retrieved timestamps differ from queried {set(closest_frames) - set(timestamps)}"
-        )
+    # convert to float32 in [0,1] range (channel first)
+    closest_frames = closest_frames.type(torch.float32) / 255
 
+    assert len(timestamps) == len(closest_frames)
     return closest_frames
 
 
@@ -309,170 +257,104 @@ def encode_video_frames(
     g: int | None = 2,
     crf: int | None = 30,
     fast_decode: int = 0,
-    log_level: int | None = av.logging.ERROR,
+    log_level: str | None = "quiet",
     overwrite: bool = False,
 ) -> None:
     """More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
-    # Check encoder availability
-    if vcodec not in ["h264", "hevc", "libsvtav1"]:
-        raise ValueError(f"Unsupported video codec: {vcodec}. Supported codecs are: h264, hevc, libsvtav1.")
-
     video_path = Path(video_path)
     imgs_dir = Path(imgs_dir)
-
-    if video_path.exists() and not overwrite:
-        logging.warning(f"Video file already exists: {video_path}. Skipping encoding.")
-        return
-
     video_path.parent.mkdir(parents=True, exist_ok=True)
 
-    # Encoders/pixel formats incompatibility check
-    if (vcodec == "libsvtav1" or vcodec == "hevc") and pix_fmt == "yuv444p":
-        logging.warning(
-            f"Incompatible pixel format 'yuv444p' for codec {vcodec}, auto-selecting format 'yuv420p'"
-        )
-        pix_fmt = "yuv420p"
-
-    # Get input frames
-    template = "frame-" + ("[0-9]" * 6) + ".png"
-    input_list = sorted(
-        glob.glob(str(imgs_dir / template)), key=lambda x: int(x.split("-")[-1].split(".")[0])
+    ffmpeg_args = OrderedDict(
+        [
+            ("-f", "image2"),
+            ("-r", str(fps)),
+            ("-i", str(imgs_dir / "frame-%06d.png")),
+            ("-vcodec", vcodec),
+            ("-pix_fmt", pix_fmt),
+        ]
     )
 
-    # Define video output frame size (assuming all input frames are the same size)
-    if len(input_list) == 0:
-        raise FileNotFoundError(f"No images found in {imgs_dir}.")
-    dummy_image = Image.open(input_list[0])
-    width, height = dummy_image.size
-
-    # Define video codec options
-    video_options = {}
-
     if g is not None:
-        video_options["g"] = str(g)
+        ffmpeg_args["-g"] = str(g)
 
     if crf is not None:
-        video_options["crf"] = str(crf)
+        ffmpeg_args["-crf"] = str(crf)
 
     if fast_decode:
-        key = "svtav1-params" if vcodec == "libsvtav1" else "tune"
+        key = "-svtav1-params" if vcodec == "libsvtav1" else "-tune"
         value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
-        video_options[key] = value
+        ffmpeg_args[key] = value
 
-    # Set logging level
     if log_level is not None:
-        # "While less efficient, it is generally preferable to modify logging with Python's logging"
-        logging.getLogger("libav").setLevel(log_level)
+        ffmpeg_args["-loglevel"] = str(log_level)
 
-    # Create and open output file (overwrite by default)
-    with av.open(str(video_path), "w") as output:
-        output_stream = output.add_stream(vcodec, fps, options=video_options)
-        output_stream.pix_fmt = pix_fmt
-        output_stream.width = width
-        output_stream.height = height
+    ffmpeg_args = [item for pair in ffmpeg_args.items() for item in pair]
+    if overwrite:
+        ffmpeg_args.append("-y")
 
-        # Loop through input frames and encode them
-        for input_data in input_list:
-            input_image = Image.open(input_data).convert("RGB")
-            input_frame = av.VideoFrame.from_image(input_image)
-            packet = output_stream.encode(input_frame)
-            if packet:
-                output.mux(packet)
-
-        # Flush the encoder
-        packet = output_stream.encode()
-        if packet:
-            output.mux(packet)
-
-    # Reset logging level
-    if log_level is not None:
-        av.logging.restore_default_callback()
+    ffmpeg_cmd = ["ffmpeg"] + ffmpeg_args + [str(video_path)]
+    # redirect stdin to subprocess.DEVNULL to prevent reading random keyboard inputs from terminal
+    subprocess.run(ffmpeg_cmd, check=True, stdin=subprocess.DEVNULL)
 
     if not video_path.exists():
-        raise OSError(f"Video encoding did not work. File not found: {video_path}.")
+        raise OSError(
+            f"Video encoding did not work. File not found: {video_path}. "
+            f"Try running the command manually to debug: `{''.join(ffmpeg_cmd)}`"
+        )
 
 
-def concatenate_video_files(
-    input_video_paths: list[Path | str], output_video_path: Path, overwrite: bool = True
-):
+def concat_video_files(paths_to_cat: list[Path], root: Path, video_key: str, chunk_idx: int, file_idx: int):
     """
-    Concatenate multiple video files into a single video file using pyav.
+    Concatenate multiple video files into a single video file using ffmpeg.
 
-    This function takes a list of video input file paths and concatenates them into a single
+    This function takes a list of video file paths and concatenates them into a single
     output video file. It uses ffmpeg's concat demuxer with stream copy mode for fast
     concatenation without re-encoding.
 
     Args:
-        input_video_paths: Ordered list of input video file paths to concatenate.
-        output_video_path: Path to the output video file.
-        overwrite: Whether to overwrite the output video file if it already exists. Default is True.
+        paths_to_cat: List of video file paths to concatenate, in order.
+        root: Root directory where temporary files and output will be created.
+        video_key: Video key identifier (e.g., camera name) used in output path.
+        chunk_idx: Chunk index for organizing output files.
+        file_idx: File index within the chunk.
 
     Note:
         - Creates a temporary directory for intermediate files that is cleaned up after use.
         - Uses ffmpeg's concat demuxer which requires all input videos to have the same
           codec, resolution, and frame rate for proper concatenation.
+        - Output path follows the DEFAULT_VIDEO_PATH pattern with video_key, chunk_idx,
+          and file_idx parameters.
     """
 
-    output_video_path = Path(output_video_path)
+    tmp_dir = Path(tempfile.mkdtemp(dir=root))
+    path_concat_video_files = tmp_dir / "concat_video_files.txt"
+    with open(path_concat_video_files, "w") as f:
+        for ep_path in paths_to_cat:
+            f.write(f"file '{str(ep_path)}'\n")
 
-    if output_video_path.exists() and not overwrite:
-        logging.warning(f"Video file already exists: {output_video_path}. Skipping concatenation.")
-        return
+    path_tmp_output = tmp_dir / "tmp_output.mp4"
+    command = [
+        "ffmpeg",
+        "-y",
+        "-f",
+        "concat",
+        "-safe",
+        "0",
+        "-i",
+        str(path_concat_video_files),
+        "-c",
+        "copy",
+        str(path_tmp_output),
+    ]
+    subprocess.run(command, check=True)
 
-    output_video_path.parent.mkdir(parents=True, exist_ok=True)
-
-    if len(input_video_paths) == 0:
-        raise FileNotFoundError("No input video paths provided.")
-
-    # Create a temporary .ffconcat file to list the input video paths
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".ffconcat", delete=False) as tmp_concatenate_file:
-        tmp_concatenate_file.write("ffconcat version 1.0\n")
-        for input_path in input_video_paths:
-            tmp_concatenate_file.write(f"file '{str(input_path)}'\n")
-        tmp_concatenate_file.flush()
-        tmp_concatenate_path = tmp_concatenate_file.name
-
-    # Create input and output containers
-    input_container = av.open(
-        tmp_concatenate_path, mode="r", format="concat", options={"safe": "0"}
-    )  # safe = 0 allows absolute paths as well as relative paths
-
-    tmp_output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    output_container = av.open(
-        tmp_output_video_path, mode="w", options={"movflags": "faststart"}
-    )  # faststart is to move the metadata to the beginning of the file to speed up loading
-
-    # Replicate input streams in output container
-    stream_map = {}
-    for input_stream in input_container.streams:
-        if input_stream.type in ("video", "audio", "subtitle"):  # only copy compatible streams
-            stream_map[input_stream.index] = output_container.add_stream_from_template(
-                template=input_stream, opaque=True
-            )
-            stream_map[
-                input_stream.index
-            ].time_base = (
-                input_stream.time_base
-            )  # set the time base to the input stream time base (missing in the codec context)
-
-    # Demux + remux packets (no re-encode)
-    for packet in input_container.demux():
-        # Skip packets from un-mapped streams
-        if packet.stream.index not in stream_map:
-            continue
-
-        # Skip demux flushing packets
-        if packet.dts is None:
-            continue
-
-        output_stream = stream_map[packet.stream.index]
-        packet.stream = output_stream
-        output_container.mux(packet)
-
-    input_container.close()
-    output_container.close()
-    shutil.move(tmp_output_video_path, output_video_path)
-    Path(tmp_concatenate_path).unlink()
+    output_path = root / DEFAULT_VIDEO_PATH.format(
+        video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
+    )
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    shutil.move(str(path_tmp_output), str(output_path))
+    shutil.rmtree(str(tmp_dir))
 
 
 @dataclass
@@ -508,68 +390,78 @@ with warnings.catch_warnings():
 
 
 def get_audio_info(video_path: Path | str) -> dict:
-    # Set logging level
-    logging.getLogger("libav").setLevel(av.logging.ERROR)
+    ffprobe_audio_cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=channels,codec_name,bit_rate,sample_rate,bit_depth,channel_layout,duration",
+        "-of",
+        "json",
+        str(video_path),
+    ]
+    result = subprocess.run(ffprobe_audio_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Error running ffprobe: {result.stderr}")
 
-    # Getting audio stream information
-    audio_info = {}
-    with av.open(str(video_path), "r") as audio_file:
-        try:
-            audio_stream = audio_file.streams.audio[0]
-        except IndexError:
-            # Reset logging level
-            av.logging.restore_default_callback()
-            return {"has_audio": False}
+    info = json.loads(result.stdout)
+    audio_stream_info = info["streams"][0] if info.get("streams") else None
+    if audio_stream_info is None:
+        return {"has_audio": False}
 
-        audio_info["audio.channels"] = audio_stream.channels
-        audio_info["audio.codec"] = audio_stream.codec.canonical_name
-        # In an ideal loseless case : bit depth x sample rate x channels = bit rate.
-        # In an actual compressed case, the bit rate is set according to the compression level : the lower the bit rate, the more compression is applied.
-        audio_info["audio.bit_rate"] = audio_stream.bit_rate
-        audio_info["audio.sample_rate"] = audio_stream.sample_rate  # Number of samples per second
-        # In an ideal loseless case : fixed number of bits per sample.
-        # In an actual compressed case : variable number of bits per sample (often reduced to match a given depth rate).
-        audio_info["audio.bit_depth"] = audio_stream.format.bits
-        audio_info["audio.channel_layout"] = audio_stream.layout.name
-        audio_info["has_audio"] = True
-
-    # Reset logging level
-    av.logging.restore_default_callback()
-
-    return audio_info
+    # Return the information, defaulting to None if no audio stream is present
+    return {
+        "has_audio": True,
+        "audio.channels": audio_stream_info.get("channels", None),
+        "audio.codec": audio_stream_info.get("codec_name", None),
+        "audio.bit_rate": int(audio_stream_info["bit_rate"]) if audio_stream_info.get("bit_rate") else None,
+        "audio.sample_rate": int(audio_stream_info["sample_rate"])
+        if audio_stream_info.get("sample_rate")
+        else None,
+        "audio.bit_depth": audio_stream_info.get("bit_depth", None),
+        "audio.channel_layout": audio_stream_info.get("channel_layout", None),
+    }
 
 
 def get_video_info(video_path: Path | str) -> dict:
-    # Set logging level
-    logging.getLogger("libav").setLevel(av.logging.ERROR)
+    ffprobe_video_cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "v:0",
+        "-show_entries",
+        "stream=r_frame_rate,width,height,codec_name,nb_frames,duration,pix_fmt",
+        "-of",
+        "json",
+        str(video_path),
+    ]
+    result = subprocess.run(ffprobe_video_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Error running ffprobe: {result.stderr}")
 
-    # Getting video stream information
-    video_info = {}
-    with av.open(str(video_path), "r") as video_file:
-        try:
-            video_stream = video_file.streams.video[0]
-        except IndexError:
-            # Reset logging level
-            av.logging.restore_default_callback()
-            return {}
+    info = json.loads(result.stdout)
+    video_stream_info = info["streams"][0]
 
-        video_info["video.height"] = video_stream.height
-        video_info["video.width"] = video_stream.width
-        video_info["video.codec"] = video_stream.codec.canonical_name
-        video_info["video.pix_fmt"] = video_stream.pix_fmt
-        video_info["video.is_depth_map"] = False
+    # Calculate fps from r_frame_rate
+    r_frame_rate = video_stream_info["r_frame_rate"]
+    num, denom = map(int, r_frame_rate.split("/"))
+    fps = num / denom
 
-        # Calculate fps from r_frame_rate
-        video_info["video.fps"] = int(video_stream.base_rate)
+    pixel_channels = get_video_pixel_channels(video_stream_info["pix_fmt"])
 
-        pixel_channels = get_video_pixel_channels(video_stream.pix_fmt)
-        video_info["video.channels"] = pixel_channels
-
-    # Reset logging level
-    av.logging.restore_default_callback()
-
-    # Adding audio stream information
-    video_info.update(**get_audio_info(video_path))
+    video_info = {
+        "video.fps": fps,
+        "video.height": video_stream_info["height"],
+        "video.width": video_stream_info["width"],
+        "video.channels": pixel_channels,
+        "video.codec": video_stream_info["codec_name"],
+        "video.pix_fmt": video_stream_info["pix_fmt"],
+        "video.is_depth_map": False,
+        **get_audio_info(video_path),
+    }
 
     return video_info
 
@@ -596,88 +488,3 @@ def get_image_pixel_channels(image: Image):
         return 4  # RGBA
     else:
         raise ValueError("Unknown format")
-
-
-def get_video_duration_in_s(video_path: Path | str) -> float:
-    """
-    Get the duration of a video file in seconds using PyAV.
-
-    Args:
-        video_path: Path to the video file.
-
-    Returns:
-        Duration of the video in seconds.
-    """
-    with av.open(str(video_path)) as container:
-        # Get the first video stream
-        video_stream = container.streams.video[0]
-        # Calculate duration: stream.duration * stream.time_base gives duration in seconds
-        if video_stream.duration is not None:
-            duration = float(video_stream.duration * video_stream.time_base)
-        else:
-            # Fallback to container duration if stream duration is not available
-            duration = float(container.duration / av.time_base)
-    return duration
-
-
-class VideoEncodingManager:
-    """
-    Context manager that ensures proper video encoding and data cleanup even if exceptions occur.
-
-    This manager handles:
-    - Batch encoding for any remaining episodes when recording interrupted
-    - Cleaning up temporary image files from interrupted episodes
-    - Removing empty image directories
-
-    Args:
-        dataset: The LeRobotDataset instance
-    """
-
-    def __init__(self, dataset):
-        self.dataset = dataset
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        # Handle any remaining episodes that haven't been batch encoded
-        if self.dataset.episodes_since_last_encoding > 0:
-            if exc_type is not None:
-                logging.info("Exception occurred. Encoding remaining episodes before exit...")
-            else:
-                logging.info("Recording stopped. Encoding remaining episodes...")
-
-            start_ep = self.dataset.num_episodes - self.dataset.episodes_since_last_encoding
-            end_ep = self.dataset.num_episodes
-            logging.info(
-                f"Encoding remaining {self.dataset.episodes_since_last_encoding} episodes, "
-                f"from episode {start_ep} to {end_ep - 1}"
-            )
-            self.dataset._batch_save_episode_video(start_ep, end_ep)
-
-        # Clean up episode images if recording was interrupted
-        if exc_type is not None:
-            interrupted_episode_index = self.dataset.num_episodes
-            for key in self.dataset.meta.video_keys:
-                img_dir = self.dataset._get_image_file_path(
-                    episode_index=interrupted_episode_index, image_key=key, frame_index=0
-                ).parent
-                if img_dir.exists():
-                    logging.debug(
-                        f"Cleaning up interrupted episode images for episode {interrupted_episode_index}, camera {key}"
-                    )
-                    shutil.rmtree(img_dir)
-
-        # Clean up any remaining images directory if it's empty
-        img_dir = self.dataset.root / "images"
-        # Check for any remaining PNG files
-        png_files = list(img_dir.rglob("*.png"))
-        if len(png_files) == 0:
-            # Only remove the images directory if no PNG files remain
-            if img_dir.exists():
-                shutil.rmtree(img_dir)
-                logging.debug("Cleaned up empty images directory")
-        else:
-            logging.debug(f"Images directory is not empty, containing {len(png_files)} PNG files")
-
-        return False  # Don't suppress the original exception

src/lerobot/envs/configs.py

@@ -30,7 +30,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
     fps: int = 30
     features: dict[str, PolicyFeature] = field(default_factory=dict)
     features_map: dict[str, str] = field(default_factory=dict)
-    max_parallel_tasks: int = 1
 
     @property
     def type(self) -> str:
@@ -272,55 +271,3 @@ class HILEnvConfig(EnvConfig):
             "use_gamepad": self.use_gamepad,
             "gripper_penalty": self.gripper_penalty,
         }
-
-
-@EnvConfig.register_subclass("libero")
-@dataclass
-class LiberoEnv(EnvConfig):
-    task: str = "libero_10"  # can also choose libero_spatial, libero_object, etc.
-    fps: int = 30
-    episode_length: int = 520
-    obs_type: str = "pixels_agent_pos"
-    render_mode: str = "rgb_array"
-    camera_name: str = "agentview_image,robot0_eye_in_hand_image"
-    init_states: bool = True
-    camera_name_mapping: dict[str, str] | None = (None,)
-    features: dict[str, PolicyFeature] = field(
-        default_factory=lambda: {
-            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-        }
-    )
-    features_map: dict[str, str] = field(
-        default_factory=lambda: {
-            "action": ACTION,
-            "agent_pos": OBS_STATE,
-            "pixels/agentview_image": f"{OBS_IMAGES}.image",
-            "pixels/robot0_eye_in_hand_image": f"{OBS_IMAGES}.image2",
-        }
-    )
-
-    def __post_init__(self):
-        if self.obs_type == "pixels":
-            self.features["pixels/agentview_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
-            )
-            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
-            )
-        elif self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(8,))
-            self.features["pixels/agentview_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
-            )
-            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
-                type=FeatureType.VISUAL, shape=(360, 360, 3)
-            )
-        else:
-            raise ValueError(f"Unsupported obs_type: {self.obs_type}")
-
-    @property
-    def gym_kwargs(self) -> dict:
-        return {
-            "obs_type": self.obs_type,
-            "render_mode": self.render_mode,
-        }

src/lerobot/envs/factory.py

@@ -17,7 +17,7 @@ import importlib
 
 import gymnasium as gym
 
-from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, LiberoEnv, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -29,15 +29,11 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
         return XarmEnv(**kwargs)
     elif env_type == "hil":
         return HILEnvConfig(**kwargs)
-    elif env_type == "libero":
-        return LiberoEnv(**kwargs)
     else:
         raise ValueError(f"Policy type '{env_type}' is not available.")
 
 
-def make_env(
-    cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False
-) -> dict[str, dict[int, gym.vector.VectorEnv]]:
+def make_env(cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False) -> gym.vector.VectorEnv | None:
     """Makes a gym vector environment according to the config.
 
     Args:
@@ -51,44 +47,25 @@ def make_env(
         ModuleNotFoundError: If the requested env package is not installed
 
     Returns:
-        dict[str, dict[int, gym.vector.VectorEnv]]:
-            A mapping from suite name to indexed vectorized environments.
-            - For multi-task benchmarks (e.g., LIBERO): one entry per suite, and one vec env per task_id.
-            - For single-task environments: a single suite entry (cfg.type) with task_id=0.
-
+        gym.vector.VectorEnv: The parallelized gym.env instance.
     """
     if n_envs < 1:
-        raise ValueError("`n_envs` must be at least 1")
-
-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-
-    if "libero" in cfg.type:
-        from lerobot.envs.libero import create_libero_envs
-
-        return create_libero_envs(
-            task=cfg.task,
-            n_envs=n_envs,
-            camera_name=cfg.camera_name,
-            init_states=cfg.init_states,
-            gym_kwargs=cfg.gym_kwargs,
-            env_cls=env_cls,
-        )
+        raise ValueError("`n_envs must be at least 1")
 
     package_name = f"gym_{cfg.type}"
+
     try:
         importlib.import_module(package_name)
     except ModuleNotFoundError as e:
-        raise ModuleNotFoundError(
-            f'{package_name} is not installed. Install with: pip install "lerobot[{cfg.type}]"'
-        ) from e
+        print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
+        raise e
 
     gym_handle = f"{package_name}/{cfg.task}"
 
-    def _make_one():
-        return gym.make(gym_handle, disable_env_checker=True, **(cfg.gym_kwargs or {}))
+    # batched version of the env that returns an observation of shape (b, c)
+    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
+    env = env_cls(
+        [lambda: gym.make(gym_handle, disable_env_checker=True, **cfg.gym_kwargs) for _ in range(n_envs)]
+    )
 
-    vec = env_cls([_make_one for _ in range(n_envs)])
-
-    # normalize to {suite: {task_id: vec_env}} for consistency
-    suite_name = cfg.type  # e.g., "pusht", "aloha"
-    return {suite_name: {0: vec}}
+    return env

src/lerobot/envs/libero.py

@@ -1,399 +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
-
-import math
-import os
-from collections import defaultdict
-from collections.abc import Callable, Iterable, Mapping, Sequence
-from functools import partial
-from pathlib import Path
-from typing import Any
-
-import gymnasium as gym
-import numpy as np
-import torch
-from gymnasium import spaces
-from libero.libero import benchmark, get_libero_path
-from libero.libero.envs import OffScreenRenderEnv
-
-
-def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
-    """Normalize camera_name into a non-empty list of strings."""
-    if isinstance(camera_name, str):
-        cams = [c.strip() for c in camera_name.split(",") if c.strip()]
-    elif isinstance(camera_name, (list, tuple)):
-        cams = [str(c).strip() for c in camera_name if str(c).strip()]
-    else:
-        raise TypeError(f"camera_name must be str or sequence[str], got {type(camera_name).__name__}")
-    if not cams:
-        raise ValueError("camera_name resolved to an empty list.")
-    return cams
-
-
-def _get_suite(name: str) -> Any:
-    """Instantiate a LIBERO suite by name with clear validation."""
-    bench = benchmark.get_benchmark_dict()
-    if name not in bench:
-        raise ValueError(f"Unknown LIBERO suite '{name}'. Available: {', '.join(sorted(bench.keys()))}")
-    suite = bench[name]()
-    if not getattr(suite, "tasks", None):
-        raise ValueError(f"Suite '{name}' has no tasks.")
-    return suite
-
-
-def _select_task_ids(total_tasks: int, task_ids: Iterable[int] | None) -> list[int]:
-    """Validate/normalize task ids. If None → all tasks."""
-    if task_ids is None:
-        return list(range(total_tasks))
-    ids = sorted({int(t) for t in task_ids})
-    for t in ids:
-        if t < 0 or t >= total_tasks:
-            raise ValueError(f"task_id {t} out of range [0, {total_tasks - 1}].")
-    return ids
-
-
-def quat2axisangle(quat: np.ndarray) -> np.ndarray:
-    """
-    Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
-
-    Converts quaternion to axis-angle format.
-    Returns a unit vector direction scaled by its angle in radians.
-
-    Args:
-        quat (np.array): (x,y,z,w) vec4 float angles
-
-    Returns:
-        np.array: (ax,ay,az) axis-angle exponential coordinates
-    """
-    # clip quaternion
-    if quat[3] > 1.0:
-        quat[3] = 1.0
-    elif quat[3] < -1.0:
-        quat[3] = -1.0
-
-    den = np.sqrt(1.0 - quat[3] * quat[3])
-    if math.isclose(den, 0.0):
-        # This is (close to) a zero degree rotation, immediately return
-        return np.zeros(3)
-
-    return (quat[:3] * 2.0 * math.acos(quat[3])) / den
-
-
-def get_task_init_states(task_suite: Any, i: int) -> np.ndarray:
-    init_states_path = (
-        Path(get_libero_path("init_states"))
-        / task_suite.tasks[i].problem_folder
-        / task_suite.tasks[i].init_states_file
-    )
-    init_states = torch.load(init_states_path, weights_only=False)  # nosec B614
-    return init_states
-
-
-def get_libero_dummy_action():
-    """Get dummy/no-op action, used to roll out the simulation while the robot does nothing."""
-    return [0, 0, 0, 0, 0, 0, -1]
-
-
-OBS_STATE_DIM = 8
-ACTION_DIM = 7
-TASK_SUITE_MAX_STEPS: dict[str, int] = {
-    "libero_spatial": 280,  # longest training demo has 193 steps
-    "libero_object": 280,  # longest training demo has 254 steps
-    "libero_goal": 300,  # longest training demo has 270 steps
-    "libero_10": 520,  # longest training demo has 505 steps
-    "libero_90": 400,  # longest training demo has 373 steps
-}
-
-
-class LiberoEnv(gym.Env):
-    metadata = {"render_modes": ["rgb_array"], "render_fps": 80}
-
-    def __init__(
-        self,
-        task_suite: Any,
-        task_id: int,
-        task_suite_name: str,
-        camera_name: str | Sequence[str] = "agentview_image,robot0_eye_in_hand_image",
-        obs_type: str = "pixels",
-        render_mode: str = "rgb_array",
-        observation_width: int = 256,
-        observation_height: int = 256,
-        visualization_width: int = 640,
-        visualization_height: int = 480,
-        init_states: bool = True,
-        episode_index: int = 0,
-        camera_name_mapping: dict[str, str] | None = None,
-        num_steps_wait: int = 10,
-    ):
-        super().__init__()
-        self.task_id = task_id
-        self.obs_type = obs_type
-        self.render_mode = render_mode
-        self.observation_width = observation_width
-        self.observation_height = observation_height
-        self.visualization_width = visualization_width
-        self.visualization_height = visualization_height
-        self.init_states = init_states
-        self.camera_name = camera_name.split(
-            ","
-        )  # agentview_image (main) or robot0_eye_in_hand_image (wrist)
-
-        # Map raw camera names to "image1" and "image2".
-        # The preprocessing step `preprocess_observation` will then prefix these with `.images.*`,
-        # following the LeRobot convention (e.g., `observation.images.image`, `observation.images.image2`).
-        # This ensures the policy consistently receives observations in the
-        # expected format regardless of the original camera naming.
-        if camera_name_mapping is None:
-            camera_name_mapping = {
-                "agentview_image": "image",
-                "robot0_eye_in_hand_image": "image2",
-            }
-        self.camera_name_mapping = camera_name_mapping
-        self.num_steps_wait = num_steps_wait
-        self.episode_index = episode_index
-        # Load once and keep
-        self._init_states = get_task_init_states(task_suite, self.task_id) if self.init_states else None
-        self._init_state_id = self.episode_index  # tie each sub-env to a fixed init state
-
-        self._env = self._make_envs_task(task_suite, self.task_id)
-        default_steps = 500
-        self._max_episode_steps = TASK_SUITE_MAX_STEPS.get(task_suite_name, default_steps)
-
-        images = {}
-        for cam in self.camera_name:
-            images[self.camera_name_mapping[cam]] = spaces.Box(
-                low=0,
-                high=255,
-                shape=(self.observation_height, self.observation_width, 3),
-                dtype=np.uint8,
-            )
-
-        if self.obs_type == "state":
-            raise NotImplementedError(
-                "The 'state' observation type is not supported in LiberoEnv. "
-                "Please switch to an image-based obs_type (e.g. 'pixels', 'pixels_agent_pos')."
-            )
-
-        elif self.obs_type == "pixels":
-            self.observation_space = spaces.Dict(
-                {
-                    "pixels": spaces.Dict(images),
-                }
-            )
-        elif self.obs_type == "pixels_agent_pos":
-            self.observation_space = spaces.Dict(
-                {
-                    "pixels": spaces.Dict(images),
-                    "agent_pos": spaces.Box(
-                        low=-1000.0,
-                        high=1000.0,
-                        shape=(OBS_STATE_DIM,),
-                        dtype=np.float64,
-                    ),
-                }
-            )
-
-        self.action_space = spaces.Box(low=-1, high=1, shape=(ACTION_DIM,), dtype=np.float32)
-
-    def render(self):
-        raw_obs = self._env.env._get_observations()
-        image = self._format_raw_obs(raw_obs)["pixels"]["image"]
-        return image
-
-    def _make_envs_task(self, task_suite: Any, task_id: int = 0):
-        task = task_suite.get_task(task_id)
-        self.task = task.name
-        self.task_description = task.language
-        task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file)
-
-        env_args = {
-            "bddl_file_name": task_bddl_file,
-            "camera_heights": self.observation_height,
-            "camera_widths": self.observation_width,
-        }
-        env = OffScreenRenderEnv(**env_args)
-        env.reset()
-        return env
-
-    def _format_raw_obs(self, raw_obs: dict[str, Any]) -> dict[str, Any]:
-        images = {}
-        for camera_name in self.camera_name:
-            image = raw_obs[camera_name]
-            image = image[::-1, ::-1]  # rotate 180 degrees
-            images[self.camera_name_mapping[camera_name]] = image
-        state = np.concatenate(
-            (
-                raw_obs["robot0_eef_pos"],
-                quat2axisangle(raw_obs["robot0_eef_quat"]),
-                raw_obs["robot0_gripper_qpos"],
-            )
-        )
-        agent_pos = state
-        if self.obs_type == "pixels":
-            return {"pixels": images.copy()}
-        if self.obs_type == "pixels_agent_pos":
-            return {
-                "pixels": images.copy(),
-                "agent_pos": agent_pos,
-            }
-        raise NotImplementedError(
-            f"The observation type '{self.obs_type}' is not supported in LiberoEnv. "
-            "Please switch to an image-based obs_type (e.g. 'pixels', 'pixels_agent_pos')."
-        )
-
-    def reset(self, seed=None, **kwargs):
-        super().reset(seed=seed)
-        self._env.seed(seed)
-        if self.init_states and self._init_states is not None:
-            self._env.set_init_state(self._init_states[self._init_state_id])
-        raw_obs = self._env.reset()
-
-        # After reset, objects may be unstable (slightly floating, intersecting, etc.).
-        # Step the simulator with a no-op action for a few frames so everything settles.
-        # Increasing this value can improve determinism and reproducibility across resets.
-        for _ in range(self.num_steps_wait):
-            raw_obs, _, _, _ = self._env.step(get_libero_dummy_action())
-        observation = self._format_raw_obs(raw_obs)
-        info = {"is_success": False}
-        return observation, info
-
-    def step(self, action: np.ndarray) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
-        if action.ndim != 1:
-            raise ValueError(
-                f"Expected action to be 1-D (shape (action_dim,)), "
-                f"but got shape {action.shape} with ndim={action.ndim}"
-            )
-        raw_obs, reward, done, info = self._env.step(action)
-
-        is_success = self._env.check_success()
-        terminated = done or is_success
-        info["is_success"] = is_success
-
-        observation = self._format_raw_obs(raw_obs)
-        if done:
-            self.reset()
-            info.update(
-                {
-                    "task": self.task,
-                    "task_id": self.task_id,
-                    "done": done,
-                    "is_success": is_success,
-                }
-            )
-        truncated = False
-        return observation, reward, terminated, truncated, info
-
-    def close(self):
-        self._env.close()
-
-
-def _make_env_fns(
-    *,
-    suite,
-    suite_name: str,
-    task_id: int,
-    n_envs: int,
-    camera_names: list[str],
-    init_states: bool,
-    gym_kwargs: Mapping[str, Any],
-) -> list[Callable[[], LiberoEnv]]:
-    """Build n_envs factory callables for a single (suite, task_id)."""
-    joined_cams = ",".join(camera_names)  # keep backward-compat: downstream expects a string
-
-    def _make_env(episode_index: int, **kwargs) -> LiberoEnv:
-        local_kwargs = dict(kwargs)
-        return LiberoEnv(
-            task_suite=suite,
-            task_id=task_id,
-            task_suite_name=suite_name,
-            camera_name=joined_cams,
-            init_states=init_states,
-            episode_index=episode_index,
-            **local_kwargs,
-        )
-
-    fns: list[Callable[[], LiberoEnv]] = []
-    for episode_index in range(n_envs):
-        fns.append(partial(_make_env, episode_index, **gym_kwargs))
-    return fns
-
-
-# ---- Main API ----------------------------------------------------------------
-
-
-def create_libero_envs(
-    task: str,
-    n_envs: int,
-    gym_kwargs: dict[str, Any] | None = None,
-    camera_name: str | Sequence[str] = "agentview_image,robot0_eye_in_hand_image",
-    init_states: bool = True,
-    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
-) -> dict[str, dict[int, Any]]:
-    """
-    Create vectorized LIBERO environments with a consistent return shape.
-
-    Returns:
-        dict[suite_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
-    Notes:
-        - n_envs is the number of rollouts *per task* (episode_index = 0..n_envs-1).
-        - `task` can be a single suite or a comma-separated list of suites.
-        - You may pass `task_ids` (list[int]) inside `gym_kwargs` to restrict tasks per suite.
-    """
-    if env_cls is None or not callable(env_cls):
-        raise ValueError("env_cls must be a callable that wraps a list of environment factory callables.")
-    if not isinstance(n_envs, int) or n_envs <= 0:
-        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
-
-    gym_kwargs = dict(gym_kwargs or {})
-    task_ids_filter = gym_kwargs.pop("task_ids", None)  # optional: limit to specific tasks
-
-    camera_names = _parse_camera_names(camera_name)
-    suite_names = [s.strip() for s in str(task).split(",") if s.strip()]
-    if not suite_names:
-        raise ValueError("`task` must contain at least one LIBERO suite name.")
-
-    print(
-        f"Creating LIBERO envs | suites={suite_names} | n_envs(per task)={n_envs} | init_states={init_states}"
-    )
-    if task_ids_filter is not None:
-        print(f"Restricting to task_ids={task_ids_filter}")
-
-    out: dict[str, dict[int, Any]] = defaultdict(dict)
-
-    for suite_name in suite_names:
-        suite = _get_suite(suite_name)
-        total = len(suite.tasks)
-        selected = _select_task_ids(total, task_ids_filter)
-
-        if not selected:
-            raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).")
-
-        for tid in selected:
-            fns = _make_env_fns(
-                suite=suite,
-                suite_name=suite_name,
-                task_id=tid,
-                n_envs=n_envs,
-                camera_names=camera_names,
-                init_states=init_states,
-                gym_kwargs=gym_kwargs,
-            )
-            out[suite_name][tid] = env_cls(fns)
-            print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}")
-
-    # return plain dicts for predictability
-    return {suite: dict(task_map) for suite, task_map in out.items()}

src/lerobot/envs/utils.py

@@ -14,8 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import warnings
-from collections.abc import Mapping, Sequence
-from functools import singledispatch
 from typing import Any
 
 import einops
@@ -99,6 +97,7 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
 
         policy_key = env_cfg.features_map[key]
         policy_features[policy_key] = feature
+
     return policy_features
 
 
@@ -128,68 +127,10 @@ 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)]
     return observation
-
-
-def _close_single_env(env: Any) -> None:
-    try:
-        env.close()
-    except Exception as exc:
-        print(f"Exception while closing env {env}: {exc}")
-
-
-@singledispatch
-def close_envs(obj: Any) -> None:
-    """Default: raise if the type is not recognized."""
-    raise NotImplementedError(f"close_envs not implemented for type {type(obj).__name__}")
-
-
-@close_envs.register
-def _(env: Mapping) -> None:
-    for v in env.values():
-        if isinstance(v, Mapping):
-            close_envs(v)
-        elif hasattr(v, "close"):
-            _close_single_env(v)
-
-
-@close_envs.register
-def _(envs: Sequence) -> None:
-    if isinstance(envs, (str, bytes)):
-        return
-    for v in envs:
-        if isinstance(v, Mapping) or isinstance(v, Sequence) and not isinstance(v, (str, bytes)):
-            close_envs(v)
-        elif hasattr(v, "close"):
-            _close_single_env(v)
-
-
-@close_envs.register
-def _(env: gym.Env) -> None:
-    _close_single_env(env)

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

@@ -15,17 +15,6 @@
 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 .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
-from .smolvla.processor_smolvla import SmolVLANewLineProcessor
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
-
-__all__ = [
-    "ACTConfig",
-    "DiffusionConfig",
-    "PI0Config",
-    "SmolVLAConfig",
-    "TDMPCConfig",
-    "VQBeTConfig",
-]

src/lerobot/policies/act/modeling_act.py

@@ -35,6 +35,7 @@ from torchvision.ops.misc import FrozenBatchNorm2d
 
 from lerobot.constants import ACTION, OBS_IMAGES
 from lerobot.policies.act.configuration_act import ACTConfig
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 
 
@@ -50,16 +51,27 @@ class ACTPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: ACTConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
         super().__init__(config)
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.model = ACT(config)
 
         if config.temporal_ensemble_coeff is not None:
@@ -125,19 +137,23 @@ class ACTPolicy(PreTrainedPolicy):
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
         actions = self.model(batch)[0]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
+        batch = self.normalize_targets(batch)
         actions_hat, (mu_hat, log_sigma_x2_hat) = self.model(batch)
 
         l1_loss = (

src/lerobot/policies/act/processor_act.py

@@ -1,85 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Tony Z. Zhao and 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 typing import Any
-
-import torch
-
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.act.configuration_act import ACTConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_act_pre_post_processors(
-    config: ACTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """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.
-
-    Returns:
-        tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
-        pre-processor pipeline and the post-processor pipeline.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-            device=config.device,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -35,6 +35,7 @@ from torch import Tensor, nn
 
 from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import (
     get_device_from_parameters,
@@ -56,6 +57,7 @@ class DiffusionPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: DiffusionConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
@@ -68,6 +70,14 @@ class DiffusionPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         # queues are populated during rollout of the policy, they contain the n latest observations and actions
         self._queues = None
 
@@ -96,6 +106,9 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         actions = self.diffusion.generate_actions(batch)
 
+        # TODO(rcadene): make above methods return output dictionary?
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
         return actions
 
     @torch.no_grad()
@@ -124,6 +137,7 @@ class DiffusionPolicy(PreTrainedPolicy):
         if ACTION in batch:
             batch.pop(ACTION)
 
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
@@ -139,9 +153,11 @@ class DiffusionPolicy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
         """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch = self.normalize_targets(batch)
         loss = self.diffusion.compute_loss(batch)
         # no output_dict so returning None
         return loss, None

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -1,92 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab,
-# and 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 typing import Any
-
-import torch
-
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_diffusion_pre_post_processors(
-    config: DiffusionConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-    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.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/factory.py

@@ -14,17 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import logging
-from typing import Any, TypedDict
 
-import torch
-from typing_extensions import Unpack
+from torch import nn
 
 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
@@ -39,32 +34,10 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla
 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 PolicyAction, PolicyProcessorPipeline
-from lerobot.processor.converters import (
-    batch_to_transition,
-    policy_action_to_transition,
-    transition_to_batch,
-    transition_to_policy_action,
-)
 
 
-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.
-    """
+def get_policy_class(name: str) -> PreTrainedPolicy:
+    """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
 
@@ -106,24 +79,6 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
 
 
 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":
@@ -146,187 +101,30 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         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_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
-
-
-def make_pre_post_processors(
-    policy_cfg: PreTrainedConfig,
-    pretrained_path: str | None = None,
-    **kwargs: Unpack[ProcessorConfigKwargs],
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    Create or load pre- and post-processor pipelines for a given policy.
-
-    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`).
-
-    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`.
-
-    Returns:
-        A tuple containing the input (pre-processor) and output (post-processor) pipelines.
-
-    Raises:
-        NotImplementedError: If a processor factory is not implemented for the given
-            policy configuration type.
-    """
-    if pretrained_path:
-        return (
-            PolicyProcessorPipeline.from_pretrained(
-                pretrained_model_name_or_path=pretrained_path,
-                config_filename=kwargs.get(
-                    "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
-                ),
-                overrides=kwargs.get("preprocessor_overrides", {}),
-                to_transition=batch_to_transition,
-                to_output=transition_to_batch,
-            ),
-            PolicyProcessorPipeline.from_pretrained(
-                pretrained_model_name_or_path=pretrained_path,
-                config_filename=kwargs.get(
-                    "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
-                ),
-                overrides=kwargs.get("postprocessor_overrides", {}),
-                to_transition=policy_action_to_transition,
-                to_output=transition_to_policy_action,
-            ),
-        )
-
-    # 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
-
-        processors = make_tdmpc_pre_post_processors(
-            config=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
-
-        processors = make_diffusion_pre_post_processors(
-            config=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
-
-        processors = make_act_pre_post_processors(
-            config=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
-
-        processors = make_vqbet_pre_post_processors(
-            config=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
-
-        processors = make_pi0_pre_post_processors(
-            config=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
-
-        processors = make_pi0fast_pre_post_processors(
-            config=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
-
-        processors = make_sac_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, RewardClassifierConfig):
-        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"),
-        )
-
-    elif isinstance(policy_cfg, SmolVLAConfig):
-        from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors
-
-        processors = make_smolvla_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    else:
-        raise NotImplementedError(f"Processor for policy type '{policy_cfg.type}' is not implemented.")
-
-    return processors
-
-
 def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,
     env_cfg: EnvConfig | 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.")
@@ -349,6 +147,7 @@ def make_policy(
     kwargs = {}
     if ds_meta is not None:
         features = dataset_to_policy_features(ds_meta.features)
+        kwargs["dataset_stats"] = ds_meta.stats
     else:
         if not cfg.pretrained_path:
             logging.warning(
@@ -356,9 +155,8 @@ def make_policy(
                 "rather than a dataset. Normalization modules inside the policy will have infinite values "
                 "by default without stats from a dataset."
             )
-        if env_cfg is None:
-            raise ValueError("env_cfg cannot be None when ds_meta is not provided")
         features = env_to_policy_features(env_cfg)
+
     cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
     cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
     kwargs["config"] = cfg
@@ -371,7 +169,10 @@ def make_policy(
     else:
         # Make a fresh 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")
+
     return policy

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

@@ -56,15 +56,18 @@ from collections import deque
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
+from transformers import AutoTokenizer
 
-from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0.paligemma_with_expert import (
     PaliGemmaWithExpertConfig,
     PaliGemmaWithExpertModel,
 )
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.utils import get_safe_dtype
+from lerobot.policies.utils import log_model_loading_keys
+from lerobot.utils.utils import get_safe_dtype, init_logging
 
 
 def create_sinusoidal_pos_embedding(
@@ -220,17 +223,28 @@ class PI0Policy(PreTrainedPolicy):
     def __init__(
         self,
         config: PI0Config,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
 
         super().__init__(config)
         config.validate_features()
         self.config = config
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
 
+        self.language_tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
         self.model = PI0FlowMatching(config)
 
         self.reset()
@@ -239,6 +253,99 @@ class PI0Policy(PreTrainedPolicy):
         """This should be called whenever the environment is reset."""
         self._action_queue = deque([], maxlen=self.config.n_action_steps)
 
+    @classmethod
+    def _transform_state_dict_keys(cls, state_dict: dict) -> dict:
+        """
+        Transform state dict keys to match expected model structure.
+
+        Transformations:
+        - model.paligemma_with_expert.paligemma.language_model.lm_head ->
+          model.paligemma_with_expert.paligemma.lm_head
+        - model.paligemma_with_expert.paligemma.language_model.model ->
+          model.paligemma_with_expert.paligemma.model.language_model
+        - model.paligemma_with_expert.paligemma.vision_tower ->
+          model.paligemma_with_expert.paligemma.model.vision_tower
+        - model.paligemma_with_expert.paligemma.multi_modal_projector ->
+          model.paligemma_with_expert.paligemma.model.multi_modal_projector
+
+        Also handles tied weights between lm_head.weight and
+        embed_tokens.weight.
+        """
+        import re
+
+        transformed_dict = {}
+
+        transformations = [
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.lm_head"),
+                ".paligemma_with_expert.paligemma.lm_head",
+            ),
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.model"),
+                ".paligemma_with_expert.paligemma.model.language_model",
+            ),
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.vision_tower"),
+                ".paligemma_with_expert.paligemma.model.vision_tower",
+            ),
+            (
+                re.compile(r"\.paligemma_with_expert\.paligemma\.multi_modal_projector"),
+                ".paligemma_with_expert.paligemma.model.multi_modal_projector",
+            ),
+        ]
+
+        for key, value in state_dict.items():
+            new_key = key
+            for pattern, replacement in transformations:
+                new_key = pattern.sub(replacement, new_key)
+            transformed_dict[new_key] = value
+
+        # Handle tied weights: lm_head.weight and embed_tokens.weight share memory
+        lm_head_key = None
+        embed_tokens_key = None
+
+        for key in transformed_dict:
+            if key.endswith(".paligemma_with_expert.paligemma.lm_head.weight"):
+                lm_head_key = key
+            elif key.endswith(".paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"):
+                embed_tokens_key = key
+            if lm_head_key and embed_tokens_key:
+                break
+
+        if lm_head_key and not embed_tokens_key:
+            embed_tokens_key = lm_head_key.replace(
+                ".lm_head.weight", ".model.language_model.embed_tokens.weight"
+            )
+            transformed_dict[embed_tokens_key] = transformed_dict[lm_head_key]
+        elif embed_tokens_key and not lm_head_key:
+            lm_head_key = embed_tokens_key.replace(
+                ".model.language_model.embed_tokens.weight", ".lm_head.weight"
+            )
+            transformed_dict[lm_head_key] = transformed_dict[embed_tokens_key]
+
+        return transformed_dict
+
+    @classmethod
+    def _load_as_safetensor(
+        cls, model: "PI0Policy", model_file: str, map_location: str, strict: bool
+    ) -> "PI0Policy":
+        """Override to apply key transformations before loading."""
+        from safetensors.torch import load_file
+
+        init_logging()
+        # Load the state dict from file safely
+        state_dict = load_file(model_file, device=map_location)
+
+        # Apply key transformations
+        transformed_state_dict = cls._transform_state_dict_keys(state_dict)
+
+        # Load the transformed state dict
+        msg = model.load_state_dict(transformed_state_dict, strict=strict)
+
+        # Log message
+        log_model_loading_keys(msg.missing_keys, msg.unexpected_keys)
+        return model
+
     def get_optim_params(self) -> dict:
         return self.parameters()
 
@@ -270,13 +377,14 @@ class PI0Policy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
+        batch = self.normalize_inputs(batch)
+
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
             images, img_masks = self.prepare_images(batch)
             state = self.prepare_state(batch)
-            lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
-            lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
+            lang_tokens, lang_masks = self.prepare_language(batch)
 
             actions = self.model.sample_actions(
                 images, img_masks, lang_tokens, lang_masks, state, noise=noise
@@ -286,6 +394,8 @@ class PI0Policy(PreTrainedPolicy):
             original_action_dim = self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
+            actions = self.unnormalize_outputs({"action": actions})["action"]
+
             if self.config.adapt_to_pi_aloha:
                 actions = self._pi_aloha_encode_actions(actions)
 
@@ -300,10 +410,12 @@ class PI0Policy(PreTrainedPolicy):
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
 
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
+
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
-        lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
+        lang_tokens, lang_masks = self.prepare_language(batch)
         actions = self.prepare_action(batch)
         actions_is_pad = batch.get("action_is_pad")
 
@@ -370,6 +482,26 @@ class PI0Policy(PreTrainedPolicy):
 
         return images, img_masks
 
+    def prepare_language(self, batch) -> tuple[Tensor, Tensor]:
+        """Tokenize the text input"""
+        device = batch[OBS_STATE].device
+        tasks = batch["task"]
+
+        # PaliGemma prompt has to end with a new line
+        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
+
+        tokenized_prompt = self.language_tokenizer.__call__(
+            tasks,
+            padding="max_length",
+            padding_side="right",
+            max_length=self.config.tokenizer_max_length,
+            return_tensors="pt",
+        )
+        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
+        lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
+
+        return lang_tokens, lang_masks
+
     def _pi_aloha_decode_state(self, state):
         # Flip the joints.
         for motor_idx in [1, 2, 8, 9]:
@@ -435,7 +567,7 @@ class PI0FlowMatching(nn.Module):
     └──────────────────────────────┘
     """
 
-    def __init__(self, config: PI0Config):
+    def __init__(self, config):
         super().__init__()
         self.config = config
 

src/lerobot/policies/pi0/processor_pi0.py

@@ -1,166 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 Physical Intelligence and 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 typing import Any
-
-import torch
-
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    ProcessorStep,
-    ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-@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.
-    """
-
-    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
-
-        task = complementary_data["task"]
-        if task is None:
-            return complementary_data
-
-        new_complementary_data = dict(complementary_data)
-
-        # 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"
-        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]
-        # If task is neither string nor list of strings, leave unchanged
-
-        return new_complementary_data
-
-    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.
-        """
-        return features
-
-
-def make_pi0_pre_post_processors(
-    config: PI0Config,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    Constructs pre-processor and post-processor pipelines for the PI0 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. 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.
-
-    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 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.
-    """
-
-    # Add remaining processors
-    input_steps: list[ProcessorStep] = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
-        Pi0NewLineProcessor(),  # Add newlines before tokenization for PaliGemma
-        TokenizerProcessorStep(
-            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,
-        ),
-    ]
-
-    output_steps: list[ProcessorStep] = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -58,6 +58,7 @@ from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING
 
 from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 
@@ -145,6 +146,14 @@ class PI0FASTPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.language_tokenizer = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224")
         self.model = PI0FAST(config)
 
@@ -212,6 +221,8 @@ class PI0FASTPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
+        batch = self.normalize_inputs(batch)
+
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
@@ -224,6 +235,8 @@ class PI0FASTPolicy(PreTrainedPolicy):
             ]  # self.config.max_action_dim  # self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
+            actions = self.unnormalize_outputs({"action": actions})["action"]
+
             if self.config.adapt_to_pi_aloha:
                 actions = self._pi_aloha_encode_actions(actions)
 
@@ -236,6 +249,8 @@ class PI0FASTPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
         loss_dict = self.model.forward(batch)
         return loss_dict["loss"], loss_dict
 

src/lerobot/policies/pi0fast/processor_pi0fast.py

@@ -1,92 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 Physical Intelligence and 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 typing import Any
-
-import torch
-
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_pi0fast_pre_post_processors(
-    config: PI0FASTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/sac/modeling_sac.py

@@ -28,6 +28,7 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
 
+from lerobot.policies.normalize import NormalizeBuffer
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
 from lerobot.policies.utils import get_device_from_parameters
@@ -44,6 +45,7 @@ class SACPolicy(
     def __init__(
         self,
         config: SACConfig | None = None,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         super().__init__(config)
         config.validate_features()
@@ -51,6 +53,7 @@ class SACPolicy(
 
         # Determine action dimension and initialize all components
         continuous_action_dim = config.output_features["action"].shape[0]
+        self._init_normalization(dataset_stats)
         self._init_encoders()
         self._init_critics(continuous_action_dim)
         self._init_actor(continuous_action_dim)
@@ -85,7 +88,8 @@ class SACPolicy(
 
         observations_features = None
         if self.shared_encoder and self.actor.encoder.has_images:
-            observations_features = self.actor.encoder.get_cached_image_features(batch)
+            # Cache and normalize image features
+            observations_features = self.actor.encoder.get_cached_image_features(batch, normalize=True)
 
         actions, _, _ = self.actor(batch, observations_features)
 
@@ -387,12 +391,28 @@ class SACPolicy(
         actor_loss = ((self.temperature * log_probs) - min_q_preds).mean()
         return actor_loss
 
+    def _init_normalization(self, dataset_stats):
+        """Initialize input/output normalization modules."""
+        self.normalize_inputs = nn.Identity()
+        self.normalize_targets = nn.Identity()
+        if self.config.dataset_stats is not None:
+            params = _convert_normalization_params_to_tensor(self.config.dataset_stats)
+            self.normalize_inputs = NormalizeBuffer(
+                self.config.input_features, self.config.normalization_mapping, params
+            )
+            stats = dataset_stats or params
+            self.normalize_targets = NormalizeBuffer(
+                self.config.output_features, self.config.normalization_mapping, stats
+            )
+
     def _init_encoders(self):
         """Initialize shared or separate encoders for actor and critic."""
         self.shared_encoder = self.config.shared_encoder
-        self.encoder_critic = SACObservationEncoder(self.config)
+        self.encoder_critic = SACObservationEncoder(self.config, self.normalize_inputs)
         self.encoder_actor = (
-            self.encoder_critic if self.shared_encoder else SACObservationEncoder(self.config)
+            self.encoder_critic
+            if self.shared_encoder
+            else SACObservationEncoder(self.config, self.normalize_inputs)
         )
 
     def _init_critics(self, continuous_action_dim):
@@ -404,7 +424,9 @@ class SACPolicy(
             )
             for _ in range(self.config.num_critics)
         ]
-        self.critic_ensemble = CriticEnsemble(encoder=self.encoder_critic, ensemble=heads)
+        self.critic_ensemble = CriticEnsemble(
+            encoder=self.encoder_critic, ensemble=heads, output_normalization=self.normalize_targets
+        )
         target_heads = [
             CriticHead(
                 input_dim=self.encoder_critic.output_dim + continuous_action_dim,
@@ -412,7 +434,9 @@ class SACPolicy(
             )
             for _ in range(self.config.num_critics)
         ]
-        self.critic_target = CriticEnsemble(encoder=self.encoder_critic, ensemble=target_heads)
+        self.critic_target = CriticEnsemble(
+            encoder=self.encoder_critic, ensemble=target_heads, output_normalization=self.normalize_targets
+        )
         self.critic_target.load_state_dict(self.critic_ensemble.state_dict())
 
         if self.config.use_torch_compile:
@@ -466,9 +490,10 @@ class SACPolicy(
 class SACObservationEncoder(nn.Module):
     """Encode image and/or state vector observations."""
 
-    def __init__(self, config: SACConfig) -> None:
+    def __init__(self, config: SACConfig, input_normalizer: nn.Module) -> None:
         super().__init__()
         self.config = config
+        self.input_normalization = input_normalizer
         self._init_image_layers()
         self._init_state_layers()
         self._compute_output_dim()
@@ -543,10 +568,11 @@ class SACObservationEncoder(nn.Module):
     def forward(
         self, obs: dict[str, Tensor], cache: dict[str, Tensor] | None = None, detach: bool = False
     ) -> Tensor:
+        obs = self.input_normalization(obs)
         parts = []
         if self.has_images:
             if cache is None:
-                cache = self.get_cached_image_features(obs)
+                cache = self.get_cached_image_features(obs, normalize=False)
             parts.append(self._encode_images(cache, detach))
         if self.has_env:
             parts.append(self.env_encoder(obs["observation.environment_state"]))
@@ -559,7 +585,7 @@ class SACObservationEncoder(nn.Module):
             "No parts to concatenate, you should have at least one image or environment state or state"
         )
 
-    def get_cached_image_features(self, obs: dict[str, Tensor]) -> dict[str, Tensor]:
+    def get_cached_image_features(self, obs: dict[str, Tensor], normalize: bool = False) -> dict[str, Tensor]:
         """Extract and optionally cache image features from observations.
 
         This function processes image observations through the vision encoder once and returns
@@ -571,17 +597,26 @@ class SACObservationEncoder(nn.Module):
         - The vision encoder forward pass is typically the main computational bottleneck during training and inference
         - Caching these features can provide 2-4x speedup in training and inference
 
+        Normalization behavior:
+        - When called from inside forward(): set normalize=False since inputs are already normalized
+        - When called from outside forward(): set normalize=True to ensure proper input normalization
+
         Usage patterns:
-        - Called in select_action()
+        - Called in select_action() with normalize=True
         - Called in learner.py's get_observation_features() to pre-compute features for all policy components
-        - Called internally by forward()
+        - Called internally by forward() with normalize=False
 
         Args:
             obs: Dictionary of observation tensors containing image keys
+            normalize: Whether to normalize observations before encoding
+                      Set to True when calling directly from outside the encoder's forward method
+                      Set to False when calling from within forward() where inputs are already normalized
 
         Returns:
             Dictionary mapping image keys to their corresponding encoded features
         """
+        if normalize:
+            obs = self.input_normalization(obs)
         batched = torch.cat([obs[k] for k in self.image_keys], dim=0)
         out = self.image_encoder(batched)
         chunks = torch.chunk(out, len(self.image_keys), dim=0)
@@ -712,6 +747,7 @@ class CriticEnsemble(nn.Module):
     Args:
         encoder (SACObservationEncoder): encoder for observations.
         ensemble (List[CriticHead]): list of critic heads.
+        output_normalization (nn.Module): normalization layer for actions.
         init_final (float | None): optional initializer scale for final layers.
 
     Forward returns a tensor of shape (num_critics, batch_size) containing Q-values.
@@ -721,11 +757,13 @@ class CriticEnsemble(nn.Module):
         self,
         encoder: SACObservationEncoder,
         ensemble: list[CriticHead],
+        output_normalization: nn.Module,
         init_final: float | None = None,
     ):
         super().__init__()
         self.encoder = encoder
         self.init_final = init_final
+        self.output_normalization = output_normalization
         self.critics = nn.ModuleList(ensemble)
 
     def forward(
@@ -737,6 +775,11 @@ class CriticEnsemble(nn.Module):
         device = get_device_from_parameters(self)
         # Move each tensor in observations to device
         observations = {k: v.to(device) for k, v in observations.items()}
+        # NOTE: We normalize actions it helps for sample efficiency
+        actions: dict[str, torch.tensor] = {"action": actions}
+        # NOTE: Normalization layer took dict in input and outputs a dict that why
+        actions = self.output_normalization(actions)["action"]
+        actions = actions.to(device)
 
         obs_enc = self.encoder(observations, cache=observation_features)
 

src/lerobot/policies/sac/processor_sac.py

@@ -1,92 +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.
-
-from typing import Any
-
-import torch
-
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.sac.configuration_sac import SACConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_sac_pre_post_processors(
-    config: SACConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-
-    # Add remaining processors
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

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

@@ -20,6 +20,7 @@ import torch
 from torch import Tensor, nn
 
 from lerobot.constants import OBS_IMAGE, REWARD
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 
@@ -107,12 +108,22 @@ class Classifier(PreTrainedPolicy):
     def __init__(
         self,
         config: RewardClassifierConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         from transformers import AutoModel
 
         super().__init__(config)
         self.config = config
 
+        # Initialize normalization (standardized with the policy framework)
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         # Set up encoder
         encoder = AutoModel.from_pretrained(self.config.model_name, trust_remote_code=True)
         # Extract vision model if we're given a multimodal model
@@ -236,6 +247,10 @@ class Classifier(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Tensor]]:
         """Standard forward pass for training compatible with train.py."""
+        # Normalize inputs if needed
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
+
         # Extract images and labels
         images, labels = self.extract_images_and_labels(batch)
 

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

@@ -1,82 +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 typing import Any
-
-import torch
-
-from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.processor import (
-    DeviceProcessorStep,
-    IdentityProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_classifier_processor(
-    config: RewardClassifierConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-    """
-
-    input_steps = [
-        NormalizerProcessorStep(
-            features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        NormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        DeviceProcessorStep(device=config.device),
-    ]
-    output_steps = [DeviceProcessorStep(device="cpu"), IdentityProcessorStep()]
-
-    return (
-        PolicyProcessorPipeline(
-            steps=input_steps,
-            name="classifier_preprocessor",
-        ),
-        PolicyProcessorPipeline(
-            steps=output_steps,
-            name="classifier_postprocessor",
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -53,13 +53,21 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 """
 
 import math
+import os
+import re
 from collections import deque
 
+import safetensors
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
+from transformers import AutoProcessor
 
-from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import (
+    Normalize,
+    Unnormalize,
+)
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
@@ -68,6 +76,102 @@ from lerobot.policies.utils import (
 )
 from lerobot.utils.utils import get_safe_dtype
 
+# Matches ".soNNN", optionally followed by "-something", up to the "_buffer_" marker
+_VARIANT_RE = re.compile(r"\.so\d+(?:-[\w]+)?_buffer_")
+
+
+def canonicalise(k: str) -> str:
+    """
+    Remove dataset-variant markers like '.so100-blue_' or '.so100_' from a
+    normalisation-buffer key.
+    """
+    return _VARIANT_RE.sub(".buffer_", k)
+
+
+def standardise_state_dict(
+    checkpoint: dict[str, torch.Tensor], ref_keys: set[str], *, verbose: bool = True
+) -> tuple[dict[str, torch.Tensor], list[str]]:
+    """
+    • Re-keys `checkpoint ` so that every entry matches the *reference* key set.
+    • If several variant keys collapse to the same canonical name we keep the
+      first one and log the collision.
+    • Returns the new dict + a list of entries that could not be matched.
+    """
+    out, collisions, unmatched = {}, {}, []
+
+    for k, v in checkpoint.items():
+        canon = canonicalise(k)
+        if canon in ref_keys:
+            if canon in out:  # duplicate after collapsing
+                collisions.setdefault(canon, []).append(k)
+            else:
+                out[canon] = v
+        else:
+            unmatched.append(k)
+
+    if verbose:
+        for canon, variants in collisions.items():
+            print(f"[standardise_state_dict] '{canon}'  ←  {variants}")
+        if unmatched:
+            print(f"[standardise_state_dict] kept {len(unmatched)} unmatched keys")
+
+    out.update({k: checkpoint[k] for k in unmatched})
+    return out, unmatched
+
+
+def rename_checkpoint_keys(checkpoint: dict, rename_str: str):
+    """
+    Renames keys in a checkpoint dictionary based on the given rename string.
+
+    Args:
+        checkpoint (dict): The checkpoint dictionary.
+        rename_str (str): A string specifying key mappings in the format "old1//new1,old2//new2".
+
+    Returns:
+        dict: The modified checkpoint with renamed keys.
+    """
+
+    rename_dict = dict(pair.split("//") for pair in rename_str.split(","))
+
+    new_checkpoint = {}
+    for k, v in checkpoint.items():
+        for old_key, new_key in rename_dict.items():
+            if old_key in k:
+                k = k.replace(old_key, new_key)
+        new_checkpoint[k] = v
+    return new_checkpoint
+
+
+def load_smolvla(
+    model: torch.nn.Module,
+    filename: str | os.PathLike,
+    *,
+    device: str = "cpu",
+    checkpoint_keys_mapping: str = "",
+) -> torch.nn.Module:
+    state_dict = safetensors.torch.load_file(filename, device=device)
+
+    # Optional user-supplied renames (e.g. "model._orig_mod.//model.")
+    if checkpoint_keys_mapping and "//" in checkpoint_keys_mapping:
+        state_dict = rename_checkpoint_keys(state_dict, checkpoint_keys_mapping)
+
+    state_dict, _ = standardise_state_dict(state_dict, set(model.state_dict().keys()))
+
+    # HACK(aliberts): to not overwrite normalization parameters as they should come from the dataset
+    norm_keys = ("normalize_inputs", "normalize_targets", "unnormalize_outputs")
+    state_dict = {k: v for k, v in state_dict.items() if not k.startswith(norm_keys)}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+
+    if not all(key.startswith(norm_keys) for key in missing) or unexpected:
+        raise RuntimeError(
+            "SmolVLA %d missing / %d unexpected keys",
+            len(missing),
+            len(unexpected),
+        )
+
+    return model
+
 
 def create_sinusoidal_pos_embedding(
     time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
@@ -222,17 +326,28 @@ class SmolVLAPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: SmolVLAConfig,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
 
         super().__init__(config)
         config.validate_features()
         self.config = config
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
 
+        self.language_tokenizer = AutoProcessor.from_pretrained(self.config.vlm_model_name).tokenizer
         self.model = VLAFlowMatching(config)
         self.reset()
 
@@ -242,6 +357,23 @@ class SmolVLAPolicy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
+    # HACK(aliberts, danaaubakirova): we overwrite this classmethod here to fix smolVLA-specific issues
+    @classmethod
+    def _load_as_safetensor(
+        cls,
+        model: "SmolVLAPolicy",
+        model_file: str,
+        map_location: str,
+        strict: bool,
+    ):
+        safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)
+        return load_smolvla(
+            model,
+            model_file,
+            device=map_location,
+            checkpoint_keys_mapping="model._orig_mod.//model.",
+        )
+
     def get_optim_params(self) -> dict:
         return self.parameters()
 
@@ -257,8 +389,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
-        lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
+        lang_tokens, lang_masks = self.prepare_language(batch)
 
         actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
 
@@ -266,6 +397,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
         original_action_dim = self.config.action_feature.shape[0]
         actions = actions[:, :, :original_action_dim]
 
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
         if self.config.adapt_to_pi_aloha:
             actions = self._pi_aloha_encode_actions(actions)
 
@@ -275,6 +408,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
+        batch = self.normalize_inputs(batch)
+
         return batch
 
     @torch.no_grad()
@@ -315,11 +450,11 @@ class SmolVLAPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
-
+        batch = self.normalize_inputs(batch)
+        batch = self.normalize_targets(batch)
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
-        lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
+        lang_tokens, lang_masks = self.prepare_language(batch)
         actions = self.prepare_action(batch)
         actions_is_pad = batch.get("actions_id_pad")
         loss_dict = {}
@@ -383,6 +518,30 @@ class SmolVLAPolicy(PreTrainedPolicy):
             img_masks.append(mask)
         return images, img_masks
 
+    def prepare_language(self, batch) -> tuple[Tensor, Tensor]:
+        """Tokenize the text input"""
+        device = batch[OBS_STATE].device
+        tasks = batch["task"]
+        if isinstance(tasks, str):
+            tasks = [tasks]
+
+        if len(tasks) == 1:
+            tasks = [tasks[0] for _ in range(batch[OBS_STATE].shape[0])]
+
+        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
+
+        tokenized_prompt = self.language_tokenizer.__call__(
+            tasks,
+            padding=self.config.pad_language_to,
+            padding_side="right",
+            max_length=self.config.tokenizer_max_length,
+            return_tensors="pt",
+        )
+        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
+        lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
+
+        return lang_tokens, lang_masks
+
     def _pi_aloha_decode_state(self, state):
         # Flip the joints.
         for motor_idx in [1, 2, 8, 9]:

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -1,141 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 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 typing import Any
-
-import torch
-
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_smolvla_pre_post_processors(
-    config: SmolVLAConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
-        SmolVLANewLineProcessor(),
-        TokenizerProcessorStep(
-            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,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )
-
-
-@ProcessorStepRegistry.register(name="smolvla_new_line_processor")
-class SmolVLANewLineProcessor(ComplementaryDataProcessorStep):
-    """
-    A processor step that ensures the 'task' description ends with a newline character.
-
-    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
-
-        task = complementary_data["task"]
-        if task is None:
-            return complementary_data
-
-        new_complementary_data = dict(complementary_data)
-
-        # 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"
-        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]
-        # If task is neither string nor list of strings, leave unchanged
-
-        return new_complementary_data
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -36,6 +36,7 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
 from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_STATE, REWARD
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
@@ -62,19 +63,26 @@ class TDMPCPolicy(PreTrainedPolicy):
     config_class = TDMPCConfig
     name = "tdmpc"
 
-    def __init__(
-        self,
-        config: TDMPCConfig,
-    ):
+    def __init__(self, config: TDMPCConfig, dataset_stats: dict[str, dict[str, Tensor]] | None = None):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                 the configuration class is used.
+            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
+                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
         super().__init__(config)
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.model = TDMPCTOLD(config)
         self.model_target = deepcopy(self.model)
         for param in self.model_target.parameters():
@@ -129,6 +137,7 @@ class TDMPCPolicy(PreTrainedPolicy):
 
         actions = torch.clamp(actions, -1, +1)
 
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     @torch.no_grad()
@@ -138,12 +147,11 @@ class TDMPCPolicy(PreTrainedPolicy):
         if ACTION in batch:
             batch.pop(ACTION)
 
+        batch = self.normalize_inputs(batch)
+
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
-        # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
-        if ACTION in batch:
-            batch.pop(ACTION)
 
         self._queues = populate_queues(self._queues, batch)
 
@@ -312,9 +320,11 @@ class TDMPCPolicy(PreTrainedPolicy):
         """
         device = get_device_from_parameters(self)
 
+        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
+        batch = self.normalize_targets(batch)
 
         info = {}
 

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -1,90 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Nicklas Hansen, Xiaolong Wang, Hao Su,
-# and 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 typing import Any
-
-import torch
-
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_tdmpc_pre_post_processors(
-    config: TDMPCConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/vqbet/modeling_vqbet.py

@@ -28,6 +28,7 @@ import torchvision
 from torch import Tensor, nn
 
 from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
@@ -47,6 +48,7 @@ class VQBeTPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: VQBeTConfig | None = None,
+        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
@@ -59,6 +61,14 @@ class VQBeTPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
+        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
+        self.normalize_targets = Normalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+        self.unnormalize_outputs = Unnormalize(
+            config.output_features, config.normalization_mapping, dataset_stats
+        )
+
         self.vqbet = VQBeTModel(config)
 
         self.reset()
@@ -118,6 +128,7 @@ class VQBeTPolicy(PreTrainedPolicy):
     def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     @torch.no_grad()
@@ -131,12 +142,10 @@ class VQBeTPolicy(PreTrainedPolicy):
         # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
         if ACTION in batch:
             batch.pop(ACTION)
+        batch = self.normalize_inputs(batch)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         # NOTE: It's important that this happens after stacking the images into a single key.
         batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
-        # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
-        if ACTION in batch:
-            batch.pop(ACTION)
 
         self._queues = populate_queues(self._queues, batch)
 
@@ -156,8 +165,10 @@ class VQBeTPolicy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
+        batch = self.normalize_inputs(batch)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch = self.normalize_targets(batch)
         # VQ-BeT discretizes action using VQ-VAE before training BeT (please refer to section 3.2 in the VQ-BeT paper https://huggingface.co/papers/2403.03181)
         if not self.vqbet.action_head.vqvae_model.discretized.item():
             # loss: total loss of training RVQ

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -1,91 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Seungjae Lee and Yibin Wang and Haritheja Etukuru
-# and H. Jin Kim and Nur Muhammad Mahi Shafiullah and Lerrel Pinto
-# and 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 typing import Any
-
-import torch
-
-from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-
-
-def make_vqbet_pre_post_processors(
-    config: VQBeTConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    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.
-
-    Returns:
-        A tuple containing the configured pre-processor and post-processor pipelines.
-    """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # Let the possibility to the user to rename the keys
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        DeviceProcessorStep(device="cpu"),
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/processor/__init__.py

@@ -14,111 +14,41 @@
 # 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,
-    transition_to_batch,
-)
-from .core import (
-    EnvAction,
-    EnvTransition,
-    PolicyAction,
-    RobotAction,
-    RobotObservation,
-    TransitionKey,
-)
-from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
-from .device_processor import DeviceProcessorStep
-from .factory import (
-    make_default_processors,
-    make_default_robot_action_processor,
-    make_default_robot_observation_processor,
-    make_default_teleop_action_processor,
-)
-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 .device_processor import DeviceProcessor
+from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
+from .observation_processor import VanillaObservationProcessor
 from .pipeline import (
-    ActionProcessorStep,
-    ComplementaryDataProcessorStep,
-    DataProcessorPipeline,
-    DoneProcessorStep,
-    IdentityProcessorStep,
-    InfoProcessorStep,
-    ObservationProcessorStep,
-    PolicyActionProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorKwargs,
+    ActionProcessor,
+    DoneProcessor,
+    EnvTransition,
+    IdentityProcessor,
+    InfoProcessor,
+    ObservationProcessor,
     ProcessorStep,
     ProcessorStepRegistry,
-    RewardProcessorStep,
-    RobotActionProcessorStep,
-    RobotProcessorPipeline,
-    TruncatedProcessorStep,
+    RewardProcessor,
+    RobotProcessor,
+    TransitionKey,
+    TruncatedProcessor,
 )
-from .rename_processor import RenameObservationsProcessorStep
-from .tokenizer_processor import TokenizerProcessorStep
+from .rename_processor import RenameProcessor
 
 __all__ = [
-    "ActionProcessorStep",
-    "AddTeleopActionAsComplimentaryDataStep",
-    "AddTeleopEventsAsInfoStep",
-    "ComplementaryDataProcessorStep",
-    "batch_to_transition",
-    "create_transition",
-    "DeviceProcessorStep",
-    "DoneProcessorStep",
-    "EnvAction",
+    "ActionProcessor",
+    "DeviceProcessor",
+    "DoneProcessor",
     "EnvTransition",
-    "GripperPenaltyProcessorStep",
-    "hotswap_stats",
-    "IdentityProcessorStep",
-    "ImageCropResizeProcessorStep",
-    "InfoProcessorStep",
-    "InterventionActionProcessorStep",
-    "JointVelocityProcessorStep",
-    "make_default_processors",
-    "make_default_teleop_action_processor",
-    "make_default_robot_action_processor",
-    "make_default_robot_observation_processor",
-    "MapDeltaActionToRobotActionStep",
-    "MapTensorToDeltaActionDictStep",
-    "MotorCurrentProcessorStep",
-    "NormalizerProcessorStep",
-    "Numpy2TorchActionProcessorStep",
-    "ObservationProcessorStep",
-    "PolicyAction",
-    "PolicyActionProcessorStep",
-    "PolicyProcessorPipeline",
-    "ProcessorKwargs",
+    "IdentityProcessor",
+    "InfoProcessor",
+    "NormalizerProcessor",
+    "UnnormalizerProcessor",
+    "ObservationProcessor",
     "ProcessorStep",
     "ProcessorStepRegistry",
-    "RobotAction",
-    "RobotActionProcessorStep",
-    "RobotObservation",
-    "RenameObservationsProcessorStep",
-    "RewardClassifierProcessorStep",
-    "RewardProcessorStep",
-    "DataProcessorPipeline",
-    "TimeLimitProcessorStep",
-    "AddBatchDimensionProcessorStep",
-    "RobotProcessorPipeline",
-    "TokenizerProcessorStep",
-    "Torch2NumpyActionProcessorStep",
-    "transition_to_batch",
+    "RenameProcessor",
+    "RewardProcessor",
+    "RobotProcessor",
     "TransitionKey",
-    "TruncatedProcessorStep",
-    "UnnormalizerProcessorStep",
-    "VanillaObservationProcessorStep",
+    "TruncatedProcessor",
+    "VanillaObservationProcessor",
 ]

src/lerobot/processor/batch_processor.py

@@ -1,254 +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.
-
-"""
-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
-
-from torch import Tensor
-
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-
-from .core import EnvTransition, PolicyAction
-from .pipeline import (
-    ComplementaryDataProcessorStep,
-    ObservationProcessorStep,
-    PolicyActionProcessorStep,
-    ProcessorStep,
-    ProcessorStepRegistry,
-    TransitionKey,
-)
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="to_batch_processor_action")
-class AddBatchDimensionActionStep(PolicyActionProcessorStep):
-    """
-    Processor step to add a batch dimension to a 1D tensor action.
-
-    This is useful for creating a batch of size 1 from a single action sample.
-    """
-
-    def action(self, action: PolicyAction) -> PolicyAction:
-        """
-        Adds a batch dimension to the action if it's a 1D tensor.
-
-        Args:
-            action: The action tensor.
-
-        Returns:
-            The action tensor with an added batch dimension.
-        """
-        if 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:
-                state_value = observation[state_key]
-                if isinstance(state_value, Tensor) and state_value.dim() == 1:
-                    observation[state_key] = state_value.unsqueeze(0)
-
-        # Process single image observation - add batch dim if 3D
-        if OBS_IMAGE in observation:
-            image_value = observation[OBS_IMAGE]
-            if isinstance(image_value, Tensor) and image_value.dim() == 3:
-                observation[OBS_IMAGE] = image_value.unsqueeze(0)
-
-        # Process multiple image observations - add batch dim if 3D
-        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.
-
-        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_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"]
-            if isinstance(task_value, str):
-                complementary_data["task"] = [task_value]
-
-        # Process index field - add batch dim if 0D
-        if "index" in complementary_data:
-            index_value = complementary_data["index"]
-            if isinstance(index_value, Tensor) and index_value.dim() == 0:
-                complementary_data["index"] = index_value.unsqueeze(0)
-
-        # Process task_index field - add batch dim if 0D
-        if "task_index" in complementary_data:
-            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.
-
-        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")
-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.
-        """
-        if transition[TransitionKey.ACTION] is not None:
-            transition = self.to_batch_action_processor(transition)
-        if transition[TransitionKey.OBSERVATION] is not None:
-            transition = self.to_batch_observation_processor(transition)
-        if transition[TransitionKey.COMPLEMENTARY_DATA] is not None:
-            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
-        return features

src/lerobot/processor/converters.py

@@ -1,393 +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 collections.abc import Sequence
-from functools import singledispatch
-from typing import Any
-
-import numpy as np
-import torch
-
-from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, 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)}")
-
-
-@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.
-    """
-    if isinstance(x, torch.Tensor):
-        return x.item() if x.numel() == 1 else x.detach().cpu().numpy()
-    return x
-
-
-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 create_transition(
-    observation: dict[str, Any] | None = None,
-    action: PolicyAction | RobotAction | 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,
-) -> 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 {},
-    }
-
-
-def robot_action_to_transition(action: RobotAction) -> EnvTransition:
-    """
-    Convert a raw robot 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.
-    """
-    if not isinstance(action, dict):
-        raise ValueError(f"Action should be a RobotAction type got {type(action)}")
-    return create_transition(action=action)
-
-
-def observation_to_transition(observation: RobotObservation) -> 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.
-    """
-    if not isinstance(observation, dict):
-        raise ValueError(f"Observation should be a RobotObservation type got {type(observation)}")
-    return create_transition(observation=observation)
-
-
-def transition_to_robot_action(transition: EnvTransition) -> RobotAction:
-    """
-    Extract a raw robot action dictionary for a robot from an `EnvTransition`.
-
-    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.
-
-    Args:
-        transition: The `EnvTransition` containing the action.
-
-    Returns:
-        A dictionary representing the raw robot action.
-    """
-    if not isinstance(transition, dict):
-        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
-
-    action = transition.get(TransitionKey.ACTION)
-    if not isinstance(action, dict):
-        raise ValueError(f"Action should be a RobotAction type (dict) got {type(action)}")
-    return transition.get(TransitionKey.ACTION)
-
-
-def transition_to_policy_action(transition: EnvTransition) -> PolicyAction:
-    """
-    Convert an `EnvTransition` to a `PolicyAction`.
-    """
-    if not isinstance(transition, dict):
-        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
-
-    action = transition.get(TransitionKey.ACTION)
-    if not isinstance(action, PolicyAction):
-        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
-    return action
-
-
-def transition_to_observation(transition: EnvTransition) -> RobotObservation:
-    """
-    Convert an `EnvTransition` to a `RobotObservation`.
-    """
-    if not isinstance(transition, dict):
-        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
-
-    observation = transition.get(TransitionKey.OBSERVATION)
-    if not isinstance(observation, dict):
-        raise ValueError(f"Observation should be a RobotObservation (dict) type got {type(observation)}")
-    return observation
-
-
-def policy_action_to_transition(action: PolicyAction) -> EnvTransition:
-    """
-    Convert a `PolicyAction` to an `EnvTransition`.
-    """
-    if not isinstance(action, PolicyAction):
-        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
-    return create_transition(action=action)
-
-
-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__}")
-
-    action = batch.get("action")
-    if action is not None and not isinstance(action, PolicyAction):
-        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
-
-    # Extract observation and complementary data keys.
-    observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
-    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.
-    """
-    if not isinstance(transition, dict):
-        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
-
-    batch = {
-        "action": transition.get(TransitionKey.ACTION),
-        "next.reward": transition.get(TransitionKey.REWARD, 0.0),
-        "next.done": transition.get(TransitionKey.DONE, False),
-        "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
-        "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(transition: 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 transition

src/lerobot/processor/core.py

@@ -1,56 +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, TypeAlias, TypedDict
-
-import numpy as np
-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"
-
-
-PolicyAction: TypeAlias = torch.Tensor
-RobotAction: TypeAlias = dict[str, Any]
-EnvAction: TypeAlias = np.ndarray
-RobotObservation: TypeAlias = dict[str, Any]
-
-
-EnvTransition = TypedDict(
-    "EnvTransition",
-    {
-        TransitionKey.OBSERVATION.value: dict[str, Any] | None,
-        TransitionKey.ACTION.value: PolicyAction | RobotAction | EnvAction | 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,145 +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 lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-
-from .core import PolicyAction, RobotAction
-from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep
-
-
-@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: PolicyAction) -> RobotAction:
-        if not isinstance(action, PolicyAction):
-            raise ValueError("Only PolicyAction is supported for this processor")
-
-        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]]:
-        for axis in ["x", "y", "z"]:
-            features[PipelineFeatureType.ACTION][f"delta_{axis}"] = 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(RobotActionProcessorStep):
-    """
-    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: RobotAction) -> RobotAction:
-        # 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")
-        delta_y = action.pop("delta_y")
-        delta_z = action.pop("delta_z")
-        gripper = action.pop("gripper")
-
-        # 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]]:
-        for axis in ["x", "y", "z", "gripper"]:
-            features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
-
-        for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
-            features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
-                type=FeatureType.ACTION, shape=(1,)
-            )
-
-        return features

src/lerobot/processor/device_processor.py

@@ -13,178 +13,70 @@
 # 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, TransitionKey
 from lerobot.utils.utils import get_safe_torch_device
 
-from .core import EnvTransition, PolicyAction, 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.
 
-    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.
     """
 
-    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 = "cpu"
 
     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.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:
-                raise ValueError(
-                    f"Invalid float_dtype '{self.float_dtype}'. Available options: {list(self.DTYPE_MAPPING.keys())}"
-                )
-            self._target_float_dtype = self.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)
-
-        # Convert float dtype if specified and tensor is floating point
-        if self._target_float_dtype is not None and tensor.is_floating_point():
-            tensor = tensor.to(dtype=self._target_float_dtype)
-
-        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()
-        action = new_transition.get(TransitionKey.ACTION)
 
-        if action is not None and not isinstance(action, PolicyAction):
-            raise ValueError(f"If action is not None should be a PolicyAction type got {type(action)}")
+        # Process observation tensors
+        observation = transition.get(TransitionKey.OBSERVATION)
+        if observation is not None:
+            new_observation = {
+                k: v.to(self.device, non_blocking=self.non_blocking) if isinstance(v, torch.Tensor) else v
+                for k, v in observation.items()
+            }
+            new_transition[TransitionKey.OBSERVATION] = new_observation
 
-        simple_tensor_keys = [
-            TransitionKey.ACTION,
-            TransitionKey.REWARD,
-            TransitionKey.DONE,
-            TransitionKey.TRUNCATED,
-        ]
+        # Process action tensor
+        action = transition.get(TransitionKey.ACTION)
+        if action is not None and isinstance(action, torch.Tensor):
+            new_transition[TransitionKey.ACTION] = action.to(self.device, non_blocking=self.non_blocking)
 
-        dict_tensor_keys = [
-            TransitionKey.OBSERVATION,
-            TransitionKey.COMPLEMENTARY_DATA,
-        ]
+        # Process reward tensor
+        reward = transition.get(TransitionKey.REWARD)
+        if reward is not None and isinstance(reward, torch.Tensor):
+            new_transition[TransitionKey.REWARD] = reward.to(self.device, non_blocking=self.non_blocking)
 
-        # 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 done tensor
+        done = transition.get(TransitionKey.DONE)
+        if done is not None and isinstance(done, torch.Tensor):
+            new_transition[TransitionKey.DONE] = done.to(self.device, non_blocking=self.non_blocking)
 
-        # 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 truncated tensor
+        truncated = transition.get(TransitionKey.TRUNCATED)
+        if truncated is not None and isinstance(truncated, torch.Tensor):
+            new_transition[TransitionKey.TRUNCATED] = truncated.to(
+                self.device, non_blocking=self.non_blocking
+            )
 
         return new_transition
 
     def get_config(self) -> dict[str, Any]:
-        """
-        Returns the serializable configuration of the processor.
+        """Return configuration for serialization."""
+        return {"device": self.device}
 
-        Returns:
-            A dictionary containing the device and float_dtype settings.
-        """
-        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.
-
-        Device and dtype transformations do not alter the fundamental definition of the features (e.g., shape).
-
-        Args:
-            features: A dictionary of policy features.
-
-        Returns:
-            The original dictionary of policy features.
-        """
+    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features

src/lerobot/processor/factory.py

@@ -1,58 +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 .converters import (
-    observation_to_transition,
-    robot_action_to_transition,
-    transition_to_observation,
-    transition_to_robot_action,
-)
-from .core import RobotAction, RobotObservation
-from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
-
-
-def make_default_teleop_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
-    teleop_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-        steps=[IdentityProcessorStep()],
-        to_transition=robot_action_to_transition,
-        to_output=transition_to_robot_action,
-    )
-    return teleop_action_processor
-
-
-def make_default_robot_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
-    robot_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
-        steps=[IdentityProcessorStep()],
-        to_transition=robot_action_to_transition,
-        to_output=transition_to_robot_action,
-    )
-    return robot_action_processor
-
-
-def make_default_robot_observation_processor() -> RobotProcessorPipeline[RobotObservation, RobotObservation]:
-    robot_observation_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
-        steps=[IdentityProcessorStep()],
-        to_transition=observation_to_transition,
-        to_output=transition_to_observation,
-    )
-    return robot_observation_processor
-
-
-def make_default_processors():
-    teleop_action_processor = make_default_teleop_action_processor()
-    robot_action_processor = make_default_robot_action_processor()
-    robot_observation_processor = make_default_robot_observation_processor()
-    return (teleop_action_processor, robot_action_processor, robot_observation_processor)

src/lerobot/processor/gym_action_processor.py

@@ -1,93 +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 lerobot.configs.types import PipelineFeatureType, PolicyFeature
-
-from .converters import to_tensor
-from .core import EnvAction, PolicyAction
-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: PolicyAction) -> EnvAction:
-        if not isinstance(action, PolicyAction):
-            raise TypeError(
-                f"Expected PolicyAction 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: EnvAction) -> PolicyAction:
-        if not isinstance(action, EnvAction):
-            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