fix: minor things on the aggregation job

* add example scripts to load behaviour 1k data in `load_behaviour1k_dataset.py`

docs/source/_toctree.yml

@@ -15,8 +15,6 @@
     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"
@@ -40,6 +38,12 @@
   - local: groot
     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
@@ -59,6 +63,8 @@
     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

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

@@ -0,0 +1,188 @@
+# 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/behavior_1k/aggregate.slurm

@@ -0,0 +1,29 @@
+#!/bin/bash
+#SBATCH -J b1k-aggregate
+#SBATCH -p hopper-cpu
+#SBATCH --qos=high
+#SBATCH -c 2
+#SBATCH -t 20:00:00
+#SBATCH --mem=4G
+#SBATCH -D /admin/home/francesco_capuano/lerobot
+#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%j.out
+#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%j.err
+
+set -euo pipefail
+set -x
+export PYTHONUNBUFFERED=1
+export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
+
+source "$HOME/.bashrc" 2>/dev/null || true
+if ! command -v conda >/dev/null 2>&1; then
+  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+fi
+conda activate lerobot
+
+python examples/behavior_1k/aggregate_tasks_datasets.py \
+  --task-datasets-dir /fsx/francesco_capuano/behavior1k-v3 \
+  --aggregated-root /fsx/francesco_capuano/behavior1k-v3/behavior1k \
+  --num-tasks 50 \
+  --hf-user fracapuano \
+  --push-to-hub

examples/behavior_1k/aggregate_tasks_datasets.py

@@ -0,0 +1,100 @@
+"""Aggregate multiple task-specific LeRobot datasets into a single combined dataset."""
+
+import argparse
+import os
+from pathlib import Path
+
+from lerobot.datasets.aggregate import aggregate_datasets
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Aggregate multiple task-specific datasets into a single LeRobot dataset"
+    )
+    parser.add_argument(
+        "--task-datasets-dir",
+        type=str,
+        required=True,
+        help="Directory containing individual task datasets (e.g., /path/to/behavior1k/)",
+    )
+    parser.add_argument(
+        "--aggregated-root",
+        type=str,
+        required=True,
+        help="Path where the aggregated dataset will be written",
+    )
+    parser.add_argument(
+        "--num-tasks",
+        type=int,
+        default=50,
+        help="Number of tasks to aggregate (default: 50)",
+    )
+    parser.add_argument(
+        "--task-start-idx",
+        type=int,
+        default=0,
+        help="Starting task index (default: 0)",
+    )
+    parser.add_argument(
+        "--hf-user",
+        type=str,
+        default=None,
+        help="HuggingFace username for repo IDs (defaults to HF_USER env var or 'lerobot')",
+    )
+    parser.add_argument(
+        "--aggregated-repo-id",
+        type=str,
+        default=None,
+        help="Repository ID for the aggregated dataset (defaults to {hf_user}/behavior1k)",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        action="store_true",
+        help="Push the aggregated dataset to the Hugging Face Hub",
+    )
+
+    args = parser.parse_args()
+
+    # Determine HF user
+    hf_user = args.hf_user or os.environ.get("HF_USER", "lerobot")
+
+    # Set default aggregated repo ID if not provided
+    aggregated_repo_id = args.aggregated_repo_id or f"{hf_user}/behavior1k"
+
+    # Generate task indices
+    task_indices = range(args.task_start_idx, args.task_start_idx + args.num_tasks)
+
+    # Generate repo IDs for individual tasks
+    repo_ids = [f"{hf_user}/behavior1k-task{i:04d}" for i in task_indices]
+
+    # Generate local paths for individual task datasets
+    task_datasets_dir = Path(args.task_datasets_dir)
+    roots = [task_datasets_dir / f"behavior1k-task{i:04d}" for i in task_indices]
+
+    # Aggregated dataset path
+    aggregated_root = Path(args.aggregated_root)
+
+    print(f"🔹 Aggregating {args.num_tasks} task datasets")
+    print(f"Task datasets directory: {task_datasets_dir}")
+    print(f"Aggregated output: {aggregated_root}")
+    print(f"Aggregated repo ID: {aggregated_repo_id}")
+
+    aggregate_datasets(
+        repo_ids=repo_ids,
+        roots=roots,
+        aggr_repo_id=aggregated_repo_id,
+        aggr_root=aggregated_root,
+    )
+
+    print("✅ Aggregation complete")
+
+    if args.push_to_hub:
+        print(f"📤 Pushing aggregated dataset to {aggregated_repo_id}")
+        ds = LeRobotDataset(repo_id=aggregated_repo_id, root=aggregated_root)
+        ds.push_to_hub()
+        print("✅ Successfully pushed to hub")
+
+
+if __name__ == "__main__":
+    main()

examples/behavior_1k/convert.slurm

@@ -0,0 +1,38 @@
+#!/bin/bash
+#SBATCH -J b1k-convert
+#SBATCH -p hopper-cpu            # pick your partition
+#SBATCH --qos=high
+#SBATCH --array=0-49%8           # 50 tasks, max 8 running concurrently (conversion is I/O bound)
+#SBATCH -c 1                     # CPUs per conversion (tune as needed)
+#SBATCH -t 2:00:00               # Time per conversion
+#SBATCH --mem=3G                 # ~1.75GB for task 0, ~doubled for safety
+#SBATCH -D /admin/home/francesco_capuano/lerobot
+#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.out
+#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.err
+
+set -euo pipefail
+set -x
+export PYTHONUNBUFFERED=1
+export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}  # avoid BLAS oversubscription
+
+DATA_PATH="/fsx/francesco_capuano/behavior1k-2025-v21"
+BASE_OUT="/fsx/francesco_capuano/behavior1k-v3"
+mkdir -p "$BASE_OUT" logs
+
+i="${SLURM_ARRAY_TASK_ID}"
+OUT_DIR="$(printf "%s/behavior1k-task%04d" "$BASE_OUT" "$i")"
+
+# activate your env if needed
+source "$HOME/.bashrc" 2>/dev/null || true
+if ! command -v conda >/dev/null 2>&1; then
+  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+fi
+conda activate lerobot
+
+python examples/behavior_1k/convert_to_lerobot_v3.py \
+  --data-path "$DATA_PATH" \
+  --new-repo "$OUT_DIR" \
+  --task-id "$i" \
+  --force-conversion \
+  --push-to-hub

examples/behavior_1k/convert_to_lerobot_v3.py

