faster canbus

.github/workflows/release.yml

@@ -83,11 +83,11 @@ jobs:
           fi
 
       - name: Remove Tags with Git dependencies
-        # TODO(Steven): Temporary patch to remove pi from PyPi 0.4.0 release due to its reliance on git dependencies.
+        # TODO(Steven): Temporary patch to remove libero and pi from PyPi 0.4.0 release due to its reliance on git dependencies.
         run: |
           echo "::info:: Checking for Git dependencies to remove from pyproject.toml..."
-          grep -E '@ git\+https|lerobot\[pi\]' pyproject.toml | sed 's/^/::warning:: Removing line: /' || true
-          sed -E -i '/@ git\+https|lerobot\[pi\]/d' pyproject.toml
+          grep -E '@ git\+https|lerobot\[pi\]|lerobot\[libero\]' pyproject.toml | sed 's/^/::warning:: Removing line: /' || true
+          sed -E -i '/@ git\+https|lerobot\[pi\]|lerobot\[libero\]/d' pyproject.toml
           echo "::info:: Git dependencies removed. Proceeding with build."
 
       - name: Install build dependencies

.github/workflows/unbound_deps_tests.yml

@@ -70,7 +70,7 @@ jobs:
           echo "Dependencies unbound:" && cat pyproject.toml
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot # TODO(Steven): Make flash-attn optional
+        run: uv sync --all-extras
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv

README.md

@@ -186,7 +186,7 @@ For a full list of optional dependencies, see:
 https://pypi.org/project/lerobot/
 
 > [!NOTE]
-> For lerobot 0.4.0, if you want to install pi tags, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
+> For lerobot 0.4.0, if you want to install libero or pi tags, you will have to do: `pip install "lerobot[pi,libero]@git+https://github.com/huggingface/lerobot.git"`.
 >
 > This will be solved in the next patch release
 

docs/source/_toctree.yml

@@ -15,6 +15,8 @@
     title: Train a Robot with RL
   - local: hilserl_sim
     title: Train RL in Simulation
+  - local: async
+    title: Use Async Inference
   - local: multi_gpu_training
     title: Multi GPU training
   title: "Tutorials"
@@ -39,14 +41,6 @@
     title: NVIDIA GR00T N1.5
   title: "Policies"
 - sections:
-  - local: async
-    title: Use Async Inference
-  - local: rtc
-    title: Real-Time Chunking (RTC)
-  title: "Inference"
-- sections:
-  - local: envhub
-    title: Environments from the Hub
   - local: il_sim
     title: Imitation Learning in Sim
   - local: libero
@@ -63,8 +57,6 @@
     title: Implement your own processor
   - local: processors_robots_teleop
     title: Processors for Robots and Teleoperators
-  - local: env_processor
-    title: Environment Processors
   title: "Robot Processors"
 - sections:
   - local: so101

docs/source/env_processor.mdx

@@ -1,418 +0,0 @@
-# Environment Processors
-
-Environment processors are a critical layer in LeRobot's data processing architecture that handle **environment-specific** transformations, separate from policy-specific processing. This separation of concerns enables cleaner code, better modularity, and easier experimentation with different environments and policies.
-
-## Why Environment Processors?
-
-When working with different robot environments (LIBERO, MetaWorld, Aloha, etc.), each environment often has unique data formats, coordinate systems, and conventions that need standardization **before** policy processing. Without environment processors, these transformations would be:
-
-1. **Hardcoded in environment code** - Making it difficult to experiment with different state representations
-2. **Duplicated across policies** - Each policy would need to handle environment-specific quirks
-3. **Mixed with policy logic** - Violating separation of concerns and making debugging harder
-
-Environment processors solve this by providing a **dedicated processing layer** between raw environment observations and policy inputs.
-
-## The Processing Pipeline
-
-Here's how data flows through the complete processing pipeline during evaluation:
-
-```python
-# In lerobot_eval.py rollout() function:
-
-# 1. Raw environment observation (numpy arrays, various formats)
-raw_observation = env.step(action)
-
-# 2. Convert numpy to torch, normalize images [0,1]
-observation = preprocess_observation(raw_observation)
-
-# 3. Add task metadata (for multi-task environments)
-observation = add_envs_task(env, observation)
-
-# 4. ENVIRONMENT-SPECIFIC preprocessing (NEW!)
-#    - Flatten robot states
-#    - Rotate images to match dataset conventions
-#    - Handle environment-specific coordinate systems
-observation = env_preprocessor(observation)
-
-# 5. POLICY-SPECIFIC preprocessing
-#    - Normalize with dataset statistics
-#    - Add batch dimensions
-#    - Move to GPU
-#    - Tokenize language instructions
-observation = preprocessor(observation)
-
-# 6. Policy inference
-action = policy.select_action(observation)
-
-# 7. POLICY-SPECIFIC postprocessing
-#    - Unnormalize actions
-#    - Remove batch dimensions
-action = postprocessor(action)
-
-# 8. ENVIRONMENT-SPECIFIC postprocessing (NEW!)
-#    - Convert action formats if needed
-#    - Apply environment-specific constraints
-action_transition = {"action": action}
-action_transition = env_postprocessor(action_transition)
-action = action_transition["action"]
-
-# 9. Execute in environment
-env.step(action)
-```
-
-## The Benefits
-
-### 1. **Separation of Concerns**
-
-Environment processors handle transformations specific to the **environment's data format**, while policy processors handle transformations specific to the **model's requirements**.
-
-```python
-# ❌ Before: Mixed concerns
-class LiberoVLAPolicy:
-    def preprocess(self, obs):
-        # Environment-specific: Flatten robot state (shouldn't be in policy!)
-        state = self._flatten_robot_state(obs["robot_state"])
-        # Policy-specific: Normalize with dataset stats
-        state = self.normalizer(state)
-        return state
-
-# ✅ After: Clear separation
-# Environment processor: Handles LIBERO's nested robot state
-env_preprocessor = LiberoProcessorStep()  # Flattens robot_state
-
-# Policy processor: Handles model requirements
-policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
-```
-
-### 2. **Flexibility and Reusability**
-
-The same policy can work with different environment processors, and the same environment processor can work with different policies:
-
-```python
-# Use SmolVLA policy with LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(libero_cfg)
-smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
-
-# Or use ACT policy with the same LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(libero_cfg)
-act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
-```
-
-### 3. **Easier Experimentation**
-
-Want to try different state representations for LIBERO? Just create a new processor:
-
-```python
-# Original: 8D state (pos + quat→axisangle + gripper)
-@ProcessorStepRegistry.register("libero_processor")
-class LiberoProcessorStep(ObservationProcessorStep):
-    def _process_observation(self, obs):
-        eef_pos = robot_state["eef"]["pos"]          # 3D
-        eef_axisangle = quat2axisangle(quat)         # 3D
-        gripper = robot_state["gripper"]["qpos"]     # 2D
-        state = torch.cat([eef_pos, eef_axisangle, gripper], dim=-1)  # 8D
-        return state
-
-# Experiment: Add velocity for better control
-@ProcessorStepRegistry.register("libero_velocity_processor")
-class LiberoVelocityProcessorStep(ObservationProcessorStep):
-    def _process_observation(self, obs):
-        # Include velocities for 14D state
-        eef_pos = robot_state["eef"]["pos"]          # 3D
-        eef_axisangle = quat2axisangle(quat)         # 3D
-        eef_vel = robot_state["eef"]["vel"]          # 3D  (NEW)
-        gripper_pos = robot_state["gripper"]["qpos"] # 2D
-        gripper_vel = robot_state["gripper"]["qvel"] # 3D  (NEW)
-        state = torch.cat([eef_pos, eef_axisangle, eef_vel,
-                          gripper_pos, gripper_vel], dim=-1)  # 14D
-        return state
-```
-
-### 4. **Cleaner Environment Code**
-
-Environments expose **all available data** without needing to know what downstream models will use:
-
-```python
-# LIBERO environment exposes full robot state
-observation = {
-    "pixels": {"image": img, "image2": img2},
-    "robot_state": {
-        "eef": {"pos": ..., "quat": ..., "vel": ..., "mat": ..., "axisangle": ...},
-        "gripper": {"qpos": ..., "qvel": ...},
-        "joints": {"pos": ..., "vel": ...}
-    }
-}
-
-# Environment processor decides what to use
-# Policy processor handles model-specific transformations
-```
-
-## Using Environment Processors
-
-### Factory Function
-
-The `make_env_pre_post_processors` function follows the same pattern as `make_pre_post_processors` for policies:
-
-```python
-from lerobot.envs.factory import make_env_pre_post_processors
-from lerobot.envs.configs import LiberoEnv, PushtEnv
-
-# For LIBERO: Returns LiberoProcessorStep in preprocessor
-libero_cfg = LiberoEnv(task="libero_spatial", camera_name=["agentview"])
-env_preprocessor, env_postprocessor = make_env_pre_post_processors(libero_cfg)
-
-# For other environments: Returns identity processors (no-op)
-pusht_cfg = PushtEnv()
-env_preprocessor, env_postprocessor = make_env_pre_post_processors(pusht_cfg)
-```
-
-### Implementation in `envs/factory.py`
-
-```python
-def make_env_pre_post_processors(
-    env_cfg: EnvConfig,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-]:
-    """
-    Create preprocessor and postprocessor pipelines for environment observations.
-
-    Args:
-        env_cfg: The configuration of the environment.
-
-    Returns:
-        A tuple containing:
-            - preprocessor: Pipeline that processes environment observations
-            - postprocessor: Pipeline that processes environment outputs
-    """
-    # For LIBERO environments, add the LiberoProcessorStep to preprocessor
-    if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
-        preprocessor = PolicyProcessorPipeline(steps=[LiberoProcessorStep()])
-    else:
-        # For all other environments, return an identity preprocessor
-        preprocessor = PolicyProcessorPipeline(steps=[])
-
-    # Postprocessor is currently identity for all environments
-    # Future: Could add environment-specific action transformations
-    postprocessor = PolicyProcessorPipeline(steps=[])
-
-    return preprocessor, postprocessor
-```
-
-### Integration in Evaluation
-
-In `lerobot_eval.py`, the environment processors are created once and used throughout:
-
-```python
-def eval_main(cfg: EvalPipelineConfig):
-    # Create environment
-    envs = make_env(cfg.env, n_envs=cfg.eval.batch_size)
-
-    # Create policy
-    policy = make_policy(cfg=cfg.policy, env_cfg=cfg.env)
-
-    # Create policy processors
-    preprocessor, postprocessor = make_pre_post_processors(
-        policy_cfg=cfg.policy,
-        pretrained_path=cfg.policy.pretrained_path,
-    )
-
-    # Create environment processors (NEW!)
-    env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env)
-
-    # Run evaluation with both processor types
-    eval_policy_all(
-        envs=envs,
-        policy=policy,
-        env_preprocessor=env_preprocessor,      # Environment-specific
-        env_postprocessor=env_postprocessor,    # Environment-specific
-        preprocessor=preprocessor,              # Policy-specific
-        postprocessor=postprocessor,            # Policy-specific
-        n_episodes=cfg.eval.n_episodes,
-    )
-```
-
-## Example: LIBERO Environment Processor
-
-The `LiberoProcessorStep` demonstrates a real-world environment processor:
-
-```python
-from lerobot.processor.pipeline import ObservationProcessorStep
-
-@dataclass
-@ProcessorStepRegistry.register(name="libero_processor")
-class LiberoProcessorStep(ObservationProcessorStep):
-    """
-    Processes LIBERO observations into the LeRobot format.
-
-    **State Processing:**
-    - Extracts end-effector position (3D)
-    - Converts quaternion to axis-angle representation (3D)
-    - Extracts gripper joint positions (2D)
-    - Concatenates into 8D state vector
-
-    **Image Processing:**
-    - Rotates images 180° to match HuggingFaceVLA/libero convention
-    """
-
-    def _process_observation(self, observation):
-        processed_obs = observation.copy()
-
-        # Process images: Flip 180° for camera convention
-        for key in list(processed_obs.keys()):
-            if key.startswith("observation.images."):
-                img = processed_obs[key]
-                img = torch.flip(img, dims=[2, 3])  # Flip H and W
-                processed_obs[key] = img
-
-        # Process robot_state: Flatten to 8D vector
-        if "observation.robot_state" in processed_obs:
-            robot_state = processed_obs.pop("observation.robot_state")
-
-            eef_pos = robot_state["eef"]["pos"]           # (B, 3)
-            eef_quat = robot_state["eef"]["quat"]         # (B, 4)
-            gripper_qpos = robot_state["gripper"]["qpos"] # (B, 2)
-
-            # Convert quaternion to axis-angle
-            eef_axisangle = self._quat2axisangle(eef_quat)  # (B, 3)
-
-            # Concatenate into single state vector
-            state = torch.cat((eef_pos, eef_axisangle, gripper_qpos), dim=-1)
-            state = state.float()
-
-            processed_obs["observation.state"] = state
-
-        return processed_obs
-```
-
-### Why These Transformations?
-
-1. **Image Rotation**: The HuggingFaceVLA/libero dataset has images rotated 180° from the raw LIBERO simulator. The processor handles this convention mismatch so policies trained on the dataset work seamlessly.
-
-2. **State Flattening**: The raw LIBERO environment exposes nested dictionaries with all available state information (position, quaternion, velocity, matrix representation, etc.). The processor:
-   - Selects the relevant components (pos, quat, gripper)
-   - Converts quaternion to axis-angle (more suitable for learning)
-   - Flattens to a single 8D vector that policies expect
-
-3. **Flexibility**: The environment still exposes **all** raw data. If you want to try different state representations (e.g., including velocities, using matrix representation instead of axis-angle), you can create a new processor without modifying the environment code.
-
-## Adding Environment Processors for New Environments
-
-To add environment processors for a new environment:
-
-### 1. Create the Processor Step
-
-```python
-# In src/lerobot/processor/env_processor.py
-
-@dataclass
-@ProcessorStepRegistry.register(name="myenv_processor")
-class MyEnvProcessorStep(ObservationProcessorStep):
-    """Process observations from MyEnv."""
-
-    def _process_observation(self, observation):
-        processed = observation.copy()
-
-        # Your environment-specific transformations
-        if "myenv.specific.state" in processed:
-            state = processed.pop("myenv.specific.state")
-            # Transform to standard format
-            processed["observation.state"] = self._transform_state(state)
-
-        return processed
-```
-
-### 2. Update the Factory
-
-```python
-# In src/lerobot/envs/factory.py
-
-def make_env_pre_post_processors(env_cfg: EnvConfig):
-    if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
-        preprocessor = PolicyProcessorPipeline(steps=[LiberoProcessorStep()])
-    elif isinstance(env_cfg, MyEnvConfig) or "myenv" in env_cfg.type:
-        preprocessor = PolicyProcessorPipeline(steps=[MyEnvProcessorStep()])
-    else:
-        preprocessor = PolicyProcessorPipeline(steps=[])
-
-    postprocessor = PolicyProcessorPipeline(steps=[])
-    return preprocessor, postprocessor
-```
-
-### 3. Use in Evaluation
-
-No changes needed! The evaluation script automatically uses the appropriate processor:
-
-```bash
-lerobot-eval \
-    --policy.path=lerobot/my_policy \
-    --env.type=myenv \  # Automatically uses MyEnvProcessorStep
-    --eval.n_episodes=10
-```
-
-## Future: Environment Postprocessors
-
-Currently, postprocessors are identity (no-op) for all environments. Future use cases include:
-
-### Action Space Transformations
-
-```python
-@dataclass
-class MyEnvActionPostprocessor(ProcessorStep):
-    """Convert policy actions to environment-specific format."""
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        action = transition["action"]
-
-        # Example: Convert from Cartesian to joint space
-        if self.action_space == "joint":
-            action = self.ik_solver(action)
-
-        # Example: Apply environment-specific safety limits
-        action = torch.clamp(action, self.min_action, self.max_action)
-
-        transition["action"] = action
-        return transition
-```
-
-### Coordinate System Conversions
-
-```python
-@dataclass
-class CoordinateTransformPostprocessor(ProcessorStep):
-    """Transform actions between coordinate systems."""
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        action = transition["action"]
-
-        # Example: Policy outputs in world frame, env expects base frame
-        action = self.world_to_base_transform(action)
-
-        transition["action"] = action
-        return transition
-```
-
-## Best Practices
-
-1. **Keep environment processors simple**: They should only handle environment-specific data format issues, not complex learning-related transformations.
-
-2. **Use policy processors for model requirements**: Normalization, batching, device placement, and tokenization belong in policy processors.
-
-3. **Expose all data from environments**: Let processors decide what to use rather than hardcoding choices in the environment.
-
-4. **Document conventions**: Clearly document any coordinate system conventions, camera orientations, or data formats that your processor handles.
-
-5. **Test independently**: Environment processors should be testable without loading full policies or environments.
-
-## Summary
-
-Environment processors provide a **clean separation** between environment-specific data transformations and policy-specific model requirements. This architecture:
-
-- ✅ Enables easy experimentation with different state representations
-- ✅ Allows policies to work seamlessly across different environments
-- ✅ Keeps environment code focused on simulation/hardware interface
-- ✅ Makes processor pipelines more maintainable and debuggable
-- ✅ Follows the single responsibility principle
-
-The key insight: **Environments define data formats, processors standardize them, policies consume standardized data.** Each layer has a clear, focused responsibility.

docs/source/envhub.mdx

@@ -1,424 +0,0 @@
-# Loading Environments from the Hub
-
-The **EnvHub** feature allows you to load simulation environments directly from the Hugging Face Hub with a single line of code. This unlocks a powerful new model for collaboration: instead of environments being locked away inside monolithic libraries, anyone can publish custom environments and share them with the community.
-
-## Overview
-
-With EnvHub, you can:
-
-- Load environments from the Hub instantly
-- Share your custom simulation tasks with the community
-- Version control your environments using Git
-- Distribute complex physics simulations without packaging hassles
-
-## Quick Start
-
-Loading an environment from the Hub is as simple as:
-
-```python
-from lerobot.envs.factory import make_env
-
-# Load a hub environment (requires explicit consent to run remote code)
-env = make_env("lerobot/cartpole-env", trust_remote_code=True)
-```
-
-<Tip warning={true}>
-  **Security Notice**: Loading environments from the Hub executes Python code
-  from third-party repositories. Only use `trust_remote_code=True` with
-  repositories you trust. We strongly recommend pinning to a specific commit
-  hash for reproducibility and security.
-</Tip>
-
-## What is EnvHub?
-
-EnvHub is a framework that allows researchers and developers to:
-
-1. **Publish environments** to the Hugging Face Hub as Git repositories
-2. **Load environments** dynamically without installing them as packages
-3. **Version and track** environment changes using Git semantics
-4. **Discover** new simulation tasks shared by the community
-
-This design means you can go from discovering an interesting environment on the Hub to running experiments in seconds, without worrying about dependency conflicts or complex installation procedures.
-
-## Repository Structure
-
-To make your environment loadable from the Hub, your repository must contain at minimum:
-
-### Required Files
-
-**`env.py`** (or custom Python file)
-
-- Must expose a `make_env(n_envs: int, use_async_envs: bool)` function
-- This function should return one of:
-  - A `gym.vector.VectorEnv` (most common)
-  - A single `gym.Env` (will be automatically wrapped)
-  - A dict mapping `{suite_name: {task_id: VectorEnv}}` (for multi-task benchmarks)
-
-### Optional Files
-
-**`requirements.txt`**
-
-- List any additional dependencies your environment needs
-- Users will need to install these manually before loading your environment
-
-**`README.md`**
-
-- Document your environment: what task it implements, observation/action spaces, rewards, etc.
-- Include usage examples and any special setup instructions
-
-**`.gitignore`**
-
-- Exclude unnecessary files from your repository
-
-### Example Repository Structure
-
-```
-my-environment-repo/
-├── env.py                 # Main environment definition (required)
-├── requirements.txt       # Dependencies (optional)
-├── README.md             # Documentation (recommended)
-├── assets/               # Images, videos, etc. (optional)
-│   └── demo.gif
-└── configs/              # Config files if needed (optional)
-    └── task_config.yaml
-```
-
-## Creating Your Environment Repository
-
-### Step 1: Define Your Environment
-
-Create an `env.py` file with a `make_env` function:
-
-```python
-# env.py
-import gymnasium as gym
-
-def make_env(n_envs: int = 1, use_async_envs: bool = False):
-    """
-    Create vectorized environments for your custom task.
-
-    Args:
-        n_envs: Number of parallel environments
-        use_async_envs: Whether to use AsyncVectorEnv or SyncVectorEnv
-
-    Returns:
-        gym.vector.VectorEnv or dict mapping suite names to vectorized envs
-    """
-    def _make_single_env():
-        # Create your custom environment
-        return gym.make("CartPole-v1")
-
-    # Choose vector environment type
-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-
-    # Create vectorized environment
-    vec_env = env_cls([_make_single_env for _ in range(n_envs)])
-
-    return vec_env
-```
-
-### Step 2: Test Locally
-
-Before uploading, test your environment locally:
-
-```python
-from lerobot.envs.utils import _load_module_from_path, _call_make_env, _normalize_hub_result
-
-# Load your module
-module = _load_module_from_path("./env.py")
-
-# Test the make_env function
-result = _call_make_env(module, n_envs=2, use_async_envs=False)
-normalized = _normalize_hub_result(result)
-
-# Verify it works
-suite_name = next(iter(normalized))
-env = normalized[suite_name][0]
-obs, info = env.reset()
-print(f"Observation shape: {obs.shape if hasattr(obs, 'shape') else type(obs)}")
-env.close()
-```
-
-### Step 3: Upload to the Hub
-
-Upload your repository to Hugging Face:
-
-```bash
-# Install huggingface_hub if needed
-pip install huggingface_hub
-
-# Login to Hugging Face
-huggingface-cli login
-
-# Create a new repository
-huggingface-cli repo create my-custom-env --type space --org my-org
-
-# Initialize git and push
-git init
-git add .
-git commit -m "Initial environment implementation"
-git remote add origin https://huggingface.co/my-org/my-custom-env
-git push -u origin main
-```
-
-Alternatively, use the `huggingface_hub` Python API:
-
-```python
-from huggingface_hub import HfApi
-
-api = HfApi()
-
-# Create repository
-api.create_repo("my-custom-env", repo_type="space")
-
-# Upload files
-api.upload_folder(
-    folder_path="./my-env-folder",
-    repo_id="username/my-custom-env",
-    repo_type="space",
-)
-```
-
-## Loading Environments from the Hub
-
-### Basic Usage
-
-```python
-from lerobot.envs.factory import make_env
-
-# Load from the hub
-envs_dict = make_env(
-    "username/my-custom-env",
-    n_envs=4,
-    trust_remote_code=True
-)
-
-# Access the environment
-suite_name = next(iter(envs_dict))
-env = envs_dict[suite_name][0]
-
-# Use it like any gym environment
-obs, info = env.reset()
-action = env.action_space.sample()
-obs, reward, terminated, truncated, info = env.step(action)
-```
-
-### Advanced: Pinning to Specific Versions
-
-For reproducibility and security, pin to a specific Git revision:
-
-```python
-# Pin to a specific branch
-env = make_env("username/my-env@main", trust_remote_code=True)
-
-# Pin to a specific commit (recommended for papers/experiments)
-env = make_env("username/my-env@abc123def456", trust_remote_code=True)
-
-# Pin to a tag
-env = make_env("username/my-env@v1.0.0", trust_remote_code=True)
-```
-
-### Custom File Paths
-
-If your environment definition is not in `env.py`:
-
-```python
-# Load from a custom file
-env = make_env("username/my-env:custom_env.py", trust_remote_code=True)
-
-# Combine with version pinning
-env = make_env("username/my-env@v1.0:envs/task_a.py", trust_remote_code=True)
-```
-
-### Async Environments
-
-For better performance with multiple environments:
-
-```python
-envs_dict = make_env(
-    "username/my-env",
-    n_envs=8,
-    use_async_envs=True,  # Use AsyncVectorEnv for parallel execution
-    trust_remote_code=True
-)
-```
-
-## URL Format Reference
-
-The hub URL format supports several patterns:
-
-| Pattern              | Description                    | Example                                |
-| -------------------- | ------------------------------ | -------------------------------------- |
-| `user/repo`          | Load `env.py` from main branch | `make_env("lerobot/pusht-env")`        |
-| `user/repo@revision` | Load from specific revision    | `make_env("lerobot/pusht-env@main")`   |
-| `user/repo:path`     | Load custom file               | `make_env("lerobot/envs:pusht.py")`    |
-| `user/repo@rev:path` | Revision + custom file         | `make_env("lerobot/envs@v1:pusht.py")` |
-
-## Multi-Task Environments
-
-For benchmarks with multiple tasks (like LIBERO), return a nested dictionary:
-
-```python
-def make_env(n_envs: int = 1, use_async_envs: bool = False):
-    env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
-
-    # Return dict: {suite_name: {task_id: VectorEnv}}
-    return {
-        "suite_1": {
-            0: env_cls([lambda: gym.make("Task1-v0") for _ in range(n_envs)]),
-            1: env_cls([lambda: gym.make("Task2-v0") for _ in range(n_envs)]),
-        },
-        "suite_2": {
-            0: env_cls([lambda: gym.make("Task3-v0") for _ in range(n_envs)]),
-        }
-    }
-```
-
-## Security Considerations
-
-<Tip warning={true}>
-  **Important**: The `trust_remote_code=True` flag is required to execute
-  environment code from the Hub. This is by design for security.
-</Tip>
-
-When loading environments from the Hub:
-
-1. **Review the code first**: Visit the repository and inspect `env.py` before loading
-2. **Pin to commits**: Use specific commit hashes for reproducibility
-3. **Check dependencies**: Review `requirements.txt` for suspicious packages
-4. **Use trusted sources**: Prefer official organizations or well-known researchers
-5. **Sandbox if needed**: Run untrusted code in isolated environments (containers, VMs)
-
-Example of safe usage:
-
-```python
-# ❌ BAD: Loading without inspection
-env = make_env("random-user/untrusted-env", trust_remote_code=True)
-
-# ✅ GOOD: Review code, then pin to specific commit
-# 1. Visit https://huggingface.co/trusted-org/verified-env
-# 2. Review the env.py file
-# 3. Copy the commit hash
-env = make_env("trusted-org/verified-env@a1b2c3d4", trust_remote_code=True)
-```
-
-## Example: CartPole from the Hub
-
-Here's a complete example using the reference CartPole environment:
-
-```python
-from lerobot.envs.factory import make_env
-import numpy as np
-
-# Load the environment
-envs_dict = make_env("lerobot/cartpole-env", n_envs=4, trust_remote_code=True)
-
-# Get the vectorized environment
-suite_name = next(iter(envs_dict))
-env = envs_dict[suite_name][0]
-
-# Run a simple episode
-obs, info = env.reset()
-done = np.zeros(env.num_envs, dtype=bool)
-total_reward = np.zeros(env.num_envs)
-
-while not done.all():
-    # Random policy
-    action = env.action_space.sample()
-    obs, reward, terminated, truncated, info = env.step(action)
-    total_reward += reward
-    done = terminated | truncated
-
-print(f"Average reward: {total_reward.mean():.2f}")
-env.close()
-```
-
-## Benefits of EnvHub
-
-### For Environment Authors
-
-- **Easy distribution**: No PyPI packaging required
-- **Version control**: Use Git for environment versioning
-- **Rapid iteration**: Push updates instantly
-- **Documentation**: Hub README renders beautifully
-- **Community**: Reach LeRobot users directly
-
-### For Researchers
-
-- **Quick experiments**: Load any environment in one line
-- **Reproducibility**: Pin to specific commits
-- **Discovery**: Browse environments on the Hub
-- **No conflicts**: No need to install conflicting packages
-
-### For the Community
-
-- **Growing ecosystem**: More diverse simulation tasks
-- **Standardization**: Common `make_env` API
-- **Collaboration**: Fork and improve existing environments
-- **Accessibility**: Lower barrier to sharing research
-
-## Troubleshooting
-
-### "Refusing to execute remote code"
-
-You must explicitly pass `trust_remote_code=True`:
-
-```python
-env = make_env("user/repo", trust_remote_code=True)
-```
-
-### "Module X not found"
-
-The hub environment has dependencies you need to install:
-
-```bash
-# Check the repo's requirements.txt and install dependencies
-pip install gymnasium numpy
-```
-
-### "make_env not found in module"
-
-Your `env.py` must expose a `make_env` function:
-
-```python
-def make_env(n_envs: int, use_async_envs: bool):
-    # Your implementation
-    pass
-```
-
-### Environment returns wrong type
-
-The `make_env` function must return:
-
-- A `gym.vector.VectorEnv`, or
-- A single `gym.Env`, or
-- A dict `{suite_name: {task_id: VectorEnv}}`
-
-## Best Practices
-
-1. **Document your environment**: Include observation/action space descriptions, reward structure, and termination conditions in your README
-2. **Add requirements.txt**: List all dependencies with versions
-3. **Test thoroughly**: Verify your environment works locally before pushing
-4. **Use semantic versioning**: Tag releases with version numbers
-5. **Add examples**: Include usage examples in your README
-6. **Keep it simple**: Minimize dependencies when possible
-7. **License your work**: Add a LICENSE file to clarify usage terms
-
-## Future Directions
-
-The EnvHub ecosystem enables exciting possibilities:
-
-- **GPU-accelerated physics**: Share Isaac Gym or Brax environments
-- **Photorealistic rendering**: Distribute environments with advanced graphics
-- **Multi-agent scenarios**: Complex interaction tasks
-- **Real-world simulators**: Digital twins of physical setups
-- **Procedural generation**: Infinite task variations
-- **Domain randomization**: Pre-configured DR pipelines
-
-As more researchers and developers contribute, the diversity and quality of available environments will grow, benefiting the entire robotics learning community.
-
-## See Also
-
-- [Hugging Face Hub Documentation](https://huggingface.co/docs/hub/en/index)
-- [Gymnasium Documentation](https://gymnasium.farama.org/index.html)
-- [Example Hub Environment](https://huggingface.co/lerobot/cartpole-env)

docs/source/groot.mdx

@@ -40,7 +40,7 @@ python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} i
 3. Install LeRobot by running:
 
 ```bash
-pip install lerobot[groot]
+pip install lerobot[groot] # consider also installing libero,dev and test tags
 ```
 
 ## Usage
@@ -83,9 +83,6 @@ accelerate launch \
 
 ### Libero Benchmark Results
 
-> [!NOTE]
-> Follow our instructions for Libero usage: [Libero](./libero)
-
 GR00T has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the Libero dataset and compared the results to the GR00T reference results.
 
 | Benchmark          | LeRobot Implementation | GR00T Reference |

docs/source/installation.mdx

@@ -82,7 +82,7 @@ For a full list of optional dependencies, see:
 https://pypi.org/project/lerobot/
 
 > [!NOTE]
-> For lerobot 0.4.0, if you want to install pi, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`
+> For lerobot 0.4.0, if you want to install libero or pi, you will have to do: `pip install "lerobot[pi,libero]@git+https://github.com/huggingface/lerobot.git"`
 
 ### Troubleshooting
 

docs/source/openarms.mdx

@@ -0,0 +1,328 @@
+# OpenArms Robot
+
+OpenArms is a 7 DOF robotic arm with a gripper, designed by [Enactic, Inc.](https://www.enactic.com/) It uses Damiao motors controlled via CAN bus communication and MIT control mode for smooth, precise motion.
+
+## Hardware Overview
+
+- **7 DOF per arm** (14 DOF total for dual arm setup)
+- **1 gripper per arm** (2 grippers total)
+- **Damiao motors** with 4 different types:
+  - **DM8009** (DM-J8009P-2EC) for shoulders (J1, J2) - high torque
+  - **DM4340** for shoulder rotation and elbow (J3, J4)
+  - **DM4310** (DM-J4310-2EC V1.1) for wrist (J5, J6, J7) and gripper (J8)
+- **24V power supply** required
+- **CAN interface device**:
+  - **Linux**: Any SocketCAN-compatible adapter
+  - **macOS**: CANable, PEAK PCAN-USB, or Kvaser USBcan
+- Proper CAN wiring (CANH, CANL, 120Ω termination)
+
+
+## Motor Configuration
+
+Each arm has the following motor configuration based on the [OpenArm setup guide](https://docs.openarm.dev/software/setup/):
+
+| Joint | Motor | Motor Type | Sender CAN ID | Receiver ID | Description |
+|-------|-------|------------|---------------|-------------|-------------|
+| J1 | joint_1 | DM8009 | 0x01 | 0x11 | Shoulder pan |
+| J2 | joint_2 | DM8009 | 0x02 | 0x12 | Shoulder lift |
+| J3 | joint_3 | DM4340 | 0x03 | 0x13 | Shoulder rotation |
+| J4 | joint_4 | DM4340 | 0x04 | 0x14 | Elbow flex |
+| J5 | joint_5 | DM4310 | 0x05 | 0x15 | Wrist roll |
+| J6 | joint_6 | DM4310 | 0x06 | 0x16 | Wrist pitch |
+| J7 | joint_7 | DM4310 | 0x07 | 0x17 | Wrist rotation |
+| J8 | gripper | DM4310 | 0x08 | 0x18 | Gripper |
+
+For dual arm setups, the left arm uses IDs 0x09-0x10 for joints 1-8 with the same motor types.
+
+## Quick Start 
+
+```bash
+# Install system dependencies
+sudo apt install can-utils iproute2
+
+# Install LeRobot with OpenArms support
+pip install -e ".[openarms]"
+```
+
+## Setup Guide
+
+### Step 1: Motor ID Configuration
+
+**IMPORTANT**: Before using the robot, motors must be configured with the correct CAN IDs.
+
+Refer to the [OpenArm Motor ID Configuration Guide](https://docs.openarm.dev/software/setup/motor-id) for detailed instructions using the Damiao Debugging Tools on Windows.
+
+Key points:
+- Each motor needs a unique **Sender CAN ID** (0x01-0x08)
+- Each motor needs a unique **Receiver/Master ID** (0x11-0x18)
+- Use the Damiao Debugging Tools to set these IDs
+
+### Step 2: Setup CAN Interface
+
+Configure your CAN interface as described in the [OpenArm CAN Setup Guide](https://docs.openarm.dev/software/setup/can-setup):
+
+#### Linux (SocketCAN)
+
+```bash
+# Find your CAN interface
+ip link show
+
+# Configure can0, 1, 2, 3
+sudo ip link set can0 down
+sudo ip link set can0 type can bitrate 1000000
+sudo ip link set can0 up
+
+sudo ip link set can1 down
+sudo ip link set can1 type can bitrate 1000000
+sudo ip link set can1 up
+
+sudo ip link set can2 down
+sudo ip link set can2 type can bitrate 1000000
+sudo ip link set can2 up
+
+sudo ip link set can3 down
+sudo ip link set can3 type can bitrate 1000000
+sudo ip link set can3 up
+
+# Verify configuration
+ip link show can0
+```
+
+or run:
+
+`examples/openarms/setup_can.sh`
+
+### Testing canbus and motor connection
+
+Please run this script to check if all motors can be found and to find your can-fd speed: `python examples/openarms/debug_can_communication.py`
+
+## Usage
+
+### Basic Setup
+
+
+```python
+from lerobot.robots.openarms import OpenArmsFollower
+from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
+
+# Configure for dual arm setup
+config = OpenArmsFollowerConfig(
+    port="can0",
+    can_interface="socketcan",  # Or "auto" for auto-detection
+    id="openarms_dual",
+    is_dual_arm=True,
+)
+
+robot = OpenArmsFollower(config)
+robot.connect()
+```
+
+### Calibration
+
+On first use, you'll need to calibrate the robot:
+
+```python
+robot.calibrate()
+```
+
+The calibration process will:
+1. Disable torque on all motors
+2. Ask you to position arms in **hanging position with grippers closed**
+3. Set this as the zero position
+4. Ask you to move each joint through its full range
+5. Record min/max positions for each joint
+6. Save calibration to file
+
+### Reading Observations
+
+The robot provides comprehensive state information:
+
+```python
+observation = robot.get_observation()
+
+# Observation includes for each motor:
+# - {motor_name}.pos: Position in degrees
+# - {motor_name}.vel: Velocity in degrees/second  
+# - {motor_name}.torque: Motor torque
+# - {camera_name}: Camera images (if configured)
+
+print(f"Right arm joint 1 position: {observation['right_joint_1.pos']:.1f}°")
+print(f"Right arm joint 1 velocity: {observation['right_joint_1.vel']:.1f}°/s")
+print(f"Right arm joint 1 torque: {observation['right_joint_1.torque']:.3f} N·m")
+```
+
+### Sending Actions
+
+```python
+# Send target positions (in degrees)
+action = {
+    "right_joint_1.pos": 45.0,
+    "right_joint_2.pos": -30.0,
+    # ... all joints
+    "right_gripper.pos": 45.0,  # Half-closed
+}
+
+actual_action = robot.send_action(action)
+```
+
+### Gripper Control
+
+```python
+# Open gripper
+robot.open_gripper(arm="right")
+
+# Close gripper  
+robot.close_gripper(arm="right")
+```
+
+## Safety Features
+
+### 1. Maximum Relative Target
+
+Limits how far a joint can move in a single command to prevent sudden movements:
+
+```python
+config = OpenArmsFollowerConfig(
+    port="can0",
+    # Limit all joints to 10 degrees per command
+    max_relative_target=10.0,
+    
+    # Or set per-motor limits
+    max_relative_target={
+        "right_joint_1": 15.0,  # Slower moving joint
+        "right_joint_2": 10.0,
+        "right_gripper": 5.0,   # Very slow gripper
+    }
+)
+```
+
+**How it works**: If current position is 50° and you command 80°, with `max_relative_target=10.0`, the robot will only move to 60° in that step.
+
+### 2. Torque Limits
+
+Control maximum torque output, especially important for grippers and teleoperation:
+
+```python
+config = OpenArmsFollowerConfig(
+    port="can0",
+    # Gripper torque limit (fraction of motor's max torque)
+    gripper_torque_limit=0.5,  # 50% of max torque
+)
+```
+
+Lower torque limits prevent damage when gripping delicate objects.
+
+### 3. MIT Control Gains
+
+Control responsiveness and stability via PID-like gains:
+
+```python
+config = OpenArmsFollowerConfig(
+    port="can0",
+    position_kp=10.0,  # Position gain (higher = more responsive)
+    position_kd=0.5,   # Velocity damping (higher = more damped)
+)
+```
+
+**Guidelines**:
+- **For following (robot)**: Higher gains for responsiveness
+  - `position_kp=10.0`, `position_kd=0.5`
+- **For teleoperation (leader)**: Lower gains or disable torque for manual movement
+  - `manual_control=True` (torque disabled)
+
+### 4. Velocity Limits
+
+Velocity limits are enforced by the Damiao motors based on motor type. For DM4310:
+- Max velocity: 30 rad/s ≈ 1718°/s
+
+The motors will automatically limit velocity to safe values.
+
+## Teleoperation
+
+### Leader Arm Setup
+
+The leader arm is moved manually (torque disabled) to generate commands:
+
+```python
+from lerobot.teleoperators.openarms import OpenArmsLeader
+from lerobot.teleoperators.openarms.config_openarms_leader import OpenArmsLeaderConfig
+
+config = OpenArmsLeaderConfig(
+    port="can1",  # Separate CAN interface for leader
+    id="openarms_leader",
+    manual_control=True,  # Torque disabled for manual movement
+    is_dual_arm=True,
+)
+
+leader = OpenArmsLeader(config)
+leader.connect()
+
+# Read current position as action
+action = leader.get_action()
+# action contains positions for all joints in degrees
+```
+
+### Safety Considerations for Teleoperation
+
+1. **Use separate CAN interfaces** for leader and follower to avoid conflicts
+2. **Enable max_relative_target** on follower to smooth abrupt movements
+3. **Lower torque limits** on follower to prevent damage from tracking errors
+4. **Test with one arm** before enabling dual arm teleoperation
+5. **Have emergency stop** ready (power switch or CAN disable)
+
+```python
+# Recommended follower config for teleoperation
+follower_config = OpenArmsFollowerConfig(
+    port="can0",
+    max_relative_target=5.0,  # Small steps for smooth following
+    gripper_torque_limit=0.3, # Low torque for safety
+    position_kp=5.0,  # Lower gains for gentler following
+    position_kd=0.3,
+)
+```
+
+## Troubleshooting
+
+### Motor Shaking/Unstable
+
+- **Lower control gains**: Reduce `position_kp` and `position_kd`
+- **Check calibration**: Re-run calibration procedure
+- **Verify power**: Insufficient current can cause instability
+- **Check mechanical**: Loose connections, binding, or damaged components
+
+### CAN Bus Errors
+
+```bash
+# Check for errors
+ip -s link show can0
+
+# Reset CAN interface
+sudo ip link set can0 down
+sudo ip link set can0 up
+```
+
+### Control Mode
+
+OpenArms uses **MIT control mode** which allows simultaneous control of:
+- Position (degrees)
+- Velocity (degrees/second)
+- Torque (N·m)
+- Position gain (Kp)
+- Velocity damping (Kd)
+
+### Communication
+
+- **Protocol**: CAN 2.0 at 1 Mbps (or CAN-FD at 5 Mbps)
+- **Frame format**: Standard 11-bit IDs
+- **Update rate**: Typically 50-100 Hz depending on motor count
+- **Latency**: ~10-20ms per motor command
+
+## References
+
+- [OpenArm Official Documentation](https://docs.openarm.dev/)
+- [OpenArm Setup Guide](https://docs.openarm.dev/software/setup/)
+- [Motor ID Configuration](https://docs.openarm.dev/software/setup/motor-id)
+- [CAN Interface Setup](https://docs.openarm.dev/software/setup/can-setup)
+- [Motor Communication Test](https://docs.openarm.dev/software/setup/configure-test)
+- [Damiao Motor Documentation](https://wiki.seeedstudio.com/damiao_series/)
+- [Enactic GitHub](https://github.com/enactic/openarm_can)

docs/source/pi0.mdx

@@ -28,11 +28,6 @@ As described by Physical Intelligence, while AI has achieved remarkable success
    pip install -e ".[pi]"
    ```
 
-   > [!NOTE]
-   > For lerobot 0.4.0, if you want to install pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
-   >
-   > This will be solved in the next patch release
-
 ## Training Data and Capabilities
 
 π₀ is trained on the largest robot interaction dataset to date, combining three key data sources:

docs/source/pi05.mdx

@@ -36,11 +36,6 @@ This diverse training mixture creates a "curriculum" that enables generalization
    pip install -e ".[pi]"
    ```
 
-   > [!NOTE]
-   > For lerobot 0.4.0, if you want to install pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
-   >
-   > This will be solved in the next patch release
-
 ## Usage
 
 To use π₀.₅ in your LeRobot configuration, specify the policy type as:

docs/source/rtc.mdx

@@ -1,188 +0,0 @@
-# Real-Time Chunking (RTC)
-
-Real-Time Chunking (RTC) is an inference-time method that allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency.
-
-These policies generate chunks of future actions (e.g., 50 steps at a time) instead of single actions.
-Because the models are large, producing each chunk takes longer than the time it takes the robot to execute it.
-Naively executing chunks leads to problems such as pauses, jerky transitions, or sudden changes in strategy whenever the next chunk arrives late or disagrees with the previously executed actions.
-
-RTC solves this by asynchronously generating the next chunk while the robot continues executing the current one, and by guiding the new chunk so it aligns smoothly with the portion of the previous chunk that has already been executed.
-
-## How RTC Works (simplified)
-
-RTC lets the robot think ahead while it’s still moving. When the robot is carrying out one chunk of actions, RTC starts creating the next chunk early.
-But since the robot has already moved a bit by the time the new chunk is ready, RTC has to make sure the new chunk still lines up smoothly with what the robot is currently doing.
-
-To do this, RTC treats the beginning of the new chunk like an inpainting or “fill-in-the-gaps” problem:
-it gently adjusts the first part of the new chunk so it blends naturally with the robot’s ongoing motion. The result is no pauses, no sudden jumps.
-
-In technical terms, RTC adds a guidance term to the flow-matching denoising process that forces the overlapping timesteps of the new chunk to stay close to the executed portion of the previous chunk, typically using a soft transition mask.
-
-## Quick Start
-
-### Installation
-
-RTC is built into LeRobot. Just install the policy dependencies you need:
-
-```bash
-# For Pi0 or Pi0.5
-pip install -e ".[pi]"
-
-# For SmolVLA
-pip install -e ".[smolvla]"
-```
-
-### Using RTC with Pi0
-
-You can find a complete reference implementation in [eval_with_real_robot.py](examples/rtc/eval_with_real_robot.py).
-The snippet below provides a simplified pseudo-example of how RTC operates with Pi0 in your pipeline:
-
-```python
-from lerobot.policies.pi0 import PI0Policy, PI0Config
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.action_queue import ActionQueue
-
-# Load Pi0 with RTC enabled
-policy_cfg = PI0Config()
-
-# Enable RTC
-policy_cfg.rtc_config = RTCConfig(
-    enabled=True,
-    execution_horizon=10,  # How many steps to blend with previous chunk
-    max_guidance_weight=10.0,  # How strongly to enforce consistency
-    prefix_attention_schedule=RTCAttentionSchedule.EXP,  # Exponential blend
-)
-
-# Load the policy
-policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
-
-# Now use predict_action_chunk with RTC parameters
-inference_delay = 4  # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
-
-# Initialize the action queue
-action_queue = ActionQueue(policy_cfg.rtc_config)
-
-# Start in a separate thread with the following function
-def get_actions():
-  while True:
-    if should_get_actions:
-
-      prev_actions = action_queue.get_left_over()
-      obs = get_robot_observations(robot)
-
-      # Generate actions WITH RTC
-      actions = policy.predict_action_chunk(
-          obs,
-          inference_delay=inference_delay,
-          prev_chunk_left_over=prev_actions,
-      )
-
-      action_queue.merge(
-          actions, actions, inference_delay
-      )
-
-for step in range(num_steps):
-    action = action_queue.get()
-
-    # Execute the first N actions
-    execute_actions(action)
-```
-
-## Key Parameters
-
-`RTCConfig` has the following parameters to tune:
-
-**`execution_horizon`**: How many timesteps from the previous chunk to maintain consistency with. Higher values mean smoother transitions but potentially less reactivity.
-
-Typical values: 8-12 steps
-
-```python
-RTCConfig(execution_horizon=10)
-```
-
-**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
-
-**`prefix_attention_schedule`**: How to weight consistency across the overlap region.
-
-- `LINEAR`: Linear decay from inference_delay to execution_horizon
-- `EXP`: Exponential decay (recommended for getting started)
-- `ONES`: Full weight across entire execution_horizon
-- `ZEROS`: Binary (full weight up to inference_delay, then zero)
-
-**`inference_delay`**: How many timesteps of inference latency your system has. This is passed to `predict_action_chunk()` rather than the config, since it may vary at runtime.
-
-## Testing RTC Offline
-
-Before running on a real robot, test RTC with dataset samples to visualize how it works:
-
-```bash
-python examples/rtc/eval_dataset.py \
-    --policy.path=lerobot/pi0_libero_finetuned \
-    --dataset.repo_id=HuggingFaceVLA/libero \
-    --rtc.execution_horizon=10 \
-    --rtc.max_guidance_weight=10.0 \
-    --device=cuda
-```
-
-The script generates a visualization of the denoising process, comparing standard generation (left) with RTC (right). In the RTC plots, you can see how the first few steps (blue/purple lines) are guided to match the red ground truth trajectory (previous chunk's tail), ensuring a smooth transition between chunks.
-
-<p align="center">
-  <img
-    src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/flow_matching.png"
-    alt="Denoising steps with and without RTC"
-    width="100%"
-  />
-</p>
-
-## Testing RTC with a Real Robot
-
-```bash
-python examples/rtc/eval_with_real_robot.py \
-    --policy.path=${HF_USERNAME}/policy_repo_id \
-    --robot.type=so100_follower \
-    --robot.port=/dev/tty.usbmodem58FA0834591 \
-    --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-    --task="Move green small object into the purple platform" \
-    --duration=120 \
-    --device=cuda
-```
-
-## How It Differs from the Async Inference in LeRobot
-
-Both RTC and [async inference](./async) improve real-time robot control, but they solve different problems.
-
-| Aspect        | Async Inference                                                            | RTC                                                 |
-| ------------- | -------------------------------------------------------------------------- | --------------------------------------------------- |
-| **Problem**   | Idle frames while waiting for inference                                    | Discontinuities between action chunks               |
-| **Solution**  | Decouple prediction from execution                                         | Guide new chunks to continue smoothly from previous |
-| **Benefit**   | No waiting, continuous action                                              | Smooth transitions, natural motion                  |
-| **Best Used** | Async inference is best used with large models with high inference latency | Flow-matching based policies                        |
-
-**Use both together** for maximum smoothness and reactivity!
-
-## Advanced: Debug Tracking
-
-RTC includes built-in debug tracking to help you understand what's happening during inference:
-
-```python
-# Enable debug tracking
-policy_cfg.rtc_config.debug = True
-policy_cfg.rtc_config.debug_maxlen = 100
-
-# After inference, access debug data
-debug_data = policy.rtc_processor.get_debug_data()
-
-# Visualize denoising steps, corrections, etc.
-from lerobot.policies.rtc.debug_visualizer import RTCDebugVisualizer
-visualizer = RTCDebugVisualizer()
-# ... create plots
-```
-
-See `examples/rtc/eval_dataset.py` for a complete example of visualization.
-
-## References
-
-- [Smooth-As-Butter Robot Policies](https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html) - Excellent technical explanation with real robot results
-- [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/research/real_time_chunking) - Original paper and research
-- [Kinetix RTC Implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix) - Reference implementation from Physical Intelligence

examples/dataset_annotation/visualize_subtask_annotations.py

@@ -1,525 +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.
-
-"""
-Visualize SARM Subtask Annotations
-
-This script creates visualizations of the subtask annotations generated by subtask_annotation.py.
-For each episode, it shows:
-- A timeline with dashed vertical lines at subtask boundaries
-- Sample frames from the episode at key points (start, middle, end of each subtask)
-- Color-coded subtask segments
-
-Usage:
-    python visualize_subtask_annotations.py --repo-id pepijn223/mydataset --video-key observation.images.top --num-episodes 5
-"""
-
-import argparse
-import random
-from pathlib import Path
-
-import cv2
-import matplotlib.pyplot as plt
-import matplotlib.patches as mpatches
-import numpy as np
-import pandas as pd
-from matplotlib.lines import Line2D
-from rich.console import Console
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import load_episodes
-from lerobot.policies.sarm.sarm_utils import SubtaskAnnotation, Subtask, Timestamp
-
-
-def timestamp_to_seconds(timestamp: str) -> float:
-    """Convert MM:SS or SS timestamp to seconds"""
-    parts = timestamp.split(":")
-    if len(parts) == 2:
-        return int(parts[0]) * 60 + int(parts[1])
-    else:
-        return int(parts[0])
-
-
-def load_annotations_from_dataset(dataset_path: Path) -> dict[int, SubtaskAnnotation]:
-    """
-    Load annotations from LeRobot dataset parquet files.
-    
-    Reads subtask annotations from the episodes metadata parquet files.
-    """
-    episodes_dataset = load_episodes(dataset_path)
-    
-    if episodes_dataset is None or len(episodes_dataset) == 0:
-        return {}
-    
-    # Check if subtask columns exist
-    if "subtask_names" not in episodes_dataset.column_names:
-        return {}
-    
-    # Convert to pandas DataFrame for easier access
-    episodes_df = episodes_dataset.to_pandas()
-    
-    annotations = {}
-    
-    for ep_idx in episodes_df.index:
-        subtask_names = episodes_df.loc[ep_idx, "subtask_names"]
-        
-        # Skip episodes without annotations
-        if subtask_names is None or (isinstance(subtask_names, float) and pd.isna(subtask_names)):
-            continue
-        
-        start_times = episodes_df.loc[ep_idx, "subtask_start_times"]
-        end_times = episodes_df.loc[ep_idx, "subtask_end_times"]
-        
-        # Reconstruct SubtaskAnnotation from stored data
-        subtasks = []
-        for i, name in enumerate(subtask_names):
-            # Convert seconds back to MM:SS format
-            start_sec = int(start_times[i])
-            end_sec = int(end_times[i])
-            start_str = f"{start_sec // 60:02d}:{start_sec % 60:02d}"
-            end_str = f"{end_sec // 60:02d}:{end_sec % 60:02d}"
-            
-            subtasks.append(
-                Subtask(
-                    name=name,
-                    timestamps=Timestamp(start=start_str, end=end_str)
-                )
-            )
-        
-        annotations[int(ep_idx)] = SubtaskAnnotation(subtasks=subtasks)
-    
-    return annotations
-
-
-# Color palette for subtasks (colorblind-friendly)
-SUBTASK_COLORS = [
-    "#E69F00",  # Orange
-    "#56B4E9",  # Sky blue
-    "#009E73",  # Bluish green
-    "#F0E442",  # Yellow
-    "#0072B2",  # Blue
-    "#D55E00",  # Vermillion
-    "#CC79A7",  # Reddish purple
-    "#999999",  # Gray
-]
-
-
-def extract_frame_from_video(video_path: Path, timestamp: float) -> np.ndarray | None:
-    """Extract a single frame from video at given timestamp."""
-    cap = cv2.VideoCapture(str(video_path))
-    if not cap.isOpened():
-        return None
-    
-    # Set position to timestamp
-    cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
-    ret, frame = cap.read()
-    cap.release()
-    
-    if ret:
-        # Convert BGR to RGB
-        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    return None
-
-
-def visualize_episode(
-    episode_idx: int,
-    annotation,
-    video_path: Path,
-    video_start_timestamp: float,
-    video_end_timestamp: float,
-    fps: int,
-    output_path: Path,
-    video_key: str,
-):
-    """
-    Create visualization for a single episode.
-    
-    Shows:
-    - Top row: Sample frames from the episode (one per subtask)
-    - Bottom: Timeline with subtask segments and boundary lines
-    """
-    subtasks = annotation.subtasks
-    num_subtasks = len(subtasks)
-    
-    if num_subtasks == 0:
-        print(f"No subtasks found for episode {episode_idx}")
-        return
-    
-    # Calculate episode duration
-    episode_duration = video_end_timestamp - video_start_timestamp
-    
-    # Extract sample frames - get frame from middle of each subtask
-    sample_frames = []
-    frame_timestamps = []
-    
-    for subtask in subtasks:
-        start_sec = timestamp_to_seconds(subtask.timestamps.start)
-        end_sec = timestamp_to_seconds(subtask.timestamps.end)
-        mid_sec = (start_sec + end_sec) / 2
-        
-        # Convert to video timestamp (add video_start_timestamp offset)
-        video_timestamp = video_start_timestamp + mid_sec
-        frame_timestamps.append(mid_sec)
-        
-        frame = extract_frame_from_video(video_path, video_timestamp)
-        sample_frames.append(frame)
-    
-    # Create figure
-    fig = plt.figure(figsize=(16, 10))
-    
-    # Use a dark background for better contrast
-    fig.patch.set_facecolor('#1a1a2e')
-    
-    # Calculate grid layout
-    # Top section: frames (variable number of columns based on subtasks)
-    # Bottom section: timeline
-    
-    # Create gridspec
-    gs = fig.add_gridspec(
-        2, max(num_subtasks, 1), 
-        height_ratios=[2, 1],
-        hspace=0.3,
-        wspace=0.1,
-        left=0.05, right=0.95,
-        top=0.88, bottom=0.1
-    )
-    
-    # Add title
-    fig.suptitle(
-        f"Episode {episode_idx} - Subtask Annotations",
-        fontsize=18,
-        fontweight='bold',
-        color='white',
-        y=0.96
-    )
-    
-    # Add subtitle with video info
-    fig.text(
-        0.5, 0.91,
-        f"Camera: {video_key} | Duration: {episode_duration:.1f}s | {num_subtasks} subtasks",
-        ha='center',
-        fontsize=11,
-        color='#888888'
-    )
-    
-    # Plot sample frames
-    for i, (frame, subtask) in enumerate(zip(sample_frames, subtasks)):
-        ax = fig.add_subplot(gs[0, i])
-        ax.set_facecolor('#16213e')
-        
-        if frame is not None:
-            ax.imshow(frame)
-        else:
-            ax.text(0.5, 0.5, "Frame\nN/A", ha='center', va='center', 
-                   fontsize=12, color='white', transform=ax.transAxes)
-        
-        ax.set_title(
-            f"{subtask.name}",
-            fontsize=10,
-            fontweight='bold',
-            color=SUBTASK_COLORS[i % len(SUBTASK_COLORS)],
-            pad=8
-        )
-        ax.axis('off')
-        
-        # Add frame timestamp below
-        ax.text(
-            0.5, -0.08,
-            f"t={frame_timestamps[i]:.1f}s",
-            ha='center',
-            fontsize=9,
-            color='#888888',
-            transform=ax.transAxes
-        )
-    
-    # Create timeline subplot spanning all columns
-    ax_timeline = fig.add_subplot(gs[1, :])
-    ax_timeline.set_facecolor('#16213e')
-    
-    # Get total duration from last subtask end time
-    total_duration = timestamp_to_seconds(subtasks[-1].timestamps.end)
-    
-    # Draw subtask segments as colored bars
-    bar_height = 0.6
-    bar_y = 0.5
-    
-    for i, subtask in enumerate(subtasks):
-        start_sec = timestamp_to_seconds(subtask.timestamps.start)
-        end_sec = timestamp_to_seconds(subtask.timestamps.end)
-        color = SUBTASK_COLORS[i % len(SUBTASK_COLORS)]
-        
-        # Draw segment bar
-        rect = mpatches.FancyBboxPatch(
-            (start_sec, bar_y - bar_height/2),
-            end_sec - start_sec,
-            bar_height,
-            boxstyle="round,pad=0.02,rounding_size=0.1",
-            facecolor=color,
-            edgecolor='white',
-            linewidth=1.5,
-            alpha=0.85
-        )
-        ax_timeline.add_patch(rect)
-        
-        # Add subtask label inside bar
-        mid_x = (start_sec + end_sec) / 2
-        duration = end_sec - start_sec
-        
-        # Only add text if segment is wide enough
-        if duration > total_duration * 0.08:
-            ax_timeline.text(
-                mid_x, bar_y,
-                subtask.name,
-                ha='center', va='center',
-                fontsize=9,
-                fontweight='bold',
-                color='black' if i in [3] else 'white',  # Yellow needs dark text
-                rotation=0 if duration > total_duration * 0.15 else 45
-            )
-    
-    # Draw boundary lines (dashed vertical lines between subtasks)
-    boundary_times = []
-    for i, subtask in enumerate(subtasks):
-        start_sec = timestamp_to_seconds(subtask.timestamps.start)
-        end_sec = timestamp_to_seconds(subtask.timestamps.end)
-        
-        # Add start boundary (except for first subtask at t=0)
-        if i == 0 and start_sec > 0:
-            boundary_times.append(start_sec)
-        elif i > 0:
-            boundary_times.append(start_sec)
-        
-        # Add end boundary for last subtask
-        if i == len(subtasks) - 1:
-            boundary_times.append(end_sec)
-    
-    # Draw dashed lines at boundaries
-    for t in boundary_times:
-        ax_timeline.axvline(
-            x=t, 
-            ymin=0.1, ymax=0.9,
-            color='white', 
-            linestyle='--', 
-            linewidth=2,
-            alpha=0.9
-        )
-        
-        # Add time label below line
-        ax_timeline.text(
-            t, 0.0,
-            f"{int(t//60):02d}:{int(t%60):02d}",
-            ha='center', va='top',
-            fontsize=8,
-            color='#cccccc'
-        )
-    
-    # Add start line at t=0
-    ax_timeline.axvline(x=0, ymin=0.1, ymax=0.9, color='#00ff00', linestyle='-', linewidth=2.5, alpha=0.9)
-    ax_timeline.text(0, 0.0, "00:00", ha='center', va='top', fontsize=8, color='#00ff00', fontweight='bold')
-    
-    # Configure timeline axes
-    ax_timeline.set_xlim(-total_duration * 0.02, total_duration * 1.02)
-    ax_timeline.set_ylim(-0.3, 1.2)
-    ax_timeline.set_xlabel("Time (seconds)", fontsize=11, color='white', labelpad=10)
-    ax_timeline.set_ylabel("")
-    
-    # Style the axes
-    ax_timeline.spines['top'].set_visible(False)
-    ax_timeline.spines['right'].set_visible(False)
-    ax_timeline.spines['left'].set_visible(False)
-    ax_timeline.spines['bottom'].set_color('#444444')
-    ax_timeline.tick_params(axis='x', colors='#888888', labelsize=9)
-    ax_timeline.tick_params(axis='y', left=False, labelleft=False)
-    
-    # Add x-axis ticks at regular intervals
-    tick_interval = max(1, int(total_duration / 10))
-    ax_timeline.set_xticks(np.arange(0, total_duration + tick_interval, tick_interval))
-    
-    # Add legend explaining line styles
-    legend_elements = [
-        Line2D([0], [0], color='#00ff00', linewidth=2.5, linestyle='-', label='Start'),
-        Line2D([0], [0], color='white', linewidth=2, linestyle='--', label='Subtask boundary'),
-    ]
-    ax_timeline.legend(
-        handles=legend_elements, 
-        loc='upper right',
-        framealpha=0.3,
-        facecolor='#16213e',
-        edgecolor='#444444',
-        fontsize=9,
-        labelcolor='white'
-    )
-    
-    # Save figure
-    plt.savefig(output_path, dpi=150, facecolor=fig.get_facecolor(), edgecolor='none', bbox_inches='tight')
-    plt.close()
-    
-    return output_path
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Visualize SARM subtask annotations",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=True,
-        help="HuggingFace dataset repository ID",
-    )
-    parser.add_argument(
-        "--num-episodes",
-        type=int,
-        default=5,
-        help="Number of random episodes to visualize (default: 5)",
-    )
-    parser.add_argument(
-        "--episodes",
-        type=int,
-        nargs="+",
-        default=None,
-        help="Specific episode indices to visualize (overrides --num-episodes)",
-    )
-    parser.add_argument(
-        "--video-key",
-        type=str,
-        default=None,
-        help="Camera/video key to use. If not specified, uses first available.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default="./subtask_viz",
-        help="Output directory for visualizations (default: ./subtask_viz)",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Random seed for reproducibility",
-    )
-    
-    args = parser.parse_args()
-    
-    console = Console()
-    
-    # Set random seed if specified
-    if args.seed is not None:
-        random.seed(args.seed)
-    
-    console.print(f"\n[cyan]Loading dataset: {args.repo_id}[/cyan]")
-    dataset = LeRobotDataset(args.repo_id, download_videos=True)
-    fps = dataset.fps
-    
-    # Get video key
-    if args.video_key:
-        if args.video_key not in dataset.meta.video_keys:
-            console.print(f"[red]Error: Video key '{args.video_key}' not found[/red]")
-            console.print(f"[yellow]Available: {', '.join(dataset.meta.video_keys)}[/yellow]")
-            return
-        video_key = args.video_key
-    else:
-        video_key = dataset.meta.video_keys[0]
-    
-    console.print(f"[cyan]Using camera: {video_key}[/cyan]")
-    console.print(f"[cyan]FPS: {fps}[/cyan]")
-    
-    # Load annotations
-    console.print(f"\n[cyan]Loading annotations...[/cyan]")
-    annotations = load_annotations_from_dataset(dataset.root)
-    
-    if not annotations:
-        console.print("[red]Error: No annotations found in dataset[/red]")
-        console.print("[yellow]Run subtask_annotation.py first to generate annotations[/yellow]")
-        return
-    
-    console.print(f"[green]Found {len(annotations)} annotated episodes[/green]")
-    
-    # Determine which episodes to visualize
-    if args.episodes:
-        episode_indices = args.episodes
-        # Validate episodes exist
-        for ep in episode_indices:
-            if ep not in annotations:
-                console.print(f"[yellow]Warning: Episode {ep} has no annotation, skipping[/yellow]")
-        episode_indices = [ep for ep in episode_indices if ep in annotations]
-    else:
-        # Random selection
-        available_episodes = list(annotations.keys())
-        num_to_select = min(args.num_episodes, len(available_episodes))
-        episode_indices = random.sample(available_episodes, num_to_select)
-        episode_indices.sort()
-    
-    if not episode_indices:
-        console.print("[red]Error: No valid episodes to visualize[/red]")
-        return
-    
-    console.print(f"[cyan]Visualizing episodes: {episode_indices}[/cyan]")
-    
-    # Create output directory
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    
-    # Generate visualizations
-    for ep_idx in episode_indices:
-        console.print(f"\n[cyan]Processing episode {ep_idx}...[/cyan]")
-        
-        annotation = annotations[ep_idx]
-        
-        # Get video path and timestamps
-        video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, video_key)
-        
-        if not video_path.exists():
-            console.print(f"[red]Video not found: {video_path}[/red]")
-            continue
-        
-        # Get episode-specific timestamps within the video file
-        video_path_key = f"videos/{video_key}/from_timestamp"
-        video_path_key_to = f"videos/{video_key}/to_timestamp"
-        
-        video_start_timestamp = float(dataset.meta.episodes[video_path_key][ep_idx])
-        video_end_timestamp = float(dataset.meta.episodes[video_path_key_to][ep_idx])
-        
-        # Create visualization
-        output_path = output_dir / f"episode_{ep_idx:04d}_subtasks.png"
-        
-        try:
-            visualize_episode(
-                episode_idx=ep_idx,
-                annotation=annotation,
-                video_path=video_path,
-                video_start_timestamp=video_start_timestamp,
-                video_end_timestamp=video_end_timestamp,
-                fps=fps,
-                output_path=output_path,
-                video_key=video_key,
-            )
-            console.print(f"[green]✓ Saved: {output_path}[/green]")
-        except Exception as e:
-            console.print(f"[red]✗ Failed to visualize episode {ep_idx}: {e}[/red]")
-    
-    # Print summary
-    console.print(f"\n[bold green]{'=' * 50}[/bold green]")
-    console.print(f"[bold green]Visualization Complete![/bold green]")
-    console.print(f"[bold green]{'=' * 50}[/bold green]")
-    console.print(f"Output directory: {output_dir.absolute()}")
-    console.print(f"Episodes visualized: {len(episode_indices)}")
-
-
-if __name__ == "__main__":
-    main()
-

examples/openarms/debug_can_communication.py

@@ -0,0 +1,416 @@
+#!/usr/bin/env python3
+"""
+Comprehensive debug script for OpenArms CAN FD communication.
+Tests all 4 CAN interfaces with CAN FD support.
+"""
+
+import can
+import time
+import sys
+import subprocess
+
+def check_can_interface(port):
+    """Check if CAN interface is UP and configured."""
+    try:
+        result = subprocess.run(['ip', 'link', 'show', port], 
+                              capture_output=True, text=True)
+        if result.returncode != 0:
+            return False, "Interface not found", None
+        
+        output = result.stdout
+        if 'UP' not in output:
+            return False, "Interface is DOWN", None
+        
+        # Check if CAN FD is enabled
+        is_fd = 'fd on' in output.lower() or 'canfd' in output.lower()
+        
+        return True, "Interface is UP", is_fd
+    except FileNotFoundError:
+        return None, "Cannot check (ip command not found)", None
+
+
+def test_motor_on_interface(bus, motor_id, timeout=2.0, use_fd=False):
+    """
+    Test a single motor and return all responses.
+    
+    Returns:
+        list of (arbitration_id, data) tuples for all responses received
+    """
+    # Send enable command
+    enable_msg = can.Message(
+        arbitration_id=motor_id,
+        data=[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFC],
+        is_extended_id=False,
+        is_fd=use_fd
+    )
+    
+    try:
+        bus.send(enable_msg)
+    except Exception as e:
+        return None, f"Send error: {e}"
+    
+    # Listen for responses
+    responses = []
+    start_time = time.time()
+    
+    while time.time() - start_time < timeout:
+        msg = bus.recv(timeout=0.1)
+        if msg:
+            responses.append((msg.arbitration_id, msg.data, msg.is_fd if hasattr(msg, 'is_fd') else False))
+    
+    # Send disable command
+    disable_msg = can.Message(
+        arbitration_id=motor_id,
+        data=[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFD],
+        is_extended_id=False,
+        is_fd=use_fd
+    )
+    try:
+        bus.send(disable_msg)
+    except:
+        pass
+    
+    return responses, None
+
+
+def test_interface(port, interface_type="socketcan", use_can_fd=True):
+    """Test all 8 motors on a single CAN interface."""
+    
+    results = {
+        'interface': port,
+        'status': None,
+        'is_fd': use_can_fd,
+        'motors': {}
+    }
+    
+    # Check interface status
+    status_ok, status_msg, interface_has_fd = check_can_interface(port)
+    
+    if interface_has_fd is not None:
+        results['interface_fd_enabled'] = interface_has_fd
+        if use_can_fd and not interface_has_fd:
+            status_msg += " (CAN FD NOT enabled on interface!)"
+        elif interface_has_fd:
+            status_msg += " (CAN FD enabled)"
+    
+    results['status'] = status_msg
+    
+    if status_ok is False:
+        return results
+    
+    # Try to connect
+    try:
+        if use_can_fd:
+            print(f"  Connecting to {port} with CAN FD (1 Mbps / 5 Mbps)...")
+            bus = can.interface.Bus(
+                channel=port,
+                interface=interface_type,
+                bitrate=1000000,
+                data_bitrate=5000000,
+                fd=True
+            )
+        else:
+            print(f"  Connecting to {port} with CAN 2.0 (1 Mbps)...")
+            bus = can.interface.Bus(
+                channel=port,
+                interface=interface_type,
+                bitrate=1000000
+            )
+    except Exception as e:
+        results['status'] = f"Connection failed: {e}"
+        return results
+    
+    try:
+        # Clear any pending messages
+        while bus.recv(timeout=0.01):
+            pass
+        
+        # Test each motor (0x01 to 0x08)
+        for motor_id in range(0x01, 0x09):
+            responses, error = test_motor_on_interface(bus, motor_id, timeout=1.0, use_fd=use_can_fd)
+            
+            if error:
+                results['motors'][motor_id] = {'error': error}
+            elif responses:
+                results['motors'][motor_id] = {
+                    'found': True,
+                    'responses': responses
+                }
+            else:
+                results['motors'][motor_id] = {
+                    'found': False,
+                    'responses': []
+                }
+            
+            time.sleep(0.05)  # Small delay between motors
+        
+    finally:
+        bus.shutdown()
+    
+    return results
+
+
+def print_results(all_results):
+    """Print formatted results for all interfaces."""
+    
+    print("SUMMARY - Motors Found on Each Interface")
+    
+    motor_names = {
+        0x01: "joint_1 (Shoulder pan)",
+        0x02: "joint_2 (Shoulder lift)",
+        0x03: "joint_3 (Shoulder rotation)",
+        0x04: "joint_4 (Elbow flex)",
+        0x05: "joint_5 (Wrist roll)",
+        0x06: "joint_6 (Wrist pitch)",
+        0x07: "joint_7 (Wrist rotation)",
+        0x08: "gripper",
+    }
+    
+    total_found = 0
+    
+    for result in all_results:
+        interface = result['interface']
+        status = result['status']
+        
+        print(f"{interface}: {status}")
+        if result.get('is_fd'):
+            print(f"  Mode: CAN FD")
+        else:
+            print(f"  Mode: CAN 2.0")
+        
+        if 'Connection failed' in status or 'DOWN' in status:
+            print(f"  ⚠ Cannot test {interface}")
+            continue
+        
+        motors_found = 0
+        
+        for motor_id in range(0x01, 0x09):
+            motor_data = result['motors'].get(motor_id, {})
+            motor_name = motor_names.get(motor_id, "Unknown")
+            
+            if motor_data.get('error'):
+                print(f"  Motor 0x{motor_id:02X} ({motor_name}): ✗ {motor_data['error']}")
+            elif motor_data.get('found'):
+                motors_found += 1
+                total_found += 1
+                responses = motor_data['responses']
+                print(f"  Motor 0x{motor_id:02X} ({motor_name}): ✓ FOUND")
+                
+                for resp_id, data, is_fd in responses:
+                    data_hex = data.hex()
+                    fd_flag = " [FD]" if is_fd else " [2.0]"
+                    print(f"    → Response from 0x{resp_id:02X}{fd_flag}: {data_hex}")
+            else:
+                print(f"  Motor 0x{motor_id:02X} ({motor_name}): ✗ No response")
+        
+        print(f"\n  Summary: {motors_found}/8 motors found on {interface}")
+    
+    # Overall summary
+    print("OVERALL SUMMARY")
+    print(f"Total motors found across all interfaces: {total_found}")
+    
+    # Analyze configuration
+    print("DIAGNOSIS")
+    
+    for result in all_results:
+        interface = result['interface']
+        motors_found = sum(1 for m in result['motors'].values() if m.get('found'))
+        
+        if motors_found == 0:
+            print(f"\n⚠ {interface}: NO MOTORS FOUND")
+            print("  Possible issues:")
+            print("    1. CAN FD mode mismatch (interface vs motor configuration)")
+            print("    2. Missing 120Ω termination resistors at BOTH cable ends")
+            print("    3. Motor timeout parameter set incorrectly (should NOT be 0)")
+            print("    4. CANH/CANL wiring issue")
+            print("    5. Cable too long (>40m for CAN FD at 5Mbps)")
+            
+            # Check FD mismatch
+            if result.get('is_fd') and not result.get('interface_fd_enabled'):
+                print("    ⚠️ CRITICAL: Trying CAN FD but interface NOT configured for FD!")
+                print(f"       Fix: sudo ip link set {interface} type can bitrate 1000000 dbitrate 5000000 fd on")
+                
+        elif motors_found < 8:
+            print(f"\n⚠ {interface}: Only {motors_found}/8 motors responding")
+            print("  Check power and connections for missing motors")
+        else:
+            print(f"\n✓ {interface}: All 8 motors responding correctly!")
+    
+    # Check for unexpected response IDs
+    print("RESPONSE ID ANALYSIS")
+    
+    for result in all_results:
+        interface = result['interface']
+        unexpected = []
+        
+        for motor_id, motor_data in result['motors'].items():
+            if motor_data.get('found'):
+                expected_id = motor_id + 0x10
+                actual_ids = [resp[0] for resp in motor_data['responses']]
+                
+                if expected_id not in actual_ids:
+                    unexpected.append((motor_id, actual_ids))
+        
+        if unexpected:
+            print(f"\n⚠ {interface}: Unexpected response IDs detected")
+            for motor_id, actual_ids in unexpected:
+                expected_id = motor_id + 0x10
+                print(f"  Motor 0x{motor_id:02X}: Expected 0x{expected_id:02X}, "
+                      f"got {[f'0x{id:02X}' for id in actual_ids]}")
+            print("  → Motor Master IDs need reconfiguration")
+        else:
+            motors_found = sum(1 for m in result['motors'].values() if m.get('found'))
+            if motors_found > 0:
+                print(f"\n✓ {interface}: All responding motors use correct IDs")
+
+
+def test_communication_speed(interface, motor_id, num_iterations=100):
+    """
+    Test communication speed with a motor.
+    
+    Returns:
+        tuple: (hz, avg_latency_ms) or (None, None) if test failed
+    """
+    try:
+        # Connect to interface
+        bus = can.interface.Bus(
+            channel=interface,
+            interface="socketcan",
+            bitrate=1000000,
+            data_bitrate=5000000,
+            fd=True
+        )
+        
+        # Send refresh commands and measure round-trip time
+        latencies = []
+        successful = 0
+        
+        for _ in range(num_iterations):
+            start = time.perf_counter()
+            
+            # Send enable command (lightweight operation)
+            enable_msg = can.Message(
+                arbitration_id=motor_id,
+                data=[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFC],
+                is_extended_id=False,
+                is_fd=True
+            )
+            bus.send(enable_msg)
+            
+            # Wait for response
+            msg = bus.recv(timeout=0.1)
+            
+            if msg:
+                latency = (time.perf_counter() - start) * 1000  # Convert to ms
+                latencies.append(latency)
+                successful += 1
+        
+        bus.shutdown()
+        
+        if successful > 0:
+            avg_latency = sum(latencies) / len(latencies)
+            hz = 1000.0 / avg_latency if avg_latency > 0 else 0
+            return hz, avg_latency
+        
+        return None, None
+        
+    except Exception as e:
+        print(f"    Speed test error: {e}")
+        return None, None
+
+
+def main():
+    """Main function to test all CAN interfaces with CAN FD."""
+    
+    print("\nThis will test all 4 CAN interfaces (can0-can3) with CAN FD")
+    print("Testing motors 0x01-0x08 on each interface")
+    print()
+    print("Make sure:")
+    print("  ✓ Motors are powered (24V)")
+    print("  ✓ CAN interfaces configured with FD mode:")
+    print("    ./examples/openarms/setup_can.sh")
+    print("  ✓ Motor 'timeout' parameter NOT set to 0 (use Damiao tools)")
+    print("  ✓ CAN wiring includes 120Ω termination at BOTH ends")
+    print()
+    
+    input("Press ENTER to start testing...")
+    
+    # Test all 4 interfaces with CAN FD
+    all_results = []
+    
+    for i in range(4):
+        interface = f"can{i}"
+        print(f"Testing {interface}...")
+        
+        result = test_interface(interface, use_can_fd=True)
+        all_results.append(result)
+        
+        # Quick status
+        if 'Connection failed' in result['status'] or 'DOWN' in result['status']:
+            print(f"  ⚠ {interface}: {result['status']}")
+        else:
+            motors_found = sum(1 for m in result['motors'].values() if m.get('found'))
+            print(f"  {interface}: {motors_found}/8 motors found")
+        
+        time.sleep(0.2)
+    
+    # Print detailed results
+    print_results(all_results)
+    
+    print("Testing Complete!")
+    
+    all_found = sum(sum(1 for m in r['motors'].values() if m.get('found')) for r in all_results)
+    
+    if all_found == 0:
+        print("\n⚠️ CRITICAL: No motors found on any interface!")
+        print("\nTop issues to check:")
+        print("  1. Motor 'timeout' parameter (use Damiao tools to set > 0)")
+        print("  2. CAN FD not enabled (run ./examples/openarms/setup_can.sh)")
+        print("  3. Missing termination resistors")
+        print("\nTry:")
+        print("  a) Check motor parameters with Damiao Debugging Tools")
+        print("  b) Verify CAN FD is enabled: ip -d link show can0 | grep fd")
+        print("  c) Run setup script: ./examples/openarms/setup_can.sh")
+    else:
+        # Run speed test on interfaces with motors
+        print("COMMUNICATION SPEED TEST")
+        print("\nTesting maximum communication frequency...")
+        
+        for result in all_results:
+            interface = result['interface']
+            
+            # Find first responding motor
+            responding_motor = None
+            for motor_id, motor_data in result['motors'].items():
+                if motor_data.get('found'):
+                    responding_motor = motor_id
+                    break
+            
+            if responding_motor:
+                print(f"\n{interface}: Testing with motor 0x{responding_motor:02X}...")
+                hz, latency = test_communication_speed(interface, responding_motor, num_iterations=100)
+                
+                if hz:
+                    print(f"  ✓ Max frequency: {hz:.1f} Hz")
+                    print(f"  ✓ Avg latency: {latency:.2f} ms")
+                    print(f"  ✓ Commands per second: ~{int(hz)}")
+                else:
+                    print(f"  ✗ Speed test failed")
+            else:
+                print(f"\n{interface}: No motors found, skipping speed test")
+
+        print()
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\n\nTesting interrupted by user.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"\nUnexpected error: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+

examples/openarms/setup_can.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+# Setup all OpenArms CAN interfaces with CAN FD
+
+set -e
+
+echo "=========================================="
+echo "OpenArms CAN FD Interface Setup"
+echo "=========================================="
+echo ""
+echo "Mode: CAN FD"
+echo "  - Nominal bitrate: 1 Mbps"
+echo "  - Data bitrate: 5 Mbps"
+echo ""
+echo "Configuring interfaces can0, can1, can2, can3..."
+echo ""
+
+# Configure each CAN interface with CAN FD
+for i in 0 1 2 3; do
+    interface="can$i"
+    
+    # Check if interface exists
+    if ! ip link show "$interface" &> /dev/null; then
+        echo "⚠ $interface: Not found, skipping"
+        continue
+    fi
+    
+    # Bring down interface
+    sudo ip link set "$interface" down 2>/dev/null
+    
+    # Configure CAN FD mode
+    sudo ip link set "$interface" type can \
+        bitrate 1000000 \
+        dbitrate 5000000 \
+        fd on
+    
+    # Bring up interface
+    sudo ip link set "$interface" up
+    
+    # Verify configuration
+    if ip link show "$interface" | grep -q "UP"; then
+        echo "✓ $interface: Configured and UP"
+    else
+        echo "✗ $interface: Failed to bring UP"
+    fi
+done
+
+echo ""
+echo "=========================================="
+echo "Verification"
+echo "=========================================="
+echo ""
+
+# Show detailed status for each interface
+for i in 0 1 2 3; do
+    interface="can$i"
+    if ip link show "$interface" &> /dev/null; then
+        echo "$interface:"
+        # Show key parameters
+        ip -d link show "$interface" | grep -E "can|state|bitrate|dbitrate" | head -3
+        echo ""
+    fi
+done
+
+echo "=========================================="
+echo "Setup Complete!"
+echo "=========================================="
+echo ""
+echo "All interfaces configured for CAN FD mode"
+echo ""
+echo "Next steps:"
+echo "  1. Test motors: python debug_can_communication.py"
+echo "  2. Run teleoperation: python examples/openarms/teleop.py"
+echo ""

examples/openarms/teleop.py

@@ -0,0 +1,148 @@
+"""
+OpenArms Teleoperation Example - Full Dual Arms
+
+This script demonstrates teleoperation of OpenArms follower robot using an OpenArms leader arm.
+It first calibrates both devices, then enters a teleoperation loop for both arms.
+"""
+
+import time
+
+from lerobot.robots.openarms.openarms_follower import OpenArmsFollower
+from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
+from lerobot.teleoperators.openarms.openarms_leader import OpenArmsLeader
+from lerobot.teleoperators.openarms.config_openarms_leader import OpenArmsLeaderConfig
+
+
+follower_config = OpenArmsFollowerConfig(
+    port_left="can0",   # CAN interface for follower left arm
+    port_right="can1",  # CAN interface for follower right arm
+    can_interface="socketcan",  # Linux SocketCAN
+    id="openarms_follower",
+    disable_torque_on_disconnect=True,
+    max_relative_target=5.0,  # Safety limit
+)
+
+
+leader_config = OpenArmsLeaderConfig(
+    port_left="can2",   # CAN interface for leader left arm
+    port_right="can3",  # CAN interface for leader right arm
+    can_interface="socketcan",  # Linux SocketCAN
+    id="openarms_leader",
+    manual_control=True,  # Enable manual control (torque disabled)
+)
+
+print("=" * 60)
+print("OpenArms Teleoperation - Full Dual Arms")
+print("=" * 60)
+
+# Initialize devices
+print("\n[1/4] Initializing devices...")
+follower = OpenArmsFollower(follower_config)
+leader = OpenArmsLeader(leader_config)
+
+# Connect and calibrate follower
+print("\n[2/4] Connecting and calibrating follower robot...")
+print("Note: If you have existing calibration, just press ENTER to use it.")
+follower.connect(calibrate=True)
+
+# Connect and calibrate leader
+print("\n[3/4] Connecting and calibrating leader arm...")
+print("Note: The leader arm will have torque disabled for manual control.")
+leader.connect(calibrate=True)
+
+# Wait for user to be ready
+print("\n[4/4] Ready for teleoperation!")
+print("\nBoth arms will be controlled (16 motors total):")
+print("  RIGHT ARM: joints 1-7 + gripper")
+print("  LEFT ARM: joints 1-7 + gripper")
+
+print("\nPress ENTER to start teleoperation...")
+input()
+
+print("\nTeleoperation started! Move both leader arms.")
+print("Press Ctrl+C to stop.\n")
+
+# All joints for both arms (16 motors total)
+all_joints = [
+    # Right arm
+    "right_joint_1",
+    "right_joint_2",
+    "right_joint_3",
+    "right_joint_4",
+    "right_joint_5",
+    "right_joint_6",
+    "right_joint_7",
+    "right_gripper",
+    # Left arm
+    "left_joint_1",
+    "left_joint_2",
+    "left_joint_3",
+    "left_joint_4",
+    "left_joint_5",
+    "left_joint_6",
+    "left_joint_7",
+    "left_gripper",
+]
+
+# Performance monitoring
+loop_times = []
+start_time = time.perf_counter()
+last_print_time = start_time
+
+try:
+    while True:
+        loop_start = time.perf_counter()
+        
+        # Get action from leader
+        leader_action = leader.get_action()
+        
+        # Filter to only position data for all joints (both arms)
+        joint_action = {}
+        for joint in all_joints:
+            pos_key = f"{joint}.pos"
+            if pos_key in leader_action:
+                joint_action[pos_key] = leader_action[pos_key]
+        
+        # Send action to follower (both arms)
+        if joint_action:
+            follower.send_action(joint_action)
+        
+        # Measure loop time
+        loop_end = time.perf_counter()
+        loop_time = loop_end - loop_start
+        loop_times.append(loop_time)
+        
+        # Print stats every 2 seconds
+        if loop_end - last_print_time >= 2.0:
+            if loop_times:
+                avg_time = sum(loop_times) / len(loop_times)
+                current_hz = 1.0 / avg_time if avg_time > 0 else 0
+                min_time = min(loop_times)
+                max_time = max(loop_times)
+                max_hz = 1.0 / min_time if min_time > 0 else 0
+                min_hz = 1.0 / max_time if max_time > 0 else 0
+                
+                print(f"[Hz Stats] Avg: {current_hz:.1f} Hz | "
+                      f"Range: {min_hz:.1f}-{max_hz:.1f} Hz | "
+                      f"Avg loop time: {avg_time*1000:.1f} ms")
+                
+                # Reset for next measurement window
+                loop_times = []
+                last_print_time = loop_end
+            
+except KeyboardInterrupt:
+    print("\n\nStopping teleoperation...")
+finally:
+    # Disconnect devices
+    print("Disconnecting devices...")
+    try:
+        follower.disconnect()
+    except Exception as e:
+        print(f"Error disconnecting follower: {e}")
+    
+    try:
+        leader.disconnect()
+    except Exception as e:
+        print(f"Error disconnecting leader: {e}")
+    
+    print("Done!")

examples/port_datasets/slurm_aggregate_shards.py

@@ -15,12 +15,16 @@
 # limitations under the License.
 
 import argparse
+import logging
 from pathlib import Path
 
 from datatrove.executor import LocalPipelineExecutor
 from datatrove.executor.slurm import SlurmPipelineExecutor
 from datatrove.pipeline.base import PipelineStep
-from port_droid import DROID_SHARDS
+from port_datasets.droid_rlds.port_droid import DROID_SHARDS
+
+from lerobot.datasets.aggregate import aggregate_datasets
+from lerobot.utils.utils import init_logging
 
 
 class AggregateDatasets(PipelineStep):
@@ -34,11 +38,6 @@ class AggregateDatasets(PipelineStep):
         self.aggr_repo_id = aggregated_repo_id
 
     def run(self, data=None, rank: int = 0, world_size: int = 1):
-        import logging
-
-        from lerobot.datasets.aggregate import aggregate_datasets
-        from lerobot.utils.utils import init_logging
-
         init_logging()
 
         # Since aggregate_datasets already handles parallel processing internally,

examples/port_datasets/slurm_port_shards.py

@@ -20,7 +20,7 @@ from pathlib import Path
 from datatrove.executor import LocalPipelineExecutor
 from datatrove.executor.slurm import SlurmPipelineExecutor
 from datatrove.pipeline.base import PipelineStep
-from port_droid import DROID_SHARDS
+from port_datasets.droid_rlds.port_droid import DROID_SHARDS
 
 
 class PortDroidShards(PipelineStep):
@@ -35,7 +35,7 @@ class PortDroidShards(PipelineStep):
 
     def run(self, data=None, rank: int = 0, world_size: int = 1):
         from datasets.utils.tqdm import disable_progress_bars
-        from port_droid import port_droid, validate_dataset
+        from port_datasets.droid_rlds.port_droid import port_droid, validate_dataset
 
         from lerobot.utils.utils import init_logging
 

examples/port_datasets/slurm_upload.py

@@ -24,7 +24,7 @@ from datatrove.executor.slurm import SlurmPipelineExecutor
 from datatrove.pipeline.base import PipelineStep
 from huggingface_hub import HfApi
 from huggingface_hub.constants import REPOCARD_NAME
-from port_droid import DROID_SHARDS
+from port_datasets.droid_rlds.port_droid import DROID_SHARDS
 
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
 from lerobot.datasets.utils import create_lerobot_dataset_card
@@ -185,11 +185,11 @@ class UploadDataset(PipelineStep):
 
 
 def make_upload_executor(
-    repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, private=False, slurm=True
+    repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
 ):
     kwargs = {
         "pipeline": [
-            UploadDataset(repo_id, private=private),
+            UploadDataset(repo_id),
         ],
         "logging_dir": str(logs_dir / job_name),
     }
@@ -267,12 +267,6 @@ def main():
         default="1950M",
         help="Memory per cpu that each worker will use.",
     )
-    parser.add_argument(
-        "--private",
-        action="store_true",
-        default=False,
-        help="Whether to create a private repository.",
-    )
 
     init_logging()
 

examples/rtc/eval_dataset.py

@@ -1,951 +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.
-
-"""
-Evaluate Real-Time Chunking (RTC) performance on dataset samples.
-
-This script takes two random samples from a dataset:
-- Uses actions from the first sample as previous chunk
-- Generates new actions for the second sample with and without RTC
-
-It compares action predictions with and without RTC on dataset samples,
-measuring consistency and ground truth alignment.
-
-Usage:
-    # Basic usage with smolvla policy
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --dataset.repo_id=helper2424/check_rtc \
-        --rtc.execution_horizon=8 \
-        --device=mps \
-        --rtc.max_guidance_weight=10.0 \
-        --rtc.prefix_attention_schedule=EXP \
-        --seed=10
-
-    # Basic usage with pi0.5 policy
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=lerobot/pi05_libero_finetuned \
-        --dataset.repo_id=HuggingFaceVLA/libero \
-        --rtc.execution_horizon=10 \
-        --device=mps
-        --seed=10
-
-    # Basic usage with pi0.5 policy with cuda device
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=lerobot/pi05_libero_finetuned \
-        --dataset.repo_id=HuggingFaceVLA/libero \
-        --rtc.execution_horizon=8 \
-        --device=cuda
-
-    # Basic usage with pi0 policy with cuda device
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=lerobot/pi0_libero_finetuned \
-        --dataset.repo_id=HuggingFaceVLA/libero \
-        --rtc.execution_horizon=8 \
-        --device=cuda
-
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=lipsop/reuben_pi0 \
-        --dataset.repo_id=ReubenLim/so101_cube_in_cup \
-        --rtc.execution_horizon=8 \
-        --device=cuda
-
-    # With torch.compile for faster inference (PyTorch 2.0+)
-    # Note: CUDA graphs disabled by default due to in-place ops in denoising loop
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --dataset.repo_id=helper2424/check_rtc \
-        --rtc.execution_horizon=8 \
-        --device=mps \
-        --use_torch_compile=true \
-        --torch_compile_mode=max-autotune
-
-    # With torch.compile on CUDA (CUDA graphs disabled by default)
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --dataset.repo_id=helper2424/check_rtc \
-        --rtc.execution_horizon=8 \
-        --device=cuda \
-        --use_torch_compile=true \
-        --torch_compile_mode=reduce-overhead
-
-    # Enable CUDA graphs (advanced - may cause tensor aliasing errors)
-    uv run python examples/rtc/eval_dataset.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --dataset.repo_id=helper2424/check_rtc \
-        --use_torch_compile=true \
-        --torch_compile_backend=inductor \
-        --torch_compile_mode=max-autotune \
-        --torch_compile_disable_cudagraphs=false
-"""
-
-import gc
-import logging
-import os
-import random
-from dataclasses import dataclass, field
-
-import numpy as np
-import torch
-
-try:
-    import matplotlib.pyplot as plt
-
-    MATPLOTLIB_AVAILABLE = True
-except ImportError:
-    MATPLOTLIB_AVAILABLE = False
-    plt = None
-
-from lerobot.configs import parser
-from lerobot.configs.default import DatasetConfig
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.datasets.factory import resolve_delta_timestamps
-from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.debug_visualizer import RTCDebugVisualizer
-from lerobot.utils.hub import HubMixin
-from lerobot.utils.utils import init_logging
-
-
-def set_seed(seed: int):
-    """Set random seed for reproducibility."""
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-    if torch.backends.mps.is_available():
-        torch.mps.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-def _check_matplotlib_available():
-    """Check if matplotlib is available, raise helpful error if not."""
-    if not MATPLOTLIB_AVAILABLE:
-        raise ImportError(
-            "matplotlib is required for RTC debug visualizations. "
-            "Please install it by running:\n"
-            "  uv pip install matplotlib"
-        )
-
-
-@dataclass
-class RTCEvalConfig(HubMixin):
-    """Configuration for RTC evaluation."""
-
-    # Policy configuration
-    policy: PreTrainedConfig | None = None
-
-    # Dataset configuration
-    dataset: DatasetConfig = field(default_factory=DatasetConfig)
-
-    # RTC configuration
-    rtc: RTCConfig = field(
-        default_factory=lambda: RTCConfig(
-            enabled=True,
-            execution_horizon=20,
-            max_guidance_weight=10.0,
-            prefix_attention_schedule=RTCAttentionSchedule.EXP,
-            debug=True,
-            debug_maxlen=1000,
-        )
-    )
-
-    # Device configuration
-    device: str | None = field(
-        default=None,
-        metadata={"help": "Device to run on (cuda, cpu, mps, auto)"},
-    )
-
-    # Output configuration
-    output_dir: str = field(
-        default="rtc_debug_output",
-        metadata={"help": "Directory to save debug visualizations"},
-    )
-
-    # Seed configuration
-    seed: int = field(
-        default=42,
-        metadata={"help": "Random seed for reproducibility"},
-    )
-
-    inference_delay: int = field(
-        default=4,
-        metadata={"help": "Inference delay for RTC"},
-    )
-
-    # Torch compile configuration
-    use_torch_compile: bool = field(
-        default=False,
-        metadata={"help": "Use torch.compile for faster inference (PyTorch 2.0+)"},
-    )
-
-    torch_compile_backend: str = field(
-        default="inductor",
-        metadata={"help": "Backend for torch.compile (inductor, aot_eager, cudagraphs)"},
-    )
-
-    torch_compile_mode: str = field(
-        default="default",
-        metadata={"help": "Compilation mode (default, reduce-overhead, max-autotune)"},
-    )
-
-    torch_compile_disable_cudagraphs: bool = field(
-        default=True,
-        metadata={
-            "help": "Disable CUDA graphs in torch.compile. Required due to in-place tensor "
-            "operations in denoising loop (x_t += dt * v_t) which cause tensor aliasing issues."
-        },
-    )
-
-    def __post_init__(self):
-        # Parse policy path
-        policy_path = parser.get_path_arg("policy")
-        if policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = policy_path
-        else:
-            raise ValueError("Policy path is required (--policy.path)")
-
-        # Auto-detect device if not specified
-        if self.device is None or self.device == "auto":
-            if torch.cuda.is_available():
-                self.device = "cuda"
-            elif torch.backends.mps.is_available():
-                self.device = "mps"
-            else:
-                self.device = "cpu"
-            logging.info(f"Auto-detected device: {self.device}")
-
-    @classmethod
-    def __get_path_fields__(cls) -> list[str]:
-        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
-        return ["policy"]
-
-
-class RTCEvaluator:
-    """Evaluator for RTC on dataset samples."""
-
-    def __init__(self, cfg: RTCEvalConfig):
-        self.cfg = cfg
-        self.device = cfg.device
-
-        # Load dataset with proper delta_timestamps based on policy configuration
-        # Calculate delta_timestamps using the same logic as make_dataset factory
-        logging.info(f"Loading dataset: {cfg.dataset.repo_id}")
-
-        # Get dataset metadata to extract FPS
-        ds_meta = LeRobotDatasetMetadata(cfg.dataset.repo_id)
-
-        # Calculate delta_timestamps from policy's delta_indices
-        delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
-
-        # Create dataset with calculated delta_timestamps
-        self.dataset = LeRobotDataset(
-            cfg.dataset.repo_id,
-            delta_timestamps=delta_timestamps,
-        )
-        logging.info(f"Dataset loaded: {len(self.dataset)} samples, {self.dataset.num_episodes} episodes")
-
-        # Create preprocessor/postprocessor
-        self.preprocessor, self.postprocessor = make_pre_post_processors(
-            policy_cfg=cfg.policy,
-            pretrained_path=cfg.policy.pretrained_path,
-            preprocessor_overrides={
-                "device_processor": {"device": self.device},
-            },
-        )
-
-        logging.info("=" * 80)
-        logging.info("Ready to run evaluation with sequential policy loading:")
-        logging.info("  1. policy_prev_chunk - Generate reference chunk, then destroy")
-        logging.info("  2. policy_no_rtc - Generate without RTC, then destroy")
-        logging.info("  3. policy_rtc - Generate with RTC, then destroy")
-        logging.info("  Note: Only one policy in memory at a time for efficient memory usage")
-        logging.info("=" * 80)
-
-    def _init_policy(self, name: str, rtc_enabled: bool, rtc_debug: bool):
-        """Initialize a single policy instance with specified RTC configuration.
-
-        Args:
-            name: Name identifier for logging purposes
-            rtc_enabled: Whether to enable RTC for this policy
-            rtc_debug: Whether to enable debug tracking for this policy
-
-        Returns:
-            Configured policy instance with optional torch.compile applied
-        """
-        logging.info(f"Initializing {name}...")
-
-        # Load policy from pretrained
-        policy_class = get_policy_class(self.cfg.policy.type)
-
-        config = PreTrainedConfig.from_pretrained(self.cfg.policy.pretrained_path)
-
-        if self.cfg.policy.type == "pi05" or self.cfg.policy.type == "pi0":
-            config.compile_model = self.cfg.use_torch_compile
-
-        policy = policy_class.from_pretrained(self.cfg.policy.pretrained_path, config=config)
-        policy = policy.to(self.device)
-        policy.eval()
-
-        # Configure RTC
-        rtc_config = RTCConfig(
-            enabled=rtc_enabled,
-            execution_horizon=self.cfg.rtc.execution_horizon,
-            max_guidance_weight=self.cfg.rtc.max_guidance_weight,
-            prefix_attention_schedule=self.cfg.rtc.prefix_attention_schedule,
-            debug=rtc_debug,
-            debug_maxlen=self.cfg.rtc.debug_maxlen,
-        )
-        policy.config.rtc_config = rtc_config
-        policy.init_rtc_processor()
-
-        logging.info(f"  RTC enabled: {rtc_enabled}")
-        logging.info(f"  RTC debug: {rtc_debug}")
-        logging.info(f"  Policy config: {config}")
-
-        # Apply torch.compile to predict_action_chunk method if enabled
-        if self.cfg.use_torch_compile:
-            policy = self._apply_torch_compile(policy, name)
-
-        logging.info(f"✓ {name} initialized successfully")
-        return policy
-
-    def _apply_torch_compile(self, policy, policy_name: str):
-        """Apply torch.compile to the policy's predict_action_chunk method.
-
-        Args:
-            policy: Policy instance to compile
-            policy_name: Name for logging purposes
-
-        Returns:
-            Policy with compiled predict_action_chunk method
-        """
-
-        # PI models handle their own compilation
-        if policy.type == "pi05" or policy.type == "pi0":
-            return policy
-
-        try:
-            # Check if torch.compile is available (PyTorch 2.0+)
-            if not hasattr(torch, "compile"):
-                logging.warning(
-                    f"  [{policy_name}] torch.compile is not available. Requires PyTorch 2.0+. "
-                    f"Current version: {torch.__version__}. Skipping compilation."
-                )
-                return policy
-
-            logging.info(f"  [{policy_name}] Applying torch.compile to predict_action_chunk...")
-            logging.info(f"    Backend: {self.cfg.torch_compile_backend}")
-            logging.info(f"    Mode: {self.cfg.torch_compile_mode}")
-            logging.info(f"    Disable CUDA graphs: {self.cfg.torch_compile_disable_cudagraphs}")
-            logging.info("    Note: Debug tracker excluded from compilation via @torch._dynamo.disable")
-
-            # Compile the predict_action_chunk method
-            # - Debug tracker is excluded from compilation via @torch._dynamo.disable
-            # - CUDA graphs disabled to prevent tensor aliasing from in-place ops (x_t += dt * v_t)
-            compile_kwargs = {
-                "backend": self.cfg.torch_compile_backend,
-                "mode": self.cfg.torch_compile_mode,
-            }
-
-            # Disable CUDA graphs if requested (prevents tensor aliasing issues)
-            if self.cfg.torch_compile_disable_cudagraphs:
-                compile_kwargs["options"] = {"triton.cudagraphs": False}
-
-            original_method = policy.predict_action_chunk
-            compiled_method = torch.compile(original_method, **compile_kwargs)
-            policy.predict_action_chunk = compiled_method
-            logging.info(f"  ✓ [{policy_name}] Successfully compiled predict_action_chunk")
-
-        except Exception as e:
-            logging.error(f"  [{policy_name}] Failed to apply torch.compile: {e}")
-            logging.warning(f"  [{policy_name}] Continuing without torch.compile")
-
-        return policy
-
-    def _destroy_policy(self, policy, policy_name: str):
-        """Explicitly destroy a policy and free all associated memory.
-
-        This method performs aggressive cleanup to ensure maximum memory is freed,
-        which is critical for large models (e.g., VLAs with billions of parameters).
-
-        Args:
-            policy: Policy instance to destroy
-            policy_name: Name for logging purposes
-        """
-        logging.info(f"  Destroying {policy_name} and freeing memory...")
-
-        try:
-            # Step 1: Move policy to CPU to free GPU/MPS memory
-            policy.cpu()
-
-            # Step 2: Delete the policy object
-            del policy
-
-            # Step 3: Force garbage collection to reclaim memory immediately
-            gc.collect()
-
-            # Step 4: Clear device-specific caches
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-                torch.cuda.synchronize()  # Ensure all operations complete
-
-            if torch.backends.mps.is_available():
-                torch.mps.empty_cache()
-
-            logging.info(f"  ✓ {policy_name} destroyed and memory freed")
-
-        except Exception as e:
-            logging.warning(f"  Warning: Error during {policy_name} cleanup: {e}")
-
-    def run_evaluation(self):
-        """Run evaluation on two random dataset samples using three separate policies.
-
-        Note: Policies are deinitalized after each step to free memory. Large models
-        (e.g., VLA models with billions of parameters) cannot fit three instances in
-        memory simultaneously. By deleting and garbage collecting after each step,
-        we ensure only one policy is loaded at a time.
-        """
-        # Create output directory
-        os.makedirs(self.cfg.output_dir, exist_ok=True)
-        logging.info(f"Output directory: {self.cfg.output_dir}")
-
-        logging.info("=" * 80)
-        logging.info("Starting RTC evaluation")
-        logging.info(f"Inference delay: {self.cfg.inference_delay}")
-        logging.info("=" * 80)
-
-        # Load two random samples from dataset
-        data_loader = torch.utils.data.DataLoader(self.dataset, batch_size=1, shuffle=True)
-        loader_iter = iter(data_loader)
-        first_sample = next(loader_iter)
-        second_sample = next(loader_iter)
-
-        preprocessed_first_sample = self.preprocessor(first_sample)
-        preprocessed_second_sample = self.preprocessor(second_sample)
-
-        # ============================================================================
-        # Step 1: Generate previous chunk using policy_prev_chunk
-        # ============================================================================
-        # This policy is only used to generate the reference chunk and then freed
-        logging.info("=" * 80)
-        logging.info("Step 1: Generating previous chunk with policy_prev_chunk")
-        logging.info("=" * 80)
-
-        # Initialize policy 1
-        policy_prev_chunk_policy = self._init_policy(
-            name="policy_prev_chunk",
-            rtc_enabled=False,
-            rtc_debug=False,
-        )
-        with torch.no_grad():
-            prev_chunk_left_over = policy_prev_chunk_policy.predict_action_chunk(
-                preprocessed_first_sample,
-            )[:, :25, :].squeeze(0)
-        logging.info(f"  Generated prev_chunk shape: {prev_chunk_left_over.shape}")
-
-        # Destroy policy_prev_chunk to free memory for large models
-        self._destroy_policy(policy_prev_chunk_policy, "policy_prev_chunk")
-
-        # ============================================================================
-        # Step 2: Generate actions WITHOUT RTC using policy_no_rtc
-        # ============================================================================
-        logging.info("=" * 80)
-        logging.info("Step 2: Generating actions WITHOUT RTC with policy_no_rtc")
-        logging.info("=" * 80)
-
-        set_seed(self.cfg.seed)
-
-        # Initialize policy 2
-        policy_no_rtc_policy = self._init_policy(
-            name="policy_no_rtc",
-            rtc_enabled=False,
-            rtc_debug=True,
-        )
-
-        # Sample noise (use same noise for both RTC and non-RTC for fair comparison)
-        noise_size = (1, policy_no_rtc_policy.config.chunk_size, policy_no_rtc_policy.config.max_action_dim)
-        noise = policy_no_rtc_policy.model.sample_noise(noise_size, self.device)
-        noise_clone = noise.clone()
-        policy_no_rtc_policy.rtc_processor.reset_tracker()
-        with torch.no_grad():
-            no_rtc_actions = policy_no_rtc_policy.predict_action_chunk(
-                preprocessed_second_sample,
-                noise=noise,
-            )
-        no_rtc_tracked_steps = policy_no_rtc_policy.rtc_processor.tracker.get_all_steps()
-        logging.info(f"  Tracked {len(no_rtc_tracked_steps)} steps without RTC")
-        logging.info(f"  Generated no_rtc_actions shape: {no_rtc_actions.shape}")
-
-        # Destroy policy_no_rtc to free memory before loading policy_rtc
-        self._destroy_policy(policy_no_rtc_policy, "policy_no_rtc")
-
-        # ============================================================================
-        # Step 3: Generate actions WITH RTC using policy_rtc
-        # ============================================================================
-        logging.info("=" * 80)
-        logging.info("Step 3: Generating actions WITH RTC with policy_rtc")
-        logging.info("=" * 80)
-
-        set_seed(self.cfg.seed)
-
-        # Initialize policy 3
-        policy_rtc_policy = self._init_policy(
-            name="policy_rtc",
-            rtc_enabled=True,
-            rtc_debug=True,
-        )
-        policy_rtc_policy.rtc_processor.reset_tracker()
-        with torch.no_grad():
-            rtc_actions = policy_rtc_policy.predict_action_chunk(
-                preprocessed_second_sample,
-                noise=noise_clone,
-                inference_delay=self.cfg.inference_delay,
-                prev_chunk_left_over=prev_chunk_left_over,
-                execution_horizon=self.cfg.rtc.execution_horizon,
-            )
-        rtc_tracked_steps = policy_rtc_policy.rtc_processor.get_all_debug_steps()
-        logging.info(f"  Tracked {len(rtc_tracked_steps)} steps with RTC")
-        logging.info(f"  Generated rtc_actions shape: {rtc_actions.shape}")
-
-        # Save num_steps before destroying policy (needed for plotting)
-        try:
-            num_steps = policy_rtc_policy.config.num_steps
-        except Exception as e:
-            logging.error(f"  Error getting num_steps: {e}")
-            num_steps = policy_rtc_policy.config.num_inference_steps
-            logging.warning(f"  Using num_inference_steps: {num_steps} instead of num_steps")
-
-        # Destroy policy_rtc after final use
-        self._destroy_policy(policy_rtc_policy, "policy_rtc")
-
-        # Plot and save results
-        logging.info("=" * 80)
-        logging.info("Plotting results...")
-        self.plot_tracked_data(rtc_tracked_steps, no_rtc_tracked_steps, prev_chunk_left_over, num_steps)
-
-        # Plot final actions comparison
-        logging.info("=" * 80)
-        logging.info("Plotting final actions comparison...")
-        self.plot_final_actions_comparison(rtc_actions, no_rtc_actions, prev_chunk_left_over)
-
-        logging.info("=" * 80)
-        logging.info("Evaluation completed successfully")
-
-    def plot_final_actions_comparison(self, rtc_actions, no_rtc_actions, prev_chunk_left_over):
-        """Plot final action predictions comparison on a single chart.
-
-        Args:
-            rtc_actions: Final actions from RTC policy
-            no_rtc_actions: Final actions from non-RTC policy
-            prev_chunk_left_over: Previous chunk used as ground truth
-        """
-        _check_matplotlib_available()
-
-        # Remove batch dimension if present
-        rtc_actions_plot = rtc_actions.squeeze(0).cpu() if len(rtc_actions.shape) == 3 else rtc_actions.cpu()
-        no_rtc_actions_plot = (
-            no_rtc_actions.squeeze(0).cpu() if len(no_rtc_actions.shape) == 3 else no_rtc_actions.cpu()
-        )
-        prev_chunk_plot = prev_chunk_left_over.cpu()
-
-        # Create figure with 6 subplots (one per action dimension)
-        fig, axes = plt.subplots(6, 1, figsize=(16, 12))
-        fig.suptitle("Final Action Predictions Comparison (Raw)", fontsize=16)
-
-        # Plot each action dimension
-        for dim_idx, ax in enumerate(axes):
-            # Plot previous chunk (ground truth) in red
-            RTCDebugVisualizer.plot_waypoints(
-                [ax],
-                prev_chunk_plot[:, dim_idx : dim_idx + 1],
-                start_from=0,
-                color="red",
-                label="Previous Chunk (Ground Truth)",
-                linewidth=2.5,
-                alpha=0.8,
-            )
-
-            # Plot no-RTC actions in blue
-            RTCDebugVisualizer.plot_waypoints(
-                [ax],
-                no_rtc_actions_plot[:, dim_idx : dim_idx + 1],
-                start_from=0,
-                color="blue",
-                label="No RTC",
-                linewidth=2,
-                alpha=0.7,
-            )
-
-            # Plot RTC actions in green
-            RTCDebugVisualizer.plot_waypoints(
-                [ax],
-                rtc_actions_plot[:, dim_idx : dim_idx + 1],
-                start_from=0,
-                color="green",
-                label="RTC",
-                linewidth=2,
-                alpha=0.7,
-            )
-
-            # Add vertical lines for inference delay and execution horizon
-            inference_delay = self.cfg.inference_delay
-            execution_horizon = self.cfg.rtc.execution_horizon
-
-            if inference_delay > 0:
-                ax.axvline(
-                    x=inference_delay - 1,
-                    color="orange",
-                    linestyle="--",
-                    alpha=0.5,
-                    label=f"Inference Delay ({inference_delay})",
-                )
-
-            if execution_horizon > 0:
-                ax.axvline(
-                    x=execution_horizon,
-                    color="purple",
-                    linestyle="--",
-                    alpha=0.5,
-                    label=f"Execution Horizon ({execution_horizon})",
-                )
-
-            ax.set_ylabel(f"Dim {dim_idx}", fontsize=10)
-            ax.grid(True, alpha=0.3)
-
-            # Set x-axis ticks to show all integer values
-            max_len = max(rtc_actions_plot.shape[0], no_rtc_actions_plot.shape[0], prev_chunk_plot.shape[0])
-            ax.set_xticks(range(0, max_len, max(1, max_len // 20)))  # Show ~20 ticks
-            ax.set_xlim(-0.5, max_len - 0.5)
-
-        axes[-1].set_xlabel("Step", fontsize=10)
-
-        # Collect legend handles and labels from first subplot
-        handles, labels = axes[0].get_legend_handles_labels()
-        # Remove duplicates while preserving order
-        seen = set()
-        unique_handles = []
-        unique_labels = []
-        for handle, label in zip(handles, labels, strict=True):
-            if label not in seen:
-                seen.add(label)
-                unique_handles.append(handle)
-                unique_labels.append(label)
-
-        # Add legend outside the plot area (to the right)
-        fig.legend(
-            unique_handles,
-            unique_labels,
-            loc="center right",
-            fontsize=9,
-            bbox_to_anchor=(1.0, 0.5),
-            framealpha=0.9,
-        )
-
-        # Save figure
-        output_path = os.path.join(self.cfg.output_dir, "final_actions_comparison.png")
-        fig.tight_layout(rect=[0, 0, 0.85, 1])  # Leave space for legend on right
-        fig.savefig(output_path, dpi=150, bbox_inches="tight")
-        logging.info(f"Saved final actions comparison to {output_path}")
-        plt.close(fig)
-
-    def plot_tracked_data(self, rtc_tracked_steps, no_rtc_tracked_steps, prev_chunk_left_over, num_steps):
-        _check_matplotlib_available()
-
-        # Create side-by-side figures for denoising visualization
-        fig_xt, axs_xt = self._create_figure("x_t Denoising: No RTC (left) vs RTC (right)")
-        fig_vt, axs_vt = self._create_figure("v_t Denoising: No RTC (left) vs RTC (right)")
-        fig_corr, axs_corr = self._create_figure("Correction: No RTC (left) vs RTC (right)")
-        fig_x1t, axs_x1t = self._create_figure(
-            "x1_t Predicted State & Error: No RTC (left - empty) vs RTC (right)"
-        )
-        self._plot_denoising_steps_from_tracker(
-            rtc_tracked_steps,
-            axs_xt[:, 1],  # Right column for x_t
-            axs_vt[:, 1],  # Right column for v_t
-            axs_corr[:, 1],  # Right column for correction
-            axs_x1t[:, 1],  # Right column for x1_t
-            num_steps,
-            add_labels=True,  # Add labels for RTC (right column)
-        )
-
-        self._plot_denoising_steps_from_tracker(
-            no_rtc_tracked_steps,
-            axs_xt[:, 0],  # Left column for x_t
-            axs_vt[:, 0],  # Left column for v_t
-            axs_corr[:, 0],  # Left column for correction
-            axs_x1t[:, 0],  # Left column for x1_t
-            num_steps,
-            add_labels=False,  # No labels for No RTC (left column)
-        )
-
-        # Plot no-RTC x_t data on right chart as orange dashed line for comparison
-        self._plot_no_rtc_xt_reference(no_rtc_tracked_steps, axs_xt[:, 1], num_steps)
-
-        # Plot ground truth on x_t axes
-        RTCDebugVisualizer.plot_waypoints(
-            axs_xt[:, 1], prev_chunk_left_over, start_from=0, color="red", label="Ground truth"
-        )
-
-        # Plot ground truth on x1_t axes
-        RTCDebugVisualizer.plot_waypoints(
-            axs_x1t[:, 1], prev_chunk_left_over, start_from=0, color="red", label="Ground truth"
-        )
-
-        # Plot ground truth on x_t axes (no labels for left column)
-        RTCDebugVisualizer.plot_waypoints(
-            axs_xt[:, 0], prev_chunk_left_over, start_from=0, color="red", label=None
-        )
-
-        RTCDebugVisualizer.plot_waypoints(
-            axs_x1t[:, 0], prev_chunk_left_over, start_from=0, color="red", label=None
-        )
-
-        # Add legends outside the plot area for each figure
-        self._add_figure_legend(fig_xt, axs_xt)
-        self._add_figure_legend(fig_vt, axs_vt)
-        self._add_figure_legend(fig_corr, axs_corr)
-        self._add_figure_legend(fig_x1t, axs_x1t)
-
-        # Save denoising plots
-        self._save_figure(fig_xt, os.path.join(self.cfg.output_dir, "denoising_xt_comparison.png"))
-        self._save_figure(fig_vt, os.path.join(self.cfg.output_dir, "denoising_vt_comparison.png"))
-        self._save_figure(fig_corr, os.path.join(self.cfg.output_dir, "denoising_correction_comparison.png"))
-        self._save_figure(fig_x1t, os.path.join(self.cfg.output_dir, "denoising_x1t_comparison.png"))
-
-    def _create_figure(self, title):
-        fig, axs = plt.subplots(6, 2, figsize=(24, 12))
-        fig.suptitle(title, fontsize=16)
-
-        for ax in axs[:, 0]:
-            ax.set_title("No RTC (N/A)" if ax == axs[0, 0] else "", fontsize=12)
-        for ax in axs[:, 1]:
-            ax.set_title("RTC" if ax == axs[0, 1] else "", fontsize=12)
-
-        return fig, axs
-
-    def _add_figure_legend(self, fig, axs):
-        """Add a legend outside the plot area on the right side.
-
-        Args:
-            fig: Matplotlib figure to add legend to
-            axs: Array of axes to collect legend handles from
-        """
-        # Collect all handles and labels from the first row of axes (right column)
-        handles, labels = axs[0, 1].get_legend_handles_labels()
-
-        # Remove duplicates while preserving order
-        seen = set()
-        unique_handles = []
-        unique_labels = []
-        for handle, label in zip(handles, labels, strict=True):
-            if label not in seen:
-                seen.add(label)
-                unique_handles.append(handle)
-                unique_labels.append(label)
-
-        # Add legend outside the plot area (to the right, close to charts)
-        if unique_handles:
-            fig.legend(
-                unique_handles,
-                unique_labels,
-                loc="center left",
-                fontsize=8,
-                bbox_to_anchor=(0.87, 0.5),
-                framealpha=0.9,
-                ncol=1,
-            )
-
-    def _save_figure(self, fig, path):
-        fig.tight_layout(rect=[0, 0, 0.85, 1])  # Leave space for legend/colorbar on right
-        fig.savefig(path, dpi=150, bbox_inches="tight")
-        logging.info(f"Saved figure to {path}")
-        plt.close(fig)
-
-    def _plot_denoising_steps_from_tracker(
-        self, tracked_steps, xt_axs, vt_axs, corr_axs, x1t_axs, num_steps, add_labels=True
-    ):
-        """Plot denoising steps from tracker data.
-
-        Args:
-            tracked_steps: List of DebugStep objects containing debug steps
-            xt_axs: Matplotlib axes for x_t plots (array of 6 axes)
-            vt_axs: Matplotlib axes for v_t plots (array of 6 axes)
-            corr_axs: Matplotlib axes for correction plots (array of 6 axes)
-            x1t_axs: Matplotlib axes for x1_t plots (array of 6 axes)
-            num_steps: Total number of denoising steps for colormap
-            add_labels: Whether to add legend labels for the plots
-        """
-
-        logging.info("=" * 80)
-        logging.info(f"Plotting {len(tracked_steps)} steps")
-
-        debug_steps = tracked_steps
-        if not debug_steps:
-            return
-
-        # Define colors for different denoise steps (using a colormap)
-        colors = plt.cm.viridis(np.linspace(0, 1, num_steps))
-
-        for step_idx, debug_step in enumerate(debug_steps):
-            color = colors[step_idx % len(colors)]
-            label = f"Step {step_idx}" if add_labels else None
-
-            # Plot x_t
-            if debug_step.x_t is not None:
-                RTCDebugVisualizer.plot_waypoints(
-                    xt_axs, debug_step.x_t, start_from=0, color=color, label=label
-                )
-
-            # Plot v_t
-            if debug_step.v_t is not None:
-                RTCDebugVisualizer.plot_waypoints(
-                    vt_axs, debug_step.v_t, start_from=0, color=color, label=label
-                )
-
-            # Plot correction on separate axes
-            if debug_step.correction is not None:
-                RTCDebugVisualizer.plot_waypoints(
-                    corr_axs,
-                    debug_step.correction,
-                    start_from=0,
-                    color=color,
-                    label=label,
-                )
-
-            # Plot x1_t (predicted state)
-            if x1t_axs is not None and debug_step.x1_t is not None:
-                x1t_label = f"x1_t Step {step_idx}" if add_labels else None
-                RTCDebugVisualizer.plot_waypoints(
-                    x1t_axs,
-                    debug_step.x1_t,
-                    start_from=0,
-                    color=color,
-                    label=x1t_label,
-                )
-
-            # Plot error in orange dashed
-            if x1t_axs is not None and debug_step.err is not None:
-                error_chunk = (
-                    debug_step.err[0].cpu().numpy()
-                    if len(debug_step.err.shape) == 3
-                    else debug_step.err.cpu().numpy()
-                )
-
-                num_dims = min(error_chunk.shape[-1], 6)
-                error_label = f"error Step {step_idx}" if add_labels else None
-                for j in range(num_dims):
-                    x1t_axs[j].plot(
-                        np.arange(0, error_chunk.shape[0]),
-                        error_chunk[:, j],
-                        color="orange",
-                        linestyle="--",
-                        alpha=0.7,
-                        label=error_label,
-                    )
-
-        # Recalculate axis limits after plotting to ensure proper scaling
-        self._rescale_axes(xt_axs)
-        self._rescale_axes(vt_axs)
-        self._rescale_axes(corr_axs)
-        self._rescale_axes(x1t_axs)
-
-    def _plot_no_rtc_xt_reference(self, no_rtc_tracked_steps, xt_axs, num_steps):
-        """Plot final no-RTC x_t data as orange dashed line on the RTC chart for comparison.
-
-        Args:
-            no_rtc_tracked_steps: List of DebugStep objects containing no-RTC debug steps
-            xt_axs: Matplotlib axes for x_t plots (array of 6 axes, right column)
-            num_steps: Total number of denoising steps for colormap
-        """
-        debug_steps = no_rtc_tracked_steps
-        if not debug_steps:
-            return
-
-        # Plot only the final x_t step as orange dashed line
-        final_step = debug_steps[-1]
-        logging.info("Plotting final no-RTC x_t step as orange dashed reference")
-
-        if final_step.x_t is not None:
-            x_t_chunk = (
-                final_step.x_t[0].cpu().numpy()
-                if len(final_step.x_t.shape) == 3
-                else final_step.x_t.cpu().numpy()
-            )
-
-            num_dims = min(x_t_chunk.shape[-1], 6)
-            for j in range(num_dims):
-                xt_axs[j].plot(
-                    np.arange(0, x_t_chunk.shape[0]),
-                    x_t_chunk[:, j],
-                    color="orange",
-                    linestyle="--",
-                    alpha=0.7,
-                    linewidth=2,
-                    label="No RTC (final)" if j == 0 else "",
-                )
-
-    def _rescale_axes(self, axes):
-        """Rescale axes to show all data with proper margins.
-
-        Args:
-            axes: Array of matplotlib axes to rescale
-        """
-        for ax in axes:
-            ax.relim()
-            ax.autoscale_view()
-
-            # Add 10% margin to y-axis for better visualization
-            ylim = ax.get_ylim()
-            y_range = ylim[1] - ylim[0]
-            if y_range > 0:  # Avoid division by zero
-                margin = y_range * 0.1
-                ax.set_ylim(ylim[0] - margin, ylim[1] + margin)
-
-            # Set x-axis ticks to show all integer values
-            xlim = ax.get_xlim()
-            max_len = int(xlim[1]) + 1
-            if max_len > 0:
-                ax.set_xticks(range(0, max_len, max(1, max_len // 20)))  # Show ~20 ticks
-                ax.set_xlim(-0.5, max_len - 0.5)
-
-
-@parser.wrap()
-def main(cfg: RTCEvalConfig):
-    """Main entry point for RTC evaluation."""
-    # Set random seed for reproducibility
-    set_seed(cfg.seed)
-
-    init_logging()
-
-    logging.info("=" * 80)
-    logging.info("RTC Dataset Evaluation")
-    logging.info(f"Config: {cfg}")
-    logging.info("=" * 80)
-
-    evaluator = RTCEvaluator(cfg)
-    evaluator.run_evaluation()
-
-
-if __name__ == "__main__":
-    main()

examples/rtc/eval_with_real_robot.py

@@ -1,549 +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.
-
-"""
-Demo script showing how to use Real-Time Chunking (RTC) with action chunking policies on real robots.
-
-This script demonstrates:
-1. Creating a robot and policy (SmolVLA, Pi0, etc.) with RTC
-2. Consuming actions from the policy while the robot executes
-3. Periodically requesting new action chunks in the background using threads
-4. Managing action buffers and timing for real-time operation
-
-For simulation environments, see eval_with_simulation.py
-
-Usage:
-    # Run RTC with Real robot with RTC
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --policy.device=mps \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with Real robot without RTC
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=helper2424/smolvla_check_rtc_last3 \
-        --policy.device=mps \
-        --rtc.enabled=false \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with Real robot with pi0.5 policy
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=helper2424/pi05_check_rtc \
-        --policy.device=mps \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-"""
-
-import logging
-import math
-import sys
-import time
-import traceback
-from dataclasses import dataclass, field
-from threading import Event, Lock, Thread
-
-import torch
-from torch import Tensor
-
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.configs import parser
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.action_queue import ActionQueue
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.latency_tracker import LatencyTracker
-from lerobot.processor.factory import (
-    make_default_robot_action_processor,
-    make_default_robot_observation_processor,
-)
-from lerobot.rl.process import ProcessSignalHandler
-from lerobot.robots import (  # noqa: F401
-    Robot,
-    RobotConfig,
-    koch_follower,
-    so100_follower,
-    so101_follower,
-)
-from lerobot.robots.utils import make_robot_from_config
-from lerobot.utils.constants import OBS_IMAGES
-from lerobot.utils.hub import HubMixin
-from lerobot.utils.utils import init_logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class RobotWrapper:
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.lock = Lock()
-
-    def get_observation(self) -> dict[str, Tensor]:
-        with self.lock:
-            return self.robot.get_observation()
-
-    def send_action(self, action: Tensor):
-        with self.lock:
-            self.robot.send_action(action)
-
-    def observation_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.observation_features
-
-    def action_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.action_features
-
-
-@dataclass
-class RTCDemoConfig(HubMixin):
-    """Configuration for RTC demo with action chunking policies and real robots."""
-
-    # Policy configuration
-    policy: PreTrainedConfig | None = None
-
-    # Robot configuration
-    robot: RobotConfig | None = None
-
-    # RTC configuration
-    rtc: RTCConfig = field(
-        default_factory=lambda: RTCConfig(
-            execution_horizon=10,
-            max_guidance_weight=1.0,
-            prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        )
-    )
-
-    # Demo parameters
-    duration: float = 30.0  # Duration to run the demo (seconds)
-    fps: float = 10.0  # Action execution frequency (Hz)
-
-    # Compute device
-    device: str | None = None  # Device to run on (cuda, cpu, auto)
-
-    # Get new actions horizon. The amount of executed steps after which will be requested new actions.
-    # It should be higher than inference delay + execution horizon.
-    action_queue_size_to_get_new_actions: int = 30
-
-    # Task to execute
-    task: str = field(default="", metadata={"help": "Task to execute"})
-
-    # Torch compile configuration
-    use_torch_compile: bool = field(
-        default=False,
-        metadata={"help": "Use torch.compile for faster inference (PyTorch 2.0+)"},
-    )
-
-    torch_compile_backend: str = field(
-        default="inductor",
-        metadata={"help": "Backend for torch.compile (inductor, aot_eager, cudagraphs)"},
-    )
-
-    torch_compile_mode: str = field(
-        default="default",
-        metadata={"help": "Compilation mode (default, reduce-overhead, max-autotune)"},
-    )
-
-    torch_compile_disable_cudagraphs: bool = field(
-        default=True,
-        metadata={
-            "help": "Disable CUDA graphs in torch.compile. Required due to in-place tensor "
-            "operations in denoising loop (x_t += dt * v_t) which cause tensor aliasing issues."
-        },
-    )
-
-    def __post_init__(self):
-        # HACK: We parse again the cli args here to get the pretrained path if there was one.
-        policy_path = parser.get_path_arg("policy")
-        if policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = policy_path
-        else:
-            raise ValueError("Policy path is required")
-
-        # Validate that robot configuration is provided
-        if self.robot is None:
-            raise ValueError("Robot configuration must be provided")
-
-    @classmethod
-    def __get_path_fields__(cls) -> list[str]:
-        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
-        return ["policy"]
-
-
-def is_image_key(k: str) -> bool:
-    return k.startswith(OBS_IMAGES)
-
-
-def get_actions(
-    policy,
-    robot: RobotWrapper,
-    robot_observation_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to request action chunks from the policy.
-
-    Args:
-        policy: The policy instance (SmolVLA, Pi0, etc.)
-        robot: The robot instance for getting observations
-        robot_observation_processor: Processor for raw robot observations
-        action_queue: Queue to put new action chunks
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[GET_ACTIONS] Starting get actions thread")
-
-        latency_tracker = LatencyTracker()  # Track latency of action chunks
-        fps = cfg.fps
-        time_per_chunk = 1.0 / fps
-
-        dataset_features = hw_to_dataset_features(robot.observation_features(), "observation")
-        policy_device = policy.config.device
-
-        # Load preprocessor and postprocessor from pretrained files
-        # The stats are embedded in the processor .safetensors files
-        logger.info(f"[GET_ACTIONS] Loading preprocessor/postprocessor from {cfg.policy.pretrained_path}")
-
-        preprocessor, postprocessor = make_pre_post_processors(
-            policy_cfg=cfg.policy,
-            pretrained_path=cfg.policy.pretrained_path,
-            dataset_stats=None,  # Will load from pretrained processor files
-            preprocessor_overrides={
-                "device_processor": {"device": cfg.policy.device},
-            },
-        )
-
-        logger.info("[GET_ACTIONS] Preprocessor/postprocessor loaded successfully with embedded stats")
-
-        get_actions_threshold = cfg.action_queue_size_to_get_new_actions
-
-        if not cfg.rtc.enabled:
-            get_actions_threshold = 0
-
-        while not shutdown_event.is_set():
-            if action_queue.qsize() <= get_actions_threshold:
-                current_time = time.perf_counter()
-                action_index_before_inference = action_queue.get_action_index()
-                prev_actions = action_queue.get_left_over()
-
-                inference_latency = latency_tracker.max()
-                inference_delay = math.ceil(inference_latency / time_per_chunk)
-
-                obs = robot.get_observation()
-
-                # Apply robot observation processor
-                obs_processed = robot_observation_processor(obs)
-
-                obs_with_policy_features = build_dataset_frame(
-                    dataset_features, obs_processed, prefix="observation"
-                )
-
-                for name in obs_with_policy_features:
-                    obs_with_policy_features[name] = torch.from_numpy(obs_with_policy_features[name])
-                    if "image" in name:
-                        obs_with_policy_features[name] = (
-                            obs_with_policy_features[name].type(torch.float32) / 255
-                        )
-                        obs_with_policy_features[name] = (
-                            obs_with_policy_features[name].permute(2, 0, 1).contiguous()
-                        )
-                    obs_with_policy_features[name] = obs_with_policy_features[name].unsqueeze(0)
-                    obs_with_policy_features[name] = obs_with_policy_features[name].to(policy_device)
-
-                obs_with_policy_features["task"] = [cfg.task]  # Task should be a list, not a string!
-                obs_with_policy_features["robot_type"] = (
-                    robot.robot.name if hasattr(robot.robot, "name") else ""
-                )
-
-                preproceseded_obs = preprocessor(obs_with_policy_features)
-
-                # Generate actions WITH RTC
-                actions = policy.predict_action_chunk(
-                    preproceseded_obs,
-                    inference_delay=inference_delay,
-                    prev_chunk_left_over=prev_actions,
-                )
-
-                # Store original actions (before postprocessing) for RTC
-                original_actions = actions.squeeze(0).clone()
-
-                postprocessed_actions = postprocessor(actions)
-
-                postprocessed_actions = postprocessed_actions.squeeze(0)
-
-                new_latency = time.perf_counter() - current_time
-                new_delay = math.ceil(new_latency / time_per_chunk)
-                latency_tracker.add(new_latency)
-
-                if cfg.action_queue_size_to_get_new_actions < cfg.rtc.execution_horizon + new_delay:
-                    logger.warning(
-                        "[GET_ACTIONS] cfg.action_queue_size_to_get_new_actions Too small, It should be higher than inference delay + execution horizon."
-                    )
-
-                action_queue.merge(
-                    original_actions, postprocessed_actions, new_delay, action_index_before_inference
-                )
-            else:
-                # Small sleep to prevent busy waiting
-                time.sleep(0.1)
-
-        logger.info("[GET_ACTIONS] get actions thread shutting down")
-    except Exception as e:
-        logger.error(f"[GET_ACTIONS] Fatal exception in get_actions thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def actor_control(
-    robot: RobotWrapper,
-    robot_action_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to execute actions on the robot.
-
-    Args:
-        robot: The robot instance
-        action_queue: Queue to get actions from
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[ACTOR] Starting actor thread")
-
-        action_count = 0
-        action_interval = 1.0 / cfg.fps
-
-        while not shutdown_event.is_set():
-            start_time = time.perf_counter()
-
-            # Try to get an action from the queue with timeout
-            action = action_queue.get()
-
-            if action is not None:
-                action = action.cpu()
-                action_dict = {key: action[i].item() for i, key in enumerate(robot.action_features())}
-                action_processed = robot_action_processor((action_dict, None))
-                robot.send_action(action_processed)
-
-                action_count += 1
-
-            dt_s = time.perf_counter() - start_time
-            time.sleep(max(0, (action_interval - dt_s) - 0.001))
-
-        logger.info(f"[ACTOR] Actor thread shutting down. Total actions executed: {action_count}")
-    except Exception as e:
-        logger.error(f"[ACTOR] Fatal exception in actor_control thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def _apply_torch_compile(policy, cfg: RTCDemoConfig):
-    """Apply torch.compile to the policy's predict_action_chunk method.
-
-    Args:
-        policy: Policy instance to compile
-        cfg: Configuration containing torch compile settings
-
-    Returns:
-        Policy with compiled predict_action_chunk method
-    """
-
-    # PI models handle their own compilation
-    if policy.type == "pi05" or policy.type == "pi0":
-        return policy
-
-    try:
-        # Check if torch.compile is available (PyTorch 2.0+)
-        if not hasattr(torch, "compile"):
-            logger.warning(
-                f"torch.compile is not available. Requires PyTorch 2.0+. "
-                f"Current version: {torch.__version__}. Skipping compilation."
-            )
-            return policy
-
-        logger.info("Applying torch.compile to predict_action_chunk...")
-        logger.info(f"  Backend: {cfg.torch_compile_backend}")
-        logger.info(f"  Mode: {cfg.torch_compile_mode}")
-        logger.info(f"  Disable CUDA graphs: {cfg.torch_compile_disable_cudagraphs}")
-
-        # Compile the predict_action_chunk method
-        # - CUDA graphs disabled to prevent tensor aliasing from in-place ops (x_t += dt * v_t)
-        compile_kwargs = {
-            "backend": cfg.torch_compile_backend,
-            "mode": cfg.torch_compile_mode,
-        }
-
-        # Disable CUDA graphs if requested (prevents tensor aliasing issues)
-        if cfg.torch_compile_disable_cudagraphs:
-            compile_kwargs["options"] = {"triton.cudagraphs": False}
-
-        original_method = policy.predict_action_chunk
-        compiled_method = torch.compile(original_method, **compile_kwargs)
-        policy.predict_action_chunk = compiled_method
-        logger.info("✓ Successfully compiled predict_action_chunk")
-
-    except Exception as e:
-        logger.error(f"Failed to apply torch.compile: {e}")
-        logger.warning("Continuing without torch.compile")
-
-    return policy
-
-
-@parser.wrap()
-def demo_cli(cfg: RTCDemoConfig):
-    """Main entry point for RTC demo with draccus configuration."""
-
-    # Initialize logging
-    init_logging()
-
-    logger.info(f"Using device: {cfg.device}")
-
-    # Setup signal handler for graceful shutdown
-    signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
-    shutdown_event = signal_handler.shutdown_event
-
-    policy = None
-    robot = None
-    get_actions_thread = None
-    actor_thread = None
-
-    policy_class = get_policy_class(cfg.policy.type)
-
-    # Load config and set compile_model for pi0/pi05 models
-    config = PreTrainedConfig.from_pretrained(cfg.policy.pretrained_path)
-
-    if cfg.policy.type == "pi05" or cfg.policy.type == "pi0":
-        config.compile_model = cfg.use_torch_compile
-
-    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
-
-    # Turn on RTC
-    policy.config.rtc_config = cfg.rtc
-
-    # Init RTC processort, as by default if RTC disabled in the config
-    # The processor won't be created
-    policy.init_rtc_processor()
-
-    assert policy.name in ["smolvla", "pi05", "pi0"], "Only smolvla, pi05, and pi0 are supported for RTC"
-
-    policy = policy.to(cfg.device)
-    policy.eval()
-
-    # Apply torch.compile to predict_action_chunk method if enabled
-    if cfg.use_torch_compile:
-        policy = _apply_torch_compile(policy, cfg)
-
-    # Create robot
-    logger.info(f"Initializing robot: {cfg.robot.type}")
-    robot = make_robot_from_config(cfg.robot)
-    robot.connect()
-    robot_wrapper = RobotWrapper(robot)
-
-    # Create robot observation processor
-    robot_observation_processor = make_default_robot_observation_processor()
-    robot_action_processor = make_default_robot_action_processor()
-
-    # Create action queue for communication between threads
-    action_queue = ActionQueue(cfg.rtc)
-
-    # Start chunk requester thread
-    get_actions_thread = Thread(
-        target=get_actions,
-        args=(policy, robot_wrapper, robot_observation_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="GetActions",
-    )
-    get_actions_thread.start()
-    logger.info("Started get actions thread")
-
-    # Start action executor thread
-    actor_thread = Thread(
-        target=actor_control,
-        args=(robot_wrapper, robot_action_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="Actor",
-    )
-    actor_thread.start()
-    logger.info("Started actor thread")
-
-    logger.info("Started stop by duration thread")
-
-    # Main thread monitors for duration or shutdown
-    logger.info(f"Running demo for {cfg.duration} seconds...")
-    start_time = time.time()
-
-    while not shutdown_event.is_set() and (time.time() - start_time) < cfg.duration:
-        time.sleep(10)
-
-        # Log queue status periodically
-        if int(time.time() - start_time) % 5 == 0:
-            logger.info(f"[MAIN] Action queue size: {action_queue.qsize()}")
-
-        if time.time() - start_time > cfg.duration:
-            break
-
-    logger.info("Demo duration reached or shutdown requested")
-
-    # Signal shutdown
-    shutdown_event.set()
-
-    # Wait for threads to finish
-    if get_actions_thread and get_actions_thread.is_alive():
-        logger.info("Waiting for chunk requester thread to finish...")
-        get_actions_thread.join()
-
-    if actor_thread and actor_thread.is_alive():
-        logger.info("Waiting for action executor thread to finish...")
-        actor_thread.join()
-
-    # Cleanup robot
-    if robot:
-        robot.disconnect()
-        logger.info("Robot disconnected")
-
-    logger.info("Cleanup completed")
-
-
-if __name__ == "__main__":
-    demo_cli()
-    logging.info("RTC demo finished")

pyproject.toml

@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
 
 [project]
 name = "lerobot"
-version = "0.4.2"
+version = "0.4.1"
 description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
 readme = "README.md"
 license = { text = "Apache-2.0" }
@@ -102,8 +102,10 @@ grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"] # TODO: Bumb dependency (com
 # Motors
 feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0"]
 dynamixel = ["dynamixel-sdk>=3.7.31,<3.9.0"]
+damiao = ["python-can>=4.2.0,<5.0.0"]
 
 # Robots
+openarms = ["lerobot[damiao]"]
 gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0,<0.15.0"]
 hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"]
 lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1,<28.0.0"]
@@ -142,12 +144,13 @@ video_benchmark = ["scikit-image>=0.23.2,<0.26.0", "pandas>=2.2.2,<2.4.0"]
 # Simulation
 aloha = ["gym-aloha>=0.1.2,<0.2.0"]
 pusht = ["gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
-libero = ["lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0"]
+libero = ["lerobot[transformers-dep]", "libero @ git+https://github.com/huggingface/lerobot-libero.git@main#egg=libero"]
 metaworld = ["metaworld==3.0.0"]
 
 # All
 all = [
     "lerobot[dynamixel]",
+    "lerobot[openarms]",
     "lerobot[gamepad]",
     "lerobot[hopejr]",
     "lerobot[lekiwi]",

scripts/visualize_sarm_predictions.py

@@ -1,761 +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.
-
-"""
-Inference script for SARM (Stage-Aware Reward Model).
-
-This script loads a trained SARM model and runs inference on a dataset episode,
-generating visualizations of the predicted task stages and progress over time.
-
-Example usage:
-    python scripts/visualize_sarm_predictions.py \
-        --model-id username/sarm-model \
-        --dataset-repo lerobot/aloha_sim_insertion_human \
-        --episode-index 0 \
-        --output-dir outputs/sarm_viz \
-        --task-description "insert the peg into the socket"
-"""
-
-import argparse
-import json
-import logging
-from pathlib import Path
-from typing import Optional
-
-import matplotlib.pyplot as plt
-import matplotlib.gridspec as gridspec
-import matplotlib.patches as mpatches
-import numpy as np
-import pandas as pd
-import torch
-from tqdm import tqdm
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.policies.sarm.modeling_sarm import SARMRewardModel
-from lerobot.policies.sarm.sarm_utils import (
-    pad_state_to_max_dim,
-    compute_tau,
-    compute_cumulative_progress_batch,
-)
-from lerobot.datasets.utils import load_stats
-
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Run SARM inference and visualize predictions")
-    
-    # Model arguments
-    parser.add_argument(
-        "--model-id",
-        type=str,
-        required=True,
-        help="HuggingFace model ID or local path to trained SARM model"
-    )
-    
-    # Dataset arguments
-    parser.add_argument(
-        "--dataset-repo",
-        type=str,
-        required=True,
-        help="HuggingFace dataset repository ID (e.g., lerobot/aloha_sim_insertion_human)"
-    )
-    parser.add_argument(
-        "--episode-index",
-        type=int,
-        default=0,
-        help="Index of the episode to visualize (default: 0)"
-    )
-    parser.add_argument(
-        "--task-description",
-        type=str,
-        default="perform the task",
-        help="Task description for the reward model (default: 'perform the task')"
-    )
-    
-    # Output arguments
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default="outputs/sarm_inference",
-        help="Directory to save visualization outputs (default: outputs/sarm_inference)"
-    )
-    parser.add_argument(
-        "--image-key",
-        type=str,
-        default=None,
-        help="Key for images in dataset (e.g., observation.images.image). If not specified, uses model config's image_key"
-    )
-    parser.add_argument(
-        "--state-key",
-        type=str,
-        default=None,
-        help="Key for joint states in dataset. If None, auto-detects from dataset"
-    )
-    
-    # Visualization options
-    parser.add_argument(
-        "--show-frames",
-        action="store_true",
-        help="Include sample frames in the visualization"
-    )
-    parser.add_argument(
-        "--num-sample-frames",
-        type=int,
-        default=8,
-        help="Number of sample frames to show (default: 8)"
-    )
-    parser.add_argument(
-        "--figsize",
-        type=int,
-        nargs=2,
-        default=[14, 8],
-        help="Figure size as width height (default: 14 8)"
-    )
-    
-    # Device
-    parser.add_argument(
-        "--device",
-        type=str,
-        default=None,
-        help="Device to run inference on (cuda/cpu, default: auto-detect)"
-    )
-    
-    return parser.parse_args()
-
-
-def load_episode_data(
-    dataset: LeRobotDataset,
-    episode_index: int,
-    image_key: str,
-    state_key: str | None = None
-) -> tuple[np.ndarray, np.ndarray, int, int, str]:
-    """
-    Load all frames and states from a specific episode.
-    
-    Args:
-        dataset: LeRobotDataset instance
-        episode_index: Index of the episode to load
-        image_key: Key for accessing images in the dataset
-        state_key: Key for accessing joint states (auto-detected if None)
-        
-    Returns:
-        Tuple of (frames, states, start_index, end_index, task_description)
-    """
-    # Get episode boundaries
-    episode_data = dataset.meta.episodes
-    start_idx = episode_data["dataset_from_index"][episode_index]
-    end_idx = episode_data["dataset_to_index"][episode_index]
-    
-    logger.info(f"Loading episode {episode_index}: frames {start_idx} to {end_idx} ({end_idx - start_idx} frames)")
-    
-    # Auto-detect state key if not provided
-    if state_key is None:
-        first_item = dataset[start_idx]
-        state_keys = [k for k in first_item.keys() if 'state' in k.lower() or 'qpos' in k.lower()]
-        if state_keys:
-            state_key = state_keys[0]
-            logger.info(f"Auto-detected state key: {state_key}")
-    
-    # Get task description from the dataset if available
-    task_description = None
-    first_item = dataset[start_idx]
-    if "task" in first_item:
-        task_description = first_item["task"]
-        logger.info(f"✓ Extracted task from episode {episode_index}: '{task_description}'")
-    
-    # Load all frames and states from the episode
-    frames = []
-    states = []
-    for idx in tqdm(range(start_idx, end_idx), desc="Loading frames"):
-        item = dataset[idx]
-        
-        # Get image
-        img = item[image_key]
-        
-        # Convert to numpy if needed
-        if isinstance(img, torch.Tensor):
-            img = img.cpu().numpy()
-        
-        # Handle different image formats (C, H, W) or (H, W, C)
-        if img.shape[0] in [1, 3]:  # Channel first
-            img = np.transpose(img, (1, 2, 0))
-        
-        # Convert to uint8 if needed
-        if img.dtype != np.uint8:
-            if img.max() <= 1.0:
-                img = (img * 255).astype(np.uint8)
-            else:
-                img = img.astype(np.uint8)
-        
-        frames.append(img)
-        
-        # Get state if available
-        if state_key and state_key in item:
-            state = item[state_key]
-            if isinstance(state, torch.Tensor):
-                state = state.cpu().numpy()
-            states.append(state)
-    
-    frames = np.array(frames)
-    states = np.array(states) if states else None
-    logger.info(f"Loaded {len(frames)} frames with shape {frames[0].shape}")
-    if states is not None:
-        logger.info(f"Loaded states with shape {states.shape}")
-    
-    return frames, states, start_idx, end_idx, task_description
-
-
-@torch.no_grad()
-def run_inference(
-    model: SARMRewardModel,
-    frames: np.ndarray,
-    states: Optional[np.ndarray],
-    task_description: str,
-    dataset_stats: dict | None = None,
-    state_key: str = "observation.state",
-    batch_size: int = 32
-) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Run SARM inference on video frames and joint states.
-    
-    (per SARM paper Section A.4):
-    - Frame 0: Initial frame of the episode (frame 0)
-    - Frames 1-8: 8 consecutive frames with frame_gap spacing ending at current frame t
-    Pattern: [frame_0, t-(7*gap), t-(6*gap), ..., t-gap, t]
-    
-    Args:
-        model: SARM model
-        frames: Video frames (num_frames, H, W, C) - all frames from ONE episode
-        states: Joint states (num_frames, state_dim)
-        task_description: Task description text
-        dataset_stats: Dataset statistics for state normalization (same as training)
-        state_key: Key for state in dataset_stats
-        batch_size: Batch size for processing slices
-        
-    Returns:
-        Tuple of (progress_predictions, stage_predictions)
-            - progress_predictions: (num_frames,)
-            - stage_predictions: (num_frames, num_stages)
-    """
-    logger.info("Encoding video frames with CLIP...")
-    video_embeddings = model.encode_images(frames)
-    
-    logger.info("Encoding task description with CLIP...")
-    text_embedding = model.encode_text(task_description)
-    
-    # Get config values
-    num_frames_model = model.config.num_frames  # 9
-    frame_gap = model.config.frame_gap  # 30
-    
-    logger.info("Creating video slices (SARM paper: initial frame + 8 consecutive)...")
-    
-    # Convert to tensors
-    video_embeddings = torch.tensor(video_embeddings, dtype=torch.float32)
-    text_embedding = torch.tensor(text_embedding, dtype=torch.float32)
-    if states is not None:
-        state_embeddings = torch.tensor(states, dtype=torch.float32)
-        
-        # Normalize states using dataset stats (same as training processor)
-        if dataset_stats is not None and state_key in dataset_stats:
-            mean = torch.tensor(dataset_stats[state_key]["mean"], dtype=torch.float32)
-            std = torch.tensor(dataset_stats[state_key]["std"], dtype=torch.float32)
-            state_embeddings = (state_embeddings - mean) / (std + 1e-8)
-            logger.info(f"✓ Applied MEAN_STD normalization to states using {state_key}")
-        else:
-            logger.warning("⚠ No dataset_stats provided - states not normalized (may differ from training)")
-    else:
-        state_embeddings = None
-    
-    video_slices = []
-    state_slices = []
-    
-    for current_frame in tqdm(range(len(video_embeddings)), desc="Creating slices"):
-        # Compute frame indices using symmetric bidirectional pattern:
-        # [initial (0), t-4*gap, t-3*gap, t-2*gap, t-gap, t, t+gap, t+2*gap, t+3*gap]
-        # Boundary handling: clamp to [0, last_valid]
-        deltas = model.config.observation_delta_indices
-        last_valid = len(video_embeddings) - 1
-        
-        frame_indices = []
-        for delta in deltas:
-            idx = current_frame + delta
-            idx = max(0, min(idx, last_valid))  # Clamp to valid range
-            frame_indices.append(idx)
-
-        video_slice = video_embeddings[frame_indices]
-        video_slices.append(video_slice)
-        
-        if state_embeddings is not None:
-            state_slice = state_embeddings[frame_indices]
-            state_slices.append(state_slice)
-    
-    video_slices = torch.stack(video_slices)  # (num_frames, num_frames_model, 512)
-    if state_embeddings is not None:
-        state_slices = torch.stack(state_slices)  # (num_frames, num_frames_model, state_dim)
-        # Pad states to max_state_dim (same as training processor)
-        state_slices = pad_state_to_max_dim(state_slices, model.config.max_state_dim)
-    else:
-        state_slices = None
-    
-    logger.info("Running SARM inference on all slices...")
-    # Process in batches
-    all_progress = []
-    all_stages = []
-    
-    for i in tqdm(range(0, len(video_slices), batch_size), desc="Inference"):
-        batch_video = video_slices[i:i + batch_size].to(model.device)
-        batch_states = state_slices[i:i + batch_size].to(model.device) if state_slices is not None else None
-        batch_size_actual = batch_video.shape[0]
-        
-        # Replicate text embedding for batch
-        batch_text = text_embedding.unsqueeze(0).repeat(batch_size_actual, 1).to(model.device)
-        
-        # Get predictions
-        stage_logits, stage_probs, progress_preds = model.sarm_transformer(
-            batch_video, batch_text, batch_states
-        )
-        
-        # Extract predictions at the "current frame" position
-        # With symmetric pattern [initial, t-4g, t-3g, t-2g, t-g, t, t+g, t+2g, t+3g],
-        # the current frame is at position 5 (0-indexed)
-        current_frame_idx = 5
-        batch_progress = progress_preds[:, current_frame_idx, 0].cpu().numpy()
-        batch_stages = stage_probs[:, current_frame_idx, :].cpu().numpy()
-        
-        all_progress.extend(batch_progress)
-        all_stages.extend(batch_stages)
-    
-    return np.array(all_progress), np.array(all_stages)
-
-
-def compute_ground_truth_progress(
-    dataset: LeRobotDataset,
-    episode_index: int,
-    temporal_proportions: dict[str, float],
-    subtask_names_ordered: list[str],
-) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]:
-    """
-    Compute ground truth progress and stage labels for an episode using annotations.
-    
-    Uses SARM Paper Formula (2):
-        y_t = P_{k-1} + ᾱ_k × τ_t
-    
-    where:
-        - τ_t = (t - s_k) / (e_k - s_k) is within-subtask progress
-        - P_{k-1} is cumulative prior (sum of previous subtask proportions)
-        - ᾱ_k is the temporal proportion for subtask k
-    
-    Args:
-        dataset: LeRobotDataset instance
-        episode_index: Index of the episode
-        temporal_proportions: Dict mapping subtask name to proportion
-        subtask_names_ordered: Ordered list of subtask names (for consistent stage indexing)
-        
-    Returns:
-        Tuple of (ground_truth_progress, ground_truth_stages) arrays, or (None, None) if no annotations
-    """
-    # Load episode metadata
-    episodes_df = dataset.meta.episodes.to_pandas()
-    
-    # Check if annotations exist
-    if "subtask_names" not in episodes_df.columns:
-        logger.warning("No subtask_names column found in episodes metadata")
-        return None, None
-    
-    ep_subtask_names = episodes_df.loc[episode_index, "subtask_names"]
-    if ep_subtask_names is None or (isinstance(ep_subtask_names, float) and pd.isna(ep_subtask_names)):
-        logger.warning(f"No annotations found for episode {episode_index}")
-        return None, None
-    
-    subtask_start_frames = episodes_df.loc[episode_index, "subtask_start_frames"]
-    subtask_end_frames = episodes_df.loc[episode_index, "subtask_end_frames"]
-    
-    # Get episode boundaries
-    ep_start = dataset.meta.episodes["dataset_from_index"][episode_index]
-    ep_end = dataset.meta.episodes["dataset_to_index"][episode_index]
-    num_frames = ep_end - ep_start
-    
-    # Get temporal proportions as ordered list
-    temporal_proportions_list = [
-        temporal_proportions.get(name, 0.0) for name in subtask_names_ordered
-    ]
-    
-    logger.info(f"Computing ground truth for {num_frames} frames using {len(ep_subtask_names)} annotated subtasks")
-    logger.info(f"Subtask names in episode: {ep_subtask_names}")
-    logger.info(f"Subtask start frames: {subtask_start_frames}")
-    logger.info(f"Subtask end frames: {subtask_end_frames}")
-    logger.info(f"Temporal proportions (ordered): {dict(zip(subtask_names_ordered, temporal_proportions_list))}")
-    
-    # Compute ground truth for each frame
-    gt_progress = np.zeros(num_frames)
-    gt_stages = np.zeros(num_frames, dtype=np.int32)
-    
-    for frame_rel in range(num_frames):
-        # Find which subtask this frame belongs to
-        found = False
-        for j, (name, start_frame, end_frame) in enumerate(zip(ep_subtask_names, subtask_start_frames, subtask_end_frames)):
-            if frame_rel >= start_frame and frame_rel <= end_frame:
-                # Found the subtask - get its global index
-                stage_idx = subtask_names_ordered.index(name) if name in subtask_names_ordered else 0
-                
-                # Compute τ_t using utility function
-                tau = compute_tau(frame_rel, start_frame, end_frame)
-                
-                # Compute cumulative progress using utility function
-                progress = compute_cumulative_progress_batch(tau, stage_idx, temporal_proportions_list)
-                
-                gt_progress[frame_rel] = progress
-                gt_stages[frame_rel] = stage_idx
-                found = True
-                break
-        
-        if not found:
-            # Handle frames outside annotated subtasks
-            if frame_rel < subtask_start_frames[0]:
-                gt_progress[frame_rel] = 0.0
-                gt_stages[frame_rel] = 0
-            elif frame_rel > subtask_end_frames[-1]:
-                gt_progress[frame_rel] = 1.0
-                gt_stages[frame_rel] = len(subtask_names_ordered) - 1
-            else:
-                # Between subtasks - find previous subtask
-                for j in range(len(ep_subtask_names) - 1):
-                    if frame_rel > subtask_end_frames[j] and frame_rel < subtask_start_frames[j + 1]:
-                        name = ep_subtask_names[j]
-                        stage_idx = subtask_names_ordered.index(name) if name in subtask_names_ordered else j
-                        progress = compute_cumulative_progress_batch(1.0, stage_idx, temporal_proportions_list)
-                        gt_progress[frame_rel] = progress
-                        gt_stages[frame_rel] = stage_idx
-                        break
-    
-    logger.info(f"✓ Ground truth computed: final={gt_progress[-1]:.3f}, max={gt_progress.max():.3f}")
-    return gt_progress, gt_stages
-
-
-def visualize_predictions(
-    frames: np.ndarray,
-    progress_predictions: np.ndarray,
-    stage_predictions: np.ndarray,
-    task_description: str,
-    output_path: Path,
-    num_sample_frames: int = 8,
-    figsize: tuple = (14, 8),
-    subtask_names: list[str] | None = None,
-    temporal_proportions: dict[str, float] | None = None,
-    ground_truth_progress: np.ndarray | None = None,
-    ground_truth_stages: np.ndarray | None = None,
-):
-    """
-    Create visualization of SARM predictions with optional ground truth comparison.
-    
-    Args:
-        frames: Video frames (num_frames, H, W, C)
-        progress_predictions: Progress predictions (num_frames,)
-        stage_predictions: Stage probabilities (num_frames, num_stages)
-        task_description: Task description
-        output_path: Path to save the figure
-        num_sample_frames: Number of frames to show
-        figsize: Figure size (width, height)
-        subtask_names: Optional list of subtask names for labeling
-        temporal_proportions: Optional dict of temporal proportions for each subtask
-        ground_truth_progress: Optional ground truth progress array (num_frames,)
-        ground_truth_stages: Optional ground truth stage indices array (num_frames,)
-    """
-    num_stages = stage_predictions.shape[1]
-    stage_colors = plt.cm.tab10(np.linspace(0, 1, num_stages))
-    
-    # Use subtask names if available, otherwise use generic labels
-    if subtask_names is not None and len(subtask_names) == num_stages:
-        stage_labels = subtask_names
-    else:
-        stage_labels = [f'Stage {i+1}' for i in range(num_stages)]
-    
-    # Create figure with progress plot, stage plot, and sample frames
-    fig = plt.figure(figsize=(figsize[0], figsize[1] + 4))
-    gs = gridspec.GridSpec(3, 1, height_ratios=[2, 1, 1], hspace=0.3)
-    
-    ax_progress = fig.add_subplot(gs[0])
-    ax_stages = fig.add_subplot(gs[1], sharex=ax_progress)
-    ax_frames = fig.add_subplot(gs[2])
-    
-    frame_indices = np.arange(len(progress_predictions))
-    
-    # Plot 1: Progress over time
-    ax_progress.plot(frame_indices, progress_predictions, linewidth=2, color='#2E86AB', label='Predicted Progress')
-    ax_progress.fill_between(frame_indices, 0, progress_predictions, alpha=0.3, color='#2E86AB')
-    
-    # Plot ground truth if available
-    if ground_truth_progress is not None:
-        ax_progress.plot(frame_indices, ground_truth_progress, linewidth=2, color='#28A745', 
-                        linestyle='--', label='Ground Truth Progress')
-        ax_progress.fill_between(frame_indices, 0, ground_truth_progress, alpha=0.15, color='#28A745')
-    
-    ax_progress.axhline(y=1.0, color='gray', linestyle='--', alpha=0.5, linewidth=1)
-    ax_progress.set_ylabel('Task Progress', fontsize=12)
-    ax_progress.set_title(f'Task: "{task_description}"', fontsize=14, fontweight='bold')
-    ax_progress.grid(True, alpha=0.3)
-    ax_progress.set_ylim(-0.05, 1.1)
-    ax_progress.legend(loc='upper left')
-    
-    # Add statistics box
-    stats_text = (
-        f'Frames: {len(progress_predictions)}\n'
-        f'Final Progress: {progress_predictions[-1]:.3f}\n'
-        f'Max Progress: {progress_predictions.max():.3f}\n'
-        f'Mean Progress: {progress_predictions.mean():.3f}'
-    )
-    if ground_truth_progress is not None:
-        mse = np.mean((progress_predictions - ground_truth_progress) ** 2)
-        stats_text += f'\nMSE vs GT: {mse:.4f}'
-        stats_text += f'\nGT Final: {ground_truth_progress[-1]:.3f}'
-    
-    ax_progress.text(0.98, 0.02, stats_text, transform=ax_progress.transAxes,
-                    fontsize=10, verticalalignment='bottom', horizontalalignment='right',
-                    bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
-    
-    # Plot 2: Stage predictions (stacked area plot)
-    ax_stages.stackplot(frame_indices, *[stage_predictions[:, i] for i in range(num_stages)],
-                        colors=stage_colors, alpha=0.8, labels=stage_labels)
-    
-    # Plot ground truth stage as vertical bands or markers
-    if ground_truth_stages is not None:
-        # Find stage transition points in ground truth
-        stage_changes = np.where(np.diff(ground_truth_stages) != 0)[0] + 1
-        for change_idx in stage_changes:
-            ax_stages.axvline(x=change_idx, color='black', linestyle='-', alpha=0.7, linewidth=1.5)
-            ax_progress.axvline(x=change_idx, color='black', linestyle='-', alpha=0.3, linewidth=1)
-        
-        # Add small markers at bottom showing GT stage
-        gt_stage_normalized = ground_truth_stages / max(num_stages - 1, 1)
-        ax_stages.scatter(frame_indices[::30], np.zeros(len(frame_indices[::30])) + 0.02, 
-                         c=[stage_colors[s] for s in ground_truth_stages[::30]], 
-                         s=20, marker='|', alpha=0.8, label='GT Stage Markers')
-    
-    ax_stages.set_xlabel('Frame Index', fontsize=12)
-    ax_stages.set_ylabel('Stage Probability', fontsize=12)
-    ax_stages.set_ylim(0, 1)
-    ax_stages.grid(True, alpha=0.3)
-    
-    # Adjust legend based on number of stages and label lengths
-    if num_stages <= 5:
-        ax_stages.legend(loc='upper left', ncol=num_stages, fontsize=8)
-    else:
-        ax_stages.legend(loc='upper left', ncol=3, fontsize=7)
-    
-    # Add vertical lines and labels for expected stage transitions (if temporal proportions available)
-    if temporal_proportions is not None and subtask_names is not None:
-        cumulative_progress = 0.0
-        for i, name in enumerate(stage_labels):
-            if name in temporal_proportions:
-                # Find approximate frame where this stage should end
-                stage_end_progress = cumulative_progress + temporal_proportions[name]
-                
-                # Find frame index closest to this progress
-                progress_diffs = np.abs(progress_predictions - stage_end_progress)
-                stage_end_frame = np.argmin(progress_diffs)
-                
-                # Draw vertical line
-                ax_progress.axvline(x=stage_end_frame, color='gray', linestyle=':', alpha=0.5, linewidth=1)
-                ax_stages.axvline(x=stage_end_frame, color='gray', linestyle=':', alpha=0.5, linewidth=1)
-                
-                cumulative_progress = stage_end_progress
-    
-    # Plot 3: Sample frames (if requested)
-    frame_indices_to_show = np.linspace(0, len(frames) - 1, num_sample_frames, dtype=int)
-    
-    ax_frames.axis('off')
-    
-    # Create grid for frames
-    frame_height = frames[0].shape[0]
-    frame_width = frames[0].shape[1]
-    
-    combined_width = frame_width * num_sample_frames
-    combined_image = np.zeros((frame_height, combined_width, 3), dtype=np.uint8)
-    
-    for i, frame_idx in enumerate(frame_indices_to_show):
-        frame = frames[frame_idx]
-        if frame.shape[-1] == 1:
-            frame = np.repeat(frame, 3, axis=-1)
-        
-        # Add frame to combined image
-        x_start = i * frame_width
-        x_end = (i + 1) * frame_width
-        combined_image[:, x_start:x_end] = frame
-        
-        # Add frame number, progress, and stage
-        progress_val = progress_predictions[frame_idx]
-        stage_idx = np.argmax(stage_predictions[frame_idx])
-        stage_name = stage_labels[stage_idx] if stage_idx < len(stage_labels) else f'{stage_idx+1}'
-        
-        # Truncate long stage names for display
-        if len(stage_name) > 15:
-            stage_name = stage_name[:12] + '...'
-        
-        label = f'Frame {frame_idx}\nProg: {progress_val:.2f}\n{stage_name}'
-        
-        # Draw label on image
-        ax_frames.text(x_start + frame_width / 2, -10, label, 
-                      ha='center', va='top', fontsize=7, 
-                      bbox=dict(boxstyle='round', facecolor='white', alpha=0.7))
-    
-    ax_frames.imshow(combined_image)
-    ax_frames.set_title('Sample Frames', fontsize=12, pad=20)
-    
-    plt.tight_layout()
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    plt.savefig(output_path, dpi=150, bbox_inches='tight')
-    logger.info(f"Saved visualization to {output_path}")
-    
-    plt.close()
-
-
-def main():
-    args = parse_args()
-    
-    # Setup device
-    if args.device is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-    else:
-        device = args.device
-    logger.info(f"Using device: {device}")
-    
-    # Load model
-    logger.info(f"Loading SARM model from {args.model_id}...")
-    model = SARMRewardModel.from_pretrained(args.model_id)
-    model.to(device)
-    model.eval()
-    logger.info("Model loaded successfully")
-    
-    # Load dataset
-    logger.info(f"Loading dataset {args.dataset_repo}...")
-    dataset = LeRobotDataset(args.dataset_repo)
-    logger.info(f"Dataset loaded: {len(dataset.meta.episodes)} episodes, {len(dataset)} frames")
-    
-    # Validate episode index
-    if args.episode_index >= len(dataset.meta.episodes):
-        raise ValueError(
-            f"Episode index {args.episode_index} out of range. "
-            f"Dataset has {len(dataset.meta.episodes)} episodes."
-        )
-    
-    image_key = args.image_key if args.image_key is not None else model.config.image_key
-    state_key = args.state_key if args.state_key is not None else model.config.state_key
-    logger.info(f"Using image key: {image_key}")
-    logger.info(f"Using state key: {state_key}")
-    
-    # Load dataset stats for state normalization (same as training)
-    dataset_stats = load_stats(dataset.root)
-    if dataset_stats:
-        logger.info(f"✓ Loaded dataset stats from {dataset.root}")
-    else:
-        logger.warning("⚠ Could not load dataset stats - states will not be normalized")
-    
-    # Load episode data
-    frames, states, start_idx, end_idx, dataset_task = load_episode_data(
-        dataset, args.episode_index, image_key, state_key
-    )
-    
-    # Use task description from dataset if available, otherwise use command-line argument
-    task_description = dataset_task if dataset_task is not None else args.task_description
-    logger.info(f"Using task description: '{task_description}'")
-    
-    # Run inference
-    progress_predictions, stage_predictions = run_inference(
-        model, frames, states, task_description, 
-        dataset_stats=dataset_stats, state_key=state_key
-    )
-    
-    # Extract subtask names and temporal proportions from model config if available
-    subtask_names = None
-    temporal_proportions = None
-    
-    if hasattr(model.config, 'subtask_names') and model.config.subtask_names is not None:
-        subtask_names = model.config.subtask_names
-        logger.info(f"✓ Found {len(subtask_names)} subtask names in model config: {subtask_names}")
-    
-    # Try to load temporal proportions from model config
-    if hasattr(model.config, 'temporal_proportions') and model.config.temporal_proportions is not None:
-        temporal_proportions = {
-            name: prop for name, prop in zip(model.config.subtask_names, model.config.temporal_proportions)
-        }
-        logger.info(f"✓ Loaded temporal proportions from model config: {temporal_proportions}")
-    
-    # Fallback: try to load from dataset meta
-    if temporal_proportions is None:
-        proportions_path = dataset.root / "meta" / "temporal_proportions.json"
-        if proportions_path.exists():
-            with open(proportions_path, 'r') as f:
-                temporal_proportions = json.load(f)
-                logger.info(f"✓ Loaded temporal proportions from dataset: {temporal_proportions}")
-                
-                # Also extract subtask names from proportions if not already set
-                if subtask_names is None:
-                    subtask_names = sorted(temporal_proportions.keys())
-                    logger.info(f"✓ Extracted subtask names from proportions: {subtask_names}")
-    
-    # Compute ground truth progress if annotations are available
-    ground_truth_progress = None
-    ground_truth_stages = None
-    
-    if temporal_proportions is not None and subtask_names is not None:
-        logger.info("Attempting to compute ground truth progress from annotations...")
-        ground_truth_progress, ground_truth_stages = compute_ground_truth_progress(
-            dataset,
-            args.episode_index,
-            temporal_proportions,
-            subtask_names
-        )
-        if ground_truth_progress is None:
-            logger.warning("⚠ Ground truth not available - annotations may be missing for this episode")
-    else:
-        logger.warning("⚠ Cannot compute ground truth - temporal_proportions or subtask_names not available")
-    
-    output_dir = Path(args.output_dir)
-    output_path = output_dir / f"sarm_prediction_ep{args.episode_index}.png"
-    
-    visualize_predictions(
-        frames,
-        progress_predictions,
-        stage_predictions,
-        task_description,
-        output_path,
-        num_sample_frames=args.num_sample_frames,
-        figsize=tuple(args.figsize),
-        subtask_names=subtask_names,
-        temporal_proportions=temporal_proportions,
-        ground_truth_progress=ground_truth_progress,
-        ground_truth_stages=ground_truth_stages,
-    )
-    
-    predictions_path = output_dir / f"predictions_ep{args.episode_index}.npz"
-    save_dict = {
-        'progress': progress_predictions, 
-        'stages': stage_predictions
-    }
-    if ground_truth_progress is not None:
-        save_dict['gt_progress'] = ground_truth_progress
-        save_dict['gt_stages'] = ground_truth_stages
-    np.savez(predictions_path, **save_dict)
-    logger.info(f"Saved predictions to {predictions_path}")
-    logger.info(f"\nVisualization: {output_path}")
-
-
-if __name__ == "__main__":
-    main()
-

src/lerobot/configs/train.py

@@ -64,26 +64,9 @@ class TrainPipelineConfig(HubMixin):
     scheduler: LRSchedulerConfig | None = None
     eval: EvalConfig = field(default_factory=EvalConfig)
     wandb: WandBConfig = field(default_factory=WandBConfig)
-    
-    # RA-BC (Reward-Aligned Behavior Cloning) parameters
-    use_rabc: bool = False  # Enable reward-weighted training
-    reward_model_path: str | None = None  # Path to pre-trained reward model (e.g., SARM)
-    rabc_kappa: float = 0.01  # Hard threshold for high-quality samples
-    rabc_epsilon: float = 1e-6  # Small constant for numerical stability
-    rabc_update_freq: int = 1  # Compute rewards every N batches (1 = every batch)
-
-    # Rename map for the observation to override the image and state keys
-    rename_map: dict[str, str] = field(default_factory=dict)       
     checkpoint_path: Path | None = field(init=False, default=None)
-        
-
-    def validate(self):
-        # Validate RA-BC configuration
-        if self.use_rabc and not self.reward_model_path:
-            raise ValueError(
-                "RA-BC is enabled (use_rabc=True) but no reward_model_path provided. "
-                "Please specify a pre-trained reward model (e.g., SARM) path."
-            )
+    # Rename map for the observation to override the image and state keys
+    rename_map: dict[str, str] = field(default_factory=dict)
 
     def validate(self) -> None:
         # HACK: We parse again the cli args here to get the pretrained paths if there was some.

src/lerobot/configs/types.py

@@ -43,10 +43,3 @@ class NormalizationMode(str, Enum):
 class PolicyFeature:
     type: FeatureType
     shape: tuple[int, ...]
-
-
-class RTCAttentionSchedule(str, Enum):
-    ZEROS = "ZEROS"
-    ONES = "ONES"
-    LINEAR = "LINEAR"
-    EXP = "EXP"

src/lerobot/datasets/dataset_tools.py

@@ -39,7 +39,6 @@ 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 (
-    DATA_DIR,
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
@@ -963,23 +962,28 @@ def _copy_data_with_feature_changes(
     remove_features: list[str] | None = None,
 ) -> None:
     """Copy data while adding or removing features."""
-    data_dir = dataset.root / DATA_DIR
-    parquet_files = sorted(data_dir.glob("*/*.parquet"))
+    if dataset.meta.episodes is None:
+        dataset.meta.episodes = load_episodes(dataset.meta.root)
 
-    if not parquet_files:
-        raise ValueError(f"No parquet files found in {data_dir}")
+    # Map file paths to episode indices to extract chunk/file indices
+    file_to_episodes: dict[Path, set[int]] = {}
+    for ep_idx in range(dataset.meta.total_episodes):
+        file_path = dataset.meta.get_data_file_path(ep_idx)
+        if file_path not in file_to_episodes:
+            file_to_episodes[file_path] = set()
+        file_to_episodes[file_path].add(ep_idx)
 
     frame_idx = 0
 
-    for src_path in tqdm(parquet_files, desc="Processing data files"):
-        df = pd.read_parquet(src_path).reset_index(drop=True)
+    for src_path in tqdm(sorted(file_to_episodes.keys()), desc="Processing data files"):
+        df = pd.read_parquet(dataset.root / src_path).reset_index(drop=True)
 
-        relative_path = src_path.relative_to(dataset.root)
-        chunk_dir = relative_path.parts[1]
-        file_name = relative_path.parts[2]
-
-        chunk_idx = int(chunk_dir.split("-")[1])
-        file_idx = int(file_name.split("-")[1].split(".")[0])
+        # Get chunk_idx and file_idx from the source file's first episode
+        episodes_in_file = file_to_episodes[src_path]
+        first_ep_idx = min(episodes_in_file)
+        src_ep = dataset.meta.episodes[first_ep_idx]
+        chunk_idx = src_ep["data/chunk_index"]
+        file_idx = src_ep["data/file_index"]
 
         if remove_features:
             df = df.drop(columns=remove_features, errors="ignore")
@@ -999,21 +1003,13 @@ def _copy_data_with_feature_changes(
                     df[feature_name] = feature_values
                 else:
                     feature_slice = values[frame_idx:end_idx]
-                    if len(feature_slice.shape) == 1:
-                        # 1D array - can assign directly
-                        df[feature_name] = feature_slice
-                    elif len(feature_slice.shape) == 2 and feature_slice.shape[1] == 1:
-                        # 2D array with single column - flatten it
+                    if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
                         df[feature_name] = feature_slice.flatten()
-                    elif len(feature_slice.shape) == 2:
-                        # 2D array with multiple columns (e.g., embeddings) - convert to list of lists
-                        df[feature_name] = feature_slice.tolist()
                     else:
-                        # Higher dimensional - convert to list
-                        df[feature_name] = [row.tolist() for row in feature_slice]
+                        df[feature_name] = feature_slice
             frame_idx = end_idx
 
-        # Write using the same chunk/file structure as source
+        # Write using the preserved chunk_idx and file_idx from source
         dst_path = new_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
         dst_path.parent.mkdir(parents=True, exist_ok=True)
 

src/lerobot/datasets/generating_embeddings/README.md

@@ -1,146 +0,0 @@
-# LeRobot Embedding Generation Script
-
-Generate embeddings for LeRobot datasets to make them more lightweight and efficient for training.
-
-## Overview
-
-This script processes v3.0 LeRobot datasets and adds pre-computed embeddings for:
-
-- **Task embeddings**: Language command embeddings using MiniLM
-- **Image embeddings**: Frame embeddings using DinoV2
-
-The resulting dataset can be used more efficiently during training by loading pre-computed embeddings instead of running encoders on-the-fly.
-
-## Supported Encoders
-
-### Image Encoders (DinoV2)
-
-DinoV2 is a self-supervised vision transformer that produces high-quality image embeddings:
-
-- **`dinov2_vits14`**: ViT-S/14 (384-dim) - Fastest, smaller model
-- **`dinov2_vitb14`**: ViT-B/14 (768-dim) - **Recommended** - Good balance
-- **`dinov2_vitl14`**: ViT-L/14 (1024-dim) - Best quality, slower
-
-### Language Encoders (MiniLM)
-
-MiniLM is a lightweight sentence transformer model:
-
-- **`minilm-l6`**: MiniLM-L6-v2 (384-dim) - Faster
-- **`minilm-l12`**: MiniLM-L12-v2 (384-dim) - **Recommended** - Better quality
-
-## Usage
-
-### Basic Command
-
-```bash
-python src/lerobot/datasets/generating_embeddings/generate_embeddings.py \
-    --repo-id lerobot/utokyo_xarm_bimanual \
-    --output-repo-id your-username/utokyo_xarm_bimanual_embeddings \
-    --image-encoder dinov2_vitb14 \
-    --language-encoder minilm-l12 \
-    --push-to-hub
-```
-
-### Lightweight Version (No Videos)
-
-Removes video files to significantly reduce storage:
-
-```bash
-python src/lerobot/datasets/generating_embeddings/generate_embeddings.py \
-    --repo-id lerobot/utokyo_xarm_bimanual \
-    --output-repo-id your-username/utokyo_xarm_bimanual_lightweight \
-    --image-encoder dinov2_vitb14 \
-    --language-encoder minilm-l12 \
-    --remove-videos \
-    --push-to-hub
-```
-
-## Output
-
-The script adds new features to your dataset:
-
-### New Features
-
-1. **`task_embedding`**: Language embedding for each frame
-   - Shape: `[384]` (MiniLM)
-   - One embedding per frame based on its task
-
-2. **`{camera_key}_embedding`**: Image embedding for each camera view
-   - Shape: `[384]`, `[768]`, or `[1024]` depending on DinoV2 model
-   - Examples: `observation.images.top_embedding`, `observation.images.wrist_embedding`
-
-### Using Embeddings in Training
-
-```python
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-# Load dataset with embeddings
-dataset = LeRobotDataset("your-username/utokyo_xarm_bimanual_embeddings")
-
-# Access embeddings
-item = dataset[0]
-task_emb = item["task_embedding"]  # Shape: [384]
-img_emb = item["observation.images.top_embedding"]  # Shape: [768]
-
-# Use in your policy
-# Instead of running encoders during training, use pre-computed embeddings
-```
-
-## Extending with New Encoders
-
-The script is designed to be easily extensible. To add a new encoder:
-
-### 1. Create Encoder Class
-
-```python
-class MyCustomImageEncoder(ImageEncoder):
-    """Your custom image encoder."""
-
-    def __init__(self, device: str = "cuda"):
-        super().__init__(device)
-        # Load your model
-        self.model = load_my_model()
-        self.model = self.model.to(self.device)
-        self.model.eval()
-
-    def encode(self, images: list[np.ndarray]) -> np.ndarray:
-        """Encode a batch of images."""
-        # Your encoding logic here
-        embeddings = []
-        for img in images:
-            emb = self.model(img)
-            embeddings.append(emb)
-        return np.array(embeddings)
-
-    @property
-    def embedding_dim(self) -> int:
-        """Return embedding dimension."""
-        return 512  # Your embedding dimension
-```
-
-### 2. Add to Factory Function
-
-```python
-def get_image_encoder(encoder_name: str, device: str = "cuda") -> ImageEncoder:
-    encoders = {
-        "dinov2_vits14": lambda: DinoV2Encoder(model_name="dinov2_vits14", device=device),
-        "dinov2_vitb14": lambda: DinoV2Encoder(model_name="dinov2_vitb14", device=device),
-        "dinov2_vitl14": lambda: DinoV2Encoder(model_name="dinov2_vitl14", device=device),
-        # Add your encoder
-        "my_custom": lambda: MyCustomImageEncoder(device=device),
-    }
-    # ... rest of function
-```
-
-## Validating Embeddings
-
-After generating embeddings, you can validate them using `validate_embeddings.py`:
-
-```bash
-python src/lerobot/datasets/generating_embeddings/validate_embeddings.py \
-    --original-repo-id lerobot/utokyo_xarm_bimanual \
-    --embeddings-repo-id pepijn223/utokyo_xarm_bimanual_embeddings \
-    --image-encoder dinov2_vitb14 \
-    --language-encoder minilm-l12 \
-    --num-samples 20
-```

src/lerobot/datasets/generating_embeddings/encoders.py

@@ -1,147 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-
-import numpy as np
-import torch
-from PIL import Image
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class ImageEncoder:
-    """Base class for image encoders."""
-
-    def __init__(self, device: str = "cuda"):
-        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
-
-    def encode(self, images: list[np.ndarray]) -> np.ndarray:
-        """Encode a batch of images."""
-        raise NotImplementedError
-
-
-class DinoV2Encoder(ImageEncoder):
-    """DinoV2 image encoder.
-
-    DinoV2 is a self-supervised vision transformer that produces high-quality image embeddings.
-    Supports multiple model sizes (ViT-S/14, ViT-B/14, ViT-L/14).
-    """
-
-    def __init__(self, model_name: str = "dinov2_vitb14", device: str = "cuda", batch_size: int = 32):
-        super().__init__(device)
-        self.batch_size = batch_size
-        self.model_name = model_name
-        logger.info(f"Loading DinoV2 model: {model_name}")
-        self.model = torch.hub.load("facebookresearch/dinov2", model_name)  # nosec B614
-        self.model = self.model.to(self.device)
-        self.model.eval()
-
-        # DinoV2 preprocessing
-        from torchvision import transforms
-
-        self.transform = transforms.Compose(
-            [
-                transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
-                transforms.CenterCrop(224),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-            ]
-        )
-
-    def encode(self, images: list[np.ndarray]) -> np.ndarray:
-        """Encode a batch of images."""
-        embeddings = []
-
-        with torch.inference_mode():
-            for i in range(0, len(images), self.batch_size):
-                batch_images = images[i : i + self.batch_size]
-                # Convert numpy arrays to PIL Images and apply transforms
-                pil_images = [Image.fromarray(img.astype(np.uint8)) for img in batch_images]
-                tensors = torch.stack([self.transform(img) for img in pil_images]).to(self.device)
-
-                # Get embeddings
-                batch_embeddings = self.model(tensors).cpu().numpy()
-                embeddings.append(batch_embeddings)
-
-        return np.concatenate(embeddings, axis=0)
-
-    @property
-    def embedding_dim(self) -> int:
-        """Return the embedding dimension based on model size."""
-        if "vits14" in self.model_name:
-            return 384  # DinoV2 ViT-S/14
-        elif "vitb14" in self.model_name:
-            return 768  # DinoV2 ViT-B/14
-        elif "vitl14" in self.model_name:
-            return 1024  # DinoV2 ViT-L/14
-        else:
-            return 768  # Default to ViT-B/14
-
-
-class LanguageEncoder:
-    """Base class for language encoders."""
-
-    def __init__(self, device: str = "cuda"):
-        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
-
-    def encode(self, texts: list[str]) -> np.ndarray:
-        """Encode a batch of texts."""
-        raise NotImplementedError
-
-
-class MiniLMEncoder(LanguageEncoder):
-    """MiniLM language encoder.
-
-    MiniLM is a lightweight sentence transformer model that produces high-quality text embeddings.
-    Supports L6 and L12 model sizes.
-    """
-
-    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L12-v2", device: str = "cuda"):
-        super().__init__(device)
-        self.model_name = model_name
-        logger.info(f"Loading MiniLM model: {model_name}")
-
-        from transformers import AutoModel, AutoTokenizer
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModel.from_pretrained(model_name).to(self.device)
-        self.model.eval()
-
-    def _mean_pooling(self, model_output, attention_mask):
-        """Mean pooling to get sentence embeddings."""
-        token_embeddings = model_output[0]
-        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
-            input_mask_expanded.sum(1), min=1e-9
-        )
-
-    def encode(self, texts: list[str]) -> np.ndarray:
-        """Encode a batch of texts."""
-        with torch.inference_mode():
-            encoded_input = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
-            encoded_input = {k: v.to(self.device) for k, v in encoded_input.items()}
-
-            model_output = self.model(**encoded_input)
-            embeddings = self._mean_pooling(model_output, encoded_input["attention_mask"])
-
-            return embeddings.cpu().numpy()
-
-    @property
-    def embedding_dim(self) -> int:
-        """Return the embedding dimension."""
-        return 384  # Both MiniLM-L6 and L12 output 384-dim embeddings

src/lerobot/datasets/generating_embeddings/generate_embeddings.py

@@ -1,329 +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.
-
-"""
-Generate embeddings for LeRobot datasets to make them more lightweight and efficient.
-
-This script:
-1. Loads a v3.0 LeRobot dataset from the hub
-2. Computes embeddings for tasks (language commands) and frames (images)
-3. Stores embeddings as new features in the dataset
-4. Optionally removes video files to reduce size
-5. Pushes the converted dataset to the hub
-
-Current supported encoders:
-- Image: DinoV2 (dinov2_vits14, dinov2_vitb14, dinov2_vitl14)
-- Language: MiniLM (minilm-l6, minilm-l12)
-
-The architecture is extensible - you can add more encoders by:
-1. Creating a new encoder class inheriting from ImageEncoder or LanguageEncoder
-2. Implementing the encode() method and embedding_dim property
-3. Adding it to the get_image_encoder() or get_language_encoder() factory function
-
-Usage example:
-    python src/lerobot/datasets/generating_embeddings/generate_embeddings.py \
-        --repo-id lerobot/utokyo_xarm_bimanual \
-        --output-repo-id lerobot/utokyo_xarm_bimanual_embeddings \
-        --image-encoder dinov2_vitb14 \
-        --language-encoder minilm-l12 \
-        --remove-videos \
-        --push-to-hub
-"""
-
-import argparse
-import shutil
-from pathlib import Path
-
-import numpy as np
-import torch
-from tqdm import tqdm
-
-from lerobot.datasets.generating_embeddings.encoders import (
-    DinoV2Encoder,
-    ImageEncoder,
-    LanguageEncoder,
-    MiniLMEncoder,
-)
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-
-def get_image_encoder(encoder_name: str, device: str = "cuda") -> ImageEncoder:
-    """Factory function to get image encoder.
-
-    To add a new encoder:
-    1. Create a new class inheriting from ImageEncoder
-    2. Implement encode() and embedding_dim property
-    3. Add it to the encoders dictionary below
-    """
-    encoders = {
-        "dinov2_vits14": lambda: DinoV2Encoder(model_name="dinov2_vits14", device=device),
-        "dinov2_vitb14": lambda: DinoV2Encoder(model_name="dinov2_vitb14", device=device),
-        "dinov2_vitl14": lambda: DinoV2Encoder(model_name="dinov2_vitl14", device=device),
-    }
-
-    if encoder_name not in encoders:
-        raise ValueError(f"Unknown image encoder: {encoder_name}. Available options: {list(encoders.keys())}")
-
-    return encoders[encoder_name]()
-
-
-def get_language_encoder(encoder_name: str, device: str = "cuda") -> LanguageEncoder:
-    """Factory function to get language encoder.
-
-    To add a new encoder:
-    1. Create a new class inheriting from LanguageEncoder
-    2. Implement encode() and embedding_dim property
-    3. Add it to the encoders dictionary below
-    """
-    encoders = {
-        "minilm-l6": lambda: MiniLMEncoder(
-            model_name="sentence-transformers/all-MiniLM-L6-v2", device=device
-        ),
-        "minilm-l12": lambda: MiniLMEncoder(
-            model_name="sentence-transformers/all-MiniLM-L12-v2", device=device
-        ),
-    }
-
-    if encoder_name not in encoders:
-        raise ValueError(
-            f"Unknown language encoder: {encoder_name}. Available options: {list(encoders.keys())}"
-        )
-
-    return encoders[encoder_name]()
-
-
-def generate_embeddings_for_dataset(
-    repo_id: str,
-    output_repo_id: str,
-    image_encoder: ImageEncoder,
-    language_encoder: LanguageEncoder,
-    remove_videos: bool = False,
-    local_dir: Path | None = None,
-    output_local_dir: Path | None = None,
-    push_to_hub: bool = False,
-):
-    """Generate embeddings for a LeRobot dataset.
-
-    Args:
-        repo_id: Source dataset repository ID
-        output_repo_id: Output dataset repository ID
-        image_encoder: Image encoder instance
-        language_encoder: Language encoder instance
-        remove_videos: Whether to remove video files
-        local_dir: Local directory for source dataset
-        output_local_dir: Local directory for output dataset
-        push_to_hub: Whether to push to hub after conversion
-    """
-    from lerobot.datasets.dataset_tools import modify_features
-
-    print(f"Loading dataset: {repo_id}")
-
-    dataset = LeRobotDataset(repo_id, root=local_dir, download_videos=True)
-    print(f"Dataset: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
-
-    print("Computing task embeddings...")
-    unique_tasks = dataset.meta.tasks.index.tolist()
-    task_embeddings = {}
-
-    for task in tqdm(unique_tasks, desc="Encoding tasks"):
-        # Clean up task text
-        task_clean = task.strip().capitalize().strip(" .,!?-_")
-        embedding = language_encoder.encode([task_clean])[0]
-        task_embeddings[task] = embedding
-
-    print(f"Computed {len(task_embeddings)} task embeddings")
-
-    print("Processing frames and computing embeddings...")
-    all_task_embeddings = []
-    all_image_embeddings_dict = {cam_key: [] for cam_key in dataset.meta.camera_keys}
-
-    for frame_idx in tqdm(range(dataset.num_frames), desc="Processing frames"):
-        item = dataset.hf_dataset[frame_idx]
-        ep_idx = item["episode_index"].item()
-
-        task = dataset.meta.tasks.iloc[item["task_index"].item()].name
-        task_emb = task_embeddings[task]
-        all_task_embeddings.append(task_emb)
-
-        for cam_key in dataset.meta.camera_keys:
-            if cam_key in dataset.meta.video_keys:
-                current_ts = item["timestamp"].item()
-                video_frames = dataset._query_videos({cam_key: [current_ts]}, ep_idx)
-                img = video_frames[cam_key]
-
-                if isinstance(img, torch.Tensor):
-                    if img.ndim == 4:
-                        img = img[0]  # (T, C, H, W) -> (C, H, W)
-                    elif img.ndim != 3:
-                        raise ValueError(f"Unexpected video frame shape {img.shape} for camera {cam_key}")
-                    img_np = (img.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
-                else:
-                    img_np = np.array(img)
-            else:
-                img = item[cam_key]
-                if isinstance(img, torch.Tensor):
-                    if img.ndim == 3:
-                        img_np = (img.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
-                    else:
-                        raise ValueError(f"Unexpected image shape {img.shape} for camera {cam_key}")
-                else:
-                    img_np = np.array(img)
-
-            all_image_embeddings_dict[cam_key].append(img_np)
-
-    print("Computing image embeddings...")
-    image_embeddings_dict = {}
-    for cam_key, images in all_image_embeddings_dict.items():
-        print(f"  {cam_key}: {len(images)} images")
-        embeddings = image_encoder.encode(images)
-        image_embeddings_dict[cam_key] = embeddings
-
-    all_task_embeddings = np.array(all_task_embeddings)
-    for cam_key in dataset.meta.camera_keys:
-        image_embeddings_dict[cam_key] = np.array(image_embeddings_dict[cam_key])
-
-    img_emb_dim = image_encoder.embedding_dim
-    lang_emb_dim = language_encoder.embedding_dim
-
-    add_features_dict = {
-        "task_embedding": (
-            all_task_embeddings,
-            {"dtype": "float32", "shape": [lang_emb_dim], "names": None},
-        ),
-    }
-
-    for cam_key in dataset.meta.camera_keys:
-        add_features_dict[f"{cam_key}_embedding"] = (
-            image_embeddings_dict[cam_key],
-            {"dtype": "float32", "shape": [img_emb_dim], "names": None},
-        )
-
-    print("Adding embeddings to dataset...")
-    remove_features_list = None
-    if remove_videos:
-        remove_features_list = dataset.meta.video_keys
-
-    output_dataset = modify_features(
-        dataset=dataset,
-        add_features=add_features_dict,
-        remove_features=remove_features_list,
-        output_dir=output_local_dir,
-        repo_id=output_repo_id,
-    )
-
-    if remove_videos:
-        print("Removing video files...")
-        videos_dir = output_dataset.root / "videos"
-        if videos_dir.exists():
-            shutil.rmtree(videos_dir)
-
-    print(f"Saved to: {output_dataset.root}")
-
-    if push_to_hub:
-        print(f"Pushing to hub: {output_repo_id}")
-        output_dataset.push_to_hub(push_videos=not remove_videos)
-        print("Done!")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Generate embeddings for LeRobot datasets",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-  # Basic usage with default encoders (DinoV2 ViT-B/14 + MiniLM-L12)
-  python src/lerobot/datasets/generating_embeddings/generate_embeddings.py \\
-      --repo-id lerobot/utokyo_xarm_bimanual \\
-      --output-repo-id your-username/utokyo_xarm_bimanual_embeddings \\
-      --image-encoder dinov2_vitb14 \\
-      --language-encoder minilm-l12 \\
-      --push-to-hub
-
-  # Generate embeddings and remove videos
-  python src/lerobot/datasets/generating_embeddings/generate_embeddings.py \\
-      --repo-id lerobot/utokyo_xarm_bimanual \\
-      --output-repo-id your-username/utokyo_xarm_bimanual_lightweight \\
-      --image-encoder dinov2_vitb14 \\
-      --language-encoder minilm-l12 \\
-      --remove-videos \\
-      --push-to-hub
-
-Available image encoders:
-  - dinov2_vits14: DinoV2 ViT-S/14 (384-dim, faster)
-  - dinov2_vitb14: DinoV2 ViT-B/14 (768-dim, recommended)
-  - dinov2_vitl14: DinoV2 ViT-L/14 (1024-dim, best quality)
-
-Available language encoders:
-  - minilm-l6: MiniLM-L6-v2 (384-dim, faster)
-  - minilm-l12: MiniLM-L12-v2 (384-dim, recommended)
-        """,
-    )
-    parser.add_argument("--repo-id", type=str, required=True, help="Source dataset repository ID")
-    parser.add_argument("--output-repo-id", type=str, required=True, help="Output dataset repository ID")
-    parser.add_argument(
-        "--image-encoder",
-        type=str,
-        default="dinov2_vitb14",
-        help="Image encoder to use (default: dinov2_vitb14)",
-    )
-    parser.add_argument(
-        "--language-encoder",
-        type=str,
-        default="minilm-l12",
-        help="Language encoder to use (default: minilm-l12)",
-    )
-    parser.add_argument(
-        "--remove-videos",
-        action="store_true",
-        help="Remove video files after generating embeddings",
-    )
-    parser.add_argument("--local-dir", type=str, default=None, help="Local directory for source dataset")
-    parser.add_argument(
-        "--output-local-dir", type=str, default=None, help="Local directory for output dataset"
-    )
-    parser.add_argument(
-        "--push-to-hub",
-        action="store_true",
-        help="Push the converted dataset to the hub",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda",
-        help="Device to use for encoding (default: cuda)",
-    )
-
-    args = parser.parse_args()
-
-    # Load encoders
-    image_encoder = get_image_encoder(args.image_encoder, device=args.device)
-    language_encoder = get_language_encoder(args.language_encoder, device=args.device)
-
-    # Generate embeddings
-    generate_embeddings_for_dataset(
-        repo_id=args.repo_id,
-        output_repo_id=args.output_repo_id,
-        image_encoder=image_encoder,
-        language_encoder=language_encoder,
-        remove_videos=args.remove_videos,
-        local_dir=Path(args.local_dir) if args.local_dir else None,
-        output_local_dir=Path(args.output_local_dir) if args.output_local_dir else None,
-        push_to_hub=args.push_to_hub,
-    )
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/datasets/generating_embeddings/validate_embeddings.py

@@ -1,222 +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.
-
-"""
-Validate pre-computed embeddings against on-the-fly computed embeddings.
-
-Usage:
-    python src/lerobot/datasets/generating_embeddings/validate_embeddings.py \
-        --original-repo-id lerobot/utokyo_xarm_bimanual \
-        --embeddings-repo-id <your_username>/utokyo_xarm_bimanual_embeddings \
-        --image-encoder dinov2_vitb14 \
-        --language-encoder minilm-l12 \
-        --num-samples 10
-"""
-
-import argparse
-
-import numpy as np
-import torch
-from tqdm import tqdm
-
-from lerobot.datasets.generating_embeddings.encoders import ImageEncoder, LanguageEncoder
-from lerobot.datasets.generating_embeddings.generate_embeddings import (
-    get_image_encoder,
-    get_language_encoder,
-)
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-
-def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
-    """Compute cosine similarity between two vectors."""
-    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
-
-
-def validate_embeddings(
-    original_repo_id: str,
-    embeddings_repo_id: str,
-    image_encoder: ImageEncoder,
-    language_encoder: LanguageEncoder,
-    num_samples: int = 10,
-    device: str = "cuda",
-):
-    """Validate pre-computed embeddings against on-the-fly embeddings.
-
-    Args:
-        original_repo_id: Original dataset repository ID
-        embeddings_repo_id: Dataset with pre-computed embeddings repository ID
-        image_encoder: Image encoder instance
-        language_encoder: Language encoder instance
-        num_samples: Number of samples to validate
-        device: Device to use for encoding
-    """
-    # Load both datasets
-    print("Loading datasets...")
-    original_dataset = LeRobotDataset(original_repo_id, download_videos=True)
-    embeddings_dataset = LeRobotDataset(embeddings_repo_id, download_videos=False)
-
-    # Verify both datasets have the same number of frames
-    assert original_dataset.num_frames == embeddings_dataset.num_frames, (
-        f"Frame count mismatch: original={original_dataset.num_frames}, "
-        f"embeddings={embeddings_dataset.num_frames}"
-    )
-
-    camera_keys = original_dataset.meta.camera_keys
-
-    # Check embedding features exist
-    expected_features = ["task_embedding"] + [f"{cam}_embedding" for cam in camera_keys]
-    for feat in expected_features:
-        if feat not in embeddings_dataset.features:
-            raise ValueError(f"Embedding feature not found: {feat}")
-
-    # Select random sample indices
-    sample_indices = np.random.choice(
-        original_dataset.num_frames, size=min(num_samples, original_dataset.num_frames), replace=False
-    )
-    print(f"Validating {len(sample_indices)} samples...")
-
-    # Track statistics
-    task_similarities = []
-    image_similarities = {cam: [] for cam in camera_keys}
-
-    for idx in tqdm(sample_indices, desc="Validating"):
-        idx = int(idx)
-
-        embeddings_item = embeddings_dataset[idx]
-        precomputed_task_emb = embeddings_item["task_embedding"].numpy()
-        precomputed_image_embs = {cam: embeddings_item[f"{cam}_embedding"].numpy() for cam in camera_keys}
-
-        original_item = original_dataset[idx]
-
-        # Get task and compute embedding
-        task = original_item["task"]
-        # Clean up task text (same as in generate_embeddings.py)
-        task_clean = task.strip().capitalize().strip(" .,!?-_")
-        onthefly_task_emb = language_encoder.encode([task_clean])[0]
-
-        # Get images and compute embeddings
-        onthefly_image_embs = {}
-        for cam in camera_keys:
-            img = original_item[cam]
-            # Convert to numpy if needed
-            if isinstance(img, torch.Tensor):
-                if img.ndim == 3:  # (C, H, W)
-                    img_np = (img.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
-                else:
-                    raise ValueError(f"Unexpected image shape: {img.shape}")
-            else:
-                img_np = np.array(img)
-
-            onthefly_image_embs[cam] = image_encoder.encode([img_np])[0]
-
-        # Task embedding comparison
-        task_sim = cosine_similarity(precomputed_task_emb, onthefly_task_emb)
-        task_similarities.append(task_sim)
-
-        # Image embedding comparison
-        for cam in camera_keys:
-            img_sim = cosine_similarity(precomputed_image_embs[cam], onthefly_image_embs[cam])
-            image_similarities[cam].append(img_sim)
-
-    # Results
-    print("\nResults:")
-    task_sim_threshold = 0.99
-    img_sim_threshold = 0.99
-
-    task_mean_sim = np.mean(task_similarities)
-    task_pass = task_mean_sim >= task_sim_threshold
-
-    print(f"  Task: {task_mean_sim:.4f} {'✓' if task_pass else '✗'}")
-
-    for cam in camera_keys:
-        cam_mean_sim = np.mean(image_similarities[cam])
-        cam_pass = cam_mean_sim >= img_sim_threshold
-        print(f"  {cam}: {cam_mean_sim:.4f} {'✓' if cam_pass else '✗'}")
-
-    image_pass = all(np.mean(image_similarities[cam]) >= img_sim_threshold for cam in camera_keys)
-
-    print()
-    if task_pass and image_pass:
-        print("✓ PASSED")
-    else:
-        print("✗ FAILED")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Validate and compare pre-computed embeddings with on-the-fly embeddings",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Example:
-  python src/lerobot/datasets/generating_embeddings/validate_embeddings.py \\
-      --original-repo-id lerobot/utokyo_xarm_bimanual \\
-      --embeddings-repo-id lerobot/utokyo_xarm_bimanual_embeddings \\
-      --image-encoder dinov2_vitb14 \\
-      --language-encoder minilm-l12 \\
-      --num-samples 20
-        """,
-    )
-    parser.add_argument("--original-repo-id", type=str, required=True, help="Original dataset repository ID")
-    parser.add_argument(
-        "--embeddings-repo-id",
-        type=str,
-        required=True,
-        help="Dataset with pre-computed embeddings repository ID",
-    )
-    parser.add_argument(
-        "--image-encoder",
-        type=str,
-        default="dinov2_vitb14",
-        help="Image encoder to use (default: dinov2_vitb14)",
-    )
-    parser.add_argument(
-        "--language-encoder",
-        type=str,
-        default="minilm-l12",
-        help="Language encoder to use (default: minilm-l12)",
-    )
-    parser.add_argument(
-        "--num-samples",
-        type=int,
-        default=10,
-        help="Number of samples to validate (default: 10)",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda",
-        help="Device to use for encoding (default: cuda)",
-    )
-
-    args = parser.parse_args()
-
-    # Load encoders
-    image_encoder = get_image_encoder(args.image_encoder, device=args.device)
-    language_encoder = get_language_encoder(args.language_encoder, device=args.device)
-
-    # Validate embeddings
-    validate_embeddings(
-        original_repo_id=args.original_repo_id,
-        embeddings_repo_id=args.embeddings_repo_id,
-        image_encoder=image_encoder,
-        language_encoder=language_encoder,
-        num_samples=args.num_samples,
-        device=args.device,
-    )
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/datasets/lerobot_dataset.py

@@ -430,7 +430,9 @@ class LeRobotDatasetMetadata:
         video_keys = [video_key] if video_key is not None else self.video_keys
         for key in video_keys:
             if not self.features[key].get("info", None):
-                video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0)
+                video_path = self.root / self.video_path.format(
+                    video_key=video_key, chunk_index=0, file_index=0
+                )
                 self.info["features"][key]["info"] = get_video_info(video_path)
 
     def update_chunk_settings(
@@ -712,15 +714,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
             self.download(download_videos)
             self.hf_dataset = self.load_hf_dataset()
 
-        # Create mapping from absolute indices to relative indices when only a subset of the episodes are loaded
-        # Build a mapping: absolute_index -> relative_index_in_filtered_dataset
-        self._absolute_to_relative_idx = None
-        if self.episodes is not None:
-            self._absolute_to_relative_idx = {
-                abs_idx.item() if isinstance(abs_idx, torch.Tensor) else abs_idx: rel_idx
-                for rel_idx, abs_idx in enumerate(self.hf_dataset["index"])
-            }
-
         # Setup delta_indices
         if self.delta_timestamps is not None:
             check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s)
@@ -839,7 +832,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
     def load_hf_dataset(self) -> datasets.Dataset:
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
         features = get_hf_features_from_features(self.features)
-        hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
+        hf_dataset = load_nested_dataset(self.root / "data", features=features)
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
@@ -856,8 +849,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         # Determine requested episodes
         if self.episodes is None:
+            # Requesting all episodes - check if we have all episodes from metadata
             requested_episodes = set(range(self.meta.total_episodes))
         else:
+            # Requesting specific episodes
             requested_episodes = set(self.episodes)
 
         # Check if all requested episodes are available in cached data
@@ -939,11 +934,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         query_timestamps = {}
         for key in self.meta.video_keys:
             if query_indices is not None and key in query_indices:
-                if self._absolute_to_relative_idx is not None:
-                    relative_indices = [self._absolute_to_relative_idx[idx] for idx in query_indices[key]]
-                    timestamps = self.hf_dataset[relative_indices]["timestamp"]
-                else:
-                    timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
+                timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
                 query_timestamps[key] = torch.stack(timestamps).tolist()
             else:
                 query_timestamps[key] = [current_ts]
@@ -951,32 +942,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
         return query_timestamps
 
     def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
-        """
-        Query dataset for indices across keys, skipping video keys.
-
-        Tries column-first [key][indices] for speed, falls back to row-first.
-
-        Args:
-            query_indices: Dict mapping keys to index lists to retrieve
-
-        Returns:
-            Dict with stacked tensors of queried data (video keys excluded)
-        """
-        result: dict = {}
-        for key, q_idx in query_indices.items():
-            if key in self.meta.video_keys:
-                continue
-            # Map absolute indices to relative indices if needed
-            relative_indices = (
-                q_idx
-                if self._absolute_to_relative_idx is None
-                else [self._absolute_to_relative_idx[idx] for idx in q_idx]
-            )
-            try:
-                result[key] = torch.stack(self.hf_dataset[key][relative_indices])
-            except (KeyError, TypeError, IndexError):
-                result[key] = torch.stack(self.hf_dataset[relative_indices][key])
-        return result
+        return {
+            key: torch.stack(self.hf_dataset[q_idx][key])
+            for key, q_idx in query_indices.items()
+            if key not in self.meta.video_keys
+        }
 
     def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
         """Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
@@ -1515,7 +1485,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.image_transforms = None
         obj.delta_timestamps = None
         obj.delta_indices = None
-        obj._absolute_to_relative_idx = None
         obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         obj.writer = None
         obj.latest_episode = None

src/lerobot/datasets/temporal_sampler.py

@@ -1,151 +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.
-
-"""
-SARM Temporal Sampler for reward model training.
-
-Samples frames uniformly from episodes for SARM's 9-frame symmetric pattern:
-- 1 initial frame + 4 frames before + current + 3 frames after
-
-Boundary handling: clamp to first/last frame when indices go out of bounds.
-This enables truly uniform sampling across entire episodes.
-"""
-
-import logging
-from typing import Iterator, Optional
-import numpy as np
-import torch
-from torch.utils.data import Sampler
-import random
-
-
-class SARMTemporalSampler(Sampler):
-    """
-    Temporal sampler for SARM reward model training with symmetric/bidirectional sampling.
-    
-    SARM uses 9 frames per sample:
-    - Frame 0: Initial frame of the episode (always frame 0)
-    - Frames 1-8: Symmetric context around current frame
-      Pattern: [t-4*gap, t-3*gap, t-2*gap, t-gap, t, t+gap, t+2*gap, t+3*gap]
-    
-    Boundary handling:
-    - Early frames: backward indices clamp to 0 (e.g., [0,0,0,5,35,65,95,125])
-    - Late frames: forward indices clamp to last frame (e.g., [850,880,910,940,970,1000,1000,1000])
-    
-    This enables truly uniform sampling across entire episodes.
-    
-    Args:
-        dataset_from_index: Start indices of episodes (global dataset indices)
-        dataset_to_index: End indices of episodes (global dataset indices)
-        frame_gap: Gap between consecutive frames (default: 30 = 1 second at 30fps)
-        shuffle: Whether to shuffle sampling order
-        seed: Random seed for reproducibility
-        samples_per_epoch: Number of samples per epoch (default: 6400)
-        min_episode_length: Minimum episode length to include (default: 1)
-    """
-    
-    def __init__(
-        self,
-        dataset_from_index: np.ndarray,
-        dataset_to_index: np.ndarray,
-        frame_gap: int = 30,
-        shuffle: bool = True,
-        seed: Optional[int] = None,
-        samples_per_epoch: int = 6400,
-        min_episode_length: int = 1,
-    ):
-        self.dataset_from_index = np.array(dataset_from_index)
-        self.dataset_to_index = np.array(dataset_to_index)
-        self.frame_gap = frame_gap
-        self.shuffle = shuffle
-        self.samples_per_epoch = samples_per_epoch
-        self.min_episode_length = min_episode_length
-        
-        if seed is not None:
-            self.seed = seed
-            random.seed(seed)
-            np.random.seed(seed)
-            self.generator = torch.Generator().manual_seed(seed)
-        else:
-            self.generator = torch.Generator()
-        
-        # Compute valid episodes and sampling positions (ALL frames for uniform sampling)
-        self._compute_valid_positions()
-        
-        logging.info(
-            f"SARMTemporalSampler: {len(self.valid_episodes)} valid episodes, "
-            f"{len(self.all_valid_positions)} positions (uniform sampling), "
-            f"{self.samples_per_epoch} samples per epoch, "
-            f"frame_gap={frame_gap}, symmetric bidirectional pattern"
-        )
-    
-    def _compute_valid_positions(self):
-        """Compute valid episodes and ALL sampling positions for uniform sampling.
-        
-        With symmetric bidirectional sampling, we can sample from ANY frame:
-        - Early frames: backward indices clamp to first frame
-        - Late frames: forward indices clamp to last frame
-        """
-        self.valid_episodes = []
-        self.all_valid_positions = []
-        
-        for ep_idx in range(len(self.dataset_from_index)):
-            ep_start = self.dataset_from_index[ep_idx]
-            ep_end = self.dataset_to_index[ep_idx]
-            episode_length = ep_end - ep_start
-            
-            # Include all episodes with at least min_episode_length frames
-            if episode_length >= self.min_episode_length:
-                self.valid_episodes.append((ep_idx, ep_start, ep_end))
-                
-                # Include ALL positions in the episode (truly uniform sampling)
-                for pos in range(ep_start, ep_end):
-                    self.all_valid_positions.append(pos)
-        
-        self.valid_episodes = np.array(self.valid_episodes)
-        self.all_valid_positions = np.array(self.all_valid_positions)
-        
-        if len(self.all_valid_positions) == 0:
-            raise ValueError(
-                f"No valid sampling positions found! "
-                f"Check that episodes have at least {self.min_episode_length} frames."
-            )
-    
-    def __len__(self) -> int:
-        return self.samples_per_epoch
-    
-    def __iter__(self) -> Iterator[int]:
-        """
-        Yields global dataset indices for uniform sampling across episodes.
-        
-        Each yielded index represents the "current frame" position.
-        The dataset's observation_delta_indices then handles loading:
-        - Frame 0: Episode initial frame (via large negative delta clamping)
-        - Frames 1-8: Symmetric context around current frame (with boundary clamping)
-        
-        For early frames: backward indices clamp to first frame (progress ~0%)
-        For late frames: forward indices clamp to last frame (progress ~100%)
-        """
-        if self.shuffle:
-            # Randomly sample from all valid positions
-            for _ in range(self.samples_per_epoch):
-                idx = np.random.randint(0, len(self.all_valid_positions))
-                yield int(self.all_valid_positions[idx])
-        else:
-            # Sequential sampling with wrap-around
-            for i in range(self.samples_per_epoch):
-                idx = i % len(self.all_valid_positions)
-                yield int(self.all_valid_positions[idx])

src/lerobot/datasets/utils.py

@@ -28,7 +28,6 @@ import numpy as np
 import packaging.version
 import pandas
 import pandas as pd
-import pyarrow.dataset as pa_ds
 import pyarrow.parquet as pq
 import torch
 from datasets import Dataset
@@ -104,9 +103,7 @@ def update_chunk_file_indices(chunk_idx: int, file_idx: int, chunks_size: int) -
     return chunk_idx, file_idx
 
 
-def load_nested_dataset(
-    pq_dir: Path, features: datasets.Features | None = None, episodes: list[int] | None = None
-) -> Dataset:
+def load_nested_dataset(pq_dir: Path, features: datasets.Features | None = None) -> Dataset:
     """Find parquet files in provided directory {pq_dir}/chunk-xxx/file-xxx.parquet
     Convert parquet files to pyarrow memory mapped in a cache folder for efficient RAM usage
     Concatenate all pyarrow references to return HF Dataset format
@@ -114,26 +111,15 @@ def load_nested_dataset(
     Args:
         pq_dir: Directory containing parquet files
         features: Optional features schema to ensure consistent loading of complex types like images
-        episodes: Optional list of episode indices to filter. Uses PyArrow predicate pushdown for efficiency.
     """
     paths = sorted(pq_dir.glob("*/*.parquet"))
     if len(paths) == 0:
         raise FileNotFoundError(f"Provided directory does not contain any parquet file: {pq_dir}")
 
+    # TODO(rcadene): set num_proc to accelerate conversion to pyarrow
     with SuppressProgressBars():
-        # When no filtering needed, Dataset uses memory-mapped loading for efficiency
-        # PyArrow loads the entire dataset into memory
-        if episodes is None:
-            return Dataset.from_parquet([str(path) for path in paths], features=features)
-
-        arrow_dataset = pa_ds.dataset(paths, format="parquet")
-        filter_expr = pa_ds.field("episode_index").isin(episodes)
-        table = arrow_dataset.to_table(filter=filter_expr)
-
-        if features is not None:
-            table = table.cast(features.arrow_schema)
-
-        return Dataset(table)
+        datasets = Dataset.from_parquet([str(path) for path in paths], features=features)
+    return datasets
 
 
 def get_parquet_num_frames(parquet_path: str | Path) -> int:

src/lerobot/envs/configs.py

@@ -21,22 +21,7 @@ import draccus
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.robots import RobotConfig
 from lerobot.teleoperators.config import TeleoperatorConfig
-from lerobot.utils.constants import (
-    ACTION,
-    LIBERO_KEY_EEF_MAT,
-    LIBERO_KEY_EEF_POS,
-    LIBERO_KEY_EEF_QUAT,
-    LIBERO_KEY_GRIPPER_QPOS,
-    LIBERO_KEY_GRIPPER_QVEL,
-    LIBERO_KEY_JOINTS_POS,
-    LIBERO_KEY_JOINTS_VEL,
-    LIBERO_KEY_PIXELS_AGENTVIEW,
-    LIBERO_KEY_PIXELS_EYE_IN_HAND,
-    OBS_ENV_STATE,
-    OBS_IMAGE,
-    OBS_IMAGES,
-    OBS_STATE,
-)
+from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
 
 @dataclass
@@ -261,61 +246,28 @@ class LiberoEnv(EnvConfig):
     features_map: dict[str, str] = field(
         default_factory=lambda: {
             ACTION: ACTION,
-            LIBERO_KEY_EEF_POS: f"{OBS_STATE}.eef_pos",
-            LIBERO_KEY_EEF_QUAT: f"{OBS_STATE}.eef_quat",
-            LIBERO_KEY_EEF_MAT: f"{OBS_STATE}.eef_mat",
-            LIBERO_KEY_GRIPPER_QPOS: f"{OBS_STATE}.gripper_qpos",
-            LIBERO_KEY_GRIPPER_QVEL: f"{OBS_STATE}.gripper_qvel",
-            LIBERO_KEY_JOINTS_POS: f"{OBS_STATE}.joint_pos",
-            LIBERO_KEY_JOINTS_VEL: f"{OBS_STATE}.joint_vel",
-            LIBERO_KEY_PIXELS_AGENTVIEW: f"{OBS_IMAGES}.image",
-            LIBERO_KEY_PIXELS_EYE_IN_HAND: f"{OBS_IMAGES}.image2",
+            "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[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
+            self.features["pixels/agentview_image"] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
-            self.features[LIBERO_KEY_PIXELS_EYE_IN_HAND] = PolicyFeature(
+            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
         elif self.obs_type == "pixels_agent_pos":
-            self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
+            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(8,))
+            self.features["pixels/agentview_image"] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
-            self.features[LIBERO_KEY_PIXELS_EYE_IN_HAND] = PolicyFeature(
+            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
-            self.features[LIBERO_KEY_EEF_POS] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(3,),
-            )
-            self.features[LIBERO_KEY_EEF_QUAT] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(4,),
-            )
-            self.features[LIBERO_KEY_EEF_MAT] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(3, 3),
-            )
-            self.features[LIBERO_KEY_GRIPPER_QPOS] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(2,),
-            )
-            self.features[LIBERO_KEY_GRIPPER_QVEL] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(2,),
-            )
-            self.features[LIBERO_KEY_JOINTS_POS] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(7,),
-            )
-            self.features[LIBERO_KEY_JOINTS_VEL] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(7,),
-            )
         else:
             raise ValueError(f"Unsupported obs_type: {self.obs_type}")
 

src/lerobot/envs/factory.py

@@ -14,16 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import importlib
-from typing import Any
 
 import gymnasium as gym
 from gymnasium.envs.registration import registry as gym_registry
 
 from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv
-from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
-from lerobot.processor import ProcessorStep
-from lerobot.processor.env_processor import LiberoProcessorStep
-from lerobot.processor.pipeline import PolicyProcessorPipeline
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -37,60 +32,16 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
         raise ValueError(f"Policy type '{env_type}' is not available.")
 
 
-def make_env_pre_post_processors(
-    env_cfg: EnvConfig,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-]:
-    """
-    Create preprocessor and postprocessor pipelines for environment observations.
-
-    This function creates processor pipelines that transform raw environment
-    observations and actions. By default, it returns identity processors that do nothing.
-    For specific environments like LIBERO, it adds environment-specific processing steps.
-
-    Args:
-        env_cfg: The configuration of the environment.
-
-    Returns:
-        A tuple containing:
-            - preprocessor: Pipeline that processes environment observations
-            - postprocessor: Pipeline that processes environment outputs (currently identity)
-    """
-    # Preprocessor and Postprocessor steps are Identity for most environments
-    preprocessor_steps: list[ProcessorStep] = []
-    postprocessor_steps: list[ProcessorStep] = []
-
-    # For LIBERO environments, add the LiberoProcessorStep to preprocessor
-    if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
-        preprocessor_steps.append(LiberoProcessorStep())
-
-    preprocessor = PolicyProcessorPipeline(steps=preprocessor_steps)
-    postprocessor = PolicyProcessorPipeline(steps=postprocessor_steps)
-
-    return preprocessor, postprocessor
-
-
 def make_env(
-    cfg: EnvConfig | str,
-    n_envs: int = 1,
-    use_async_envs: bool = False,
-    hub_cache_dir: str | None = None,
-    trust_remote_code: bool = False,
+    cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False
 ) -> dict[str, dict[int, gym.vector.VectorEnv]]:
-    """Makes a gym vector environment according to the config or Hub reference.
+    """Makes a gym vector environment according to the config.
 
     Args:
-        cfg (EnvConfig | str): Either an `EnvConfig` object describing the environment to build locally,
-            or a Hugging Face Hub repository identifier (e.g. `"username/repo"`). In the latter case,
-            the repo must include a Python file (usually `env.py`).
+        cfg (EnvConfig): the config of the environment to instantiate.
         n_envs (int, optional): The number of parallelized env to return. Defaults to 1.
         use_async_envs (bool, optional): Whether to return an AsyncVectorEnv or a SyncVectorEnv. Defaults to
             False.
-        hub_cache_dir (str | None): Optional cache path for downloaded hub files.
-        trust_remote_code (bool): **Explicit consent** to execute remote code from the Hub.
-            Default False — must be set to True to import/exec hub `env.py`.
 
     Raises:
         ValueError: if n_envs < 1
@@ -103,21 +54,6 @@ def make_env(
             - For single-task environments: a single suite entry (cfg.type) with task_id=0.
 
     """
-    # if user passed a hub id string (e.g., "username/repo", "username/repo@main:env.py")
-    # simplified: only support hub-provided `make_env`
-    if isinstance(cfg, str):
-        # _download_hub_file will raise the same RuntimeError if trust_remote_code is False
-        repo_id, file_path, local_file, revision = _download_hub_file(cfg, trust_remote_code, hub_cache_dir)
-
-        # import and surface clear import errors
-        module = _import_hub_module(local_file, repo_id)
-
-        # call the hub-provided make_env
-        raw_result = _call_make_env(module, n_envs=n_envs, use_async_envs=use_async_envs)
-
-        # normalize the return into {suite: {task_id: vec_env}}
-        return _normalize_hub_result(raw_result)
-
     if n_envs < 1:
         raise ValueError("`n_envs` must be at least 1")
 

src/lerobot/envs/libero.py

@@ -28,6 +28,7 @@ import torch
 from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv
+from robosuite.utils.transform_utils import quat2axisangle
 
 
 def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
@@ -174,36 +175,11 @@ class LiberoEnv(gym.Env):
             self.observation_space = spaces.Dict(
                 {
                     "pixels": spaces.Dict(images),
-                    "robot_state": spaces.Dict(
-                        {
-                            "eef": spaces.Dict(
-                                {
-                                    "pos": spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float64),
-                                    "quat": spaces.Box(
-                                        low=-np.inf, high=np.inf, shape=(4,), dtype=np.float64
-                                    ),
-                                    "mat": spaces.Box(
-                                        low=-np.inf, high=np.inf, shape=(3, 3), dtype=np.float64
-                                    ),
-                                }
-                            ),
-                            "gripper": spaces.Dict(
-                                {
-                                    "qpos": spaces.Box(
-                                        low=-np.inf, high=np.inf, shape=(2,), dtype=np.float64
-                                    ),
-                                    "qvel": spaces.Box(
-                                        low=-np.inf, high=np.inf, shape=(2,), dtype=np.float64
-                                    ),
-                                }
-                            ),
-                            "joints": spaces.Dict(
-                                {
-                                    "pos": spaces.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float64),
-                                    "vel": spaces.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float64),
-                                }
-                            ),
-                        }
+                    "agent_pos": spaces.Box(
+                        low=AGENT_POS_LOW,
+                        high=AGENT_POS_HIGH,
+                        shape=(OBS_STATE_DIM,),
+                        dtype=np.float64,
                     ),
                 }
             )
@@ -215,7 +191,6 @@ class LiberoEnv(gym.Env):
     def render(self):
         raw_obs = self._env.env._get_observations()
         image = self._format_raw_obs(raw_obs)["pixels"]["image"]
-        image = image[::-1, ::-1]  # flip both H and W for visualization
         return image
 
     def _make_envs_task(self, task_suite: Any, task_id: int = 0):
@@ -237,48 +212,23 @@ class LiberoEnv(gym.Env):
         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
-
-        eef_pos = raw_obs.get("robot0_eef_pos")
-        eef_quat = raw_obs.get("robot0_eef_quat")
-
-        # rotation matrix from controller
-        eef_mat = self._env.robots[0].controller.ee_ori_mat if eef_pos is not None else None
-        gripper_qpos = raw_obs.get("robot0_gripper_qpos")
-        gripper_qvel = raw_obs.get("robot0_gripper_qvel")
-        joint_pos = raw_obs.get("robot0_joint_pos")
-        joint_vel = raw_obs.get("robot0_joint_vel")
-        obs = {
-            "pixels": images,
-            "robot_state": {
-                "eef": {
-                    "pos": eef_pos,  # (3,)
-                    "quat": eef_quat,  # (4,)
-                    "mat": eef_mat,  # (3, 3)
-                },
-                "gripper": {
-                    "qpos": gripper_qpos,  # (2,)
-                    "qvel": gripper_qvel,  # (2,)
-                },
-                "joints": {
-                    "pos": joint_pos,  # (7,)
-                    "vel": joint_vel,  # (7,)
-                },
-            },
-        }
+        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":
-            # Validate required fields are present
-            if eef_pos is None or eef_quat is None or gripper_qpos is None:
-                raise ValueError(
-                    f"Missing required robot state fields in raw observation. "
-                    f"Got eef_pos={eef_pos is not None}, eef_quat={eef_quat is not None}, "
-                    f"gripper_qpos={gripper_qpos is not None}"
-                )
-            return obs
-
+            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')."
@@ -405,10 +355,12 @@ def create_libero_envs(
         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}).")
 

src/lerobot/envs/utils.py

@@ -13,8 +13,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import importlib.util
-import os
 import warnings
 from collections.abc import Mapping, Sequence
 from functools import singledispatch
@@ -24,27 +22,14 @@ import einops
 import gymnasium as gym
 import numpy as np
 import torch
-from huggingface_hub import hf_hub_download, snapshot_download
 from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.envs.configs import EnvConfig
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
+from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.utils.utils import get_channel_first_image_shape
 
 
-def _convert_nested_dict(d):
-    result = {}
-    for k, v in d.items():
-        if isinstance(v, dict):
-            result[k] = _convert_nested_dict(v)
-        elif isinstance(v, np.ndarray):
-            result[k] = torch.from_numpy(v)
-        else:
-            result[k] = v
-    return result
-
-
 def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]:
     # TODO(aliberts, rcadene): refactor this to use features from the environment (no hardcoding)
     """Convert environment observation to LeRobot format observation.
@@ -90,14 +75,12 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
 
         return_observations[OBS_ENV_STATE] = env_state
 
-    if "agent_pos" in observations:
-        agent_pos = torch.from_numpy(observations["agent_pos"]).float()
-        if agent_pos.dim() == 1:
-            agent_pos = agent_pos.unsqueeze(0)
-        return_observations[OBS_STATE] = agent_pos
+    # TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
+    agent_pos = torch.from_numpy(observations["agent_pos"]).float()
+    if agent_pos.dim() == 1:
+        agent_pos = agent_pos.unsqueeze(0)
+    return_observations[OBS_STATE] = agent_pos
 
-    if "robot_state" in observations:
-        return_observations[f"{OBS_STR}.robot_state"] = _convert_nested_dict(observations["robot_state"])
     return return_observations
 
 
@@ -212,132 +195,3 @@ def _(envs: Sequence) -> None:
 @close_envs.register
 def _(env: gym.Env) -> None:
     _close_single_env(env)
-
-
-# helper to safely load a python file as a module
-def _load_module_from_path(path: str, module_name: str | None = None):
-    module_name = module_name or f"hub_env_{os.path.basename(path).replace('.', '_')}"
-    spec = importlib.util.spec_from_file_location(module_name, path)
-    if spec is None:
-        raise ImportError(f"Could not load module spec for {module_name} from {path}")
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-
-
-# helper to parse hub string (supports "user/repo", "user/repo@rev", optional path)
-# examples:
-#   "user/repo" -> will look for env.py at repo root
-#   "user/repo@main:envs/my_env.py" -> explicit revision and path
-def _parse_hub_url(hub_uri: str):
-    # very small parser: [repo_id][@revision][:path]
-    # repo_id is required (user/repo or org/repo)
-    revision = None
-    file_path = "env.py"
-    if "@" in hub_uri:
-        repo_and_rev, *rest = hub_uri.split(":", 1)
-        repo_id, rev = repo_and_rev.split("@", 1)
-        revision = rev
-        if rest:
-            file_path = rest[0]
-    else:
-        repo_id, *rest = hub_uri.split(":", 1)
-        if rest:
-            file_path = rest[0]
-    return repo_id, revision, file_path
-
-
-def _download_hub_file(
-    cfg_str: str,
-    trust_remote_code: bool,
-    hub_cache_dir: str | None,
-) -> tuple[str, str, str, str]:
-    """
-    Parse `cfg_str` (hub URL), enforce `trust_remote_code`, and return
-    (repo_id, file_path, local_file, revision).
-    """
-    if not trust_remote_code:
-        raise RuntimeError(
-            f"Refusing to execute remote code from the Hub for '{cfg_str}'. "
-            "Executing hub env modules runs arbitrary Python code from third-party repositories. "
-            "If you trust this repo and understand the risks, call `make_env(..., trust_remote_code=True)` "
-            "and prefer pinning to a specific revision: 'user/repo@<commit-hash>:env.py'."
-        )
-
-    repo_id, revision, file_path = _parse_hub_url(cfg_str)
-
-    try:
-        local_file = hf_hub_download(
-            repo_id=repo_id, filename=file_path, revision=revision, cache_dir=hub_cache_dir
-        )
-    except Exception as e:
-        # fallback to snapshot download
-        snapshot_dir = snapshot_download(repo_id=repo_id, revision=revision, cache_dir=hub_cache_dir)
-        local_file = os.path.join(snapshot_dir, file_path)
-        if not os.path.exists(local_file):
-            raise FileNotFoundError(
-                f"Could not find {file_path} in repository {repo_id}@{revision or 'main'}"
-            ) from e
-
-    return repo_id, file_path, local_file, revision
-
-
-def _import_hub_module(local_file: str, repo_id: str) -> Any:
-    """
-    Import the downloaded file as a module and surface helpful import error messages.
-    """
-    module_name = f"hub_env_{repo_id.replace('/', '_')}"
-    try:
-        module = _load_module_from_path(local_file, module_name=module_name)
-    except ModuleNotFoundError as e:
-        missing = getattr(e, "name", None) or str(e)
-        raise ModuleNotFoundError(
-            f"Hub env '{repo_id}:{os.path.basename(local_file)}' failed to import because the dependency "
-            f"'{missing}' is not installed locally.\n\n"
-        ) from e
-    except ImportError as e:
-        raise ImportError(
-            f"Failed to load hub env module '{repo_id}:{os.path.basename(local_file)}'. Import error: {e}\n\n"
-        ) from e
-    return module
-
-
-def _call_make_env(module: Any, n_envs: int, use_async_envs: bool) -> Any:
-    """
-    Ensure module exposes make_env and call it.
-    """
-    if not hasattr(module, "make_env"):
-        raise AttributeError(
-            f"The hub module {getattr(module, '__name__', 'hub_module')} must expose `make_env(n_envs=int, use_async_envs=bool)`."
-        )
-    entry_fn = module.make_env
-    return entry_fn(n_envs=n_envs, use_async_envs=use_async_envs)
-
-
-def _normalize_hub_result(result: Any) -> dict[str, dict[int, gym.vector.VectorEnv]]:
-    """
-    Normalize possible return types from hub `make_env` into the mapping:
-      { suite_name: { task_id: vector_env } }
-    Accepts:
-      - dict (assumed already correct)
-      - gym.vector.VectorEnv
-      - gym.Env (will be wrapped into SyncVectorEnv)
-    """
-    if isinstance(result, dict):
-        return result
-
-    # VectorEnv: use its spec.id if available
-    if isinstance(result, gym.vector.VectorEnv):
-        suite_name = getattr(result, "spec", None) and getattr(result.spec, "id", None) or "hub_env"
-        return {suite_name: {0: result}}
-
-    # Single Env: wrap into SyncVectorEnv
-    if isinstance(result, gym.Env):
-        vec = gym.vector.SyncVectorEnv([lambda: result])
-        suite_name = getattr(result, "spec", None) and getattr(result.spec, "id", None) or "hub_env"
-        return {suite_name: {0: vec}}
-
-    raise ValueError(
-        "Hub `make_env` must return either a mapping {suite: {task_id: vec_env}}, "
-        "a gym.vector.VectorEnv, or a single gym.Env."
-    )

src/lerobot/motors/__init__.py

@@ -14,4 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .motors_bus import Motor, MotorCalibration, MotorNormMode, MotorsBus
+from .motors_bus import (
+    Motor,
+    MotorCalibration,
+    MotorNormMode,
+    MotorsBus,  # Backward compatibility (alias for SerialMotorsBus)
+    MotorsBusBase,
+    SerialMotorsBus,
+)

src/lerobot/motors/damiao/__init__.py

@@ -0,0 +1,18 @@
+#!/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 .damiao import DamiaoMotorsBus
+from .tables import *

src/lerobot/motors/damiao/damiao.py

@@ -0,0 +1,787 @@
+# 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.
+
+# TODO(pepijn): add license of: https://github.com/cmjang/DM_Control_Python?tab=MIT-1-ov-file#readme
+
+import logging
+import time
+from contextlib import contextmanager
+from copy import deepcopy
+from functools import cached_property
+from typing import Dict, List, Optional, Tuple, Union
+
+import can
+import numpy as np
+
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode, MotorsBusBase
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.utils.utils import enter_pressed, move_cursor_up
+
+from .tables import (
+    AVAILABLE_BAUDRATES,
+    CAN_CMD_DISABLE,
+    CAN_CMD_ENABLE,
+    CAN_CMD_REFRESH,
+    CAN_CMD_SET_ZERO,
+    CAN_PARAM_ID,
+    DEFAULT_BAUDRATE,
+    DEFAULT_TIMEOUT_MS,
+    MODEL_RESOLUTION,
+    MOTOR_LIMIT_PARAMS,
+    NORMALIZED_DATA,
+    MotorType,
+)
+
+logger = logging.getLogger(__name__)
+
+NameOrID = Union[str, int]
+Value = Union[int, float]
+
+
+class DamiaoMotorsBus(MotorsBusBase):
+    """
+    The Damiao implementation for a MotorsBus using CAN bus communication.
+    
+    This class uses python-can for CAN bus communication with Damiao motors.
+    For more info, see:
+    - python-can documentation: https://python-can.readthedocs.io/en/stable/
+    - Seedstudio documentation: https://wiki.seeedstudio.com/damiao_series/
+    - DM_Control_Python repo: https://github.com/cmjang/DM_Control_Python
+    """
+
+    # CAN-specific settings
+    available_baudrates = deepcopy(AVAILABLE_BAUDRATES)
+    default_baudrate = DEFAULT_BAUDRATE
+    default_timeout = DEFAULT_TIMEOUT_MS
+    
+    # Motor configuration
+    model_resolution_table = deepcopy(MODEL_RESOLUTION)
+    normalized_data = deepcopy(NORMALIZED_DATA)
+
+    def __init__(
+        self,
+        port: str,
+        motors: dict[str, Motor],
+        calibration: dict[str, MotorCalibration] | None = None,
+        can_interface: str = "auto",
+        use_can_fd: bool = True,
+        bitrate: int = 1000000,
+        data_bitrate: int | None = 5000000,
+    ):
+        """
+        Initialize the Damiao motors bus.
+
+        Args:
+            port: CAN interface name (e.g., "can0" for Linux, "/dev/cu.usbmodem*" for macOS)
+            motors: Dictionary mapping motor names to Motor objects
+            calibration: Optional calibration data
+            can_interface: CAN interface type - "auto" (default), "socketcan" (Linux), or "slcan" (macOS/serial)
+            use_can_fd: Whether to use CAN FD mode (default: True for OpenArms)
+            bitrate: Nominal bitrate in bps (default: 1000000 = 1 Mbps)
+            data_bitrate: Data bitrate for CAN FD in bps (default: 5000000 = 5 Mbps), ignored if use_can_fd is False
+        """
+        super().__init__(port, motors, calibration)
+        self.port = port
+        self.can_interface = can_interface
+        self.use_can_fd = use_can_fd
+        self.bitrate = bitrate
+        self.data_bitrate = data_bitrate
+        self.canbus = None
+        self._is_connected = False
+        
+        # Map motor names to CAN IDs
+        self._motor_can_ids = {}
+        self._recv_id_to_motor = {}
+        
+        # Store motor types and recv IDs
+        self._motor_types = {}
+        for name, motor in self.motors.items():
+            if hasattr(motor, "motor_type"):
+                self._motor_types[name] = motor.motor_type
+            else:
+                # Default to DM4310 if not specified
+                self._motor_types[name] = MotorType.DM4310
+            
+            # Map recv_id to motor name for filtering responses
+            if hasattr(motor, "recv_id"):
+                self._recv_id_to_motor[motor.recv_id] = name
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the CAN bus is connected."""
+        return self._is_connected and self.canbus is not None
+
+    def connect(self, handshake: bool = True) -> None:
+        """
+        Open the CAN bus and initialize communication.
+
+        Args:
+            handshake: If True, ping all motors to verify they're present
+        """
+        if self.is_connected:
+            raise DeviceAlreadyConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is already connected."
+            )
+
+        try:
+            # Auto-detect interface type based on port name
+            if self.can_interface == "auto":
+                if self.port.startswith("/dev/"):
+                    # Serial device (macOS/Windows)
+                    self.can_interface = "slcan"
+                    logger.info(f"Auto-detected slcan interface for port {self.port}")
+                else:
+                    # Network interface (Linux)
+                    self.can_interface = "socketcan"
+                    logger.info(f"Auto-detected socketcan interface for port {self.port}")
+            
+            # Connect to CAN bus
+            if self.can_interface == "socketcan":
+                # Linux SocketCAN with CAN FD support
+                if self.use_can_fd and self.data_bitrate is not None:
+                    self.canbus = can.interface.Bus(
+                        channel=self.port,
+                        interface="socketcan",
+                        bitrate=self.bitrate,
+                        data_bitrate=self.data_bitrate,
+                        fd=True
+                    )
+                    logger.info(f"Connected to {self.port} with CAN FD (bitrate={self.bitrate}, data_bitrate={self.data_bitrate})")
+                else:
+                    self.canbus = can.interface.Bus(
+                        channel=self.port,
+                        interface="socketcan",
+                        bitrate=self.bitrate
+                    )
+                    logger.info(f"Connected to {self.port} with CAN 2.0 (bitrate={self.bitrate})")
+            elif self.can_interface == "slcan":
+                # Serial Line CAN (macOS, Windows, or USB adapters)
+                # Note: SLCAN typically doesn't support CAN FD
+                self.canbus = can.interface.Bus(
+                    channel=self.port,
+                    interface="slcan",
+                    bitrate=self.bitrate
+                )
+                logger.info(f"Connected to {self.port} with SLCAN (bitrate={self.bitrate})")
+            else:
+                # Generic interface (vector, pcan, etc.)
+                if self.use_can_fd and self.data_bitrate is not None:
+                    self.canbus = can.interface.Bus(
+                        channel=self.port,
+                        interface=self.can_interface,
+                        bitrate=self.bitrate,
+                        data_bitrate=self.data_bitrate,
+                        fd=True
+                    )
+                else:
+                    self.canbus = can.interface.Bus(
+                        channel=self.port,
+                        interface=self.can_interface,
+                        bitrate=self.bitrate
+                    )
+            
+            self._is_connected = True
+            
+            if handshake:
+                self._handshake()
+                
+            logger.debug(f"{self.__class__.__name__} connected via {self.can_interface}.")
+        except Exception as e:
+            self._is_connected = False
+            raise ConnectionError(f"Failed to connect to CAN bus: {e}")
+
+    def _handshake(self) -> None:
+        """Verify all motors are present by refreshing their status."""
+        for motor_name in self.motors:
+            self._refresh_motor(motor_name)
+            time.sleep(0.01)  # Small delay between motors
+
+    def disconnect(self, disable_torque: bool = True) -> None:
+        """
+        Close the CAN bus connection.
+
+        Args:
+            disable_torque: If True, disable torque on all motors before disconnecting
+        """
+        if not self.is_connected:
+            raise DeviceNotConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is not connected."
+            )
+
+        if disable_torque:
+            try:
+                self.disable_torque()
+            except Exception as e:
+                logger.warning(f"Failed to disable torque during disconnect: {e}")
+
+        if self.canbus:
+            self.canbus.shutdown()
+            self.canbus = None
+        self._is_connected = False
+        logger.debug(f"{self.__class__.__name__} disconnected.")
+
+    def configure_motors(self) -> None:
+        """Configure all motors with default settings."""
+        # Damiao motors don't require much configuration in MIT mode
+        # Just ensure they're enabled
+        for motor in self.motors:
+            self._enable_motor(motor)
+            time.sleep(0.01)
+
+    def _enable_motor(self, motor: NameOrID) -> None:
+        """Enable a single motor."""
+        motor_id = self._get_motor_id(motor)
+        recv_id = self._get_motor_recv_id(motor)
+        data = [0xFF] * 7 + [CAN_CMD_ENABLE]
+        msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
+        self.canbus.send(msg)
+        self._recv_motor_response(expected_recv_id=recv_id)
+
+    def _disable_motor(self, motor: NameOrID) -> None:
+        """Disable a single motor."""
+        motor_id = self._get_motor_id(motor)
+        recv_id = self._get_motor_recv_id(motor)
+        data = [0xFF] * 7 + [CAN_CMD_DISABLE]
+        msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
+        self.canbus.send(msg)
+        self._recv_motor_response(expected_recv_id=recv_id)
+
+    def enable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        """Enable torque on selected motors."""
+        motors = self._get_motors_list(motors)
+        for motor in motors:
+            for _ in range(num_retry + 1):
+                try:
+                    self._enable_motor(motor)
+                    break
+                except Exception as e:
+                    if _ == num_retry:
+                        raise e
+                    time.sleep(0.01)
+
+    def disable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        """Disable torque on selected motors."""
+        motors = self._get_motors_list(motors)
+        for motor in motors:
+            for _ in range(num_retry + 1):
+                try:
+                    self._disable_motor(motor)
+                    break
+                except Exception as e:
+                    if _ == num_retry:
+                        raise e
+                    time.sleep(0.01)
+
+    @contextmanager
+    def torque_disabled(self, motors: str | list[str] | None = None):
+        """
+        Context manager that guarantees torque is re-enabled.
+        
+        This helper is useful to temporarily disable torque when configuring motors.
+        
+        Examples:
+            >>> with bus.torque_disabled():
+            ...     # Safe operations here with torque disabled
+            ...     pass
+        """
+        self.disable_torque(motors)
+        try:
+            yield
+        finally:
+            self.enable_torque(motors)
+
+    def set_zero_position(self, motors: str | list[str] | None = None) -> None:
+        """Set current position as zero for selected motors."""
+        motors = self._get_motors_list(motors)
+        for motor in motors:
+            motor_id = self._get_motor_id(motor)
+            recv_id = self._get_motor_recv_id(motor)
+            data = [0xFF] * 7 + [CAN_CMD_SET_ZERO]
+            msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
+            self.canbus.send(msg)
+            self._recv_motor_response(expected_recv_id=recv_id)
+            time.sleep(0.01)
+
+    def _refresh_motor(self, motor: NameOrID) -> Optional[can.Message]:
+        """Refresh motor status and return the response."""
+        motor_id = self._get_motor_id(motor)
+        recv_id = self._get_motor_recv_id(motor)
+        data = [motor_id & 0xFF, (motor_id >> 8) & 0xFF, CAN_CMD_REFRESH, 0, 0, 0, 0, 0]
+        msg = can.Message(arbitration_id=CAN_PARAM_ID, data=data, is_extended_id=False)
+        self.canbus.send(msg)
+        return self._recv_motor_response(expected_recv_id=recv_id)
+
+    def _recv_motor_response(self, expected_recv_id: Optional[int] = None, timeout: float = 0.001) -> Optional[can.Message]:
+        """
+        Receive a response from a motor.
+        
+        Args:
+            expected_recv_id: If provided, only return messages from this CAN ID
+            timeout: Timeout in seconds (default: 1ms for high-speed operation)
+        
+        Returns:
+            CAN message if received, None otherwise
+        """
+        try:
+            start_time = time.time()
+            messages_seen = []
+            while time.time() - start_time < timeout:
+                msg = self.canbus.recv(timeout=0.0001)  # 100us timeout for fast polling
+                if msg:
+                    messages_seen.append(f"0x{msg.arbitration_id:02X}")
+                    # If no filter specified, return any message
+                    if expected_recv_id is None:
+                        return msg
+                    # Otherwise, only return if it matches the expected recv_id
+                    if msg.arbitration_id == expected_recv_id:
+                        return msg
+                    else:
+                        logger.debug(f"Ignoring message from CAN ID 0x{msg.arbitration_id:02X}, expected 0x{expected_recv_id:02X}")
+            
+            # Only log warnings if we're in debug mode to reduce overhead
+            if logger.isEnabledFor(logging.DEBUG):
+                if messages_seen:
+                    logger.debug(f"Received {len(messages_seen)} message(s) from IDs {set(messages_seen)}, but expected 0x{expected_recv_id:02X}")
+                else:
+                    logger.debug(f"No CAN messages received (expected from 0x{expected_recv_id:02X})")
+        except Exception as e:
+            logger.debug(f"Failed to receive CAN message: {e}")
+        return None
+    
+    def _recv_all_responses(self, expected_recv_ids: list[int], timeout: float = 0.002) -> dict[int, can.Message]:
+        """
+        Efficiently receive responses from multiple motors at once.
+        Uses the OpenArms pattern: collect all available messages within timeout.
+        
+        Args:
+            expected_recv_ids: List of CAN IDs we expect responses from
+            timeout: Total timeout in seconds (default: 2ms)
+        
+        Returns:
+            Dictionary mapping recv_id to CAN message
+        """
+        responses = {}
+        expected_set = set(expected_recv_ids)
+        start_time = time.time()
+        
+        try:
+            while len(responses) < len(expected_recv_ids) and (time.time() - start_time) < timeout:
+                msg = self.canbus.recv(timeout=0.0001)  # 100us poll timeout
+                if msg and msg.arbitration_id in expected_set:
+                    responses[msg.arbitration_id] = msg
+                    if len(responses) == len(expected_recv_ids):
+                        break  # Got all responses, exit early
+        except Exception as e:
+            logger.debug(f"Error receiving responses: {e}")
+        
+        return responses
+
+    def _mit_control(
+        self,
+        motor: NameOrID,
+        kp: float,
+        kd: float,
+        position_degrees: float,
+        velocity_deg_per_sec: float,
+        torque: float,
+    ) -> None:
+        """
+        Send MIT control command to a motor.
+
+        Args:
+            motor: Motor name or ID
+            kp: Position gain
+            kd: Velocity gain
+            position_degrees: Target position (degrees)
+            velocity_deg_per_sec: Target velocity (degrees/s)
+            torque: Target torque (N·m)
+        """
+        motor_id = self._get_motor_id(motor)
+        motor_name = self._get_motor_name(motor)
+        motor_type = self._motor_types.get(motor_name, MotorType.DM4310)
+        
+        # Convert degrees to radians for motor control
+        position_rad = np.radians(position_degrees)
+        velocity_rad_per_sec = np.radians(velocity_deg_per_sec)
+        
+        # Get motor limits
+        pmax, vmax, tmax = MOTOR_LIMIT_PARAMS[motor_type]
+        
+        # Encode parameters
+        kp_uint = self._float_to_uint(kp, 0, 500, 12)
+        kd_uint = self._float_to_uint(kd, 0, 5, 12)
+        q_uint = self._float_to_uint(position_rad, -pmax, pmax, 16)
+        dq_uint = self._float_to_uint(velocity_rad_per_sec, -vmax, vmax, 12)
+        tau_uint = self._float_to_uint(torque, -tmax, tmax, 12)
+        
+        # Pack data
+        data = [0] * 8
+        data[0] = (q_uint >> 8) & 0xFF
+        data[1] = q_uint & 0xFF
+        data[2] = dq_uint >> 4
+        data[3] = ((dq_uint & 0xF) << 4) | ((kp_uint >> 8) & 0xF)
+        data[4] = kp_uint & 0xFF
+        data[5] = kd_uint >> 4
+        data[6] = ((kd_uint & 0xF) << 4) | ((tau_uint >> 8) & 0xF)
+        data[7] = tau_uint & 0xFF
+        
+        msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
+        self.canbus.send(msg)
+        recv_id = self._get_motor_recv_id(motor)
+        self._recv_motor_response(expected_recv_id=recv_id)
+
+    def _float_to_uint(self, x: float, x_min: float, x_max: float, bits: int) -> int:
+        """Convert float to unsigned integer for CAN transmission."""
+        x = max(x_min, min(x_max, x))  # Clamp to range
+        span = x_max - x_min
+        data_norm = (x - x_min) / span
+        return int(data_norm * ((1 << bits) - 1))
+
+    def _uint_to_float(self, x: int, x_min: float, x_max: float, bits: int) -> float:
+        """Convert unsigned integer from CAN to float."""
+        span = x_max - x_min
+        data_norm = float(x) / ((1 << bits) - 1)
+        return data_norm * span + x_min
+
+    def _decode_motor_state(self, data: bytes, motor_type: MotorType) -> Tuple[float, float, float, int, int]:
+        """
+        Decode motor state from CAN data.
+
+        Returns:
+            Tuple of (position_degrees, velocity_deg_per_sec, torque, temp_mos, temp_rotor)
+        """
+        if len(data) < 8:
+            raise ValueError("Invalid motor state data")
+        
+        # Extract encoded values
+        q_uint = (data[1] << 8) | data[2]
+        dq_uint = (data[3] << 4) | (data[4] >> 4)
+        tau_uint = ((data[4] & 0x0F) << 8) | data[5]
+        t_mos = data[6]
+        t_rotor = data[7]
+        
+        # Get motor limits
+        pmax, vmax, tmax = MOTOR_LIMIT_PARAMS[motor_type]
+        
+        # Decode to physical values (radians)
+        position_rad = self._uint_to_float(q_uint, -pmax, pmax, 16)
+        velocity_rad_per_sec = self._uint_to_float(dq_uint, -vmax, vmax, 12)
+        torque = self._uint_to_float(tau_uint, -tmax, tmax, 12)
+        
+        # Convert to degrees
+        position_degrees = np.degrees(position_rad)
+        velocity_deg_per_sec = np.degrees(velocity_rad_per_sec)
+        
+        return position_degrees, velocity_deg_per_sec, torque, t_mos, t_rotor
+
+    def read(
+        self,
+        data_name: str,
+        motor: str,
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> Value:
+        """Read a value from a single motor. Positions are always in degrees."""
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        # Refresh motor to get latest state
+        msg = self._refresh_motor(motor)
+        if msg is None:
+            motor_id = self._get_motor_id(motor)
+            recv_id = self._get_motor_recv_id(motor)
+            raise ConnectionError(
+                f"No response from motor '{motor}' (send ID: 0x{motor_id:02X}, recv ID: 0x{recv_id:02X}). "
+                f"Check that: 1) Motor is powered (24V), 2) CAN wiring is correct, "
+                f"3) Motor IDs are configured correctly using Damiao Debugging Tools"
+            )
+        
+        motor_type = self._motor_types.get(motor, MotorType.DM4310)
+        position_degrees, velocity_deg_per_sec, torque, t_mos, t_rotor = self._decode_motor_state(msg.data, motor_type)
+        
+        # Return requested data (already in degrees for position/velocity)
+        if data_name == "Present_Position":
+            value = position_degrees
+        elif data_name == "Present_Velocity":
+            value = velocity_deg_per_sec
+        elif data_name == "Present_Torque":
+            value = torque
+        elif data_name == "Temperature_MOS":
+            value = t_mos
+        elif data_name == "Temperature_Rotor":
+            value = t_rotor
+        else:
+            raise ValueError(f"Unknown data_name: {data_name}")
+        
+        # For Damiao, positions are always in degrees, no normalization needed
+        # We keep the normalize parameter for compatibility but don't use it
+        return value
+
+    def write(
+        self,
+        data_name: str,
+        motor: str,
+        value: Value,
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> None:
+        """Write a value to a single motor. Positions are always in degrees."""
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        # Value is expected to be in degrees for positions
+        if data_name == "Goal_Position":
+            # Use MIT control with position in degrees
+            self._mit_control(motor, 10.0, 0.5, value, 0, 0)
+        else:
+            raise ValueError(f"Writing {data_name} not supported in MIT mode")
+
+    def sync_read(
+        self,
+        data_name: str,
+        motors: str | list[str] | None = None,
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> Dict[str, Value]:
+        """
+        Read the same value from multiple motors simultaneously.
+        Uses batched operations: sends all refresh commands, then collects all responses.
+        This is MUCH faster than sequential reads (OpenArms pattern).
+        """
+        motors = self._get_motors_list(motors)
+        result = {}
+
+        # Step 1: Send refresh commands to ALL motors first (no waiting)
+        for motor in motors:
+            motor_id = self._get_motor_id(motor)
+            data = [motor_id & 0xFF, (motor_id >> 8) & 0xFF, CAN_CMD_REFRESH, 0, 0, 0, 0, 0]
+            msg = can.Message(arbitration_id=CAN_PARAM_ID, data=data, is_extended_id=False)
+            self.canbus.send(msg)
+        
+        # Step 2: Collect all responses at once (batch receive)
+        expected_recv_ids = [self._get_motor_recv_id(motor) for motor in motors]
+        responses = self._recv_all_responses(expected_recv_ids, timeout=0.003)  # 3ms total timeout
+        
+        # Step 3: Parse responses
+        for motor in motors:
+            try:
+                recv_id = self._get_motor_recv_id(motor)
+                msg = responses.get(recv_id)
+                
+                if msg is None:
+                    logger.warning(f"No response from motor '{motor}' (recv ID: 0x{recv_id:02X})")
+                    result[motor] = 0.0
+                    continue
+                
+                motor_type = self._motor_types.get(motor, MotorType.DM4310)
+                position_degrees, velocity_deg_per_sec, torque, t_mos, t_rotor = self._decode_motor_state(msg.data, motor_type)
+                
+                # Return requested data
+                if data_name == "Present_Position":
+                    value = position_degrees
+                elif data_name == "Present_Velocity":
+                    value = velocity_deg_per_sec
+                elif data_name == "Present_Torque":
+                    value = torque
+                elif data_name == "Temperature_MOS":
+                    value = t_mos
+                elif data_name == "Temperature_Rotor":
+                    value = t_rotor
+                else:
+                    raise ValueError(f"Unknown data_name: {data_name}")
+                
+                result[motor] = value
+                
+            except Exception as e:
+                logger.warning(f"Failed to read {data_name} from {motor}: {e}")
+                result[motor] = 0.0
+        
+        return result
+
+    def sync_write(
+        self,
+        data_name: str,
+        values: Dict[str, Value],
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> None:
+        """
+        Write different values to multiple motors simultaneously. Positions are always in degrees.
+        Uses batched operations: sends all commands first, then collects responses (OpenArms pattern).
+        """
+        if data_name == "Goal_Position":
+            # Step 1: Send all MIT control commands first (no waiting)
+            for motor, value_degrees in values.items():
+                motor_id = self._get_motor_id(motor)
+                motor_name = self._get_motor_name(motor)
+                motor_type = self._motor_types.get(motor_name, MotorType.DM4310)
+                
+                # Convert degrees to radians
+                position_rad = np.radians(value_degrees)
+                
+                # Default gains for position control
+                kp, kd = 10.0, 0.5
+                
+                # Get motor limits and encode parameters
+                pmax, vmax, tmax = MOTOR_LIMIT_PARAMS[motor_type]
+                kp_uint = self._float_to_uint(kp, 0, 500, 12)
+                kd_uint = self._float_to_uint(kd, 0, 5, 12)
+                q_uint = self._float_to_uint(position_rad, -pmax, pmax, 16)
+                dq_uint = self._float_to_uint(0, -vmax, vmax, 12)
+                tau_uint = self._float_to_uint(0, -tmax, tmax, 12)
+                
+                # Pack data
+                data = [0] * 8
+                data[0] = (q_uint >> 8) & 0xFF
+                data[1] = q_uint & 0xFF
+                data[2] = dq_uint >> 4
+                data[3] = ((dq_uint & 0xF) << 4) | ((kp_uint >> 8) & 0xF)
+                data[4] = kp_uint & 0xFF
+                data[5] = kd_uint >> 4
+                data[6] = ((kd_uint & 0xF) << 4) | ((tau_uint >> 8) & 0xF)
+                data[7] = tau_uint & 0xFF
+                
+                msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
+                self.canbus.send(msg)
+            
+            # Step 2: Collect all responses at once
+            expected_recv_ids = [self._get_motor_recv_id(motor) for motor in values.keys()]
+            self._recv_all_responses(expected_recv_ids, timeout=0.002)  # 2ms timeout
+        else:
+            # Fall back to individual writes for other data types
+            for motor, value in values.items():
+                self.write(data_name, motor, value, normalize=normalize, num_retry=num_retry)
+
+    def read_calibration(self) -> dict[str, MotorCalibration]:
+        """Read calibration data from motors."""
+        # Damiao motors don't store calibration internally
+        # Return existing calibration or empty dict
+        return self.calibration if self.calibration else {}
+
+    def write_calibration(self, calibration_dict: dict[str, MotorCalibration], cache: bool = True) -> None:
+        """Write calibration data to motors."""
+        # Damiao motors don't store calibration internally
+        # Just cache it in memory
+        if cache:
+            self.calibration = calibration_dict
+
+    def record_ranges_of_motion(
+        self, motors: NameOrID | list[NameOrID] | None = None, display_values: bool = True
+    ) -> tuple[dict[NameOrID, Value], dict[NameOrID, Value]]:
+        """
+        Interactively record the min/max values of each motor in degrees.
+
+        Move the joints by hand (with torque disabled) while the method streams live positions.
+        Press Enter to finish.
+        """
+        if motors is None:
+            motors = list(self.motors.keys())
+        elif isinstance(motors, (str, int)):
+            motors = [motors]
+
+        # Disable torque for manual movement
+        self.disable_torque(motors)
+        time.sleep(0.1)
+
+        # Get initial positions (already in degrees)
+        start_positions = self.sync_read("Present_Position", motors, normalize=False)
+        mins = start_positions.copy()
+        maxes = start_positions.copy()
+
+        print("\nMove joints through their full range of motion. Press ENTER when done.")
+        user_pressed_enter = False
+        
+        while not user_pressed_enter:
+            positions = self.sync_read("Present_Position", motors, normalize=False)
+            
+            for motor in motors:
+                if motor in positions:
+                    mins[motor] = min(positions[motor], mins.get(motor, positions[motor]))
+                    maxes[motor] = max(positions[motor], maxes.get(motor, positions[motor]))
+
+            if display_values:
+                print("\n" + "=" * 50)
+                print(f"{'MOTOR':<20} | {'MIN (deg)':>12} | {'POS (deg)':>12} | {'MAX (deg)':>12}")
+                print("-" * 50)
+                for motor in motors:
+                    if motor in positions:
+                        print(f"{motor:<20} | {mins[motor]:>12.1f} | {positions[motor]:>12.1f} | {maxes[motor]:>12.1f}")
+
+            if enter_pressed():
+                user_pressed_enter = True
+
+            if display_values and not user_pressed_enter:
+                # Move cursor up to overwrite the previous output
+                move_cursor_up(len(motors) + 4)
+
+            time.sleep(0.05)
+
+        # Re-enable torque
+        self.enable_torque(motors)
+        
+        # Validate ranges
+        for motor in motors:
+            if motor in mins and motor in maxes:
+                if abs(maxes[motor] - mins[motor]) < 5.0:  # At least 5 degrees of range
+                    raise ValueError(f"Motor {motor} has insufficient range of motion (< 5 degrees)")
+
+        return mins, maxes
+
+    def _get_motors_list(self, motors: str | list[str] | None) -> list[str]:
+        """Convert motor specification to list of motor names."""
+        if motors is None:
+            return list(self.motors.keys())
+        elif isinstance(motors, str):
+            return [motors]
+        elif isinstance(motors, list):
+            return motors
+        else:
+            raise TypeError(f"Invalid motors type: {type(motors)}")
+
+    def _get_motor_id(self, motor: NameOrID) -> int:
+        """Get CAN ID for a motor."""
+        if isinstance(motor, str):
+            if motor in self.motors:
+                return self.motors[motor].id
+            else:
+                raise ValueError(f"Unknown motor: {motor}")
+        else:
+            return motor
+
+    def _get_motor_name(self, motor: NameOrID) -> str:
+        """Get motor name from name or ID."""
+        if isinstance(motor, str):
+            return motor
+        else:
+            for name, m in self.motors.items():
+                if m.id == motor:
+                    return name
+            raise ValueError(f"Unknown motor ID: {motor}")
+    
+    def _get_motor_recv_id(self, motor: NameOrID) -> Optional[int]:
+        """Get motor recv_id from name or ID."""
+        motor_name = self._get_motor_name(motor)
+        motor_obj = self.motors.get(motor_name)
+        if motor_obj and hasattr(motor_obj, "recv_id"):
+            return motor_obj.recv_id
+        return None
+
+    @cached_property
+    def is_calibrated(self) -> bool:
+        """Check if motors are calibrated."""
+        return bool(self.calibration)

src/lerobot/motors/damiao/seed_damiao.py

@@ -0,0 +1,833 @@
+## This is a derivative of the following software.
+## https://github.com/cmjang/DM_Control_Python/blob/main/DM_CAN.py
+
+import can
+from time import sleep, time
+import numpy as np
+from enum import IntEnum
+from struct import unpack
+from struct import pack
+
+class Motor:
+    def __init__(self, MotorType, SlaveID, MasterID):
+        """
+        define Motor object 定义电机对象
+        :param MotorType: Motor type 电机类型
+        :param SlaveID: CANID 电机ID
+        :param MasterID: MasterID 主机ID 建议不要设为0
+        """
+        self.Pd = float(0)
+        self.Vd = float(0)
+        self.goal_position = float(0)
+        self.goal_tau = float(0)
+        self.state_q = float(0)
+        self.state_dq = float(0)
+        self.state_tau = float(0)
+        self.state_tmos = int(0)
+        self.state_trotor = int(0)
+        self.SlaveID = SlaveID
+        self.MasterID = MasterID
+        self.MotorType = MotorType
+        self.isEnable = False
+        self.NowControlMode = Control_Type.MIT
+        self.temp_param_dict = {}
+
+    def recv_data(self, q: float, dq: float, tau: float, tmos: int, trotor: int):
+        self.state_q = q
+        self.state_dq = dq
+        self.state_tau = tau
+        self.state_tmos = tmos
+        self.state_trotor = trotor
+
+    def getPosition(self):
+        """
+        get the position of the motor 获取电机位置
+        :return: the position of the motor 电机位置
+        """
+        return self.state_q
+
+    def getVelocity(self):
+        """
+        get the velocity of the motor 获取电机速度
+        :return: the velocity of the motor 电机速度
+        """
+        return self.state_dq
+
+    def getTorque(self):
+        """
+        get the torque of the motor 获取电机力矩
+        :return: the torque of the motor 电机力矩
+        """
+        return self.state_tau
+
+    def getParam(self, RID):
+        """
+        get the parameter of the motor 获取电机内部的参数，需要提前读取
+        :param RID: DM_variable 电机参数
+        :return: the parameter of the motor 电机参数
+        """
+        if RID in self.temp_param_dict:
+            return self.temp_param_dict[RID]
+        else:
+            return None
+
+
+class MotorControl:
+    #send_data_frame = np.array(
+    #    [0x55, 0xAA, 0x1e, 0x03, 0x01, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x00, 0, 0, 0, 0, 0x00, 0x08, 0x00,
+    #     0x00, 0, 0, 0, 0, 0, 0, 0, 0, 0x00], np.uint8)
+    #                4310           4310_48        4340           4340_48
+    Limit_Param = [[12.5, 30, 10], [12.5, 50, 10], [12.5, 8, 28], [12.5, 10, 28],
+                   # 6006           8006           8009            10010L         10010
+                   [12.5, 45, 20], [12.5, 45, 40], [12.5, 45, 54], [12.5, 25, 200], [12.5, 20, 200],
+                   # H3510            DMG62150      DMH6220
+                   [12.5 , 280 , 1],[12.5 , 45 , 10],[12.5 , 45 , 10]]
+
+    def __init__(self, channel: str, bitrate: int = 1000000):
+        """
+        define MotorControl object 定义电机控制对象
+        :param serial_device: serial object 串口对象
+        """
+        #self.serial_ = serial_device
+        self.motors_map = dict()
+        self.data_save = bytes()  # save data
+        #if self.serial_.is_open:  # open the serial port
+        #    print("Serial port is open")
+        #    serial_device.close()
+        #self.serial_.open()
+        self.canbus = can.interface.Bus(channel=channel, interface='socketcan', bitrate=bitrate)
+
+        #print("can is open")
+        
+
+    def controlMIT(self, DM_Motor, kp: float, kd: float, q: float, dq: float, tau: float):
+        """
+        MIT Control Mode Function 达妙电机MIT控制模式函数
+        :param DM_Motor: Motor object 电机对象
+        :param kp: kp
+        :param kd:  kd
+        :param q:  position  期望位置
+        :param dq:  velocity  期望速度
+        :param tau: torque  期望力矩
+        :return: None
+        """
+        if DM_Motor.SlaveID not in self.motors_map:
+            print("controlMIT ERROR : Motor ID not found")
+            return
+        kp_uint = float_to_uint(kp, 0, 500, 12)
+        kd_uint = float_to_uint(kd, 0, 5, 12)
+        MotorType = DM_Motor.MotorType
+        Q_MAX = self.Limit_Param[MotorType][0]
+        DQ_MAX = self.Limit_Param[MotorType][1]
+        TAU_MAX = self.Limit_Param[MotorType][2]
+        q_uint = float_to_uint(q, -Q_MAX, Q_MAX, 16)
+        dq_uint = float_to_uint(dq, -DQ_MAX, DQ_MAX, 12)
+        tau_uint = float_to_uint(tau, -TAU_MAX, TAU_MAX, 12)
+        data_buf = np.array([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00], np.uint8)
+        data_buf[0] = (q_uint >> 8) & 0xff
+        data_buf[1] = q_uint & 0xff
+        data_buf[2] = dq_uint >> 4
+        data_buf[3] = ((dq_uint & 0xf) << 4) | ((kp_uint >> 8) & 0xf)
+        data_buf[4] = kp_uint & 0xff
+        data_buf[5] = kd_uint >> 4
+        data_buf[6] = ((kd_uint & 0xf) << 4) | ((tau_uint >> 8) & 0xf)
+        data_buf[7] = tau_uint & 0xff
+        self.__send_data(DM_Motor.SlaveID, data_buf)
+        self.recv()  # receive the data from serial port
+
+    def control_delay(self, DM_Motor, kp: float, kd: float, q: float, dq: float, tau: float, delay: float):
+        """
+        MIT Control Mode Function with delay 达妙电机MIT控制模式函数带延迟
+        :param DM_Motor: Motor object 电机对象
+        :param kp: kp
+        :param kd: kd
+        :param q:  position  期望位置
+        :param dq:  velocity  期望速度
+        :param tau: torque  期望力矩
+        :param delay: delay time 延迟时间 单位秒
+        """
+        self.controlMIT(DM_Motor, kp, kd, q, dq, tau)
+        sleep(delay)
+
+    def control_Pos_Vel(self, Motor, P_desired: float, V_desired: float):
+        """
+        control the motor in position and velocity control mode 电机位置速度控制模式
+        :param Motor: Motor object 电机对象
+        :param P_desired: desired position 期望位置
+        :param V_desired: desired velocity 期望速度
+        :return: None
+        """
+        if Motor.SlaveID not in self.motors_map:
+            print("Control Pos_Vel Error : Motor ID not found")
+            return
+        motorid = 0x100 + Motor.SlaveID
+        data_buf = np.array([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00], np.uint8)
+        P_desired_uint8s = float_to_uint8s(P_desired)
+        V_desired_uint8s = float_to_uint8s(V_desired)
+        data_buf[0:4] = P_desired_uint8s
+        data_buf[4:8] = V_desired_uint8s
+        self.__send_data(motorid, data_buf)
+        self.recv()  # receive the data from serial port
+
+    def control_Vel(self, Motor, Vel_desired):
+        """
+        control the motor in velocity control mode 电机速度控制模式
+        :param Motor: Motor object 电机对象
+        :param Vel_desired: desired velocity 期望速度
+        """
+        if Motor.SlaveID not in self.motors_map:
+            print("control_VEL ERROR : Motor ID not found")
+            return
+        motorid = 0x200 + Motor.SlaveID
+        data_buf = np.array([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00], np.uint8)
+        Vel_desired_uint8s = float_to_uint8s(Vel_desired)
+        data_buf[0:4] = Vel_desired_uint8s
+        self.__send_data(motorid, data_buf)
+        self.recv()  # receive the data from serial port
+
+    def control_pos_force(self, Motor, Pos_des: float, Vel_des, i_des):
+        """
+        control the motor in EMIT control mode 电机力位混合模式
+        :param Pos_des: desired position rad  期望位置 单位为rad
+        :param Vel_des: desired velocity rad/s  期望速度 为放大100倍
+        :param i_des: desired current rang 0-10000 期望电流标幺值放大10000倍
+        电流标幺值：实际电流值除以最大电流值，最大电流见上电打印
+        """
+        if Motor.SlaveID not in self.motors_map:
+            print("control_pos_vel ERROR : Motor ID not found")
+            return
+        motorid = 0x300 + Motor.SlaveID
+        data_buf = np.array([0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00], np.uint8)
+        Pos_desired_uint8s = float_to_uint8s(Pos_des)
+        data_buf[0:4] = Pos_desired_uint8s
+        Vel_uint = np.uint16(Vel_des)
+        ides_uint = np.uint16(i_des)
+        data_buf[4] = Vel_uint & 0xff
+        data_buf[5] = Vel_uint >> 8
+        data_buf[6] = ides_uint & 0xff
+        data_buf[7] = ides_uint >> 8
+        self.__send_data(motorid, data_buf)
+        self.recv()  # receive the data from serial port
+
+    def enable(self, Motor):
+        """
+        enable motor 使能电机
+        最好在上电后几秒后再使能电机
+        :param Motor: Motor object 电机对象
+        """
+        self.__control_cmd(Motor, np.uint8(0xFC))
+        sleep(0.1)
+        self.recv()  # receive the data from serial port
+
+    def enable_old(self, Motor ,ControlMode):
+        """
+        enable motor old firmware 使能电机旧版本固件，这个是为了旧版本电机固件的兼容性
+        可恶的旧版本固件使能需要加上偏移量
+        最好在上电后几秒后再使能电机
+        :param Motor: Motor object 电机对象
+        """
+        data_buf = np.array([0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xfc], np.uint8)
+        enable_id = ((int(ControlMode)-1) << 2) + Motor.SlaveID
+        self.__send_data(enable_id, data_buf)
+        sleep(0.1)
+        self.recv()  # receive the data from serial port
+
+    def disable(self, Motor):
+        """
+        disable motor 失能电机
+        :param Motor: Motor object 电机对象
+        """
+        self.__control_cmd(Motor, np.uint8(0xFD))
+        sleep(0.1)
+        self.recv()  # receive the data from serial port
+
+    def set_zero_position(self, Motor):
+        """
+        set the zero position of the motor 设置电机0位
+        :param Motor: Motor object 电机对象
+        """
+        self.__control_cmd(Motor, np.uint8(0xFE))
+        sleep(0.1)
+        self.recv()  # receive the data from serial port
+
+    def recv(self):
+        # 把上次没有解析完的剩下的也放进来
+        # data_recv = b''.join([self.data_save, self.serial_.read_all()])
+        #data_recv = b''.join([self.data_save, self.canbus.recv()])
+
+
+        # packets = self.__extract_packets(data_recv)
+        # for packet in packets:
+        #     data = packet[7:15]
+        #     CANID = (packet[6] << 24) | (packet[5] << 16) | (packet[4] << 8) | packet[3]
+        #     CMD = packet[1]
+        #     self.__process_packet(data, CANID, CMD)
+        
+        data_recv = self.canbus.recv(0.1)
+
+        if data_recv is not None:
+            # data = data_recv.data
+            # err = data[0] >> 12
+            # id = data[0] & 0x7f
+            # pos = (data[1] << 8) + data[2]
+            # vel = (data[3] << 4) + (data[4] >> 4)
+            # tau = ((data[4] & 0x0f) << 8) + data[5]
+            # t_mos = data[6]
+            # t_rotor = data[7]
+            # print(hex(id), err, id, pos, vel, tau, goal_tau, t_mos, t_rotor)
+            # CANID = data_recv.arbitration_id
+            CANID = data_recv.data[0]
+            # CMD = data_recv.data[3]
+            CMD = 0x11                  # 飯田：修正の必要あり
+            self.__process_packet(data_recv.data, CANID, CMD)
+
+            # 飯田：Debug print
+            # print(hex(CANID),hex(CMD))
+            # print(hex(data_recv.data[0]),hex(data_recv.data[1]),hex(data_recv.data[2]),hex(data_recv.data[3]),hex(data_recv.data[4]),hex(data_recv.data[5]),hex(data_recv.data[6]),hex(data_recv.data[7]))
+            #return data
+
+    def recv_set_param_data(self):
+        #data_recv = self.serial_.read_all()
+
+        # packets = self.__extract_packets(data_recv)
+        # for packet in packets:
+        #     data = packet[7:15]
+        #     CANID = (packet[6] << 24) | (packet[5] << 16) | (packet[4] << 8) | packet[3]
+        #     CMD = packet[1]
+        #     self.__process_set_param_packet(data, CANID, CMD)
+        
+        data_recv = self.canbus.recv(0.1)
+
+
+        if data_recv is not None:
+            data = data_recv.data
+            CANID = data_recv.arbitration_id
+            # CANID = data_recv.data[0]
+            # CMD = data_recv.data[3]  
+            CMD = 0x11                  # 飯田：修正の必要あり
+            self.__process_packet(data, CANID, CMD)
+
+
+            # 飯田：Debug print
+            print(hex(CANID),hex(CMD))
+            print(hex(data_recv.data[0]),hex(data_recv.data[1]),hex(data_recv.data[2]),hex(data_recv.data[3]),hex(data_recv.data[4]),hex(data_recv.data[5]),hex(data_recv.data[6]),hex(data_recv.data[7]))
+
+    def __process_packet(self, data, CANID, CMD):
+        if CMD == 0x11:
+            if CANID != 0x00:
+                if CANID in self.motors_map:
+                    q_uint = np.uint16((np.uint16(data[1]) << 8) | data[2])
+                    dq_uint = np.uint16((np.uint16(data[3]) << 4) | (data[4] >> 4))
+                    tau_uint = np.uint16(((data[4] & 0xf) << 8) | data[5])
+                    t_mos = data[6]
+                    t_rotor = data[7]
+                    MotorType_recv = self.motors_map[CANID].MotorType
+                    Q_MAX = self.Limit_Param[MotorType_recv][0]
+                    DQ_MAX = self.Limit_Param[MotorType_recv][1]
+                    TAU_MAX = self.Limit_Param[MotorType_recv][2]
+                    recv_q = uint_to_float(q_uint, -Q_MAX, Q_MAX, 16)
+                    recv_dq = uint_to_float(dq_uint, -DQ_MAX, DQ_MAX, 12)
+                    recv_tau = uint_to_float(tau_uint, -TAU_MAX, TAU_MAX, 12)
+                    self.motors_map[CANID].recv_data(recv_q, recv_dq, recv_tau, t_mos, t_rotor)
+            else:
+                MasterID=data[0] & 0x0f
+                if MasterID in self.motors_map:
+                    q_uint = np.uint16((np.uint16(data[1]) << 8) | data[2])
+                    dq_uint = np.uint16((np.uint16(data[3]) << 4) | (data[4] >> 4))
+                    tau_uint = np.uint16(((data[4] & 0xf) << 8) | data[5])
+                    t_mos = data[6]
+                    t_rotor = data[7]
+                    MotorType_recv = self.motors_map[MasterID].MotorType
+                    Q_MAX = self.Limit_Param[MotorType_recv][0]
+                    DQ_MAX = self.Limit_Param[MotorType_recv][1]
+                    TAU_MAX = self.Limit_Param[MotorType_recv][2]
+                    recv_q = uint_to_float(q_uint, -Q_MAX, Q_MAX, 16)
+                    recv_dq = uint_to_float(dq_uint, -DQ_MAX, DQ_MAX, 12)
+                    recv_tau = uint_to_float(tau_uint, -TAU_MAX, TAU_MAX, 12)
+                    self.motors_map[MasterID].recv_data(recv_q, recv_dq, recv_tau, t_mos, t_rotor)
+
+
+    def __process_set_param_packet(self, data, CANID, CMD):
+        if CMD == 0x11 and (data[2] == 0x33 or data[2] == 0x55):
+            masterid=CANID
+            slaveId = ((data[1] << 8) | data[0])
+            if CANID==0x00:  #防止有人把MasterID设为0稳一手
+                masterid=slaveId
+
+            if masterid not in self.motors_map:
+                if slaveId not in self.motors_map:
+                    return
+                else:
+                    masterid=slaveId
+
+            RID = data[3]
+            # 读取参数得到的数据
+            if is_in_ranges(RID):
+                #uint32类型
+                num = uint8s_to_uint32(data[4], data[5], data[6], data[7])
+                self.motors_map[masterid].temp_param_dict[RID] = num
+
+            else:
+                #float类型
+                num = uint8s_to_float(data[4], data[5], data[6], data[7])
+                self.motors_map[masterid].temp_param_dict[RID] = num
+
+
+    def addMotor(self, Motor):
+        """
+        add motor to the motor control object 添加电机到电机控制对象
+        :param Motor: Motor object 电机对象
+        """
+        self.motors_map[Motor.SlaveID] = Motor
+        if Motor.MasterID != 0:
+            self.motors_map[Motor.MasterID] = Motor
+        return True
+
+    def __control_cmd(self, Motor, cmd: np.uint8):     # 飯田：コマンドは通ります
+        data_buf = np.array([0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, cmd], np.uint8)
+        self.__send_data(Motor.SlaveID, data_buf)
+
+    def __send_data(self, motor_id, data):
+        """
+        send data to the motor 发送数据到电机
+        :param motor_id:
+        :param data:
+        :return:
+        """
+        #self.send_data_frame[13] = motor_id & 0xff
+        #self.send_data_frame[14] = (motor_id >> 8)& 0xff  #id high 8 bits
+        #self.send_data_frame[21:29] = data
+        #self.serial_.write(bytes(self.send_data_frame.T))
+
+        msg =can.Message(is_extended_id=False,arbitration_id=motor_id,data=data,is_remote_frame = False)
+        self.canbus.send(msg)
+
+    def __read_RID_param(self, Motor, RID):             # 飯田：修正の必要あり?
+        can_id_l = Motor.SlaveID & 0xff #id low 8 bits
+        can_id_h = (Motor.SlaveID >> 8)& 0xff  #id high 8 bits
+        data_buf = np.array([np.uint8(can_id_l), np.uint8(can_id_h), 0x33, np.uint8(RID), 0x00, 0x00, 0x00, 0x00], np.uint8)
+        self.__send_data(0x7FF, data_buf)
+        
+
+
+    def __write_motor_param(self, Motor, RID, data):             # 飯田：修正の必要あり?
+        can_id_l = Motor.SlaveID & 0xff #id low 8 bits
+        can_id_h = (Motor.SlaveID >> 8)& 0xff  #id high 8 bits
+        data_buf = np.array([np.uint8(can_id_l), np.uint8(can_id_h), 0x55, np.uint8(RID), 0x00, 0x00, 0x00, 0x00], np.uint8)
+        if not is_in_ranges(RID):
+            # data is float
+            data_buf[4:8] = float_to_uint8s(data)
+        else:
+            # data is int
+            data_buf[4:8] = data_to_uint8s(int(data))
+        self.__send_data(0x7FF, data_buf)
+
+    def switchControlMode(self, Motor, ControlMode):
+        """
+        switch the control mode of the motor 切换电机控制模式
+        :param Motor: Motor object 电机对象
+        :param ControlMode: Control_Type 电机控制模式 example:MIT:Control_Type.MIT MIT模式
+        """
+        max_retries = 20
+        retry_interval = 0.1  #retry times
+        RID = 10
+        self.__write_motor_param(Motor, RID, np.uint8(ControlMode))
+        for _ in range(max_retries):
+            sleep(retry_interval)
+            self.recv_set_param_data()
+            if Motor.SlaveID in self.motors_map:
+                if RID in self.motors_map[Motor.SlaveID].temp_param_dict:
+                    if abs(self.motors_map[Motor.SlaveID].temp_param_dict[RID] - ControlMode) < 0.1:
+                        return True
+                    else:
+                        return False
+        return False
+
+    def save_motor_param(self, Motor):
+        """
+        save the all parameter  to flash 保存所有电机参数
+        :param Motor: Motor object 电机对象
+        :return:
+        """
+        can_id_l = Motor.SlaveID & 0xff #id low 8 bits
+        can_id_h = (Motor.SlaveID >> 8)& 0xff  #id high 8 bits
+        data_buf = np.array([np.uint8(can_id_l), np.uint8(can_id_h), 0xAA, 0x00, 0x00, 0x00, 0x00, 0x00], np.uint8)
+        self.disable(Motor)  # before save disable the motor
+        self.__send_data(0x7FF, data_buf)
+        sleep(0.001)
+
+    def change_limit_param(self, Motor_Type, PMAX, VMAX, TMAX):
+        """
+        change the PMAX VMAX TMAX of the motor 改变电机的PMAX VMAX TMAX
+        :param Motor_Type:
+        :param PMAX: 电机的PMAX
+        :param VMAX: 电机的VMAX
+        :param TMAX: 电机的TMAX
+        :return:
+        """
+        self.Limit_Param[Motor_Type][0] = PMAX
+        self.Limit_Param[Motor_Type][1] = VMAX
+        self.Limit_Param[Motor_Type][2] = TMAX
+
+    def refresh_motor_status(self,Motor):
+        """
+        get the motor status 获得电机状态
+        """
+        can_id_l = Motor.SlaveID & 0xff #id low 8 bits
+        can_id_h = (Motor.SlaveID >> 8) & 0xff  #id high 8 bits
+        data_buf = np.array([np.uint8(can_id_l), np.uint8(can_id_h), 0xCC, 0x00, 0x00, 0x00, 0x00, 0x00], np.uint8)
+        self.__send_data(0x7FF, data_buf)
+        self.recv()  # receive the data from serial port
+
+    def change_motor_param(self, Motor, RID, data):
+        """
+        change the RID of the motor 改变电机的参数
+        :param Motor: Motor object 电机对象
+        :param RID: DM_variable 电机参数
+        :param data: 电机参数的值
+        :return: True or False ,True means success, False means fail
+        """
+        max_retries = 20
+        retry_interval = 0.05  #retry times
+
+        self.__write_motor_param(Motor, RID, data)
+        for _ in range(max_retries):
+            self.recv_set_param_data()
+            if Motor.SlaveID in self.motors_map and RID in self.motors_map[Motor.SlaveID].temp_param_dict:
+                if abs(self.motors_map[Motor.SlaveID].temp_param_dict[RID] - data) < 0.1:
+                    return True
+                else:
+                    return False
+            sleep(retry_interval)
+        return False
+
+    def read_motor_param(self, Motor, RID):
+        """
+        read only the RID of the motor 读取电机的内部信息例如 版本号等
+        :param Motor: Motor object 电机对象
+        :param RID: DM_variable 电机参数
+        :return: 电机参数的值
+        """
+        max_retries = 5
+        retry_interval = 0.05  #retry times
+        self.__read_RID_param(Motor, RID)
+        for _ in range(max_retries):
+            sleep(retry_interval)
+            self.recv_set_param_data()
+            if Motor.SlaveID in self.motors_map:
+                if RID in self.motors_map[Motor.SlaveID].temp_param_dict:
+                    return self.motors_map[Motor.SlaveID].temp_param_dict[RID]
+        return None
+
+    # -------------------------------------------------
+    # Extract packets from the serial data
+    def __extract_packets(self, data):          
+        frames = []
+        header = 0xAA
+        tail = 0x55
+        frame_length = 16
+        i = 0
+        remainder_pos = 0
+
+        while i <= len(data) - frame_length:
+            if data[i] == header and data[i + frame_length - 1] == tail:
+                frame = data[i:i + frame_length]
+                frames.append(frame)
+                i += frame_length
+                remainder_pos = i
+            else:
+                i += 1
+        self.data_save = data[remainder_pos:]
+        return frames
+
+
+def LIMIT_MIN_MAX(x, min, max):
+    if x <= min:
+        x = min
+    elif x > max:
+        x = max
+
+
+def float_to_uint(x: float, x_min: float, x_max: float, bits):
+    LIMIT_MIN_MAX(x, x_min, x_max)
+    span = x_max - x_min
+    data_norm = (x - x_min) / span
+    return np.uint16(data_norm * ((1 << bits) - 1))
+
+
+def uint_to_float(x: np.uint16, min: float, max: float, bits):
+    span = max - min
+    data_norm = float(x) / ((1 << bits) - 1)
+    temp = data_norm * span + min
+    return np.float32(temp)
+
+
+def float_to_uint8s(value):
+    # Pack the float into 4 bytes
+    packed = pack('f', value)
+    # Unpack the bytes into four uint8 values
+    return unpack('4B', packed)
+
+
+def data_to_uint8s(value):
+    # Check if the value is within the range of uint32
+    if isinstance(value, int) and (0 <= value <= 0xFFFFFFFF):
+        # Pack the uint32 into 4 bytes
+        packed = pack('I', value)
+    else:
+        raise ValueError("Value must be an integer within the range of uint32")
+
+    # Unpack the bytes into four uint8 values
+    return unpack('4B', packed)
+
+
+def is_in_ranges(number):
+    """
+    check if the number is in the range of uint32
+    :param number:
+    :return:
+    """
+    if (7 <= number <= 10) or (13 <= number <= 16) or (35 <= number <= 36):
+        return True
+    return False
+
+
+def uint8s_to_uint32(byte1, byte2, byte3, byte4):
+    # Pack the four uint8 values into a single uint32 value in little-endian order
+    packed = pack('<4B', byte1, byte2, byte3, byte4)
+    # Unpack the packed bytes into a uint32 value
+    return unpack('<I', packed)[0]
+
+
+def uint8s_to_float(byte1, byte2, byte3, byte4):
+    # Pack the four uint8 values into a single float value in little-endian order
+    packed = pack('<4B', byte1, byte2, byte3, byte4)
+    # Unpack the packed bytes into a float value
+    return unpack('<f', packed)[0]
+
+
+def print_hex(data):
+    hex_values = [f'{byte:02X}' for byte in data]
+    print(' '.join(hex_values))
+
+
+def get_enum_by_index(index, enum_class):
+    try:
+        return enum_class(index)
+    except ValueError:
+        return None
+
+
+class DM_Motor_Type(IntEnum):
+    DM4310 = 0
+    DM4310_48V = 1
+    DM4340 = 2
+    DM4340_48V = 3
+    DM6006 = 4
+    DM8006 = 5
+    DM8009 = 6
+    DM10010L = 7
+    DM10010 = 8
+    DMH3510 = 9
+    DMH6215 = 10
+    DMG6220 = 11
+
+
+class DM_variable(IntEnum):
+    UV_Value = 0
+    KT_Value = 1
+    OT_Value = 2
+    OC_Value = 3
+    ACC = 4
+    DEC = 5
+    MAX_SPD = 6
+    MST_ID = 7
+    ESC_ID = 8
+    TIMEOUT = 9
+    CTRL_MODE = 10
+    Damp = 11
+    Inertia = 12
+    hw_ver = 13
+    sw_ver = 14
+    SN = 15
+    NPP = 16
+    Rs = 17
+    LS = 18
+    Flux = 19
+    Gr = 20
+    PMAX = 21
+    VMAX = 22
+    TMAX = 23
+    I_BW = 24
+    KP_ASR = 25
+    KI_ASR = 26
+    KP_APR = 27
+    KI_APR = 28
+    OV_Value = 29
+    GREF = 30
+    Deta = 31
+    V_BW = 32
+    IQ_c1 = 33
+    VL_c1 = 34
+    can_br = 35
+    sub_ver = 36
+    u_off = 50
+    v_off = 51
+    k1 = 52
+    k2 = 53
+    m_off = 54
+    dir = 55
+    p_m = 80
+    xout = 81
+
+
+class Control_Type(IntEnum):
+    MIT = 1
+    POS_VEL = 2
+    VEL = 3
+    Torque_Pos = 4
+
+class DamiaoPort:
+    def __init__(self, device, types, can_ids, master_ids, motor_with_torque, control_mode=Control_Type.MIT):
+        self.device = device
+        self.types = types
+        self.can_ids = can_ids
+        self.master_ids = master_ids
+        self.control = MotorControl(self.device, bitrate=4000000)
+        self.motors = [Motor(type, can_id, master_id) for type, can_id, master_id in zip(types, can_ids, master_ids)]
+        self.stat_data = []
+        self.stat_time = []
+        for motor in self.motors:
+            self.control.addMotor(motor)
+            self.control.enable(motor)
+
+    def get_present_status(self):
+        self.stat_time.append(time())
+        stat = [[
+            motor.goal_position,
+            motor.goal_tau,
+            motor.getPosition(),
+            motor.getVelocity(),
+            motor.getTorque(),
+            motor.state_tmos,
+            motor.state_trotor,
+        ] for motor in self.motors]
+        self.stat_data.append(stat)
+
+        return stat
+
+    def save_status(self, filename):
+        np.savez(filename, np.array(self.stat_time), np.array(self.stat_data))
+
+    def disable(self):
+        for motor in self.motors:
+            self.control.disable(motor)
+
+    def shutdown(self):
+        for motor in self.motors:
+            self.control.controlMIT(motor, 0, 0, 0, 0, 0)
+        self.control.canbus.shutdown()
+
+    def set_zero_position(self):
+        for motor in self.motors:
+            self.control.disable(motor)
+        sleep(1)
+        for motor in self.motors:
+            self.control.set_zero_position(motor)
+        sleep(1)
+        for motor in self.motors:
+            self.control.enable(motor)
+        return 0
+
+    async def move_towards(self, goal_positions, kps, kds):
+        for motor, goal_position, kp, kd in zip(self.motors, goal_positions, kps, kds):
+            delta = goal_position - motor.getPosition()
+            v = motor.getVelocity()
+            tau = kp * delta - kd * v
+            motor.goal_position = goal_position
+            motor.goal_tau = tau
+            self.control.controlMIT(motor, 0, 0, 0, 0, tau)
+            await asyncio.sleep(0.00003)
+
+    def move_regressor_sync(self, regs, search_range, search_step, goal_positions, kps, kds):
+        TORQUE_SCALER=30
+        if len(self.stat_data) == 0:
+            return self.move_towards_sync(goal_positions, kps, kds)
+        for motor, reg, goal_position, kp, kd, stat in zip(
+                self.motors, regs, goal_positions, kps, kds, self.stat_data[-1]):
+            pos = motor.getPosition()
+            vel = motor.getVelocity()
+            delta = goal_position - pos
+            goal_tau = kp * delta - kd * vel
+            _goal_pos, _goal_tau, _pos, _vel, _tau = stat
+            x = np.array([[_pos, _vel, _tau, _goal_pos, _goal_tau],
+                          [pos, vel, motor.getTorque(), goal_position, goal_tau]])
+            x /= np.array([[np.pi, 10, TORQUE_SCALER, np.pi, TORQUE_SCALER]])
+            xs = []
+            for tau in np.linspace(goal_tau/TORQUE_SCALER - search_range,
+                                   goal_tau/TORQUE_SCALER + search_range,
+                                   num=search_step):
+                x_ = x.copy()
+                x_[0,4] = tau
+                xs.append(x_.flatten())
+            h = reg.predict(xs)
+            diff = h - goal_position
+            tau = TORQUE_SCALER * xs[np.argmin(diff ** 2)][4]
+            goal_tau = tau
+            motor.goal_position = goal_position
+            motor.goal_tau = goal_tau
+            self.control.controlMIT(motor, 0, 0, 0, 0, goal_tau)
+            sleep(0.00003)
+
+    def move_towards_sync(self, goal_positions, kps, kds):
+        for motor, goal_position, kp, kd in zip(self.motors, goal_positions, kps, kds):
+            delta = goal_position - motor.getPosition()
+            v = motor.getVelocity()
+            tau = kp * delta - kd * v
+            motor.goal_position = goal_position
+            motor.goal_tau = tau
+            self.control.controlMIT(motor, 0, 0, 0, 0, tau)
+            # sleep(0.00003)
+
+    def set_goal_torque_sync(self, goal_taus):
+        for motor, goal_tau in zip(self.motors, goal_taus):
+            motor.goal_position = 0
+            motor.goal_tau = goal_tau
+            self.control.controlMIT(motor, 0, 0, 0, 0, motor.goal_tau)
+            sleep(0.00003)
+
+    def move_torque_sync(self, taus):
+        for motor,tau in zip(self.motors, taus):
+            motor.goal_position = 0
+            motor.goal_tau = tau
+            self.control.controlMIT(motor, 0, 0, 0, 0, motor.goal_tau)
+            sleep(0.00003)
+
+    def keep_torque_sync(self):
+        for motor in self.motors:
+            self.control.controlMIT(motor, 0, 0, 0, 0, motor.goal_tau)
+            sleep(0.00003)
+
+    async def set_goal_positions(self, goal_positions, kps):
+        for motor, goal_position, kp in zip(self.motors, goal_positions, kps):
+            motor.goal_position = goal_position
+            motor.goal_tau = 0
+            self.control.controlMIT(motor, kp, 1.2, goal_position, 0, 0)
+            await asyncio.sleep(0.00003)
+
+    def set_goal_positions_sync(self, goal_positions, kps, kds):
+        for motor, goal_position, kp, kd in zip(self.motors, goal_positions, kps, kds):
+            motor.goal_position = goal_position
+            motor.goal_tau = 0
+            self.control.controlMIT(motor, kp, kd, goal_position, 0, 0)
+            sleep(0.00003)
+
+    def set_goal_posvel(self, goal_positions):
+        for motor, goal_position in zip(self.motors, goal_positions):
+            motor.goal_position = goal_position
+            motor.goal_tau = 0
+            self.control.control_pos_force(motor, goal_position, 1, 1)
+
+    def controlMIT(self, motor, kp, kd, q, dq, tau):
+        self.control.controlMIT(self.motors[motor], kp, kd, q, dq, tau)
+

src/lerobot/motors/damiao/tables.py

@@ -0,0 +1,209 @@
+# 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.
+
+"""Configuration tables for Damiao motors."""
+
+from enum import IntEnum
+from typing import Dict, List, Tuple
+
+# Motor type definitions
+class MotorType(IntEnum):
+    DM3507 = 0
+    DM4310 = 1
+    DM4310_48V = 2
+    DM4340 = 3
+    DM4340_48V = 4
+    DM6006 = 5
+    DM8006 = 6
+    DM8009 = 7
+    DM10010L = 8
+    DM10010 = 9
+    DMH3510 = 10
+    DMH6215 = 11
+    DMG6220 = 12
+
+# Control modes
+class ControlMode(IntEnum):
+    MIT = 1
+    POS_VEL = 2
+    VEL = 3
+    TORQUE_POS = 4
+
+# Motor variable IDs (RID)
+class MotorVariable(IntEnum):
+    UV_VALUE = 0
+    KT_VALUE = 1
+    OT_VALUE = 2
+    OC_VALUE = 3
+    ACC = 4
+    DEC = 5
+    MAX_SPD = 6
+    MST_ID = 7
+    ESC_ID = 8
+    TIMEOUT = 9
+    CTRL_MODE = 10
+    DAMP = 11
+    INERTIA = 12
+    HW_VER = 13
+    SW_VER = 14
+    SN = 15
+    NPP = 16
+    RS = 17
+    LS = 18
+    FLUX = 19
+    GR = 20
+    PMAX = 21
+    VMAX = 22
+    TMAX = 23
+    I_BW = 24
+    KP_ASR = 25
+    KI_ASR = 26
+    KP_APR = 27
+    KI_APR = 28
+    OV_VALUE = 29
+    GREF = 30
+    DETA = 31
+    V_BW = 32
+    IQ_C1 = 33
+    VL_C1 = 34
+    CAN_BR = 35
+    SUB_VER = 36
+    U_OFF = 50
+    V_OFF = 51
+    K1 = 52
+    K2 = 53
+    M_OFF = 54
+    DIR = 55
+    P_M = 80
+    XOUT = 81
+
+# Motor limit parameters [PMAX, VMAX, TMAX] 
+# PMAX: Maximum position (rad)
+# VMAX: Maximum velocity (rad/s)
+# TMAX: Maximum torque (N·m)
+MOTOR_LIMIT_PARAMS = {
+    MotorType.DM3507: (12.5, 30, 10),
+    MotorType.DM4310: (12.5, 30, 10),
+    MotorType.DM4310_48V: (12.5, 50, 10),
+    MotorType.DM4340: (12.5, 8, 28),
+    MotorType.DM4340_48V: (12.5, 10, 28),
+    MotorType.DM6006: (12.5, 45, 20),
+    MotorType.DM8006: (12.5, 45, 40),
+    MotorType.DM8009: (12.5, 45, 54),
+    MotorType.DM10010L: (12.5, 25, 200),
+    MotorType.DM10010: (12.5, 20, 200),
+    MotorType.DMH3510: (12.5, 280, 1),
+    MotorType.DMH6215: (12.5, 45, 10),
+    MotorType.DMG6220: (12.5, 45, 10),
+}
+
+# Motor model names
+MODEL_NAMES = {
+    MotorType.DM3507: "dm3507",
+    MotorType.DM4310: "dm4310",
+    MotorType.DM4310_48V: "dm4310_48v",
+    MotorType.DM4340: "dm4340",
+    MotorType.DM4340_48V: "dm4340_48v",
+    MotorType.DM6006: "dm6006",
+    MotorType.DM8006: "dm8006",
+    MotorType.DM8009: "dm8009",
+    MotorType.DM10010L: "dm10010l",
+    MotorType.DM10010: "dm10010",
+    MotorType.DMH3510: "dmh3510",
+    MotorType.DMH6215: "dmh6215",
+    MotorType.DMG6220: "dmg6220",
+}
+
+# Motor resolution table (encoder counts per revolution)
+MODEL_RESOLUTION = {
+    "dm3507": 65536,
+    "dm4310": 65536,
+    "dm4310_48v": 65536,
+    "dm4340": 65536,
+    "dm4340_48v": 65536,
+    "dm6006": 65536,
+    "dm8006": 65536,
+    "dm8009": 65536,
+    "dm10010l": 65536,
+    "dm10010": 65536,
+    "dmh3510": 65536,
+    "dmh6215": 65536,
+    "dmg6220": 65536,
+}
+
+# CAN baudrates supported by Damiao motors
+AVAILABLE_BAUDRATES = [
+    125000,   # 0: 125 kbps
+    200000,   # 1: 200 kbps
+    250000,   # 2: 250 kbps
+    500000,   # 3: 500 kbps
+    1000000,  # 4: 1 mbps (default for OpenArms)
+    2000000,  # 5: 2 mbps
+    2500000,  # 6: 2.5 mbps
+    3200000,  # 7: 3.2 mbps
+    4000000,  # 8: 4 mbps
+    5000000,  # 9: 5 mbps
+]
+DEFAULT_BAUDRATE = 1000000  # 1 Mbps is standard for OpenArms
+
+# Default timeout in milliseconds
+DEFAULT_TIMEOUT_MS = 1000
+
+# Data that should be normalized
+NORMALIZED_DATA = ["Present_Position", "Goal_Position"]
+
+# OpenArms specific configurations
+# Based on: https://docs.openarm.dev/software/setup/configure-test
+# OpenArms has 7 DOF per arm (14 total for dual arm)
+OPENARMS_ARM_MOTOR_IDS = {
+    "joint_1": {"send": 0x01, "recv": 0x11},  # J1 - Shoulder pan
+    "joint_2": {"send": 0x02, "recv": 0x12},  # J2 - Shoulder lift
+    "joint_3": {"send": 0x03, "recv": 0x13},  # J3 - Elbow flex
+    "joint_4": {"send": 0x04, "recv": 0x14},  # J4 - Wrist flex
+    "joint_5": {"send": 0x05, "recv": 0x15},  # J5 - Wrist roll
+    "joint_6": {"send": 0x06, "recv": 0x16},  # J6 - Wrist pitch
+    "joint_7": {"send": 0x07, "recv": 0x17},  # J7 - Wrist rotation
+}
+
+OPENARMS_GRIPPER_MOTOR_IDS = {
+    "gripper": {"send": 0x08, "recv": 0x18},  # J8 - Gripper
+}
+
+# Default motor types for OpenArms
+OPENARMS_DEFAULT_MOTOR_TYPES = {
+    "joint_1": MotorType.DM8009,   # Shoulder pan - high torque
+    "joint_2": MotorType.DM8009,   # Shoulder lift - high torque
+    "joint_3": MotorType.DM4340,   # Shoulder rotation
+    "joint_4": MotorType.DM4340,   # Elbow flex
+    "joint_5": MotorType.DM4310,   # Wrist roll
+    "joint_6": MotorType.DM4310,   # Wrist pitch
+    "joint_7": MotorType.DM4310,   # Wrist rotation
+    "gripper": MotorType.DM4310,   # Gripper
+}
+
+# MIT control parameter ranges
+MIT_KP_RANGE = (0.0, 500.0)
+MIT_KD_RANGE = (0.0, 5.0)
+
+# CAN frame command IDs
+CAN_CMD_ENABLE = 0xFC
+CAN_CMD_DISABLE = 0xFD
+CAN_CMD_SET_ZERO = 0xFE
+CAN_CMD_REFRESH = 0xCC
+CAN_CMD_QUERY_PARAM = 0x33
+CAN_CMD_WRITE_PARAM = 0x55
+CAN_CMD_SAVE_PARAM = 0xAA
+
+# CAN ID for parameter operations
+CAN_PARAM_ID = 0x7FF

src/lerobot/motors/dynamixel/dynamixel.py

@@ -24,7 +24,7 @@ from enum import Enum
 
 from lerobot.motors.encoding_utils import decode_twos_complement, encode_twos_complement
 
-from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
+from ..motors_bus import Motor, MotorCalibration, NameOrID, SerialMotorsBus, Value, get_address
 from .tables import (
     AVAILABLE_BAUDRATES,
     MODEL_BAUDRATE_TABLE,
@@ -100,7 +100,7 @@ def _split_into_byte_chunks(value: int, length: int) -> list[int]:
     return data
 
 
-class DynamixelMotorsBus(MotorsBus):
+class DynamixelMotorsBus(SerialMotorsBus):
     """
     The Dynamixel implementation for a MotorsBus. It relies on the python dynamixel sdk to communicate with
     the motors. For more info, see the Dynamixel SDK Documentation:

src/lerobot/motors/feetech/feetech.py

@@ -19,7 +19,7 @@ from pprint import pformat
 
 from lerobot.motors.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
 
-from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
+from ..motors_bus import Motor, MotorCalibration, NameOrID, SerialMotorsBus, Value, get_address
 from .tables import (
     FIRMWARE_MAJOR_VERSION,
     FIRMWARE_MINOR_VERSION,
@@ -96,7 +96,7 @@ def patch_setPacketTimeout(self, packet_length):  # noqa: N802
     self.packet_timeout = (self.tx_time_per_byte * packet_length) + (self.tx_time_per_byte * 3.0) + 50
 
 
-class FeetechMotorsBus(MotorsBus):
+class FeetechMotorsBus(SerialMotorsBus):
     """
     The FeetechMotorsBus class allows to efficiently read and write to the attached motors. It relies on the
     python feetech sdk to communicate with the motors, which is itself based on the dynamixel sdk.

src/lerobot/motors/motors_bus.py

@@ -19,6 +19,8 @@
 # TODO(aliberts): Add block noqa when feature below is available
 # https://github.com/astral-sh/ruff/issues/3711
 
+from __future__ import annotations
+
 import abc
 import logging
 from contextlib import contextmanager
@@ -41,6 +43,92 @@ Value: TypeAlias = int | float
 logger = logging.getLogger(__name__)
 
 
+class MotorsBusBase(abc.ABC):
+    """
+    Base class for all motor bus implementations.
+    
+    This is a minimal interface that all motor buses must implement, regardless of their
+    communication protocol (serial, CAN, etc.).
+    """
+
+    def __init__(
+        self,
+        port: str,
+        motors: dict[str, Motor],
+        calibration: dict[str, MotorCalibration] | None = None,
+    ):
+        self.port = port
+        self.motors = motors
+        self.calibration = calibration if calibration else {}
+
+    @abc.abstractmethod
+    def connect(self, handshake: bool = True) -> None:
+        """Establish connection to the motors."""
+        pass
+
+    @abc.abstractmethod
+    def disconnect(self, disable_torque: bool = True) -> None:
+        """Disconnect from the motors."""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def is_connected(self) -> bool:
+        """Check if connected to the motors."""
+        pass
+
+    @abc.abstractmethod
+    def read(self, data_name: str, motor: str, *, normalize: bool = True, num_retry: int = 0) -> Value:
+        """Read a value from a single motor."""
+        pass
+
+    @abc.abstractmethod
+    def write(
+        self, data_name: str, motor: str, value: Value, *, normalize: bool = True, num_retry: int = 0
+    ) -> None:
+        """Write a value to a single motor."""
+        pass
+
+    @abc.abstractmethod
+    def sync_read(
+        self, data_name: str, motors: str | list[str] | None = None, *, normalize: bool = True
+    ) -> dict[str, Value]:
+        """Read a value from multiple motors."""
+        pass
+
+    @abc.abstractmethod
+    def sync_write(
+        self,
+        data_name: str,
+        values: Value | dict[str, Value],
+        motors: str | list[str] | None = None,
+        *,
+        normalize: bool = True,
+    ) -> None:
+        """Write values to multiple motors."""
+        pass
+
+    @abc.abstractmethod
+    def enable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        """Enable torque on selected motors."""
+        pass
+
+    @abc.abstractmethod
+    def disable_torque(self, motors: int | str | list[str] | None = None, num_retry: int = 0) -> None:
+        """Disable torque on selected motors."""
+        pass
+
+    @abc.abstractmethod
+    def read_calibration(self) -> dict[str, MotorCalibration]:
+        """Read calibration parameters from the motors."""
+        pass
+
+    @abc.abstractmethod
+    def write_calibration(self, calibration_dict: dict[str, MotorCalibration], cache: bool = True) -> None:
+        """Write calibration parameters to the motors."""
+        pass
+
+
 def get_ctrl_table(model_ctrl_table: dict[str, dict], model: str) -> dict[str, tuple[int, int]]:
     ctrl_table = model_ctrl_table.get(model)
     if ctrl_table is None:
@@ -203,15 +291,15 @@ class GroupSyncWrite(Protocol):
     def txPacket(self): ...
 
 
-class MotorsBus(abc.ABC):
+class SerialMotorsBus(MotorsBusBase):
     """
-    A MotorsBus allows to efficiently read and write to the attached motors.
+    A SerialMotorsBus allows to efficiently read and write to motors connected via serial communication.
     It represents several motors daisy-chained together and connected through a serial port.
-    There are currently two implementations of this abstract class:
+    There are currently two implementations of this class:
         - DynamixelMotorsBus
         - FeetechMotorsBus
 
-    Note: This class may evolve in the future should we add support for other types of bus.
+    This class is specifically for serial-based motor protocols (Dynamixel, Feetech, etc.).
 
     A MotorsBus subclass instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
     To find the port, you can run our utility script:
@@ -1212,3 +1300,7 @@ class MotorsBus(abc.ABC):
         for id_, value in ids_values.items():
             data = self._serialize_data(value, length)
             self.sync_writer.addParam(id_, data)
+
+
+# Backward compatibility alias
+MotorsBus = SerialMotorsBus

src/lerobot/policies/factory.py

@@ -35,11 +35,9 @@ from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi05.configuration_pi05 import PI05Config
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.policies.utils import validate_visual_features_consistency
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 from lerobot.processor.converters import (
@@ -104,10 +102,6 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
-    elif name == "sarm":
-        from lerobot.policies.sarm.modeling_sarm import SARMRewardModel
-
-        return SARMRewardModel
     elif name == "groot":
         from lerobot.policies.groot.modeling_groot import GrootPolicy
 
@@ -327,14 +321,6 @@ def make_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
         )
 
-    elif isinstance(policy_cfg, SARMConfig):
-        from lerobot.policies.sarm.processor_sarm import make_sarm_pre_post_processors
-
-        processors = make_sarm_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            dataset_meta=kwargs.get("dataset_meta"),
-        )
     elif isinstance(policy_cfg, GrootConfig):
         from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
 
@@ -418,13 +404,6 @@ def make_policy(
     if not cfg.input_features:
         cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
     kwargs["config"] = cfg
-    
-    # Pass dataset_stats to the policy if available (needed for some policies like SARM)
-    if ds_meta is not None and hasattr(ds_meta, 'stats'):
-        kwargs["dataset_stats"] = ds_meta.stats
-    
-    if ds_meta is not None:
-        kwargs["dataset_meta"] = ds_meta
 
     if cfg.pretrained_path:
         # Load a pretrained policy and override the config if needed (for example, if there are inference-time
@@ -441,7 +420,20 @@ def make_policy(
     # policy = torch.compile(policy, mode="reduce-overhead")
 
     if not rename_map:
-        validate_visual_features_consistency(cfg, features)
-        # TODO: (jadechoghari) - add a check_state(cfg, features) and check_action(cfg, features)
-
+        expected_features = set(cfg.input_features.keys()) | set(cfg.output_features.keys())
+        provided_features = set(features.keys())
+        if expected_features and provided_features != expected_features:
+            missing = expected_features - provided_features
+            extra = provided_features - expected_features
+            # TODO (jadechoghari): provide a dynamic rename map suggestion to the user.
+            raise ValueError(
+                f"Feature mismatch between dataset/environment and policy config.\n"
+                f"- Missing features: {sorted(missing) if missing else 'None'}\n"
+                f"- Extra features: {sorted(extra) if extra else 'None'}\n\n"
+                f"Please ensure your dataset and policy use consistent feature names.\n"
+                f"If your dataset uses different observation keys (e.g., cameras named differently), "
+                f"use the `--rename_map` argument, for example:\n"
+                f'  --rename_map=\'{{"observation.images.left": "observation.images.camera1", '
+                f'"observation.images.top": "observation.images.camera2"}}\''
+            )
     return policy

src/lerobot/policies/pi0/configuration_pi0.py

@@ -20,7 +20,6 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.utils.constants import OBS_IMAGES
 
 
@@ -48,9 +47,6 @@ class PI0Config(PreTrainedConfig):
     min_period: float = 4e-3
     max_period: float = 4.0
 
-    # Real-Time Chunking (RTC) configuration
-    rtc_config: RTCConfig | None = None
-
     image_resolution: tuple[int, int] = (224, 224)  # see openpi `preprocessing_pytorch.py`
 
     # Add empty images. Used to add empty cameras when no image features are present.

src/lerobot/policies/pi0/modeling_pi0.py

@@ -19,12 +19,11 @@ import logging
 import math
 from collections import deque
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, TypedDict
+from typing import TYPE_CHECKING, Literal
 
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
-from typing_extensions import Unpack
 
 from lerobot.utils.import_utils import _transformers_available
 
@@ -43,7 +42,6 @@ else:
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pretrained import PreTrainedPolicy, T
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
@@ -53,12 +51,6 @@ from lerobot.utils.constants import (
 )
 
 
-class ActionSelectKwargs(TypedDict, total=False):
-    inference_delay: int | None
-    prev_chunk_left_over: Tensor | None
-    execution_horizon: int | None
-
-
 def get_safe_dtype(target_dtype, device_type):
     """Get a safe dtype for the given device type."""
     if device_type == "mps" and target_dtype == torch.float64:
@@ -511,10 +503,9 @@ class PaliGemmaWithExpertModel(
 class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
     """Core PI0 PyTorch model."""
 
-    def __init__(self, config: PI0Config, rtc_processor: RTCProcessor | None = None):
+    def __init__(self, config: PI0Config):
         super().__init__()
         self.config = config
-        self.rtc_processor = rtc_processor
 
         paligemma_config = get_gemma_config(config.paligemma_variant)
         action_expert_config = get_gemma_config(config.action_expert_variant)
@@ -569,9 +560,6 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI0Pytorch model")
 
-    def _rtc_enabled(self):
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
     def _apply_checkpoint(self, func, *args, **kwargs):
         """Helper method to apply gradient checkpointing if enabled."""
         if self.gradient_checkpointing_enabled and self.training:
@@ -768,15 +756,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
     @torch.no_grad()  # see openpi `sample_actions` (slightly adapted)
     def sample_actions(
-        self,
-        images,
-        img_masks,
-        lang_tokens,
-        lang_masks,
-        state,
-        noise=None,
-        num_steps=None,
-        **kwargs: Unpack[ActionSelectKwargs],
+        self, images, img_masks, lang_tokens, lang_masks, state, noise=None, num_steps=None
     ) -> Tensor:
         """Do a full inference forward and compute the action."""
         if num_steps is None:
@@ -818,41 +798,14 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         time = torch.tensor(1.0, dtype=torch.float32, device=device)
         while time >= -dt / 2:
             expanded_time = time.expand(bsize)
-
-            # Define a closure function to properly capture expanded_time
-            # This avoids the lambda expression (E731) and loop variable binding (B023) issues
-            def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
-                return self.denoise_step(
-                    state=state,
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    x_t=input_x_t,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
-            # Euler step
-            x_t += dt * v_t
-
-            # Record x_t and v_t after Euler step
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
+            v_t = self.denoise_step(
+                state,
+                prefix_pad_masks,
+                past_key_values,
+                x_t,
+                expanded_time,
+            )
+            x_t = x_t + dt * v_t
             time += dt
 
         return x_t
@@ -916,8 +869,7 @@ class PI0Policy(PreTrainedPolicy):
         self.config = config
 
         # Initialize the core PI0 model
-        self.init_rtc_processor()
-        self.model = PI0Pytorch(config, rtc_processor=self.rtc_processor)
+        self.model = PI0Pytorch(config)
 
         # Enable gradient checkpointing if requested
         if config.gradient_checkpointing:
@@ -1107,22 +1059,6 @@ class PI0Policy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
-    def init_rtc_processor(self):
-        """Initialize RTC processor if RTC is enabled in config."""
-        self.rtc_processor = None
-
-        # Create processor if config provided
-        # If RTC is not enabled - we can still track the denoising data
-        if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
-
-            model_value = getattr(self, "model", None)
-            if model_value is not None:
-                model_value.rtc_processor = self.rtc_processor
-
-    def _rtc_enabled(self) -> bool:
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
     def _preprocess_images(self, batch: dict[str, Tensor]) -> tuple[list[Tensor], list[Tensor]]:
         """Preprocess images for the model.
 
@@ -1201,10 +1137,6 @@ class PI0Policy(PreTrainedPolicy):
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations."""
-        assert not self._rtc_enabled(), (
-            "RTC is not supported for select_action, use it with predict_action_chunk"
-        )
-
         self.eval()
 
         # Action queue logic for n_action_steps > 1
@@ -1216,7 +1148,7 @@ class PI0Policy(PreTrainedPolicy):
         return self._action_queue.popleft()
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
@@ -1225,8 +1157,8 @@ class PI0Policy(PreTrainedPolicy):
         lang_tokens, lang_masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
         state = self.prepare_state(batch)
 
-        # Sample actions using the model (pass through RTC kwargs)
-        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, **kwargs)
+        # Sample actions using the model
+        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state)
 
         # Unpad actions to actual action dimension
         original_action_dim = self.config.output_features[ACTION].shape[0]

src/lerobot/policies/pi05/configuration_pi05.py

@@ -20,7 +20,6 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 
 
 @PreTrainedConfig.register_subclass("pi05")
@@ -47,9 +46,6 @@ class PI05Config(PreTrainedConfig):
     min_period: float = 4e-3
     max_period: float = 4.0
 
-    # Real-Time Chunking (RTC) configuration
-    rtc_config: RTCConfig | None = None
-
     image_resolution: tuple[int, int] = (224, 224)  # see openpi `preprocessing_pytorch.py`
 
     # Add empty images. Used to add empty cameras when no image features are present.

src/lerobot/policies/pi05/modeling_pi05.py

@@ -19,12 +19,11 @@ import logging
 import math
 from collections import deque
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, TypedDict
+from typing import TYPE_CHECKING, Literal
 
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
-from typing_extensions import Unpack
 
 from lerobot.utils.import_utils import _transformers_available
 
@@ -43,7 +42,6 @@ else:
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.policies.pi05.configuration_pi05 import PI05Config
 from lerobot.policies.pretrained import PreTrainedPolicy, T
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
@@ -52,12 +50,6 @@ from lerobot.utils.constants import (
 )
 
 
-class ActionSelectKwargs(TypedDict, total=False):
-    inference_delay: int | None
-    prev_chunk_left_over: Tensor | None
-    execution_horizon: int | None
-
-
 def get_safe_dtype(target_dtype, device_type):
     """Get a safe dtype for the given device type."""
     if device_type == "mps" and target_dtype == torch.float64:
@@ -510,10 +502,9 @@ class PaliGemmaWithExpertModel(
 class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
     """Core PI05 PyTorch model."""
 
-    def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
+    def __init__(self, config: PI05Config):
         super().__init__()
         self.config = config
-        self.rtc_processor = rtc_processor
 
         paligemma_config = get_gemma_config(config.paligemma_variant)
         action_expert_config = get_gemma_config(config.action_expert_variant)
@@ -565,9 +556,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI05Pytorch model")
 
-    def _rtc_enabled(self):
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
     def _apply_checkpoint(self, func, *args, **kwargs):
         """Helper method to apply gradient checkpointing if enabled."""
         if self.gradient_checkpointing_enabled and self.training:
@@ -743,16 +731,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         return F.mse_loss(u_t, v_t, reduction="none")
 
     @torch.no_grad()  # see openpi `sample_actions` (slightly adapted)
-    def sample_actions(
-        self,
-        images,
-        img_masks,
-        tokens,
-        masks,
-        noise=None,
-        num_steps=None,
-        **kwargs: Unpack[ActionSelectKwargs],
-    ) -> Tensor:
+    def sample_actions(self, images, img_masks, tokens, masks, noise=None, num_steps=None) -> Tensor:
         """Do a full inference forward and compute the action."""
         if num_steps is None:
             num_steps = self.config.num_inference_steps
@@ -791,40 +770,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         time = torch.tensor(1.0, dtype=torch.float32, device=device)
         while time >= -dt / 2:
             expanded_time = time.expand(bsize)
-
-            # Define a closure function to properly capture expanded_time
-            # This avoids the lambda expression (E731) and loop variable binding (B023) issues
-            def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
-                return self.denoise_step(
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    x_t=input_x_t,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
-            # Euler step
-            x_t += dt * v_t
-
-            # Record x_t and v_t after Euler step
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
+            v_t = self.denoise_step(
+                prefix_pad_masks,
+                past_key_values,
+                x_t,
+                expanded_time,
+            )
+            x_t = x_t + dt * v_t
             time += dt
 
         return x_t
@@ -887,8 +839,7 @@ class PI05Policy(PreTrainedPolicy):
         self.config = config
 
         # Initialize the core PI05 model
-        self.init_rtc_processor()
-        self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
+        self.model = PI05Pytorch(config)
 
         # Enable gradient checkpointing if requested
         if config.gradient_checkpointing:
@@ -1084,22 +1035,6 @@ class PI05Policy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
-    def init_rtc_processor(self):
-        """Initialize RTC processor if RTC is enabled in config."""
-        self.rtc_processor = None
-
-        # Create processor if config provided
-        # If RTC is not enabled - we can still track the denoising data
-        if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
-
-            model_value = getattr(self, "model", None)
-            if model_value is not None:
-                model_value.rtc_processor = self.rtc_processor
-
-    def _rtc_enabled(self) -> bool:
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
     def _preprocess_images(self, batch: dict[str, Tensor]) -> tuple[list[Tensor], list[Tensor]]:
         """Preprocess images for the model.
 
@@ -1174,10 +1109,6 @@ class PI05Policy(PreTrainedPolicy):
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations."""
-        assert not self._rtc_enabled(), (
-            "RTC is not supported for select_action, use it with predict_action_chunk"
-        )
-
         self.eval()
 
         # Action queue logic for n_action_steps > 1
@@ -1189,7 +1120,7 @@ class PI05Policy(PreTrainedPolicy):
         return self._action_queue.popleft()
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
@@ -1197,8 +1128,8 @@ class PI05Policy(PreTrainedPolicy):
         images, img_masks = self._preprocess_images(batch)
         tokens, masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
 
-        # Sample actions using the model (pass through RTC kwargs, no separate state needed for PI05)
-        actions = self.model.sample_actions(images, img_masks, tokens, masks, **kwargs)
+        # Sample actions using the model (no separate state needed for PI05)
+        actions = self.model.sample_actions(images, img_masks, tokens, masks)
 
         # Unpad actions to actual action dimension
         original_action_dim = self.config.output_features[ACTION].shape[0]

src/lerobot/policies/rtc/README.md

@@ -1,38 +0,0 @@
-# Real-Time Chunking (RTC)
-
-This module contains the LeRobot implementation of **Real-Time Chunking (RTC)**, an inference-time technique for flow-matching based policies.
-
-**Note**: RTC is not a policy itself, but rather an inference enhancement that works with flow-matching based policies including [π₀](../pi0/), [π₀.₅](../pi05/), and [SmolVLA](../smolvla/).
-
----
-
-## Citation
-
-If you use Real-Time Chunking in your work, please cite:
-
-```bibtex
-@misc{openpi2024,
-  author       = {Physical Intelligence Lab},
-  title        = {OpenPI: PyTorch Implementation of π0 and π0.5 Policies},
-  year         = {2024},
-  publisher    = {GitHub},
-  howpublished = {\url{https://github.com/Physical-Intelligence/openpi}},
-  license      = {Apache-2.0}
-}
-
-@misc{black2025realtimeexecutionactionchunking,
-      title={Real-Time Execution of Action Chunking Flow Policies},
-      author={Kevin Black and Manuel Y. Galliker and Sergey Levine},
-      year={2025},
-      eprint={2506.07339},
-      archivePrefix={arXiv},
-      primaryClass={cs.RO},
-      url={https://arxiv.org/abs/2506.07339},
-}
-```
-
----
-
-## License
-
-This implementation follows the **Apache 2.0 License**, consistent with the LeRobot project.

src/lerobot/policies/rtc/action_queue.py

@@ -1,219 +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.
-
-"""Action queue management for Real-Time Chunking (RTC).
-
-This module provides ActionQueue, a thread-safe queue for managing action chunks
-in real-time control scenarios. It supports both RTC-enabled and non-RTC modes,
-handling action merging and leftover tracking.
-"""
-
-import logging
-from threading import Lock
-
-import torch
-from torch import Tensor
-
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-
-logger = logging.getLogger(__name__)
-
-
-class ActionQueue:
-    """Thread-safe queue for managing action chunks in real-time control.
-
-    This queue handles two types of action sequences:
-    - Original actions: Used for RTC to compute leftovers from previous chunks
-    - Processed actions: Post-processed actions ready for robot execution
-
-    The queue operates in two modes:
-    1. RTC-enabled: Replaces the entire queue with new actions, accounting for inference delay
-    2. RTC-disabled: Appends new actions to the queue, maintaining continuity
-
-    Args:
-        cfg (RTCConfig): Configuration for Real-Time Chunking behavior.
-
-    Attributes:
-        queue (Tensor | None): Processed actions for robot rollout (time_steps, action_dim).
-        original_queue (Tensor | None): Original actions for RTC computation (time_steps, action_dim).
-        last_index (int): Current consumption index in the queue.
-    """
-
-    def __init__(self, cfg: RTCConfig):
-        """Initialize the action queue.
-
-        Args:
-            cfg: RTC configuration controlling queue behavior.
-        """
-        self.queue = None  # Processed actions for robot rollout
-        self.original_queue = None  # Original actions for RTC
-        self.lock = Lock()
-        self.last_index = 0
-        self.cfg = cfg
-
-    def get(self) -> Tensor | None:
-        """Get the next action from the queue.
-
-        Returns:
-            Tensor | None: The next action (action_dim,) or None if queue is empty.
-                          Returns a clone to prevent external modifications.
-        """
-        with self.lock:
-            if self.queue is None or self.last_index >= len(self.queue):
-                return None
-
-            action = self.queue[self.last_index]
-            self.last_index += 1
-            return action.clone()
-
-    def qsize(self) -> int:
-        """Get the number of remaining actions in the queue.
-
-        Returns:
-            int: Number of unconsumed actions.
-        """
-        if self.queue is None:
-            return 0
-        length = len(self.queue)
-        return length - self.last_index
-
-    def empty(self) -> bool:
-        """Check if the queue is empty.
-
-        Returns:
-            bool: True if no actions remain, False otherwise.
-        """
-        if self.queue is None:
-            return True
-
-        length = len(self.queue)
-        return length - self.last_index <= 0
-
-    def get_action_index(self) -> int:
-        """Get the current action consumption index.
-
-        Returns:
-            int: Index of the next action to be consumed.
-        """
-        return self.last_index
-
-    def get_left_over(self) -> Tensor | None:
-        """Get leftover original actions for RTC prev_chunk_left_over.
-
-        These are the unconsumed actions from the current chunk, which will be
-        used by RTC to compute corrections for the next chunk.
-
-        Returns:
-            Tensor | None: Remaining original actions (remaining_steps, action_dim),
-                          or None if no original queue exists.
-        """
-        with self.lock:
-            if self.original_queue is None:
-                return None
-            return self.original_queue[self.last_index :]
-
-    def merge(
-        self,
-        original_actions: Tensor,
-        processed_actions: Tensor,
-        real_delay: int,
-        action_index_before_inference: int | None = 0,
-    ):
-        """Merge new actions into the queue.
-
-        This method operates differently based on RTC mode:
-        - RTC enabled: Replaces the queue, accounting for inference delay
-        - RTC disabled: Appends to the queue, maintaining continuity
-
-        Args:
-            original_actions: Unprocessed actions from policy (time_steps, action_dim).
-            processed_actions: Post-processed actions for robot (time_steps, action_dim).
-            real_delay: Number of time steps of inference delay.
-            action_index_before_inference: Index before inference started, for validation.
-        """
-        with self.lock:
-            self._check_delays(real_delay, action_index_before_inference)
-
-            if self.cfg.enabled:
-                self._replace_actions_queue(original_actions, processed_actions, real_delay)
-                return
-
-            self._append_actions_queue(original_actions, processed_actions)
-
-    def _replace_actions_queue(self, original_actions: Tensor, processed_actions: Tensor, real_delay: int):
-        """Replace the queue with new actions (RTC mode).
-
-        Discards the first `real_delay` actions since they correspond to the time
-        spent during inference, when the robot was executing previous actions.
-
-        Args:
-            original_actions: Unprocessed actions from policy.
-            processed_actions: Post-processed actions for robot.
-            real_delay: Number of time steps to skip due to inference delay.
-        """
-        self.original_queue = original_actions[real_delay:].clone()
-        self.queue = processed_actions[real_delay:].clone()
-
-        logger.debug(f"original_actions shape: {self.original_queue.shape}")
-        logger.debug(f"processed_actions shape: {self.queue.shape}")
-        logger.debug(f"real_delay: {real_delay}")
-
-        self.last_index = 0
-
-    def _append_actions_queue(self, original_actions: Tensor, processed_actions: Tensor):
-        """Append new actions to the queue (non-RTC mode).
-
-        Removes already-consumed actions and appends new ones, maintaining
-        queue continuity without replacement.
-
-        Args:
-            original_actions: Unprocessed actions from policy.
-            processed_actions: Post-processed actions for robot.
-        """
-        if self.queue is None:
-            self.original_queue = original_actions.clone()
-            self.queue = processed_actions.clone()
-            return
-
-        self.original_queue = torch.cat([self.original_queue, original_actions.clone()])
-        self.original_queue = self.original_queue[self.last_index :]
-
-        self.queue = torch.cat([self.queue, processed_actions.clone()])
-        self.queue = self.queue[self.last_index :]
-
-        self.last_index = 0
-
-    def _check_delays(self, real_delay: int, action_index_before_inference: int | None = None):
-        """Validate that computed delays match expectations.
-
-        Compares the delay computed from inference latency with the actual
-        number of actions consumed during inference.
-
-        Args:
-            real_delay: Delay computed from inference latency.
-            action_index_before_inference: Action index when inference started.
-        """
-        if action_index_before_inference is None:
-            return
-
-        indexes_diff = self.last_index - action_index_before_inference
-        if indexes_diff != real_delay:
-            # Let's check that action index difference (real delay calculated based on action queue)
-            # is the same as delay calculated based on inference latency
-            logger.warning(
-                f"[ACTION_QUEUE] Indexes diff is not equal to real delay. "
-                f"Indexes diff: {indexes_diff}, real delay: {real_delay}"
-            )

src/lerobot/policies/rtc/configuration_rtc.py

@@ -1,55 +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.
-
-"""
-Real Time Chunking (RTC) and Bidirectional Decoding (BID) configuration classes.
-
-Based on:
-- Real Time Chunking: https://www.physicalintelligence.company/research/real_time_chunking
-"""
-
-from dataclasses import dataclass
-
-from lerobot.configs.types import RTCAttentionSchedule
-
-
-@dataclass
-class RTCConfig:
-    """Configuration for Real Time Chunking (RTC) inference.
-
-    RTC improves real-time inference by treating chunk generation as an inpainting problem,
-    strategically handling overlapping timesteps between action chunks using prefix attention.
-    """
-
-    # Infrastructure
-    enabled: bool = False
-
-    # Core RTC settings
-    # Todo change to exp
-    prefix_attention_schedule: RTCAttentionSchedule = RTCAttentionSchedule.LINEAR
-    max_guidance_weight: float = 10.0
-    execution_horizon: int = 10
-
-    # Debug settings
-    debug: bool = False
-    debug_maxlen: int = 100
-
-    def __post_init__(self):
-        """Validate RTC configuration parameters."""
-        if self.max_guidance_weight <= 0:
-            raise ValueError(f"max_guidance_weight must be positive, got {self.max_guidance_weight}")
-        if self.debug_maxlen <= 0:
-            raise ValueError(f"debug_maxlen must be positive, got {self.debug_maxlen}")

src/lerobot/policies/rtc/debug_tracker.py

@@ -1,233 +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.
-
-"""Debug information handler for Real-Time Chunking (RTC)."""
-
-from dataclasses import dataclass, field
-from typing import Any
-
-import torch
-from torch import Tensor
-
-
-@dataclass
-class DebugStep:
-    """Container for debug information from a single denoising step.
-
-    Attributes:
-        step_idx (int): Step index/counter.
-        x_t (Tensor | None): Current latent/state tensor.
-        v_t (Tensor | None): Velocity from denoiser.
-        x1_t (Tensor | None): Denoised prediction (x_t - time * v_t).
-        correction (Tensor | None): Correction gradient tensor.
-        err (Tensor | None): Weighted error term.
-        weights (Tensor | None): Prefix attention weights.
-        guidance_weight (float | Tensor | None): Applied guidance weight.
-        time (float | Tensor | None): Time parameter.
-        inference_delay (int | None): Inference delay parameter.
-        execution_horizon (int | None): Execution horizon parameter.
-        metadata (dict[str, Any]): Additional metadata.
-    """
-
-    step_idx: int = 0
-    x_t: Tensor | None = None
-    v_t: Tensor | None = None
-    x1_t: Tensor | None = None
-    correction: Tensor | None = None
-    err: Tensor | None = None
-    weights: Tensor | None = None
-    guidance_weight: float | Tensor | None = None
-    time: float | Tensor | None = None
-    inference_delay: int | None = None
-    execution_horizon: int | None = None
-    metadata: dict[str, Any] = field(default_factory=dict)
-
-    def to_dict(self, include_tensors: bool = False) -> dict[str, Any]:
-        """Convert debug step to dictionary.
-
-        Args:
-            include_tensors (bool): If True, include tensor values. If False, only include
-                tensor statistics (shape, mean, std, min, max).
-
-        Returns:
-            Dictionary representation of the debug step.
-        """
-        result = {
-            "step_idx": self.step_idx,
-            "guidance_weight": (
-                self.guidance_weight.item()
-                if isinstance(self.guidance_weight, Tensor)
-                else self.guidance_weight
-            ),
-            "time": self.time.item() if isinstance(self.time, Tensor) else self.time,
-            "inference_delay": self.inference_delay,
-            "execution_horizon": self.execution_horizon,
-            "metadata": self.metadata.copy(),
-        }
-
-        # Add tensor information
-        tensor_fields = ["x_t", "v_t", "x1_t", "correction", "err", "weights"]
-        for field_name in tensor_fields:
-            tensor = getattr(self, field_name)
-            if tensor is not None:
-                if include_tensors:
-                    result[field_name] = tensor.detach().cpu()
-                else:
-                    result[f"{field_name}_stats"] = {
-                        "shape": tuple(tensor.shape),
-                        "mean": tensor.mean().item(),
-                        "std": tensor.std().item(),
-                        "min": tensor.min().item(),
-                        "max": tensor.max().item(),
-                    }
-
-        return result
-
-
-class Tracker:
-    """Collects and manages debug information for RTC processing.
-
-    This tracker stores debug information from recent denoising steps in a dictionary,
-    using time as the key for efficient lookups and updates.
-
-    Args:
-        enabled (bool): Whether debug collection is enabled.
-        maxlen (int | None): Optional sliding window size. If provided, only the
-            most recent ``maxlen`` debug steps are kept. If ``None``, keeps all.
-    """
-
-    def __init__(self, enabled: bool = False, maxlen: int = 100):
-        self.enabled = enabled
-        self._steps = {} if enabled else None  # Dictionary with time as key
-        self._maxlen = maxlen
-        self._step_counter = 0
-
-    def reset(self) -> None:
-        """Clear all recorded debug information."""
-        if self.enabled and self._steps is not None:
-            self._steps.clear()
-        self._step_counter = 0
-
-    @torch._dynamo.disable
-    def track(
-        self,
-        time: float | Tensor,
-        x_t: Tensor | None = None,
-        v_t: Tensor | None = None,
-        x1_t: Tensor | None = None,
-        correction: Tensor | None = None,
-        err: Tensor | None = None,
-        weights: Tensor | None = None,
-        guidance_weight: float | Tensor | None = None,
-        inference_delay: int | None = None,
-        execution_horizon: int | None = None,
-        **metadata,
-    ) -> None:
-        """Track debug information for a denoising step at a given time.
-
-        If a step with the given time already exists, it will be updated with the new data.
-        Otherwise, a new step will be created. Only non-None fields are updated/set.
-
-        Note: This method is excluded from torch.compile to avoid graph breaks from
-        operations like .item() which are incompatible with compiled graphs.
-
-        Args:
-            time (float | Tensor): Time parameter - used as the key to identify the step.
-            x_t (Tensor | None): Current latent/state tensor.
-            v_t (Tensor | None): Velocity from denoiser.
-            x1_t (Tensor | None): Denoised prediction.
-            correction (Tensor | None): Correction gradient tensor.
-            err (Tensor | None): Weighted error term.
-            weights (Tensor | None): Prefix attention weights.
-            guidance_weight (float | Tensor | None): Applied guidance weight.
-            inference_delay (int | None): Inference delay parameter.
-            execution_horizon (int | None): Execution horizon parameter.
-            **metadata: Additional metadata to store.
-        """
-        if not self.enabled:
-            return
-
-        # Convert time to float and round to avoid float precision issues
-        time_value = time.item() if isinstance(time, Tensor) else time
-        time_key = round(time_value, 6)  # Use rounded time as dictionary key
-
-        # Check if step with this time already exists
-        if time_key in self._steps:
-            # Update existing step with non-None fields
-            existing_step = self._steps[time_key]
-            if x_t is not None:
-                existing_step.x_t = x_t.detach().clone()
-            if v_t is not None:
-                existing_step.v_t = v_t.detach().clone()
-            if x1_t is not None:
-                existing_step.x1_t = x1_t.detach().clone()
-            if correction is not None:
-                existing_step.correction = correction.detach().clone()
-            if err is not None:
-                existing_step.err = err.detach().clone()
-            if weights is not None:
-                existing_step.weights = weights.detach().clone()
-            if guidance_weight is not None:
-                existing_step.guidance_weight = guidance_weight
-            if inference_delay is not None:
-                existing_step.inference_delay = inference_delay
-            if execution_horizon is not None:
-                existing_step.execution_horizon = execution_horizon
-            if metadata:
-                existing_step.metadata.update(metadata)
-        else:
-            # Create new step
-            step = DebugStep(
-                step_idx=self._step_counter,
-                x_t=x_t.detach().clone() if x_t is not None else None,
-                v_t=v_t.detach().clone() if v_t is not None else None,
-                x1_t=x1_t.detach().clone() if x1_t is not None else None,
-                correction=correction.detach().clone() if correction is not None else None,
-                err=err.detach().clone() if err is not None else None,
-                weights=weights.detach().clone() if weights is not None else None,
-                guidance_weight=guidance_weight,
-                time=time_value,
-                inference_delay=inference_delay,
-                execution_horizon=execution_horizon,
-                metadata=metadata,
-            )
-
-            # Add to dictionary
-            self._steps[time_key] = step
-            self._step_counter += 1
-
-            # Enforce maxlen if set
-            if self._maxlen is not None and len(self._steps) > self._maxlen:
-                # Remove oldest entry (first key in dict - Python 3.7+ preserves insertion order)
-                oldest_key = next(iter(self._steps))
-                del self._steps[oldest_key]
-
-    def get_all_steps(self) -> list[DebugStep]:
-        """Get all recorded debug steps.
-
-        Returns:
-            List of all DebugStep objects (may be empty if disabled).
-        """
-        if not self.enabled or self._steps is None:
-            return []
-
-        return list(self._steps.values())
-
-    def __len__(self) -> int:
-        """Return the number of recorded debug steps."""
-        if not self.enabled or self._steps is None:
-            return 0
-        return len(self._steps)

src/lerobot/policies/rtc/debug_visualizer.py

@@ -1,113 +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.
-
-"""Visualization utilities for RTC debug information."""
-
-import torch
-
-
-class RTCDebugVisualizer:
-    """Visualizer for RTC debug information.
-
-    This class provides methods to visualize debug information collected by the Tracker,
-    including corrections, errors, weights, and guidance weights over denoising steps.
-    """
-
-    @staticmethod
-    def plot_waypoints(
-        axes,
-        tensor,
-        start_from: int = 0,
-        color: str = "blue",
-        label: str = "",
-        alpha: float = 0.7,
-        linewidth: float = 2,
-        marker: str | None = None,
-        markersize: int = 4,
-    ):
-        """Plot trajectories across multiple dimensions.
-
-        This function plots a tensor's values across time for multiple dimensions,
-        with each dimension plotted on a separate axis.
-
-        Args:
-            axes: Array of matplotlib axes (one for each dimension).
-            tensor: The tensor to plot (can be torch.Tensor or numpy array).
-                   Shape should be (time_steps, num_dims) or (batch, time_steps, num_dims).
-            start_from: Starting index for the x-axis.
-            color: Color for the plot lines.
-            label: Label for the plot legend.
-            alpha: Transparency level for the plot.
-            linewidth: Width of the plot lines.
-            marker: Marker style for data points (e.g., 'o', 's', '^').
-            markersize: Size of the markers.
-        """
-        import numpy as np
-
-        # Handle None tensor
-        if tensor is None:
-            return
-
-        # Convert tensor to numpy if needed
-        tensor_np = tensor.detach().cpu().numpy() if isinstance(tensor, torch.Tensor) else tensor
-
-        # Handle different tensor shapes
-        if tensor_np.ndim == 3:
-            # If batch dimension present, take first batch
-            tensor_np = tensor_np[0]
-        elif tensor_np.ndim == 1:
-            # If 1D, reshape to (time_steps, 1)
-            tensor_np = tensor_np.reshape(-1, 1)
-
-        # Get dimensions
-        time_steps, num_dims = tensor_np.shape
-
-        # Create x-axis indices
-        x_indices = np.arange(start_from, start_from + time_steps)
-
-        # Plot each dimension on its corresponding axis
-        num_axes = len(axes) if hasattr(axes, "__len__") else 1
-        for dim_idx in range(min(num_dims, num_axes)):
-            ax = axes[dim_idx] if hasattr(axes, "__len__") else axes
-
-            # Plot the trajectory
-            if marker:
-                ax.plot(
-                    x_indices,
-                    tensor_np[:, dim_idx],
-                    color=color,
-                    label=label if dim_idx == 0 else "",  # Only show label once
-                    alpha=alpha,
-                    linewidth=linewidth,
-                    marker=marker,
-                    markersize=markersize,
-                )
-            else:
-                ax.plot(
-                    x_indices,
-                    tensor_np[:, dim_idx],
-                    color=color,
-                    label=label if dim_idx == 0 else "",  # Only show label once
-                    alpha=alpha,
-                    linewidth=linewidth,
-                )
-
-            # Add grid and labels if not already present
-            if not ax.xaxis.get_label().get_text():
-                ax.set_xlabel("Step", fontsize=10)
-            if not ax.yaxis.get_label().get_text():
-                ax.set_ylabel(f"Dim {dim_idx}", fontsize=10)
-            ax.grid(True, alpha=0.3)

src/lerobot/policies/rtc/latency_tracker.py

@@ -1,72 +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.
-
-"""Latency tracking utilities for Real-Time Chunking (RTC)."""
-
-from collections import deque
-
-import numpy as np
-
-
-class LatencyTracker:
-    """Tracks recent latencies and provides max/percentile queries.
-
-    Args:
-        maxlen (int | None): Optional sliding window size. If provided, only the
-            most recent ``maxlen`` latencies are kept. If ``None``, keeps all.
-    """
-
-    def __init__(self, maxlen: int = 100):
-        self._values = deque(maxlen=maxlen)
-        self.reset()
-
-    def reset(self) -> None:
-        """Clear all recorded latencies."""
-        self._values.clear()
-        self.max_latency = 0.0
-
-    def add(self, latency: float) -> None:
-        """Add a latency sample (seconds)."""
-        # Ensure numeric and non-negative
-        val = float(latency)
-
-        if val < 0:
-            return
-        self._values.append(val)
-        self.max_latency = max(self.max_latency, val)
-
-    def __len__(self) -> int:
-        return len(self._values)
-
-    def max(self) -> float | None:
-        """Return the maximum latency or None if empty."""
-        return self.max_latency
-
-    def percentile(self, q: float) -> float | None:
-        """Return the q-quantile (q in [0,1]) of recorded latencies or None if empty."""
-        if not self._values:
-            return 0.0
-        q = float(q)
-        if q <= 0.0:
-            return min(self._values)
-        if q >= 1.0:
-            return self.max_latency
-        vals = np.array(list(self._values), dtype=np.float32)
-        return float(np.quantile(vals, q))
-
-    def p95(self) -> float | None:
-        """Return the 95th percentile latency or None if empty."""
-        return self.percentile(0.95)

src/lerobot/policies/rtc/modeling_rtc.py

@@ -1,297 +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.
-
-"""
-Real-Time Chunking (RTC) implementation for LeRobot.
-
-Based on Physical Intelligence's Kinetix implementation:
-https://github.com/Physical-Intelligence/real-time-chunking-kinetix/blob/main/src/model.py#L214
-"""
-
-import logging
-import math
-
-import torch
-from torch import Tensor
-
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.debug_tracker import Tracker
-
-logger = logging.getLogger(__name__)
-
-
-class RTCProcessor:
-    """Real-Time Chunking processor for action chunking policies.
-
-    This class implements RTC techniques including velocity calculation,
-    prefix attention, and adaptive chunk processing.
-    """
-
-    def __init__(self, rtc_config: RTCConfig):
-        self.rtc_config = rtc_config
-
-        self.tracker = None
-
-        if rtc_config.debug:
-            self.tracker = Tracker(
-                enabled=rtc_config.debug,
-                maxlen=rtc_config.debug_maxlen,
-            )
-
-    # ====================== Tracker Proxy Methods ======================
-    def track(
-        self,
-        time: float | Tensor,
-        x_t: Tensor | None = None,
-        v_t: Tensor | None = None,
-        x1_t: Tensor | None = None,
-        correction: Tensor | None = None,
-        err: Tensor | None = None,
-        weights: Tensor | None = None,
-        guidance_weight: float | Tensor | None = None,
-        inference_delay: int | None = None,
-        execution_horizon: int | None = None,
-        **metadata,
-    ) -> None:
-        """Proxy method to track debug information.
-
-        If tracker is None or disabled, this method does nothing.
-        Otherwise, it forwards the call to tracker.track().
-        """
-        if self.tracker is not None:
-            self.tracker.track(
-                time=time,
-                x_t=x_t,
-                v_t=v_t,
-                x1_t=x1_t,
-                correction=correction,
-                err=err,
-                weights=weights,
-                guidance_weight=guidance_weight,
-                inference_delay=inference_delay,
-                execution_horizon=execution_horizon,
-                **metadata,
-            )
-
-    def get_all_debug_steps(self) -> list:
-        """Get all debug steps from tracker.
-
-        Returns empty list if tracker is disabled or None.
-        """
-        if self.tracker is not None:
-            return self.tracker.get_all_steps()
-        return []
-
-    def is_debug_enabled(self) -> bool:
-        """Check if debug tracking is enabled.
-
-        Returns True if tracker exists and is enabled.
-        """
-        return self.tracker is not None and self.tracker.enabled
-
-    def reset_tracker(self) -> None:
-        """Reset the tracker, clearing all recorded steps.
-
-        Does nothing if tracker is None.
-        """
-        if self.tracker is not None:
-            self.tracker.reset()
-
-    # ====================== End Tracker Proxy Methods ======================
-
-    def denoise_step(
-        self,
-        x_t,
-        prev_chunk_left_over,
-        inference_delay,
-        time,
-        original_denoise_step_partial,
-        execution_horizon=None,
-    ) -> Tensor:
-        """RTC guidance wrapper around an existing denoiser.
-
-        This method wraps an original denoising callable that only takes ``x_t`` and
-        returns a base denoised velocity ``v_t``. It then applies Real-Time Chunking
-        (RTC) prefix guidance using the leftover prefix from the previous chunk.
-
-        Args:
-            x_t (Tensor): Current latent/state to denoise. Shape ``(B, T, A)`` or ``(T, A)``.
-            prev_chunk_left_over (Tensor | None): Unexecuted prefix from the previous
-                chunk. Shape ``(B, T_prev, A)`` or ``(T_prev, A)``. If ``None``, no guidance
-                is applied and the method returns ``v_t`` from the original denoiser.
-            inference_delay (int): Number of timesteps from the prefix to use for guidance.
-            time (float | Tensor): Scalar in [0, 1] indicating normalized time. Must be
-                broadcastable with ``x_t``.
-            original_denoise_step_partial (Callable[[Tensor], Tensor]): Callable that
-                computes the base denoised velocity given only ``x_t``.
-            execution_horizon (int | None): Horizon used to build prefix weights. If
-                ``None``, defaults to ``self.rtc_config.execution_horizon``.
-
-        Returns:
-            Tensor: Guided velocity with the same shape as ``v_t``.
-
-        Notes:
-            - If inputs are 2D, a batch dimension is temporarily added and removed at the end.
-            - If ``prev_chunk_left_over`` is shorter than the current chunk length ``T``, it is
-              right-padded with zeros to match ``T``.
-            - Prefix weights are constructed via ``get_prefix_weights(inference_delay, execution_horizon, T)``
-              and broadcast to ``(B, T, A)``.
-            - Guidance correction is computed via autograd using ``x1_t = x_t + time * v_t`` and
-              ``error = (prev_chunk_left_over - x1_t) * weights``.
-            - The final guidance weight is clamped by ``max_guidance_weight`` from the config.
-
-        Reference:
-            https://www.physicalintelligence.company/download/real_time_chunking.pdf
-        """
-
-        # In the original implementation, the time goes from 0 to 1 and
-        # In our implementation, the time goes from 1 to 0
-        # So we need to invert the time
-        tau = 1 - time
-
-        if prev_chunk_left_over is None:
-            # First step, no guidance - return v_t
-            v_t = original_denoise_step_partial(x_t)
-            return v_t
-
-        x_t = x_t.clone().detach()
-
-        squeezed = False
-        if len(x_t.shape) < 3:
-            # Add batch dimension
-            x_t = x_t.unsqueeze(0)
-            squeezed = True
-
-        if len(prev_chunk_left_over.shape) < 3:
-            # Add batch dimension
-            prev_chunk_left_over = prev_chunk_left_over.unsqueeze(0)
-
-        if execution_horizon is None:
-            execution_horizon = self.rtc_config.execution_horizon
-
-        # If the previous action chunk is to short then it doesn't make sense to use long execution horizon
-        # because there is nothing to merge
-        if execution_horizon > prev_chunk_left_over.shape[1]:
-            execution_horizon = prev_chunk_left_over.shape[1]
-
-        batch_size = x_t.shape[0]
-        action_chunk_size = x_t.shape[1]
-        action_dim = x_t.shape[2]
-
-        if prev_chunk_left_over.shape[1] < action_chunk_size or prev_chunk_left_over.shape[2] < action_dim:
-            padded = torch.zeros(batch_size, action_chunk_size, action_dim).to(x_t.device)
-            padded[:, : prev_chunk_left_over.shape[1], : prev_chunk_left_over.shape[2]] = prev_chunk_left_over
-            prev_chunk_left_over = padded
-
-        assert prev_chunk_left_over.shape == x_t.shape, (
-            "The padded previous chunk must be the same size as the input tensor"
-        )
-
-        weights = (
-            self.get_prefix_weights(inference_delay, execution_horizon, action_chunk_size)
-            .to(x_t.device)
-            .unsqueeze(0)
-            .unsqueeze(-1)
-        )
-
-        with torch.enable_grad():
-            v_t = original_denoise_step_partial(x_t)
-            x_t.requires_grad_(True)
-
-            x1_t = x_t - time * v_t  # noqa: N806
-            err = (prev_chunk_left_over - x1_t) * weights
-            grad_outputs = err.clone().detach()
-            correction = torch.autograd.grad(x1_t, x_t, grad_outputs, retain_graph=False)[0]
-
-        max_guidance_weight = torch.as_tensor(self.rtc_config.max_guidance_weight)
-        tau_tensor = torch.as_tensor(tau)
-        squared_one_minus_tau = (1 - tau_tensor) ** 2
-        inv_r2 = (squared_one_minus_tau + tau_tensor**2) / (squared_one_minus_tau)
-        c = torch.nan_to_num((1 - tau_tensor) / tau_tensor, posinf=max_guidance_weight)
-        guidance_weight = torch.nan_to_num(c * inv_r2, posinf=max_guidance_weight)
-        guidance_weight = torch.minimum(guidance_weight, max_guidance_weight)
-
-        result = v_t - guidance_weight * correction
-
-        # Remove the batch dimension if it was added
-        if squeezed:
-            result = result.squeeze(0)
-            correction = correction.squeeze(0)
-            x1_t = x1_t.squeeze(0)
-            err = err.squeeze(0)
-
-        self.track(
-            time=time,
-            x1_t=x1_t,
-            correction=correction,
-            err=err,
-            weights=weights,
-            guidance_weight=guidance_weight,
-            inference_delay=inference_delay,
-            execution_horizon=execution_horizon,
-        )
-
-        return result
-
-    def get_prefix_weights(self, start, end, total):
-        start = min(start, end)
-
-        if self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.ZEROS:
-            weights = torch.zeros(total)
-            weights[:start] = 1.0
-        elif self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.ONES:
-            weights = torch.ones(total)
-            weights[end:] = 0.0
-        elif self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR:
-            lin_weights = self._linweights(start, end, total)
-            weights = self._add_trailing_zeros(lin_weights, total, end)
-            weights = self._add_leading_ones(weights, start, total)
-        elif self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.EXP:
-            lin_weights = self._linweights(start, end, total)
-            lin_weights = lin_weights * torch.expm1(lin_weights).div(math.e - 1)
-            weights = self._add_trailing_zeros(lin_weights, total, end)
-            weights = self._add_leading_ones(weights, start, total)
-
-        return weights
-
-    def _linweights(self, start, end, total):
-        skip_steps_at_end = max(total - end, 0)
-
-        linspace_steps = total - skip_steps_at_end - start
-
-        if end <= start or linspace_steps <= 0:
-            return torch.tensor([])
-
-        return torch.linspace(1, 0, linspace_steps + 2)[1:-1]
-
-    def _add_trailing_zeros(self, weights, total, end):
-        zeros_len = total - end
-
-        if zeros_len <= 0:
-            return weights
-
-        zeros = torch.zeros(zeros_len)
-        return torch.cat([weights, zeros])
-
-    def _add_leading_ones(self, weights, start, total):
-        ones_len = min(start, total)
-
-        if ones_len <= 0:
-            return weights
-
-        ones = torch.ones(ones_len)
-        return torch.cat([ones, weights])

src/lerobot/policies/sarm/configuration_sarm.py

@@ -1,186 +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, field
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import PolicyFeature, FeatureType, NormalizationMode
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-
-
-@PreTrainedConfig.register_subclass("sarm")
-@dataclass
-class SARMConfig(PreTrainedConfig):
-    """Configuration class for SARM (Stage-Aware Reward Modeling)"""
-    
-    # CLIP params
-    image_dim: int = 512 
-    text_dim: int = 512
-    num_frames: int = 9  # 1 initial + 8 consecutive frames
-    frame_gap: int = 30  # Frame gap between frames (at 30 fps = 1 second)
-    
-    # Architecture params
-    hidden_dim: int = 768  
-    num_heads: int = 12  
-    num_layers: int = 8  
-    max_state_dim: int = 32
-    num_stages: int = 5  # Number of task stages (auto-updated from annotations if available)
-    subtask_names: list | None = None  # List of subtask names (auto-populated from annotations)
-    temporal_proportions: list | None = None  # Temporal proportions for each stage (auto-computed from annotations)
-    max_length: int = num_frames  # Maximum video sequence length (matches num_frames)
-    use_temporal_sampler: bool = True  # Always enable temporal sequence loading
-    
-    # Training params
-    batch_size: int = 64
-    clip_batch_size: int = 64  # Batch size for CLIP encoding
-    dropout: float = 0.1
-    stage_loss_weight: float = 1.0  # Weight for stage classification loss when using subtask annotations
-    
-    pretrained_model_path: str | None = None
-    device: str | None = None
-    
-    # Processor settings
-    image_key: str = "observation.images.top"  # Key for image used from the dataset
-
-    # State key in the dataset (for normalization)
-    state_key: str = "observation.state"
-    
-    # Populated by the processor (video_features, state_features, text_features)
-    input_features: dict = field(default_factory=lambda: {})
-    
-    # Output features
-    output_features: dict = field(default_factory=lambda: {
-        "stage": PolicyFeature(shape=(9, 5), type=FeatureType.REWARD),
-        "progress": PolicyFeature(shape=(9, 1), type=FeatureType.REWARD),
-    })
-
-    normalization_mapping: dict[str, NormalizationMode] = field(
-        default_factory=lambda: {
-            "VISUAL": NormalizationMode.IDENTITY,
-            "STATE": NormalizationMode.MEAN_STD,
-            "LANGUAGE": NormalizationMode.IDENTITY,
-            "REWARD": NormalizationMode.IDENTITY,
-        }
-    )
-    
-    def __post_init__(self):
-        super().__post_init__()
-
-        # Add the image_key as VISUAL
-        if self.image_key:
-            self.input_features[self.image_key] = PolicyFeature(
-                shape=(480, 640, 3),
-                type=FeatureType.VISUAL
-            )
-        
-        # Add state_key as STATE
-        self.input_features[self.state_key] = PolicyFeature(
-            shape=(self.max_state_dim,),  # Single frame state, temporal sampling handles sequence
-            type=FeatureType.STATE
-        )
-        
-        # Update output features with actual dimensions
-        self.output_features["stage"] = PolicyFeature(
-            shape=(self.num_frames, self.num_stages), 
-            type=FeatureType.REWARD
-        )
-        self.output_features["progress"] = PolicyFeature(
-            shape=(self.num_frames, 1), 
-            type=FeatureType.REWARD
-        )
-        
-        # Validate configuration
-        if self.hidden_dim % self.num_heads != 0:
-            raise ValueError(
-                f"hidden_dim ({self.hidden_dim}) must be divisible by num_heads ({self.num_heads})"
-            )
-        
-        if self.max_length != self.num_frames:
-            raise ValueError(
-                f"max_length ({self.max_length}) must equal num_frames ({self.num_frames})"
-            )
-        
-        if self.num_stages < 2:
-            raise ValueError(f"num_stages must be at least 2, got {self.num_stages}")
-    
-    def get_optimizer_preset(self) -> AdamWConfig:
-        """Get default optimizer configuration for SARM training."""
-        return AdamWConfig(
-            lr=5e-5,
-            weight_decay=1e-3,
-            betas=(0.9, 0.999),
-            eps=1e-8,
-        )
-    
-    def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
-        """Get default learning rate scheduler configuration."""
-        return CosineDecayWithWarmupSchedulerConfig(
-            peak_lr=5e-5,
-            decay_lr=5e-6,
-            num_warmup_steps=500,
-            num_decay_steps=50000,
-        )
-    
-    def validate_features(self) -> None:
-        """Validate input and output features."""
-        pass
-    
-    @property
-    def observation_delta_indices(self) -> list[int]:
-        """Load frames for SARM temporal sampling with SYMMETRIC/BIDIRECTIONAL pattern.
-        
-        The model uses 9 frames with symmetric context around current frame:
-        - Frame 0: Initial frame of the episode (clamped via large negative delta)
-        - Frames 1-8: Symmetric context: 4 before + current + 3 after
-        
-        Pattern: [initial, t-4*gap, t-3*gap, t-2*gap, t-gap, t, t+gap, t+2*gap, t+3*gap]
-        
-        Boundary handling (done by dataset loader):
-        - Early frames: backward indices clamp to 0 (first frame)
-        - Late frames: forward indices clamp to episode end (last frame)
-        
-        This enables truly uniform sampling across entire episodes.
-        
-        Returns:
-            9 delta indices: [-1_000_000, -4*gap, -3*gap, -2*gap, -gap, 0, gap, 2*gap, 3*gap]
-        """
-        initial_frame_delta = -1_000_000
-        
-        # Symmetric pattern: 4 frames before, current (0), 3 frames after = 8 context frames
-        symmetric_deltas = [
-            -4 * self.frame_gap,
-            -3 * self.frame_gap,
-            -2 * self.frame_gap,
-            -1 * self.frame_gap,
-            0,  # current frame
-            1 * self.frame_gap,
-            2 * self.frame_gap,
-            3 * self.frame_gap,
-        ]
-        
-        return [initial_frame_delta] + symmetric_deltas
-    
-    @property
-    def action_delta_indices(self) -> None:
-        """SARM is a reward model, not an action policy."""
-        return None
-    
-    @property
-    def reward_delta_indices(self) -> None:
-        """SARM doesn't use delta rewards."""
-        return None
-

src/lerobot/policies/sarm/modeling_sarm.py

@@ -1,650 +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 logging
-from typing import List, Union, Optional
-import random
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from PIL import Image
-from transformers import CLIPModel, CLIPProcessor
-from torch import Tensor
-
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
-from lerobot.policies.sarm.sarm_utils import compute_cumulative_progress_batch, pad_state_to_max_dim
-from lerobot.policies.pretrained import PreTrainedPolicy
-
-class SARMTransformer(nn.Module):
-    """
-    SARM Transformer model for stage-aware reward prediction.
-    
-    This model has a dual-head architecture:
-    1. Stage estimator: Predicts the high-level task stage (classification)
-    2. Subtask estimator: Predicts fine-grained progress within the stage (regression)
-    """
-    
-    def __init__(
-        self,
-        video_dim: int = 512,  
-        text_dim: int = 512, 
-        max_state_dim: int = 32, 
-        hidden_dim: int = 768,
-        num_heads: int = 12,
-        num_layers: int = 8,
-        num_stages: int = 5,
-        max_length: int = 9,
-        dropout: float = 0.1,
-        temporal_proportions: list[float] | None = None
-    ):
-        super().__init__()
-        self.hidden_dim = hidden_dim
-        self.max_length = max_length
-        self.num_stages = num_stages
-        self.max_state_dim = max_state_dim
-        
-        if temporal_proportions is None:
-            raise ValueError(
-                "temporal_proportions is required for SARM. "
-                "Provide subtask annotations in your dataset or set temporal_proportions in config."
-            )
-        
-        # ᾱ_k: proportion for each stage
-        alpha = torch.tensor(temporal_proportions, dtype=torch.float32)
-        
-        # P_k: cumulative proportion up to stage k (P_0 = 0)
-        cumulative = torch.zeros(num_stages + 1, dtype=torch.float32)
-        cumulative[1:] = torch.cumsum(alpha, dim=0)
-        self.register_buffer('alpha', alpha)
-        self.register_buffer('cumulative_prior', cumulative)
-        
-        self.video_proj = nn.Linear(video_dim, hidden_dim)
-        self.text_proj = nn.Linear(text_dim, hidden_dim)
-        self.state_proj = nn.Linear(max_state_dim, hidden_dim) 
-        
-        # Position embedding only for the first frame
-        self.first_pos_embed = nn.Parameter(torch.randn(1, hidden_dim))
-        
-        encoder_layer = nn.TransformerEncoderLayer(
-            d_model=hidden_dim,
-            nhead=num_heads,
-            dim_feedforward=hidden_dim * 4,
-            dropout=dropout,
-            batch_first=True
-        )
-        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
-        
-        # Stage estimator head (classification)
-        self.stage_head = nn.Sequential(
-            nn.Linear(hidden_dim, 512),
-            nn.LayerNorm(512),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(512, num_stages)
-        )
-        
-        # Subtask estimator head (regression)
-        self.stage_embedding = nn.Embedding(num_stages, hidden_dim // 4)
-        subtask_input_dim = hidden_dim + hidden_dim // 4
-        self.subtask_head = nn.Sequential(
-            nn.Linear(subtask_input_dim, 512),
-            nn.LayerNorm(512),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(512, 1),
-            nn.Sigmoid()
-        )
-        
-        # Attention mask
-        self.register_buffer("attention_mask", None, persistent=False)
-    
-    def _get_attention_mask(self, seq_length: int, device: torch.device) -> torch.Tensor:
-        """Generate or retrieve cached causal attention mask."""
-        if self.attention_mask is None or self.attention_mask.shape[0] != seq_length:
-            # Create causal mask
-            mask = nn.Transformer.generate_square_subsequent_mask(seq_length, device=device)
-            self.attention_mask = mask
-        return self.attention_mask
-    
-    def forward(
-        self, 
-        video_frames: torch.Tensor, 
-        text_embed: torch.Tensor,
-        state_features: Optional[torch.Tensor] = None
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Forward pass through the SARM transformer.
-        
-        Args:
-            video_frames: Video frame embeddings (batch_size, seq_len, video_dim)
-            text_embed: Text embeddings (batch_size, text_dim)
-            state_features: Joint state features (batch_size, seq_len, state_dim)
-            
-        Returns:
-            Tuple of:
-                - Stage logits for each frame (batch_size, seq_len, num_stages)
-                - Stage probabilities (batch_size, seq_len, num_stages)
-                - Progress predictions for each frame (batch_size, seq_len, 1)
-        """        
-        # Project inputs to common dimension
-        video_embed = self.video_proj(video_frames)  # [batch_size, seq_len, hidden_dim]
-        text_embed = self.text_proj(text_embed).unsqueeze(1)  # [batch_size, 1, hidden_dim]
-
-        # Pad state features to max_state_dim before projection
-        state_features_padded = pad_state_to_max_dim(state_features, self.max_state_dim)
-
-        state_embed = self.state_proj(state_features_padded)  # [batch_size, seq_len, hidden_dim]
-
-        # Fuse video and state features 
-        video_embed = video_embed + state_embed
-        
-        # Add positional embedding to first video frame
-        video_embed[:, 0] += self.first_pos_embed
-        
-        # Combine sequence: [text, video_frames]
-        sequence = torch.cat([text_embed, video_embed], dim=1)
-        
-        # Get causal attention mask
-        seq_length = sequence.shape[1]
-        attention_mask = self._get_attention_mask(seq_length, sequence.device)
-        
-        # Pass through transformer with causal masking
-        transformed = self.transformer(sequence, mask=attention_mask, is_causal=True)
-        
-        # Get frame features
-        frame_features = transformed[:, 1:]  # [batch_size, seq_len, hidden_dim]
-        
-        # Stage estimation
-        stage_logits = self.stage_head(frame_features)  # [batch_size, seq_len, num_stages]
-        stage_probs = F.softmax(stage_logits, dim=-1)  # [batch_size, seq_len, num_stages]
-        
-        # Get predicted stage indices
-        stage_indices = torch.argmax(stage_probs, dim=-1)  # [batch_size, seq_len]
-        
-        # Get stage embeddings for conditioning
-        stage_embeds = self.stage_embedding(stage_indices) 
-        
-        # Concatenate frame features with stage embeddings
-        conditioned_features = torch.cat([frame_features, stage_embeds], dim=-1)
-        
-        # Subtask progress estimation (conditioned on stage)
-        # τ̂ = within-subtask progress (0-1)
-        tau_preds = self.subtask_head(conditioned_features)  # [batch_size, seq_len, 1]
-        
-        # Convert τ̂ to cumulative progress ŷ using Paper Formula (2):
-        # ŷ = P_{k-1} + ᾱ_k × τ̂
-        progress_preds = compute_cumulative_progress_batch(
-            tau_preds, stage_indices, self.alpha, self.cumulative_prior
-        )
-        
-        return stage_logits, stage_probs, progress_preds
-
-
-class SARMRewardModel(PreTrainedPolicy):
-    """
-    SARM Reward Model for stage-aware task completion rewards.
-    
-    Per SARM paper (Appendix A.4): "We employ a frozen clip-vit-base-patch32 encoder 
-    to process both RGB image sequences and task descriptions."
-    
-    This model combines:
-    - CLIP for encoding video frames AND text descriptions
-    - SARMTransformer for predicting task stage and progress
-    - Optional RA-BC (Reward-Aligned Behavior Cloning) for weighted training
-    """
-    
-    name = "sarm"
-    config_class = SARMConfig
-    
-    def __init__(self, config: SARMConfig, dataset_stats: dict | None = None, dataset_meta=None):
-        super().__init__(config, dataset_stats)
-        config.validate_features() 
-        self.config = config
-        self.dataset_stats = dataset_stats
-        self.device = torch.device(config.device if config.device else "cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Load temporal proportions from dataset
-        if config.temporal_proportions is None and dataset_meta is not None:
-            self._load_temporal_proportions(dataset_meta)
-        
-        logging.info("Loading CLIP encoder")
-        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", use_fast=True)
-        self.clip_model.to(self.device)
-        self.clip_model.eval()
-        
-        self.sarm_transformer = SARMTransformer(
-            video_dim=config.image_dim,
-            text_dim=config.text_dim,
-            max_state_dim=config.max_state_dim,
-            hidden_dim=config.hidden_dim,
-            num_heads=config.num_heads,
-            num_layers=config.num_layers,
-            num_stages=config.num_stages,
-            max_length=config.max_length,
-            dropout=config.dropout,
-            temporal_proportions=config.temporal_proportions
-        )
-        self.sarm_transformer.to(self.device)
-        logging.info(f"SARM initialized on {self.device}")
-    
-    def _load_temporal_proportions(self, dataset_meta) -> None:
-        """
-        Load pre-computed temporal proportions from dataset metadata JSON file.
-
-        The temporal proportions are computed during dataset annotation using SARM Paper Formula (1):
-            ᾱ_k = (1/M) × Σ_i (L_{i,k} / T_i)
-        """
-        import json
-        
-        proportions_path = dataset_meta.root / "meta" / "temporal_proportions.json"
-        
-        if not proportions_path.exists():
-            raise ValueError(
-                f"Temporal proportions not found at {proportions_path}. "
-                "Run the subtask annotation tool first to compute and save temporal proportions."
-            )
-        
-        with open(proportions_path, "r") as f:
-            temporal_proportions_dict = json.load(f)
-        
-        # Sort subtask names for consistent ordering
-        subtask_names = sorted(temporal_proportions_dict.keys())
-        
-        self.config.num_stages = len(subtask_names)
-        self.config.subtask_names = subtask_names
-        self.config.temporal_proportions = [temporal_proportions_dict[name] for name in subtask_names]
-        
-        logging.info(f"Loaded {len(subtask_names)} subtasks: {subtask_names}")
-        logging.info(f"Temporal proportions: {temporal_proportions_dict}")
-    
-    def to(self, device):
-        """Override to method to ensure all components move together."""
-        super().to(device)
-        self.device = device if isinstance(device, torch.device) else torch.device(device)
-        self.clip_model.to(device)
-        self.sarm_transformer.to(device)
-        return self
-    
-    @torch.no_grad()
-    def encode_images(self, images: np.ndarray) -> np.ndarray:
-        """
-        Encode video frames using CLIP.
-        
-        Args:
-            images: Video frames with shape (num_videos, num_frames, H, W, C) in uint8.
-                   Can also be (num_frames, H, W, C) for a single video.
-                   
-        Returns:
-            Encoded image features (num_videos, num_frames, 512) or (num_frames, 512).
-        """
-        # Handle single video case
-        single_video = False
-        if len(images.shape) == 4:
-            images = images[np.newaxis, ...]
-            single_video = True
-        
-        assert len(images.shape) == 5, f"Expected 5D input (num_videos, num_frames, H, W, C), got {images.shape}"
-        
-        all_embeddings = []
-        
-        for video in images:
-            video_embeddings = []
-            
-            # Convert frames to PIL images for CLIP processor
-            frames = []
-            for frame in video:
-                if frame.shape[0] == 3:  # Channel first
-                    frame = frame.transpose(1, 2, 0)
-                if frame.dtype != np.uint8:
-                    frame = (frame * 255).astype(np.uint8) if frame.max() <= 1.0 else frame.astype(np.uint8)
-                frames.append(Image.fromarray(frame))
-            
-            # Batch process frames with CLIP
-            for i in range(0, len(frames), self.config.clip_batch_size):
-                batch = frames[i:i + self.config.clip_batch_size]
-                inputs = self.clip_processor(images=batch, return_tensors="pt")
-                inputs = {k: v.to(self.device) for k, v in inputs.items()}
-                
-                # Get image embeddings from CLIP
-                embeddings = self.clip_model.get_image_features(**inputs).detach().cpu()
-                
-                # Handle single frame case
-                if embeddings.dim() == 1:
-                    embeddings = embeddings.unsqueeze(0)
-                
-                video_embeddings.append(embeddings)
-            
-            video_embeddings = torch.cat(video_embeddings)
-            all_embeddings.append(video_embeddings)
-        
-        result = torch.stack(all_embeddings).numpy()
-        
-        if single_video:
-            result = result[0]
-        
-        return result
-    
-    @torch.no_grad()
-    def encode_text(self, text: Union[str, List[str]]) -> np.ndarray:
-        """
-        Encode text using CLIP text encoder (per SARM paper A.4).
-        
-        Args:
-            text: Text string or list of text strings.
-            
-        Returns:
-            Encoded text features (batch_size, 512) or (512,) for single text.
-        """
-        if isinstance(text, str):
-            text = [text]
-            single_text = True
-        else:
-            single_text = False
-        
-        # Use CLIP's tokenizer directly (avoids image processor validation issues)
-        tokenizer = self.clip_processor.tokenizer
-        
-        # Process in batches
-        all_embeddings = []
-        for i in range(0, len(text), self.config.batch_size):
-            batch_text = text[i:i + self.config.batch_size]
-            
-            inputs = tokenizer(batch_text, return_tensors="pt", padding=True, truncation=True)
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
-            
-            text_embeddings = self.clip_model.get_text_features(**inputs)
-            all_embeddings.append(text_embeddings.cpu())
-        
-        result = torch.cat(all_embeddings).numpy()
-        
-        if single_text:
-            result = result[0]
-        
-        return result
-    
-    @torch.no_grad()
-    def calculate_rewards(
-        self,
-        text_embeddings: Union[np.ndarray, torch.Tensor],
-        video_embeddings: Union[np.ndarray, torch.Tensor],
-        state_features: Optional[Union[np.ndarray, torch.Tensor]] = None,
-        return_all_frames: bool = False,
-        return_stages: bool = False
-    ) -> Union[np.ndarray, tuple]:
-        """
-        Calculate rewards for given text, video, and state representations.
-        
-        Args:
-            text_embeddings: Encoded text representations (batch_size, 512)
-            video_embeddings: Encoded video representations (batch_size, num_frames, 512)
-            state_features: Joint state features (batch_size, num_frames, state_dim)
-            return_all_frames: If True, return rewards for all frames
-            return_stages: If True, also return stage predictions
-            
-        Returns:
-            If return_stages=False:
-                Reward values (batch_size,) or (batch_size, num_frames)
-            If return_stages=True:
-                Tuple of (rewards, stage_probs)
-        """
-        if isinstance(text_embeddings, np.ndarray):
-            text_embeddings = torch.tensor(text_embeddings, dtype=torch.float32)
-        if isinstance(video_embeddings, np.ndarray):
-            video_embeddings = torch.tensor(video_embeddings, dtype=torch.float32)
-        if state_features is not None and isinstance(state_features, np.ndarray):
-            state_features = torch.tensor(state_features, dtype=torch.float32)
-        
-        # Handle single sample case
-        if text_embeddings.dim() == 1:
-            text_embeddings = text_embeddings.unsqueeze(0)
-            video_embeddings = video_embeddings.unsqueeze(0)
-            if state_features is not None:
-                state_features = state_features.unsqueeze(0)
-            single_sample = True
-        else:
-            single_sample = False
-        
-        # Process in batches
-        all_rewards = []
-        all_stage_probs = []
-        
-        for i in range(0, len(video_embeddings), self.config.batch_size):
-            batch_texts = text_embeddings[i:i + self.config.batch_size].to(self.device)
-            batch_videos = video_embeddings[i:i + self.config.batch_size].to(self.device)
-            batch_states = None
-            if state_features is not None:
-                batch_states = state_features[i:i + self.config.batch_size].to(self.device)
-            
-            # Get predictions
-            stage_logits, stage_probs, progress_preds = self.sarm_transformer(
-                batch_videos.float(), batch_texts.float(), batch_states.float() if batch_states is not None else None
-            )
-            
-            if return_all_frames:
-                all_rewards.append(progress_preds.squeeze(-1).cpu())
-            else:
-                # Return only last frame reward
-                all_rewards.append(progress_preds[:, -1, 0].cpu())
-            
-            if return_stages:
-                all_stage_probs.append(stage_probs.cpu())
-        
-        rewards = torch.cat(all_rewards).numpy()
-        
-        if single_sample:
-            rewards = rewards[0] if not return_all_frames else rewards[0]
-        
-        if return_stages:
-            stage_probs = torch.cat(all_stage_probs).numpy()
-            if single_sample:
-                stage_probs = stage_probs[0]
-            return rewards, stage_probs
-        
-        return rewards
-    
-    def train(self, mode: bool = True):
-        """Overwrite train method to ensure CLIP encoder stays frozen during training"""
-        super().train(mode)
-        self.clip_model.eval()
-        self.sarm_transformer.train(mode)
-        return self
-    
-    def eval(self):
-        """Overwrite eval method to ensure CLIP encoder stays frozen during evaluation"""
-        return self.train(False)
-    
-    def parameters(self):
-        """Override to return trainable parameters (only SARM transformer, not CLIP encoder)."""
-        return self.sarm_transformer.parameters()
-    
-    def get_optim_params(self):
-        """Override to return optimizer parameters (only SARM transformer, not CLIP encoder)."""
-        return self.parameters()
-    
-    def reset(self):
-        """Required by PreTrainedPolicy but not used for reward models."""
-        pass
-    
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
-        """Required by PreTrainedPolicy but not used for reward models."""
-        raise NotImplementedError("SARM model does not predict action chunks")
-    
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
-        """Required by PreTrainedPolicy but not used for SARM."""
-        raise NotImplementedError("SARM model does not select actions")
-    
-    def _apply_temporal_augmentation(
-        self, 
-        video: torch.Tensor, 
-        progress: torch.Tensor, 
-        state: torch.Tensor | None,
-        max_length: int,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        """Apply temporal augmentation by appending reversed frames (SARM paper A.4).
-        
-        This helps the model learn to handle non-monotonic progress (failures, recoveries).
-        Appends 1-4 reversed frames to simulate going backwards in task progress.
-        """
-        num_reverse = random.randint(1, min(4, max_length - 1))
-        
-        # Reverse and take frames (skip first which is last of original)
-        reversed_video = video.flip(0)[1:num_reverse + 1]
-        reversed_progress = progress.flip(0)[1:num_reverse + 1]
-        
-        # Concatenate and trim
-        video = torch.cat([video, reversed_video], dim=0)[:max_length]
-        progress = torch.cat([progress, reversed_progress], dim=0)[:max_length]
-        
-        if state is not None:
-            reversed_state = state.flip(0)[1:num_reverse + 1]
-            state = torch.cat([state, reversed_state], dim=0)[:max_length]
-        
-        return video, progress, state
-    
-    def _ensure_sequence_length(self, tensor: torch.Tensor, target_len: int) -> torch.Tensor:
-        """Pad or trim tensor to target length."""
-        current_len = tensor.shape[0]
-        if current_len == target_len:
-            return tensor
-        if current_len < target_len:
-            padding = target_len - current_len
-            return torch.cat([tensor, tensor[-1:].expand(padding, *tensor.shape[1:])])
-        return tensor[:target_len]
-    
-    def forward(self, batch):
-        """
-        Forward pass for SARM reward model training.
-        
-        Uses annotation-based progress targets following SARM paper Eq. 2:
-        yt = Pk-1 + α̅k × τt
-        where:
-        - τt = (t - sk) / (ek - sk) is within-subtask normalized time
-        - Pk-1 is cumulative prior (sum of previous subtask proportions)
-        - α̅k is the temporal proportion for subtask k
-        
-        Args:
-            batch: Dictionary with 'observation' containing:
-                - 'video_features': (B, T, 512) pre-encoded video features
-                - 'text_features': (B, 512) pre-encoded text features (CLIP)
-                - 'state_features': (B, T, state_dim) joint state features
-                - 'stage_labels': (B, T) stage labels from annotations
-                - 'progress_targets': (B, T, 1) progress targets from annotations
-        
-        Returns:
-            Tuple of (total_loss, output_dict with loss components)
-        """
-        observation = batch.get('observation', batch)
-        
-        # Extract required features
-        video_features = observation['video_features'].to(self.device)
-        text_features = observation['text_features'].to(self.device)
-        state_features = observation.get('state_features').to(self.device)
-        
-        batch_size = video_features.shape[0]
-        max_length = self.config.num_frames
-        
-        # Ensure 3D video features (B, T, D)
-        if video_features.dim() == 2:
-            video_features = video_features.unsqueeze(1).expand(-1, max_length, -1)
-        if state_features is not None and state_features.dim() == 2:
-            state_features = state_features.unsqueeze(1).expand(-1, max_length, -1)
-        
-        # Get annotation-based progress targets (required for SARM paper formula)
-        progress_from_annotations = observation.get('progress_targets')
-        if progress_from_annotations is None:
-            raise ValueError("progress_targets from annotations is required for SARM training")
-        
-        progress_from_annotations = progress_from_annotations.to(self.device)
-        if progress_from_annotations.dim() == 2:
-            progress_from_annotations = progress_from_annotations.unsqueeze(-1)
-        if progress_from_annotations.dim() == 3 and progress_from_annotations.shape[0] == 1:
-            progress_from_annotations = progress_from_annotations.expand(batch_size, -1, -1)
-        
-        # Process each sample: apply temporal REWIND augmentation 
-        processed_videos = []
-        processed_states = []
-        progress_targets = []
-        
-        for i in range(batch_size):
-            video = video_features[i]
-            state = state_features[i] if state_features is not None else None
-            progress = progress_from_annotations[i].squeeze(-1)  # (T,)
-            
-            # Apply temporal REWIND augmentation with 50% probability: appends up to 4 reversed frames to simulate failures/recoveries
-            if random.random() < 0.5:
-                video, progress, state = self._apply_temporal_augmentation(video, progress, state, max_length)
-            
-            # Ensure correct sequence length
-            video = self._ensure_sequence_length(video, max_length)
-            progress = self._ensure_sequence_length(progress.unsqueeze(-1), max_length).squeeze(-1)
-            if state is not None:
-                state = self._ensure_sequence_length(state, max_length)
-            
-            processed_videos.append(video)
-            progress_targets.append(progress)
-            if state is not None:
-                processed_states.append(state)
-        
-        # Stack into batches
-        processed_videos = torch.stack(processed_videos)
-        progress_targets = torch.stack(progress_targets).unsqueeze(-1)  # (B, T, 1)
-        processed_states = torch.stack(processed_states) if processed_states else None
-        
-        # Get model predictions
-        stage_logits, stage_probs, progress_preds = self.sarm_transformer(
-            processed_videos, text_features, processed_states
-        )
-        
-        # Compute progress loss (MSE)
-        progress_loss = F.mse_loss(progress_preds, progress_targets)
-        output_dict = {'progress_loss': progress_loss.item()}
-        total_loss = progress_loss
-        
-        # Compute stage loss (cross-entropy)
-        stage_labels = observation.get('stage_labels')
-        if stage_labels is None:
-            raise ValueError("stage_labels from annotations is required for SARM training")
-        
-        stage_labels = stage_labels.to(self.device)
-        if stage_labels.dim() == 1:
-            stage_labels = stage_labels.unsqueeze(0).expand(batch_size, -1)
-        stage_loss = compute_stage_loss(stage_logits, stage_labels)
-        total_loss = total_loss + self.config.stage_loss_weight * stage_loss
-        output_dict['stage_loss'] = stage_loss.item()
-        
-        # Misaligned loss: 20% probability
-        if random.random() < 0.2:
-            shuffle_idx = torch.randperm(batch_size, device=self.device)
-            _, _, misaligned_preds = self.sarm_transformer(
-                processed_videos, text_features[shuffle_idx], processed_states
-            )
-            misaligned_loss = F.mse_loss(misaligned_preds, torch.zeros_like(misaligned_preds))
-            total_loss = total_loss + misaligned_loss
-            output_dict['misaligned_loss'] = misaligned_loss.item()
-        
-        output_dict['total_loss'] = total_loss.item()
-        return total_loss, output_dict
-
-def compute_stage_loss(stage_logits: torch.Tensor, target_stages: torch.Tensor) -> torch.Tensor:
-    _, _, num_stages = stage_logits.shape
-    stage_logits_flat = stage_logits.reshape(-1, num_stages)
-    target_stages_flat = target_stages.reshape(-1)
-    
-    loss = F.cross_entropy(stage_logits_flat, target_stages_flat)
-    return loss

src/lerobot/policies/sarm/processor_sarm.py

@@ -1,644 +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 numpy as np
-import torch
-from PIL import Image
-import pandas as pd
-from transformers import CLIPModel, CLIPProcessor
-
-from lerobot.processor.core import TransitionKey
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
-from lerobot.policies.sarm.sarm_utils import compute_tau, compute_cumulative_progress_batch, pad_state_to_max_dim
-from lerobot.processor import (
-    ProcessorStep,
-    PolicyProcessorPipeline,
-    PolicyAction,
-    DeviceProcessorStep,
-    AddBatchDimensionProcessorStep,
-    NormalizerProcessorStep,
-)
-from lerobot.processor.converters import (
-    policy_action_to_transition,
-    transition_to_policy_action,
-    from_tensor_to_numpy,
-)
-from lerobot.processor.pipeline import PipelineFeatureType
-from lerobot.processor.core import EnvTransition, TransitionKey
-from lerobot.configs.types import PolicyFeature, FeatureType
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-
-
-class SARMEncodingProcessorStep(ProcessorStep):
-    """ProcessorStep that encodes images and text with CLIP."""
-    def __init__(
-        self,
-        config: SARMConfig,
-        image_key: str | None = None,
-        dataset_meta = None,
-        dataset_stats: dict | None = None,
-    ):
-        super().__init__()
-        self.config = config
-        self.image_key = image_key or config.image_key
-        self.dataset_meta = dataset_meta
-        self.dataset_stats = dataset_stats
-        self.temporal_proportions = {name: prop for name, prop in zip(self.config.subtask_names, self.config.temporal_proportions)}
-        self.subtask_names = self.config.subtask_names
-
-        self.device = torch.device(
-            self.config.device if self.config.device 
-            else "cuda" if torch.cuda.is_available() else "cpu"
-        )
-        
-        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", use_fast=True)
-        self.clip_model.to(self.device)
-        self.clip_model.eval()
-
-    def _find_episode_for_frame(self, frame_idx: int) -> int:
-        """Find the episode index for a given frame index."""
-        for ep_idx in range(len(self.dataset_meta.episodes)):
-            ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
-            ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
-            if ep_start <= frame_idx < ep_end:
-                return ep_idx
-        return 0  
-    
-    def _get_episode_indices(self, frame_indices: np.ndarray, episode_index) -> np.ndarray:
-        """Get episode indices for each frame index."""
-        if episode_index is None:
-            return np.array([self._find_episode_for_frame(int(f)) for f in frame_indices])
-        
-        episode_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(episode_index)))
-        
-        # If single episode but multiple frames, compute episode for each frame
-        if len(episode_indices) == 1 and len(frame_indices) > 1:
-            return np.array([self._find_episode_for_frame(int(f)) for f in frame_indices])
-        
-        return episode_indices
-    
-    def _compute_absolute_indices(self, frame_idx: int, ep_start: int, ep_end: int, num_frames: int) -> torch.Tensor:
-        """Compute absolute frame indices for symmetric bidirectional pattern.
-        
-        Pattern: [ep_start, t-4*gap, t-3*gap, t-2*gap, t-gap, t, t+gap, t+2*gap, t+3*gap]
-        
-        Boundary handling:
-        - Backward indices clamp to ep_start (first frame)
-        - Forward indices clamp to ep_end - 1 (last frame)
-        """
-        indices = []
-        indices.append(ep_start)  # Initial frame is always episode start
-        
-        # Symmetric context: 4 before, current, 3 after
-        num_before = 4
-        num_after = 3
-        last_valid_frame = ep_end - 1
-        
-        # Frames before current (clamp to first frame)
-        for i in range(num_before, 0, -1):
-            idx = max(ep_start, frame_idx - i * self.config.frame_gap)
-            indices.append(idx)
-        
-        # Current frame
-        indices.append(frame_idx)
-        
-        # Frames after current (clamp to last frame)
-        for i in range(1, num_after + 1):
-            idx = min(last_valid_frame, frame_idx + i * self.config.frame_gap)
-            indices.append(idx)
-        
-        return torch.tensor(indices)
-    
-    def _compute_episode_metadata(
-        self, 
-        frame_indices: np.ndarray, 
-        episode_indices: np.ndarray,
-        num_frames: int,
-    ) -> tuple[list | torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Compute episode metadata for all samples.
-        
-        Returns:
-            Tuple of (absolute_frame_indices, remaining_lengths, episode_lengths)
-        """
-        absolute_indices_list = []
-        remaining_lengths = []
-        episode_lengths = []
-        
-        for ep_idx, frame_idx in zip(episode_indices.tolist(), frame_indices.tolist()):
-            ep_idx, frame_idx = int(ep_idx), int(frame_idx)
-            ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
-            ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
-            
-            episode_lengths.append(ep_end - ep_start)
-            abs_indices = self._compute_absolute_indices(frame_idx, ep_start, ep_end, num_frames)
-            absolute_indices_list.append(abs_indices)
-            remaining_lengths.append(ep_end - abs_indices[0].item())
-        
-        return absolute_indices_list, torch.tensor(remaining_lengths), torch.tensor(episode_lengths)
-    
-    def _compute_stage_and_progress_for_frame(
-        self, 
-        current_frame: int,
-        subtask_names: list,
-        subtask_start_frames: list,
-        subtask_end_frames: list,
-        transition_smoothing_frames: int = 15,
-    ) -> tuple[int, float, dict[int, float] | None]:
-        """Compute stage index, cumulative progress, and soft stage labels for a single frame.
-        
-        Implements SARM Paper Formula (2):
-            y_t = P_{k-1} + ᾱ_k × τ_t
-        
-        where:
-            - τ_t = (t - s_k) / (e_k - s_k) is within-subtask progress
-            - P_{k-1} is cumulative prior (sum of previous subtask proportions)
-            - ᾱ_k is the temporal proportion for subtask k
-        
-        Additionally computes soft stage labels near transitions to mitigate discrete jumps
-        in the stage classifier. Near stage boundaries, labels are blended between adjacent
-        stages to encourage smoother predictions.
-        
-        Args:
-            current_frame: Frame index relative to episode start
-            subtask_names: List of subtask names for this episode
-            subtask_start_frames: List of subtask start frames
-            subtask_end_frames: List of subtask end frames
-            transition_smoothing_frames: Number of frames over which to smooth labels near transitions
-            
-        Returns:
-            Tuple of (stage_idx, cumulative_progress, soft_stage_labels)
-            - stage_idx: Hard stage index (for compatibility)
-            - cumulative_progress: Progress value in [0, 1]
-            - soft_stage_labels: Dict mapping stage_idx -> probability, or None if not near transition
-        """
-        # Get temporal proportions as list for compute_cumulative_progress
-        temporal_proportions_list = [
-            self.temporal_proportions.get(name, 0.0) for name in self.subtask_names
-        ]
-        num_stages = len(self.subtask_names)
-        
-        # Find which subtask this frame belongs to
-        for j, (name, start_frame, end_frame) in enumerate(zip(subtask_names, subtask_start_frames, subtask_end_frames)):
-            if current_frame >= start_frame and current_frame <= end_frame:
-                # Found the subtask, get its global index
-                stage_idx = self.subtask_names.index(name) if name in self.subtask_names else 0
-                
-                # Compute τ_t using utility function (Paper Formula 2)
-                tau = compute_tau(current_frame, start_frame, end_frame)
-                
-                # Compute cumulative progress using utility function (Paper Formula 2)
-                cumulative_progress = compute_cumulative_progress_batch(
-                    tau, stage_idx, temporal_proportions_list
-                )
-                
-                # Compute soft stage labels near transitions
-                soft_stage_labels = None
-                frames_from_start = current_frame - start_frame
-                frames_to_end = end_frame - current_frame
-                
-                if frames_from_start < transition_smoothing_frames and j > 0:
-                    # Near start of stage - blend with previous stage
-                    blend = frames_from_start / transition_smoothing_frames
-                    prev_name = subtask_names[j - 1]
-                    prev_stage_idx = self.subtask_names.index(prev_name) if prev_name in self.subtask_names else max(0, stage_idx - 1)
-                    soft_stage_labels = {prev_stage_idx: 1.0 - blend, stage_idx: blend}
-                    
-                elif frames_to_end < transition_smoothing_frames and j < len(subtask_names) - 1:
-                    # Near end of stage - blend with next stage
-                    blend = frames_to_end / transition_smoothing_frames
-                    next_name = subtask_names[j + 1]
-                    next_stage_idx = self.subtask_names.index(next_name) if next_name in self.subtask_names else min(num_stages - 1, stage_idx + 1)
-                    soft_stage_labels = {stage_idx: blend, next_stage_idx: 1.0 - blend}
-                
-                return stage_idx, cumulative_progress, soft_stage_labels
-        
-        # No matching subtask found
-        if current_frame < subtask_start_frames[0]:
-            return 0, 0.0, None
-        elif current_frame > subtask_end_frames[-1]:
-            return len(self.subtask_names) - 1, 1.0, None
-        else:
-            # Between subtasks - use previous subtask's end state (tau = 1.0)
-            for j in range(len(subtask_names) - 1):
-                if current_frame > subtask_end_frames[j] and current_frame < subtask_start_frames[j + 1]:
-                    name = subtask_names[j]
-                    stage_idx = self.subtask_names.index(name) if name in self.subtask_names else j
-                    
-                    # Completed subtask, so tau = 1.0
-                    cumulative_progress = compute_cumulative_progress_batch(
-                        1.0, stage_idx, temporal_proportions_list
-                    )
-                    return stage_idx, cumulative_progress, None
-        
-        return 0, 0.0, None
-    
-    def _compute_labels_for_sample(
-        self,
-        frame_idx: int,
-        ep_idx: int,
-        seq_len: int,
-        episodes_df: pd.DataFrame,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] | tuple[None, None, None]:
-        """Compute stage labels, progress targets, and soft stage labels for symmetric bidirectional pattern.
-        
-        Pattern: [initial, t-4*gap, t-3*gap, t-2*gap, t-gap, t, t+gap, t+2*gap, t+3*gap]
-        
-        Boundary handling:
-        - Before episode start: clamp to frame 0 (progress ~0%)
-        - After episode end: clamp to last frame (progress ~100%)
-        
-        Soft stage labels are computed near stage transitions to mitigate discrete jumps.
-        
-        Args:
-            frame_idx: The frame index for this sample
-            ep_idx: The episode index
-            seq_len: Number of frames in the sequence
-            episodes_df: DataFrame with episode metadata
-            
-        Returns:
-            Tuple of (stage_labels, progress_targets, soft_stage_labels):
-            - stage_labels: (T,) hard stage indices
-            - progress_targets: (T, 1) progress values
-            - soft_stage_labels: (T, num_stages) soft probability labels, or None if no transitions nearby
-        """
-        # Check if episode has valid annotations
-        if ep_idx >= len(episodes_df):
-            return None, None, None
-        
-        subtask_names = episodes_df.loc[ep_idx, 'subtask_names']
-        if subtask_names is None or (isinstance(subtask_names, float) and pd.isna(subtask_names)):
-            return None, None, None
-        
-        subtask_start_frames = episodes_df.loc[ep_idx, 'subtask_start_frames']
-        subtask_end_frames = episodes_df.loc[ep_idx, 'subtask_end_frames']
-        ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
-        ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
-        ep_length = ep_end - ep_start
-        last_valid_frame = ep_length - 1
-        
-        num_stages = len(self.subtask_names)
-        
-        # Generate labels for each frame in the sequence
-        stage_labels = []
-        progress_targets = []
-        soft_labels_list = []  # List of soft label dicts (or None)
-        has_any_soft_labels = False
-        
-        # Symmetric pattern: initial + 4 before + current + 3 after = 9 frames
-        num_before = 4
-        num_after = 3
-        
-        for i in range(seq_len):
-            if i == 0:
-                # Position 0: Initial frame of the episode
-                current_frame = 0  # Relative to episode start
-            elif i <= num_before:
-                # Positions 1-4: frames before current (with clamping to first frame)
-                offset = -(num_before - i + 1) * self.config.frame_gap
-                current_frame = max(0, frame_idx + offset - ep_start)
-            elif i == num_before + 1:
-                # Position 5: current frame
-                current_frame = frame_idx - ep_start
-            else:
-                # Positions 6-8: frames after current (with clamping to last frame)
-                offset = (i - num_before - 1) * self.config.frame_gap
-                current_frame = min(last_valid_frame, frame_idx + offset - ep_start)
-            
-            stage_idx, cumulative_progress, soft_stage_labels = self._compute_stage_and_progress_for_frame(
-                current_frame, subtask_names, subtask_start_frames, subtask_end_frames
-            )
-            
-            stage_labels.append(stage_idx)
-            progress_targets.append(cumulative_progress)
-            soft_labels_list.append(soft_stage_labels)
-            if soft_stage_labels is not None:
-                has_any_soft_labels = True
-        
-        stage_labels = torch.tensor(stage_labels, dtype=torch.long)
-        progress_targets = torch.tensor(progress_targets, dtype=torch.float32).unsqueeze(-1)
-        
-        # Convert soft labels to tensor if any exist
-        soft_stage_labels_tensor = None
-        if has_any_soft_labels:
-            soft_stage_labels_tensor = torch.zeros(seq_len, num_stages, dtype=torch.float32)
-            for i, soft_dict in enumerate(soft_labels_list):
-                if soft_dict is not None:
-                    for stage_idx, prob in soft_dict.items():
-                        soft_stage_labels_tensor[i, stage_idx] = prob
-                else:
-                    # Use hard one-hot label
-                    soft_stage_labels_tensor[i, stage_labels[i]] = 1.0
-        
-        return stage_labels, progress_targets, soft_stage_labels_tensor
-    
-    def _generate_stage_and_progress_labels(self, frame_index, episode_index, video_features):
-        """Generate stage labels, progress targets, and soft stage labels from subtask annotations.
-        
-        Args:
-            frame_index: Current frame index or tensor of indices
-            episode_index: Episode index or tensor of indices  
-            video_features: Video features tensor to determine sequence length
-            
-        Returns:
-            Tuple of (stage_labels, progress_targets, soft_stage_labels) or (None, None, None) if no annotations.
-            - stage_labels: (B, T) hard stage indices
-            - progress_targets: (B, T, 1) progress values
-            - soft_stage_labels: (B, T, num_stages) soft probability labels, or None
-        """
-        if self.temporal_proportions is None or episode_index is None:
-            return None, None, None
-        
-        # Normalize inputs to numpy arrays
-        frame_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(frame_index)))
-        episode_indices = self._get_episode_indices(frame_indices, episode_index)
-        
-        # Determine sequence length
-        if video_features is not None and video_features.dim() >= 2:
-            seq_len = video_features.shape[1]
-        else:
-            seq_len = 1
-        
-        episodes_df = self.dataset_meta.episodes.to_pandas()
-        num_stages = len(self.subtask_names)
-        
-        all_stage_labels = []
-        all_progress_targets = []
-        all_soft_stage_labels = []
-        has_any_soft_labels = False
-        
-        for ep_idx, frame_idx in zip(episode_indices.tolist(), frame_indices.tolist()):
-            stage_labels, progress_targets, soft_labels = self._compute_labels_for_sample(
-                int(frame_idx), int(ep_idx), seq_len, episodes_df
-            )
-            
-            if stage_labels is None:
-                all_stage_labels.append(torch.zeros(seq_len, dtype=torch.long))
-                all_progress_targets.append(torch.zeros(seq_len, 1, dtype=torch.float32))
-                all_soft_stage_labels.append(None)
-            else:
-                all_stage_labels.append(stage_labels)
-                all_progress_targets.append(progress_targets)
-                all_soft_stage_labels.append(soft_labels)
-                if soft_labels is not None:
-                    has_any_soft_labels = True
-        
-        stacked_stage_labels = torch.stack(all_stage_labels, dim=0)
-        stacked_progress_targets = torch.stack(all_progress_targets, dim=0)
-        
-        # Stack soft labels if any exist
-        stacked_soft_labels = None
-        if has_any_soft_labels:
-            soft_labels_tensors = []
-            for i, soft_labels in enumerate(all_soft_stage_labels):
-                if soft_labels is not None:
-                    soft_labels_tensors.append(soft_labels)
-                else:
-                    # Create one-hot from hard labels
-                    one_hot = torch.zeros(seq_len, num_stages, dtype=torch.float32)
-                    for t in range(seq_len):
-                        one_hot[t, all_stage_labels[i][t]] = 1.0
-                    soft_labels_tensors.append(one_hot)
-            stacked_soft_labels = torch.stack(soft_labels_tensors, dim=0)
-        
-        return stacked_stage_labels, stacked_progress_targets, stacked_soft_labels
-    
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        """Encode images, text, and normalize states in the transition."""
-
-        new_transition = transition.copy() if hasattr(transition, 'copy') else dict(transition)
-        observation = new_transition.get(TransitionKey.OBSERVATION)
-        
-        image = observation.get(self.image_key)
-        
-        if isinstance(image, torch.Tensor):
-            image = image.cpu().numpy()
-        video_features = self._encode_images_batch(image)
-        observation['video_features'] = video_features
-        
-        # Extract state and pad to max_state_dim (already normalized by NormalizerProcessorStep)
-        state_key = self.config.state_key
-        state_data = observation.get(state_key)
-        
-        if isinstance(state_data, torch.Tensor):
-            state_tensor = state_data.float()
-        else:
-            state_tensor = torch.tensor(state_data, dtype=torch.float32)
-        
-        observation['state_features'] = pad_state_to_max_dim(state_tensor, self.config.max_state_dim)
-        
-        comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
-        
-        # Get task description from dataset (complementary_data["task"])
-        task = comp_data.get('task')
-        if isinstance(task, list):
-            # If batch, take first task (assuming same task for all items in batch)
-            task = task[0] if task else ""
-        
-        # Encode text with CLIP
-        batch_size = video_features.shape[0]
-        observation['text_features'] = self._encode_text_clip(task, batch_size)
-        
-        frame_index = comp_data.get('index')
-        episode_index = comp_data.get('episode_index')
-        
-        if frame_index is None:
-            raise ValueError("Frame index ('index') not found in COMPLEMENTARY_DATA")
-        if episode_index is None:
-            raise ValueError("Episode index ('episode_index') not found in COMPLEMENTARY_DATA")
-        
-        # Compute episode metadata if dataset_meta is available
-        if self.dataset_meta is not None:
-            frame_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(frame_index)))
-            episode_indices = self._get_episode_indices(frame_indices, episode_index)
-            
-            # Determine number of frames from video features
-            if video_features.dim() >= 2:
-                num_frames = video_features.shape[1]
-            else:
-                num_frames = 1
-            
-            abs_indices, remaining, ep_lengths = self._compute_episode_metadata(
-                frame_indices, episode_indices, num_frames
-            )
-            observation['absolute_frame_indices'] = abs_indices
-            observation['remaining_length'] = remaining
-            observation['episode_length'] = ep_lengths
-        
-        # Generate stage labels, progress targets, and soft stage labels from subtask annotations
-        if self.temporal_proportions is not None and self.dataset_meta is not None:
-            stage_labels, progress_targets, soft_stage_labels = self._generate_stage_and_progress_labels(
-                frame_index, episode_index, video_features
-            )
-            if stage_labels is not None:
-                observation['stage_labels'] = stage_labels
-                observation['progress_targets'] = progress_targets
-                if soft_stage_labels is not None:
-                    observation['soft_stage_labels'] = soft_stage_labels
-        
-        new_transition[TransitionKey.OBSERVATION] = observation
-        return new_transition
-    
-    @torch.no_grad()
-    def _encode_images_batch(self, images: np.ndarray) -> torch.Tensor:
-        """Encode a batch of images using CLIP.
-        
-        Args:
-            images: Batched images with shape: (B, T, C, H, W)
-            
-        Returns:
-            Encoded feature vectors with shape (B, T, 512)
-        """
-
-        batch_size, seq_length = images.shape[0], images.shape[1] 
-        images = images.reshape(batch_size * seq_length, *images.shape[2:]) 
-        
-        # Convert to list of PIL images
-        num_frames = images.shape[0]
-        images_list = []
-        for i in range(num_frames):
-            img = images[i]
-            if img.shape[0] in [1, 3]:  # Channel first (C, H, W)
-                img = img.transpose(1, 2, 0)
-            
-            # Handle single channel
-            if img.shape[-1] == 1:
-                img = np.repeat(img, 3, axis=-1)
-            
-            # Convert to uint8
-            if img.dtype != np.uint8:
-                img = (img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8)
-            
-            images_list.append(Image.fromarray(img))
-        
-        # Encode each batch
-        all_embeddings = []
-        for i in range(0, num_frames, self.config.clip_batch_size):
-            batch_imgs = images_list[i:i + self.config.clip_batch_size]
-            
-            # Process with CLIP
-            inputs = self.clip_processor(images=batch_imgs, return_tensors="pt")
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
-            
-            # Get image embeddings
-            embeddings = self.clip_model.get_image_features(**inputs).detach().cpu()
-            
-            # Handle single frame case
-            if embeddings.dim() == 1:
-                embeddings = embeddings.unsqueeze(0)
-            
-            all_embeddings.append(embeddings)
-        
-        # Concatenate all embeddings
-        all_embeddings = torch.cat(all_embeddings)  # (B*T, 512)
-        
-        # Reshape back 
-        all_embeddings = all_embeddings.reshape(batch_size, seq_length, -1)  # (B, T, 512)
-        
-        return all_embeddings
-    
-    @torch.no_grad()
-    def _encode_text_clip(self, text: str, batch_size: int) -> torch.Tensor:
-        """Encode text using CLIP text encoder (per SARM paper A.4).
-        
-        Args:
-            text: Task description text to encode
-            batch_size: Batch size to replicate for
-            
-        Returns:
-            Encoded text features with shape (B, 512)
-        """
-        # Use CLIP's tokenizer directly for text
-        tokenizer = self.clip_processor.tokenizer
-        inputs = tokenizer([text], return_tensors="pt", padding=True, truncation=True)
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
-        
-        # Get text features from CLIP
-        text_embedding = self.clip_model.get_text_features(**inputs).detach().cpu()
-        
-        # Replicate for batch (B, 512)
-        text_embedding = text_embedding.expand(batch_size, -1)
-        
-        return text_embedding
-    
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """Add encoded features to the observation features."""
-        features[PipelineFeatureType.OBSERVATION]['video_features'] = PolicyFeature(
-            type=FeatureType.VISUAL, 
-            shape=(self.config.num_frames, self.config.image_dim)
-        )
-        features[PipelineFeatureType.OBSERVATION]['text_features'] = PolicyFeature(
-            type=FeatureType.LANGUAGE, 
-            shape=(self.config.text_dim,)
-        )
-        features[PipelineFeatureType.OBSERVATION]['state_features'] = PolicyFeature(
-            type=FeatureType.STATE, 
-            shape=(self.config.num_frames, self.config.max_state_dim)
-        )
-        return features
-
-
-def make_sarm_pre_post_processors(
-    config: SARMConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    dataset_meta = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """
-    Create pre-processor and post-processor pipelines for SARM.
-    
-    The pre-processing pipeline:
-    1. Adds batch dimension
-    2. Normalizes observation.state using NormalizerProcessorStep (MEAN_STD)
-    3. SARMEncodingProcessorStep:
-       - Encodes images with CLIP 
-       - Pads states to max_state_dim
-       - Encodes text with CLIP 
-    4. Moves data to device
-    
-    The post-processing pipeline:
-    1. Moves data to CPU
-    """
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=[
-                    AddBatchDimensionProcessorStep(),
-                    NormalizerProcessorStep(
-                        features={**config.input_features, **config.output_features},
-                        norm_map=config.normalization_mapping,
-                        stats=dataset_stats,
-                    ),
-                    SARMEncodingProcessorStep(
-                        config=config,
-                        dataset_meta=dataset_meta,
-                        dataset_stats=dataset_stats
-                    ),
-                    DeviceProcessorStep(device=config.device),
-                ],
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=[DeviceProcessorStep(device="cpu")],
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )

src/lerobot/policies/sarm/sarm_utils.py

@@ -1,257 +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 numpy as np
-import torch
-import torch.nn.functional as F
-from typing import Sequence, Any
-from pydantic import BaseModel, Field
-
-# Pydantic Models for SARM-style Annotation
-class Timestamp(BaseModel):
-    """Timestamp in MM:SS or SS format"""
-    start: str = Field(description="Start timestamp (MM:SS or just seconds)")
-    end: str = Field(description="End timestamp (MM:SS or just seconds)")
-
-
-class Subtask(BaseModel):
-    """Individual subtask/stage - must use EXACT names from provided list"""
-    name: str = Field(description="Subtask name - MUST match one from the predefined list exactly")
-    timestamps: Timestamp
-
-
-class SubtaskAnnotation(BaseModel):
-    """Complete annotation for a robot manipulation episode"""
-    subtasks: list[Subtask] = Field(description="List of all subtasks in temporal order")
-
-
-def compute_temporal_proportions(annotations: dict[int, Any], fps: int = 30) -> dict[str, float]:
-    """
-    Compute dataset-level temporal proportions (priors) for each subtask.
-    
-    Implements SARM Paper Formula (1):
-        ᾱ_k = (1/M) × Σ_i (L_{i,k} / T_i)
-    
-    where:
-        - M is the number of trajectories (episodes)
-        - L_{i,k} is the duration of subtask k in trajectory i
-        - T_i is the total duration of trajectory i
-    
-    This averages the PROPORTION of each subtask within each trajectory,
-    giving equal weight to all trajectories regardless of their absolute length.
-    
-    Args:
-        annotations: Dict mapping episode index to SubtaskAnnotation object.
-            Each annotation has a .subtasks list where each subtask has:
-            - .name: subtask name
-            - .timestamps.start: start time as "MM:SS" string
-            - .timestamps.end: end time as "MM:SS" string
-        fps: Frames per second (unused, kept for API compatibility)
-        
-    Returns:
-        Dict mapping subtask name to its temporal proportion (ᾱ_k).
-        Proportions are normalized to sum to 1.0.
-    """
-    subtask_proportions: dict[str, list[float]] = {}
-    
-    for annotation in annotations.values():
-        total_duration = 0
-        durations: dict[str, int] = {}
-        
-        for subtask in annotation.subtasks:
-            start_parts = subtask.timestamps.start.split(":")
-            end_parts = subtask.timestamps.end.split(":")
-            
-            start_seconds = int(start_parts[0]) * 60 + int(start_parts[1]) if len(start_parts) == 2 else int(start_parts[0])
-            end_seconds = int(end_parts[0]) * 60 + int(end_parts[1]) if len(end_parts) == 2 else int(end_parts[0])
-            
-            duration = end_seconds - start_seconds
-            durations[subtask.name] = duration
-            total_duration += duration
-        
-        # Calculate L_{i,k} / T_i for each subtask in this trajectory
-        if total_duration > 0:
-            for name, duration in durations.items():
-                if name not in subtask_proportions:
-                    subtask_proportions[name] = []
-                subtask_proportions[name].append(duration / total_duration)
-    
-    if not subtask_proportions:
-        return {}
-    
-    # Average across trajectories: (1/M) × Σ_i (L_{i,k} / T_i)
-    avg_proportions = {
-        name: sum(props) / len(props)
-        for name, props in subtask_proportions.items()
-    }
-    
-    # Normalize to ensure sum = 1
-    total = sum(avg_proportions.values())
-    if total > 0:
-        avg_proportions = {name: prop / total for name, prop in avg_proportions.items()}
-    
-    return avg_proportions
-
-
-def compute_tau(
-    current_frame: int | float,
-    subtask_start: int | float,
-    subtask_end: int | float,
-) -> float:
-    """
-    Compute within-subtask normalized time τ_t.
-    
-    Implements part of SARM Paper Formula (2):
-        τ_t = (t - s_k) / (e_k - s_k) ∈ [0, 1]
-    
-    where:
-        - t is the current frame
-        - s_k is the start frame of subtask k
-        - e_k is the end frame of subtask k
-    
-    Args:
-        current_frame: Current frame index (t)
-        subtask_start: Start frame of the subtask (s_k)
-        subtask_end: End frame of the subtask (e_k)
-        
-    Returns:
-        Within-subtask progress τ_t ∈ [0, 1]
-    """
-    subtask_duration = subtask_end - subtask_start
-    
-    if subtask_duration <= 0:
-        return 1.0
-    
-    tau = (current_frame - subtask_start) / subtask_duration
-    
-    return float(np.clip(tau, 0.0, 1.0))
-
-
-def compute_cumulative_progress_batch(
-    tau: torch.Tensor | float,
-    stage_indices: torch.Tensor | int,
-    alpha: torch.Tensor | Sequence[float],
-    cumulative_prior: torch.Tensor | None = None,
-) -> torch.Tensor | float:
-    """
-    Compute cumulative normalized progress from within-subtask progress.
-    
-    This function implements the core formula used in SARM for both:
-    
-    **Formula 2 (Training labels):**
-        y_t = P_{k-1} + ᾱ_k × τ_t ∈ [0, 1]
-        
-        Used to compute ground-truth progress labels from subtask annotations.
-        - τ_t comes from annotated frame position: τ_t = (t - s_k) / (e_k - s_k)
-        - k is the known subtask from annotations
-        
-    **Formula 4 (Inference predictions):**
-        ŷ_{1:N} = P̂_{k-1, 1:N} + ᾱ_{k, 1:N} × τ̂_{1:N} ∈ [0, 1]
-        
-        Used to convert model outputs to cumulative progress during inference.
-        - τ̂ comes from the subtask MLP head (conditioned on predicted stage)
-        - k = Ŝ is the predicted stage from Formula 3: Ŝ = argmax(softmax(Ψ))
-    
-    The formulas are mathematically identical; only the source of inputs differs:
-    - Training: τ and k from annotations → ground-truth labels
-    - Inference: τ̂ and Ŝ from model → predicted progress
-    
-    where:
-        - P_{k-1} = Σ_{j=1}^{k-1} ᾱ_j is the cumulative prior (sum of previous proportions)
-        - ᾱ_k is the temporal proportion for subtask k (from Formula 1)
-        - τ is within-subtask progress ∈ [0, 1]
-    
-    This ensures:
-        - y at start of subtask k = P_{k-1}
-        - y at end of subtask k = P_k
-    
-    Supports both scalar and batched tensor inputs:
-        - Scalar: tau (float), stage_indices (int), alpha (list/sequence)
-        - Batch: tau (Tensor), stage_indices (Tensor), alpha (Tensor), cumulative_prior (Tensor)
-    
-    Args:
-        tau: Within-subtask progress τ ∈ [0, 1]. 
-             For training: computed from frame position in annotated subtask.
-             For inference: predicted by subtask MLP head.
-             Scalar float or Tensor with shape (..., 1)
-        stage_indices: Index of current subtask k (0-indexed).
-             For training: known from annotations.
-             For inference: predicted via argmax(stage_probs) (Formula 3).
-             Scalar int or Tensor with shape (...)
-        alpha: Temporal proportions ᾱ with shape (num_stages,) or Sequence[float].
-             Computed from dataset annotations using Formula 1.
-        cumulative_prior: Optional. Cumulative priors P with shape (num_stages + 1,)
-             where cumulative_prior[k] = P_k = Σ_{j=1}^{k} ᾱ_j.
-             If None, will be computed from alpha.
-        
-    Returns:
-        Cumulative progress y ∈ [0, 1]. 
-        Scalar float if inputs are scalar, otherwise Tensor with shape (..., 1)
-    """    
-    if not isinstance(tau, torch.Tensor):
-        if not alpha:
-            raise ValueError("alpha (temporal_proportions) cannot be empty")
-        
-        if isinstance(alpha, torch.Tensor):
-            alpha_list = alpha.tolist()
-        else:
-            alpha_list = list(alpha)
-        
-        if stage_indices < 0 or stage_indices >= len(alpha_list):
-            raise ValueError(
-                f"stage_indices {stage_indices} out of range "
-                f"for {len(alpha_list)} subtasks"
-            )
-        
-        # P_{k-1} = sum of proportions for subtasks 0 to k-1
-        P_k_minus_1 = sum(alpha_list[:stage_indices])
-        
-        # ᾱ_k = proportion for current subtask
-        alpha_k = alpha_list[stage_indices]
-        
-        # y_t = P_{k-1} + ᾱ_k × τ_t
-        y_t = P_k_minus_1 + alpha_k * tau
-        
-        return float(np.clip(y_t, 0.0, 1.0))
-    
-    if not isinstance(alpha, torch.Tensor):
-        alpha = torch.tensor(alpha, dtype=torch.float32)
-    
-    # Compute cumulative_prior if not provided
-    if cumulative_prior is None:
-        cumulative_prior = torch.zeros(len(alpha) + 1, dtype=alpha.dtype, device=alpha.device)
-        cumulative_prior[1:] = torch.cumsum(alpha, dim=0)
-    
-    # P_{k-1} for each predicted stage
-    P_k_minus_1 = cumulative_prior[stage_indices]
-    
-    # ᾱ_k for each predicted stage
-    alpha_k = alpha[stage_indices]
-    
-    # ŷ = P_{k-1} + ᾱ_k × τ̂
-    progress = P_k_minus_1.unsqueeze(-1) + alpha_k.unsqueeze(-1) * tau
-    
-    return progress
-
-def pad_state_to_max_dim(state: torch.Tensor, max_state_dim: int) -> torch.Tensor:
-    """Pad the state tensor's last dimension to max_state_dim with zeros."""
-    current_dim = state.shape[-1]
-    if current_dim >= max_state_dim:
-        return state[..., :max_state_dim]  # Truncate if larger
-    
-    # Pad with zeros on the right
-    padding = (0, max_state_dim - current_dim)  # (left, right) for last dim
-    return F.pad(state, padding, mode='constant', value=0)

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.utils.constants import OBS_IMAGES
 
 
@@ -103,9 +102,6 @@ class SmolVLAConfig(PreTrainedConfig):
     min_period: float = 4e-3  # sensitivity range for the timestep used in sine-cosine positional encoding
     max_period: float = 4.0
 
-    # Real-Time Chunking (RTC) configuration
-    rtc_config: RTCConfig | None = None
-
     def __post_init__(self):
         super().__post_init__()
 

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -54,15 +54,12 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 
 import math
 from collections import deque
-from typing import TypedDict
 
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
-from typing_extensions import Unpack
 
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.policies.utils import (
@@ -72,12 +69,6 @@ from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LAN
 from lerobot.utils.utils import get_safe_dtype
 
 
-class ActionSelectKwargs(TypedDict, total=False):
-    inference_delay: int | None
-    prev_chunk_left_over: Tensor | None
-    execution_horizon: int | None
-
-
 def create_sinusoidal_pos_embedding(
     time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
 ) -> Tensor:
@@ -241,8 +232,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
         super().__init__(config)
         config.validate_features()
         self.config = config
-        self.init_rtc_processor()
-        self.model = VLAFlowMatching(config, rtc_processor=self.rtc_processor)
+
+        self.model = VLAFlowMatching(config)
         self.reset()
 
     def reset(self):
@@ -251,28 +242,10 @@ class SmolVLAPolicy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
-    def init_rtc_processor(self):
-        """Initialize RTC processor if RTC is enabled in config."""
-        self.rtc_processor = None
-
-        # Lets create processor if the config provided
-        # If RTC is not enabled - we still can track the denoising data
-        if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
-
-            # In case of calling init_rtc_processor after the model is created
-            # We need to set the rtc_processor to the model
-            # During the normal initialization process the model is not created yet
-            model_value = getattr(self, "model", None)
-            if model_value is not None:
-                model_value.rtc_processor = self.rtc_processor
-
     def get_optim_params(self) -> dict:
         return self.parameters()
 
-    def _get_action_chunk(
-        self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
-    ) -> Tensor:
+    def _get_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
         # TODO: Check if this for loop is needed.
         # Context: In fact, self.queues contains only ACTION field, and in inference, we don't have action in the batch
         # In the case of offline inference, we have the action in the batch
@@ -287,9 +260,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
         lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
         lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
 
-        actions = self.model.sample_actions(
-            images, img_masks, lang_tokens, lang_masks, state, noise=noise, **kwargs
-        )
+        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
 
         # Unpad actions
         original_action_dim = self.config.action_feature.shape[0]
@@ -307,37 +278,30 @@ class SmolVLAPolicy(PreTrainedPolicy):
         return batch
 
     @torch.no_grad()
-    def predict_action_chunk(
-        self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
-    ) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
         self.eval()
 
         batch = self._prepare_batch(batch)
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        actions = self._get_action_chunk(batch, noise, **kwargs)
+        actions = self._get_action_chunk(batch, noise)
         return actions
 
     @torch.no_grad()
-    def select_action(
-        self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
-    ) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
         """Select a single action given environment observations.
 
         This method wraps `select_actions` in order to return one action at a time for execution in the
         environment. It works by managing the actions in a queue and only calling `select_actions` when the
         queue is empty.
         """
-
-        assert not self._rtc_enabled(), (
-            "RTC is not supported for select_action, use it with predict_action_chunk"
-        )
-
         self.eval()
         batch = self._prepare_batch(batch)
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        if self._check_get_actions_condition():
+        # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
+        # querying the policy.
+        if len(self._queues[ACTION]) == 0:
             actions = self._get_action_chunk(batch, noise)
 
             # `self.predict_action_chunk` returns a (batch_size, n_action_steps, action_dim) tensor, but the queue
@@ -346,12 +310,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         return self._queues[ACTION].popleft()
 
-    def _check_get_actions_condition(self) -> bool:
-        return len(self._queues[ACTION]) == 0
-
-    def _rtc_enabled(self) -> bool:
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
     def forward(self, batch: dict[str, Tensor], noise=None, time=None) -> dict[str, Tensor]:
         """Do a full training forward pass to compute the loss"""
         if self.config.adapt_to_pi_aloha:
@@ -513,7 +471,7 @@ class VLAFlowMatching(nn.Module):
     └──────────────────────────────┘
     """
 
-    def __init__(self, config: SmolVLAConfig, rtc_processor: RTCProcessor | None = None):
+    def __init__(self, config: SmolVLAConfig):
         super().__init__()
         self.config = config
 
@@ -527,6 +485,7 @@ class VLAFlowMatching(nn.Module):
             num_vlm_layers=self.config.num_vlm_layers,
             self_attn_every_n_layers=self.config.self_attn_every_n_layers,
             expert_width_multiplier=self.config.expert_width_multiplier,
+            device=self.config.device,
         )
         self.state_proj = nn.Linear(
             self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size
@@ -551,10 +510,6 @@ class VLAFlowMatching(nn.Module):
         self.add_image_special_tokens = self.config.add_image_special_tokens
         self.image_end_token = torch.tensor([self.fake_image_token], dtype=torch.long)
         self.prefix_length = self.config.prefix_length
-        self.rtc_processor = rtc_processor
-
-    def _rtc_enabled(self):
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
 
     def set_requires_grad(self):
         for params in self.state_proj.parameters():
@@ -751,16 +706,7 @@ class VLAFlowMatching(nn.Module):
         losses = F.mse_loss(u_t, v_t, reduction="none")
         return losses
 
-    def sample_actions(
-        self,
-        images,
-        img_masks,
-        lang_tokens,
-        lang_masks,
-        state,
-        noise=None,
-        **kwargs: Unpack[ActionSelectKwargs],
-    ) -> Tensor:
+    def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state, noise=None) -> Tensor:
         """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)"""
         bsize = state.shape[0]
         device = state.device
@@ -788,45 +734,17 @@ class VLAFlowMatching(nn.Module):
 
         x_t = noise
         time = torch.tensor(1.0, dtype=torch.float32, device=device)
-
         while time >= -dt / 2:
             expanded_time = time.expand(bsize)
-
-            # Define a closure function to properly capture expanded_time
-            # This avoids the lambda expression (E731) and loop variable binding (B023) issues
-            def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
-                return self.denoise_step(
-                    x_t=input_x_t,
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
+            v_t = self.denoise_step(
+                prefix_pad_masks,
+                past_key_values,
+                x_t,
+                expanded_time,
+            )
             # Euler step
             x_t += dt * v_t
-
-            # Record x_t and v_t after Euler step (other params are recorded in rtc_processor.denoise_step)
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
             time += dt
-
         return x_t
 
     def denoise_step(

src/lerobot/policies/utils.py

@@ -22,8 +22,6 @@ import numpy as np
 import torch
 from torch import nn
 
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.datasets.utils import build_dataset_frame
 from lerobot.processor import PolicyAction, RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
@@ -200,52 +198,3 @@ def make_robot_action(action_tensor: PolicyAction, ds_features: dict[str, dict])
         f"{name}": float(action_tensor[i]) for i, name in enumerate(action_names)
     }
     return act_processed_policy
-
-
-def raise_feature_mismatch_error(
-    provided_features: set[str],
-    expected_features: set[str],
-) -> None:
-    """
-    Raises a standardized ValueError for feature mismatches between dataset/environment and policy config.
-    """
-    missing = expected_features - provided_features
-    extra = provided_features - expected_features
-    # TODO (jadechoghari): provide a dynamic rename map suggestion to the user.
-    raise ValueError(
-        f"Feature mismatch between dataset/environment and policy config.\n"
-        f"- Missing features: {sorted(missing) if missing else 'None'}\n"
-        f"- Extra features: {sorted(extra) if extra else 'None'}\n\n"
-        f"Please ensure your dataset and policy use consistent feature names.\n"
-        f"If your dataset uses different observation keys (e.g., cameras named differently), "
-        f"use the `--rename_map` argument, for example:\n"
-        f'  --rename_map=\'{{"observation.images.left": "observation.images.camera1", '
-        f'"observation.images.top": "observation.images.camera2"}}\''
-    )
-
-
-def validate_visual_features_consistency(
-    cfg: PreTrainedConfig,
-    features: dict[str, PolicyFeature],
-) -> None:
-    """
-    Validates visual feature consistency between a policy config and provided dataset/environment features.
-    
-    Validation passes if EITHER:
-    - Policy's expected visuals are a subset of dataset (policy uses some cameras, dataset has more)
-    - Dataset's provided visuals are a subset of policy (policy declares extras for flexibility)
-
-    Args:
-        cfg (PreTrainedConfig): The model or policy configuration containing input_features and type.
-        features (Dict[str, PolicyFeature]): A mapping of feature names to PolicyFeature objects.
-    """
-    expected_visuals = {k for k, v in cfg.input_features.items() if v.type == FeatureType.VISUAL}
-    provided_visuals = {k for k, v in features.items() if v.type == FeatureType.VISUAL}
-    
-    # Accept if either direction is a subset 
-    policy_subset_of_dataset = expected_visuals.issubset(provided_visuals)
-    dataset_subset_of_policy = provided_visuals.issubset(expected_visuals)
-    
-    if not (policy_subset_of_dataset or dataset_subset_of_policy):
-        raise_feature_mismatch_error(provided_visuals, expected_visuals)
-        

src/lerobot/processor/converters.py

@@ -170,9 +170,8 @@ def _extract_complementary_data(batch: dict[str, Any]) -> dict[str, Any]:
     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 {}
-    episode_index_key = {"episode_index": batch["episode_index"]} if "episode_index" in batch else {}
 
-    return {**pad_keys, **task_key, **index_key, **task_index_key, **episode_index_key}
+    return {**pad_keys, **task_key, **index_key, **task_index_key}
 
 
 def create_transition(

src/lerobot/processor/env_processor.py

@@ -1,154 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass
-
-import torch
-
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
-
-from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="libero_processor")
-class LiberoProcessorStep(ObservationProcessorStep):
-    """
-    Processes LIBERO observations into the LeRobot format.
-
-    This step handles the specific observation structure from LIBERO environments,
-    which includes nested robot_state dictionaries and image observations.
-
-    **State Processing:**
-    -   Processes the `robot_state` dictionary which contains nested end-effector,
-        gripper, and joint information.
-    -   Extracts and concatenates:
-        - End-effector position (3D)
-        - End-effector quaternion converted to axis-angle (3D)
-        - Gripper joint positions (2D)
-    -   Maps the concatenated state to `"observation.state"`.
-
-    **Image Processing:**
-    -   Rotates images by 180 degrees by flipping both height and width dimensions.
-    -   This accounts for the HuggingFaceVLA/libero camera orientation convention.
-    """
-
-    def _process_observation(self, observation):
-        """
-        Processes both image and robot_state observations from LIBERO.
-        """
-        processed_obs = observation.copy()
-        for key in list(processed_obs.keys()):
-            if key.startswith(f"{OBS_IMAGES}."):
-                img = processed_obs[key]
-
-                # Flip both H and W
-                img = torch.flip(img, dims=[2, 3])
-
-                processed_obs[key] = img
-        # Process robot_state into a flat state vector
-        if "observation.robot_state" in processed_obs:
-            robot_state = processed_obs.pop("observation.robot_state")
-
-            # Extract components
-            eef_pos = robot_state["eef"]["pos"]  # (B, 3,)
-            eef_quat = robot_state["eef"]["quat"]  # (B, 4,)
-            gripper_qpos = robot_state["gripper"]["qpos"]  # (B, 2,)
-
-            # Convert quaternion to axis-angle
-            eef_axisangle = self._quat2axisangle(eef_quat)  # (B, 3)
-            # Concatenate into a single state vector
-            state = torch.cat((eef_pos, eef_axisangle, gripper_qpos), dim=-1)
-
-            # ensure float32
-            state = state.float()
-            if state.dim() == 1:
-                state = state.unsqueeze(0)
-
-            processed_obs[OBS_STATE] = state
-        return processed_obs
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """
-        Transforms feature keys from the LIBERO format to the LeRobot standard.
-        """
-        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = {}
-
-        # copy over non-STATE features
-        for ft, feats in features.items():
-            if ft != PipelineFeatureType.STATE:
-                new_features[ft] = feats.copy()
-
-        # rebuild STATE features
-        state_feats = {}
-
-        # add our new flattened state
-        state_feats["observation.state"] = PolicyFeature(
-            key="observation.state",
-            shape=(8,),  # [eef_pos(3), axis_angle(3), gripper(2)]
-            dtype="float32",
-            description=("Concatenated end-effector position (3), axis-angle (3), and gripper qpos (2)."),
-        )
-
-        new_features[PipelineFeatureType.STATE] = state_feats
-
-        return new_features
-
-    def observation(self, observation):
-        return self._process_observation(observation)
-
-    def _quat2axisangle(self, quat: torch.Tensor) -> torch.Tensor:
-        """
-        Convert batched quaternions to axis-angle format.
-        Only accepts torch tensors of shape (B, 4).
-
-        Args:
-            quat (Tensor): (B, 4) tensor of quaternions in (x, y, z, w) format
-
-        Returns:
-            Tensor: (B, 3) axis-angle vectors
-
-        Raises:
-            TypeError: if input is not a torch tensor
-            ValueError: if shape is not (B, 4)
-        """
-
-        if not isinstance(quat, torch.Tensor):
-            raise TypeError(f"_quat2axisangle expected a torch.Tensor, got {type(quat)}")
-
-        if quat.ndim != 2 or quat.shape[1] != 4:
-            raise ValueError(f"_quat2axisangle expected shape (B, 4), got {tuple(quat.shape)}")
-
-        quat = quat.to(dtype=torch.float32)
-        device = quat.device
-        batch_size = quat.shape[0]
-
-        w = quat[:, 3].clamp(-1.0, 1.0)
-
-        den = torch.sqrt(torch.clamp(1.0 - w * w, min=0.0))
-
-        result = torch.zeros((batch_size, 3), device=device)
-
-        mask = den > 1e-10
-
-        if mask.any():
-            angle = 2.0 * torch.acos(w[mask])  # (M,)
-            axis = quat[mask, :3] / den[mask].unsqueeze(1)
-            result[mask] = axis * angle.unsqueeze(1)
-
-        return result

src/lerobot/robots/openarms/__init__.py

@@ -14,21 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
-from lerobot.policies.sarm.modeling_sarm import (
-    SARMRewardModel,
-    SARMTransformer,
-)
-from lerobot.policies.sarm.processor_sarm import (
-    SARMEncodingProcessorStep,
-    make_sarm_pre_post_processors,
-)
-
-__all__ = [
-    "SARMConfig",
-    "SARMRewardModel",
-    "SARMTransformer",
-    "SARMEncodingProcessorStep",
-    "make_sarm_pre_post_processors",
-]
+from .config_openarms_follower import OpenArmsFollowerConfig
+from .openarms_follower import OpenArmsFollower
 
+__all__ = ["OpenArmsFollower", "OpenArmsFollowerConfig"]

src/lerobot/robots/openarms/config_openarms_follower.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+from lerobot.cameras import CameraConfig
+from lerobot.motors.damiao.tables import MotorType
+
+from ..config import RobotConfig
+
+
+@RobotConfig.register_subclass("openarms_follower")
+@dataclass
+class OpenArmsFollowerConfig(RobotConfig):
+    """Configuration for the OpenArms follower robot with Damiao motors."""
+    
+    # CAN interfaces - one per arm
+    # Right arm CAN interface (e.g., "can0")
+    # Left arm CAN interface (e.g., "can1")
+    # Linux: "can0", "can1", etc.
+    # macOS: "/dev/cu.usbmodem*" (serial device)
+    port_right: str = "can0"    # CAN interface for right arm
+    port_left: str = "can1"     # CAN interface for left arm
+    
+    # CAN interface type: "socketcan" (Linux), "slcan" (macOS/serial), or "auto" (auto-detect)
+    can_interface: str = "socketcan"
+    
+    # CAN FD settings (OpenArms uses CAN FD by default)
+    use_can_fd: bool = True
+    can_bitrate: int = 1000000      # Nominal bitrate (1 Mbps)
+    can_data_bitrate: int = 5000000  # Data bitrate for CAN FD (5 Mbps)
+    
+    # Whether to disable torque when disconnecting
+    disable_torque_on_disconnect: bool = True
+    
+    # Safety limit for relative target positions
+    # Set to a positive scalar for all motors, or a dict mapping motor names to limits
+    max_relative_target: Optional[float | Dict[str, float]] = None
+    
+    # Camera configurations
+    cameras: Dict[str, CameraConfig] = field(default_factory=dict)
+    
+    # Motor configuration for OpenArms (7 DOF per arm)
+    # Maps motor names to (send_can_id, recv_can_id, motor_type)
+    # Based on: https://docs.openarm.dev/software/setup/configure-test
+    # OpenArms uses 4 types of motors:
+    # - DM8009 (DM-J8009P-2EC) for shoulders (high torque)
+    # - DM4340P and DM4340 for shoulder rotation and elbow
+    # - DM4310 (DM-J4310-2EC V1.1) for wrist and gripper
+    motor_config: Dict[str, tuple[int, int, str]] = field(default_factory=lambda: {
+        "joint_1": (0x01, 0x11, "dm8009"),   # J1 - Shoulder pan (DM8009)
+        "joint_2": (0x02, 0x12, "dm8009"),   # J2 - Shoulder lift (DM8009)
+        "joint_3": (0x03, 0x13, "dm4340"),   # J3 - Shoulder rotation (DM4340)
+        "joint_4": (0x04, 0x14, "dm4340"),   # J4 - Elbow flex (DM4340)
+        "joint_5": (0x05, 0x15, "dm4310"),   # J5 - Wrist roll (DM4310)
+        "joint_6": (0x06, 0x16, "dm4310"),   # J6 - Wrist pitch (DM4310)
+        "joint_7": (0x07, 0x17, "dm4310"),   # J7 - Wrist rotation (DM4310)
+        "gripper": (0x08, 0x18, "dm4310"),   # J8 - Gripper (DM4310)
+    })
+    
+    # MIT control parameters for position control (per motor)
+    # Values: [joint_1, joint_2, joint_3, joint_4, joint_5, joint_6, joint_7, gripper]
+    position_kp: list[float] = field(default_factory=lambda: [240.0, 240.0, 240.0, 240.0, 24.0, 31.0, 25.0, 16.0])
+    position_kd: list[float] = field(default_factory=lambda: [3.0, 3.0, 3.0, 3.0, 0.2, 0.2, 0.2, 0.2])
+    
+    # Calibration parameters
+    calibration_mode: str = "manual"  # "manual" or "auto"
+    zero_position_on_connect: bool = False  # Set zero position on connect

src/lerobot/robots/openarms/openarms_follower.py

@@ -0,0 +1,509 @@
+#!/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 logging
+import time
+from functools import cached_property
+from typing import Any, Dict
+
+import numpy as np
+import pinocchio as pin
+
+from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.damiao import DamiaoMotorsBus
+from lerobot.motors.damiao.tables import MotorType
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+
+from ..robot import Robot
+from ..utils import ensure_safe_goal_position
+from .config_openarms_follower import OpenArmsFollowerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class OpenArmsFollower(Robot):
+    """
+    OpenArms Follower Robot which uses CAN bus communication to control 7 DOF arm with a gripper.
+    The arm uses Damiao motors in MIT control mode.
+    """
+
+    config_class = OpenArmsFollowerConfig
+    name = "openarms_follower"
+
+    def __init__(self, config: OpenArmsFollowerConfig):
+        super().__init__(config)
+        self.config = config
+        
+        norm_mode_body = MotorNormMode.DEGREES  # Always use degrees for Damiao motors
+        
+        # Right arm motors (on port_right)
+        # Each arm uses the same CAN IDs since they're on separate buses
+        motors_right = {}
+        for motor_name, (send_id, recv_id, motor_type_str) in config.motor_config.items():
+            motor = Motor(send_id, motor_type_str, norm_mode_body)
+            motor.recv_id = recv_id
+            motor.motor_type = getattr(MotorType, motor_type_str.upper().replace("-", "_"))
+            motors_right[motor_name] = motor
+        
+        # Left arm motors (on port_left, same IDs as right since separate bus)
+        motors_left = {}
+        for motor_name, (send_id, recv_id, motor_type_str) in config.motor_config.items():
+            motor = Motor(send_id, motor_type_str, norm_mode_body)
+            motor.recv_id = recv_id
+            motor.motor_type = getattr(MotorType, motor_type_str.upper().replace("-", "_"))
+            motors_left[motor_name] = motor
+        
+        # Initialize separate Damiao motors buses (one per arm) with CAN FD support
+        self.bus_right = DamiaoMotorsBus(
+            port=self.config.port_right,
+            motors=motors_right,
+            calibration={k.replace("right_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("right_")},
+            can_interface=self.config.can_interface,
+            use_can_fd=self.config.use_can_fd,
+            bitrate=self.config.can_bitrate,
+            data_bitrate=self.config.can_data_bitrate if self.config.use_can_fd else None,
+        )
+        
+        self.bus_left = DamiaoMotorsBus(
+            port=self.config.port_left,
+            motors=motors_left,
+            calibration={k.replace("left_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("left_")},
+            can_interface=self.config.can_interface,
+            use_can_fd=self.config.use_can_fd,
+            bitrate=self.config.can_bitrate,
+            data_bitrate=self.config.can_data_bitrate if self.config.use_can_fd else None,
+        )
+        
+        # Initialize cameras
+        self.cameras = make_cameras_from_configs(config.cameras)
+        
+        # Initialize Pinocchio robot model for dynamics (optional)
+        self.pin_robot = None
+        try:
+            # Try to load URDF if available
+            # TODO: Add OpenArms URDF file to repository
+            self.pin_robot = pin.RobotWrapper.BuildFromURDF("urdf/openarms.urdf", "urdf")
+            logger.info("Loaded OpenArms URDF for dynamics computation")
+        except Exception as e:
+            logger.warning(f"Could not load URDF for dynamics: {e}. Gravity compensation will not be available.")
+
+    @property
+    def _motors_ft(self) -> Dict[str, type]:
+        """Motor features for observation and action spaces."""
+        features = {}
+        # Right arm motors
+        for motor in self.bus_right.motors:
+            features[f"right_{motor}.pos"] = float
+            features[f"right_{motor}.vel"] = float
+            features[f"right_{motor}.torque"] = float
+        # Left arm motors
+        for motor in self.bus_left.motors:
+            features[f"left_{motor}.pos"] = float
+            features[f"left_{motor}.vel"] = float
+            features[f"left_{motor}.torque"] = float
+        return features
+
+    @property
+    def _cameras_ft(self) -> Dict[str, tuple]:
+        """Camera features for observation space."""
+        return {
+            cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3)
+            for cam in self.cameras
+        }
+
+    @cached_property
+    def observation_features(self) -> Dict[str, type | tuple]:
+        """Combined observation features from motors and cameras."""
+        return {**self._motors_ft, **self._cameras_ft}
+
+    @cached_property
+    def action_features(self) -> Dict[str, type]:
+        """Action features (motor positions only)."""
+        return self._motors_ft
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if robot is connected."""
+        return (self.bus_right.is_connected and 
+                self.bus_left.is_connected and 
+                all(cam.is_connected for cam in self.cameras.values()))
+
+    def connect(self, calibrate: bool = True) -> None:
+        """
+        Connect to the robot and optionally calibrate.
+        
+        We assume that at connection time, the arms are in a safe rest position,
+        and torque can be safely disabled to run calibration if needed.
+        """
+        if self.is_connected:
+            raise DeviceAlreadyConnectedError(f"{self} already connected")
+
+        # Connect to both CAN buses
+        logger.info(f"Connecting right arm on {self.config.port_right}...")
+        self.bus_right.connect()
+        logger.info(f"Connecting left arm on {self.config.port_left}...")
+        self.bus_left.connect()
+        
+        # Run calibration if needed
+        if not self.is_calibrated and calibrate:
+            logger.info(
+                "No calibration found or calibration mismatch. Running calibration..."
+            )
+            self.calibrate()
+        
+        # Connect cameras
+        for cam in self.cameras.values():
+            cam.connect()
+        
+        # Configure motors
+        self.configure()
+        
+        # Optionally set zero position
+        if self.config.zero_position_on_connect:
+            logger.info("Setting current position as zero...")
+            self.bus_right.set_zero_position()
+            self.bus_left.set_zero_position()
+        
+        logger.info(f"{self} connected.")
+
+    @property
+    def is_calibrated(self) -> bool:
+        """Check if robot is calibrated."""
+        return self.bus_right.is_calibrated and self.bus_left.is_calibrated
+
+    def calibrate(self) -> None:
+        """
+        Run calibration procedure for OpenArms robot.
+        
+        The calibration procedure:
+        1. Disable torque
+        2. Ask user to position arms in hanging position with grippers closed
+        3. Set this as zero position
+        4. Record range of motion for each joint
+        5. Save calibration
+        """
+        if self.calibration:
+            # Ask user whether to use existing calibration
+            user_input = input(
+                f"Press ENTER to use existing calibration for {self.id}, "
+                f"or type 'c' and press ENTER to run new calibration: "
+            )
+            if user_input.strip().lower() != "c":
+                logger.info(f"Using existing calibration for {self.id}")
+                # Split calibration for each bus
+                cal_right = {k.replace("right_", ""): v for k, v in self.calibration.items() if k.startswith("right_")}
+                cal_left = {k.replace("left_", ""): v for k, v in self.calibration.items() if k.startswith("left_")}
+                self.bus_right.write_calibration(cal_right)
+                self.bus_left.write_calibration(cal_left)
+                return
+        
+        logger.info(f"\nRunning calibration for {self}")
+        
+        # Calibrate each arm separately
+        self._calibrate_arm("right", self.bus_right)
+        self._calibrate_arm("left", self.bus_left)
+        
+        print(f"\nCalibration complete and saved to {self.calibration_fpath}")
+    
+    def _calibrate_arm(self, arm_name: str, bus: DamiaoMotorsBus) -> None:
+        """Calibrate a single arm."""
+        logger.info(f"\n=== Calibrating {arm_name.upper()} arm ===")
+        
+        # Disable torque for manual positioning
+        bus.disable_torque()
+        time.sleep(0.1)
+        
+        # Step 1: Set zero position
+        input(
+            f"\nCalibration: Zero Position ({arm_name.upper()} arm)\n"
+            "Position the arm in the following configuration:\n"
+            "  - Arm hanging straight down\n"
+            "  - Gripper closed\n"
+            "Press ENTER when ready..."
+        )
+        
+        # Set current position as zero for all motors
+        bus.set_zero_position()
+        logger.info(f"{arm_name.capitalize()} arm zero position set.")
+        
+        # Automatically set range to -90° to +90° for all joints
+        print(
+            f"\nAutomatically setting range: -90° to +90° for all joints"
+        )
+        
+        # Create calibration data with fixed ranges
+        if self.calibration is None:
+            self.calibration = {}
+            
+        for motor_name, motor in bus.motors.items():
+            # Prefix motor name with arm name for storage
+            prefixed_name = f"{arm_name}_{motor_name}"
+            
+            # Use -90 to +90 for all joints and gripper (integers required)
+            self.calibration[prefixed_name] = MotorCalibration(
+                id=motor.id,
+                drive_mode=0,  # Normal direction
+                homing_offset=0,  # Already set via set_zero_position
+                range_min=-90,  # -90 degrees (integer)
+                range_max=90,   # +90 degrees (integer)
+            )
+            logger.info(f"  {prefixed_name}: range set to [-90°, +90°]")
+        
+        # Write calibration to this arm's motors
+        cal_for_bus = {k.replace(f"{arm_name}_", ""): v for k, v in self.calibration.items() if k.startswith(f"{arm_name}_")}
+        bus.write_calibration(cal_for_bus)
+        
+        # Re-enable torque
+        bus.enable_torque()
+        
+        # Save calibration after each arm
+        self._save_calibration()
+
+    def configure(self) -> None:
+        """Configure motors with appropriate settings."""
+        # Configure right arm
+        with self.bus_right.torque_disabled():
+            self.bus_right.configure_motors()
+        
+        # Configure left arm
+        with self.bus_left.torque_disabled():
+            self.bus_left.configure_motors()
+
+    def setup_motors(self) -> None:
+        raise NotImplementedError("Motor ID configuration is typically done via manufacturer tools for CAN motors.")
+
+    def get_observation(self) -> Dict[str, Any]:
+        """Get current observation from robot including position, velocity, and torque."""
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        obs_dict = {}
+        
+        # Read motor positions, velocities, and torques from right arm
+        start = time.perf_counter()
+        positions_right = self.bus_right.sync_read("Present_Position")
+        velocities_right = self.bus_right.sync_read("Present_Velocity")
+        torques_right = self.bus_right.sync_read("Present_Torque")
+        
+        for motor in self.bus_right.motors:
+            obs_dict[f"right_{motor}.pos"] = positions_right.get(motor, 0.0)
+            obs_dict[f"right_{motor}.vel"] = velocities_right.get(motor, 0.0)
+            obs_dict[f"right_{motor}.torque"] = torques_right.get(motor, 0.0)
+        
+        # Read motor positions, velocities, and torques from left arm
+        positions_left = self.bus_left.sync_read("Present_Position")
+        velocities_left = self.bus_left.sync_read("Present_Velocity")
+        torques_left = self.bus_left.sync_read("Present_Torque")
+        
+        for motor in self.bus_left.motors:
+            obs_dict[f"left_{motor}.pos"] = positions_left.get(motor, 0.0)
+            obs_dict[f"left_{motor}.vel"] = velocities_left.get(motor, 0.0)
+            obs_dict[f"left_{motor}.torque"] = torques_left.get(motor, 0.0)
+        
+        dt_ms = (time.perf_counter() - start) * 1e3
+        logger.debug(f"{self} read state: {dt_ms:.1f}ms")
+        
+        # Capture images from cameras
+        for cam_key, cam in self.cameras.items():
+            start = time.perf_counter()
+            obs_dict[cam_key] = cam.async_read()
+            dt_ms = (time.perf_counter() - start) * 1e3
+            logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
+        
+        return obs_dict
+
+    def send_action(self, action: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Send action command to robot.
+        
+        The action magnitude may be clipped based on safety limits.
+        
+        Args:
+            action: Dictionary with motor positions (e.g., "right_joint_1.pos", "left_joint_2.pos")
+            
+        Returns:
+            The action actually sent (potentially clipped)
+        """
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        # Extract motor positions from action and split by arm
+        goal_pos_right = {}
+        goal_pos_left = {}
+        
+        for key, val in action.items():
+            if key.endswith(".pos"):
+                motor_name = key.removesuffix(".pos")
+                if motor_name.startswith("right_"):
+                    # Remove "right_" prefix for bus access
+                    goal_pos_right[motor_name.removeprefix("right_")] = val
+                elif motor_name.startswith("left_"):
+                    # Remove "left_" prefix for bus access
+                    goal_pos_left[motor_name.removeprefix("left_")] = val
+        
+        # Apply safety limits if configured
+        if self.config.max_relative_target is not None:
+            # Get current positions
+            present_pos_right = self.bus_right.sync_read("Present_Position")
+            present_pos_left = self.bus_left.sync_read("Present_Position")
+            
+            # Apply safety limits to right arm
+            if goal_pos_right:
+                goal_present_pos_right = {
+                    key: (g_pos, present_pos_right.get(key, 0.0)) 
+                    for key, g_pos in goal_pos_right.items()
+                }
+                goal_pos_right = ensure_safe_goal_position(
+                    goal_present_pos_right, 
+                    self.config.max_relative_target
+                )
+            
+            # Apply safety limits to left arm
+            if goal_pos_left:
+                goal_present_pos_left = {
+                    key: (g_pos, present_pos_left.get(key, 0.0)) 
+                    for key, g_pos in goal_pos_left.items()
+                }
+                goal_pos_left = ensure_safe_goal_position(
+                    goal_present_pos_left, 
+                    self.config.max_relative_target
+                )
+        
+        # Motor name to index mapping for gains
+        motor_index = {
+            "joint_1": 0,
+            "joint_2": 1,
+            "joint_3": 2,
+            "joint_4": 3,
+            "joint_5": 4,
+            "joint_6": 5,
+            "joint_7": 6,
+            "gripper": 7,
+        }
+        
+        # Send MIT control commands to right arm
+        for motor_name, position_degrees in goal_pos_right.items():
+            # Get per-motor gains from config
+            idx = motor_index.get(motor_name, 0)
+            kp = self.config.position_kp[idx]
+            kd = self.config.position_kd[idx]
+            
+            # Send MIT control command (position is in degrees)
+            self.bus_right._mit_control(
+                motor_name,
+                kp=kp,
+                kd=kd,
+                position_degrees=position_degrees,
+                velocity_deg_per_sec=0.0,
+                torque=0.0
+            )
+        
+        # Send MIT control commands to left arm
+        for motor_name, position_degrees in goal_pos_left.items():
+            # Get per-motor gains from config
+            idx = motor_index.get(motor_name, 0)
+            kp = self.config.position_kp[idx]
+            kd = self.config.position_kd[idx]
+            
+            # Send MIT control command (position is in degrees)
+            self.bus_left._mit_control(
+                motor_name,
+                kp=kp,
+                kd=kd,
+                position_degrees=position_degrees,
+                velocity_deg_per_sec=0.0,
+                torque=0.0
+            )
+        
+        # Return the actions that were actually sent
+        result = {}
+        for motor, val in goal_pos_right.items():
+            result[f"right_{motor}.pos"] = val
+        for motor, val in goal_pos_left.items():
+            result[f"left_{motor}.pos"] = val
+        return result
+
+    def disconnect(self):
+        """Disconnect from robot."""
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        # Disconnect from CAN buses
+        self.bus_right.disconnect(self.config.disable_torque_on_disconnect)
+        self.bus_left.disconnect(self.config.disable_torque_on_disconnect)
+        
+        # Disconnect cameras
+        for cam in self.cameras.values():
+            cam.disconnect()
+        
+        logger.info(f"{self} disconnected.")
+    
+    def _deg_to_rad(self, deg: Dict[str, float | int]) -> Dict[str, float]:
+        """Convert degrees to radians for all motors."""
+        return {m: np.deg2rad(float(v)) for m, v in deg.items()}
+    
+    def _gravity_from_q(self, q_rad: Dict[str, float]) -> Dict[str, float]:
+        """
+        Compute g(q) [N·m] for all joints in the robot.
+        The order of joints in the URDF matches the concatenated motor lists (right then left).
+        
+        Args:
+            q_rad: Dictionary mapping motor names (with arm prefix) to positions in radians
+            
+        Returns:
+            Dictionary mapping motor names to gravity torques in N·m
+            
+        Raises:
+            RuntimeError: If URDF model is not loaded
+        """
+        if self.pin_robot is None:
+            raise RuntimeError(
+                "Cannot compute gravity: URDF model not loaded. "
+                "Ensure urdf/openarms.urdf exists and is valid."
+            )
+        
+        # Build position vector in the order of motors (right arm, then left arm)
+        q = np.zeros(self.pin_robot.model.nq)
+        idx = 0
+        
+        # Right arm motors
+        for motor_name in self.bus_right.motors:
+            full_name = f"right_{motor_name}"
+            q[idx] = q_rad.get(full_name, 0.0)
+            idx += 1
+        
+        # Left arm motors
+        for motor_name in self.bus_left.motors:
+            full_name = f"left_{motor_name}"
+            q[idx] = q_rad.get(full_name, 0.0)
+            idx += 1
+        
+        # Compute generalized gravity vector
+        g = pin.computeGeneralizedGravity(self.pin_robot.model, self.pin_robot.data, q)
+        
+        # Map back to motor names
+        result = {}
+        idx = 0
+        for motor_name in self.bus_right.motors:
+            result[f"right_{motor_name}"] = float(g[idx])
+            idx += 1
+        for motor_name in self.bus_left.motors:
+            result[f"left_{motor_name}"] = float(g[idx])
+            idx += 1
+        
+        return result
+    

src/lerobot/scripts/lerobot_eval.py

@@ -71,7 +71,7 @@ from tqdm import trange
 
 from lerobot.configs import parser
 from lerobot.configs.eval import EvalPipelineConfig
-from lerobot.envs.factory import make_env, make_env_pre_post_processors
+from lerobot.envs.factory import make_env
 from lerobot.envs.utils import (
     add_envs_task,
     check_env_attributes_and_types,
@@ -94,8 +94,6 @@ from lerobot.utils.utils import (
 def rollout(
     env: gym.vector.VectorEnv,
     policy: PreTrainedPolicy,
-    env_preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    env_postprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
     seeds: list[int] | None = None,
@@ -167,19 +165,11 @@ def rollout(
         # Infer "task" from attributes of environments.
         # TODO: works with SyncVectorEnv but not AsyncVectorEnv
         observation = add_envs_task(env, observation)
-
-        # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO)
-        observation = env_preprocessor(observation)
-
         observation = preprocessor(observation)
         with torch.inference_mode():
             action = policy.select_action(observation)
         action = postprocessor(action)
 
-        action_transition = {"action": action}
-        action_transition = env_postprocessor(action_transition)
-        action = action_transition["action"]
-
         # Convert to CPU / numpy.
         action_numpy: np.ndarray = action.to("cpu").numpy()
         assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)"
@@ -249,8 +239,6 @@ def rollout(
 def eval_policy(
     env: gym.vector.VectorEnv,
     policy: PreTrainedPolicy,
-    env_preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    env_postprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
     n_episodes: int,
@@ -331,8 +319,6 @@ def eval_policy(
         rollout_data = rollout(
             env=env,
             policy=policy,
-            env_preprocessor=env_preprocessor,
-            env_postprocessor=env_postprocessor,
             preprocessor=preprocessor,
             postprocessor=postprocessor,
             seeds=list(seeds) if seeds else None,
@@ -531,16 +517,10 @@ def eval_main(cfg: EvalPipelineConfig):
         pretrained_path=cfg.policy.pretrained_path,
         preprocessor_overrides=preprocessor_overrides,
     )
-
-    # Create environment-specific preprocessor and postprocessor (e.g., for LIBERO environments)
-    env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env)
-
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         info = eval_policy_all(
             envs=envs,
             policy=policy,
-            env_preprocessor=env_preprocessor,
-            env_postprocessor=env_postprocessor,
             preprocessor=preprocessor,
             postprocessor=postprocessor,
             n_episodes=cfg.eval.n_episodes,
@@ -581,8 +561,6 @@ def eval_one(
     env: gym.vector.VectorEnv,
     *,
     policy: PreTrainedPolicy,
-    env_preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    env_postprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
     n_episodes: int,
@@ -598,8 +576,6 @@ def eval_one(
     task_result = eval_policy(
         env=env,
         policy=policy,
-        env_preprocessor=env_preprocessor,
-        env_postprocessor=env_postprocessor,
         preprocessor=preprocessor,
         postprocessor=postprocessor,
         n_episodes=n_episodes,
@@ -624,8 +600,6 @@ def run_one(
     env,
     *,
     policy,
-    env_preprocessor,
-    env_postprocessor,
     preprocessor,
     postprocessor,
     n_episodes: int,
@@ -648,8 +622,6 @@ def run_one(
     metrics = eval_one(
         env,
         policy=policy,
-        env_preprocessor=env_preprocessor,
-        env_postprocessor=env_postprocessor,
         preprocessor=preprocessor,
         postprocessor=postprocessor,
         n_episodes=n_episodes,
@@ -667,8 +639,6 @@ def run_one(
 def eval_policy_all(
     envs: dict[str, dict[int, gym.vector.VectorEnv]],
     policy,
-    env_preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    env_postprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
     n_episodes: int,
@@ -724,8 +694,6 @@ def eval_policy_all(
     task_runner = partial(
         run_one,
         policy=policy,
-        env_preprocessor=env_preprocessor,
-        env_postprocessor=env_postprocessor,
         preprocessor=preprocessor,
         postprocessor=postprocessor,
         n_episodes=n_episodes,

src/lerobot/scripts/lerobot_train.py

@@ -29,7 +29,7 @@ from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets.factory import make_dataset
 from lerobot.datasets.sampler import EpisodeAwareSampler
 from lerobot.datasets.utils import cycle
-from lerobot.envs.factory import make_env, make_env_pre_post_processors
+from lerobot.envs.factory import make_env
 from lerobot.envs.utils import close_envs
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies.factory import make_policy, make_pre_post_processors
@@ -61,7 +61,6 @@ def update_policy(
     accelerator: Accelerator,
     lr_scheduler=None,
     lock=None,
-    rabc_weight_computer=None,
 ) -> tuple[MetricsTracker, dict]:
     """
     Performs a single training step to update the policy's weights.
@@ -86,22 +85,10 @@ def update_policy(
     """
     start_time = time.perf_counter()
     policy.train()
-    
-    # Compute RA-BC weights if enabled
-    rabc_weights = None
-    if rabc_weight_computer is not None:
-        rabc_weights = rabc_weight_computer.compute_batch_weights(batch)
 
     # Let accelerator handle mixed precision
     with accelerator.autocast():
         loss, output_dict = policy.forward(batch)
-        
-        # Apply RA-BC weights if enabled
-        if rabc_weights is not None:
-            # Weight the loss
-            loss = loss * rabc_weights.mean()
-            output_dict['rabc_mean_weight'] = rabc_weights.mean().item()
-        
         # TODO(rcadene): policy.unnormalize_outputs(out_dict)
 
     # Use accelerator's backward method
@@ -153,6 +140,8 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
         cfg: A `TrainPipelineConfig` object containing all training configurations.
         accelerator: Optional Accelerator instance. If None, one will be created automatically.
     """
+    cfg.validate()
+
     # Create Accelerator if not provided
     # It will automatically detect if running in distributed mode or single-process mode
     # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes
@@ -169,8 +158,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
     # When using accelerate, only the main process should log to avoid duplicate outputs
     is_main_process = accelerator.is_main_process
 
-    cfg.validate()
-
     # Only log on main process
     if is_main_process:
         logging.info(pformat(cfg.to_dict()))
@@ -228,10 +215,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
     if (cfg.policy.pretrained_path and not cfg.resume) or not cfg.policy.pretrained_path:
         # Only provide dataset_stats when not resuming from saved processor state
         processor_kwargs["dataset_stats"] = dataset.meta.stats
-    
-    # For SARM, always provide dataset_meta for progress normalization
-    if cfg.policy.type == "sarm":
-        processor_kwargs["dataset_meta"] = dataset.meta
 
     if cfg.policy.pretrained_path is not None:
         processor_kwargs["preprocessor_overrides"] = {
@@ -263,28 +246,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
     if is_main_process:
         logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
-    
-    # Load reward model for RA-BC if enabled
-    rabc_weight_computer = None
-    if cfg.use_rabc:
-        logging.info(f"Loading reward model for RA-BC from {cfg.reward_model_path}")
-        from lerobot.policies.factory import get_policy_class
-        from lerobot.utils.rabc import RABCWeightComputer
-        
-        # Detect reward model type from path
-        # For now, assume SARM if not specified
-        reward_model_class = get_policy_class("sarm")
-        reward_model = reward_model_class.from_pretrained(cfg.reward_model_path)
-        reward_model.to(device)
-        reward_model.eval()
-        
-        rabc_weight_computer = RABCWeightComputer(
-            reward_model=reward_model,
-            kappa=cfg.rabc_kappa,
-            epsilon=cfg.rabc_epsilon,
-            device=device,
-        )
-        logging.info("RA-BC weight computer initialized")
 
     step = 0  # number of policy updates (forward + backward + optim)
 
@@ -298,8 +259,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
         logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
         if cfg.env is not None:
             logging.info(f"{cfg.env.task=}")
-            logging.info("Creating environment processors")
-            env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env)
         logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
         logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
         logging.info(f"{dataset.num_episodes=}")
@@ -315,22 +274,9 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
         sampler = EpisodeAwareSampler(
             dataset.meta.episodes["dataset_from_index"],
             dataset.meta.episodes["dataset_to_index"],
-            episode_indices_to_use=dataset.episodes,
             drop_n_last_frames=cfg.policy.drop_n_last_frames,
             shuffle=True,
         )
-    elif cfg.policy.type == "sarm" and getattr(cfg.policy, "use_temporal_sampler", False):
-        # Use SARM temporal sampler for reward model training
-        from lerobot.datasets.temporal_sampler import SARMTemporalSampler
-        
-        shuffle = False
-        sampler = SARMTemporalSampler(
-            dataset_from_index=dataset.meta.episodes["dataset_from_index"],
-            dataset_to_index=dataset.meta.episodes["dataset_to_index"],
-            frame_gap=getattr(cfg.policy, "frame_gap", 30),
-            shuffle=True,
-            seed=cfg.seed,
-        )
     else:
         shuffle = True
         sampler = None
@@ -375,7 +321,7 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
     )
 
     if is_main_process:
-        logging.info(f"Start offline training on a fixed dataset, with effective batch size: {effective_batch_size}")
+        logging.info("Start offline training on a fixed dataset")
 
     for _ in range(step, cfg.steps):
         start_time = time.perf_counter()
@@ -391,7 +337,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
             cfg.optimizer.grad_clip_norm,
             accelerator=accelerator,
             lr_scheduler=lr_scheduler,
-            rabc_weight_computer=rabc_weight_computer,
         )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -408,14 +353,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
                 wandb_log_dict = train_tracker.to_dict()
                 if output_dict:
                     wandb_log_dict.update(output_dict)
-                # Log RA-BC statistics if enabled
-                if rabc_weight_computer is not None:
-                    rabc_stats = rabc_weight_computer.get_stats()
-                    wandb_log_dict.update({
-                        'rabc_progress_mean': rabc_stats['mean'],
-                        'rabc_progress_std': rabc_stats['std'],
-                        'rabc_samples_seen': rabc_stats['count'],
-                    })
                 wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 
@@ -447,8 +384,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
                     eval_info = eval_policy_all(
                         envs=eval_env,  # dict[suite][task_id] -> vec_env
                         policy=accelerator.unwrap_model(policy),
-                        env_preprocessor=env_preprocessor,
-                        env_postprocessor=env_postprocessor,
                         preprocessor=preprocessor,
                         postprocessor=postprocessor,
                         n_episodes=cfg.eval.n_episodes,

src/lerobot/teleoperators/openarms/__init__.py

@@ -0,0 +1,20 @@
+#!/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 .config_openarms_leader import OpenArmsLeaderConfig
+from .openarms_leader import OpenArmsLeader
+
+__all__ = ["OpenArmsLeader", "OpenArmsLeaderConfig"]

src/lerobot/teleoperators/openarms/config_openarms_leader.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Dict
+
+from ..config import TeleoperatorConfig
+
+
+@TeleoperatorConfig.register_subclass("openarms_leader")
+@dataclass
+class OpenArmsLeaderConfig(TeleoperatorConfig):
+    """Configuration for the OpenArms leader/teleoperator with Damiao motors."""
+    
+    # CAN interfaces - one per arm
+    # Right arm CAN interface (e.g., "can2")
+    # Left arm CAN interface (e.g., "can3")
+    # Linux: "can0", "can1", etc.
+    # macOS: "/dev/cu.usbmodem*" (serial device)
+    port_right: str = "can2"    # CAN interface for right arm
+    port_left: str = "can3"     # CAN interface for left arm
+    
+    # CAN interface type: "socketcan" (Linux), "slcan" (macOS/serial), or "auto" (auto-detect)
+    can_interface: str = "socketcan"
+    
+    # CAN FD settings (OpenArms uses CAN FD by default)
+    use_can_fd: bool = True
+    can_bitrate: int = 1000000      # Nominal bitrate (1 Mbps)
+    can_data_bitrate: int = 5000000  # Data bitrate for CAN FD (5 Mbps)
+    
+    # Motor configuration for OpenArms (7 DOF per arm)
+    # Maps motor names to (send_can_id, recv_can_id, motor_type)
+    # Based on: https://docs.openarm.dev/software/setup/configure-test
+    # OpenArms uses 4 types of motors:
+    # - DM8009 (DM-J8009P-2EC) for shoulders (high torque)
+    # - DM4340P and DM4340 for shoulder rotation and elbow
+    # - DM4310 (DM-J4310-2EC V1.1) for wrist and gripper
+    motor_config: Dict[str, tuple[int, int, str]] = field(default_factory=lambda: {
+        "joint_1": (0x01, 0x11, "dm8009"),   # J1 - Shoulder pan (DM8009)
+        "joint_2": (0x02, 0x12, "dm8009"),   # J2 - Shoulder lift (DM8009)
+        "joint_3": (0x03, 0x13, "dm4340"),   # J3 - Shoulder rotation (DM4340)
+        "joint_4": (0x04, 0x14, "dm4340"),   # J4 - Elbow flex (DM4340)
+        "joint_5": (0x05, 0x15, "dm4310"),   # J5 - Wrist roll (DM4310)
+        "joint_6": (0x06, 0x16, "dm4310"),   # J6 - Wrist pitch (DM4310)
+        "joint_7": (0x07, 0x17, "dm4310"),   # J7 - Wrist rotation (DM4310)
+        "gripper": (0x08, 0x18, "dm4310"),   # J8 - Gripper (DM4310)
+    })
+    
+    # Torque mode settings for manual control
+    # When enabled, motors have torque disabled for manual movement
+    manual_control: bool = True

src/lerobot/teleoperators/openarms/openarms_leader.py

@@ -0,0 +1,310 @@
+#!/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 logging
+import time
+from typing import Any, Dict
+
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.damiao import DamiaoMotorsBus
+from lerobot.motors.damiao.tables import MotorType
+from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+
+from ..teleoperator import Teleoperator
+from .config_openarms_leader import OpenArmsLeaderConfig
+
+logger = logging.getLogger(__name__)
+
+
+class OpenArmsLeader(Teleoperator):
+    """
+    OpenArms Leader/Teleoperator Arm with Damiao motors.
+    
+    This teleoperator uses CAN bus communication to read positions from 
+    Damiao motors that are manually moved (torque disabled).
+    """
+
+    config_class = OpenArmsLeaderConfig
+    name = "openarms_leader"
+
+    def __init__(self, config: OpenArmsLeaderConfig):
+        super().__init__(config)
+        self.config = config
+        
+        norm_mode_body = MotorNormMode.DEGREES  # Always use degrees for Damiao motors
+        
+        # Right arm motors (on port_right)
+        # Each arm uses the same CAN IDs since they're on separate buses
+        motors_right = {}
+        for motor_name, (send_id, recv_id, motor_type_str) in config.motor_config.items():
+            motor = Motor(send_id, motor_type_str, norm_mode_body)
+            motor.recv_id = recv_id
+            motor.motor_type = getattr(MotorType, motor_type_str.upper().replace("-", "_"))
+            motors_right[motor_name] = motor
+        
+        # Left arm motors (on port_left, same IDs as right since separate bus)
+        motors_left = {}
+        for motor_name, (send_id, recv_id, motor_type_str) in config.motor_config.items():
+            motor = Motor(send_id, motor_type_str, norm_mode_body)
+            motor.recv_id = recv_id
+            motor.motor_type = getattr(MotorType, motor_type_str.upper().replace("-", "_"))
+            motors_left[motor_name] = motor
+        
+        # Initialize separate Damiao motors buses (one per arm) with CAN FD support
+        self.bus_right = DamiaoMotorsBus(
+            port=self.config.port_right,
+            motors=motors_right,
+            calibration={k.replace("right_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("right_")},
+            can_interface=self.config.can_interface,
+            use_can_fd=self.config.use_can_fd,
+            bitrate=self.config.can_bitrate,
+            data_bitrate=self.config.can_data_bitrate if self.config.use_can_fd else None,
+        )
+        
+        self.bus_left = DamiaoMotorsBus(
+            port=self.config.port_left,
+            motors=motors_left,
+            calibration={k.replace("left_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("left_")},
+            can_interface=self.config.can_interface,
+            use_can_fd=self.config.use_can_fd,
+            bitrate=self.config.can_bitrate,
+            data_bitrate=self.config.can_data_bitrate if self.config.use_can_fd else None,
+        )
+
+    @property
+    def action_features(self) -> Dict[str, type]:
+        """Features produced by this teleoperator."""
+        features = {}
+        # Right arm motors
+        for motor in self.bus_right.motors:
+            features[f"right_{motor}.pos"] = float
+            features[f"right_{motor}.vel"] = float
+            features[f"right_{motor}.torque"] = float
+        # Left arm motors
+        for motor in self.bus_left.motors:
+            features[f"left_{motor}.pos"] = float
+            features[f"left_{motor}.vel"] = float
+            features[f"left_{motor}.torque"] = float
+        return features
+
+    @property
+    def feedback_features(self) -> Dict[str, type]:
+        """Feedback features (not implemented for OpenArms)."""
+        return {}
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if teleoperator is connected."""
+        return self.bus_right.is_connected and self.bus_left.is_connected
+
+    def connect(self, calibrate: bool = True) -> None:
+        """
+        Connect to the teleoperator.
+        
+        For manual control, we disable torque after connecting so the
+        arm can be moved by hand.
+        """
+        if self.is_connected:
+            raise DeviceAlreadyConnectedError(f"{self} already connected")
+
+        # Connect to CAN buses
+        logger.info(f"Connecting right arm on {self.config.port_right}...")
+        self.bus_right.connect()
+        logger.info(f"Connecting left arm on {self.config.port_left}...")
+        self.bus_left.connect()
+        
+        # Run calibration if needed
+        if not self.is_calibrated and calibrate:
+            logger.info(
+                "No calibration found or calibration mismatch. Running calibration..."
+            )
+            self.calibrate()
+        
+        # Configure for manual control
+        self.configure()
+        
+        logger.info(f"{self} connected.")
+
+    @property
+    def is_calibrated(self) -> bool:
+        """Check if teleoperator is calibrated."""
+        return self.bus_right.is_calibrated and self.bus_left.is_calibrated
+
+    def calibrate(self) -> None:
+        """
+        Run calibration procedure for OpenArms leader.
+        
+        The calibration procedure:
+        1. Disable torque (if not already disabled)
+        2. Ask user to position arm in zero position (hanging with gripper closed)
+        3. Set this as zero position
+        4. Record range of motion for each joint
+        5. Save calibration
+        """
+        if self.calibration:
+            # Ask user whether to use existing calibration
+            user_input = input(
+                f"Press ENTER to use existing calibration for {self.id}, "
+                f"or type 'c' and press ENTER to run new calibration: "
+            )
+            if user_input.strip().lower() != "c":
+                logger.info(f"Using existing calibration for {self.id}")
+                # Split calibration for each bus
+                cal_right = {k.replace("right_", ""): v for k, v in self.calibration.items() if k.startswith("right_")}
+                cal_left = {k.replace("left_", ""): v for k, v in self.calibration.items() if k.startswith("left_")}
+                self.bus_right.write_calibration(cal_right)
+                self.bus_left.write_calibration(cal_left)
+                return
+        
+        logger.info(f"\nRunning calibration for {self}")
+        
+        # Calibrate each arm separately
+        self._calibrate_arm("right", self.bus_right)
+        self._calibrate_arm("left", self.bus_left)
+        
+        print(f"\nCalibration complete and saved to {self.calibration_fpath}")
+    
+    def _calibrate_arm(self, arm_name: str, bus: DamiaoMotorsBus) -> None:
+        """Calibrate a single arm."""
+        logger.info(f"\n=== Calibrating {arm_name.upper()} arm ===")
+        
+        # Ensure torque is disabled for manual positioning
+        bus.disable_torque()
+        time.sleep(0.1)
+        
+        # Step 1: Set zero position
+        input(
+            f"\nCalibration: Zero Position ({arm_name.upper()} arm)\n"
+            "Position the arm in the following configuration:\n"
+            "  - Arm hanging straight down\n"
+            "  - Gripper closed\n"
+            "Press ENTER when ready..."
+        )
+        
+        # Set current position as zero for all motors
+        bus.set_zero_position()
+        logger.info(f"{arm_name.capitalize()} arm zero position set.")
+        
+        # Automatically set range to -90° to +90° for all joints
+        print(
+            f"\nAutomatically setting range: -90° to +90° for all joints"
+        )
+        
+        # Create calibration data with fixed ranges
+        if self.calibration is None:
+            self.calibration = {}
+            
+        for motor_name, motor in bus.motors.items():
+            # Prefix motor name with arm name for storage
+            prefixed_name = f"{arm_name}_{motor_name}"
+            
+            # Use -90 to +90 for all joints and gripper (integers required)
+            self.calibration[prefixed_name] = MotorCalibration(
+                id=motor.id,
+                drive_mode=0,  # Normal direction
+                homing_offset=0,  # Already set via set_zero_position
+                range_min=-90,  # -90 degrees (integer)
+                range_max=90,   # +90 degrees (integer)
+            )
+            logger.info(f"  {prefixed_name}: range set to [-90°, +90°]")
+        
+        # Write calibration to this arm's motors
+        cal_for_bus = {k.replace(f"{arm_name}_", ""): v for k, v in self.calibration.items() if k.startswith(f"{arm_name}_")}
+        bus.write_calibration(cal_for_bus)
+        
+        # Save calibration after each arm
+        self._save_calibration()
+
+    def configure(self) -> None:
+        """
+        Configure motors for manual teleoperation.
+        
+        For manual control, we disable torque so the arm can be moved by hand.
+        """
+        if self.config.manual_control:
+            # Disable torque for manual control
+            logger.info("Disabling torque for manual control...")
+            self.bus_right.disable_torque()
+            self.bus_left.disable_torque()
+        else:
+            # Configure motors normally
+            self.bus_right.configure_motors()
+            self.bus_left.configure_motors()
+
+    def setup_motors(self) -> None:
+        raise NotImplementedError("Motor ID configuration is typically done via manufacturer tools for CAN motors.")
+
+
+    def get_action(self) -> Dict[str, Any]:
+        """
+        Get current action from the leader arm.
+        
+        This is the main method for teleoperators - it reads the current state
+        of the leader arm and returns it as an action that can be sent to a follower.
+        """
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        action_dict = {}
+        
+        # Read motor positions, velocities, and torques from right arm
+        start = time.perf_counter()
+        positions_right = self.bus_right.sync_read("Present_Position")
+        velocities_right = self.bus_right.sync_read("Present_Velocity")
+        torques_right = self.bus_right.sync_read("Present_Torque")
+        
+        for motor in self.bus_right.motors:
+            action_dict[f"right_{motor}.pos"] = positions_right.get(motor, 0.0)
+            action_dict[f"right_{motor}.vel"] = velocities_right.get(motor, 0.0)
+            action_dict[f"right_{motor}.torque"] = torques_right.get(motor, 0.0)
+        
+        # Read motor positions, velocities, and torques from left arm
+        positions_left = self.bus_left.sync_read("Present_Position")
+        velocities_left = self.bus_left.sync_read("Present_Velocity")
+        torques_left = self.bus_left.sync_read("Present_Torque")
+        
+        for motor in self.bus_left.motors:
+            action_dict[f"left_{motor}.pos"] = positions_left.get(motor, 0.0)
+            action_dict[f"left_{motor}.vel"] = velocities_left.get(motor, 0.0)
+            action_dict[f"left_{motor}.torque"] = torques_left.get(motor, 0.0)
+        
+        dt_ms = (time.perf_counter() - start) * 1e3
+        logger.debug(f"{self} read state: {dt_ms:.1f}ms")
+        
+        return action_dict
+
+    def send_feedback(self, feedback: Dict[str, float]) -> None:
+        raise NotImplementedError("Feedback is not yet implemented for OpenArms leader.")
+
+    def disconnect(self) -> None:
+        """Disconnect from teleoperator."""
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+        
+        # For manual control, ensure torque is disabled before disconnecting
+        if self.config.manual_control:
+            try:
+                self.bus_right.disable_torque()
+                self.bus_left.disable_torque()
+            except Exception as e:
+                logger.warning(f"Failed to disable torque during disconnect: {e}")
+        
+        # Disconnect from CAN buses
+        self.bus_right.disconnect(disable_torque=False)  # Already disabled above if needed
+        self.bus_left.disconnect(disable_torque=False)
+        
+        logger.info(f"{self} disconnected.")
+

src/lerobot/utils/constants.py

@@ -70,15 +70,3 @@ LOOKAHEAD_BACKTRACKTABLE = 100
 
 # openpi
 OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38  # TODO(pepijn): Modify this when extending support to fp8 models
-
-# Constants for LIBERO observation keys
-LIBERO_KEY_EEF_POS = "robot_state/eef/pos"
-LIBERO_KEY_EEF_QUAT = "robot_state/eef/quat"
-LIBERO_KEY_EEF_MAT = "robot_state/eef/mat"
-LIBERO_KEY_EEF_AXISANGLE = "robot_state/eef/axisangle"
-LIBERO_KEY_GRIPPER_QPOS = "robot_state/gripper/qpos"
-LIBERO_KEY_GRIPPER_QVEL = "robot_state/gripper/qvel"
-LIBERO_KEY_JOINTS_POS = "robot_state/joints/pos"
-LIBERO_KEY_JOINTS_VEL = "robot_state/joints/vel"
-LIBERO_KEY_PIXELS_AGENTVIEW = "pixels/agentview_image"
-LIBERO_KEY_PIXELS_EYE_IN_HAND = "pixels/robot0_eye_in_hand_image"

src/lerobot/utils/rabc.py

@@ -1,183 +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.
-
-"""
-Reward-Aligned Behavior Cloning (RA-BC) utilities.
-
-RA-BC uses a pre-trained reward model (e.g., SARM) to compute progress-based weights
-for training samples, emphasizing high-quality demonstrations and down-weighting
-suboptimal ones.
-"""
-
-import logging
-import torch
-import torch.nn as nn
-
-
-class RABCWeightComputer:
-    """
-    Computes RA-BC weights for training batches using a pre-trained reward model.
-    
-    Uses Welford's online algorithm for numerically stable running statistics
-    and applies soft weighting based on progress deltas.
-    
-    Args:
-        reward_model: Pre-trained reward model (e.g., SARM)
-        kappa: Hard threshold for high-quality samples (default: 0.01)
-        epsilon: Small constant for numerical stability (default: 1e-6)
-        device: Device to run reward model on
-    """
-    
-    def __init__(
-        self,
-        reward_model: nn.Module,
-        kappa: float = 0.01,
-        epsilon: float = 1e-6,
-        device: torch.device = None,
-    ):
-        self.reward_model = reward_model
-        self.reward_model.eval()  # Always in eval mode
-        self.kappa = kappa
-        self.epsilon = epsilon
-        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Running statistics (Welford's algorithm)
-        self.mean = 0.0
-        self.m2 = 0.0
-        self.count = 0
-        
-        logging.info(f"RA-BC WeightComputer initialized with kappa={kappa}, epsilon={epsilon}")
-    
-    def _update_stats(self, deltas: torch.Tensor):
-        """Update running statistics using Welford's online algorithm."""
-        for delta in deltas:
-            self.count += 1
-            delta_val = delta.item()
-            delta_mean = delta_val - self.mean
-            self.mean += delta_mean / self.count
-            delta_m2 = delta_val - self.mean
-            self.m2 += delta_mean * delta_m2
-    
-    def _compute_weights(self, deltas: torch.Tensor) -> torch.Tensor:
-        """Compute RA-BC weights from progress deltas."""
-        if self.count < 2:
-            # Not enough data, use uniform weights
-            return torch.ones_like(deltas)
-        
-        # Get running statistics
-        mean = max(self.mean, 0.0)  # Clamp mean to non-negative
-        variance = self.m2 / (self.count - 1)
-        std = torch.tensor(variance).sqrt().item()
-        
-        # Compute soft weights
-        lower_bound = mean - 2 * std
-        upper_bound = mean + 2 * std
-        weights = (deltas - lower_bound) / (4 * std + self.epsilon)
-        weights = torch.clamp(weights, 0.0, 1.0)
-        
-        # Apply hard threshold
-        high_quality_mask = deltas > self.kappa
-        weights = torch.where(high_quality_mask, torch.ones_like(weights), weights)
-        
-        return weights
-    
-    @torch.no_grad()
-    def compute_batch_weights(self, batch: dict, chunk_size: int = 1) -> torch.Tensor:
-        """
-        Compute RA-BC weights for a training batch.
-        
-        This function:
-        1. Extracts current and next observations from the batch
-        2. Computes rewards using the reward model
-        3. Calculates progress deltas
-        4. Updates running statistics
-        5. Returns normalized weights
-        
-        Args:
-            batch: Training batch containing observations
-            chunk_size: Size of action chunks for computing deltas (default: 1)
-            
-        Returns:
-            Weights tensor (batch_size,) normalized to sum to batch_size
-        """
-        observation = batch.get('observation', batch)
-        batch_size = next(iter(observation.values())).shape[0]
-        
-        # Extract features needed for reward computation
-        # These should already be encoded by the preprocessor
-        if 'video_features' not in observation or 'text_features' not in observation:
-            logging.warning("RA-BC: Missing video/text features, using uniform weights")
-            return torch.ones(batch_size, device=self.device)
-        
-        video_features = observation['video_features'].to(self.device)
-        text_features = observation['text_features'].to(self.device)
-        state_features = observation.get('state_features', None)
-        if state_features is not None:
-            state_features = state_features.to(self.device)
-        
-        # Compute rewards for current observations
-        # Handle both single-frame and multi-frame features
-        if video_features.dim() == 3:  # (B, T, D)
-            # Multi-frame: use last frame reward
-            if hasattr(self.reward_model, 'calculate_rewards'):
-                current_rewards = self.reward_model.calculate_rewards(
-                    text_features, video_features, state_features,
-                    return_all_frames=False
-                )
-            else:
-                # Fallback for models without calculate_rewards
-                current_rewards = torch.zeros(batch_size, device=self.device)
-        else:  # (B, D)
-            # Single frame
-            if hasattr(self.reward_model, 'calculate_rewards'):
-                current_rewards = self.reward_model.calculate_rewards(
-                    text_features, video_features.unsqueeze(1), state_features,
-                    return_all_frames=False
-                )
-            else:
-                current_rewards = torch.zeros(batch_size, device=self.device)
-        
-        if isinstance(current_rewards, tuple):
-            current_rewards = current_rewards[0]
-        
-        current_rewards = torch.tensor(current_rewards, device=self.device) if isinstance(current_rewards, (list, tuple)) else current_rewards
-        
-        # For simplicity, assume progress delta is proportional to reward
-        # In practice, you'd want to compute next_frame rewards and take differences
-        # For now, use current reward as a proxy for progress delta
-        progress_deltas = current_rewards
-        
-        # Update running statistics
-        self._update_stats(progress_deltas)
-        
-        # Compute weights
-        weights = self._compute_weights(progress_deltas)
-        
-        # Normalize weights to sum to batch_size (maintains effective batch size)
-        weight_sum = weights.sum() + self.epsilon
-        weights = weights * batch_size / weight_sum
-        
-        return weights
-    
-    def get_stats(self) -> dict:
-        """Get current running statistics."""
-        std = torch.tensor(self.m2 / (self.count - 1)).sqrt().item() if self.count > 1 else 0.0
-        return {
-            'mean': self.mean,
-            'std': std,
-            'count': self.count,
-        }
-

tests/conftest.py

@@ -32,6 +32,39 @@ pytest_plugins = [
 ]
 
 
+def pytest_addoption(parser):
+    """Add custom command line option for hardware tests."""
+    parser.addoption(
+        "--run-hardware",
+        action="store_true",
+        default=False,
+        help="Run hardware tests that require actual motors connected",
+    )
+    parser.addoption(
+        "--can-port",
+        action="store",
+        default=None,
+        help="CAN interface port (e.g., 'can0' for Linux, '/dev/cu.usbmodem*' for macOS)",
+    )
+
+
+def pytest_configure(config):
+    """Register custom marker for hardware tests."""
+    config.addinivalue_line("markers", "hardware: mark test as requiring hardware")
+
+
+def pytest_collection_modifyitems(config, items):
+    """Skip hardware tests unless --run-hardware flag is provided."""
+    if config.getoption("--run-hardware"):
+        # --run-hardware given in cli: do not skip hardware tests
+        return
+    
+    skip_hardware = pytest.mark.skip(reason="need --run-hardware option to run")
+    for item in items:
+        if "hardware" in item.keywords:
+            item.add_marker(skip_hardware)
+
+
 def pytest_collection_finish():
     print(f"\nTesting with {DEVICE=}")
 

tests/envs/test_envs.py

@@ -17,7 +17,6 @@ import importlib
 from dataclasses import dataclass, field
 
 import gymnasium as gym
-import numpy as np
 import pytest
 import torch
 from gymnasium.envs.registration import register, registry as gym_registry
@@ -27,11 +26,7 @@ import lerobot
 from lerobot.configs.types import PolicyFeature
 from lerobot.envs.configs import EnvConfig
 from lerobot.envs.factory import make_env, make_env_config
-from lerobot.envs.utils import (
-    _normalize_hub_result,
-    _parse_hub_url,
-    preprocess_observation,
-)
+from lerobot.envs.utils import preprocess_observation
 from tests.utils import require_env
 
 OBS_TYPES = ["state", "pixels", "pixels_agent_pos"]
@@ -113,156 +108,3 @@ def test_factory_custom_gym_id():
     finally:
         if gym_id in gym_registry:
             del gym_registry[gym_id]
-
-
-# Hub environment loading tests
-
-
-def test_make_env_hub_url_parsing():
-    """Test URL parsing for hub environment references."""
-    # simple repo_id
-    repo_id, revision, file_path = _parse_hub_url("user/repo")
-    assert repo_id == "user/repo"
-    assert revision is None
-    assert file_path == "env.py"
-
-    # repo with revision
-    repo_id, revision, file_path = _parse_hub_url("user/repo@main")
-    assert repo_id == "user/repo"
-    assert revision == "main"
-    assert file_path == "env.py"
-
-    # repo with custom file path
-    repo_id, revision, file_path = _parse_hub_url("user/repo:custom_env.py")
-    assert repo_id == "user/repo"
-    assert revision is None
-    assert file_path == "custom_env.py"
-
-    # repo with revision and custom file path
-    repo_id, revision, file_path = _parse_hub_url("user/repo@v1.0:envs/my_env.py")
-    assert repo_id == "user/repo"
-    assert revision == "v1.0"
-    assert file_path == "envs/my_env.py"
-
-    # repo with commit hash
-    repo_id, revision, file_path = _parse_hub_url("org/repo@abc123def456")
-    assert repo_id == "org/repo"
-    assert revision == "abc123def456"
-    assert file_path == "env.py"
-
-
-def test_normalize_hub_result():
-    """Test normalization of different return types from hub make_env."""
-    # test with VectorEnv (most common case)
-    mock_vec_env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1")])
-    result = _normalize_hub_result(mock_vec_env)
-    assert isinstance(result, dict)
-    assert len(result) == 1
-    suite_name = next(iter(result))
-    assert 0 in result[suite_name]
-    assert isinstance(result[suite_name][0], gym.vector.VectorEnv)
-    mock_vec_env.close()
-
-    # test with single Env
-    mock_env = gym.make("CartPole-v1")
-    result = _normalize_hub_result(mock_env)
-    assert isinstance(result, dict)
-    suite_name = next(iter(result))
-    assert 0 in result[suite_name]
-    assert isinstance(result[suite_name][0], gym.vector.VectorEnv)
-    result[suite_name][0].close()
-
-    # test with dict (already normalized)
-    mock_vec_env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1")])
-    input_dict = {"my_suite": {0: mock_vec_env}}
-    result = _normalize_hub_result(input_dict)
-    assert result == input_dict
-    assert "my_suite" in result
-    assert 0 in result["my_suite"]
-    mock_vec_env.close()
-
-    # test with invalid type
-    with pytest.raises(ValueError, match="Hub `make_env` must return"):
-        _normalize_hub_result("invalid_type")
-
-
-def test_make_env_from_hub_requires_trust_remote_code():
-    """Test that loading from hub requires explicit trust_remote_code=True."""
-    hub_id = "lerobot/cartpole-env"
-
-    # Should raise RuntimeError when trust_remote_code=False (default)
-    with pytest.raises(RuntimeError, match="Refusing to execute remote code"):
-        make_env(hub_id, trust_remote_code=False)
-
-    # Should also raise when not specified (defaults to False)
-    with pytest.raises(RuntimeError, match="Refusing to execute remote code"):
-        make_env(hub_id)
-
-
-@pytest.mark.parametrize(
-    "hub_id",
-    [
-        "lerobot/cartpole-env",
-        "lerobot/cartpole-env@main",
-        "lerobot/cartpole-env:env.py",
-    ],
-)
-def test_make_env_from_hub_with_trust(hub_id):
-    """Test loading environment from Hugging Face Hub with trust_remote_code=True."""
-    # load environment from hub
-    envs_dict = make_env(hub_id, n_envs=2, trust_remote_code=True)
-
-    # verify structure
-    assert isinstance(envs_dict, dict)
-    assert len(envs_dict) >= 1
-
-    # get the first suite and task
-    suite_name = next(iter(envs_dict))
-    task_id = next(iter(envs_dict[suite_name]))
-    env = envs_dict[suite_name][task_id]
-
-    # verify it's a vector environment
-    assert isinstance(env, gym.vector.VectorEnv)
-    assert env.num_envs == 2
-
-    # test basic environment interaction
-    obs, info = env.reset()
-    assert obs is not None
-    assert isinstance(obs, (dict, np.ndarray))
-
-    # take a random action
-    action = env.action_space.sample()
-    obs, reward, terminated, truncated, info = env.step(action)
-    assert obs is not None
-    assert isinstance(reward, np.ndarray)
-    assert len(reward) == 2
-
-    # clean up
-    env.close()
-
-
-def test_make_env_from_hub_async():
-    """Test loading hub environment with async vector environments."""
-    hub_id = "lerobot/cartpole-env"
-
-    # load with async envs
-    envs_dict = make_env(hub_id, n_envs=2, use_async_envs=True, trust_remote_code=True)
-
-    suite_name = next(iter(envs_dict))
-    task_id = next(iter(envs_dict[suite_name]))
-    env = envs_dict[suite_name][task_id]
-
-    # verify it's an async vector environment
-    assert isinstance(env, gym.vector.AsyncVectorEnv)
-    assert env.num_envs == 2
-
-    # test basic interaction
-    obs, info = env.reset()
-    assert obs is not None
-
-    action = env.action_space.sample()
-    obs, reward, terminated, truncated, info = env.step(action)
-    assert len(reward) == 2
-
-    # clean up
-    env.close()

tests/motors/test_damiao.py

@@ -0,0 +1,338 @@
+#!/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.
+
+"""
+Test script for Damiao motor communication and control.
+
+This script tests basic functionality of a single Damiao motor via CAN bus:
+1. Connects to CAN interface
+2. Discovers and enables the motor
+3. Reads current position
+4. Sets zero position
+5. Writes target positions
+6. Disables torque
+
+Requirements:
+- Motor must be connected and powered (24V)
+- CAN interface must be configured (e.g., can0)
+- Motor ID must be set to 0x01 (send) and 0x11 (receive)
+
+Setup CAN interface:
+    sudo ip link set can0 type can bitrate 1000000
+    sudo ip link set can0 up
+    
+Verify connection:
+    candump can0  # In another terminal
+    cansend can0 001#FFFFFFFFFFFFFFFC  # Should enable motor and LED turns green
+"""
+
+import time
+
+import pytest
+
+from lerobot.motors import Motor, MotorNormMode
+from lerobot.motors.damiao import DamiaoMotorsBus
+from lerobot.motors.damiao.tables import MotorType
+
+@pytest.fixture
+def can_port(request):
+    """Get CAN port from command line or raise error if not provided."""
+    port = request.config.getoption("--can-port")
+    if port is None:
+        pytest.skip("CAN port not specified. Use --can-port to specify the CAN interface.")
+    return port
+
+
+@pytest.mark.hardware
+def test_single_motor_basic_operations(can_port):
+    """
+    Test basic operations with a single Damiao motor.
+    
+    This test requires actual hardware and is skipped by default.
+    To run with hardware, use: pytest tests/motors/test_damiao.py --run-hardware --can-port PORT
+    """
+    
+    # Configuration
+    MOTOR_ID = 0x01  # Sender CAN ID
+    MOTOR_RECV_ID = 0x11  # Receiver/Master ID
+    MOTOR_TYPE = "dm4310"
+    MOTOR_NAME = "test_motor"
+    
+    print(f"\n{'='*60}")
+    print("Damiao Motor Test - Single Motor Basic Operations")
+    print(f"{'='*60}\n")
+    
+    # Step 1: Create motor configuration
+    print(f"Step 1: Creating motor configuration...")
+    print(f"  - Motor ID: 0x{MOTOR_ID:02X} (send) / 0x{MOTOR_RECV_ID:02X} (recv)")
+    print(f"  - Motor Type: {MOTOR_TYPE}")
+    print(f"  - CAN Port: {can_port}")
+    
+    motor = Motor(MOTOR_ID, MOTOR_TYPE, MotorNormMode.DEGREES)
+    motor.recv_id = MOTOR_RECV_ID
+    motor.motor_type = MotorType.DM4310
+    
+    motors = {MOTOR_NAME: motor}
+    
+    # Step 2: Connect to CAN bus
+    print(f"\nStep 2: Connecting to CAN bus...")
+    bus = DamiaoMotorsBus(port=can_port, motors=motors)
+    
+    try:
+        bus.connect(handshake=True)
+        print(f"  ✓ Connected to {can_port}")
+    except Exception as e:
+        print(f"  ✗ Failed to connect: {e}")
+        print("\nTroubleshooting:")
+        print(f"  1. Check CAN interface is up: ip link show {can_port}")
+        print(f"  2. Setup if needed: sudo ip link set {can_port} type can bitrate 1000000")
+        print(f"  3. Bring up: sudo ip link set {can_port} up")
+        print(f"  4. Test with: cansend {can_port} 001#FFFFFFFFFFFFFFFC")
+        return
+    
+    try:
+        # Step 3: Enable motor (torque on)
+        print(f"\nStep 3: Enabling motor...")
+        bus.enable_torque(MOTOR_NAME)
+        time.sleep(0.1)
+        print(f"  ✓ Motor enabled (LED should be green)")
+        
+        # Step 4: Read current position
+        print(f"\nStep 4: Reading current position...")
+        current_pos = bus.read("Present_Position", MOTOR_NAME, normalize=False)
+        current_vel = bus.read("Present_Velocity", MOTOR_NAME, normalize=False)
+        current_torque = bus.read("Present_Torque", MOTOR_NAME, normalize=False)
+        
+        print(f"  Current State:")
+        print(f"    Position: {current_pos:8.2f}°")
+        print(f"    Velocity: {current_vel:8.2f}°/s")
+        print(f"    Torque:   {current_torque:8.3f} N·m")
+        
+        # Step 5: Set zero position
+        print(f"\nStep 5: Setting current position as zero...")
+        bus.set_zero_position([MOTOR_NAME])
+        time.sleep(0.2)
+        
+        new_pos = bus.read("Present_Position", MOTOR_NAME, normalize=False)
+        print(f"  Position after zero: {new_pos:8.2f}°")
+        print(f"  ✓ Zero position set")
+        
+        # Step 6: Test position commands
+        print(f"\nStep 6: Testing position control...")
+        
+        test_positions = [0.0, 45.0, -45.0, 0.0]
+        
+        for target_pos in test_positions:
+            print(f"\n  Moving to {target_pos:6.1f}°...")
+            bus.write("Goal_Position", MOTOR_NAME, target_pos, normalize=False)
+            time.sleep(1.0)  # Allow motor to move
+            
+            actual_pos = bus.read("Present_Position", MOTOR_NAME, normalize=False)
+            error = abs(actual_pos - target_pos)
+            
+            print(f"    Target:   {target_pos:8.2f}°")
+            print(f"    Actual:   {actual_pos:8.2f}°")
+            print(f"    Error:    {error:8.2f}°")
+            
+            if error > 10.0:
+                print(f"    ⚠ Large position error!")
+            else:
+                print(f"    ✓ Position reached")
+        
+        # Step 7: Test MIT control with custom gains
+        print(f"\nStep 7: Testing MIT control with custom gains...")
+        print(f"  Using lower gains for gentler movement...")
+        
+        # Lower gains for smoother motion
+        bus._mit_control(
+            MOTOR_NAME,
+            kp=5.0,   # Lower position gain
+            kd=0.3,   # Lower damping
+            position_degrees=30.0,
+            velocity_deg_per_sec=0.0,
+            torque=0.0
+        )
+        time.sleep(1.5)
+        
+        final_pos = bus.read("Present_Position", MOTOR_NAME, normalize=False)
+        print(f"  Final position: {final_pos:8.2f}°")
+        print(f"  ✓ MIT control test complete")
+        
+        # Step 8: Return to zero
+        print(f"\nStep 8: Returning to zero position...")
+        bus.write("Goal_Position", MOTOR_NAME, 0.0, normalize=False)
+        time.sleep(1.0)
+        
+        final_pos = bus.read("Present_Position", MOTOR_NAME, normalize=False)
+        print(f"  Final position: {final_pos:8.2f}°")
+        
+    finally:
+        # Step 9: Disable motor
+        print(f"\nStep 9: Disabling motor...")
+        if bus.is_connected:
+            bus.disable_torque(MOTOR_NAME)
+            time.sleep(0.1)
+            print(f"  ✓ Motor disabled (torque off)")
+        
+        # Step 10: Disconnect
+        print(f"\nStep 10: Disconnecting...")
+        if bus.is_connected:
+            bus.disconnect(disable_torque=False)  # Already disabled
+            print(f"  ✓ Disconnected from {can_port}")
+    
+    print(f"\n{'='*60}")
+    print("Test completed successfully!")
+    print(f"{'='*60}\n")
+
+
+@pytest.mark.hardware
+def test_motor_discovery_and_setup(can_port):
+    """
+    Test motor discovery and ID configuration.
+    
+    Note: This test requires the Damiao Debugging Tools for actual ID changes.
+    This test only demonstrates the bus scan functionality.
+    """
+    
+    print(f"\n{'='*60}")
+    print("Damiao Motor Discovery Test")
+    print(f"{'='*60}\n")
+    
+    print("Note: Motor ID configuration must be done via Damiao Debugging Tools")
+    print("See: https://docs.openarm.dev/software/setup/motor-id")
+    print()
+    
+    # Test if CAN interface is accessible
+    print(f"Testing CAN interface: {can_port}")
+    
+    # Create a minimal motor bus for testing
+    test_motor = Motor(0x01, "dm4310", MotorNormMode.DEGREES)
+    test_motor.recv_id = 0x11
+    test_motor.motor_type = MotorType.DM4310
+    
+    bus = DamiaoMotorsBus(port=can_port, motors={"test": test_motor})
+    
+    try:
+        bus.connect(handshake=False)
+        print(f"✓ CAN interface {can_port} is accessible")
+        
+        # Try to communicate with motor at 0x01
+        print(f"\nLooking for motor at ID 0x01...")
+        try:
+            bus._refresh_motor("test")
+            msg = bus._recv_motor_response(timeout=0.5)
+            if msg:
+                print(f"✓ Motor found at ID 0x01, response ID: 0x{msg.arbitration_id:02X}")
+            else:
+                print(f"✗ No response from motor")
+                print("\nTroubleshooting:")
+                print("  1. Verify motor is powered (24V)")
+                print("  2. Check CAN wiring (CANH, CANL)")
+                print("  3. Verify motor ID is set to 0x01")
+                print("  4. Enable with: cansend can0 001#FFFFFFFFFFFFFFFC")
+        except Exception as e:
+                print(f"✗ Error communicating with motor: {e}")
+        
+    except Exception as e:
+        print(f"✗ Failed to access CAN interface: {e}")
+        print("\nSetup CAN interface:")
+        print(f"  sudo ip link set {can_port} type can bitrate 1000000")
+        print(f"  sudo ip link set {can_port} up")
+    
+    finally:
+        if bus.is_connected:
+            bus.disconnect(disable_torque=True)
+    
+    print(f"\n{'='*60}\n")
+
+
+@pytest.mark.hardware
+def test_multi_motor_sync_operations(can_port):
+    """
+    Test synchronized read/write with multiple motors.
+    
+    This demonstrates how to control multiple motors simultaneously.
+    """
+    
+    print(f"\n{'='*60}")
+    print("Damiao Multi-Motor Sync Test")
+    print(f"{'='*60}\n")
+    
+    # Setup motors (adjust IDs as needed)
+    motors = {
+        "joint_1": Motor(0x01, "dm4310", MotorNormMode.DEGREES),
+        "joint_2": Motor(0x02, "dm4310", MotorNormMode.DEGREES),
+    }
+    
+    motors["joint_1"].recv_id = 0x11
+    motors["joint_1"].motor_type = MotorType.DM4310
+    motors["joint_2"].recv_id = 0x12
+    motors["joint_2"].motor_type = MotorType.DM4310
+    
+    bus = DamiaoMotorsBus(port=can_port, motors=motors)
+    
+    try:
+        bus.connect()
+        bus.enable_torque()
+        
+        print("Reading all motor positions...")
+        positions = bus.sync_read("Present_Position")
+        for motor, pos in positions.items():
+            print(f"  {motor}: {pos:.2f}°")
+        
+        print("\nMoving all motors to 45°...")
+        target_positions = {motor: 45.0 for motor in motors}
+        bus.sync_write("Goal_Position", target_positions)
+        time.sleep(2.0)
+        
+        positions = bus.sync_read("Present_Position")
+        print("Final positions:")
+        for motor, pos in positions.items():
+            print(f"  {motor}: {pos:.2f}°")
+        
+    except Exception as e:
+        print(f"✗ Test failed: {e}")
+        print("\nThis is expected on macOS without proper CAN hardware.")
+        print("macOS does not support SocketCAN natively (Linux-only feature).")
+        print("For macOS, you need a USB-CAN adapter with SLCAN support.")
+    finally:
+        if bus.is_connected:
+            bus.disable_torque()
+            bus.disconnect()
+    
+    print(f"\n{'='*60}\n")
+
+
+if __name__ == "__main__":
+    print("Damiao Motor Test Suite")
+    print("=" * 60)
+    print("\nThese tests require actual hardware to run.")
+    print("Please ensure:")
+    print("  1. Motor is connected and powered (24V)")
+    print("  2. CAN interface is configured")
+    print("  3. Motor ID is set to 0x01/0x11")
+    print("\nTo run tests with hardware:")
+    print("\n  Linux (SocketCAN):")
+    print("    sudo ip link set can0 type can bitrate 1000000")
+    print("    sudo ip link set can0 up")
+    print("    pytest tests/motors/test_damiao.py --run-hardware --can-port can0")
+    print("\n  macOS (USB-CAN adapter with SLCAN):")
+    print("    pytest tests/motors/test_damiao.py --run-hardware --can-port /dev/cu.usbmodem00000000050C1")
+    print("\nTo run without hardware (tests will be skipped):")
+    print("  pytest tests/motors/test_damiao.py")
+    print("\nNote: The --run-hardware and --can-port flags are configured in tests/conftest.py")
+    print("=" * 60)
+

tests/policies/pi0_pi05/test_pi05_rtc.py

@@ -1,336 +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.
-
-"""Test PI0.5 policy with Real-Time Chunking (RTC) enabled during inference."""
-
-import os
-
-import pytest
-import torch
-
-# Skip this entire module in CI
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local OpenPI installation and is not meant for CI",
-)
-
-from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule  # noqa: E402
-from lerobot.policies.pi05 import PI05Config, PI05Policy, make_pi05_pre_post_processors  # noqa: E402
-from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402
-from lerobot.utils.random_utils import set_seed  # noqa: E402
-from tests.utils import require_cuda  # noqa: E402
-
-
-@require_cuda
-def test_pi05_rtc_initialization():
-    """Test PI0.5 policy can initialize RTC processor."""
-    set_seed(42)
-
-    config = PI05Config(max_action_dim=7, max_state_dim=14, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Instantiate policy
-    policy = PI05Policy(config)
-
-    # Verify RTC processor is initialized
-    assert hasattr(policy, "rtc_processor")
-    assert policy.rtc_processor is not None
-    assert policy.rtc_processor.rtc_config.enabled is True
-
-    print("✓ PI0.5 RTC initialization: Test passed")
-
-
-@require_cuda
-def test_pi05_rtc_initialization_without_rtc_config():
-    """Test PI0.5 policy can initialize without RTC config."""
-    set_seed(42)
-
-    config = PI05Config(max_action_dim=7, max_state_dim=14, dtype="float32")
-
-    # Instantiate policy
-    policy = PI05Policy(config)
-
-    # Verify RTC processor is not initialized
-    assert hasattr(policy, "rtc_processor")
-    assert policy.rtc_processor is None
-    assert policy.model.rtc_processor is None
-    assert policy._rtc_enabled() is False
-
-    print("✓ PI0.5 RTC initialization without RTC config: Test passed")
-
-
-@require_cuda
-def test_pi05_rtc_inference_with_prev_chunk():
-    """Test PI0.5 policy inference with RTC and previous chunk."""
-    set_seed(42)
-
-    config = PI05Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats (PI0.5 uses QUANTILES normalization)
-    dataset_stats = {
-        "observation.state": {
-            "mean": torch.zeros(14),
-            "std": torch.ones(14),
-            "q01": -torch.ones(14),
-            "q99": torch.ones(14),
-        },
-        "action": {
-            "mean": torch.zeros(7),
-            "std": torch.ones(7),
-            "q01": -torch.ones(7),
-            "q99": torch.ones(7),
-        },
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and preprocessor
-    policy = PI05Policy(config)
-    policy.eval()
-    preprocessor, _ = make_pi05_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    # Create previous chunk
-    prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC and previous chunk
-        actions_with_rtc = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=prev_chunk,
-            inference_delay=4,
-            execution_horizon=10,
-        )
-
-        # Test without RTC for comparison
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    # Verify shapes
-    assert actions_with_rtc.shape == (1, config.chunk_size, 7)
-    assert actions_without_rtc.shape == (1, config.chunk_size, 7)
-
-    # With previous chunk, actions should be different (RTC guidance applied)
-    assert not torch.allclose(actions_with_rtc, actions_without_rtc, rtol=1e-3)
-
-    print("✓ PI0.5 RTC inference with prev_chunk: Test passed")
-
-
-@require_cuda
-def test_pi05_rtc_inference_without_prev_chunk():
-    """Test PI0.5 policy inference with RTC but no previous chunk (RTC should have no effect)."""
-    set_seed(42)
-
-    config = PI05Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats (PI0.5 uses QUANTILES normalization)
-    dataset_stats = {
-        "observation.state": {
-            "mean": torch.zeros(14),
-            "std": torch.ones(14),
-            "q01": -torch.ones(14),
-            "q99": torch.ones(14),
-        },
-        "action": {
-            "mean": torch.zeros(7),
-            "std": torch.ones(7),
-            "q01": -torch.ones(7),
-            "q99": torch.ones(7),
-        },
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and preprocessor
-    policy = PI05Policy(config)
-    policy.eval()
-    preprocessor, _ = make_pi05_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC enabled but no previous chunk
-        actions_with_rtc_no_prev = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=None,
-        )
-
-        # Test without RTC
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    # Without previous chunk, RTC should have no effect
-    assert torch.allclose(actions_with_rtc_no_prev, actions_without_rtc, rtol=1e-5)
-
-    print("✓ PI0.5 RTC inference without prev_chunk: Test passed")
-
-
-@require_cuda
-def test_pi05_rtc_validation_rules():
-    """Test PI0.5 policy with RTC follows all three validation rules."""
-    set_seed(42)
-
-    config = PI05Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats (PI0.5 uses QUANTILES normalization)
-    dataset_stats = {
-        "observation.state": {
-            "mean": torch.zeros(14),
-            "std": torch.ones(14),
-            "q01": -torch.ones(14),
-            "q99": torch.ones(14),
-        },
-        "action": {
-            "mean": torch.zeros(7),
-            "std": torch.ones(7),
-            "q01": -torch.ones(7),
-            "q99": torch.ones(7),
-        },
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and preprocessor
-    policy = PI05Policy(config)
-    policy.eval()
-    preprocessor, _ = make_pi05_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    # Create previous chunk
-    prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-    inference_delay = 4
-    execution_horizon = 10
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC
-        actions_with_rtc = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=prev_chunk,
-            inference_delay=inference_delay,
-            execution_horizon=execution_horizon,
-        )
-
-        # Test without RTC
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    assert not torch.allclose(actions_with_rtc, actions_without_rtc, rtol=1e-3)

tests/policies/pi0_pi05/test_pi0_rtc.py

@@ -1,378 +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.
-
-"""Test PI0 policy with Real-Time Chunking (RTC) enabled during inference."""
-
-import os
-
-import pytest
-import torch
-
-# Skip this entire module in CI
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local OpenPI installation and is not meant for CI",
-)
-
-from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule  # noqa: E402
-from lerobot.policies.pi0 import PI0Config, PI0Policy, make_pi0_pre_post_processors  # noqa: E402
-from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402
-from lerobot.utils.random_utils import set_seed  # noqa: E402
-from tests.utils import require_cuda  # noqa: E402
-
-
-@require_cuda
-def test_pi0_rtc_initialization():
-    """Test PI0 policy can initialize RTC processor."""
-    set_seed(42)
-
-    config = PI0Config(max_action_dim=7, max_state_dim=14, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Instantiate policy
-    policy = PI0Policy(config)
-
-    # Verify RTC processor is initialized
-    assert hasattr(policy, "rtc_processor")
-    assert policy.rtc_processor is not None
-    assert policy.rtc_processor.rtc_config.enabled is True
-
-    print("✓ PI0 RTC initialization: Test passed")
-
-
-@require_cuda
-def test_pi0_rtc_initialization_without_rtc_config():
-    """Test PI0 policy can initialize without RTC config."""
-    set_seed(42)
-
-    config = PI0Config(max_action_dim=7, max_state_dim=14, dtype="float32")
-
-    # Instantiate policy
-    policy = PI0Policy(config)
-
-    # Verify RTC processor is not initialized
-    assert hasattr(policy, "rtc_processor")
-    assert policy.rtc_processor is None
-    assert policy.model.rtc_processor is None
-    assert policy._rtc_enabled() is False
-
-    print("✓ PI0 RTC initialization without RTC config: Test passed")
-
-
-def test_pi0_rtc_inference_with_prev_chunk():
-    """Test PI0 policy inference with RTC and previous chunk."""
-    set_seed(42)
-
-    config = PI0Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and preprocessor
-    policy = PI0Policy(config)
-    policy.eval()
-    preprocessor, _ = make_pi0_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    # Create previous chunk
-    prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC and previous chunk
-        actions_with_rtc = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=prev_chunk,
-            inference_delay=4,
-            execution_horizon=10,
-        )
-
-        # Test without RTC for comparison
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    # Verify shapes
-    assert actions_with_rtc.shape == (1, config.chunk_size, 7)
-    assert actions_without_rtc.shape == (1, config.chunk_size, 7)
-
-    # With previous chunk, actions should be different (RTC guidance applied)
-    assert not torch.allclose(actions_with_rtc, actions_without_rtc, rtol=1e-3)
-
-    print("✓ PI0 RTC inference with prev_chunk: Test passed")
-
-
-@require_cuda
-def test_pi0_rtc_inference_without_prev_chunk():
-    """Test PI0 policy inference with RTC but no previous chunk (RTC should have no effect)."""
-    set_seed(42)
-
-    config = PI0Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and preprocessor
-    policy = PI0Policy(config)
-    policy.eval()
-    preprocessor, _ = make_pi0_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC enabled but no previous chunk
-        actions_with_rtc_no_prev = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=None,
-        )
-
-        # Test without RTC
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    # Without previous chunk, RTC should have no effect
-    assert torch.allclose(actions_with_rtc_no_prev, actions_without_rtc, rtol=1e-5)
-
-    print("✓ PI0 RTC inference without prev_chunk: Test passed")
-
-
-@require_cuda
-def test_pi0_rtc_validation_rules():
-    """Test PI0 policy with RTC follows all three validation rules."""
-    set_seed(42)
-
-    config = PI0Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and preprocessor
-    policy = PI0Policy(config)
-    policy.eval()
-    preprocessor, _ = make_pi0_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    # Create previous chunk
-    prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-    inference_delay = 4
-    execution_horizon = 10
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC
-        actions_with_rtc = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=prev_chunk,
-            inference_delay=inference_delay,
-            execution_horizon=execution_horizon,
-        )
-
-        # Test without RTC
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    assert not torch.allclose(actions_with_rtc, actions_without_rtc, rtol=1e-3)
-
-    """Test PI0 with different RTC attention schedules."""
-    set_seed(42)
-
-    schedules = [
-        RTCAttentionSchedule.ZEROS,
-        RTCAttentionSchedule.ONES,
-        RTCAttentionSchedule.LINEAR,
-        RTCAttentionSchedule.EXP,
-    ]
-
-    config = PI0Config(max_action_dim=7, max_state_dim=14, chunk_size=50, dtype="float32")
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    device = config.device
-
-    for schedule in schedules:
-        print(f"Testing schedule: {schedule}")
-
-        # Add RTC config with specific schedule
-        config.rtc_config = RTCConfig(
-            enabled=True,
-            execution_horizon=10,
-            max_guidance_weight=5.0,
-            prefix_attention_schedule=schedule,
-            debug=False,
-        )
-
-        # Instantiate policy
-        policy = PI0Policy(config)
-        policy.eval()
-        preprocessor, _ = make_pi0_pre_post_processors(config=config, dataset_stats=dataset_stats)
-
-        # Create dummy batch
-        batch = {
-            "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-            "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-            "task": ["Pick up the object"],
-        }
-        batch = preprocessor(batch)
-
-        # Create previous chunk
-        prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-        with torch.no_grad():
-            noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-            actions = policy.predict_action_chunk(
-                batch,
-                noise=noise,
-                prev_chunk_left_over=prev_chunk,
-                inference_delay=4,
-                execution_horizon=10,
-            )
-
-        # Verify shape
-        assert actions.shape == (1, config.chunk_size, 7)
-        print(f"  ✓ Schedule {schedule}: Test passed")
-
-    print("✓ PI0 RTC different schedules: All schedules tested")

tests/policies/rtc/test_action_queue.py

@@ -1,825 +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.
-
-"""Tests for RTC ActionQueue module."""
-
-import threading
-import time
-
-import pytest
-import torch
-
-from lerobot.policies.rtc.action_queue import ActionQueue
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-
-# ====================== Fixtures ======================
-
-
-@pytest.fixture
-def rtc_config_enabled():
-    """Create an RTC config with RTC enabled."""
-    return RTCConfig(enabled=True, execution_horizon=10, max_guidance_weight=1.0)
-
-
-@pytest.fixture
-def rtc_config_disabled():
-    """Create an RTC config with RTC disabled."""
-    return RTCConfig(enabled=False, execution_horizon=10, max_guidance_weight=1.0)
-
-
-@pytest.fixture
-def sample_actions():
-    """Create sample action tensors for testing."""
-    return {
-        "original": torch.randn(50, 6),  # (time_steps, action_dim)
-        "processed": torch.randn(50, 6),
-        "short": torch.randn(10, 6),
-        "longer": torch.randn(100, 6),
-    }
-
-
-@pytest.fixture
-def action_queue_rtc_enabled(rtc_config_enabled):
-    """Create an ActionQueue with RTC enabled."""
-    return ActionQueue(rtc_config_enabled)
-
-
-@pytest.fixture
-def action_queue_rtc_disabled(rtc_config_disabled):
-    """Create an ActionQueue with RTC disabled."""
-    return ActionQueue(rtc_config_disabled)
-
-
-# ====================== Initialization Tests ======================
-
-
-def test_action_queue_initialization_rtc_enabled(rtc_config_enabled):
-    """Test ActionQueue initializes correctly with RTC enabled."""
-    queue = ActionQueue(rtc_config_enabled)
-    assert queue.queue is None
-    assert queue.original_queue is None
-    assert queue.last_index == 0
-    assert queue.cfg.enabled is True
-
-
-def test_action_queue_initialization_rtc_disabled(rtc_config_disabled):
-    """Test ActionQueue initializes correctly with RTC disabled."""
-    queue = ActionQueue(rtc_config_disabled)
-    assert queue.queue is None
-    assert queue.original_queue is None
-    assert queue.last_index == 0
-    assert queue.cfg.enabled is False
-
-
-# ====================== get() Tests ======================
-
-
-def test_get_returns_none_when_empty(action_queue_rtc_enabled):
-    """Test get() returns None when queue is empty."""
-    action = action_queue_rtc_enabled.get()
-    assert action is None
-
-
-def test_get_returns_actions_sequentially(action_queue_rtc_enabled, sample_actions):
-    """Test get() returns actions in sequence."""
-    # Initialize queue with actions
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-
-    # Get first action
-    action1 = action_queue_rtc_enabled.get()
-    assert action1 is not None
-    assert action1.shape == (6,)
-    assert torch.equal(action1, sample_actions["processed"][0])
-
-    # Get second action
-    action2 = action_queue_rtc_enabled.get()
-    assert action2 is not None
-    assert torch.equal(action2, sample_actions["processed"][1])
-
-
-def test_get_returns_none_after_exhaustion(action_queue_rtc_enabled, sample_actions):
-    """Test get() returns None after all actions are consumed."""
-    # Use short action sequence
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume all actions
-    for _ in range(10):
-        action = action_queue_rtc_enabled.get()
-        assert action is not None
-
-    # Next get should return None
-    action = action_queue_rtc_enabled.get()
-    assert action is None
-
-
-def test_get_increments_last_index(action_queue_rtc_enabled, sample_actions):
-    """Test get() increments last_index correctly."""
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-
-    assert action_queue_rtc_enabled.last_index == 0
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.last_index == 1
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.last_index == 2
-
-
-# ====================== qsize() Tests ======================
-
-
-def test_qsize_returns_zero_when_empty(action_queue_rtc_enabled):
-    """Test qsize() returns 0 when queue is empty."""
-    assert action_queue_rtc_enabled.qsize() == 0
-
-
-def test_qsize_returns_correct_size(action_queue_rtc_enabled, sample_actions):
-    """Test qsize() returns correct number of remaining actions."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    assert action_queue_rtc_enabled.qsize() == 10
-
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.qsize() == 9
-
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.qsize() == 8
-
-
-def test_qsize_after_exhaustion(action_queue_rtc_enabled, sample_actions):
-    """Test qsize() returns 0 after queue is exhausted."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume all actions
-    for _ in range(10):
-        action_queue_rtc_enabled.get()
-
-    assert action_queue_rtc_enabled.qsize() == 0
-
-
-# ====================== empty() Tests ======================
-
-
-def test_empty_returns_true_when_empty(action_queue_rtc_enabled):
-    """Test empty() returns True when queue is empty."""
-    assert action_queue_rtc_enabled.empty() is True
-
-
-def test_empty_returns_false_when_not_empty(action_queue_rtc_enabled, sample_actions):
-    """Test empty() returns False when queue has actions."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    assert action_queue_rtc_enabled.empty() is False
-
-
-def test_empty_after_partial_consumption(action_queue_rtc_enabled, sample_actions):
-    """Test empty() returns False after partial consumption."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    action_queue_rtc_enabled.get()
-    action_queue_rtc_enabled.get()
-
-    assert action_queue_rtc_enabled.empty() is False
-
-
-def test_empty_after_full_consumption(action_queue_rtc_enabled, sample_actions):
-    """Test empty() returns True after all actions consumed."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume all
-    for _ in range(10):
-        action_queue_rtc_enabled.get()
-
-    assert action_queue_rtc_enabled.empty() is True
-
-
-# ====================== get_action_index() Tests ======================
-
-
-def test_get_action_index_initial_value(action_queue_rtc_enabled):
-    """Test get_action_index() returns 0 initially."""
-    assert action_queue_rtc_enabled.get_action_index() == 0
-
-
-def test_get_action_index_after_consumption(action_queue_rtc_enabled, sample_actions):
-    """Test get_action_index() tracks consumption correctly."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    assert action_queue_rtc_enabled.get_action_index() == 0
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.get_action_index() == 1
-    action_queue_rtc_enabled.get()
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.get_action_index() == 3
-
-
-# ====================== get_left_over() Tests ======================
-
-
-def test_get_left_over_returns_none_when_empty(action_queue_rtc_enabled):
-    """Test get_left_over() returns None when queue is empty."""
-    leftover = action_queue_rtc_enabled.get_left_over()
-    assert leftover is None
-
-
-def test_get_left_over_returns_all_when_unconsumed(action_queue_rtc_enabled, sample_actions):
-    """Test get_left_over() returns all original actions when none consumed."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    leftover = action_queue_rtc_enabled.get_left_over()
-    assert leftover is not None
-    assert leftover.shape == (10, 6)
-    assert torch.equal(leftover, sample_actions["short"])
-
-
-def test_get_left_over_returns_remaining_after_consumption(action_queue_rtc_enabled, sample_actions):
-    """Test get_left_over() returns only remaining original actions."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume 3 actions
-    action_queue_rtc_enabled.get()
-    action_queue_rtc_enabled.get()
-    action_queue_rtc_enabled.get()
-
-    leftover = action_queue_rtc_enabled.get_left_over()
-    assert leftover is not None
-    assert leftover.shape == (7, 6)
-    assert torch.equal(leftover, sample_actions["short"][3:])
-
-
-def test_get_left_over_returns_empty_after_exhaustion(action_queue_rtc_enabled, sample_actions):
-    """Test get_left_over() returns empty tensor after all consumed."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume all
-    for _ in range(10):
-        action_queue_rtc_enabled.get()
-
-    leftover = action_queue_rtc_enabled.get_left_over()
-    assert leftover is not None
-    assert leftover.shape == (0, 6)
-
-
-# ====================== merge() with RTC Enabled Tests ======================
-
-
-def test_merge_replaces_queue_when_rtc_enabled(action_queue_rtc_enabled, sample_actions):
-    """Test merge() replaces queue when RTC is enabled."""
-    # Add initial actions
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    assert action_queue_rtc_enabled.qsize() == 10
-
-    # Consume some actions
-    action_queue_rtc_enabled.get()
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.qsize() == 8
-
-    # Merge new actions - should replace, not append
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=5)
-
-    # Queue should be replaced with new actions minus delay
-    # Original has 50 actions, delay is 5, so remaining is 45
-    assert action_queue_rtc_enabled.qsize() == 45
-    assert action_queue_rtc_enabled.get_action_index() == 0
-
-
-def test_merge_respects_real_delay(action_queue_rtc_enabled, sample_actions):
-    """Test merge() correctly applies real_delay when RTC is enabled."""
-    delay = 10
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=delay)
-
-    # Queue should have original length minus delay
-    expected_size = len(sample_actions["original"]) - delay
-    assert action_queue_rtc_enabled.qsize() == expected_size
-
-    # First action should be the one at index [delay]
-    first_action = action_queue_rtc_enabled.get()
-    assert torch.equal(first_action, sample_actions["processed"][delay])
-
-
-def test_merge_resets_last_index_when_rtc_enabled(action_queue_rtc_enabled, sample_actions):
-    """Test merge() resets last_index to 0 when RTC is enabled."""
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    action_queue_rtc_enabled.get()
-    action_queue_rtc_enabled.get()
-    assert action_queue_rtc_enabled.last_index == 2
-
-    # Merge new actions
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=5)
-
-    assert action_queue_rtc_enabled.last_index == 0
-
-
-def test_merge_with_zero_delay(action_queue_rtc_enabled, sample_actions):
-    """Test merge() with zero delay keeps all actions."""
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-
-    assert action_queue_rtc_enabled.qsize() == len(sample_actions["original"])
-
-
-def test_merge_with_large_delay(action_queue_rtc_enabled, sample_actions):
-    """Test merge() with delay larger than action sequence."""
-    # Delay is larger than sequence length
-    delay = 100
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=delay)
-
-    # Queue should be empty (delay >= length)
-    assert action_queue_rtc_enabled.qsize() == 0
-
-
-# ====================== merge() with RTC Disabled Tests ======================
-
-
-def test_merge_appends_when_rtc_disabled(action_queue_rtc_disabled, sample_actions):
-    """Test merge() appends actions when RTC is disabled."""
-    # Add initial actions
-    action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    initial_size = action_queue_rtc_disabled.qsize()
-    assert initial_size == 10
-
-    # Merge more actions
-    action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Should have appended
-    assert action_queue_rtc_disabled.qsize() == initial_size + 10
-
-
-def test_merge_removes_consumed_actions_when_appending(action_queue_rtc_disabled, sample_actions):
-    """Test merge() removes consumed actions before appending when RTC is disabled."""
-    # Add initial actions
-    action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    assert action_queue_rtc_disabled.qsize() == 10
-
-    # Consume 3 actions
-    action_queue_rtc_disabled.get()
-    action_queue_rtc_disabled.get()
-    action_queue_rtc_disabled.get()
-    assert action_queue_rtc_disabled.qsize() == 7
-
-    # Merge more actions
-    action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Should have 7 remaining + 10 new = 17
-    assert action_queue_rtc_disabled.qsize() == 17
-
-
-def test_merge_resets_last_index_after_append(action_queue_rtc_disabled, sample_actions):
-    """Test merge() resets last_index after appending when RTC is disabled."""
-    action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-    action_queue_rtc_disabled.get()
-    action_queue_rtc_disabled.get()
-    assert action_queue_rtc_disabled.last_index == 2
-
-    # Merge more actions
-    action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # last_index should be reset to 0
-    assert action_queue_rtc_disabled.last_index == 0
-
-
-def test_merge_ignores_delay_when_rtc_disabled(action_queue_rtc_disabled, sample_actions):
-    """Test merge() ignores real_delay parameter when RTC is disabled."""
-    action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=10)
-
-    # All actions should be in queue (delay ignored)
-    assert action_queue_rtc_disabled.qsize() == len(sample_actions["original"])
-
-
-def test_merge_first_call_with_rtc_disabled(action_queue_rtc_disabled, sample_actions):
-    """Test merge() on first call with RTC disabled."""
-    action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-
-    assert action_queue_rtc_disabled.qsize() == len(sample_actions["original"])
-    assert action_queue_rtc_disabled.last_index == 0
-
-
-# ====================== merge() with Different Action Shapes Tests ======================
-
-
-def test_merge_with_different_action_dims():
-    """Test merge() handles actions with different dimensions."""
-    cfg = RTCConfig(enabled=True, execution_horizon=10)
-    queue = ActionQueue(cfg)
-
-    # Actions with 4 dimensions instead of 6
-    actions_4d = torch.randn(20, 4)
-    queue.merge(actions_4d, actions_4d, real_delay=5)
-
-    action = queue.get()
-    assert action.shape == (4,)
-
-
-def test_merge_with_different_lengths():
-    """Test merge() handles action sequences of varying lengths."""
-    cfg = RTCConfig(enabled=False, execution_horizon=10)
-    queue = ActionQueue(cfg)
-
-    # Add sequences of different lengths
-    queue.merge(torch.randn(10, 6), torch.randn(10, 6), real_delay=0)
-    assert queue.qsize() == 10
-
-    queue.merge(torch.randn(25, 6), torch.randn(25, 6), real_delay=0)
-    assert queue.qsize() == 35
-
-
-# ====================== merge() Delay Validation Tests ======================
-
-
-def test_merge_validates_delay_consistency(action_queue_rtc_enabled, sample_actions, caplog):
-    """Test merge() validates that real_delay matches action index difference."""
-    import logging
-
-    caplog.set_level(logging.WARNING)
-
-    # Initialize queue
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume 5 actions
-    for _ in range(5):
-        action_queue_rtc_enabled.get()
-
-    # Merge with mismatched delay (should log warning)
-    # We consumed 5 actions, so index is 5. If we pass action_index_before_inference=0,
-    # then indexes_diff=5, but if real_delay=3, it will warn
-    action_queue_rtc_enabled.merge(
-        sample_actions["original"],
-        sample_actions["processed"],
-        real_delay=3,
-        action_index_before_inference=0,
-    )
-
-    # Check warning was logged
-    assert "Indexes diff is not equal to real delay" in caplog.text
-
-
-def test_merge_no_warning_when_delays_match(action_queue_rtc_enabled, sample_actions, caplog):
-    """Test merge() doesn't warn when delays are consistent."""
-    import logging
-
-    caplog.set_level(logging.WARNING)
-
-    # Initialize queue
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    # Consume 5 actions
-    for _ in range(5):
-        action_queue_rtc_enabled.get()
-
-    # Merge with matching delay
-    action_queue_rtc_enabled.merge(
-        sample_actions["original"],
-        sample_actions["processed"],
-        real_delay=5,
-        action_index_before_inference=0,
-    )
-
-    # Should not have warning
-    assert "Indexes diff is not equal to real delay" not in caplog.text
-
-
-def test_merge_skips_validation_when_action_index_none(action_queue_rtc_enabled, sample_actions, caplog):
-    """Test merge() skips delay validation when action_index_before_inference is None."""
-    import logging
-
-    caplog.set_level(logging.WARNING)
-
-    action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-    for _ in range(5):
-        action_queue_rtc_enabled.get()
-
-    # Pass None for action_index_before_inference
-    action_queue_rtc_enabled.merge(
-        sample_actions["original"],
-        sample_actions["processed"],
-        real_delay=999,  # Doesn't matter
-        action_index_before_inference=None,
-    )
-
-    # Should not warn (validation skipped)
-    assert "Indexes diff is not equal to real delay" not in caplog.text
-
-
-# ====================== Thread Safety Tests ======================
-
-
-def test_get_is_thread_safe(action_queue_rtc_enabled, sample_actions):
-    """Test get() is thread-safe with multiple consumers."""
-    action_queue_rtc_enabled.merge(sample_actions["longer"], sample_actions["longer"], real_delay=0)
-
-    results = []
-    errors = []
-
-    def consumer():
-        try:
-            for _ in range(25):
-                action = action_queue_rtc_enabled.get()
-                if action is not None:
-                    results.append(action)
-                time.sleep(0.001)
-        except Exception as e:
-            errors.append(e)
-
-    threads = [threading.Thread(target=consumer) for _ in range(4)]
-
-    for t in threads:
-        t.start()
-
-    for t in threads:
-        t.join()
-
-    # Should not have errors
-    assert len(errors) == 0
-
-    # Should have consumed all actions (100 total, 4 threads * 25 each)
-    assert len(results) == 100
-
-    # All results should be unique (no duplicate consumption)
-    # We can verify by checking that indices are not duplicated
-    # Since we don't track indices in results, we check total count is correct
-    assert action_queue_rtc_enabled.qsize() == 0
-
-
-def test_merge_is_thread_safe(action_queue_rtc_disabled, sample_actions):
-    """Test merge() is thread-safe with multiple producers."""
-    errors = []
-
-    def producer():
-        try:
-            for _ in range(5):
-                action_queue_rtc_disabled.merge(
-                    sample_actions["short"], sample_actions["short"], real_delay=0
-                )
-                time.sleep(0.001)
-        except Exception as e:
-            errors.append(e)
-
-    threads = [threading.Thread(target=producer) for _ in range(3)]
-
-    for t in threads:
-        t.start()
-
-    for t in threads:
-        t.join()
-
-    # Should not have errors
-    assert len(errors) == 0
-
-    # Should have accumulated all actions (3 threads * 5 merges * 10 actions = 150)
-    assert action_queue_rtc_disabled.qsize() == 150
-
-
-def test_concurrent_get_and_merge(action_queue_rtc_disabled, sample_actions):
-    """Test concurrent get() and merge() operations."""
-    errors = []
-    consumed_count = [0]
-
-    def consumer():
-        try:
-            for _ in range(50):
-                action = action_queue_rtc_disabled.get()
-                if action is not None:
-                    consumed_count[0] += 1
-                time.sleep(0.001)
-        except Exception as e:
-            errors.append(e)
-
-    def producer():
-        try:
-            for _ in range(10):
-                action_queue_rtc_disabled.merge(
-                    sample_actions["short"], sample_actions["short"], real_delay=0
-                )
-                time.sleep(0.005)
-        except Exception as e:
-            errors.append(e)
-
-    consumer_threads = [threading.Thread(target=consumer) for _ in range(2)]
-    producer_threads = [threading.Thread(target=producer) for _ in range(2)]
-
-    for t in consumer_threads + producer_threads:
-        t.start()
-
-    for t in consumer_threads + producer_threads:
-        t.join()
-
-    # Should not have errors
-    assert len(errors) == 0
-
-    # Should have consumed some or all actions (non-deterministic due to timing)
-    # Total produced: 2 producers * 10 merges * 10 actions = 200
-    # Total consumed attempts: 2 consumers * 50 = 100
-    assert consumed_count[0] <= 200
-
-
-# ====================== get_left_over() Thread Safety Tests ======================
-
-
-def test_get_left_over_is_thread_safe(action_queue_rtc_enabled, sample_actions):
-    """Test get_left_over() is thread-safe with concurrent access."""
-    action_queue_rtc_enabled.merge(sample_actions["longer"], sample_actions["longer"], real_delay=0)
-
-    errors = []
-    leftovers = []
-
-    def reader():
-        try:
-            for _ in range(20):
-                leftover = action_queue_rtc_enabled.get_left_over()
-                if leftover is not None:
-                    leftovers.append(leftover.shape[0])
-                time.sleep(0.001)
-        except Exception as e:
-            errors.append(e)
-
-    threads = [threading.Thread(target=reader) for _ in range(3)]
-
-    # Also consume some actions concurrently
-    def consumer():
-        try:
-            for _ in range(10):
-                action_queue_rtc_enabled.get()
-                time.sleep(0.002)
-        except Exception as e:
-            errors.append(e)
-
-    consumer_thread = threading.Thread(target=consumer)
-
-    all_threads = threads + [consumer_thread]
-
-    for t in all_threads:
-        t.start()
-
-    for t in all_threads:
-        t.join()
-
-    # Should not have errors
-    assert len(errors) == 0
-
-    # Leftovers should be monotonically decreasing or stable
-    # (as actions are consumed, leftover size decreases)
-    assert len(leftovers) > 0
-
-
-# ====================== Edge Cases Tests ======================
-
-
-def test_queue_with_single_action(action_queue_rtc_enabled):
-    """Test queue behavior with a single action."""
-    single_action_original = torch.randn(1, 6)
-    single_action_processed = torch.randn(1, 6)
-
-    action_queue_rtc_enabled.merge(single_action_original, single_action_processed, real_delay=0)
-
-    assert action_queue_rtc_enabled.qsize() == 1
-    action = action_queue_rtc_enabled.get()
-    assert action is not None
-    assert action.shape == (6,)
-    assert action_queue_rtc_enabled.qsize() == 0
-
-
-def test_queue_behavior_after_multiple_merge_cycles(action_queue_rtc_enabled, sample_actions):
-    """Test queue maintains correct state through multiple merge cycles."""
-    for _ in range(5):
-        action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
-
-        # Consume half
-        for _ in range(5):
-            action_queue_rtc_enabled.get()
-
-        # Merge again
-        action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=3)
-
-        assert action_queue_rtc_enabled.qsize() > 0
-
-
-def test_queue_with_all_zeros_actions(action_queue_rtc_enabled):
-    """Test queue handles all-zero action tensors."""
-    zeros_actions = torch.zeros(20, 6)
-    action_queue_rtc_enabled.merge(zeros_actions, zeros_actions, real_delay=0)
-
-    action = action_queue_rtc_enabled.get()
-    assert torch.all(action == 0)
-
-
-def test_queue_clones_input_tensors(action_queue_rtc_enabled, sample_actions):
-    """Test that merge() clones input tensors, not storing references."""
-    original_copy = sample_actions["original"].clone()
-    processed_copy = sample_actions["processed"].clone()
-
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-
-    # Modify original tensors
-    sample_actions["original"].fill_(999.0)
-    sample_actions["processed"].fill_(-999.0)
-
-    # Queue should have cloned values
-    action = action_queue_rtc_enabled.get()
-    assert not torch.equal(action, sample_actions["processed"][0])
-    assert torch.equal(action, processed_copy[0])
-
-    leftover = action_queue_rtc_enabled.get_left_over()
-    assert not torch.equal(leftover, sample_actions["original"][1:])
-    assert torch.equal(leftover, original_copy[1:])
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_queue_handles_gpu_tensors():
-    """Test queue correctly handles GPU tensors."""
-    cfg = RTCConfig(enabled=True, execution_horizon=10)
-    queue = ActionQueue(cfg)
-
-    actions_gpu = torch.randn(20, 6, device="cuda")
-    queue.merge(actions_gpu, actions_gpu, real_delay=0)
-
-    action = queue.get()
-    assert action.device.type == "cuda"
-
-    leftover = queue.get_left_over()
-    assert leftover.device.type == "cuda"
-
-
-def test_queue_handles_different_dtypes():
-    """Test queue handles actions with different dtypes."""
-    cfg = RTCConfig(enabled=True, execution_horizon=10)
-    queue = ActionQueue(cfg)
-
-    # Use float64 instead of default float32
-    actions_f64 = torch.randn(20, 6, dtype=torch.float64)
-    queue.merge(actions_f64, actions_f64, real_delay=0)
-
-    action = queue.get()
-    assert action.dtype == torch.float64
-
-
-def test_empty_with_none_queue(action_queue_rtc_enabled):
-    """Test empty() correctly handles None queue."""
-    assert action_queue_rtc_enabled.queue is None
-    assert action_queue_rtc_enabled.empty() is True
-
-
-def test_qsize_with_none_queue(action_queue_rtc_enabled):
-    """Test qsize() correctly handles None queue."""
-    assert action_queue_rtc_enabled.queue is None
-    assert action_queue_rtc_enabled.qsize() == 0
-
-
-# ====================== Integration Tests ======================
-
-
-def test_typical_rtc_workflow(action_queue_rtc_enabled, sample_actions):
-    """Test a typical RTC workflow: merge, consume, merge with delay."""
-    # First inference
-    action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-    initial_size = action_queue_rtc_enabled.qsize()
-    assert initial_size == 50
-
-    # Consume 10 actions (execution_horizon)
-    for _ in range(10):
-        action = action_queue_rtc_enabled.get()
-        assert action is not None
-
-    assert action_queue_rtc_enabled.qsize() == 40
-
-    # Second inference with delay
-    action_index_before = action_queue_rtc_enabled.get_action_index()
-
-    action_queue_rtc_enabled.merge(
-        sample_actions["original"],
-        sample_actions["processed"],
-        real_delay=5,
-        action_index_before_inference=action_index_before,
-    )
-
-    # Queue should be replaced, minus delay
-    assert action_queue_rtc_enabled.qsize() == 45
-    assert action_queue_rtc_enabled.get_action_index() == 0
-
-
-def test_typical_non_rtc_workflow(action_queue_rtc_disabled, sample_actions):
-    """Test a typical non-RTC workflow: merge, consume, merge again."""
-    # First inference
-    action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-    assert action_queue_rtc_disabled.qsize() == 50
-
-    # Consume 40 actions
-    for _ in range(40):
-        action = action_queue_rtc_disabled.get()
-        assert action is not None
-
-    assert action_queue_rtc_disabled.qsize() == 10
-
-    # Second inference (should append)
-    action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
-
-    # Should have 10 remaining + 50 new = 60
-    assert action_queue_rtc_disabled.qsize() == 60

tests/policies/rtc/test_configuration_rtc.py

@@ -1,65 +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.
-
-"""Tests for RTC configuration module."""
-
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-
-# ====================== Initialization Tests ======================
-
-
-def test_rtc_config_default_initialization():
-    """Test RTCConfig initializes with default values."""
-    config = RTCConfig()
-
-    assert config.enabled is False
-    assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
-    assert config.max_guidance_weight == 10.0
-    assert config.execution_horizon == 10
-    assert config.debug is False
-    assert config.debug_maxlen == 100
-
-
-def test_rtc_config_custom_initialization():
-    """Test RTCConfig initializes with custom values."""
-    config = RTCConfig(
-        enabled=True,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        max_guidance_weight=5.0,
-        execution_horizon=20,
-        debug=True,
-        debug_maxlen=200,
-    )
-
-    assert config.enabled is True
-    assert config.prefix_attention_schedule == RTCAttentionSchedule.EXP
-    assert config.max_guidance_weight == 5.0
-    assert config.execution_horizon == 20
-    assert config.debug is True
-    assert config.debug_maxlen == 200
-
-
-def test_rtc_config_partial_initialization():
-    """Test RTCConfig with partial custom values."""
-    config = RTCConfig(enabled=True, max_guidance_weight=15.0)
-
-    assert config.enabled is True
-    assert config.max_guidance_weight == 15.0
-    # Other values should be defaults
-    assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
-    assert config.execution_horizon == 10
-    assert config.debug is False

tests/policies/rtc/test_debug_tracker.py

@@ -1,488 +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.
-
-"""Tests for RTC debug tracker module."""
-
-import pytest
-import torch
-
-from lerobot.policies.rtc.debug_tracker import DebugStep, Tracker
-
-# ====================== Fixtures ======================
-
-
-@pytest.fixture
-def sample_tensors():
-    """Create sample tensors for testing."""
-    return {
-        "x_t": torch.randn(1, 50, 6),
-        "v_t": torch.randn(1, 50, 6),
-        "x1_t": torch.randn(1, 50, 6),
-        "correction": torch.randn(1, 50, 6),
-        "err": torch.randn(1, 50, 6),
-        "weights": torch.randn(1, 50, 1),
-    }
-
-
-@pytest.fixture
-def enabled_tracker():
-    """Create an enabled tracker with default settings."""
-    return Tracker(enabled=True, maxlen=100)
-
-
-@pytest.fixture
-def disabled_tracker():
-    """Create a disabled tracker."""
-    return Tracker(enabled=False)
-
-
-# ====================== DebugStep Tests ======================
-
-
-def test_debug_step_initialization():
-    """Test that DebugStep can be initialized with default values."""
-    step = DebugStep()
-    assert step.step_idx == 0
-    assert step.x_t is None
-    assert step.v_t is None
-    assert step.x1_t is None
-    assert step.correction is None
-    assert step.err is None
-    assert step.weights is None
-    assert step.guidance_weight is None
-    assert step.time is None
-    assert step.inference_delay is None
-    assert step.execution_horizon is None
-    assert step.metadata == {}
-
-
-def test_debug_step_with_values(sample_tensors):
-    """Test DebugStep initialization with actual values."""
-    step = DebugStep(
-        step_idx=5,
-        x_t=sample_tensors["x_t"],
-        v_t=sample_tensors["v_t"],
-        x1_t=sample_tensors["x1_t"],
-        correction=sample_tensors["correction"],
-        err=sample_tensors["err"],
-        weights=sample_tensors["weights"],
-        guidance_weight=2.5,
-        time=0.8,
-        inference_delay=4,
-        execution_horizon=8,
-        metadata={"custom_key": "custom_value"},
-    )
-
-    assert step.step_idx == 5
-    assert torch.equal(step.x_t, sample_tensors["x_t"])
-    assert torch.equal(step.v_t, sample_tensors["v_t"])
-    assert torch.equal(step.x1_t, sample_tensors["x1_t"])
-    assert torch.equal(step.correction, sample_tensors["correction"])
-    assert torch.equal(step.err, sample_tensors["err"])
-    assert torch.equal(step.weights, sample_tensors["weights"])
-    assert step.guidance_weight == 2.5
-    assert step.time == 0.8
-    assert step.inference_delay == 4
-    assert step.execution_horizon == 8
-    assert step.metadata == {"custom_key": "custom_value"}
-
-
-def test_debug_step_to_dict_without_tensors(sample_tensors):
-    """Test converting DebugStep to dictionary without tensor values."""
-    step = DebugStep(
-        step_idx=3,
-        x_t=sample_tensors["x_t"],
-        v_t=sample_tensors["v_t"],
-        guidance_weight=torch.tensor(3.0),
-        time=torch.tensor(0.5),
-        inference_delay=2,
-        execution_horizon=10,
-    )
-
-    result = step.to_dict(include_tensors=False)
-
-    assert result["step_idx"] == 3
-    assert result["guidance_weight"] == 3.0
-    assert result["time"] == 0.5
-    assert result["inference_delay"] == 2
-    assert result["execution_horizon"] == 10
-
-    # Check tensor statistics are included
-    assert "x_t_stats" in result
-    assert "v_t_stats" in result
-    assert "x1_t_stats" not in result  # x1_t was None
-
-    # Verify statistics structure
-    assert "shape" in result["x_t_stats"]
-    assert "mean" in result["x_t_stats"]
-    assert "std" in result["x_t_stats"]
-    assert "min" in result["x_t_stats"]
-    assert "max" in result["x_t_stats"]
-
-    # Verify shape matches original tensor
-    assert result["x_t_stats"]["shape"] == tuple(sample_tensors["x_t"].shape)
-
-
-def test_debug_step_to_dict_with_tensors(sample_tensors):
-    """Test converting DebugStep to dictionary with tensor values."""
-    step = DebugStep(
-        step_idx=1,
-        x_t=sample_tensors["x_t"],
-        v_t=sample_tensors["v_t"],
-        guidance_weight=1.5,
-        time=0.9,
-    )
-
-    result = step.to_dict(include_tensors=True)
-
-    assert result["step_idx"] == 1
-    assert result["guidance_weight"] == 1.5
-    assert result["time"] == 0.9
-
-    # Check tensors are included (as CPU tensors)
-    assert "x_t" in result
-    assert "v_t" in result
-    assert isinstance(result["x_t"], torch.Tensor)
-    assert isinstance(result["v_t"], torch.Tensor)
-    assert result["x_t"].device.type == "cpu"
-    assert result["v_t"].device.type == "cpu"
-
-
-def test_debug_step_to_dict_with_none_guidance_weight():
-    """Test to_dict handles None guidance_weight correctly."""
-    step = DebugStep(step_idx=0, time=1.0, guidance_weight=None)
-    result = step.to_dict(include_tensors=False)
-    assert result["guidance_weight"] is None
-
-
-def test_tracker_initialization_enabled():
-    """Test tracker initialization when enabled."""
-    tracker = Tracker(enabled=True, maxlen=50)
-    assert tracker.enabled is True
-    assert tracker._steps == {}
-    assert tracker._maxlen == 50
-    assert tracker._step_counter == 0
-    assert len(tracker) == 0
-
-
-def test_tracker_reset_when_enabled(enabled_tracker, sample_tensors):
-    """Test reset clears all steps when tracker is enabled."""
-    # Add some steps
-    enabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
-    enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
-    assert len(enabled_tracker) == 2
-
-    # Reset
-    enabled_tracker.reset()
-    assert len(enabled_tracker) == 0
-    assert enabled_tracker._step_counter == 0
-    assert enabled_tracker._steps == {}
-
-
-def test_tracker_reset_when_disabled(disabled_tracker):
-    """Test reset on disabled tracker doesn't cause errors."""
-    disabled_tracker.reset()
-    assert len(disabled_tracker) == 0
-
-
-# ====================== Tracker.track() Tests ======================
-
-
-def test_track_creates_new_step(enabled_tracker, sample_tensors):
-    """Test that track creates a new step when time doesn't exist."""
-    enabled_tracker.track(
-        time=1.0,
-        x_t=sample_tensors["x_t"],
-        v_t=sample_tensors["v_t"],
-        guidance_weight=5.0,
-        inference_delay=4,
-        execution_horizon=8,
-    )
-
-    assert len(enabled_tracker) == 1
-    steps = enabled_tracker.get_all_steps()
-    assert len(steps) == 1
-    assert steps[0].step_idx == 0
-    assert steps[0].time == 1.0
-    assert torch.equal(steps[0].x_t, sample_tensors["x_t"])
-    assert torch.equal(steps[0].v_t, sample_tensors["v_t"])
-    assert steps[0].guidance_weight == 5.0
-    assert steps[0].inference_delay == 4
-    assert steps[0].execution_horizon == 8
-
-
-def test_track_updates_existing_step(enabled_tracker, sample_tensors):
-    """Test that track updates an existing step at the same time."""
-    # Create initial step
-    enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
-    assert len(enabled_tracker) == 1
-    steps = enabled_tracker.get_all_steps()
-    assert steps[0].v_t is None
-
-    # Update the same timestep with v_t
-    enabled_tracker.track(time=0.9, v_t=sample_tensors["v_t"])
-    assert len(enabled_tracker) == 1  # Still only one step
-    steps = enabled_tracker.get_all_steps()
-    assert torch.equal(steps[0].x_t, sample_tensors["x_t"])  # Original x_t preserved
-    assert torch.equal(steps[0].v_t, sample_tensors["v_t"])  # New v_t added
-
-
-def test_track_with_tensor_time(enabled_tracker, sample_tensors):
-    """Test track handles tensor time values correctly."""
-    time_tensor = torch.tensor(0.8)
-    enabled_tracker.track(time=time_tensor, x_t=sample_tensors["x_t"])
-
-    steps = enabled_tracker.get_all_steps()
-    assert len(steps) == 1
-    assert abs(steps[0].time - 0.8) < 1e-6  # Use approximate comparison for floating point
-
-
-def test_track_time_rounding(enabled_tracker, sample_tensors):
-    """Test that track rounds time to avoid floating point precision issues."""
-    # These times should be treated as the same after rounding to 6 decimals
-    enabled_tracker.track(time=0.9000001, x_t=sample_tensors["x_t"])
-    enabled_tracker.track(time=0.9000002, v_t=sample_tensors["v_t"])
-
-    # Should still be one step (times rounded to same value)
-    assert len(enabled_tracker) == 1
-    steps = enabled_tracker.get_all_steps()
-    assert torch.equal(steps[0].x_t, sample_tensors["x_t"])
-    assert torch.equal(steps[0].v_t, sample_tensors["v_t"])
-
-
-def test_track_does_nothing_when_disabled(disabled_tracker, sample_tensors):
-    """Test that track does nothing when tracker is disabled."""
-    disabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
-    assert len(disabled_tracker) == 0
-
-
-def test_track_with_metadata(enabled_tracker, sample_tensors):
-    """Test track stores custom metadata."""
-    enabled_tracker.track(time=0.7, x_t=sample_tensors["x_t"], custom_field="custom_value", count=42)
-
-    steps = enabled_tracker.get_all_steps()
-    assert steps[0].metadata["custom_field"] == "custom_value"
-    assert steps[0].metadata["count"] == 42
-
-
-def test_track_updates_metadata(enabled_tracker):
-    """Test that track updates metadata for existing steps."""
-    enabled_tracker.track(time=0.6, meta1="value1")
-    enabled_tracker.track(time=0.6, meta2="value2")
-
-    steps = enabled_tracker.get_all_steps()
-    assert steps[0].metadata["meta1"] == "value1"
-    assert steps[0].metadata["meta2"] == "value2"
-
-
-def test_track_clones_tensors(enabled_tracker, sample_tensors):
-    """Test that track clones tensors instead of storing references."""
-    x_t_original = sample_tensors["x_t"].clone()
-    enabled_tracker.track(time=0.5, x_t=sample_tensors["x_t"])
-
-    # Modify original tensor
-    sample_tensors["x_t"].fill_(999.0)
-
-    # Tracked tensor should not be affected
-    steps = enabled_tracker.get_all_steps()
-    assert not torch.equal(steps[0].x_t, sample_tensors["x_t"])
-    assert torch.equal(steps[0].x_t, x_t_original)
-
-
-def test_track_with_none_values(enabled_tracker):
-    """Test track handles None values correctly."""
-    enabled_tracker.track(
-        time=0.4,
-        x_t=None,
-        v_t=None,
-        guidance_weight=None,
-        inference_delay=None,
-    )
-
-    steps = enabled_tracker.get_all_steps()
-    assert len(steps) == 1
-    assert steps[0].x_t is None
-    assert steps[0].v_t is None
-    assert steps[0].guidance_weight is None
-    assert steps[0].inference_delay is None
-
-
-def test_track_updates_only_non_none_fields(enabled_tracker, sample_tensors):
-    """Test that update preserves existing values when None is passed."""
-    # Create step with x_t
-    enabled_tracker.track(time=0.3, x_t=sample_tensors["x_t"], guidance_weight=2.0)
-
-    # Update with v_t only (pass None for other fields)
-    enabled_tracker.track(time=0.3, v_t=sample_tensors["v_t"], x_t=None, guidance_weight=None)
-
-    # Original values should be preserved
-    steps = enabled_tracker.get_all_steps()
-    assert torch.equal(steps[0].x_t, sample_tensors["x_t"])  # Still has x_t
-    assert torch.equal(steps[0].v_t, sample_tensors["v_t"])  # Now has v_t
-    assert steps[0].guidance_weight == 2.0  # Still has guidance_weight
-
-
-# ====================== Tracker.maxlen Tests ======================
-
-
-def test_tracker_enforces_maxlen():
-    """Test that tracker enforces maxlen limit."""
-    tracker = Tracker(enabled=True, maxlen=3)
-
-    # Add 5 steps
-    for i in range(5):
-        time = 1.0 - i * 0.1  # 1.0, 0.9, 0.8, 0.7, 0.6
-        tracker.track(time=time, x_t=torch.randn(1, 10, 6))
-
-    # Should only keep the last 3
-    assert len(tracker) == 3
-
-    # Verify oldest steps were removed (should have 0.6, 0.7, 0.8)
-    steps = tracker.get_all_steps()
-    times = sorted([step.time for step in steps])
-    assert times == [0.6, 0.7, 0.8]
-
-
-def test_tracker_step_idx_increments_despite_maxlen():
-    """Test that step_idx continues incrementing even when maxlen is enforced."""
-    tracker = Tracker(enabled=True, maxlen=2)
-
-    # Add 4 steps
-    for i in range(4):
-        time = 1.0 - i * 0.1
-        tracker.track(time=time, x_t=torch.randn(1, 10, 6))
-
-    # Should have 2 steps with step_idx 2 and 3 (oldest removed)
-    steps = sorted(tracker.get_all_steps(), key=lambda s: s.step_idx)
-    assert len(steps) == 2
-    assert steps[0].step_idx == 2
-    assert steps[1].step_idx == 3
-
-
-def test_tracker_without_maxlen_keeps_all():
-    """Test that tracker without maxlen keeps all steps."""
-    tracker = Tracker(enabled=True, maxlen=None)
-
-    # Add 100 steps
-    for i in range(100):
-        time = 1.0 - i * 0.01
-        tracker.track(time=time, x_t=torch.randn(1, 10, 6))
-
-    assert len(tracker) == 100
-
-
-def test_get_all_steps_returns_empty_when_disabled(disabled_tracker):
-    """Test get_all_steps returns empty list when disabled."""
-    steps = disabled_tracker.get_all_steps()
-    assert steps == []
-    assert isinstance(steps, list)
-
-
-def test_get_all_steps_returns_empty_when_no_steps(enabled_tracker):
-    """Test get_all_steps returns empty list when no steps tracked."""
-    steps = enabled_tracker.get_all_steps()
-    assert steps == []
-
-
-def test_get_all_steps_returns_all_tracked_steps(enabled_tracker, sample_tensors):
-    """Test get_all_steps returns all tracked steps."""
-    # Track 5 steps
-    for i in range(5):
-        time = 1.0 - i * 0.1
-        enabled_tracker.track(time=time, x_t=sample_tensors["x_t"])
-
-    steps = enabled_tracker.get_all_steps()
-    assert len(steps) == 5
-
-    # Verify all are DebugStep instances
-    for step in steps:
-        assert isinstance(step, DebugStep)
-
-
-def test_get_all_steps_preserves_insertion_order(enabled_tracker):
-    """Test that get_all_steps preserves insertion order (Python 3.7+)."""
-    times = [0.9, 0.8, 0.7, 0.6, 0.5]
-    for time in times:
-        enabled_tracker.track(time=time, x_t=torch.randn(1, 10, 6))
-
-    steps = enabled_tracker.get_all_steps()
-    retrieved_times = [step.time for step in steps]
-
-    # Should be in insertion order
-    assert retrieved_times == times
-
-
-# ====================== Tracker.__len__() Tests ======================
-
-
-def test_len_returns_zero_when_disabled(disabled_tracker):
-    """Test __len__ returns 0 when tracker is disabled."""
-    assert len(disabled_tracker) == 0
-
-
-def test_len_returns_zero_when_empty(enabled_tracker):
-    """Test __len__ returns 0 when no steps are tracked."""
-    assert len(enabled_tracker) == 0
-
-
-def test_len_returns_correct_count(enabled_tracker, sample_tensors):
-    """Test __len__ returns correct number of tracked steps."""
-    assert len(enabled_tracker) == 0
-
-    enabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
-    assert len(enabled_tracker) == 1
-
-    enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
-    assert len(enabled_tracker) == 2
-
-    enabled_tracker.track(time=0.8, x_t=sample_tensors["x_t"])
-    assert len(enabled_tracker) == 3
-
-
-def test_len_after_reset(enabled_tracker, sample_tensors):
-    """Test __len__ returns 0 after reset."""
-    enabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
-    enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
-    assert len(enabled_tracker) == 2
-
-    enabled_tracker.reset()
-    assert len(enabled_tracker) == 0
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_tracker_handles_gpu_tensors():
-    """Test tracker correctly handles GPU tensors."""
-    tracker = Tracker(enabled=True, maxlen=10)
-    x_t_gpu = torch.randn(1, 50, 6, device="cuda")
-
-    tracker.track(time=1.0, x_t=x_t_gpu)
-
-    steps = tracker.get_all_steps()
-    # Tracker should clone and detach tensors
-    assert steps[0].x_t.device.type == "cuda"
-
-
-def test_tracker_with_varying_tensor_shapes(enabled_tracker):
-    """Test tracker handles varying tensor shapes across steps."""
-    enabled_tracker.track(time=1.0, x_t=torch.randn(1, 50, 6))
-    enabled_tracker.track(time=0.9, x_t=torch.randn(1, 25, 6))
-    enabled_tracker.track(time=0.8, x_t=torch.randn(2, 50, 8))
-
-    steps = enabled_tracker.get_all_steps()
-    assert len(steps) == 3
-    assert steps[0].x_t.shape == (1, 50, 6)
-    assert steps[1].x_t.shape == (1, 25, 6)
-    assert steps[2].x_t.shape == (2, 50, 8)

tests/policies/rtc/test_latency_tracker.py

@@ -1,322 +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.
-
-"""Tests for RTC LatencyTracker module."""
-
-import pytest
-
-from lerobot.policies.rtc.latency_tracker import LatencyTracker
-
-# ====================== Fixtures ======================
-
-
-@pytest.fixture
-def tracker():
-    """Create a LatencyTracker with default maxlen."""
-    return LatencyTracker(maxlen=100)
-
-
-@pytest.fixture
-def small_tracker():
-    """Create a LatencyTracker with small maxlen for overflow testing."""
-    return LatencyTracker(maxlen=5)
-
-
-# ====================== Initialization Tests ======================
-
-
-def test_latency_tracker_initialization():
-    """Test LatencyTracker initializes correctly."""
-    tracker = LatencyTracker(maxlen=50)
-    assert len(tracker) == 0
-    assert tracker.max_latency == 0.0
-    assert tracker.max() == 0.0
-
-
-def test_latency_tracker_default_maxlen():
-    """Test LatencyTracker uses default maxlen."""
-    tracker = LatencyTracker()
-    # Should accept default maxlen=100
-    assert len(tracker) == 0
-
-
-# ====================== add() Tests ======================
-
-
-def test_add_single_latency(tracker):
-    """Test adding a single latency value."""
-    tracker.add(0.5)
-    assert len(tracker) == 1
-    assert tracker.max() == 0.5
-
-
-def test_add_multiple_latencies(tracker):
-    """Test adding multiple latency values."""
-    latencies = [0.1, 0.5, 0.3, 0.8, 0.2]
-    for lat in latencies:
-        tracker.add(lat)
-
-    assert len(tracker) == 5
-    assert tracker.max() == 0.8
-
-
-def test_add_negative_latency_ignored(tracker):
-    """Test that negative latencies are ignored."""
-    tracker.add(0.5)
-    tracker.add(-0.1)
-    tracker.add(0.3)
-
-    # Should only have 2 valid latencies
-    assert len(tracker) == 2
-    assert tracker.max() == 0.5
-
-
-def test_add_zero_latency(tracker):
-    """Test adding zero latency."""
-    tracker.add(0.0)
-    assert len(tracker) == 1
-    assert tracker.max() == 0.0
-
-
-def test_add_converts_to_float(tracker):
-    """Test add() converts input to float."""
-    tracker.add(5)  # Integer
-    tracker.add("3.5")  # String
-
-    assert len(tracker) == 2
-    assert tracker.max() == 5.0
-
-
-def test_add_updates_max_latency(tracker):
-    """Test that max_latency is updated correctly."""
-    tracker.add(0.5)
-    assert tracker.max_latency == 0.5
-
-    tracker.add(0.3)
-    assert tracker.max_latency == 0.5  # Should not decrease
-
-    tracker.add(0.9)
-    assert tracker.max_latency == 0.9  # Should increase
-
-
-# ====================== reset() Tests ======================
-
-
-def test_reset_clears_values(tracker):
-    """Test reset() clears all values."""
-    tracker.add(0.5)
-    tracker.add(0.8)
-    tracker.add(0.3)
-    assert len(tracker) == 3
-
-    tracker.reset()
-    assert len(tracker) == 0
-    assert tracker.max_latency == 0.0
-
-
-def test_reset_clears_max_latency(tracker):
-    """Test reset() resets max_latency."""
-    tracker.add(1.5)
-    assert tracker.max_latency == 1.5
-
-    tracker.reset()
-    assert tracker.max_latency == 0.0
-
-
-def test_reset_allows_new_values(tracker):
-    """Test that tracker works correctly after reset."""
-    tracker.add(0.5)
-    tracker.reset()
-
-    tracker.add(0.3)
-    assert len(tracker) == 1
-    assert tracker.max() == 0.3
-
-
-# ====================== max() Tests ======================
-
-
-def test_max_returns_zero_when_empty(tracker):
-    """Test max() returns 0.0 when tracker is empty."""
-    assert tracker.max() == 0.0
-
-
-def test_max_returns_maximum_value(tracker):
-    """Test max() returns the maximum latency."""
-    latencies = [0.2, 0.8, 0.3, 0.5, 0.1]
-    for lat in latencies:
-        tracker.add(lat)
-
-    assert tracker.max() == 0.8
-
-
-def test_max_persists_after_sliding_window(small_tracker):
-    """Test max() persists even after values slide out of window."""
-    # Add values that will exceed maxlen=5
-    small_tracker.add(0.1)
-    small_tracker.add(0.9)  # This is max
-    small_tracker.add(0.2)
-    small_tracker.add(0.3)
-    small_tracker.add(0.4)
-    small_tracker.add(0.5)  # This pushes out 0.1
-
-    # Max should still be 0.9 even though only last 5 values kept
-    assert small_tracker.max() == 0.9
-
-
-def test_max_after_reset(tracker):
-    """Test max() returns 0.0 after reset."""
-    tracker.add(1.5)
-    tracker.reset()
-    assert tracker.max() == 0.0
-
-
-# ====================== p95() Tests ======================
-
-
-def test_p95_returns_zero_when_empty(tracker):
-    """Test p95() returns 0.0 when tracker is empty."""
-    assert tracker.p95() == 0.0
-
-
-def test_p95_returns_95th_percentile(tracker):
-    """Test p95() returns the 95th percentile."""
-    # Add 100 values
-    for i in range(100):
-        tracker.add(i / 100.0)
-
-    p95 = tracker.p95()
-    assert 0.93 <= p95 <= 0.96
-
-
-def test_p95_equals_percentile_95(tracker):
-    """Test p95() equals percentile(0.95)."""
-    for i in range(50):
-        tracker.add(i / 50.0)
-
-    assert tracker.p95() == tracker.percentile(0.95)
-
-
-# ====================== Edge Cases Tests ======================
-
-
-def test_single_value(tracker):
-    """Test tracker behavior with single value."""
-    tracker.add(0.75)
-
-    assert len(tracker) == 1
-    assert tracker.max() == 0.75
-    assert tracker.percentile(0.0) == 0.75
-    assert tracker.percentile(0.5) == 0.75
-    assert tracker.percentile(1.0) == 0.75
-
-
-def test_all_same_values(tracker):
-    """Test tracker with all identical values."""
-    for _ in range(10):
-        tracker.add(0.5)
-
-    assert len(tracker) == 10
-    assert tracker.max() == 0.5
-    assert tracker.percentile(0.0) == 0.5
-    assert tracker.percentile(0.5) == 0.5
-    assert tracker.percentile(1.0) == 0.5
-
-
-def test_very_small_values(tracker):
-    """Test tracker with very small float values."""
-    tracker.add(1e-10)
-    tracker.add(2e-10)
-    tracker.add(3e-10)
-
-    assert len(tracker) == 3
-    assert tracker.max() == pytest.approx(3e-10)
-
-
-def test_very_large_values(tracker):
-    """Test tracker with very large float values."""
-    tracker.add(1e10)
-    tracker.add(2e10)
-    tracker.add(3e10)
-
-    assert len(tracker) == 3
-    assert tracker.max() == pytest.approx(3e10)
-
-
-# ====================== Integration Tests ======================
-
-
-def test_typical_usage_pattern(tracker):
-    """Test a typical usage pattern of the tracker."""
-    # Simulate adding latencies over time
-    latencies = [0.05, 0.08, 0.12, 0.07, 0.15, 0.09, 0.11, 0.06, 0.14, 0.10]
-
-    for lat in latencies:
-        tracker.add(lat)
-
-    # Check statistics
-    assert len(tracker) == 10
-    assert tracker.max() == 0.15
-
-    # p95 should be close to max since we have only 10 values
-    p95 = tracker.p95()
-    assert p95 >= tracker.percentile(0.5)  # p95 should be >= median
-    assert p95 <= tracker.max()  # p95 should be <= max
-
-
-def test_reset_and_reuse(tracker):
-    """Test resetting and reusing tracker."""
-    # First batch
-    tracker.add(1.0)
-    tracker.add(2.0)
-    assert tracker.max() == 2.0
-
-    # Reset
-    tracker.reset()
-
-    # Second batch
-    tracker.add(0.5)
-    tracker.add(0.8)
-    assert len(tracker) == 2
-    assert tracker.max() == 0.8
-    assert tracker.percentile(0.5) <= 0.8
-
-
-# ====================== Type Conversion Tests ======================
-
-
-def test_add_with_integer(tracker):
-    """Test adding integer values."""
-    tracker.add(5)
-    assert len(tracker) == 1
-    assert tracker.max() == 5.0
-
-
-def test_add_with_string_number(tracker):
-    """Test adding string representation of number."""
-    tracker.add("3.14")
-    assert len(tracker) == 1
-    assert tracker.max() == pytest.approx(3.14)
-
-
-def test_percentile_converts_q_to_float(tracker):
-    """Test percentile converts q parameter to float."""
-    tracker.add(0.5)
-    tracker.add(0.8)
-
-    # Pass integer q
-    result = tracker.percentile(1)
-    assert result == 0.8

tests/policies/rtc/test_modeling_rtc.py

@@ -1,773 +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.
-
-"""Tests for RTC modeling module (RTCProcessor)."""
-
-import pytest
-import torch
-
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
-
-# ====================== Fixtures ======================
-
-
-@pytest.fixture
-def rtc_config_debug_enabled():
-    """Create RTC config with debug enabled."""
-    return RTCConfig(
-        enabled=True,
-        prefix_attention_schedule=RTCAttentionSchedule.LINEAR,
-        max_guidance_weight=10.0,
-        execution_horizon=10,
-        debug=True,
-        debug_maxlen=100,
-    )
-
-
-@pytest.fixture
-def rtc_config_debug_disabled():
-    """Create RTC config with debug disabled."""
-    return RTCConfig(
-        enabled=True,
-        prefix_attention_schedule=RTCAttentionSchedule.LINEAR,
-        max_guidance_weight=10.0,
-        execution_horizon=10,
-        debug=False,
-    )
-
-
-@pytest.fixture
-def rtc_processor_debug_enabled(rtc_config_debug_enabled):
-    """Create RTCProcessor with debug enabled."""
-    return RTCProcessor(rtc_config_debug_enabled)
-
-
-@pytest.fixture
-def rtc_processor_debug_disabled(rtc_config_debug_disabled):
-    """Create RTCProcessor with debug disabled."""
-    return RTCProcessor(rtc_config_debug_disabled)
-
-
-@pytest.fixture
-def sample_x_t():
-    """Create sample x_t tensor (batch, time, action_dim)."""
-    return torch.randn(1, 50, 6)
-
-
-@pytest.fixture
-def sample_prev_chunk():
-    """Create sample previous chunk tensor."""
-    return torch.randn(1, 50, 6)
-
-
-# ====================== Initialization Tests ======================
-
-
-def test_rtc_processor_initialization_with_debug(rtc_config_debug_enabled):
-    """Test RTCProcessor initializes with debug tracker."""
-    processor = RTCProcessor(rtc_config_debug_enabled)
-    assert processor.rtc_config == rtc_config_debug_enabled
-    assert processor.tracker is not None
-    assert processor.tracker.enabled is True
-
-
-def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled):
-    """Test RTCProcessor initializes without debug tracker."""
-    processor = RTCProcessor(rtc_config_debug_disabled)
-    assert processor.rtc_config == rtc_config_debug_disabled
-    assert processor.tracker is None
-
-
-# ====================== Tracker Proxy Methods Tests ======================
-
-
-def test_track_when_tracker_enabled(rtc_processor_debug_enabled, sample_x_t):
-    """Test track() forwards to tracker when enabled."""
-    rtc_processor_debug_enabled.track(
-        time=torch.tensor(0.5),
-        x_t=sample_x_t,
-        v_t=sample_x_t,
-        guidance_weight=2.0,
-    )
-
-    # Should have tracked one step
-    steps = rtc_processor_debug_enabled.get_all_debug_steps()
-    assert len(steps) == 1
-    assert steps[0].time == 0.5
-
-
-def test_track_when_tracker_disabled(rtc_processor_debug_disabled, sample_x_t):
-    """Test track() does nothing when tracker disabled."""
-    # Should not raise error
-    rtc_processor_debug_disabled.track(
-        time=torch.tensor(0.5),
-        x_t=sample_x_t,
-        v_t=sample_x_t,
-    )
-
-    # Should return empty list
-    steps = rtc_processor_debug_disabled.get_all_debug_steps()
-    assert len(steps) == 0
-
-
-def test_get_all_debug_steps_when_enabled(rtc_processor_debug_enabled, sample_x_t):
-    """Test get_all_debug_steps() returns tracked steps."""
-    rtc_processor_debug_enabled.track(time=torch.tensor(0.5), x_t=sample_x_t)
-    rtc_processor_debug_enabled.track(time=torch.tensor(0.4), x_t=sample_x_t)
-
-    steps = rtc_processor_debug_enabled.get_all_debug_steps()
-    assert len(steps) == 2
-
-
-def test_get_all_debug_steps_when_disabled(rtc_processor_debug_disabled):
-    """Test get_all_debug_steps() returns empty list when disabled."""
-    steps = rtc_processor_debug_disabled.get_all_debug_steps()
-    assert steps == []
-    assert isinstance(steps, list)
-
-
-def test_is_debug_enabled_when_tracker_exists(rtc_processor_debug_enabled):
-    """Test is_debug_enabled() returns True when tracker enabled."""
-    assert rtc_processor_debug_enabled.is_debug_enabled() is True
-
-
-def test_is_debug_enabled_when_tracker_disabled(rtc_processor_debug_disabled):
-    """Test is_debug_enabled() returns False when tracker disabled."""
-    assert rtc_processor_debug_disabled.is_debug_enabled() is False
-
-
-def test_reset_tracker_when_enabled(rtc_processor_debug_enabled, sample_x_t):
-    """Test reset_tracker() clears tracked steps."""
-    rtc_processor_debug_enabled.track(time=torch.tensor(0.5), x_t=sample_x_t)
-    rtc_processor_debug_enabled.track(time=torch.tensor(0.4), x_t=sample_x_t)
-    assert len(rtc_processor_debug_enabled.get_all_debug_steps()) == 2
-
-    rtc_processor_debug_enabled.reset_tracker()
-    assert len(rtc_processor_debug_enabled.get_all_debug_steps()) == 0
-
-
-def test_reset_tracker_when_disabled(rtc_processor_debug_disabled):
-    """Test reset_tracker() doesn't error when tracker disabled."""
-    rtc_processor_debug_disabled.reset_tracker()  # Should not raise
-
-
-# ====================== get_prefix_weights Tests ======================
-
-
-def test_get_prefix_weights_zeros_schedule():
-    """Test get_prefix_weights with ZEROS schedule."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.ZEROS)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=5, end=10, total=20)
-
-    # First 5 should be 1.0, rest should be 0.0
-    assert weights.shape == (20,)
-    assert torch.all(weights[:5] == 1.0)
-    assert torch.all(weights[5:] == 0.0)
-
-
-def test_get_prefix_weights_ones_schedule():
-    """Test get_prefix_weights with ONES schedule."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.ONES)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=5, end=15, total=20)
-
-    # First 15 should be 1.0, rest should be 0.0
-    assert weights.shape == (20,)
-    assert torch.all(weights[:15] == 1.0)
-    assert torch.all(weights[15:] == 0.0)
-
-
-def test_get_prefix_weights_linear_schedule():
-    """Test get_prefix_weights with LINEAR schedule."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=5, end=14, total=25)
-
-    # Should have shape (20,)
-    assert weights.shape == (25,)
-
-    # First 5 should be 1.0 (leading ones)
-    assert torch.all(weights[:5] == 1.0)
-
-    # Middle section (5:15) should be linearly decreasing from 1 to 0
-    middle_weights = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
-    assert torch.allclose(weights[5:14], middle_weights)
-
-    # Last 5 should be 0.0 (trailing zeros)
-    assert torch.all(weights[14:] == 0.0)
-
-
-def test_get_prefix_weights_exp_schedule():
-    """Test get_prefix_weights with EXP schedule."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.EXP)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=5, end=14, total=25)
-
-    # Should have shape (20,)
-    assert weights.shape == (25,)
-
-    # First 5 should be 1.0 (leading ones)
-    assert torch.all(weights[:5] == 1.0)
-
-    # Middle section should be exponentially weighted
-    middle_weights = torch.tensor([0.7645, 0.5706, 0.4130, 0.2871, 0.1888, 0.1145, 0.0611, 0.0258, 0.0061])
-    assert torch.allclose(weights[5:14], middle_weights, atol=1e-4)
-
-    # Last 5 should be 0.0 (trailing zeros)
-    assert torch.all(weights[14:] == 0.0)
-
-
-def test_get_prefix_weights_with_start_equals_end():
-    """Test get_prefix_weights when start equals end."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=10, end=10, total=20)
-
-    # Should have ones up to start, then zeros
-    assert torch.all(weights[:10] == 1.0)
-    assert torch.all(weights[10:] == 0.0)
-
-
-def test_get_prefix_weights_with_start_greater_than_end():
-    """Test get_prefix_weights when start > end (gets clamped)."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
-    processor = RTCProcessor(config)
-
-    # start > end should use min(start, end) = end
-    weights = processor.get_prefix_weights(start=15, end=10, total=20)
-
-    # Should have ones up to end (10), then zeros
-    assert torch.all(weights[:10] == 1.0)
-    assert torch.all(weights[10:] == 0.0)
-
-
-# ====================== Helper Method Tests ======================
-
-
-def test_linweights_with_end_equals_start():
-    """Test _linweights when end equals start."""
-    config = RTCConfig()
-    processor = RTCProcessor(config)
-
-    weights = processor._linweights(start=10, end=10, total=20)
-
-    # Should return empty tensor
-    assert len(weights) == 0
-
-
-def test_linweights_with_end_less_than_start():
-    """Test _linweights when end < start."""
-    config = RTCConfig()
-    processor = RTCProcessor(config)
-
-    weights = processor._linweights(start=15, end=10, total=20)
-
-    # Should return empty tensor
-    assert len(weights) == 0
-
-
-def test_add_trailing_zeros_normal():
-    """Test _add_trailing_zeros adds zeros correctly."""
-    config = RTCConfig()
-    processor = RTCProcessor(config)
-
-    weights = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2])
-    result = processor._add_trailing_zeros(weights, total=10, end=5)
-
-    # Should add 5 zeros (total - end = 10 - 5 = 5)
-    assert len(result) == 10
-    assert torch.all(result[:5] == weights)
-    assert torch.all(result[5:] == 0.0)
-
-
-def test_add_trailing_zeros_no_zeros_needed():
-    """Test _add_trailing_zeros when no zeros needed."""
-    config = RTCConfig()
-    processor = RTCProcessor(config)
-
-    weights = torch.tensor([1.0, 0.8, 0.6])
-    result = processor._add_trailing_zeros(weights, total=3, end=5)
-
-    # zeros_len = 3 - 5 = -2 <= 0, so no zeros added
-    assert torch.equal(result, weights)
-
-
-def test_add_leading_ones_normal():
-    """Test _add_leading_ones adds ones correctly."""
-    config = RTCConfig()
-    processor = RTCProcessor(config)
-
-    weights = torch.tensor([0.8, 0.6, 0.4, 0.2, 0.0])
-    result = processor._add_leading_ones(weights, start=3, total=10)
-
-    # Should add 3 ones at the start
-    assert len(result) == 8
-    assert torch.all(result[:3] == 1.0)
-    assert torch.all(result[3:] == weights)
-
-
-def test_add_leading_ones_no_ones_needed():
-    """Test _add_leading_ones when no ones needed."""
-    config = RTCConfig()
-    processor = RTCProcessor(config)
-
-    weights = torch.tensor([0.8, 0.6, 0.4])
-    result = processor._add_leading_ones(weights, start=0, total=10)
-
-    # ones_len = 0, so no ones added
-    assert torch.equal(result, weights)
-
-
-def test_get_prefix_weights_with_start_equals_total():
-    """Test get_prefix_weights when start equals total."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=10, end=10, total=20)
-
-    # Should have ones up to start, then zeros
-    assert len(weights) == 20
-    assert torch.all(weights[:10] == 1.0)
-    assert torch.all(weights[10:] == 0.0)
-
-
-def test_get_prefix_weights_with_total_less_than_start():
-    """Test get_prefix_weights when total less than start."""
-    config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
-    processor = RTCProcessor(config)
-
-    weights = processor.get_prefix_weights(start=10, end=10, total=5)
-
-    # Should have ones up to start, then zeros
-    assert len(weights) == 5
-    assert torch.all(weights == 1.0)
-
-
-# ====================== denoise_step Tests ======================
-
-
-def test_denoise_step_without_prev_chunk(rtc_processor_debug_disabled):
-    """Test denoise_step without previous chunk (no guidance)."""
-    x_t = torch.randn(1, 50, 6)
-
-    # Mock denoiser that returns fixed velocity
-    def mock_denoiser(x):
-        return torch.ones_like(x) * 0.5
-
-    result = rtc_processor_debug_disabled.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=None,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Should return v_t unchanged (no guidance)
-    expected = mock_denoiser(x_t)
-    assert torch.allclose(result, expected)
-
-
-def test_denoise_step_with_prev_chunk(rtc_processor_debug_disabled):
-    """Test denoise_step with previous chunk applies guidance."""
-    x_t = torch.ones(1, 20, 1)
-    prev_chunk = torch.full((1, 20, 1), 0.1)
-
-    def mock_denoiser(x):
-        return x * 0.5
-
-    result = rtc_processor_debug_disabled.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    expected_result = torch.tensor(
-        [
-            [
-                [1.8000],
-                [1.8000],
-                [1.8000],
-                [1.8000],
-                [1.8000],
-                [1.5833],
-                [1.3667],
-                [1.1500],
-                [0.9333],
-                [0.7167],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-            ]
-        ]
-    )
-
-    assert torch.allclose(result, expected_result, atol=1e-4)
-
-
-def test_denoise_step_adds_batch_dimension():
-    """Test denoise_step handles 2D input by adding batch dimension."""
-    config = RTCConfig(execution_horizon=10, max_guidance_weight=5.0)
-    processor = RTCProcessor(config)
-
-    # 2D input (no batch dimension)
-    x_t = torch.randn(10, 6)
-    prev_chunk = torch.randn(5, 6)
-
-    def mock_denoiser(x):
-        return x * 0.5
-
-    result = processor.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Output should be 2D (batch dimension removed)
-    assert result.ndim == 2
-    assert result.shape == (10, 6)
-
-
-def test_denoise_step_uses_custom_execution_horizon():
-    """Test denoise_step uses custom execution_horizon parameter."""
-    config = RTCConfig(execution_horizon=10)
-    processor = RTCProcessor(config)
-
-    x_t = torch.ones(1, 20, 1)
-    prev_chunk = torch.full((1, 15, 1), 0.1)
-
-    def mock_denoiser(x):
-        return x * 0.5
-
-    result = processor.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-        execution_horizon=15,
-    )
-
-    expected_result = torch.tensor(
-        [
-            [
-                [1.8000],
-                [1.8000],
-                [1.8000],
-                [1.8000],
-                [1.8000],
-                [1.6818],
-                [1.5636],
-                [1.4455],
-                [1.3273],
-                [1.2091],
-                [1.0909],
-                [0.9727],
-                [0.8545],
-                [0.7364],
-                [0.6182],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-            ]
-        ]
-    )
-
-    assert torch.allclose(result, expected_result, atol=1e-4)
-
-
-def test_denoise_step_guidance_weight_at_time_zero():
-    """Test denoise_step handles time=0 (tau=1) without NaN/Inf."""
-    config = RTCConfig(max_guidance_weight=10.0)
-    processor = RTCProcessor(config)
-
-    x_t = torch.ones(1, 20, 1)
-    prev_chunk = torch.full((1, 20, 1), 0.1)
-
-    def mock_denoiser(x):
-        return x * 0.5
-
-    result = processor.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.0),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    expected_result = torch.tensor(
-        [
-            [
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-                [0.5000],
-            ]
-        ]
-    )
-
-    assert torch.allclose(result, expected_result, atol=1e-4)
-
-
-def test_denoise_step_with_real_denoise_step_partial():
-    """Test denoise_step with a real denoiser."""
-    config = RTCConfig(max_guidance_weight=10.0)
-    processor = RTCProcessor(config)
-
-    batch_size = 10
-    action_dim = 6
-    chunk_size = 20
-
-    x_t = torch.ones(batch_size, chunk_size, action_dim)
-    prev_chunk = torch.full((batch_size, chunk_size, action_dim), 0.1)
-
-    velocity_function = torch.nn.Sequential(
-        torch.nn.Linear(action_dim, 1000),
-        torch.nn.ReLU(),
-        torch.nn.Linear(1000, 256),
-        torch.nn.ReLU(),
-        torch.nn.Linear(256, action_dim),
-    )
-
-    def mock_denoiser(x):
-        return velocity_function(x)
-
-    result = processor.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    assert result.shape == (batch_size, chunk_size, action_dim)
-
-
-def test_denoise_step_guidance_weight_at_time_one():
-    """Test denoise_step handles time=1 (tau=0) with max_guidance_weight clamping."""
-    config = RTCConfig(max_guidance_weight=10.0)
-    processor = RTCProcessor(config)
-
-    x_t = torch.randn(1, 50, 6)
-    prev_chunk = torch.randn(1, 50, 6)
-
-    def mock_denoiser(x):
-        return torch.ones_like(x) * 0.5
-
-    # Time = 1 => tau = 0, c = (1-tau)/tau = 1/0 = inf (clamped to max_guidance_weight)
-    result = processor.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(1.0),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Should clamp to max_guidance_weight (no Inf)
-    assert not torch.any(torch.isinf(result))
-
-
-def test_denoise_step_tracks_debug_info(rtc_processor_debug_enabled):
-    """Test denoise_step tracks debug information when enabled."""
-    x_t = torch.randn(1, 50, 6)
-    prev_chunk = torch.randn(1, 50, 6)
-
-    def mock_denoiser(x):
-        return torch.ones_like(x) * 0.5
-
-    rtc_processor_debug_enabled.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Should have tracked one step
-    steps = rtc_processor_debug_enabled.get_all_debug_steps()
-    assert len(steps) == 1
-
-    # Check tracked values
-    step = steps[0]
-    assert step.time == 0.5
-    assert step.x1_t is not None
-    assert step.correction is not None
-    assert step.err is not None
-    assert step.weights is not None
-    assert step.guidance_weight is not None
-    assert step.inference_delay == 5
-
-
-def test_denoise_step_doesnt_track_without_debug(rtc_processor_debug_disabled):
-    """Test denoise_step doesn't track when debug disabled."""
-    x_t = torch.randn(1, 50, 6)
-    prev_chunk = torch.randn(1, 50, 6)
-
-    def mock_denoiser(x):
-        return torch.ones_like(x) * 0.5
-
-    rtc_processor_debug_disabled.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Should not track
-    steps = rtc_processor_debug_disabled.get_all_debug_steps()
-    assert len(steps) == 0
-
-
-# ====================== Integration Tests ======================
-
-
-def test_denoise_step_full_workflow():
-    """Test complete denoise_step workflow."""
-    config = RTCConfig(
-        enabled=True,
-        prefix_attention_schedule=RTCAttentionSchedule.LINEAR,
-        max_guidance_weight=5.0,
-        execution_horizon=10,
-        debug=True,
-    )
-    processor = RTCProcessor(config)
-
-    # Simulate two denoising steps
-    x_t1 = torch.randn(1, 50, 6)
-    x_t2 = torch.randn(1, 50, 6)
-
-    def mock_denoiser(x):
-        return torch.randn_like(x) * 0.1
-
-    # First step - no guidance
-    result1 = processor.denoise_step(
-        x_t=x_t1,
-        prev_chunk_left_over=None,
-        inference_delay=5,
-        time=torch.tensor(0.8),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Second step - with guidance
-    result2 = processor.denoise_step(
-        x_t=x_t2,
-        prev_chunk_left_over=result1,
-        inference_delay=5,
-        time=torch.tensor(0.6),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Both should complete successfully
-    assert result1.shape == (1, 50, 6)
-    assert result2.shape == (1, 50, 6)
-
-    # Should have tracked one step (second one, first had no prev_chunk)
-    steps = processor.get_all_debug_steps()
-    assert len(steps) == 1
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_denoise_step_with_cuda_tensors():
-    """Test denoise_step works with CUDA tensors."""
-    config = RTCConfig(execution_horizon=10, max_guidance_weight=5.0)
-    processor = RTCProcessor(config)
-
-    x_t = torch.randn(1, 50, 6, device="cuda")
-    prev_chunk = torch.randn(1, 50, 6, device="cuda")
-
-    def mock_denoiser(x):
-        return torch.ones_like(x) * 0.5
-
-    result = processor.denoise_step(
-        x_t=x_t,
-        prev_chunk_left_over=prev_chunk,
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=mock_denoiser,
-    )
-
-    # Result should be on CUDA
-    assert result.device.type == "cuda"
-    assert result.shape == x_t.shape
-
-
-def test_denoise_step_deterministic_with_same_inputs():
-    """Test denoise_step produces same output with same inputs."""
-    config = RTCConfig(execution_horizon=10, max_guidance_weight=5.0)
-    processor = RTCProcessor(config)
-
-    torch.manual_seed(42)
-    x_t = torch.randn(1, 50, 6)
-    prev_chunk = torch.randn(1, 50, 6)
-
-    def deterministic_denoiser(x):
-        return torch.ones_like(x) * 0.5
-
-    result1 = processor.denoise_step(
-        x_t=x_t.clone(),
-        prev_chunk_left_over=prev_chunk.clone(),
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=deterministic_denoiser,
-    )
-
-    result2 = processor.denoise_step(
-        x_t=x_t.clone(),
-        prev_chunk_left_over=prev_chunk.clone(),
-        inference_delay=5,
-        time=torch.tensor(0.5),
-        original_denoise_step_partial=deterministic_denoiser,
-    )
-
-    # Should produce identical results
-    assert torch.allclose(result1, result2)

tests/policies/smolvla/test_smolvla_rtc.py

@@ -1,323 +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.
-
-"""Test SmolVLA policy with Real-Time Chunking (RTC) enabled during inference."""
-
-import pytest
-import torch
-
-from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule  # noqa: E402
-from lerobot.policies.factory import make_pre_post_processors  # noqa: E402
-from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig  # noqa: F401
-from lerobot.utils.random_utils import set_seed  # noqa: E402
-from tests.utils import require_cuda, require_package  # noqa: E402
-
-
-@require_package("transformers")
-@require_cuda
-def test_smolvla_rtc_initialization():
-    from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
-
-    """Test SmolVLA policy can initialize RTC processor."""
-    set_seed(42)
-
-    config = SmolVLAConfig(max_action_dim=7, chunk_size=50)
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Instantiate policy
-    policy = SmolVLAPolicy(config)
-
-    # Verify RTC processor is initialized
-    assert hasattr(policy, "rtc_processor")
-    assert policy.rtc_processor is not None
-    assert policy.rtc_processor.rtc_config.enabled is True
-
-    print("✓ SmolVLA RTC initialization: Test passed")
-
-
-@require_package("transformers")
-@require_cuda
-def test_smolvla_rtc_initialization_without_rtc_config():
-    from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
-
-    """Test SmolVLA policy can initialize without RTC config."""
-    set_seed(42)
-
-    config = SmolVLAConfig(max_action_dim=7, chunk_size=50)
-
-    # Instantiate policy
-    policy = SmolVLAPolicy(config)
-
-    # Verify RTC processor is not initialized
-    assert hasattr(policy, "rtc_processor")
-    assert policy.rtc_processor is None
-    assert policy.model.rtc_processor is None
-    assert policy._rtc_enabled() is False
-
-    print("✓ SmolVLA RTC initialization without RTC config: Test passed")
-
-
-@require_package("transformers")
-@require_cuda
-@pytest.mark.skipif(True, reason="Requires pretrained SmolVLA model weights")
-def test_smolvla_rtc_inference_with_prev_chunk():
-    from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
-
-    """Test SmolVLA policy inference with RTC and previous chunk."""
-    set_seed(42)
-
-    config = SmolVLAConfig(max_action_dim=7, chunk_size=50)
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and create preprocessor
-    policy = SmolVLAPolicy(config)
-    policy.eval()
-    preprocessor, _ = make_pre_post_processors(
-        policy_cfg=config, pretrained_path=None, dataset_stats=dataset_stats
-    )
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    # Create previous chunk
-    prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC and previous chunk
-        actions_with_rtc = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=prev_chunk,
-            inference_delay=4,
-            execution_horizon=10,
-        )
-
-        # Test without RTC for comparison
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    # Verify shapes
-    assert actions_with_rtc.shape == (1, config.chunk_size, 7)
-    assert actions_without_rtc.shape == (1, config.chunk_size, 7)
-
-    # With previous chunk, actions should be different (RTC guidance applied)
-    assert not torch.allclose(actions_with_rtc, actions_without_rtc, rtol=1e-3)
-
-    print("✓ SmolVLA RTC inference with prev_chunk: Test passed")
-
-
-@require_package("transformers")
-@require_cuda
-@pytest.mark.skipif(True, reason="Requires pretrained SmolVLA model weights")
-def test_smolvla_rtc_inference_without_prev_chunk():
-    from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
-
-    """Test SmolVLA policy inference with RTC but no previous chunk (RTC should have no effect)."""
-    set_seed(42)
-
-    config = SmolVLAConfig(max_action_dim=7, chunk_size=50)
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and create preprocessor
-    policy = SmolVLAPolicy(config)
-    policy.eval()
-    preprocessor, _ = make_pre_post_processors(
-        policy_cfg=config, pretrained_path=None, dataset_stats=dataset_stats
-    )
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC enabled but no previous chunk
-        actions_with_rtc_no_prev = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=None,
-        )
-
-        # Test without RTC
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    # Without previous chunk, RTC should have no effect
-    assert torch.allclose(actions_with_rtc_no_prev, actions_without_rtc, rtol=1e-5)
-
-    print("✓ SmolVLA RTC inference without prev_chunk: Test passed")
-
-
-@require_package("transformers")
-@require_cuda
-@pytest.mark.skipif(True, reason="Requires pretrained SmolVLA model weights")
-def test_smolvla_rtc_validation_rules():
-    from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
-
-    """Test SmolVLA policy with RTC follows all three validation rules."""
-    set_seed(42)
-
-    config = SmolVLAConfig(max_action_dim=7, chunk_size=50)
-
-    # Add RTC config
-    config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=10,
-        max_guidance_weight=5.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        debug=False,
-    )
-
-    config.input_features = {
-        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(14,)),
-        "observation.images.base_0_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
-    }
-    config.output_features = {
-        "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-    }
-
-    # Create dataset stats
-    dataset_stats = {
-        "observation.state": {"mean": torch.zeros(14), "std": torch.ones(14)},
-        "action": {"mean": torch.zeros(7), "std": torch.ones(7)},
-        "observation.images.base_0_rgb": {"mean": torch.zeros(3, 224, 224), "std": torch.ones(3, 224, 224)},
-    }
-
-    # Instantiate policy and create preprocessor
-    policy = SmolVLAPolicy(config)
-    policy.eval()
-    preprocessor, _ = make_pre_post_processors(
-        policy_cfg=config, pretrained_path=None, dataset_stats=dataset_stats
-    )
-
-    device = config.device
-
-    # Create dummy batch
-    batch = {
-        "observation.state": torch.randn(1, 14, dtype=torch.float32, device=device),
-        "observation.images.base_0_rgb": torch.rand(1, 3, 224, 224, dtype=torch.float32, device=device),
-        "task": ["Pick up the object"],
-    }
-    batch = preprocessor(batch)
-
-    # Create previous chunk
-    prev_chunk = torch.randn(1, 25, 7, dtype=torch.float32, device=device)
-
-    inference_delay = 4
-    execution_horizon = 10
-
-    with torch.no_grad():
-        # Use same noise for fair comparison
-        noise = policy.model.sample_noise((1, config.chunk_size, 7), device)
-
-        # Test with RTC
-        actions_with_rtc = policy.predict_action_chunk(
-            batch,
-            noise=noise.clone(),
-            prev_chunk_left_over=prev_chunk,
-            inference_delay=inference_delay,
-            execution_horizon=execution_horizon,
-        )
-
-        # Test without RTC
-        policy.config.rtc_config.enabled = False
-        actions_without_rtc = policy.predict_action_chunk(batch, noise=noise.clone())
-        policy.config.rtc_config.enabled = True
-
-    assert not torch.allclose(actions_with_rtc, actions_without_rtc, rtol=1e-3)

tests/policies/test_sarm_utils.py

@@ -1,378 +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.
-
-"""
-Tests for SARM utility functions.
-
-Tests the implementation of SARM paper formulas:
-- Formula (1): compute_temporal_proportions - dataset-level temporal proportions
-- Formula (2): compute_tau, compute_cumulative_progress - progress labels
-"""
-
-import pytest
-import numpy as np
-import torch
-
-from lerobot.policies.sarm.sarm_utils import SubtaskAnnotation, Subtask, Timestamp
-from lerobot.policies.sarm.sarm_utils import (
-    compute_temporal_proportions,
-    compute_tau,
-    compute_cumulative_progress_batch,
-)
-
-def make_annotation(subtasks: list[tuple[str, int, int]]) -> SubtaskAnnotation:
-    """Helper to create SubtaskAnnotation from list of (name, start_sec, end_sec)."""
-    return SubtaskAnnotation(
-        subtasks=[
-            Subtask(
-                name=name,
-                timestamps=Timestamp(
-                    start=f"{start // 60:02d}:{start % 60:02d}",
-                    end=f"{end // 60:02d}:{end % 60:02d}"
-                )
-            )
-            for name, start, end in subtasks
-        ]
-    )
-
-
-class TestComputeTemporalProportions:
-    """Tests for compute_temporal_proportions (SARM Paper Formula 1).
-    
-    Formula: ᾱ_k = (1/M) × Σ_i (L_{i,k} / T_i)
-    
-    Key insight: This averages the PROPORTION of each subtask within each trajectory,
-    giving equal weight to all trajectories regardless of absolute length.
-    """
-    
-    def test_basic_two_trajectories_equal_proportions(self):
-        """Test with two trajectories that have equal proportions."""
-        # Both trajectories: subtask1 = 50%, subtask2 = 50%
-        # Traj 1: T=100s, subtask1=50s, subtask2=50s
-        # Traj 2: T=200s, subtask1=100s, subtask2=100s
-        annotations = {
-            0: make_annotation([('subtask1', 0, 50), ('subtask2', 50, 100)]),
-            1: make_annotation([('subtask1', 0, 100), ('subtask2', 100, 200)]),
-        }
-        
-        result = compute_temporal_proportions(annotations)
-        
-        # Both should be 0.5
-        assert abs(result['subtask1'] - 0.5) < 1e-6
-        assert abs(result['subtask2'] - 0.5) < 1e-6
-    
-    def test_paper_example_different_from_avg_durations(self):
-        """Test that compute_temporal_proportions differs from naive average duration approach.
-        
-        This is the key test showing the difference between:
-        - Paper formula: average of (L_i,k / T_i)
-        - Naive approach: mean(L_i,k) / sum(mean(L_i,j))
-        """
-        # Episode 1: T=100s, subtask1=80s, subtask2=20s (proportions: 0.8, 0.2)
-        # Episode 2: T=200s, subtask1=40s, subtask2=160s (proportions: 0.2, 0.8)
-        annotations = {
-            0: make_annotation([('subtask1', 0, 80), ('subtask2', 80, 100)]),
-            1: make_annotation([('subtask1', 0, 40), ('subtask2', 40, 200)]),
-        }
-        
-        result = compute_temporal_proportions(annotations)
-        
-        # Paper formula: 
-        # ᾱ_1 = (1/2) × (80/100 + 40/200) = (1/2) × (0.8 + 0.2) = 0.5
-        # ᾱ_2 = (1/2) × (20/100 + 160/200) = (1/2) × (0.2 + 0.8) = 0.5
-        assert abs(result['subtask1'] - 0.5) < 1e-6
-        assert abs(result['subtask2'] - 0.5) < 1e-6
-    
-    def test_single_trajectory(self):
-        """Test with a single trajectory."""
-        # T=100s, reach=30s, grasp=20s, lift=50s
-        annotations = {
-            0: make_annotation([('reach', 0, 30), ('grasp', 30, 50), ('lift', 50, 100)]),
-        }
-        
-        result = compute_temporal_proportions(annotations)
-        
-        assert abs(result['reach'] - 0.3) < 1e-6
-        assert abs(result['grasp'] - 0.2) < 1e-6
-        assert abs(result['lift'] - 0.5) < 1e-6
-    
-    def test_sum_to_one(self):
-        """Test that proportions always sum to 1."""
-        # Three episodes with varying proportions
-        annotations = {
-            0: make_annotation([('a', 0, 10), ('b', 10, 50), ('c', 50, 100)]),  # 0.1, 0.4, 0.5
-            1: make_annotation([('a', 0, 20), ('b', 20, 70), ('c', 70, 100)]),  # 0.2, 0.5, 0.3
-            2: make_annotation([('a', 0, 30), ('b', 30, 90), ('c', 90, 100)]),  # 0.3, 0.6, 0.1
-        }
-        
-        result = compute_temporal_proportions(annotations)
-        
-        total = sum(result.values())
-        assert abs(total - 1.0) < 1e-6
-    
-    def test_empty_annotations_returns_empty(self):
-        """Test that empty annotations returns empty dict."""
-        result = compute_temporal_proportions({})
-        assert result == {}
-    
-    def test_uniform_proportions(self):
-        """Test with uniform proportions across subtasks."""
-        # Each subtask takes 25% of each episode
-        annotations = {
-            0: make_annotation([('a', 0, 25), ('b', 25, 50), ('c', 50, 75), ('d', 75, 100)]),
-            1: make_annotation([('a', 0, 50), ('b', 50, 100), ('c', 100, 150), ('d', 150, 200)]),
-        }
-        
-        result = compute_temporal_proportions(annotations)
-        
-        for name in ['a', 'b', 'c', 'd']:
-            assert abs(result[name] - 0.25) < 1e-6
-
-
-class TestComputeTau:
-    """Tests for compute_tau (within-subtask progress).
-    
-    Formula: τ_t = (t - s_k) / (e_k - s_k) ∈ [0, 1]
-    """
-    
-    def test_at_start(self):
-        """τ should be 0 at subtask start."""
-        tau = compute_tau(current_frame=10, subtask_start=10, subtask_end=50)
-        assert tau == 0.0
-    
-    def test_at_end(self):
-        """τ should be 1 at subtask end."""
-        tau = compute_tau(current_frame=50, subtask_start=10, subtask_end=50)
-        assert tau == 1.0
-    
-    def test_at_middle(self):
-        """τ should be 0.5 at subtask midpoint."""
-        tau = compute_tau(current_frame=30, subtask_start=10, subtask_end=50)
-        assert abs(tau - 0.5) < 1e-6
-    
-    def test_quarter_progress(self):
-        """Test τ at 25% through subtask."""
-        tau = compute_tau(current_frame=20, subtask_start=0, subtask_end=80)
-        assert abs(tau - 0.25) < 1e-6
-    
-    def test_zero_duration_subtask(self):
-        """τ should be 1.0 for zero-duration subtask."""
-        tau = compute_tau(current_frame=10, subtask_start=10, subtask_end=10)
-        assert tau == 1.0
-    
-    def test_clamps_below_zero(self):
-        """τ should be clamped to 0 if frame is before subtask."""
-        tau = compute_tau(current_frame=5, subtask_start=10, subtask_end=50)
-        assert tau == 0.0
-    
-    def test_clamps_above_one(self):
-        """τ should be clamped to 1 if frame is after subtask."""
-        tau = compute_tau(current_frame=60, subtask_start=10, subtask_end=50)
-        assert tau == 1.0
-    
-    def test_float_inputs(self):
-        """Test with float frame indices (from interpolation)."""
-        tau = compute_tau(current_frame=25.5, subtask_start=10.0, subtask_end=50.0)
-        expected = (25.5 - 10.0) / (50.0 - 10.0)
-        assert abs(tau - expected) < 1e-6
-
-
-class TestComputeCumulativeProgressBatchScalar:
-    """Tests for compute_cumulative_progress_batch with scalar inputs (normalized progress y_t).
-    
-    Formula: y_t = P_{k-1} + ᾱ_k × τ_t ∈ [0, 1]
-    """
-    
-    def test_first_subtask_start(self):
-        """y should be 0 at start of first subtask."""
-        proportions = [0.3, 0.5, 0.2]
-        y = compute_cumulative_progress_batch(tau=0.0, stage_indices=0, alpha=proportions)
-        assert y == 0.0
-    
-    def test_first_subtask_end(self):
-        """y should equal ᾱ_1 at end of first subtask."""
-        proportions = [0.3, 0.5, 0.2]
-        y = compute_cumulative_progress_batch(tau=1.0, stage_indices=0, alpha=proportions)
-        assert abs(y - 0.3) < 1e-6
-    
-    def test_second_subtask_start(self):
-        """y should equal P_1 at start of second subtask."""
-        proportions = [0.3, 0.5, 0.2]
-        y = compute_cumulative_progress_batch(tau=0.0, stage_indices=1, alpha=proportions)
-        assert abs(y - 0.3) < 1e-6
-    
-    def test_second_subtask_end(self):
-        """y should equal P_2 at end of second subtask."""
-        proportions = [0.3, 0.5, 0.2]
-        y = compute_cumulative_progress_batch(tau=1.0, stage_indices=1, alpha=proportions)
-        assert abs(y - 0.8) < 1e-6  # 0.3 + 0.5
-    
-    def test_third_subtask_end(self):
-        """y should be 1.0 at end of last subtask."""
-        proportions = [0.3, 0.5, 0.2]
-        y = compute_cumulative_progress_batch(tau=1.0, stage_indices=2, alpha=proportions)
-        assert abs(y - 1.0) < 1e-6
-    
-    def test_midpoint_of_subtask(self):
-        """Test progress at midpoint of a subtask."""
-        proportions = [0.4, 0.6]
-        # At τ=0.5 in subtask 1: y = P_0 + ᾱ_1 × 0.5 = 0 + 0.4 × 0.5 = 0.2
-        y = compute_cumulative_progress_batch(tau=0.5, stage_indices=0, alpha=proportions)
-        assert abs(y - 0.2) < 1e-6
-        
-        # At τ=0.5 in subtask 2: y = P_1 + ᾱ_2 × 0.5 = 0.4 + 0.6 × 0.5 = 0.7
-        y = compute_cumulative_progress_batch(tau=0.5, stage_indices=1, alpha=proportions)
-        assert abs(y - 0.7) < 1e-6
-    
-    def test_uniform_proportions(self):
-        """Test with uniform proportions."""
-        proportions = [0.25, 0.25, 0.25, 0.25]
-        
-        # At end of each subtask, progress should be 0.25, 0.5, 0.75, 1.0
-        for i in range(4):
-            y = compute_cumulative_progress_batch(tau=1.0, stage_indices=i, alpha=proportions)
-            expected = (i + 1) * 0.25
-            assert abs(y - expected) < 1e-6
-
-
-class TestComputeCumulativeProgressBatchTensor:
-    """Tests for compute_cumulative_progress_batch with tensor inputs (GPU batch version)."""
-    
-    def test_tensor_matches_scalar_version(self):
-        """Test that tensor version matches scalar version."""
-        proportions = [0.3, 0.5, 0.2]
-        alpha = torch.tensor(proportions, dtype=torch.float32)
-        cumulative = torch.zeros(len(proportions) + 1, dtype=torch.float32)
-        cumulative[1:] = torch.cumsum(alpha, dim=0)
-        
-        test_cases = [
-            (0.0, 0),  # start of subtask 0
-            (1.0, 0),  # end of subtask 0
-            (0.0, 1),  # start of subtask 1
-            (0.5, 1),  # middle of subtask 1
-            (1.0, 2),  # end of subtask 2
-        ]
-        
-        for tau_val, stage_idx in test_cases:
-            # Scalar version
-            expected = compute_cumulative_progress_batch(tau_val, stage_idx, proportions)
-            
-            # Tensor version (single element)
-            tau = torch.tensor([[[tau_val]]])  # (1, 1, 1)
-            stages = torch.tensor([[stage_idx]])  # (1, 1)
-            result = compute_cumulative_progress_batch(tau, stages, alpha, cumulative)
-            
-            assert abs(result[0, 0, 0].item() - expected) < 1e-6
-    
-    def test_batch_processing(self):
-        """Test batch processing with multiple samples."""
-        proportions = [0.4, 0.6]
-        alpha = torch.tensor(proportions, dtype=torch.float32)
-        cumulative = torch.zeros(3, dtype=torch.float32)
-        cumulative[1:] = torch.cumsum(alpha, dim=0)
-        
-        # Batch of 2 samples, sequence length 3
-        tau = torch.tensor([
-            [[0.0], [0.5], [1.0]],  # sample 1
-            [[0.0], [0.5], [1.0]],  # sample 2
-        ])
-        stages = torch.tensor([
-            [0, 0, 0],  # sample 1: all in subtask 0
-            [1, 1, 1],  # sample 2: all in subtask 1
-        ])
-        
-        result = compute_cumulative_progress_batch(tau, stages, alpha, cumulative)
-        
-        # Sample 1: subtask 0 with tau 0, 0.5, 1.0 -> y = 0, 0.2, 0.4
-        assert abs(result[0, 0, 0].item() - 0.0) < 1e-6
-        assert abs(result[0, 1, 0].item() - 0.2) < 1e-6
-        assert abs(result[0, 2, 0].item() - 0.4) < 1e-6
-        
-        # Sample 2: subtask 1 with tau 0, 0.5, 1.0 -> y = 0.4, 0.7, 1.0
-        assert abs(result[1, 0, 0].item() - 0.4) < 1e-6
-        assert abs(result[1, 1, 0].item() - 0.7) < 1e-6
-        assert abs(result[1, 2, 0].item() - 1.0) < 1e-6
-    
-    def test_auto_compute_cumulative_prior(self):
-        """Test that cumulative_prior is auto-computed when not provided."""
-        proportions = [0.3, 0.5, 0.2]
-        alpha = torch.tensor(proportions, dtype=torch.float32)
-        
-        tau = torch.tensor([[[0.5]]])
-        stages = torch.tensor([[1]])
-        
-        # Without cumulative_prior (should auto-compute)
-        result = compute_cumulative_progress_batch(tau, stages, alpha)
-        
-        # Expected: P_0 + alpha_1 * 0.5 = 0.3 + 0.5 * 0.5 = 0.55
-        assert abs(result[0, 0, 0].item() - 0.55) < 1e-6
-
-
-class TestEndToEndProgressLabeling:
-    """End-to-end tests for progress label computation."""
-    
-    def test_consistent_semantic_meaning(self):
-        """Test that same subtask completion maps to same progress across trajectories.
-        
-        This is the key semantic property: "end of subtask 1" should always 
-        mean the same progress value regardless of trajectory speed.
-        """
-        proportions = [0.3, 0.5, 0.2]
-        
-        # Fast trajectory: subtask 1 ends at frame 30 (of 100)
-        tau_fast = compute_tau(30, 0, 30)  # = 1.0
-        y_fast = compute_cumulative_progress_batch(tau_fast, 0, proportions)
-        
-        # Slow trajectory: subtask 1 ends at frame 90 (of 300)
-        tau_slow = compute_tau(90, 0, 90)  # = 1.0
-        y_slow = compute_cumulative_progress_batch(tau_slow, 0, proportions)
-        
-        # Both should map to same progress (0.3 = end of subtask 1)
-        assert abs(y_fast - y_slow) < 1e-6
-        assert abs(y_fast - 0.3) < 1e-6
-    
-    def test_monotonic_within_subtask(self):
-        """Test that progress is monotonically increasing within a subtask."""
-        proportions = [0.4, 0.6]
-        
-        prev_y = -1
-        for tau in np.linspace(0, 1, 11):
-            y = compute_cumulative_progress_batch(tau, 0, proportions)
-            assert y > prev_y or (tau == 0 and y == 0)
-            prev_y = y
-    
-    def test_continuous_across_subtasks(self):
-        """Test that progress is continuous at subtask boundaries."""
-        proportions = [0.3, 0.5, 0.2]
-        
-        # End of subtask 0 (tau=1.0)
-        y_end_0 = compute_cumulative_progress_batch(1.0, 0, proportions)
-        
-        # Start of subtask 1 (tau=0.0)
-        y_start_1 = compute_cumulative_progress_batch(0.0, 1, proportions)
-        
-        # Should be equal (P_1 = 0.3)
-        assert abs(y_end_0 - y_start_1) < 1e-6
-        
-        # End of subtask 1
-        y_end_1 = compute_cumulative_progress_batch(1.0, 1, proportions)
-        
-        # Start of subtask 2
-        y_start_2 = compute_cumulative_progress_batch(0.0, 2, proportions)
-        
-        # Should be equal (P_2 = 0.8)
-        assert abs(y_end_1 - y_start_2) < 1e-6
-

tests/processor/test_libero_processor.py

@@ -1,72 +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 numpy as np
-import torch
-
-from lerobot.envs.utils import preprocess_observation
-from lerobot.processor.env_processor import LiberoProcessorStep
-from lerobot.processor.pipeline import PolicyProcessorPipeline
-
-seed = 42
-np.random.seed(seed)
-
-B = 5
-obs1 = {
-    "pixels": {
-        "image": (np.random.rand(B, 256, 256, 3) * 255).astype(np.uint8),
-        "image2": (np.random.rand(B, 256, 256, 3) * 255).astype(np.uint8),
-    },
-    "robot_state": {
-        "eef": {
-            "pos": np.random.randn(B, 3),
-            "quat": np.random.randn(B, 4),
-            "mat": np.random.randn(B, 3, 3),
-        },
-        "gripper": {
-            "qpos": np.random.randn(B, 2),
-            "qvel": np.random.randn(B, 2),
-        },
-        "joints": {
-            "pos": np.random.randn(B, 7),
-            "vel": np.random.randn(B, 7),
-        },
-    },
-}
-
-observation = preprocess_observation(obs1)
-libero_preprocessor = PolicyProcessorPipeline(
-    steps=[
-        LiberoProcessorStep(),
-    ]
-)
-processed_obs = libero_preprocessor(observation)
-assert "observation.state" in processed_obs
-state = processed_obs["observation.state"]
-assert isinstance(state, torch.Tensor)
-assert state.dtype == torch.float32
-
-assert state.shape[0] == B
-assert state.shape[1] == 8
-
-assert "observation.images.image" in processed_obs
-assert "observation.images.image2" in processed_obs
-
-assert isinstance(processed_obs["observation.images.image"], torch.Tensor)
-assert isinstance(processed_obs["observation.images.image2"], torch.Tensor)
-
-assert processed_obs["observation.images.image"].shape == (B, 3, 256, 256)
-assert processed_obs["observation.images.image2"].shape == (B, 3, 256, 256)

tests/training/test_visual_validation.py

@@ -1,157 +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.
-
-"""
-Visual Feature Consistency Tests
-
-This module tests the `validate_visual_features_consistency` function,
-which ensures that visual features (camera observations) in a dataset/env
-match the expectations defined in a policy configuration.
-
-The purpose of this check is to prevent mismatches between what a policy expects
-(e.g., `observation.images.camera1`, `camera2`, `camera3`) and what a dataset or
-environment actually provides (e.g., `observation.images.top`, `side`, or fewer cameras).
-"""
-
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from lerobot.configs.default import DatasetConfig
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.policies.factory import make_policy_config
-from lerobot.scripts.lerobot_train import train
-from lerobot.utils.utils import auto_select_torch_device
-
-pytest.importorskip("transformers")
-
-DUMMY_REPO_ID = "dummy/repo"
-
-
-@pytest.fixture
-def temp_dir(tmp_path):
-    return tmp_path
-
-
-DUMMY_STATE_DIM = 6
-DUMMY_ACTION_DIM = 6
-IMAGE_SIZE = 8
-DEVICE = auto_select_torch_device()
-
-
-def make_dummy_dataset(camera_keys, tmp_path):
-    """Creates a minimal dummy dataset for testing rename_mapping logic."""
-    features = {
-        "action": {"dtype": "float32", "shape": (DUMMY_ACTION_DIM,), "names": None},
-        "observation.state": {"dtype": "float32", "shape": (DUMMY_STATE_DIM,), "names": None},
-    }
-    for cam in camera_keys:
-        features[f"observation.images.{cam}"] = {
-            "dtype": "image",
-            "shape": (IMAGE_SIZE, IMAGE_SIZE, 3),
-            "names": ["height", "width", "channel"],
-        }
-    dataset = LeRobotDataset.create(
-        repo_id=DUMMY_REPO_ID,
-        fps=30,
-        features=features,
-        root=tmp_path / "_dataset",
-    )
-    root = tmp_path / "_dataset"
-    for ep_idx in range(2):
-        for _ in range(3):
-            frame = {
-                "action": np.random.randn(DUMMY_ACTION_DIM).astype(np.float32),
-                "observation.state": np.random.randn(DUMMY_STATE_DIM).astype(np.float32),
-            }
-            for cam in camera_keys:
-                frame[f"observation.images.{cam}"] = np.random.randint(
-                    0, 255, size=(IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8
-                )
-            frame["task"] = f"task_{ep_idx}"
-            dataset.add_frame(frame)
-        dataset.save_episode()
-
-    dataset.finalize()
-    return dataset, root
-
-
-def custom_validate(train_config: TrainPipelineConfig, policy_path: str, empty_cameras: int):
-    train_config.policy = PreTrainedConfig.from_pretrained(policy_path)
-    train_config.policy.pretrained_path = Path(policy_path)
-    # override empty_cameras and push_to_hub for testing
-    train_config.policy.empty_cameras = empty_cameras
-    train_config.policy.push_to_hub = False
-    if train_config.use_policy_training_preset:
-        train_config.optimizer = train_config.policy.get_optimizer_preset()
-        train_config.scheduler = train_config.policy.get_scheduler_preset()
-    return train_config
-
-
-@pytest.mark.skip(reason="Skipping this test as it results OOM")
-@pytest.mark.parametrize(
-    "camera_keys, empty_cameras, rename_map, expect_success",
-    [
-        # case 1: dataset has fewer cameras than policy (3 instead of 4), but we specify empty_cameras=1 for smolvla, pi0, pi05
-        (["camera1", "camera2", "camera3"], 1, {}, True),
-        # case 2: dataset has 2 cameras with different names, rename_mapping provided
-        (
-            ["top", "side"],
-            0,
-            {
-                "observation.images.top": "observation.images.camera1",
-                "observation.images.side": "observation.images.camera2",
-            },
-            True,
-        ),
-        # case 3: dataset has 2 cameras, policy expects 3, names do not match, no empty_cameras
-        (["top", "side"], 0, {}, False),
-        # TODO: case 4: dataset has 2 cameras, policy expects 3, no rename_map, no empty_cameras, should raise for smolvla
-        # (["camera1", "camera2"], 0, {}, False),
-    ],
-)
-def test_train_with_camera_mismatch(camera_keys, empty_cameras, rename_map, expect_success, tmp_path):
-    """Tests that training works or fails depending on camera/feature alignment."""
-
-    _dataset, root = make_dummy_dataset(camera_keys, tmp_path)
-    pretrained_path = "lerobot/smolvla_base"
-    dataset_config = DatasetConfig(repo_id=DUMMY_REPO_ID, root=root)
-    policy_config = make_policy_config(
-        "smolvla",
-        optimizer_lr=0.01,
-        push_to_hub=False,
-        pretrained_path=pretrained_path,
-        device=DEVICE,
-    )
-    policy_config.empty_cameras = empty_cameras
-    train_config = TrainPipelineConfig(
-        dataset=dataset_config,
-        policy=policy_config,
-        rename_map=rename_map,
-        output_dir=tmp_path / "_output",
-        steps=1,
-    )
-    train_config = custom_validate(train_config, policy_path=pretrained_path, empty_cameras=empty_cameras)
-    # HACK: disable the internal CLI validation step for tests, we did it with custom_validate
-    train_config.validate = lambda: None
-    if expect_success:
-        train(train_config)
-    else:
-        with pytest.raises(ValueError):
-            train(train_config)