@@ -0,0 +1,667 @@
+#!/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.
+"""Convert Behavior Dataset to LeRobotDataset v3.0 format"""
+
+import argparse
+import json
+import logging
+import os
+import shutil
+from pathlib import Path
+
+import jsonlines
+import pandas as pd
+import pyarrow as pa
+import tqdm
+from datasets import Dataset, Features, Image
+
+from lerobot.datasets.compute_stats import aggregate_stats
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_DATA_PATH,
+    DEFAULT_FEATURES,
+    DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    DEFAULT_VIDEO_PATH,
+    LEGACY_EPISODES_PATH,
+    LEGACY_EPISODES_STATS_PATH,
+    LEGACY_TASKS_PATH,
+    cast_stats_to_numpy,
+    flatten_dict,
+    get_file_size_in_mb,
+    get_parquet_file_size_in_mb,
+    get_parquet_num_frames,
+    load_info,
+    update_chunk_file_indices,
+    write_episodes,
+    write_info,
+    write_stats,
+    write_tasks,
+)
+from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
+from lerobot.utils.utils import init_logging
+
+# script to convert one single task to v3.1
+# TASK = 1
+NEW_ROOT = Path("/fsx/jade_choghari/tmp/bb")
+
+
+def fix_episode_dataframe(df: pd.DataFrame) -> pd.DataFrame:
+    """Performs several fixes to an underlying dataframe to make it LeRobotDataset-v3 compatible"""
+    # Inject per-episode frame_index if missing (0..N-1 within each episode)
+    if "frame_index" not in df.columns:
+        df["frame_index"] = range(len(df))
+
+    # Remove variable-length task_info feature (NOTE(fracapuano): change to padding at some point?)
+    if "observation.task_info" in df.columns:
+        df = df.drop(columns=["observation.task_info"])
+
+    # NOTE(fracapuano): tasks are ordered (and there is one task per file/dataset)
+    if "task_index" in df.columns:
+        df["task_index"] = 0
+
+    return df
+
+
+def get_total_episodes_task(local_dir: Path, task_id: int, task_ranges: dict, step) -> int:
+    """
+    Calculates the total number of episodes for a single, specified task.
+    """
+    # Simply load the episodes for the task and count them.
+    episodes = legacy_load_episodes_task(
+        local_dir=local_dir, task_id=task_id, task_ranges=task_ranges, step=step
+    )
+    return len(episodes)
+
+
+NUM_CAMERAS = 9
+
+
+def get_total_frames_task(local_dir, meta_path, task_id: int, task_ranges: dict, step: int) -> int:
+    episodes_metadata = legacy_load_episodes_task(
+        local_dir=local_dir, task_id=task_id, task_ranges=task_ranges, step=step
+    )
+    total_frames = 0
+    # like 'duration'
+    for ep in episodes_metadata.values():
+        duration_s = ep["length"]
+        total_frames += int(duration_s)
+    return total_frames
+
+
+def convert_info(
+    root, new_root, data_file_size_in_mb, video_file_size_in_mb, meta_path, task_id: int, task_ranges, step
+):
+    info = load_info(root)
+    features = {**info["features"], **DEFAULT_FEATURES}
+    del features[
+        "observation.task_info"
+    ]  # variable-length task_info is not supported in LeRobotDataset v3.0!
+
+    info["codebase_version"] = "v3.0"
+    info["features"] = features
+    del info["total_videos"]
+    info["data_files_size_in_mb"] = data_file_size_in_mb
+    info["video_files_size_in_mb"] = video_file_size_in_mb
+    info["data_path"] = DEFAULT_DATA_PATH
+    info["video_path"] = DEFAULT_VIDEO_PATH if info["video_path"] is not None else None
+    info["fps"] = int(info["fps"])
+    for key in info["features"]:
+        if info["features"][key]["dtype"] == "video":
+            # already has fps in video_info
+            continue
+        info["features"][key]["fps"] = info["fps"]
+
+    info["total_episodes"] = get_total_episodes_task(root, task_id, task_ranges, step)
+    info["total_videos"] = info["total_episodes"] * NUM_CAMERAS
+    info["total_frames"] = get_total_frames_task(root, meta_path, task_id, task_ranges, step)
+    info["total_tasks"] = 1
+    write_info(info, new_root)
+
+
+def load_jsonlines(fpath: Path) -> list[any]:
+    with jsonlines.open(fpath, "r") as reader:
+        return list(reader)
+
+
+def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
+    tasks = load_jsonlines(local_dir / LEGACY_TASKS_PATH)
+    # return tasks dict such that
+    tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
+    task_to_task_index = {task: task_index for task_index, task in tasks.items()}
+    return tasks, task_to_task_index
+
+
+def convert_tasks(root, new_root, task_id: int):
+    tasks, _ = legacy_load_tasks(root)
+    if task_id not in tasks:
+        raise ValueError(f"Task ID {task_id} not found in tasks (available: {list(tasks.keys())})")
+    tasks = {task_id: tasks[task_id]}
+    # Tasks are ordered with 0..ntasks-1 in the converted dataset
+    task_indices = range(len(tasks.keys()))
+    task_strings = tasks.values()
+    df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
+    write_tasks(df_tasks, new_root)
+
+
+def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys):
+    # TODO(rcadene): to save RAM use Dataset.from_parquet(file) and concatenate_datasets
+    dataframes = []
+    for file in paths_to_cat:
+        df = pd.read_parquet(file)
+        df = fix_episode_dataframe(df)
+        dataframes.append(df)
+
+    # Concatenate all DataFrames along rows
+    concatenated_df = pd.concat(dataframes, ignore_index=True)
+
+    path = new_root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if len(image_keys) > 0:
+        schema = pa.Schema.from_pandas(concatenated_df)
+        features = Features.from_arrow_schema(schema)
+        for key in image_keys:
+            features[key] = Image()
+        schema = features.arrow_schema
+    else:
+        schema = None
+
+    concatenated_df.to_parquet(path, index=False, schema=schema)
+
+
+def get_image_keys(root):
+    info = load_info(root)
+    features = info["features"]
+    image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"]
+    return image_keys
+
+
+def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int, task_index: int):
+    task_dir_name = f"task-{task_index:04d}"
+    data_dir = root / "data" / task_dir_name
+    ep_paths = sorted(data_dir.glob("*.parquet"))
+    image_keys = get_image_keys(root)
+
+    ep_idx = 0
+    chunk_idx = 0
+    file_idx = 0
+    size_in_mb = 0
+    num_frames = 0
+    paths_to_cat = []
+    episodes_metadata = []
+
+    logging.info(f"Converting data files from {len(ep_paths)} episodes")
+
+    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
+        ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
+        ep_num_frames = get_parquet_num_frames(ep_path)
+        ep_metadata = {
+            "episode_index": ep_idx,
+            "data/chunk_index": chunk_idx,
+            "data/file_index": file_idx,
+            "dataset_from_index": num_frames,
+            "dataset_to_index": num_frames + ep_num_frames,
+        }
+        size_in_mb += ep_size_in_mb
+        num_frames += ep_num_frames
+        episodes_metadata.append(ep_metadata)
+
+        # write 0-based episode index instead of custom episode index (otherwise breaks compatibility with LeRobotDataset)
+        tmp_df = pd.read_parquet(ep_path)
+        tmp_df["episode_index"] = ep_idx
+        tmp_df.to_parquet(ep_path)
+
+        ep_idx += 1
+
+        if size_in_mb < data_file_size_in_mb:
+            paths_to_cat.append(ep_path)
+            continue
+
+        if paths_to_cat:
+            concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
+
+        # Reset for the next file
+        size_in_mb = ep_size_in_mb
+        paths_to_cat = [ep_path]
+
+        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+
+    # Write remaining data if any
+    if paths_to_cat:
+        concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
+
+    return episodes_metadata
+
+
+def convert_videos_of_camera(
+    root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int, task_index: int
+):
+    # Access old paths to mp4
+    # videos_dir = root / "videos"
+    # ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4"))
+    task_dir_name = f"task-{task_index:04d}"
+    videos_dir = root / "videos" / task_dir_name / video_key
+    ep_paths = sorted(videos_dir.glob("*.mp4"))
+    ep_idx = 0
+    chunk_idx = 0
+    file_idx = 0
+    size_in_mb = 0
+    duration_in_s = 0.0
+    paths_to_cat = []
+    episodes_metadata = []
+
+    for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
+        ep_size_in_mb = get_file_size_in_mb(ep_path)
+        ep_duration_in_s = get_video_duration_in_s(ep_path)
+
+        # Check if adding this episode would exceed the limit
+        if size_in_mb + ep_size_in_mb >= video_file_size_in_mb and len(paths_to_cat) > 0:
+            # Size limit would be exceeded, save current accumulation WITHOUT this episode
+            concatenate_video_files(
+                paths_to_cat,
+                new_root
+                / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
+            )
+
+            # Update episodes metadata for the file we just saved
+            for i, _ in enumerate(paths_to_cat):
+                past_ep_idx = ep_idx - len(paths_to_cat) + i
+                episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
+                episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
+
+            # Move to next file and start fresh with current episode
+            chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
+            size_in_mb = 0
+            duration_in_s = 0.0
+            paths_to_cat = []
+
+        # Add current episode metadata
+        ep_metadata = {
+            "episode_index": ep_idx,
+            f"videos/{video_key}/chunk_index": chunk_idx,  # Will be updated when file is saved
+            f"videos/{video_key}/file_index": file_idx,  # Will be updated when file is saved
+            f"videos/{video_key}/from_timestamp": duration_in_s,
+            f"videos/{video_key}/to_timestamp": duration_in_s + ep_duration_in_s,
+        }
+        episodes_metadata.append(ep_metadata)
+
+        # Add current episode to accumulation
+        paths_to_cat.append(ep_path)
+        size_in_mb += ep_size_in_mb
+        duration_in_s += ep_duration_in_s
+        ep_idx += 1
+
+    # Write remaining videos if any
+    if paths_to_cat:
+        concatenate_video_files(
+            paths_to_cat,
+            new_root
+            / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx),
+        )
+
+        # Update episodes metadata for the final file
+        for i, _ in enumerate(paths_to_cat):
+            past_ep_idx = ep_idx - len(paths_to_cat) + i
+            episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx
+            episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx
+
+    return episodes_metadata
+
+
+def get_video_keys(root):
+    info = load_info(root)
+    features = info["features"]
+    video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"]
+    return video_keys
+
+
+def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int, task_id: int):
+    logging.info(f"Converting videos from {root} to {new_root}")
+
+    video_keys = get_video_keys(root)
+    if len(video_keys) == 0:
+        return None
+
+    video_keys = sorted(video_keys)
+
+    eps_metadata_per_cam = []
+    for camera in video_keys:
+        eps_metadata = convert_videos_of_camera(root, new_root, camera, video_file_size_in_mb, task_id)
+        eps_metadata_per_cam.append(eps_metadata)
+
+    num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam]
+    if len(set(num_eps_per_cam)) != 1:
+        raise ValueError(f"All cams dont have same number of episodes ({num_eps_per_cam}).")
+
+    episodes_metadata = []
+    num_cameras = len(video_keys)
+    num_episodes = num_eps_per_cam[0]
+    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
+        # Sanity check
+        ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
+        ep_ids += [ep_idx]
+        if len(set(ep_ids)) != 1:
+            raise ValueError(f"All episode indices need to match ({ep_ids}).")
+
+        ep_dict = {}
+        for cam_idx in range(num_cameras):
+            ep_dict.update(eps_metadata_per_cam[cam_idx][ep_idx])
+        episodes_metadata.append(ep_dict)
+
+    return episodes_metadata
+
+
+def infer_task_episode_ranges(episodes_jsonl_path: Path) -> dict:
+    """
+    Parse the Behavior-1K episodes.jsonl metadata and infer contiguous episode ranges per unique task.
+    Returns a dict:
+      { task_id: { "task_string": ..., "ep_start": ..., "ep_end": ... } }
+    """
+    task_ranges = {}
+    task_id = 0
+    current_task_str = None
+    ep_start = None
+    ep_end = None
+
+    with open(episodes_jsonl_path) as f:
+        for line in f:
+            if not line.strip():
+                continue
+            ep = json.loads(line)
+            ep_idx = ep["episode_index"]
+            task_str = ep["tasks"][0] if ep["tasks"] else "UNKNOWN"
+
+            if current_task_str is None:
+                current_task_str = task_str
+                ep_start = ep_idx
+                ep_end = ep_idx
+            elif task_str == current_task_str:
+                ep_end = ep_idx
+            else:
+                # close previous task group
+                task_ranges[task_id] = {
+                    "task_string": current_task_str,
+                    "ep_start": ep_start,
+                    "ep_end": ep_end,
+                }
+                task_id += 1
+                # start new one
+                current_task_str = task_str
+                ep_start = ep_idx
+                ep_end = ep_idx
+
+    # store last task
+    if current_task_str is not None:
+        task_ranges[task_id] = {
+            "task_string": current_task_str,
+            "ep_start": ep_start,
+            "ep_end": ep_end,
+        }
+
+    return task_ranges
+
+
+def legacy_load_episodes_task(local_dir: Path, task_id: int, task_ranges: dict, step: int = 10) -> dict:
+    """
+    Load only the episodes belonging to a specific task, inferred automatically from episode ranges.
+
+    Args:
+        local_dir (Path): Root path containing legacy meta/episodes.jsonl
+        task_id (int): Which task to load (key from the inferred task_ranges dict)
+        task_ranges (dict): Mapping from infer_task_episode_ranges()
+        step (int): Episode index step (Behavior-1K = 10)
+    """
+    all_episodes = legacy_load_episodes(local_dir)
+
+    # get the range for this task
+    if task_id not in task_ranges:
+        raise ValueError(f"Task id {task_id} not found in task_ranges")
+
+    ep_start = task_ranges[task_id]["ep_start"]
+    ep_end = task_ranges[task_id]["ep_end"]
+
+    task_episode_indices = range(ep_start, ep_end + step, step)
+    return {i: all_episodes[i] for i in task_episode_indices if i in all_episodes}
+
+
+def legacy_load_episodes(local_dir: Path) -> dict:
+    episodes = load_jsonlines(local_dir / LEGACY_EPISODES_PATH)
+    return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
+
+
+def legacy_load_episodes_stats(local_dir: Path) -> dict:
+    episodes_stats = load_jsonlines(local_dir / LEGACY_EPISODES_STATS_PATH)
+    return {
+        item["episode_index"]: cast_stats_to_numpy(item["stats"])
+        for item in sorted(episodes_stats, key=lambda x: x["episode_index"])
+    }
+
+
+def legacy_load_episodes_stats_task(local_dir: Path, task_id: int, task_ranges: dict, step: int = 10) -> dict:
+    all_stats = legacy_load_episodes_stats(local_dir)
+
+    if task_id not in task_ranges:
+        raise ValueError(f"Task id {task_id} not found in task_ranges")
+
+    ep_start = task_ranges[task_id]["ep_start"]
+    ep_end = task_ranges[task_id]["ep_end"]
+
+    task_episode_indices = range(ep_start, ep_end + step, step)
+    return {i: all_stats[i] for i in task_episode_indices if i in all_stats}
+
+
+def generate_episode_metadata_dict(
+    episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_videos=None
+):
+    num_episodes = len(episodes_metadata)
+    episodes_legacy_metadata_vals = list(episodes_legacy_metadata.values())
+    episodes_stats_vals = list(episodes_stats.values())
+    episodes_stats_keys = list(episodes_stats.keys())
+
+    for i in range(num_episodes):
+        ep_legacy_metadata = episodes_legacy_metadata_vals[i]
+        ep_metadata = episodes_metadata[i]
+        ep_stats = episodes_stats_vals[i]
+
+        ep_ids_set = {
+            ep_legacy_metadata["episode_index"],
+            ep_metadata["episode_index"],
+            episodes_stats_keys[i],
+        }
+
+        if episodes_videos is None:
+            ep_video = {}
+        else:
+            ep_video = episodes_videos[i]
+            ep_ids_set.add(ep_video["episode_index"])
+
+        ep_dict = {
+            **ep_legacy_metadata,
+            **ep_video,
+            **ep_metadata,
+            **flatten_dict({"stats": ep_stats}),
+        }
+
+        # enforce contiguous indexing 0..n-1, but also stores the legacy episode index
+        ep_dict["episode_index"] = i
+
+        yield ep_dict
+
+
+def convert_episodes_metadata(
+    root, new_root, episodes_metadata, task_id: int, task_ranges, episodes_video_metadata=None
+):
+    logging.info(f"Converting episodes metadata from {root} to {new_root}")
+
+    # filter by task
+    episodes_legacy_metadata = legacy_load_episodes_task(root, task_id=task_id, task_ranges=task_ranges)
+    episodes_stats = legacy_load_episodes_stats_task(root, task_id=task_id, task_ranges=task_ranges)
+
+    num_eps_set = {len(episodes_legacy_metadata), len(episodes_metadata)}
+    if episodes_video_metadata is not None:
+        num_eps_set.add(len(episodes_video_metadata))
+
+    if len(num_eps_set) != 1:
+        raise ValueError(f"Number of episodes is not the same ({num_eps_set}).")
+
+    # Single file approach: set meta indices to 0 for all rows and write once
+    ds_episodes = Dataset.from_generator(
+        lambda: generate_episode_metadata_dict(
+            episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_video_metadata
+        )
+    )
+    num_eps = len(ds_episodes)
+    # NOTE(fracapuano): for the size of the average dataset this is fine!
+    ds_episodes = ds_episodes.add_column("meta/episodes/chunk_index", [0] * num_eps)
+    ds_episodes = ds_episodes.add_column("meta/episodes/file_index", [0] * num_eps)
+    write_episodes(ds_episodes, new_root)
+
+    stats = aggregate_stats(list(episodes_stats.values()))
+    write_stats(stats, new_root)
+
+
+def convert_dataset_local(
+    data_path: Path,
+    new_repo: Path,
+    task_id: int,
+    data_file_size_in_mb: int = DEFAULT_DATA_FILE_SIZE_IN_MB,
+    video_file_size_in_mb: int = DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+    force_conversion: bool = False,
+):
+    """
+    Convert a local dataset to v3.x format, task-by-task, without using the Hugging Face Hub.
+
+    Args:
+        data_path (Path): path to local dataset root (e.g. /fsx/.../2025-challenge-demos)
+        new_repo (Path): path where converted dataset will be written (e.g. /fsx/.../behavior1k_v3)
+        task_id (int): which task to convert (index)
+        data_file_size_in_mb (int): max size per data chunk
+        video_file_size_in_mb (int): max size per video chunk
+        force_conversion (bool): overwrite existing conversion if True
+    """
+
+    root = Path(data_path)
+    new_root = Path(new_repo)
+
+    # Clean up if needed
+    if new_root.exists() and force_conversion:
+        shutil.rmtree(new_root)
+    new_root.mkdir(parents=True, exist_ok=True)
+
+    print(f"🔹 Starting conversion for task {task_id}")
+    print(f"Input root: {root}")
+    print(f"Output root: {new_root}")
+    # Infer task episode ranges
+    episodes_meta_path = root / "meta" / "episodes.jsonl"
+    task_ranges = infer_task_episode_ranges(episodes_meta_path)
+    convert_info(
+        root,
+        new_root,
+        data_file_size_in_mb,
+        video_file_size_in_mb,
+        episodes_meta_path,
+        task_id,
+        task_ranges,
+        step=10,
+    )
+    convert_tasks(root, new_root, task_id)
+    episodes_metadata = convert_data(root, new_root, data_file_size_in_mb, task_index=task_id)
+    episodes_videos_metadata = convert_videos(root, new_root, video_file_size_in_mb, task_id=task_id)
+    convert_episodes_metadata(
+        root,
+        new_root,
+        episodes_metadata,
+        task_id=task_id,
+        task_ranges=task_ranges,
+        episodes_video_metadata=episodes_videos_metadata,
+    )
+
+    print(f"✅ Conversion complete for task {task_id}")
+    print(f"Converted dataset written to: {new_root}")
+
+
+if __name__ == "__main__":
+    import argparse
+    from pathlib import Path
+
+    init_logging()
+
+    parser = argparse.ArgumentParser(
+        description="Convert Behavior-1K tasks to LeRobot v3 format (local only)"
+    )
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        required=True,
+        help="Path to the local Behavior-1K dataset (e.g. /fsx/francesco_capuano/.cache/behavior-1k/2025-challenge-demos)",
+    )
+    parser.add_argument(
+        "--new-repo",
+        type=str,
+        required=True,
+        help="Path to the output directory for the converted dataset",
+    )
+    parser.add_argument(
+        "--task-id",
+        type=int,
+        required=True,
+        help="Task index to convert (e.g. 0, 1, 2, ...)",
+    )
+    parser.add_argument(
+        "--data-file-size-in-mb",
+        type=int,
+        default=DEFAULT_DATA_FILE_SIZE_IN_MB,
+        help=f"Maximum size per data chunk (default: {DEFAULT_DATA_FILE_SIZE_IN_MB})",
+    )
+    parser.add_argument(
+        "--video-file-size-in-mb",
+        type=int,
+        default=DEFAULT_VIDEO_FILE_SIZE_IN_MB,
+        help=f"Maximum size per video chunk (default: {DEFAULT_VIDEO_FILE_SIZE_IN_MB})",
+    )
+    parser.add_argument(
+        "--force-conversion",
+        action="store_true",
+        help="Force overwrite of existing conversion output if present.",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        action="store_true",
+        help="Push the (converted) dataset to the hub.",
+    )
+
+    args = parser.parse_args()
+
+    if args.push_to_hub:
+        HF_USER = os.environ.get("HF_USER", "fracapuano")
+        if HF_USER is None:
+            raise ValueError(
+                "HF_USER environment variable is not set! Set before converting and pushing to hub."
+            )
+
+    convert_dataset_local(
+        data_path=Path(args.data_path),
+        new_repo=Path(args.new_repo),
+        task_id=args.task_id,
+        data_file_size_in_mb=args.data_file_size_in_mb,
+        video_file_size_in_mb=args.video_file_size_in_mb,
+        force_conversion=args.force_conversion,
+    )
+
+    if args.push_to_hub:
+        ds = LeRobotDataset(repo_id=f"{HF_USER}/behavior1k-task{args.task_id:04d}", root=args.new_repo)
+        ds.push_to_hub()

examples/behavior_1k/download.slurm

@@ -0,0 +1,27 @@
+#!/bin/bash
+#SBATCH -J b1k-download
+#SBATCH -p hopper-cpu
+#SBATCH --qos=high
+#SBATCH -c 32                    # CPUs per conversion (tune as needed)
+#SBATCH -t 20:00:00              # Time per conversion
+#SBATCH -D /admin/home/francesco_capuano/lerobot
+#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A.out
+#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A.err
+
+set -euo pipefail
+set -x
+export PYTHONUNBUFFERED=1
+export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
+
+# activate your env if needed
+source "$HOME/.bashrc" 2>/dev/null || true
+if ! command -v conda >/dev/null 2>&1; then
+  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+fi
+conda activate lerobot
+
+python examples/behavior_1k/download_data.py \
+  --repo-id "behavior-1k/2025-challenge-demos" \
+  --local-dir "/fsx/francesco_capuano/behavior1k-2025-v21" \
+  --max-workers 32

examples/behavior_1k/download_data.py

@@ -0,0 +1,26 @@
+import shutil
+
+from huggingface_hub import snapshot_download
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-id", type=str, required=True)
+    parser.add_argument("--max-workers", type=int, default=8)
+    parser.add_argument("--local-dir", type=str, required=True)
+    parser.add_argument("--force-download", action="store_true")
+
+    args = parser.parse_args()
+
+    if args.force_download:
+        shutil.rmtree(args.local_dir, ignore_errors=True)
+
+    snapshot_download(
+        repo_id=args.repo_id,
+        repo_type="dataset",
+        force_download=args.force_download,
+        max_workers=args.max_workers,
+        local_dir=args.local_dir,
+        ignore_patterns=["annotations/*"],  # NOTE(fracapuano): Dropping textual annotations right now
+    )

examples/behavior_1k/upload.slurm

@@ -0,0 +1,41 @@
+#!/bin/bash
+#SBATCH -J b1k-upload
+#SBATCH -p hopper-cpu
+#SBATCH --qos=high
+#SBATCH -c 1
+#SBATCH -t 48:00:00
+#SBATCH --mem=4G
+#SBATCH --array=0-49%2
+#SBATCH -D /admin/home/francesco_capuano/lerobot
+#SBATCH -o /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.out
+#SBATCH -e /admin/home/francesco_capuano/lerobot/examples/behavior_1k/logs/%x-%A_%a.err
+
+set -euo pipefail
+set -x
+export PYTHONUNBUFFERED=1
+export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
+
+source "$HOME/.bashrc" 2>/dev/null || true
+if ! command -v conda >/dev/null 2>&1; then
+  source "$HOME/miniconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+  source "$HOME/anaconda3/etc/profile.d/conda.sh" 2>/dev/null || true
+fi
+conda activate lerobot
+
+# The SLURM_ARRAY_TASK_ID will be used as the task-id
+TASK_ID=${SLURM_ARRAY_TASK_ID}
+
+# Configuration
+ROOT_PATH="/fsx/francesco_capuano/behavior1k-v3"
+HF_USER="fracapuano"
+# Limit upload workers to reduce network contention (default in HF Hub is 4)
+# For I/O-bound uploads, 2-4 workers per task is optimal
+NUM_WORKERS=2
+
+echo "Task ${TASK_ID}: uploading with ${NUM_WORKERS} workers from ${ROOT_PATH}"
+
+python examples/behavior_1k/upload_folders.py \
+  --task-id ${TASK_ID} \
+  --root-path ${ROOT_PATH} \
+  --hf-user ${HF_USER} \
+  --num-workers ${NUM_WORKERS}

examples/behavior_1k/upload_folders.py

@@ -0,0 +1,108 @@
+import argparse
+from pathlib import Path
+
+from huggingface_hub import HfApi, upload_large_folder
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Upload a folder to Hugging Face Hub using upload_large_folder"
+    )
+    parser.add_argument(
+        "--folder-path",
+        type=str,
+        required=False,
+        help="Path to the folder to upload (used if task-id is not provided)",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=False,
+        help="Repository ID on Hugging Face Hub (e.g., 'username/repo-name'). If task-id is provided, will be constructed as '{hf-user}/behavior1k-task{task_id:04d}'",
+    )
+    parser.add_argument(
+        "--task-id",
+        type=int,
+        required=False,
+        help="Task index to upload (e.g., 0, 1, 2, ...). When provided, folder-path is constructed from root-path.",
+    )
+    parser.add_argument(
+        "--root-path",
+        type=str,
+        required=False,
+        help="Root path containing task folders (e.g., /fsx/user/behavior1k-v3). Used with --task-id to construct folder path.",
+    )
+    parser.add_argument(
+        "--hf-user",
+        type=str,
+        default=None,
+        help="Hugging Face username for constructing repo-id with task-id (default: from HF_USER env var or 'fracapuano')",
+    )
+    parser.add_argument(
+        "--create-repo", action="store_true", help="Create the repository if it doesn't exist"
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=2,
+        help="Number of parallel workers for upload (default: 2). For I/O-bound uploads, use 1-4 to avoid network contention.",
+    )
+
+    args = parser.parse_args()
+
+    # Construct folder path and repo ID based on task-id or use provided values
+    if args.task_id is not None:
+        if not args.root_path:
+            raise ValueError("--root-path is required when --task-id is provided")
+
+        task_folder_name = f"behavior1k-task{args.task_id:04d}"
+        folder_path = Path(args.root_path) / task_folder_name
+        repo_id = f"{args.hf_user}/{task_folder_name}"
+
+        print(f"Task mode: uploading task {args.task_id}")
+    else:
+        if not args.folder_path or not args.repo_id:
+            raise ValueError(
+                "Either --task-id with --root-path, or both --folder-path and --repo-id must be provided"
+            )
+
+        folder_path = Path(args.folder_path)
+        repo_id = args.repo_id
+
+    # Validate folder path
+    if not folder_path.exists():
+        raise ValueError(f"Folder path does not exist: {folder_path}")
+    if not folder_path.is_dir():
+        raise ValueError(f"Path is not a directory: {folder_path}")
+
+    print(f"Uploading folder: {folder_path}")
+    print(f"Repository: {repo_id}")
+
+    # Create repository if requested
+    if args.create_repo:
+        api = HfApi()
+        print(f"Creating repository {repo_id}...")
+        try:
+            api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
+            print("Repository created or already exists. Updating its contents")
+        except Exception as e:
+            print(f"Warning: Could not create repository: {e}")
+
+    # Upload the folder
+    print(f"Starting upload with {args.num_workers} parallel workers...")
+    try:
+        result = upload_large_folder(
+            folder_path=str(folder_path),
+            repo_id=repo_id,
+            repo_type="dataset",
+            num_workers=args.num_workers,
+        )
+        print("✓ Upload completed successfully!")
+        print(f"Commit URL: {result}")
+    except Exception as e:
+        print(f"✗ Upload failed: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    main()

examples/port_datasets/slurm_aggregate_shards.py

@@ -15,16 +15,12 @@
 # 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_datasets.droid_rlds.port_droid import DROID_SHARDS
-
-from lerobot.datasets.aggregate import aggregate_datasets
-from lerobot.utils.utils import init_logging
+from port_droid import DROID_SHARDS
 
 
 class AggregateDatasets(PipelineStep):
@@ -38,6 +34,11 @@ 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_datasets.droid_rlds.port_droid import DROID_SHARDS
+from 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_datasets.droid_rlds.port_droid import port_droid, validate_dataset
+        from 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_datasets.droid_rlds.port_droid import DROID_SHARDS
+from 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, slurm=True
+    repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, private=False, slurm=True
 ):
     kwargs = {
         "pipeline": [
-            UploadDataset(repo_id),
+            UploadDataset(repo_id, private=private),
         ],
         "logging_dir": str(logs_dir / job_name),
     }
@@ -267,6 +267,12 @@ 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/PROFILING_GUIDE.md

@@ -1,263 +0,0 @@
-# RTC Profiling Guide
-
-This guide explains how to profile RTC (Real-Time Chunking) performance to identify bottlenecks and understand why RTC might be slower than expected.
-
-## Quick Start
-
-### 1. Profile with Real Robot (Profiled Version)
-
-Use `eval_with_real_robot_profiled.py` to profile actual robot execution:
-
-```bash
-# With RTC enabled
-uv run examples/rtc/eval_with_real_robot_profiled.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=30
-
-# Without RTC for comparison
-uv run examples/rtc/eval_with_real_robot_profiled.py \
-    --policy.path=helper2424/pi05_check_rtc \
-    --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: 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=30
-```
-
-**Output**: At the end of execution, you'll see a detailed breakdown of timing for each component:
-- `get_actions.policy_inference` - Time spent in policy inference
-- `get_actions.preprocessing` - Time spent preprocessing observations
-- `get_actions.postprocessing` - Time spent postprocessing actions
-- `get_actions.action_queue_merge` - Time spent merging actions with RTC
-- `robot.get_observation` - Time to get observations from robot
-- `robot.send_action` - Time to send actions to robot
-- And more...
-
-### 2. Profile Without Robot (Comparison Script)
-
-Use `profile_rtc_comparison.py` to profile just the policy inference without needing a robot:
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20
-```
-
-**Output**: Side-by-side comparison of performance with and without RTC, including:
-- Mean/min/max inference times
-- Throughput (iterations per second)
-- Verdict on whether RTC is faster or slower
-
-### 3. Enable Detailed Method-Level Profiling
-
-For even more granular profiling, add the `--enable_detailed_profiling` flag:
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20 \
-    --enable_detailed_profiling
-```
-
-This will show timing for individual methods within the policy.
-
-## Understanding the Output
-
-### Key Metrics to Look At
-
-1. **get_actions.policy_inference** - This should be the largest component
-   - If RTC is enabled, this includes the RTC guidance overhead
-   - Compare this with/without RTC to see the overhead
-
-2. **get_actions.preprocessing** - Image preprocessing and normalization
-   - Should be relatively fast
-   - If slow, consider optimizing image processing
-
-3. **get_actions.postprocessing** - Action denormalization
-   - Should be minimal
-   - If slow, check postprocessor implementation
-
-4. **get_actions.action_queue_merge** - RTC-specific merging logic
-   - Only present when RTC is enabled
-   - If this is taking significant time, the RTC algorithm may need optimization
-
-5. **robot.get_observation** - Robot communication overhead
-   - If slow, check camera/sensor latency
-   - Consider reducing image resolution
-
-6. **robot.send_action** - Action execution overhead
-   - Should be very fast
-   - If slow, check robot communication
-
-### Expected Performance
-
-For a typical Pi0 policy on Apple Silicon (MPS):
-- **Without RTC**: ~100-200ms per inference
-- **With RTC**: Should be similar or slightly faster due to action reuse
-- **Preprocessing**: ~5-20ms depending on number of cameras
-- **Postprocessing**: ~1-5ms
-
-If RTC is significantly slower, likely causes:
-1. **RTC overhead exceeds benefits** - The guidance computation is expensive
-2. **Execution horizon too small** - Not reusing enough actions to amortize overhead
-3. **No compilation** - Try with `--use_torch_compile`
-4. **Large prev_actions buffer** - Copying/processing previous actions is slow
-
-## Profiling Your Own Code
-
-### Using the Profiling Decorator
-
-Add profiling to your own methods:
-
-```python
-from lerobot.utils.profiling import profile_method, enable_profiling, print_profiling_summary
-
-# Enable profiling
-enable_profiling()
-
-# Decorate methods you want to profile
-@profile_method
-def my_slow_function(x):
-    # ... your code ...
-    return result
-
-# At end of execution
-print_profiling_summary()
-```
-
-### Using Profile Context Manager
-
-For profiling specific code blocks:
-
-```python
-from lerobot.utils.profiling import profile_section, enable_profiling
-
-enable_profiling()
-
-with profile_section("data_loading"):
-    data = load_data()
-
-with profile_section("model_inference"):
-    output = model(data)
-```
-
-### Adding Profiling to Policy Methods
-
-To profile specific parts of the Pi0 policy, you can add decorators:
-
-```python
-# In src/lerobot/policies/pi0/modeling_pi0.py
-from lerobot.utils.profiling import profile_method, profile_section
-
-class Pi0Policy:
-    @profile_method
-    def predict_action_chunk(self, obs, inference_delay=0, prev_chunk_left_over=None):
-        # ... existing code ...
-        pass
-
-    def _generate_actions_with_rtc(self, ...):
-        with profile_section("rtc.guidance_computation"):
-            # ... guidance code ...
-            pass
-        
-        with profile_section("rtc.action_merging"):
-            # ... merging code ...
-            pass
-```
-
-## Analyzing Results
-
-### Comparison Checklist
-
-When comparing RTC vs non-RTC performance, check:
-
-- [ ] Is `policy_inference` time higher with RTC?
-- [ ] Is `action_queue_merge` taking significant time?
-- [ ] Are you running enough iterations to amortize warmup?
-- [ ] Is torch.compile enabled for fair comparison?
-- [ ] Is the execution horizon large enough? (should be >= 10-20)
-- [ ] Are you testing on the same hardware/device?
-
-### Common Bottlenecks
-
-1. **Image preprocessing dominates** 
-   - Solution: Reduce image resolution, use fewer cameras, or optimize preprocessing
-
-2. **Action queue operations are slow**
-   - Solution: Review queue implementation, consider using ring buffer
-
-3. **RTC guidance is expensive**
-   - Solution: Reduce guidance weight, simplify guidance computation, use torch.compile
-
-4. **Robot communication is slow**
-   - Solution: Increase baud rate, reduce action frequency, optimize protocol
-
-5. **Memory allocation overhead**
-   - Solution: Pre-allocate buffers, reuse tensors, avoid unnecessary copies
-
-## Advanced: Adding Custom Metrics
-
-You can add custom timing metrics to the profiled script:
-
-```python
-from lerobot.utils.profiling import record_timing
-
-start = time.perf_counter()
-# ... your code ...
-duration = time.perf_counter() - start
-record_timing("my_custom_metric", duration)
-```
-
-## Troubleshooting
-
-### Profiling shows RTC is slower by >50%
-
-1. Check if torch.compile is enabled: `--use_torch_compile`
-2. Increase execution horizon: `--rtc.execution_horizon=30`
-3. Verify inference_delay is calculated correctly
-4. Profile with `--enable_detailed_profiling` to find exact bottleneck
-
-### Profiling output is empty
-
-1. Make sure profiling is enabled with `enable_profiling()`
-2. Verify you're running enough iterations (at least 10)
-3. Check that code is actually executing (not short-circuited)
-
-### Inconsistent results between runs
-
-1. Run more iterations: `--num_iterations=100`
-2. Increase warmup iterations
-3. Check for thermal throttling on device
-4. Ensure no other processes competing for resources
-
-## Next Steps
-
-1. Run both profiling scripts (with/without robot)
-2. Compare timing breakdowns
-3. Identify the largest bottleneck
-4. Focus optimization efforts on that component
-5. Re-run profiling to verify improvements
-
-## Questions?
-
-If profiling reveals unexpected bottlenecks or you need help interpreting results, please share:
-- The full profiling output
-- Your configuration (RTC enabled/disabled, execution horizon, etc.)
-- Hardware specs (device type, memory, etc.)
-- Policy type and size
-

examples/rtc/PROFILING_QUICK_START.md

@@ -1,208 +0,0 @@
-# RTC Profiling - Quick Start
-
-Quick reference for profiling Pi0 with RTC to identify performance bottlenecks.
-
-## 🚀 Quick Commands
-
-### 1. Profile with Real Robot
-
-```bash
-# With RTC enabled (profiled version)
-uv run examples/rtc/eval_with_real_robot_profiled.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.cameras="{ gripper: {type: opencv, index_or_path: 0}, front: {type: opencv, index_or_path: 1}}" \
-    --task="Pick up object" \
-    --duration=30
-```
-
-### 2. Compare RTC vs No-RTC (No Robot Needed)
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20
-```
-
-### 3. Detailed RTC Method Profiling
-
-```bash
-uv run examples/rtc/profile_pi0_rtc_detailed.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=20 \
-    --execution_horizon=20 \
-    --enable_rtc_profiling
-```
-
-## 📊 What Each Tool Does
-
-| Tool | Purpose | Needs Robot? |
-|------|---------|--------------|
-| `eval_with_real_robot_profiled.py` | Profile actual robot execution with RTC | ✅ Yes |
-| `profile_rtc_comparison.py` | Compare RTC vs no-RTC side-by-side | ❌ No |
-| `profile_pi0_rtc_detailed.py` | Deep dive into RTC internals | ❌ No |
-
-## 🔍 Key Metrics to Watch
-
-### Overall Performance
-- **iteration.policy_inference** - Total policy inference time
-- **iteration.preprocessing** - Image preprocessing time
-- **iteration.postprocessing** - Action denormalization time
-
-### RTC-Specific (with `--enable_rtc_profiling`)
-- **rtc.denoise_step.base_denoising** - Time without RTC overhead
-- **rtc.denoise_step.autograd_correction** - Gradient computation time
-- **rtc.denoise_step.guidance_computation** - Total RTC guidance overhead
-
-### Robot Communication
-- **robot.get_observation** - Time to get robot state
-- **robot.send_action** - Time to send action command
-
-## 🎯 Quick Diagnosis
-
-### RTC is slower than expected?
-
-1. **Check if torch.compile is enabled**
-   ```bash
-   # Add this flag
-   --use_torch_compile
-   ```
-
-2. **Try larger execution horizon**
-   ```bash
-   # Increase to amortize RTC overhead
-   --rtc.execution_horizon=30
-   ```
-
-3. **Profile to find bottleneck**
-   ```bash
-   uv run examples/rtc/profile_pi0_rtc_detailed.py \
-       --policy_path=helper2424/pi05_check_rtc \
-       --device=mps \
-       --enable_rtc_profiling
-   ```
-
-### Preprocessing is slow?
-
-- Reduce image resolution in robot config
-- Use fewer cameras
-- Check camera FPS settings
-
-### Policy inference is slow?
-
-- Enable torch.compile
-- Check device (MPS vs CUDA vs CPU)
-- Try smaller model if available
-
-## 📈 Expected Performance
-
-### Typical timings on Apple Silicon (MPS):
-
-| Component | Time (ms) | Notes |
-|-----------|-----------|-------|
-| Policy inference | 100-200 | Depends on model size |
-| Preprocessing | 5-20 | Depends on #cameras |
-| Postprocessing | 1-5 | Usually fast |
-| RTC overhead | 10-50 | Should be < 50% of base |
-
-### When RTC helps:
-- ✅ Execution horizon ≥ 10
-- ✅ Inference time > action execution rate
-- ✅ Using torch.compile
-- ✅ Proper inference_delay calculation
-
-### When RTC might not help:
-- ❌ Very fast inference already
-- ❌ Small execution horizon (< 5)
-- ❌ No compilation (interpreted mode)
-- ❌ Inference delay not accounted for
-
-## 🛠️ Adding Profiling to Your Code
-
-### Quick snippet:
-
-```python
-from lerobot.utils.profiling import enable_profiling, print_profiling_summary, profile_section
-
-# Enable at start
-enable_profiling()
-
-# Profile sections
-with profile_section("my_operation"):
-    # ... your code ...
-    pass
-
-# Print at end
-print_profiling_summary()
-```
-
-### Profile specific methods:
-
-```python
-from lerobot.utils.profiling import profile_method
-
-@profile_method
-def my_slow_function():
-    # ... your code ...
-    pass
-```
-
-## 📝 Example Output
-
-```
-PROFILING SUMMARY
-================================================================================
-Function                                                    Count    Mean (ms)
---------------------------------------------------------------------------------
-iteration.policy_inference                                    20       150.23
-iteration.preprocessing                                       20        12.45
-rtc.denoise_step.guidance_computation                        200        15.67
-rtc.denoise_step.autograd_correction                         200         8.23
-rtc.denoise_step.base_denoising                             200       120.45
-================================================================================
-```
-
-## 🚨 Common Issues
-
-### "No profiling data available"
-- Did you call `enable_profiling()`?
-- Running enough iterations?
-
-### Inconsistent results
-- Increase `--num_iterations`
-- Check for thermal throttling
-- Close other applications
-
-### Can't find bottleneck
-- Enable `--enable_rtc_profiling` for detailed breakdown
-- Check both preprocessing and inference
-- Compare with and without RTC
-
-## 📖 More Details
-
-See `PROFILING_GUIDE.md` for comprehensive documentation.
-
-## 🤔 Still Slow?
-
-1. Run comparison: `profile_rtc_comparison.py`
-2. Run detailed profiling: `profile_pi0_rtc_detailed.py --enable_rtc_profiling`
-3. Share output for help (include device, model, settings)
-
-## ✅ Quick Checklist
-
-Before asking for help, verify:
-
-- [ ] Ran comparison script (with/without RTC)
-- [ ] Tried torch.compile
-- [ ] Tested different execution horizons (10, 20, 30)
-- [ ] Profiled with detailed RTC profiling
-- [ ] Checked preprocessing vs inference split
-- [ ] Verified hardware (device type, thermal state)
-

examples/rtc/PROFILING_README.md

@@ -1,352 +0,0 @@
-# RTC Profiling Toolkit
-
-Complete toolkit for profiling Pi0 with RTC to identify performance bottlenecks.
-
-## 📦 What's Included
-
-### Scripts
-
-1. **`eval_with_real_robot_profiled.py`**
-   - Profiled version of the real robot eval script
-   - Adds timing measurements throughout execution
-   - Works with actual robot hardware
-   - Same usage as original but with profiling output
-
-2. **`profile_rtc_comparison.py`**
-   - Side-by-side comparison of RTC vs no-RTC
-   - No robot needed (uses mock observations)
-   - Shows clear verdict on whether RTC is helping
-   - Great for quick performance checks
-
-3. **`profile_pi0_rtc_detailed.py`**
-   - Most detailed profiling available
-   - Can enable RTC method-level profiling
-   - Provides insights and recommendations
-   - Perfect for deep-dive investigations
-
-4. **`add_rtc_profiling.py`**
-   - Monkey-patching utility for RTC internals
-   - Profiles individual RTC operations
-   - Can be applied without modifying source
-   - Shows exactly where RTC spends time
-
-### Utilities
-
-5. **`src/lerobot/utils/profiling.py`**
-   - Core profiling utilities
-   - Decorators for method profiling
-   - Context managers for code blocks
-   - Statistics collection and reporting
-
-### Documentation
-
-6. **`PROFILING_GUIDE.md`** - Comprehensive guide
-7. **`PROFILING_QUICK_START.md`** - Quick reference
-
-## 🚀 Quick Start
-
-### Step 1: Compare Performance
-
-Run this first to see if RTC is actually slower:
-
-```bash
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20
-```
-
-**Expected output:**
-```
-COMPARISON SUMMARY
-================================================================================
-Metric                         Without RTC        With RTC      Difference
---------------------------------------------------------------------------------
-Mean time (ms)                       150.23         165.45          +15.22
-Throughput (iter/s)                    6.66           6.05           -0.61
-================================================================================
-VERDICT
-✗ RTC is SLOWER by 10.1%
-  Mean time increased by 15.22 ms
-  
-  Possible reasons:
-  - RTC overhead exceeds benefits at current execution horizon
-  - No torch.compile enabled
-```
-
-### Step 2: Identify Bottleneck
-
-If RTC is slower, find out why:
-
-```bash
-uv run examples/rtc/profile_pi0_rtc_detailed.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=20 \
-    --execution_horizon=20 \
-    --enable_rtc_profiling
-```
-
-**Expected output:**
-```
-PROFILING SUMMARY
-================================================================================
-Function                                             Count    Mean (ms)    Total (s)
-------------------------------------------------------------------------------------
-iteration.policy_inference                              20      150.23         3.00
-rtc.denoise_step.guidance_computation                  200       15.67         3.13
-rtc.denoise_step.autograd_correction                   200        8.23         1.65
-iteration.preprocessing                                 20       12.45         0.25
-================================================================================
-
-KEY INSIGHTS
-================================================================================
-Time breakdown:
-  Policy inference:  150.23 ms (87.2%)
-  Preprocessing:     12.45 ms (7.2%)
-  Postprocessing:    2.10 ms (1.2%)
-
-RTC breakdown:
-  Base denoising:    120.45 ms
-  Guidance compute:  15.67 ms
-  Autograd correct:  8.23 ms
-  RTC overhead:      23.90 ms (19.8% of base)
-
-Recommendations:
-  ⚠ RTC autograd overhead is significant
-    → This is expected, but consider increasing execution_horizon
-    → Try torch.compile if not already enabled
-  💡 torch.compile not enabled
-    → Try --use_torch_compile for potential speedup
-================================================================================
-```
-
-### Step 3: Try Optimizations
-
-Based on recommendations:
-
-```bash
-# Try with torch.compile
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=20 \
-    --use_torch_compile
-
-# Try larger execution horizon
-uv run examples/rtc/profile_rtc_comparison.py \
-    --policy_path=helper2424/pi05_check_rtc \
-    --device=mps \
-    --num_iterations=50 \
-    --execution_horizon=30
-```
-
-### Step 4: Profile Real Robot (Optional)
-
-Test with actual hardware:
-
-```bash
-uv run examples/rtc/eval_with_real_robot_profiled.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.cameras="{...}" \
-    --task="Pick up object" \
-    --duration=30
-```
-
-## 🎯 Common Scenarios
-
-### "RTC is 2x slower!"
-
-This usually means:
-- RTC overhead is high but not getting benefits
-- Need to enable torch.compile
-- Execution horizon too small
-- Inference delay not calculated correctly
-
-**Try:**
-1. `--use_torch_compile`
-2. Increase `--execution_horizon` to 30+
-3. Check inference_delay calculation
-
-### "RTC is only slightly slower"
-
-This is expected! RTC overhead is about 10-30% typically.
-The benefit comes during **execution**, not single inference:
-- Actions are reused across chunks
-- Overall system latency is reduced
-- Robot gets smoother actions
-
-### "Want to optimize specific part"
-
-Use the profiling utilities:
-
-```python
-from lerobot.utils.profiling import enable_profiling, profile_section, print_profiling_summary
-
-enable_profiling()
-
-with profile_section("my_custom_operation"):
-    # Your code here
-    pass
-
-print_profiling_summary()
-```
-
-## 📊 Understanding Results
-
-### Key Metrics
-
-**Policy Inference Time**
-- Time for forward pass through model
-- Should be largest component (70-90%)
-- Includes RTC guidance if enabled
-
-**Preprocessing Time**
-- Image normalization, resizing
-- Should be < 20% of total
-- If high: reduce image resolution
-
-**RTC Guidance Overhead**
-- Extra time for RTC guidance computation
-- Typically 10-30% of base inference
-- If > 50%: RTC may not be beneficial at current settings
-
-**Autograd Correction**
-- Time computing gradients for RTC
-- Usually 5-15% of base inference
-- Can be reduced with torch.compile
-
-### Expected Ranges (Apple Silicon MPS)
-
-| Metric | Good | Acceptable | Poor |
-|--------|------|------------|------|
-| Policy inference | 100-150ms | 150-250ms | >250ms |
-| Preprocessing | <20ms | 20-50ms | >50ms |
-| RTC overhead | 10-30% | 30-50% | >50% |
-
-## 🔧 Optimization Guide
-
-### If RTC overhead is too high:
-
-1. **Enable compilation:**
-   ```bash
-   --use_torch_compile
-   ```
-   Expected improvement: 20-40% faster
-
-2. **Increase execution horizon:**
-   ```bash
-   --execution_horizon=30  # or higher
-   ```
-   Amortizes RTC cost over more actions
-
-3. **Check guidance weight:**
-   ```python
-   # In config
-   rtc.max_guidance_weight=1.0  # try 0.5 for less overhead
-   ```
-
-### If preprocessing is slow:
-
-1. **Reduce image resolution:**
-   ```python
-   # In robot config
-   cameras={
-       "gripper": {"width": 320, "height": 240}  # instead of 640x480
-   }
-   ```
-
-2. **Use fewer cameras:**
-   - Profile which cameras are essential
-   - Remove unnecessary views
-
-### If inference is generally slow:
-
-1. Use torch.compile (if not already)
-2. Check device is correct (MPS vs CUDA)
-3. Verify model is in eval mode
-4. Check for unnecessary gradient tracking
-
-## 🐛 Troubleshooting
-
-### Empty profiling output
-```python
-# Make sure to enable profiling!
-from lerobot.utils.profiling import enable_profiling
-enable_profiling()
-```
-
-### Inconsistent timings
-- Run more iterations (50-100)
-- Check thermal throttling
-- Close background apps
-- Use `--warmup_iterations=10`
-
-### Can't find bottleneck
-1. Start with `profile_rtc_comparison.py`
-2. Then run `profile_pi0_rtc_detailed.py --enable_rtc_profiling`
-3. Compare with/without RTC
-4. Check each component separately
-
-## 📖 Full Documentation
-
-- **`PROFILING_GUIDE.md`** - Complete reference with examples
-- **`PROFILING_QUICK_START.md`** - Quick commands and tips
-
-## 🤝 Getting Help
-
-If you're still experiencing issues:
-
-1. Run comparison script and save output
-2. Run detailed profiling and save output
-3. Include:
-   - Policy path
-   - Device type
-   - RTC settings (execution_horizon, etc.)
-   - Hardware specs
-   - Full profiling output
-
-## 🎓 Learning More
-
-### Profiling your own code:
-
-```python
-from lerobot.utils.profiling import profile_method, enable_profiling
-
-enable_profiling()
-
-@profile_method
-def my_function():
-    # Automatically profiled
-    pass
-```
-
-### RTC internals:
-
-```python
-from examples.rtc.add_rtc_profiling import monkey_patch_rtc_profiling
-
-enable_profiling()
-monkey_patch_rtc_profiling()
-
-# Now RTC methods are profiled
-policy.predict_action_chunk(...)
-```
-
-## ✨ Next Steps
-
-1. Run `profile_rtc_comparison.py` to establish baseline
-2. Use `profile_pi0_rtc_detailed.py` to find bottlenecks
-3. Apply optimizations (torch.compile, larger horizon)
-4. Re-run comparison to verify improvements
-5. Test with real robot using profiled version
-
-Happy profiling! 🚀
-

examples/rtc/README.md

@@ -1,251 +0,0 @@
-# Real-Time Chunking (RTC) Examples
-
-This directory contains examples and evaluation scripts for Real-Time Chunking (RTC), a technique for improving action chunking policies in real-time robot control.
-
-## Overview
-
-Real-Time Chunking addresses the challenge of maintaining consistency and reactivity when using action chunking policies with non-negligible inference latency. It uses a guidance technique during diffusion sampling to blend new action predictions with previously planned actions.
-
-**Key Benefits:**
-
-- Maintains consistency between consecutive action chunks
-- Reduces jitter and improves smoothness
-- Adapts to inference delays dynamically
-
-**Reference:** [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/download/real_time_chunking.pdf)
-
-## Scripts
-
-### 1. `eval_dataset.py`
-
-Offline evaluation on dataset samples with detailed visualization and validation.
-
-**Features:**
-
-- Compare RTC vs non-RTC predictions on two random dataset samples
-- Validate RTC behavior (delay region, blend region, post-horizon region)
-- Generate debug visualizations:
-  - Denoising step comparisons (x_t, v_t, x1_t, corrections)
-  - Final action predictions comparison
-- Support for torch.compile() optimization
-- Memory-efficient sequential policy loading for large models
-
-**Usage:**
-
-```bash
-# 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 \
-    --seed=10
-
-# With Pi0.5 policy on CUDA
-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
-
-# With Pi0 policy
-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
-
-# With torch.compile for faster inference
-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=max-autotune
-
-# 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
-```
-
-**Key Parameters:**
-
-- `--policy.path`: Path to pretrained policy
-- `--dataset.repo_id`: Dataset to evaluate on
-- `--rtc.execution_horizon`: Number of steps to maintain consistency (default: 20)
-- `--rtc.max_guidance_weight`: Maximum guidance weight (default: 10.0)
-- `--rtc.prefix_attention_schedule`: Schedule type (ZEROS, ONES, LINEAR, EXP)
-- `--inference_delay`: Inference delay for RTC (default: 4)
-- `--seed`: Random seed for reproducibility (default: 42)
-- `--output_dir`: Directory to save visualizations (default: rtc_debug_output)
-- `--device`: Device to use (cuda, cpu, mps, auto)
-- `--use_torch_compile`: Enable torch.compile() for faster inference
-
-**Output:**
-
-The script generates several visualization files in `rtc_debug_output/`:
-
-- `denoising_xt_comparison.png` - Noisy state evolution during denoising
-- `denoising_vt_comparison.png` - Velocity predictions during denoising
-- `denoising_x1t_comparison.png` - Predicted final states during denoising
-- `denoising_correction_comparison.png` - RTC guidance corrections applied
-- `final_actions_comparison.png` - Final action predictions (prev_chunk, no_rtc, rtc)
-
-The script also validates RTC behavior and reports:
-
-- ✅ Delay region [0:inference_delay]: RTC = prev_chunk
-- ✅ Blend region [inference_delay:execution_horizon]: prev_chunk ≤ RTC ≤ no_rtc
-- ✅ Post-horizon [execution_horizon:]: RTC = no_rtc
-
-### 2. `eval_with_real_robot.py`
-
-Real-time evaluation on physical robots or simulation environments.
-
-**Features:**
-
-- Run policy with RTC on real robot or simulation
-- Multi-threaded action execution and inference
-- Action queue management with proper timing
-- Latency tracking and adaptive inference delay
-- Support for both robots and gym environments
-- Support for torch.compile() optimization
-
-**Usage:**
-
-```bash
-# With real robot
-uv run python examples/rtc/eval_with_real_robot.py \
-    --policy.path=lerobot/smolvla_base \
-    --robot.type=so100 \
-    --task="pick up the cup" \
-    --duration=30.0
-
-# With simulation environment
-uv run python examples/rtc/eval_with_real_robot.py \
-    --policy.path=lerobot/smolvla_base \
-    --env.type=pusht \
-    --duration=60.0
-
-# With policy compilation (CUDA only, not MPS)
-uv run python examples/rtc/eval_with_real_robot.py \
-    --policy.path=lerobot/smolvla_base \
-    --robot.type=so100 \
-    --use_torch_compile=true \
-    --torch_compile_mode=max-autotune
-```
-
-**Key Parameters:**
-
-- `--policy.path`: Path to pretrained policy
-- `--robot.type` or `--env.type`: Robot or environment to use
-- `--task`: Task description (for VLA models)
-- `--rtc.execution_horizon`: Number of steps to maintain consistency (default: 10)
-- `--rtc.max_guidance_weight`: Maximum guidance weight (default: 1.0)
-- `--rtc.prefix_attention_schedule`: Schedule type (ZEROS, ONES, LINEAR, EXP)
-- `--duration`: How long to run (seconds, default: 30.0)
-- `--fps`: Action execution frequency (Hz, default: 10.0)
-- `--action_queue_size_to_get_new_actions`: Queue size threshold to request new actions (default: 30)
-- `--device`: Device to use (cuda, cpu, mps, auto)
-- `--use_torch_compile`: Enable torch.compile() for faster inference
-
-## Understanding RTC Parameters
-
-### `execution_horizon`
-
-Number of timesteps from previous chunk to maintain consistency with. Higher values mean more consistency but potentially less reactivity.
-
-**Typical values:** 8-12 steps for dataset evaluation, 10 steps for real-time execution
-
-### `max_guidance_weight`
-
-Upper bound on guidance strength. Higher values give stronger consistency but may over-constrain new predictions.
-
-**Typical values:**
-
-- Dataset evaluation: 10.0-100.0 (can be higher for analysis)
-- Real-time execution: 1.0-10.0 (more conservative)
-
-### `prefix_attention_schedule`
-
-How to weight consistency across the overlap region:
-
-- `ZEROS`: Binary (full weight up to inference_delay, then zero)
-- `ONES`: Full weight across entire execution_horizon
-- `LINEAR`: Linear decay from inference_delay to execution_horizon
-- `EXP`: Exponential decay (recommended)
-
-**Recommended:** `EXP`
-
-### `inference_delay`
-
-Number of timesteps from the prefix to use for guidance. Typically calculated dynamically based on inference latency in real-time execution, but fixed for dataset evaluation.
-
-**Typical values:** 3-5 steps for dataset evaluation
-
-### `action_queue_size_to_get_new_actions` (real-time only)
-
-Threshold for requesting new action chunks. Should be higher than `inference_delay + execution_horizon` to ensure smooth operation.
-
-**Typical values:** 20-30 steps
-
-## Validation Rules (Dataset Evaluation)
-
-The dataset evaluation script validates that RTC behavior matches expectations:
-
-1. **Delay Region [0:inference_delay]**: RTC actions should equal previous chunk
-   - Ensures consistency during the inference delay period
-
-2. **Blend Region [inference_delay:execution_horizon]**: RTC should be between prev_chunk and no_rtc
-   - Smooth transition from previous plan to new predictions
-
-3. **Post-Horizon [execution_horizon:]**: RTC should equal no_rtc
-   - Full adoption of new predictions after execution horizon
-
-## Tips
-
-1. **Start with dataset evaluation** (`eval_dataset.py`) to understand RTC behavior and tune parameters before running on robot
-2. **Use visualizations** to debug unexpected behavior - check denoising steps and final actions
-3. **Tune execution_horizon** based on your inference latency and action frequency
-4. **Monitor validation output** - failures indicate potential implementation issues or misconfigured parameters
-5. **Compare different schedules** - EXP usually works best but LINEAR can be more interpretable
-
-## Troubleshooting
-
-### Validation fails in delay region
-
-- Check that `prev_chunk_left_over` is properly passed to the policy
-- Verify RTC guidance is being applied during denoising
-- Look at denoising visualizations to see where guidance diverges
-
-### Validation fails in post-horizon region
-
-- RTC and no_rtc use different noise - verify same noise is being used for comparison
-- Check that weights are correctly zeroed out after execution horizon
-- Review prefix_attention_schedule visualization
-
-### Poor performance on real robot
-
-- Increase `action_queue_size_to_get_new_actions` if you see warnings
-- Reduce `max_guidance_weight` if robot is too conservative
-- Try different `prefix_attention_schedule` values
-- Enable torch.compile() for faster inference (CUDA only)
-
-### Memory issues with large models
-
-- The dataset evaluation script loads policies sequentially to minimize memory
-- For real-time execution, only one policy is loaded
-- Use smaller batch sizes if needed
-
-## Related Documentation
-
-- [RTC Implementation](../../src/lerobot/policies/rtc/modeling_rtc.py)
-- [RTC Configuration](../../src/lerobot/policies/rtc/configuration_rtc.py)
-- [Action Queue](../../src/lerobot/policies/rtc/action_queue.py)
-- [Physical Intelligence Paper](https://www.physicalintelligence.company/download/real_time_chunking.pdf)

examples/rtc/add_rtc_profiling.py

@@ -1,202 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Script to add profiling instrumentation to RTCProcessor.
-
-This script shows which methods to profile in the RTC code to identify bottlenecks.
-You can either:
-1. Apply these changes directly to modeling_rtc.py
-2. Use monkey patching to add profiling without modifying source
-3. Use as reference for manual instrumentation
-
-Usage:
-    # Option 1: Monkey patch (no source changes)
-    python examples/rtc/add_rtc_profiling.py
-
-    # Option 2: Apply changes to source
-    # Copy the profiled methods below into src/lerobot/policies/rtc/modeling_rtc.py
-"""
-
-import logging
-
-import torch
-from torch import Tensor
-
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
-from lerobot.utils.profiling import ProfileContext, enable_profiling, is_profiling_enabled
-
-logger = logging.getLogger(__name__)
-
-
-def profile_denoise_step(self, x_t, prev_chunk_left_over, inference_delay, time, original_denoise_step_partial, execution_horizon=None) -> Tensor:
-    """Profiled version of denoise_step."""
-    
-    if not is_profiling_enabled():
-        # Call original implementation if profiling disabled
-        return self._original_denoise_step(x_t, prev_chunk_left_over, inference_delay, time, original_denoise_step_partial, execution_horizon)
-    
-    with ProfileContext("rtc.denoise_step.total"):
-        # 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
-            with ProfileContext("rtc.denoise_step.base_denoising"):
-                v_t = original_denoise_step_partial(x_t)
-            return v_t
-
-        with ProfileContext("rtc.denoise_step.setup"):
-            x_t = x_t.clone().detach()
-
-            squeezed = False
-            if len(x_t.shape) < 3:
-                x_t = x_t.unsqueeze(0)
-                squeezed = True
-
-            if len(prev_chunk_left_over.shape) < 3:
-                prev_chunk_left_over = prev_chunk_left_over.unsqueeze(0)
-
-            if execution_horizon is None:
-                execution_horizon = self.rtc_config.execution_horizon
-
-            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]
-
-        # Padding
-        with ProfileContext("rtc.denoise_step.padding"):
-            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
-
-        # Get prefix weights
-        with ProfileContext("rtc.denoise_step.get_prefix_weights"):
-            weights = (
-                self.get_prefix_weights(inference_delay, execution_horizon, action_chunk_size)
-                .to(x_t.device)
-                .unsqueeze(0)
-                .unsqueeze(-1)
-            )
-
-        # Main RTC guidance computation
-        with ProfileContext("rtc.denoise_step.guidance_computation"):
-            with torch.enable_grad():
-                # Base denoising
-                with ProfileContext("rtc.denoise_step.base_denoising"):
-                    v_t = original_denoise_step_partial(x_t)
-                
-                x_t.requires_grad_(True)
-
-                # Compute x1_t
-                with ProfileContext("rtc.denoise_step.compute_x1_t"):
-                    x1_t = x_t - time * v_t
-
-                # Compute error
-                with ProfileContext("rtc.denoise_step.compute_error"):
-                    err = (prev_chunk_left_over - x1_t) * weights
-                    grad_outputs = err.clone().detach()
-
-                # Compute correction via autograd
-                with ProfileContext("rtc.denoise_step.autograd_correction"):
-                    correction = torch.autograd.grad(x1_t, x_t, grad_outputs, retain_graph=False)[0]
-
-        # Compute guidance weight
-        with ProfileContext("rtc.denoise_step.compute_guidance_weight"):
-            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)
-
-        # Apply guidance
-        with ProfileContext("rtc.denoise_step.apply_guidance"):
-            result = v_t - guidance_weight * correction
-
-        # Cleanup
-        with ProfileContext("rtc.denoise_step.cleanup"):
-            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 monkey_patch_rtc_profiling():
-    """Apply profiling to RTCProcessor via monkey patching.
-    
-    This modifies the RTCProcessor class at runtime to add profiling
-    without changing source files.
-    """
-    logger.info("Applying RTC profiling monkey patch...")
-    
-    # Save original method
-    RTCProcessor._original_denoise_step = RTCProcessor.denoise_step
-    
-    # Replace with profiled version
-    RTCProcessor.denoise_step = profile_denoise_step
-    
-    logger.info("✓ RTC profiling enabled")
-
-
-def print_usage():
-    """Print usage instructions."""
-    print("\n" + "="*80)
-    print("RTC PROFILING INSTRUMENTATION")
-    print("="*80)
-    print("\nThis script provides profiling for RTCProcessor methods.")
-    print("\nOption 1: Monkey Patch (Recommended)")
-    print("-" * 40)
-    print("Add to your script:")
-    print("""
-    from lerobot.utils.profiling import enable_profiling, print_profiling_summary
-    from examples.rtc.add_rtc_profiling import monkey_patch_rtc_profiling
-
-    # Enable profiling
-    enable_profiling()
-    monkey_patch_rtc_profiling()
-
-    # ... run your code ...
-
-    # Print results
-    print_profiling_summary()
-    """)
-    
-    print("\nOption 2: Manual Source Modification")
-    print("-" * 40)
-    print("1. Copy profile_denoise_step() from this file")
-    print("2. Replace denoise_step() in src/lerobot/policies/rtc/modeling_rtc.py")
-    print("3. Add profiling imports at top of file")
-    
-    print("\nKey Metrics to Watch:")
-    print("-" * 40)
-    print("- rtc.denoise_step.base_denoising     - Time for base policy inference")
-    print("- rtc.denoise_step.autograd_correction - Time computing gradients")
-    print("- rtc.denoise_step.guidance_computation - Total guidance overhead")
-    print("- rtc.denoise_step.get_prefix_weights  - Time computing weights")
-    print("="*80 + "\n")
-
-
-if __name__ == "__main__":
-    print_usage()
-

examples/rtc/eval_dataset.py

@@ -39,8 +39,9 @@ Usage:
     uv run python examples/rtc/eval_dataset.py \
         --policy.path=lerobot/pi05_libero_finetuned \
         --dataset.repo_id=HuggingFaceVLA/libero \
-        --rtc.execution_horizon=8 \
+        --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 \
@@ -141,7 +142,7 @@ def _check_matplotlib_available():
         raise ImportError(
             "matplotlib is required for RTC debug visualizations. "
             "Please install it by running:\n"
-            "  uv pip install -e '.[matplotlib-dep]'"
+            "  uv pip install matplotlib"
         )
 
 
@@ -543,11 +544,6 @@ class RTCEvaluator:
         logging.info("Plotting results...")
         self.plot_tracked_data(rtc_tracked_steps, no_rtc_tracked_steps, prev_chunk_left_over, num_steps)
 
-        # Validate RTC behavior
-        # logging.info("=" * 80)
-        # logging.info("Validating RTC behavior...")
-        # self.validate_rtc_behavior(rtc_actions, no_rtc_actions, prev_chunk_left_over)
-
         # Plot final actions comparison
         logging.info("=" * 80)
         logging.info("Plotting final actions comparison...")
@@ -556,159 +552,6 @@ class RTCEvaluator:
         logging.info("=" * 80)
         logging.info("Evaluation completed successfully")
 
-    def validate_rtc_behavior(self, rtc_actions, no_rtc_actions, prev_chunk_left_over):
-        """Validate RTC behavior by comparing final action predictions with expected values.
-
-        Validation rules:
-        1. During delay [0:inference_delay]: RTC should equal prev_chunk
-        2. After delay, within execution horizon [inference_delay:execution_horizon]:
-           RTC should be between prev_chunk and no_rtc
-        3. After execution horizon [execution_horizon:]: RTC should equal no_rtc
-
-        Args:
-            rtc_actions: Final actions from RTC policy (batch, time, action_dim)
-            no_rtc_actions: Final actions from non-RTC policy (batch, time, action_dim)
-            prev_chunk_left_over: Previous chunk used as ground truth (time, action_dim)
-        """
-        # Remove batch dimension if present and move to CPU
-        rtc_actions_t = rtc_actions.squeeze(0).cpu() if len(rtc_actions.shape) == 3 else rtc_actions.cpu()
-        no_rtc_actions_t = (
-            no_rtc_actions.squeeze(0).cpu() if len(no_rtc_actions.shape) == 3 else no_rtc_actions.cpu()
-        )
-        prev_chunk = prev_chunk_left_over.cpu()
-
-        logging.info(f"  rtc_actions shape: {rtc_actions_t.shape}")
-        logging.info(f"  no_rtc_actions shape: {no_rtc_actions_t.shape}")
-        logging.info(f"  prev_chunk shape: {prev_chunk.shape}")
-
-        # Determine chunk length for comparison
-        chunk_len = min(rtc_actions_t.shape[0], no_rtc_actions_t.shape[0], prev_chunk.shape[0])
-        inference_delay = self.cfg.inference_delay
-        execution_horizon = self.cfg.rtc.execution_horizon
-
-        # Tolerance for floating point comparison
-        rtol = 1e-2  # Relative tolerance
-
-        validation_passed = True
-        warnings = []
-
-        logging.info("  Validating RTC behavior:")
-        logging.info(f"    Chunk length: {chunk_len}")
-        logging.info(f"    Inference delay: {inference_delay}")
-        logging.info(f"    Execution horizon: {execution_horizon}")
-        logging.info(f"    Tolerance: rtol={rtol}")
-
-        # ============================================================================
-        # Rule 1: During delay [0:inference_delay], RTC should equal prev_chunk
-        # ============================================================================
-        if inference_delay > 0:
-            delay_end = min(inference_delay, chunk_len)
-            rtc_delay = rtc_actions_t[:delay_end]
-            prev_delay = prev_chunk[:delay_end]
-
-            logging.info(f"  rtc_delay: {rtc_delay.shape}")
-            logging.info(f"  prev_delay: {prev_delay.shape}")
-
-            if not torch.allclose(rtc_delay, prev_delay, rtol=rtol):
-                max_diff = torch.max(torch.abs(rtc_delay - prev_delay)).item()
-                mean_diff = torch.mean(torch.abs(rtc_delay - prev_delay)).item()
-                logging.info(f"  rtc_delay: {rtc_delay}")
-                logging.info(f"  prev_delay: {prev_delay}")
-                logging.info(f"  max_diff: {max_diff}")
-                logging.info(f"  mean_diff: {mean_diff}")
-                warnings.append(
-                    f"    ⚠ VALIDATION FAILED: During delay [0:{delay_end}], "
-                    f"RTC does NOT equal prev_chunk!\n"
-                    f"      Max difference: {max_diff:.6f}\n"
-                    f"      Mean difference: {mean_diff:.6f}"
-                )
-                validation_passed = False
-            else:
-                logging.info(f"    ✓ During delay [0:{delay_end}]: RTC equals prev_chunk")
-
-        # ============================================================================
-        # Rule 2: After delay, within execution horizon [inference_delay:execution_horizon]
-        #         RTC should be between prev_chunk and no_rtc
-        # ============================================================================
-        blend_start = inference_delay
-        blend_end = min(execution_horizon, chunk_len)
-
-        if blend_end > blend_start:
-            rtc_blend = rtc_actions_t[blend_start:blend_end]
-            prev_blend = prev_chunk[blend_start:blend_end]
-            no_rtc_blend = no_rtc_actions_t[blend_start:blend_end]
-
-            # Check if RTC is between prev_chunk and no_rtc (element-wise)
-            # For each element, check if it's between the min and max of prev_chunk and no_rtc
-            min_bound = torch.minimum(prev_blend, no_rtc_blend)
-            max_bound = torch.maximum(prev_blend, no_rtc_blend)
-
-            within_bounds = torch.logical_and(rtc_blend >= min_bound, rtc_blend <= max_bound)
-
-            if not torch.all(within_bounds):
-                violations = torch.sum(~within_bounds).item()
-                total_elements = within_bounds.numel()
-                violation_pct = 100.0 * violations / total_elements
-
-                # Find max violation
-                lower_violations = torch.maximum(torch.tensor(0.0), min_bound - rtc_blend)
-                upper_violations = torch.maximum(torch.tensor(0.0), rtc_blend - max_bound)
-                max_violation = torch.max(torch.maximum(lower_violations, upper_violations)).item()
-
-                warnings.append(
-                    f"    ⚠ VALIDATION FAILED: In blend region [{blend_start}:{blend_end}], "
-                    f"RTC is NOT always between prev_chunk and no_rtc!\n"
-                    f"      Violations: {violations}/{total_elements} elements ({violation_pct:.1f}%)\n"
-                    f"      Max violation distance: {max_violation:.6f}"
-                )
-                validation_passed = False
-            else:
-                logging.info(
-                    f"    ✓ Blend region [{blend_start}:{blend_end}]: RTC is between prev_chunk and no_rtc"
-                )
-
-        # ============================================================================
-        # Rule 3: After execution horizon [execution_horizon:], RTC should equal no_rtc
-        # ============================================================================
-        if execution_horizon < chunk_len:
-            rtc_after = rtc_actions_t[execution_horizon:chunk_len]
-            no_rtc_after = no_rtc_actions_t[execution_horizon:chunk_len]
-
-            logging.info(f"  rtc_after: {rtc_after}")
-            logging.info(f"  no_rtc_after: {no_rtc_after}")
-
-            if not torch.allclose(rtc_after, no_rtc_after, rtol=rtol):
-                max_diff = torch.max(torch.abs(rtc_after - no_rtc_after)).item()
-                mean_diff = torch.mean(torch.abs(rtc_after - no_rtc_after)).item()
-                warnings.append(
-                    f"    ⚠ VALIDATION FAILED: After execution horizon [{execution_horizon}:{chunk_len}], "
-                    f"RTC does NOT equal no_rtc!\n"
-                    f"      Max difference: {max_diff:.6f}\n"
-                    f"      Mean difference: {mean_diff:.6f}"
-                )
-                validation_passed = False
-            else:
-                logging.info(
-                    f"    ✓ After execution horizon [{execution_horizon}:{chunk_len}]: RTC equals no_rtc"
-                )
-
-        # ============================================================================
-        # Report results
-        # ============================================================================
-        logging.info("=" * 80)
-        if validation_passed:
-            logging.info("  ✅ VALIDATION PASSED: All RTC behavior checks passed!")
-            logging.info("    • During delay: RTC = prev_chunk ✓")
-            logging.info("    • Blend region: prev_chunk ≤ RTC ≤ no_rtc ✓")
-            logging.info("    • After execution horizon: RTC = no_rtc ✓")
-        else:
-            logging.error("  ❌ VALIDATION FAILED: RTC behavior does not match expected!")
-            logging.error("")
-            for warning in warnings:
-                logging.error(warning)
-            logging.error("")
-            logging.error("  Please check the implementation of RTC guidance.")
-
     def plot_final_actions_comparison(self, rtc_actions, no_rtc_actions, prev_chunk_left_over):
         """Plot final action predictions comparison on a single chart.
 
@@ -795,16 +638,34 @@ class RTCEvaluator:
             ax.set_xticks(range(0, max_len, max(1, max_len // 20)))  # Show ~20 ticks
             ax.set_xlim(-0.5, max_len - 0.5)
 
-            # Add legend only to first subplot
-            if dim_idx == 0:
-                ax.legend(loc="best", fontsize=9)
-
         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()
-        fig.savefig(output_path, dpi=150)
+        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)
 
@@ -825,6 +686,7 @@ class RTCEvaluator:
             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(
@@ -834,6 +696,7 @@ class RTCEvaluator:
             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
@@ -849,15 +712,21 @@ class RTCEvaluator:
             axs_x1t[:, 1], prev_chunk_left_over, start_from=0, color="red", label="Ground truth"
         )
 
-        # Plot ground truth on x_t axes
+        # 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="Ground truth"
+            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="Ground truth"
+            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"))
@@ -875,13 +744,47 @@ class RTCEvaluator:
 
         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()
-        fig.savefig(path, dpi=150)
+        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):
+    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:
@@ -891,6 +794,7 @@ class RTCEvaluator:
             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)
@@ -905,17 +809,18 @@ class RTCEvaluator:
 
         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=f"Step {step_idx}"
+                    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=f"Step {step_idx}"
+                    vt_axs, debug_step.v_t, start_from=0, color=color, label=label
                 )
 
             # Plot correction on separate axes
@@ -925,17 +830,18 @@ class RTCEvaluator:
                     debug_step.correction,
                     start_from=0,
                     color=color,
-                    label=f"Step {step_idx}",
+                    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=f"x1_t Step {step_idx}",
+                    label=x1t_label,
                 )
 
             # Plot error in orange dashed
@@ -947,6 +853,7 @@ class RTCEvaluator:
                 )
 
                 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]),
@@ -954,7 +861,7 @@ class RTCEvaluator:
                         color="orange",
                         linestyle="--",
                         alpha=0.7,
-                        label=f"error Step {step_idx}",
+                        label=error_label,
                     )
 
         # Recalculate axis limits after plotting to ensure proper scaling

examples/rtc/eval_with_real_robot_profiled.py

@@ -1,631 +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.
-
-"""
-Profiled version of eval_with_real_robot.py for performance analysis.
-
-This version adds detailed timing measurements for:
-- Policy inference
-- Preprocessing
-- Postprocessing
-- Action queue operations
-- Robot communication
-- Thread execution times
-
-Usage: Same as eval_with_real_robot.py but with profiling output.
-"""
-
-import logging
-import math
-import sys
-import time
-import traceback
-from collections import defaultdict
-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 ProfileTimer:
-    """Context manager and utility class for timing code sections."""
-
-    def __init__(self, name: str, stats_dict: dict):
-        self.name = name
-        self.stats_dict = stats_dict
-        self.start_time = None
-
-    def __enter__(self):
-        self.start_time = time.perf_counter()
-        return self
-
-    def __exit__(self, *args):
-        elapsed = time.perf_counter() - self.start_time
-        if self.name not in self.stats_dict:
-            self.stats_dict[self.name] = []
-        self.stats_dict[self.name].append(elapsed)
-
-
-class ProfilingStats:
-    """Global profiling statistics collector."""
-
-    def __init__(self):
-        self.stats = defaultdict(list)
-        self.lock = Lock()
-
-    def record(self, name: str, duration: float):
-        with self.lock:
-            self.stats[name].append(duration)
-
-    def timer(self, name: str):
-        """Return a context manager for timing."""
-        return ProfileTimer(name, self.stats)
-
-    def get_summary(self) -> dict[str, dict[str, float]]:
-        """Get summary statistics for all timings."""
-        with self.lock:
-            summary = {}
-            for name, times in self.stats.items():
-                if times:
-                    summary[name] = {
-                        "count": len(times),
-                        "mean": sum(times) / len(times),
-                        "min": min(times),
-                        "max": max(times),
-                        "total": sum(times),
-                    }
-            return summary
-
-    def print_summary(self):
-        """Print formatted summary of all timings."""
-        summary = self.get_summary()
-        
-        logger.info("\n" + "=" * 80)
-        logger.info("PROFILING SUMMARY")
-        logger.info("=" * 80)
-        
-        # Sort by total time (descending)
-        sorted_items = sorted(summary.items(), key=lambda x: x[1]["total"], reverse=True)
-        
-        for name, stats in sorted_items:
-            logger.info(f"\n{name}:")
-            logger.info(f"  Count:     {stats['count']}")
-            logger.info(f"  Mean:      {stats['mean']*1000:.2f} ms")
-            logger.info(f"  Min:       {stats['min']*1000:.2f} ms")
-            logger.info(f"  Max:       {stats['max']*1000:.2f} ms")
-            logger.info(f"  Total:     {stats['total']:.2f} s")
-            logger.info(f"  Hz:        {stats['count']/stats['total']:.2f}")
-        
-        logger.info("\n" + "=" * 80)
-
-
-# Global profiling stats
-profiling_stats = ProfilingStats()
-
-
-class RobotWrapper:
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.lock = Lock()
-
-    def get_observation(self) -> dict[str, Tensor]:
-        with profiling_stats.timer("robot.get_observation"):
-            with self.lock:
-                return self.robot.get_observation()
-
-    def send_action(self, action: Tensor):
-        with profiling_stats.timer("robot.send_action"):
-            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 with profiling.
-
-    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
-        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
-
-        inference_count = 0
-
-        while not shutdown_event.is_set():
-            if action_queue.qsize() <= get_actions_threshold:
-                with profiling_stats.timer("get_actions.total_iteration"):
-                    inference_count += 1
-                    logger.info(f"[GET_ACTIONS] Starting inference #{inference_count}")
-
-                    current_time = time.perf_counter()
-                    action_index_before_inference = action_queue.get_action_index()
-                    
-                    with profiling_stats.timer("get_actions.get_prev_actions"):
-                        prev_actions = action_queue.get_left_over()
-
-                    inference_latency = latency_tracker.max()
-                    inference_delay = math.ceil(inference_latency / time_per_chunk)
-
-                    # Get observation
-                    obs = robot.get_observation()
-
-                    # Apply robot observation processor
-                    with profiling_stats.timer("get_actions.robot_obs_processing"):
-                        obs_processed = robot_observation_processor(obs)
-
-                    # Build dataset frame
-                    with profiling_stats.timer("get_actions.build_dataset_frame"):
-                        obs_with_policy_features = build_dataset_frame(
-                            dataset_features, obs_processed, prefix="observation"
-                        )
-
-                    # Convert to tensors and normalize
-                    with profiling_stats.timer("get_actions.tensor_conversion"):
-                        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]
-                        obs_with_policy_features["robot_type"] = (
-                            robot.robot.name if hasattr(robot.robot, "name") else ""
-                        )
-
-                    # Preprocessing
-                    with profiling_stats.timer("get_actions.preprocessing"):
-                        preproceseded_obs = preprocessor(obs_with_policy_features)
-
-                    # Policy inference
-                    with profiling_stats.timer("get_actions.policy_inference"):
-                        actions = policy.predict_action_chunk(
-                            preproceseded_obs,
-                            inference_delay=inference_delay,
-                            prev_chunk_left_over=prev_actions,
-                        )
-
-                    # Clone for RTC
-                    with profiling_stats.timer("get_actions.clone_actions"):
-                        original_actions = actions.squeeze(0).clone()
-
-                    # Postprocessing
-                    with profiling_stats.timer("get_actions.postprocessing"):
-                        postprocessed_actions = postprocessor(actions)
-                        postprocessed_actions = postprocessed_actions.squeeze(0)
-
-                    # Update latency tracker
-                    new_latency = time.perf_counter() - current_time
-                    new_delay = math.ceil(new_latency / time_per_chunk)
-                    latency_tracker.add(new_latency)
-
-                    logger.info(
-                        f"[GET_ACTIONS] Inference #{inference_count} completed in {new_latency*1000:.2f}ms "
-                        f"(delay={new_delay} chunks)"
-                    )
-
-                    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."
-                        )
-
-                    # Merge into action queue
-                    with profiling_stats.timer("get_actions.action_queue_merge"):
-                        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 with profiling.
-
-    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()
-
-            with profiling_stats.timer("actor.total_iteration"):
-                # Get action from queue
-                with profiling_stats.timer("actor.queue_get"):
-                    action = action_queue.get()
-
-                if action is not None:
-                    # Process action
-                    with profiling_stats.timer("actor.action_processing"):
-                        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))
-                    
-                    # Send to robot (includes robot.send_action timing)
-                    robot.send_action(action_processed)
-                    action_count += 1
-
-            # Sleep to maintain target FPS
-            dt_s = time.perf_counter() - start_time
-            sleep_time = max(0, (action_interval - dt_s) - 0.001)
-            if sleep_time > 0:
-                time.sleep(sleep_time)
-
-        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
-        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 profiling."""
-
-    # Initialize logging
-    init_logging()
-
-    logger.info(f"Using device: {cfg.device}")
-    logger.info("=" * 80)
-    logger.info("PROFILING MODE ENABLED")
-    logger.info("=" * 80)
-
-    # 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 processor
-    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")
-
-    # Print profiling summary
-    profiling_stats.print_summary()
-
-    logger.info("Cleanup completed")
-
-
-if __name__ == "__main__":
-    demo_cli()
-    logging.info("RTC demo finished")
-

examples/rtc/profile_pi0_rtc_detailed.py

@@ -1,358 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Comprehensive profiling script for Pi0 with RTC.
-
-This script demonstrates how to use all the profiling tools to identify
-bottlenecks in Pi0 policy inference with RTC enabled.
-
-It profiles:
-1. Overall inference time
-2. RTC-specific operations (guidance, weights, etc.)
-3. Preprocessing/postprocessing
-4. Individual method timings
-
-Usage:
-    uv run examples/rtc/profile_pi0_rtc_detailed.py \
-        --policy_path=helper2424/pi05_check_rtc \
-        --device=mps \
-        --num_iterations=20 \
-        --execution_horizon=20 \
-        --enable_rtc_profiling
-"""
-
-import argparse
-import logging
-import sys
-import time
-
-import numpy as np
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.utils.profiling import (
-    ProfileContext,
-    clear_profiling_stats,
-    enable_profiling,
-    get_profiling_stats,
-    print_profiling_summary,
-)
-
-# Import monkey patching for RTC profiling
-try:
-    from examples.rtc.add_rtc_profiling import monkey_patch_rtc_profiling
-except ImportError:
-    logging.warning("Could not import add_rtc_profiling, detailed RTC profiling disabled")
-    monkey_patch_rtc_profiling = None
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-def create_mock_observation(policy_config, device: str) -> dict:
-    """Create a mock observation matching policy requirements.
-    
-    Args:
-        policy_config: Policy configuration
-        device: Device to create tensors on
-        
-    Returns:
-        Mock observation dictionary
-    """
-    obs = {}
-    
-    # Create mock state observation
-    state_dim = 10  # Typical robot state dimension
-    obs["observation.state"] = torch.randn(1, state_dim, device=device)
-    
-    # Create mock images if needed
-    # For Pi0, we typically need at least one image
-    image_height = 224
-    image_width = 224
-    
-    # Common image keys for Pi0
-    image_keys = ["observation.images.gripper", "observation.images.front"]
-    
-    for key in image_keys:
-        # Images should be [B, C, H, W] and normalized to [0, 1]
-        obs[key] = torch.rand(1, 3, image_height, image_width, device=device)
-    
-    # Add task
-    obs["task"] = ["Pick up the object"]
-    
-    # Add language tokens and attention mask (required for Pi0)
-    # These are mock values - in real usage they come from tokenizer
-    max_seq_len = 32
-    obs["observation.language_tokens"] = torch.randint(0, 1000, (1, max_seq_len), device=device)
-    obs["observation.language_attention_mask"] = torch.ones(1, max_seq_len, device=device)
-    
-    return obs
-
-
-def profile_single_iteration(
-    policy,
-    preprocessor,
-    postprocessor,
-    observation: dict,
-    prev_actions: torch.Tensor | None,
-    use_rtc: bool,
-    inference_delay: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor | None, dict]:
-    """Profile a single inference iteration.
-    
-    Args:
-        policy: Policy instance
-        preprocessor: Observation preprocessor
-        postprocessor: Action postprocessor
-        observation: Input observation
-        prev_actions: Previous action chunk (for RTC)
-        use_rtc: Whether RTC is enabled
-        inference_delay: Inference delay in timesteps
-        
-    Returns:
-        Tuple of (actions, new_prev_actions, timings)
-    """
-    timings = {}
-    
-    with ProfileContext("iteration.total"):
-        # Preprocessing
-        with ProfileContext("iteration.preprocessing"):
-            preprocessed_obs = preprocessor(observation)
-        
-        # Policy inference
-        with ProfileContext("iteration.policy_inference"):
-            if use_rtc:
-                actions = policy.predict_action_chunk(
-                    preprocessed_obs,
-                    inference_delay=inference_delay,
-                    prev_chunk_left_over=prev_actions,
-                )
-            else:
-                actions = policy.predict_action_chunk(preprocessed_obs)
-        
-        # Clone for next iteration (if RTC)
-        new_prev_actions = None
-        if use_rtc:
-            with ProfileContext("iteration.prepare_prev_actions"):
-                execution_horizon = policy.config.rtc_config.execution_horizon
-                if actions.shape[1] > execution_horizon:
-                    new_prev_actions = actions[:, execution_horizon:].clone()
-        
-        # Postprocessing
-        with ProfileContext("iteration.postprocessing"):
-            processed_actions = postprocessor(actions)
-    
-    return processed_actions, new_prev_actions, timings
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Detailed profiling for Pi0 with RTC")
-    parser.add_argument("--policy_path", type=str, required=True, help="Path to pretrained policy")
-    parser.add_argument("--device", type=str, default="cuda", help="Device (cuda/cpu/mps)")
-    parser.add_argument("--num_iterations", type=int, default=20, help="Number of iterations")
-    parser.add_argument("--execution_horizon", type=int, default=10, help="RTC execution horizon")
-    parser.add_argument("--warmup_iterations", type=int, default=5, help="Warmup iterations")
-    parser.add_argument("--enable_rtc_profiling", action="store_true", help="Enable detailed RTC profiling")
-    parser.add_argument("--use_torch_compile", action="store_true", help="Use torch.compile")
-    
-    args = parser.parse_args()
-    
-    logger.info("="*80)
-    logger.info("DETAILED PI0 RTC PROFILING")
-    logger.info("="*80)
-    logger.info(f"Policy: {args.policy_path}")
-    logger.info(f"Device: {args.device}")
-    logger.info(f"Iterations: {args.num_iterations}")
-    logger.info(f"Execution Horizon: {args.execution_horizon}")
-    logger.info(f"RTC Profiling: {args.enable_rtc_profiling}")
-    logger.info("="*80 + "\n")
-    
-    # Enable profiling
-    enable_profiling()
-    
-    # Apply RTC profiling if requested
-    if args.enable_rtc_profiling:
-        if monkey_patch_rtc_profiling is not None:
-            monkey_patch_rtc_profiling()
-            logger.info("✓ Detailed RTC profiling enabled\n")
-        else:
-            logger.warning("⚠ Could not enable detailed RTC profiling\n")
-    
-    # Load policy
-    logger.info("Loading policy...")
-    config = PreTrainedConfig.from_pretrained(args.policy_path)
-    
-    if hasattr(config, "compile_model"):
-        config.compile_model = args.use_torch_compile
-    
-    policy_class = get_policy_class(config.type)
-    policy = policy_class.from_pretrained(args.policy_path, config=config)
-    
-    # Configure RTC
-    policy.config.rtc_config = RTCConfig(
-        enabled=True,
-        execution_horizon=args.execution_horizon,
-        max_guidance_weight=1.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-    )
-    policy.init_rtc_processor()
-    
-    policy = policy.to(args.device)
-    policy.eval()
-    
-    logger.info(f"✓ Policy loaded: {config.type}\n")
-    
-    # Create preprocessor and postprocessor
-    logger.info("Loading preprocessor/postprocessor...")
-    preprocessor, postprocessor = make_pre_post_processors(
-        policy_cfg=config,
-        pretrained_path=args.policy_path,
-        dataset_stats=None,
-        preprocessor_overrides={
-            "device_processor": {"device": args.device},
-        },
-    )
-    logger.info("✓ Preprocessor/postprocessor loaded\n")
-    
-    # Create mock observation
-    logger.info("Creating mock observation...")
-    observation = create_mock_observation(config, args.device)
-    logger.info("✓ Mock observation created\n")
-    
-    # Warmup
-    logger.info(f"Warming up ({args.warmup_iterations} iterations)...")
-    prev_actions = None
-    for i in range(args.warmup_iterations):
-        with torch.no_grad():
-            _, prev_actions, _ = profile_single_iteration(
-                policy=policy,
-                preprocessor=preprocessor,
-                postprocessor=postprocessor,
-                observation=observation,
-                prev_actions=prev_actions,
-                use_rtc=True,
-                inference_delay=0,
-            )
-    
-    # Clear warmup stats
-    clear_profiling_stats()
-    logger.info("✓ Warmup complete\n")
-    
-    # Profiled run WITH RTC
-    logger.info(f"Running profiled iterations WITH RTC ({args.num_iterations} iterations)...")
-    prev_actions = None
-    iteration_times = []
-    
-    for i in range(args.num_iterations):
-        start = time.perf_counter()
-        
-        with torch.no_grad():
-            _, prev_actions, _ = profile_single_iteration(
-                policy=policy,
-                preprocessor=preprocessor,
-                postprocessor=postprocessor,
-                observation=observation,
-                prev_actions=prev_actions,
-                use_rtc=True,
-                inference_delay=0,
-            )
-        
-        # Sync CUDA if needed
-        if args.device.startswith("cuda"):
-            torch.cuda.synchronize()
-        
-        elapsed = time.perf_counter() - start
-        iteration_times.append(elapsed)
-        
-        if (i + 1) % 5 == 0:
-            logger.info(f"  Completed {i+1}/{args.num_iterations}")
-    
-    logger.info("✓ Profiling complete\n")
-    
-    # Print summary statistics
-    logger.info("\n" + "="*80)
-    logger.info("ITERATION TIMING SUMMARY")
-    logger.info("="*80)
-    
-    times_arr = np.array(iteration_times)
-    logger.info(f"Mean time:       {np.mean(times_arr)*1000:.2f} ms")
-    logger.info(f"Median time:     {np.median(times_arr)*1000:.2f} ms")
-    logger.info(f"Std dev:         {np.std(times_arr)*1000:.2f} ms")
-    logger.info(f"Min time:        {np.min(times_arr)*1000:.2f} ms")
-    logger.info(f"Max time:        {np.max(times_arr)*1000:.2f} ms")
-    logger.info(f"Total time:      {np.sum(times_arr):.2f} s")
-    logger.info(f"Throughput:      {len(times_arr)/np.sum(times_arr):.2f} iter/s")
-    logger.info("="*80 + "\n")
-    
-    # Print detailed profiling breakdown
-    print_profiling_summary(sort_by="total")
-    
-    # Print key insights
-    stats = get_profiling_stats()
-    
-    logger.info("\n" + "="*80)
-    logger.info("KEY INSIGHTS")
-    logger.info("="*80)
-    
-    # Find bottlenecks
-    if stats:
-        policy_inference_time = stats.get("iteration.policy_inference", {}).get("mean", 0)
-        preprocessing_time = stats.get("iteration.preprocessing", {}).get("mean", 0)
-        postprocessing_time = stats.get("iteration.postprocessing", {}).get("mean", 0)
-        
-        total_time = policy_inference_time + preprocessing_time + postprocessing_time
-        
-        if total_time > 0:
-            logger.info(f"\nTime breakdown:")
-            logger.info(f"  Policy inference:  {policy_inference_time*1000:.2f} ms ({policy_inference_time/total_time*100:.1f}%)")
-            logger.info(f"  Preprocessing:     {preprocessing_time*1000:.2f} ms ({preprocessing_time/total_time*100:.1f}%)")
-            logger.info(f"  Postprocessing:    {postprocessing_time*1000:.2f} ms ({postprocessing_time/total_time*100:.1f}%)")
-        
-        # RTC-specific insights
-        if args.enable_rtc_profiling:
-            rtc_guidance = stats.get("rtc.denoise_step.guidance_computation", {}).get("mean", 0)
-            rtc_autograd = stats.get("rtc.denoise_step.autograd_correction", {}).get("mean", 0)
-            rtc_base = stats.get("rtc.denoise_step.base_denoising", {}).get("mean", 0)
-            
-            if rtc_guidance > 0:
-                logger.info(f"\nRTC breakdown:")
-                logger.info(f"  Base denoising:    {rtc_base*1000:.2f} ms")
-                logger.info(f"  Guidance compute:  {rtc_guidance*1000:.2f} ms")
-                logger.info(f"  Autograd correct:  {rtc_autograd*1000:.2f} ms")
-                logger.info(f"  RTC overhead:      {(rtc_guidance - rtc_base)*1000:.2f} ms")
-        
-        # Recommendations
-        logger.info("\nRecommendations:")
-        
-        if preprocessing_time > policy_inference_time * 0.3:
-            logger.info("  ⚠ Preprocessing is taking >30% of time")
-            logger.info("    → Consider reducing image resolution")
-            logger.info("    → Consider using fewer cameras")
-        
-        if args.enable_rtc_profiling and rtc_autograd > rtc_base * 0.5:
-            logger.info("  ⚠ RTC autograd overhead is significant")
-            logger.info("    → This is expected, but consider increasing execution_horizon")
-            logger.info("    → Try torch.compile if not already enabled")
-        
-        if not args.use_torch_compile:
-            logger.info("  💡 torch.compile not enabled")
-            logger.info("    → Try --use_torch_compile for potential speedup")
-    
-    logger.info("="*80 + "\n")
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        logger.info("\n\nProfiling interrupted by user")
-        sys.exit(0)
-    except Exception as e:
-        logger.error(f"\n\nError during profiling: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
-

examples/rtc/profile_rtc_comparison.py

@@ -1,347 +0,0 @@
-#!/usr/bin/env python
-
-"""
-Script to compare performance with and without RTC enabled.
-
-This script helps identify whether RTC is actually improving or degrading performance
-by running multiple inference passes and collecting detailed timing statistics.
-
-Usage:
-    # Profile with mock data (no robot needed)
-    uv run examples/rtc/profile_rtc_comparison.py \
-        --policy_path=helper2424/pi05_check_rtc \
-        --device=mps \
-        --num_iterations=50
-
-    # Profile with specific RTC config
-    uv run examples/rtc/profile_rtc_comparison.py \
-        --policy_path=helper2424/pi05_check_rtc \
-        --device=mps \
-        --num_iterations=50 \
-        --execution_horizon=20
-"""
-
-import argparse
-import logging
-import time
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.types import RTCAttentionSchedule
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.utils.profiling import (
-    clear_profiling_stats,
-    enable_profiling,
-    get_profiling_stats,
-    print_profiling_summary,
-)
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ProfileResults:
-    """Results from profiling run."""
-
-    mode: str  # "with_rtc" or "without_rtc"
-    mean_time: float
-    std_time: float
-    min_time: float
-    max_time: float
-    times: list[float]
-    throughput: float  # iterations per second
-
-
-def create_mock_observation(policy, device: str) -> dict:
-    """Create a mock observation for testing.
-
-    Args:
-        policy: Policy instance
-        device: Device to create tensors on
-
-    Returns:
-        Mock observation dictionary
-    """
-    # Get expected input shapes from policy config
-    # This is a simplified version - adjust based on actual policy requirements
-    obs = {}
-
-    # Mock image observations (if needed)
-    if hasattr(policy.config, "input_shapes"):
-        for key, shape in policy.config.input_shapes.items():
-            if "image" in key:
-                # Typical image shape: (batch, channels, height, width)
-                obs[key] = torch.randn(1, *shape, device=device)
-            else:
-                obs[key] = torch.randn(1, *shape, device=device)
-
-    # Add task if needed
-    if "task" in policy.config.__dict__ or hasattr(policy, "accepts_task"):
-        obs["task"] = ["Pick up the object"]
-
-    # Mock state observation
-    obs["observation.state"] = torch.randn(1, 10, device=device)  # Adjust size as needed
-
-    return obs
-
-
-def profile_inference(
-    policy, observation: dict, num_iterations: int, use_rtc: bool, execution_horizon: int = 10
-) -> ProfileResults:
-    """Profile policy inference with or without RTC.
-
-    Args:
-        policy: Policy instance
-        observation: Observation dictionary
-        num_iterations: Number of inference iterations to run
-        use_rtc: Whether to enable RTC
-        execution_horizon: Execution horizon for RTC
-
-    Returns:
-        ProfileResults with timing statistics
-    """
-    mode = "with_rtc" if use_rtc else "without_rtc"
-    logger.info(f"\n{'='*80}")
-    logger.info(f"Profiling: {mode.upper()}")
-    logger.info(f"{'='*80}")
-
-    # Configure RTC
-    if use_rtc:
-        policy.config.rtc_config.enabled = True
-        policy.config.rtc_config.execution_horizon = execution_horizon
-        policy.init_rtc_processor()
-    else:
-        policy.config.rtc_config.enabled = False
-
-    times = []
-    prev_actions = None
-
-    # Warmup
-    logger.info("Warming up (5 iterations)...")
-    for _ in range(5):
-        with torch.no_grad():
-            if use_rtc:
-                _ = policy.predict_action_chunk(
-                    observation, inference_delay=0, prev_chunk_left_over=prev_actions
-                )
-            else:
-                _ = policy.predict_action_chunk(observation)
-
-    # Actual profiling
-    logger.info(f"Running {num_iterations} profiled iterations...")
-    for i in range(num_iterations):
-        start = time.perf_counter()
-
-        with torch.no_grad():
-            if use_rtc:
-                actions = policy.predict_action_chunk(
-                    observation, inference_delay=0, prev_chunk_left_over=prev_actions
-                )
-                # Simulate consuming some actions for next iteration
-                if actions.shape[1] > execution_horizon:
-                    prev_actions = actions[:, execution_horizon:].clone()
-                else:
-                    prev_actions = None
-            else:
-                actions = policy.predict_action_chunk(observation)
-
-        # Synchronize if using CUDA
-        if observation["observation.state"].device.type == "cuda":
-            torch.cuda.synchronize()
-
-        elapsed = time.perf_counter() - start
-        times.append(elapsed)
-
-        if (i + 1) % 10 == 0:
-            logger.info(f"  Completed {i+1}/{num_iterations} iterations")
-
-    # Calculate statistics
-    times_arr = np.array(times)
-    results = ProfileResults(
-        mode=mode,
-        mean_time=float(np.mean(times_arr)),
-        std_time=float(np.std(times_arr)),
-        min_time=float(np.min(times_arr)),
-        max_time=float(np.max(times_arr)),
-        times=times,
-        throughput=num_iterations / sum(times),
-    )
-
-    logger.info(f"\nResults for {mode}:")
-    logger.info(f"  Mean time: {results.mean_time*1000:.2f} ms")
-    logger.info(f"  Std dev:   {results.std_time*1000:.2f} ms")
-    logger.info(f"  Min time:  {results.min_time*1000:.2f} ms")
-    logger.info(f"  Max time:  {results.max_time*1000:.2f} ms")
-    logger.info(f"  Throughput: {results.throughput:.2f} iter/s")
-
-    return results
-
-
-def compare_results(results_without_rtc: ProfileResults, results_with_rtc: ProfileResults):
-    """Compare and print results from both runs.
-
-    Args:
-        results_without_rtc: Results from run without RTC
-        results_with_rtc: Results from run with RTC
-    """
-    logger.info(f"\n{'='*80}")
-    logger.info("COMPARISON SUMMARY")
-    logger.info(f"{'='*80}")
-
-    mean_diff = results_with_rtc.mean_time - results_without_rtc.mean_time
-    mean_diff_pct = (mean_diff / results_without_rtc.mean_time) * 100
-
-    throughput_diff = results_with_rtc.throughput - results_without_rtc.throughput
-    throughput_diff_pct = (throughput_diff / results_without_rtc.throughput) * 100
-
-    logger.info(f"\n{'Metric':<30} {'Without RTC':>15} {'With RTC':>15} {'Difference':>15}")
-    logger.info("-" * 80)
-    logger.info(
-        f"{'Mean time (ms)':<30} "
-        f"{results_without_rtc.mean_time*1000:>15.2f} "
-        f"{results_with_rtc.mean_time*1000:>15.2f} "
-        f"{mean_diff*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Std dev (ms)':<30} "
-        f"{results_without_rtc.std_time*1000:>15.2f} "
-        f"{results_with_rtc.std_time*1000:>15.2f} "
-        f"{(results_with_rtc.std_time - results_without_rtc.std_time)*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Min time (ms)':<30} "
-        f"{results_without_rtc.min_time*1000:>15.2f} "
-        f"{results_with_rtc.min_time*1000:>15.2f} "
-        f"{(results_with_rtc.min_time - results_without_rtc.min_time)*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Max time (ms)':<30} "
-        f"{results_without_rtc.max_time*1000:>15.2f} "
-        f"{results_with_rtc.max_time*1000:>15.2f} "
-        f"{(results_with_rtc.max_time - results_without_rtc.max_time)*1000:>+15.2f}"
-    )
-    logger.info(
-        f"{'Throughput (iter/s)':<30} "
-        f"{results_without_rtc.throughput:>15.2f} "
-        f"{results_with_rtc.throughput:>15.2f} "
-        f"{throughput_diff:>+15.2f}"
-    )
-
-    logger.info(f"\n{'='*80}")
-    logger.info("VERDICT")
-    logger.info(f"{'='*80}")
-
-    if mean_diff_pct < -5:
-        logger.info(f"✓ RTC is FASTER by {abs(mean_diff_pct):.1f}%")
-        logger.info(f"  Mean time reduced by {abs(mean_diff)*1000:.2f} ms")
-    elif mean_diff_pct > 5:
-        logger.info(f"✗ RTC is SLOWER by {mean_diff_pct:.1f}%")
-        logger.info(f"  Mean time increased by {mean_diff*1000:.2f} ms")
-        logger.info("\n  Possible reasons:")
-        logger.info("  - RTC overhead exceeds benefits at current execution horizon")
-        logger.info("  - Inference delay calculation not accounting for RTC processing")
-        logger.info("  - Additional tensor operations in RTC guidance")
-    else:
-        logger.info(f"≈ Performance is SIMILAR (difference: {mean_diff_pct:+.1f}%)")
-
-    logger.info(f"{'='*80}\n")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Profile RTC performance")
-    parser.add_argument(
-        "--policy_path", type=str, required=True, help="Path to pretrained policy"
-    )
-    parser.add_argument(
-        "--device", type=str, default="cuda", help="Device to run on (cuda/cpu/mps)"
-    )
-    parser.add_argument(
-        "--num_iterations", type=int, default=50, help="Number of inference iterations"
-    )
-    parser.add_argument(
-        "--execution_horizon", type=int, default=10, help="RTC execution horizon"
-    )
-    parser.add_argument(
-        "--enable_detailed_profiling",
-        action="store_true",
-        help="Enable detailed method-level profiling",
-    )
-    parser.add_argument(
-        "--use_torch_compile", action="store_true", help="Use torch.compile for faster inference"
-    )
-
-    args = parser.parse_args()
-
-    # Load policy
-    logger.info(f"Loading policy from {args.policy_path}")
-    config = PreTrainedConfig.from_pretrained(args.policy_path)
-    policy_class = get_policy_class(config.type)
-
-    # Set compile flag if needed
-    if hasattr(config, "compile_model"):
-        config.compile_model = args.use_torch_compile
-
-    policy = policy_class.from_pretrained(args.policy_path, config=config)
-
-    # Initialize RTC config
-    policy.config.rtc_config = RTCConfig(
-        execution_horizon=args.execution_horizon,
-        max_guidance_weight=1.0,
-        prefix_attention_schedule=RTCAttentionSchedule.EXP,
-    )
-
-    policy = policy.to(args.device)
-    policy.eval()
-
-    logger.info(f"Policy loaded: {config.type}")
-    logger.info(f"Device: {args.device}")
-    logger.info(f"Execution horizon: {args.execution_horizon}")
-
-    # Create mock observation
-    logger.info("Creating mock observation...")
-    observation = create_mock_observation(policy, args.device)
-
-    # Enable detailed profiling if requested
-    if args.enable_detailed_profiling:
-        enable_profiling()
-        logger.info("Detailed profiling enabled")
-
-    # Profile without RTC
-    results_without_rtc = profile_inference(
-        policy=policy,
-        observation=observation,
-        num_iterations=args.num_iterations,
-        use_rtc=False,
-        execution_horizon=args.execution_horizon,
-    )
-
-    if args.enable_detailed_profiling:
-        logger.info("\nDetailed profiling stats (WITHOUT RTC):")
-        print_profiling_summary()
-        clear_profiling_stats()
-
-    # Profile with RTC
-    results_with_rtc = profile_inference(
-        policy=policy,
-        observation=observation,
-        num_iterations=args.num_iterations,
-        use_rtc=True,
-        execution_horizon=args.execution_horizon,
-    )
-
-    if args.enable_detailed_profiling:
-        logger.info("\nDetailed profiling stats (WITH RTC):")
-        print_profiling_summary()
-
-    # Compare results
-    compare_results(results_without_rtc, results_with_rtc)
-
-
-if __name__ == "__main__":
-    main()
-

pyproject.toml

@@ -98,7 +98,6 @@ pygame-dep = ["pygame>=2.5.1,<2.7.0"]
 placo-dep = ["placo>=0.9.6,<0.10.0"]
 transformers-dep = ["transformers>=4.53.0,<5.0.0"]
 grpcio-dep = ["grpcio==1.73.1", "protobuf==6.31.0"] # TODO: Bumb dependency (compatible with wandb)
-matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0"]
 
 # Motors
 feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0"]
@@ -133,7 +132,7 @@ groot = [
 hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
 
 # Features
-async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"]
+async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3,<4.0.0"]
 
 # Development
 dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1"]

src/lerobot/datasets/lerobot_dataset.py

@@ -712,6 +712,15 @@ 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)
@@ -830,7 +839,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)
+        hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
@@ -847,10 +856,8 @@ 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
@@ -932,7 +939,11 @@ 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:
-                timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
+                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"]
                 query_timestamps[key] = torch.stack(timestamps).tolist()
             else:
                 query_timestamps[key] = [current_ts]
@@ -955,10 +966,16 @@ class LeRobotDataset(torch.utils.data.Dataset):
         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][q_idx])
+                result[key] = torch.stack(self.hf_dataset[key][relative_indices])
             except (KeyError, TypeError, IndexError):
-                result[key] = torch.stack(self.hf_dataset[q_idx][key])
+                result[key] = torch.stack(self.hf_dataset[relative_indices][key])
         return result
 
     def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
@@ -1498,6 +1515,7 @@ 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/utils.py

@@ -28,6 +28,7 @@ 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
@@ -103,7 +104,9 @@ 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) -> Dataset:
+def load_nested_dataset(
+    pq_dir: Path, features: datasets.Features | None = None, episodes: list[int] | 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
@@ -111,15 +114,26 @@ def load_nested_dataset(pq_dir: Path, features: datasets.Features | None = None)
     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():
-        datasets = Dataset.from_parquet([str(path) for path in paths], features=features)
-    return datasets
+        # 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)
 
 
 def get_parquet_num_frames(parquet_path: str | Path) -> int:

src/lerobot/envs/configs.py

@@ -21,7 +21,22 @@ 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, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+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,
+)
 
 
 @dataclass
@@ -246,28 +261,61 @@ class LiberoEnv(EnvConfig):
     features_map: dict[str, str] = field(
         default_factory=lambda: {
             ACTION: ACTION,
-            "agent_pos": OBS_STATE,
-            "pixels/agentview_image": f"{OBS_IMAGES}.image",
-            "pixels/robot0_eye_in_hand_image": f"{OBS_IMAGES}.image2",
+            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",
         }
     )
 
     def __post_init__(self):
         if self.obs_type == "pixels":
-            self.features["pixels/agentview_image"] = PolicyFeature(
+            self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
-            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
+            self.features[LIBERO_KEY_PIXELS_EYE_IN_HAND] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
         elif self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(8,))
-            self.features["pixels/agentview_image"] = PolicyFeature(
+            self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
             )
-            self.features["pixels/robot0_eye_in_hand_image"] = PolicyFeature(
+            self.features[LIBERO_KEY_PIXELS_EYE_IN_HAND] = 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,12 +14,16 @@
 # 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:
@@ -33,6 +37,41 @@ 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,

src/lerobot/envs/libero.py

@@ -28,7 +28,6 @@ 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]:
@@ -175,11 +174,36 @@ class LiberoEnv(gym.Env):
             self.observation_space = spaces.Dict(
                 {
                     "pixels": spaces.Dict(images),
-                    "agent_pos": spaces.Box(
-                        low=AGENT_POS_LOW,
-                        high=AGENT_POS_HIGH,
-                        shape=(OBS_STATE_DIM,),
-                        dtype=np.float64,
+                    "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),
+                                }
+                            ),
+                        }
                     ),
                 }
             )
@@ -191,6 +215,7 @@ 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):
@@ -212,23 +237,48 @@ 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
-        state = np.concatenate(
-            (
-                raw_obs["robot0_eef_pos"],
-                quat2axisangle(raw_obs["robot0_eef_quat"]),
-                raw_obs["robot0_gripper_qpos"],
-            )
-        )
-        agent_pos = state
+
+        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,)
+                },
+            },
+        }
         if self.obs_type == "pixels":
             return {"pixels": images.copy()}
+
         if self.obs_type == "pixels_agent_pos":
-            return {
-                "pixels": images.copy(),
-                "agent_pos": 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
+
         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')."
@@ -355,12 +405,10 @@ 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

@@ -29,10 +29,22 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.envs.configs import EnvConfig
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
 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.
@@ -78,12 +90,14 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
 
         return_observations[OBS_ENV_STATE] = env_state
 
-    # TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
-    agent_pos = torch.from_numpy(observations["agent_pos"]).float()
-    if agent_pos.dim() == 1:
-        agent_pos = agent_pos.unsqueeze(0)
-    return_observations[OBS_STATE] = agent_pos
+    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
 
+    if "robot_state" in observations:
+        return_observations[f"{OBS_STR}.robot_state"] = _convert_nested_dict(observations["robot_state"])
     return return_observations
 
 

src/lerobot/policies/rtc/README.md

@@ -1,49 +1,38 @@
-# Real-Time Chunking (RTC) Module
+# Real-Time Chunking (RTC)
 
-This module implements Real-Time Chunking and related adaptive inference techniques for robotics policies in LeRobot.
+This module contains the LeRobot implementation of **Real-Time Chunking (RTC)**, an inference-time technique for flow-matching based policies.
 
-## Overview
+**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/).
 
-Real-Time Chunking (RTC) addresses the challenge of real-time inference in action chunking policies by treating chunk generation as an inpainting problem. It strategically handles overlapping timesteps between action chunks using prefix attention mechanisms.
+---
 
-It is particularly effective for handling long-horizon inference in robotics policies.
+## Citation
 
-## Integration with Policies
+If you use Real-Time Chunking in your work, please cite:
 
-RTC can be integrated with any policy that supports flow mathicng for chunking:
+```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}
+}
 
-- **SmolVLA**: Vision-language-action model with RTC support
-- **Pi0**: Action prediction model with adaptive chunking
-- **Pi05**: Action prediction model with adaptive chunking
-
-## Original Implementation
-
-This implementation is based on Physical Intelligence's Kinetix RTC:
-
-- [Original RTC implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix/blob/main/src/model.py#L214)
-- [Kinetix GitHub Repository](https://github.com/Physical-Intelligence/real-time-chunking-kinetix)
-
-## References
-
-- [Real Time Chunking Paper](https://www.physicalintelligence.company/research/real_time_chunking)
-- [Physical Intelligence Kinetix](https://github.com/Physical-Intelligence/real-time-chunking-kinetix)
-
-## How to run
-
-### Check with data from the dataset
-
-```bash
-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 \
---seed=42
+@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},
+}
 ```
 
-This script will evaluate RTC on a data from a dataset and save the results to a file, u can check the results in the `rtc_debug_output` directory.
+---
 
-The example output should look like this:
-![Flow Matching with RTC](./flow_matching.png)
+## License
 
-It shows how flow matching works with RTC and without it. The chart shows values of action predictions for each timestep. The colour shows the the generation progress. The blue ones - earlier timesteps, the yellow ones - later timesteps. The red line is the ground truth (previous action chunk).
+This implementation follows the **Apache 2.0 License**, consistent with the LeRobot project.

src/lerobot/policies/rtc/debug_visualizer.py

@@ -111,7 +111,3 @@ class RTCDebugVisualizer:
             if not ax.yaxis.get_label().get_text():
                 ax.set_ylabel(f"Dim {dim_idx}", fontsize=10)
             ax.grid(True, alpha=0.3)
-
-            # Add legend if label provided and this is the first dimension
-            if label and dim_idx == 0:
-                ax.legend(loc="best", fontsize=8)

src/lerobot/processor/env_processor.py

@@ -0,0 +1,154 @@
+#!/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/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
+from lerobot.envs.factory import make_env, make_env_pre_post_processors
 from lerobot.envs.utils import (
     add_envs_task,
     check_env_attributes_and_types,
@@ -94,6 +94,8 @@ 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,
@@ -165,11 +167,19 @@ 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)"
@@ -239,6 +249,8 @@ 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,
@@ -319,6 +331,8 @@ 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,
@@ -517,10 +531,16 @@ 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,
@@ -561,6 +581,8 @@ 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,
@@ -576,6 +598,8 @@ 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,
@@ -600,6 +624,8 @@ def run_one(
     env,
     *,
     policy,
+    env_preprocessor,
+    env_postprocessor,
     preprocessor,
     postprocessor,
     n_episodes: int,
@@ -622,6 +648,8 @@ 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,
@@ -639,6 +667,8 @@ 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,
@@ -694,6 +724,8 @@ 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
+from lerobot.envs.factory import make_env, make_env_pre_post_processors
 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
@@ -259,6 +259,8 @@ 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=}")
@@ -274,6 +276,7 @@ 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,
         )
@@ -384,6 +387,8 @@ 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/utils/constants.py

@@ -70,3 +70,15 @@ 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/profiling.py

@@ -1,206 +0,0 @@
-"""
-Profiling utilities for performance analysis.
-
-Usage:
-    from lerobot.utils.profiling import profile_method, get_profiling_stats, print_profiling_summary
-
-    @profile_method
-    def my_slow_function(x):
-        return x * 2
-
-    # At end of execution:
-    print_profiling_summary()
-"""
-
-import functools
-import logging
-import time
-from collections import defaultdict
-from threading import Lock
-from typing import Any, Callable
-
-logger = logging.getLogger(__name__)
-
-# Global profiling statistics storage
-_profiling_stats: dict[str, list[float]] = defaultdict(list)
-_profiling_lock = Lock()
-_profiling_enabled = False
-
-
-def enable_profiling():
-    """Enable profiling globally."""
-    global _profiling_enabled
-    _profiling_enabled = True
-    logger.info("Profiling enabled")
-
-
-def disable_profiling():
-    """Disable profiling globally."""
-    global _profiling_enabled
-    _profiling_enabled = False
-    logger.info("Profiling disabled")
-
-
-def is_profiling_enabled() -> bool:
-    """Check if profiling is enabled."""
-    return _profiling_enabled
-
-
-def record_timing(name: str, duration: float):
-    """Record a timing measurement.
-
-    Args:
-        name: Name/identifier for this timing
-        duration: Duration in seconds
-    """
-    if not _profiling_enabled:
-        return
-
-    with _profiling_lock:
-        _profiling_stats[name].append(duration)
-
-
-def profile_method(func: Callable) -> Callable:
-    """Decorator to profile a method or function.
-
-    Args:
-        func: Function to profile
-
-    Returns:
-        Wrapped function that records execution time
-    """
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs) -> Any:
-        if not _profiling_enabled:
-            return func(*args, **kwargs)
-
-        start = time.perf_counter()
-        try:
-            result = func(*args, **kwargs)
-            return result
-        finally:
-            duration = time.perf_counter() - start
-            # Use fully qualified name
-            name = f"{func.__module__}.{func.__qualname__}"
-            record_timing(name, duration)
-
-    return wrapper
-
-
-class ProfileContext:
-    """Context manager for profiling code blocks.
-
-    Usage:
-        with ProfileContext("my_operation"):
-            # ... code to profile ...
-    """
-
-    def __init__(self, name: str):
-        self.name = name
-        self.start = None
-
-    def __enter__(self):
-        if _profiling_enabled:
-            self.start = time.perf_counter()
-        return self
-
-    def __exit__(self, *args):
-        if _profiling_enabled and self.start is not None:
-            duration = time.perf_counter() - self.start
-            record_timing(self.name, duration)
-
-
-def get_profiling_stats() -> dict[str, dict[str, float]]:
-    """Get summary statistics for all profiled functions.
-
-    Returns:
-        Dictionary mapping function names to their stats (count, mean, min, max, total)
-    """
-    with _profiling_lock:
-        summary = {}
-        for name, times in _profiling_stats.items():
-            if times:
-                summary[name] = {
-                    "count": len(times),
-                    "mean": sum(times) / len(times),
-                    "min": min(times),
-                    "max": max(times),
-                    "total": sum(times),
-                    "mean_ms": (sum(times) / len(times)) * 1000,
-                    "min_ms": min(times) * 1000,
-                    "max_ms": max(times) * 1000,
-                }
-        return summary
-
-
-def clear_profiling_stats():
-    """Clear all profiling statistics."""
-    with _profiling_lock:
-        _profiling_stats.clear()
-    logger.info("Profiling stats cleared")
-
-
-def print_profiling_summary(sort_by: str = "total"):
-    """Print formatted summary of profiling statistics.
-
-    Args:
-        sort_by: Sort key ('total', 'mean', 'count', 'max')
-    """
-    summary = get_profiling_stats()
-
-    if not summary:
-        logger.info("No profiling data available")
-        return
-
-    logger.info("\n" + "=" * 100)
-    logger.info("PROFILING SUMMARY")
-    logger.info("=" * 100)
-
-    # Sort by requested key
-    sorted_items = sorted(summary.items(), key=lambda x: x[1].get(sort_by, 0), reverse=True)
-
-    # Print header
-    logger.info(
-        f"{'Function':<60} {'Count':>8} {'Mean (ms)':>12} {'Min (ms)':>12} {'Max (ms)':>12} {'Total (s)':>12}"
-    )
-    logger.info("-" * 100)
-
-    # Print each function's stats
-    for name, stats in sorted_items:
-        # Shorten long names
-        display_name = name if len(name) <= 60 else "..." + name[-57:]
-
-        logger.info(
-            f"{display_name:<60} "
-            f"{stats['count']:>8} "
-            f"{stats['mean_ms']:>12.2f} "
-            f"{stats['min_ms']:>12.2f} "
-            f"{stats['max_ms']:>12.2f} "
-            f"{stats['total']:>12.2f}"
-        )
-
-    logger.info("=" * 100)
-
-    # Print summary
-    total_time = sum(s["total"] for s in summary.values())
-    total_calls = sum(s["count"] for s in summary.values())
-    logger.info(f"\nTotal profiled time: {total_time:.2f}s across {total_calls} calls")
-    logger.info("=" * 100 + "\n")
-
-
-def profile_section(name: str):
-    """Return a context manager for profiling a code section.
-
-    Args:
-        name: Name for this section
-
-    Returns:
-        ProfileContext instance
-
-    Usage:
-        with profile_section("data_loading"):
-            data = load_data()
-    """
-    return ProfileContext(name)
-

tests/policies/smolvla/test_smolvla_rtc.py

@@ -23,13 +23,15 @@ from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedu
 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.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
 from lerobot.utils.random_utils import set_seed  # noqa: E402
-from tests.utils import require_cuda  # 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)
 
@@ -63,8 +65,11 @@ def test_smolvla_rtc_initialization():
     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)
 
@@ -82,9 +87,12 @@ def test_smolvla_rtc_initialization_without_rtc_config():
     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)
 
@@ -162,9 +170,12 @@ def test_smolvla_rtc_inference_with_prev_chunk():
     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)
 
@@ -233,9 +244,12 @@ def test_smolvla_rtc_inference_without_prev_chunk():
     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)
 

tests/processor/test_libero_processor.py

@@ -0,0 +1,72 @@
+#!/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)