chore(dependencies): update uv.lock

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

docs/source/_toctree.yml

@@ -73,8 +73,6 @@
 - sections:
   - local: sarm
     title: SARM
-  - local: topreward
-    title: TOPReward
   title: "Reward Models"
 - sections:
   - local: inference

docs/source/act.mdx

@@ -79,17 +79,13 @@ If your local computer doesn't have a powerful GPU, you can utilize Google Colab
 Once training is complete, you can evaluate your ACT policy using the `lerobot-record` command with your trained policy. This will run inference and record evaluation episodes:
 
 ```bash
-lerobot-record \
-  --robot.type=so100_follower \
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=${HF_USER}/act_policy \
+  --robot.type=so101_follower \
   --robot.port=/dev/ttyACM0 \
-  --robot.id=my_robot \
   --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
   --display_data=true \
-  --dataset.repo_id=${HF_USER}/eval_act_your_dataset \
-  --dataset.num_episodes=10 \
-  --dataset.single_task="Your task description" \
-  --dataset.streaming_encoding=true \
-  --dataset.encoder_threads=2 \
-  # --dataset.camera_encoder.vcodec=auto \
-  --policy.path=${HF_USER}/act_policy
+  --task="Your task description" \ # can be skipped for ACT
+  --duration=60
 ```

docs/source/groot.mdx

@@ -105,10 +105,12 @@ These results demonstrate GR00T's strong generalization capabilities across dive
 
 ### Evaluate in your hardware setup
 
-Once you have trained your model using your parameters you can run inference in your downstream task. Follow the instructions in [Imitation Learning for Robots](./il_robots). For example:
+Once you have trained your model using your parameters you can run inference in your downstream task. Follow the instructions in [Policy Deployment (lerobot-rollout)](./inference). For example:
 
 ```bash
-lerobot-record \
+lerobot-rollout\
+  --strategy.type=sentry \
+  --strategy.upload_every_n_episodes=5 \
   --robot.type=bi_so_follower \
   --robot.left_arm_port=/dev/ttyACM1 \
   --robot.right_arm_port=/dev/ttyACM0 \
@@ -119,14 +121,12 @@ lerobot-record \
   }' \
   --display_data=true \
   --dataset.repo_id=<user>/eval_groot-bimanual  \
-  --dataset.num_episodes=10 \
   --dataset.single_task="Grab and handover the red cube to the other arm" \
   --dataset.streaming_encoding=true \
   --dataset.encoder_threads=2 \
   # --dataset.camera_encoder.vcodec=auto \
   --policy.path=<user>/groot-bimanual \ # your trained model
-  --dataset.episode_time_s=30 \
-  --dataset.reset_time_s=10
+  --duration=600
 ```
 
 ## License

docs/source/il_robots.mdx

@@ -68,13 +68,13 @@ from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig
 from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
 
 robot_config = SO101FollowerConfig(
-    port="/dev/tty.usbmodem58760431541",
-    id="my_red_robot_arm",
+    port="/dev/tty.usbmodem5AB90687491",
+    id="my_follower_arm",
 )
 
 teleop_config = SO101LeaderConfig(
-    port="/dev/tty.usbmodem58760431551",
-    id="my_blue_leader_arm",
+    port="/dev/tty.usbmodem5AB90689011",
+    id="my_leader_arm",
 )
 
 robot = SO101Follower(robot_config)
@@ -108,13 +108,13 @@ With `rerun`, you can teleoperate again while simultaneously visualizing the cam
 <hfoption id="Command">
 ```bash
 lerobot-teleoperate \
-    --robot.type=koch_follower \
-    --robot.port=/dev/tty.usbmodem58760431541 \
-    --robot.id=my_awesome_follower_arm \
-    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \
-    --teleop.type=koch_leader \
-    --teleop.port=/dev/tty.usbmodem58760431551 \
-    --teleop.id=my_awesome_leader_arm \
+    --robot.type=so101_follower \
+    --robot.port=/dev/tty.usbmodem5AB90687491 \
+    --robot.id=my_follower_arm \
+    --robot.cameras="{front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
+    --teleop.type=so101_leader \
+    --teleop.port=/dev/tty.usbmodem5AB90689011 \
+    --teleop.id=my_leader_arm \
     --display_data=true
 ```
 </hfoption>
@@ -122,34 +122,48 @@ lerobot-teleoperate \
 
 <!-- prettier-ignore-start -->
 ```python
+import time
+from lerobot.teleoperators.so_leader import SO101Leader, SO101LeaderConfig
+from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
 from lerobot.cameras.opencv import OpenCVCameraConfig
-from lerobot.teleoperators.koch_leader import KochLeader, KochLeaderConfig
-from lerobot.robots.koch_follower import KochFollower, KochFollowerConfig
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data, shutdown_rerun
 
-camera_config = {
-    "front": OpenCVCameraConfig(index_or_path=0, width=1920, height=1080, fps=30)
-}
-
-robot_config = KochFollowerConfig(
-    port="/dev/tty.usbmodem585A0076841",
-    id="my_red_robot_arm",
-    cameras=camera_config
+robot_config = SO101FollowerConfig(
+    port="/dev/tty.usbmodem5AB90687491",
+    id="my_follower_arm",
+    cameras={
+        "wrist": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+        "top": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30)
+    }
 )
 
-teleop_config = KochLeaderConfig(
-    port="/dev/tty.usbmodem58760431551",
-    id="my_blue_leader_arm",
+teleop_config = SO101LeaderConfig(
+    port="/dev/tty.usbmodem5AB90689011",
+    id="my_leader_arm",
 )
 
-robot = KochFollower(robot_config)
-teleop_device = KochLeader(teleop_config)
+init_rerun(session_name="teleoperation")
+
+robot = SO101Follower(robot_config)
+teleop_device = SO101Leader(teleop_config)
 robot.connect()
 teleop_device.connect()
 
+TARGET_HZ = 30
+TIME_PER_FRAME = 1.0 / TARGET_HZ
+
 while True:
+    start_time = time.perf_counter()
+
     observation = robot.get_observation()
     action = teleop_device.get_action()
     robot.send_action(action)
+    log_rerun_data(observation=observation, action=action)
+
+    elapsed_time = time.perf_counter() - start_time
+    sleep_time = TIME_PER_FRAME - elapsed_time
+    if sleep_time > 0:
+        time.sleep(sleep_time)
 ```
 <!-- prettier-ignore-end -->
 
@@ -202,10 +216,11 @@ lerobot-record \
 <!-- prettier-ignore-start -->
 ```python
 from lerobot.cameras.opencv import OpenCVCameraConfig
-from lerobot.datasets import LeRobotDataset
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.utils.feature_utils import hw_to_dataset_features
-from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
-from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
+from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
+from lerobot.teleoperators.so_leader.config_so_leader import SO101LeaderConfig
+from lerobot.teleoperators.so_leader.so_leader import SO101Leader
 from lerobot.common.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import init_rerun
@@ -218,71 +233,56 @@ EPISODE_TIME_SEC = 60
 RESET_TIME_SEC = 10
 TASK_DESCRIPTION = "My task description"
 
-# Create robot configuration
-robot_config = SO100FollowerConfig(
-    id="my_awesome_follower_arm",
-    cameras={
-        "front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS) # Optional: fourcc="MJPG" for troubleshooting OpenCV async error.
-    },
-    port="/dev/tty.usbmodem58760434471",
-)
-
-teleop_config = SO100LeaderConfig(
-    id="my_awesome_leader_arm",
-    port="/dev/tty.usbmodem585A0077581",
-)
-
-# Initialize the robot and teleoperator
-robot = SO100Follower(robot_config)
-teleop = SO100Leader(teleop_config)
-
-# Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, "action")
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id="<hf_username>/<dataset_repo_id>",
-    fps=FPS,
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Initialize the keyboard listener and rerun visualization
-_, events = init_keyboard_listener()
-init_rerun(session_name="recording")
-
-# Connect the robot and teleoperator
-robot.connect()
-teleop.connect()
-
-# Create the required processors
-teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
-
-episode_idx = 0
-while episode_idx < NUM_EPISODES and not events["stop_recording"]:
-    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        teleop_action_processor=teleop_action_processor,
-        robot_action_processor=robot_action_processor,
-        robot_observation_processor=robot_observation_processor,
-        teleop=teleop,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
+def main():
+    # Create robot configuration
+    robot_config = SO101FollowerConfig(
+        port="/dev/tty.usbmodem5AB90687491",
+        id="my_follower_arm",
+        cameras={
+            "wrist": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+            "top": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30)
+        }
     )
 
-    # Reset the environment if not stopping or re-recording
-    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
-        log_say("Reset the environment")
+    teleop_config = SO101LeaderConfig(
+        port="/dev/tty.usbmodem5AB90689011",
+        id="my_leader_arm",
+    )
+
+    # Initialize the robot and teleoperator
+    robot = SO101Follower(robot_config)
+    teleop = SO101Leader(teleop_config)
+
+    # Configure the dataset features
+    action_features = hw_to_dataset_features(robot.action_features, "action")
+    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+    dataset_features = {**action_features, **obs_features}
+
+    # Create the dataset
+    dataset = LeRobotDataset.create(
+        repo_id="<hf_username>/<dataset_repo_id>",
+        fps=FPS,
+        features=dataset_features,
+        robot_type=robot.name,
+        use_videos=True,
+        image_writer_threads=4,
+    )
+
+    # Initialize the keyboard listener and rerun visualization
+    _, events = init_keyboard_listener()
+    init_rerun(session_name="recording")
+
+    # Connect the robot and teleoperator
+    robot.connect()
+    teleop.connect()
+
+    # Create the required processors
+    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+
+    episode_idx = 0
+    while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+        log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
+
         record_loop(
             robot=robot,
             events=events,
@@ -291,26 +291,50 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
             robot_action_processor=robot_action_processor,
             robot_observation_processor=robot_observation_processor,
             teleop=teleop,
-            control_time_s=RESET_TIME_SEC,
+            dataset=dataset,
+            control_time_s=EPISODE_TIME_SEC,
             single_task=TASK_DESCRIPTION,
             display_data=True,
         )
 
-    if events["rerecord_episode"]:
-        log_say("Re-recording episode")
-        events["rerecord_episode"] = False
-        events["exit_early"] = False
-        dataset.clear_episode_buffer()
-        continue
+        # Reset the environment if not stopping or re-recording
+        if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+            log_say("Reset the environment")
+            record_loop(
+                robot=robot,
+                events=events,
+                fps=FPS,
+                teleop_action_processor=teleop_action_processor,
+                robot_action_processor=robot_action_processor,
+                robot_observation_processor=robot_observation_processor,
+                teleop=teleop,
+                control_time_s=RESET_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+            )
 
-    dataset.save_episode()
-    episode_idx += 1
+        if events["rerecord_episode"]:
+            log_say("Re-recording episode")
+            events["rerecord_episode"] = False
+            events["exit_early"] = False
+            dataset.clear_episode_buffer()
+            continue
 
-# Clean up
-log_say("Stop recording")
-robot.disconnect()
-teleop.disconnect()
-dataset.push_to_hub()
+        dataset.save_episode()
+        episode_idx += 1
+
+    # finalize dataset
+    log_say("Finalizing dataset...")
+    dataset.finalize()
+    # Clean up
+    log_say("Stop recording")
+    robot.disconnect()
+    teleop.disconnect()
+    dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()
 ```
 <!-- prettier-ignore-end -->
 
@@ -348,7 +372,7 @@ The `record` function provides a suite of tools for capturing and managing data
 ##### 2. Checkpointing and Resuming
 
 - Checkpoints are automatically created during recording.
-- If an issue occurs, you can resume by re-running the same command with `--resume=true`. When resuming a recording, `--dataset.num_episodes` must be set to the **number of additional episodes to be recorded**, and not to the targeted total number of episodes in the dataset !
+- If an issue occurs or you want to record additional episodes in the same dataset, you can resume by re-running the same command with `--resume=true`. When resuming a recording, `--dataset.num_episodes` must be set to the **number of additional episodes to be recorded**, and not to the targeted total number of episodes in the dataset! Make sure that you also set `--dataset.root="local_path"`, it's a local path to save the new part of the dataset and is required to resume.
 - To start recording from scratch, **manually delete** the dataset directory.
 
 ##### 3. Recording Parameters
@@ -422,7 +446,7 @@ from lerobot.utils.utils import log_say
 
 episode_idx = 0
 
-robot_config = SO100FollowerConfig(port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm")
+robot_config = SO100FollowerConfig(port="/dev/tty.usbmodem5AB90687491", id="my_follower_arm")
 
 robot = SO100Follower(robot_config)
 robot.connect()
@@ -490,6 +514,83 @@ Additionally you can provide extra `tags` or specify a `license` for your model
 
 If your local computer doesn't have a powerful GPU you could utilize Google Colab to train your model by following the [ACT training notebook](./notebooks#training-act).
 
+#### Train using Hugging Face Jobs
+
+Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs).
+
+To run the training use this command:
+
+<hfoptions id="train_with_hf_jobs">
+<hfoption id="Command">
+```bash
+hf jobs run \
+  --flavor a10g-small \
+  --timeout 4h \
+  --secrets HF_TOKEN \
+  huggingface/lerobot-gpu:latest \
+  -- \
+  python -m lerobot.scripts.lerobot_train \
+    --dataset.repo_id=username/dataset \
+    --policy.type=act \
+    --steps=5000 \
+    --batch_size=16 \
+    --policy.device=cuda \
+    --policy.repo_id=username/your_policy \
+    --log_freq=100
+```
+</hfoption>
+<hfoption id="API example">
+
+<!-- prettier-ignore-start -->
+```python
+from huggingface_hub import run_job, get_token
+
+run_name = "act_so101_hf_jobs"
+dataset_id = "username/dataset"
+user_hub_id = "username"
+
+command_args = [
+    "python", "-m", "lerobot.scripts.lerobot_train",
+    "--dataset.repo_id", dataset_id,
+    "--policy.type", "act",
+    "--steps", "5000",
+    "--batch_size", "16",
+    "--num_workers", "4",
+    "--policy.device", "cuda",
+    "--log_freq", "100",
+    "--save_freq", "1000",
+    "--save_checkpoint", "true",
+    "--wandb.enable", "false",
+    "--policy.repo_id", f"{user_hub_id}/{run_name}"
+]
+
+print(f"Submitting job '{run_name}' to Hugging Face Infrastructure...")
+
+job_info = run_job(
+    image="huggingface/lerobot-gpu:latest",
+    command=command_args,
+    flavor="a10g-small",
+    timeout="4h",
+    secrets={"HF_TOKEN": get_token()}
+)
+
+print("\n🚀 Job successfully launched!")
+print(f"🔹 Job ID: {job_info.id}")
+print(f"🔗 Live UI Dashboard & Logs: {job_info.url}")
+```
+<!-- prettier-ignore-end -->
+
+</hfoption>
+</hfoptions>
+
+You can modify the `--flavor` to use different hardware, for example: `t4-small`, `a100-large`, `h200`. Use `hf jobs hardware` to see the full list with pricing.
+Depending on the model you want to train and the hardware you selected you can also modify the `--batch_size` and `--number_of_workers`.
+For longer training sessions increase the timeout.
+
+Once the training is started you can go to [Jobs](https://huggingface.co/settings/jobs) and see if your jobs is running as well as all the outputs. Sometimes it takes a few minutes to schedule your job so be patient.
+
+After training the model will be pushed to hub and you can use it as any other model with LeRobot.
+
 #### Upload policy checkpoints
 
 Once training is done, upload the latest checkpoint with:

docs/source/smolvla.mdx

@@ -97,22 +97,22 @@ Similarly for when recording an episode, it is recommended that you are logged i
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash
-lerobot-record \
+lerobot-rollout \
+  --strategy.type=base \
   --robot.type=so101_follower \
   --robot.port=/dev/ttyACM0 \ # <- Use your port
   --robot.id=my_blue_follower_arm \ # <- Use your robot id
   --robot.cameras="{ front: {type: opencv, index_or_path: 8, width: 640, height: 480, fps: 30}}" \ # <- Use your cameras
-  --dataset.single_task="Grasp a lego block and put it in the bin." \ # <- Use the same task description you used in your dataset recording
-  --dataset.repo_id=${HF_USER}/eval_DATASET_NAME_test \  # <- This will be the dataset name on HF Hub
-  --dataset.episode_time_s=50 \
-  --dataset.num_episodes=10 \
-  --dataset.streaming_encoding=true \
-  --dataset.encoder_threads=2 \
-  # --dataset.camera_encoder.vcodec=auto \
+  --task="Grasp a lego block and put it in the bin." \ # <- Use the same task description you used in your dataset recording
+  # <- RTC optional, use when running on low power hardware \
+  # --inference.type=rtc \
+  # --inference.rtc.execution_horizon=10 \
+  # --inference.rtc.max_guidance_weight=10.0 \
   # <- Teleop optional if you want to teleoperate in between episodes \
   # --teleop.type=so100_leader \
   # --teleop.port=/dev/ttyACM0 \
   # --teleop.id=my_red_leader_arm \
+  # --display_data=true #optional use if you want to see the camera stream \
   --policy.path=HF_USER/FINETUNE_MODEL_NAME # <- Use your fine-tuned model
 ```
 

docs/source/topreward.mdx

@@ -1,177 +0,0 @@
-# TOPReward
-
-TOPReward is a **zero-shot reward model** that extracts token log-probabilities from an off-the-shelf vision-language model (VLM) as a robotic reward signal. Given a video trajectory and a task instruction, it returns the VLM's log-likelihood that the instruction is true — no fine-tuning required.
-
-**Paper**: [TOPReward: Token Probabilities as Hidden Zero-Shot Rewards for Robotics](https://arxiv.org/abs/2602.19313)
-**Project**: [topreward.github.io](https://topreward.github.io/webpage/)
-**Original code**: [github.com/TOPReward/TOPReward](https://github.com/TOPReward/TOPReward)
-**Default backbone**: [Qwen/Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct)
-
-## Overview
-
-TOPReward asks a generic VLM how likely a task instruction is, **conditioned on the video** of a robot trying to complete that task. Concretely, given:
-
-- A trajectory video (a sequence of frames).
-- A task instruction (e.g. _"open the drawer"_).
-
-it builds a chat prompt of the form
-
-```text
-<video>
-"The above video shows a robot manipulation trajectory that completes the
- following task: <instruction> Decide whether the above statement is True
- or not. The answer is: True"
-```
-
-forwards it through the VLM, label-masks everything except the very last token, and reads back the log-probability of that token — by default the literal `"True"` that closes the suffix template. The resulting `log P("True" | video + prompt + instruction)` is the reward.
-
-Because the method only depends on a frozen VLM, TOPReward is **zero-shot**: there are no fine-tuned weights to host. The "model" in LeRobot is a small wrapper around `transformers`' `Qwen3VLForConditionalGeneration` plus the label-masking logic. The processor owns the tokeniser and builds the full chat prompt (EO-1/Robometer pattern).
-
-## What the LeRobot integration covers
-
-- Standard `reward_model.type=topreward` configuration through LeRobot.
-- VLM loading via the `transformers` `Qwen3VLForConditionalGeneration` API.
-- Prompt assembly + tokenisation in the processor (matching upstream `QwenClient.compute_instruction_reward`).
-- `compute_reward()` returns one scalar log-prob per sample.
-- LeRobot reward-model save/load — `save_pretrained` writes only `config.json` (the VLM is identified by `vlm_name`).
-- An offline labeling script that writes a `topreward_progress.parquet` (SARM-compatible schema) for RA-BC and overlay.
-
-The current LeRobot port supports the **Qwen3-VL client only**. Other upstream clients (Gemini, OpenAI, Gemma, Molmo) can be added as follow-up extras.
-
-## Installation Requirements
-
-1. Install LeRobot following the [Installation Guide](./installation).
-2. Install the TOPReward optional extra:
-
-```bash
-pip install -e ".[topreward]"
-```
-
-or, with `uv` from a source checkout:
-
-```bash
-uv sync --extra topreward
-```
-
-This pulls in `transformers`. The first time you run TOPReward, Hugging Face will also download the VLM weights from the Hub (~16 GB for Qwen3-VL-8B-Instruct). A GPU is strongly recommended.
-
-## Model Inputs and Outputs
-
-TOPReward expects:
-
-- A trajectory video or sequence of frames.
-- A natural-language task description.
-
-In LeRobot datasets the preprocessor reads:
-
-| Config field              | Default                     | Meaning                                       |
-| ------------------------- | --------------------------- | --------------------------------------------- |
-| `reward_model.image_key`  | `observation.images.top`    | Camera observation used by TOPReward          |
-| `reward_model.task_key`   | `task`                      | Key in complementary data for the task string |
-| `reward_model.max_frames` | `16`                        | Cap on frames per sample                      |
-| `reward_model.fps`        | `2.0`                       | Metadata passed to the Qwen video processor   |
-| `reward_model.vlm_name`   | `Qwen/Qwen3-VL-8B-Instruct` | Hugging Face Hub id of the underlying VLM     |
-
-The model returns:
-
-- `compute_reward(batch)`: one log-probability per sample. Higher = better task-video alignment. When `success_threshold` is finite, returns the binary thresholded value instead.
-
-## Usage
-
-### Load the reward model directly
-
-```python
-from lerobot.rewards.topreward import TOPRewardConfig, TOPRewardModel
-
-cfg = TOPRewardConfig(
-    vlm_name="Qwen/Qwen3-VL-8B-Instruct",
-    device="cuda",
-)
-reward_model = TOPRewardModel(cfg)
-```
-
-### Use the reward factory
-
-```python
-from lerobot.rewards import make_reward_model, make_reward_model_config, make_reward_pre_post_processors
-
-cfg = make_reward_model_config(
-    "topreward",
-    vlm_name="Qwen/Qwen3-VL-8B-Instruct",
-    device="cuda",
-    image_key="observation.images.top",
-)
-reward_model = make_reward_model(cfg)
-preprocessor, postprocessor = make_reward_pre_post_processors(cfg)
-```
-
-The preprocessor tokenises the full prompt (video + prefix + instruction suffix), writes Qwen-VL tensors + `prompt_length` under `observation.topreward.*`. The model reads those tensors, label-masks based on `prompt_length`, and extracts the log-prob reward.
-
-### Offline dataset labeling
-
-Write a `topreward_progress.parquet` for RA-BC training and overlay videos:
-
-```bash
-# Sparse-dense (15 anchors per episode, matches upstream)
-uv run python -m lerobot.rewards.topreward.compute_rabc_weights \
-    --dataset-repo-id lerobot/libero_10_image \
-    --num-samples 15 \
-    --device cuda
-```
-
-Then render the progress overlay for any episode:
-
-```bash
-uv run examples/dataset/create_progress_videos.py \
-    --repo-id lerobot/libero_10_image \
-    --episode 0 \
-    --progress-file topreward_progress.parquet \
-    --gif
-```
-
-## Configuration Notes
-
-### Prompt knobs
-
-The default prompt mirrors the upstream paper:
-
-```text
-prompt_prefix = "The above video shows a robot manipulation trajectory that completes the following task: "
-prompt_suffix_template = "{instruction} Decide whether the above statement is True or not. The answer is: True"
-```
-
-Both are exposed on `TOPRewardConfig` for ablation. The suffix template **must** contain `{instruction}`.
-
-### Chat template
-
-`add_chat_template=True` wraps the full prompt (including instruction) with the tokenizer's chat template before tokenisation. Default is `False`, matching the upstream paper's main experiments.
-
-## Limitations
-
-- The current LeRobot port is **inference-only and zero-shot**; `forward()` is not overridden and `is_trainable` returns `False`.
-- Only the **Qwen3-VL family** is supported; other upstream clients are out of scope.
-- TOPReward inherits the underlying VLM's biases.
-
-## References
-
-- [TOPReward project page](https://topreward.github.io/webpage/)
-- [TOPReward paper](https://arxiv.org/abs/2602.19313)
-- [Original TOPReward code](https://github.com/TOPReward/TOPReward)
-- [Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct)
-
-## Citation
-
-```bibtex
-@article{chen2026topreward,
-  title={TOPReward: Token Probabilities as Hidden Zero-Shot Rewards for Robotics},
-  author={Chen, Shirui and Harrison, Cole and Lee, Ying-Chun and Yang, Angela Jin and
-          Ren, Zhongzheng and Ratliff, Lillian J and Duan, Jiafei and Fox, Dieter and
-          Krishna, Ranjay},
-  journal={arXiv preprint arXiv:2602.19313},
-  year={2026}
-}
-```
-
-## License
-
-The original TOPReward codebase is MIT-licensed. The LeRobot port follows the LeRobot Apache 2.0 license; the wrapped Qwen3-VL weights are subject to the original Qwen license.

examples/dataset/create_progress_videos.py

@@ -15,10 +15,12 @@
 # limitations under the License.
 
 """
-Create MP4 (or GIF) videos with sarm_progress overlay for specified episodes.
+Create MP4 (or GIF) videos with per-frame progress overlay for specified episodes.
 
 Downloads datasets from HuggingFace, seeks directly into the episode segment
 of the source video, draws a progress line on each frame, and writes the result.
+The progress data is read from a parquet file that lives alongside the dataset
+(configurable via ``--progress-file``).
 
 Usage:
     python examples/dataset/create_progress_videos.py \
@@ -56,22 +58,26 @@ SCORE_FONT_SCALE = 0.8
 TASK_FONT_SCALE = 0.55
 
 
-def download_episode_metadata(repo_id: str, episode: int) -> Path:
-    """Download only the metadata and sarm_progress files for a dataset.
+def download_episode_metadata(
+    repo_id: str, episode: int, progress_file: str = "sarm_progress.parquet"
+) -> Path:
+    """Download only the metadata and per-frame progress file for a dataset.
 
     Args:
         repo_id: HuggingFace dataset repository ID.
         episode: Episode index (used for logging only; all meta is fetched).
+        progress_file: Filename of the per-frame progress parquet inside the
+            dataset repo.
 
     Returns:
         Local cache path for the downloaded snapshot.
     """
-    logging.info("[1/4] Downloading metadata for %s (episode %d) ...", repo_id, episode)
+    logging.info("[1/4] Downloading metadata + %s for %s (episode %d) ...", progress_file, repo_id, episode)
     local_path = Path(
         snapshot_download(
             repo_id=repo_id,
             repo_type="dataset",
-            allow_patterns=["meta/**", "sarm_progress.parquet"],
+            allow_patterns=["meta/**", progress_file],
             ignore_patterns=["*.mp4"],
         )
     )
@@ -215,25 +221,28 @@ def download_video_file(repo_id: str, local_path: Path, video_rel: str) -> Path:
     return video_path
 
 
-def load_progress_data(local_path: Path, episode: int) -> np.ndarray | None:
-    """Load sarm_progress values for an episode.
+def load_progress_data(
+    local_path: Path, episode: int, progress_file: str = "sarm_progress.parquet"
+) -> np.ndarray | None:
+    """Load per-frame progress values for an episode.
 
     Args:
         local_path: Dataset cache root.
         episode: Episode index.
+        progress_file: Filename of the per-frame progress parquet.
 
     Returns:
         Sorted (N, 2) array of (frame_index, progress), or None if unavailable.
     """
-    parquet_path = local_path / "sarm_progress.parquet"
+    parquet_path = local_path / progress_file
     if not parquet_path.exists():
-        logging.warning("sarm_progress.parquet not found")
+        logging.warning("%s not found", progress_file)
         return None
     df = pd.read_parquet(parquet_path)
-    logging.info("   sarm_progress.parquet columns: %s", list(df.columns))
+    logging.info("   %s columns: %s", progress_file, list(df.columns))
     episode_df = df[df["episode_index"] == episode].copy()
     if episode_df.empty:
-        logging.warning("No sarm_progress rows for episode %d", episode)
+        logging.warning("No progress rows for episode %d in %s", episode, progress_file)
         return None
     episode_df = episode_df.sort_values("frame_index")
 
@@ -576,6 +585,7 @@ def process_dataset(
     camera_key: str | None,
     output_dir: Path,
     create_gif: bool = False,
+    progress_file: str = "sarm_progress.parquet",
 ) -> Path | None:
     """Full pipeline: download, extract metadata, composite progress, write output.
 
@@ -585,6 +595,8 @@ def process_dataset(
         camera_key: Camera key to use, or None for auto-selection.
         output_dir: Directory to write output files.
         create_gif: If True, also generate a GIF from the MP4.
+        progress_file: Filename of the per-frame progress parquet inside the
+            dataset repo.
 
     Returns:
         Path to the final output file, or None on failure.
@@ -592,7 +604,7 @@ def process_dataset(
     safe_name = repo_id.replace("/", "_")
     logging.info("Processing: %s  |  episode %d", repo_id, episode)
 
-    local_path = download_episode_metadata(repo_id, episode)
+    local_path = download_episode_metadata(repo_id, episode, progress_file)
     logging.info("   Local cache: %s", local_path)
 
     episode_meta = load_episode_meta(local_path, episode, camera_key)
@@ -600,9 +612,9 @@ def process_dataset(
 
     video_path = download_video_file(repo_id, local_path, episode_meta["video_rel"])
 
-    progress_data = load_progress_data(local_path, episode)
+    progress_data = load_progress_data(local_path, episode, progress_file)
     if progress_data is None:
-        logging.error("Could not load sarm_progress data. Skipping overlay.")
+        logging.error("Could not load progress data from %s. Skipping overlay.", progress_file)
         return None
 
     logging.info("   Progress frames: %d", len(progress_data))
@@ -627,7 +639,7 @@ def process_dataset(
 
 def main() -> None:
     parser = argparse.ArgumentParser(
-        description="Create MP4/GIF videos with sarm_progress overlay for dataset episodes."
+        description="Create MP4/GIF videos with per-frame progress overlay for dataset episodes."
     )
     parser.add_argument(
         "--repo-id",
@@ -658,6 +670,15 @@ def main() -> None:
         action="store_true",
         help="Also generate a GIF from the MP4 output.",
     )
+    parser.add_argument(
+        "--progress-file",
+        type=str,
+        default="sarm_progress.parquet",
+        help=(
+            "Filename of the per-frame progress parquet inside the dataset repo "
+            "(default: 'sarm_progress.parquet')."
+        ),
+    )
     args = parser.parse_args()
 
     logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
@@ -670,6 +691,7 @@ def main() -> None:
         camera_key=args.camera_key,
         output_dir=args.output_dir,
         create_gif=args.gif,
+        progress_file=args.progress_file,
     )
 
     if result:

pyproject.toml

@@ -138,7 +138,9 @@ dataset_viz = ["lerobot[dataset]", "lerobot[viz]"]
 # Common
 av-dep = ["av>=15.0.0,<16.0.0"]
 pygame-dep = ["pygame>=2.5.1,<2.7.0"]
-placo-dep = ["placo>=0.9.6,<0.9.17"]
+# NOTE: 0.9.16 links against liburdfdom_sensor.so.4, which is unavailable on Ubuntu 24.04
+# (noble ships urdfdom 3.x). Cap below 0.9.16 until system urdfdom 4.x is broadly available.
+placo-dep = ["placo>=0.9.6,<0.9.16"]
 transformers-dep = ["transformers>=5.4.0,<5.6.0"]
 grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.0"]
 can-dep = ["python-can>=4.2.0,<5.0.0"]
@@ -209,7 +211,6 @@ groot = [
     "flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'"
 ]
 sarm = ["lerobot[transformers-dep]", "pydantic>=2.0.0,<3.0.0", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"]
-topreward = ["lerobot[transformers-dep]"]
 xvla = ["lerobot[transformers-dep]"]
 eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"]
 hilserl = ["lerobot[transformers-dep]", "lerobot[dataset]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
@@ -287,7 +288,6 @@ all = [
     "lerobot[libero]; sys_platform == 'linux'",
     "lerobot[metaworld]",
     "lerobot[sarm]",
-    "lerobot[topreward]",
     "lerobot[peft]",
     # "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2
 ]

src/lerobot/configs/parser.py

@@ -255,8 +255,7 @@ def extract_path_fields_from_config(config_path: str, path_fields: list[str]) ->
             remaining = config_data[field]
             if remaining:
                 _config_yaml_overrides[field] = _flatten_to_cli_args(remaining)
-            else:
-                del config_data[field]
+            del config_data[field]
             modified = True
 
     if not modified:
@@ -311,7 +310,13 @@ def wrap(config_path: Path | None = None) -> Callable[[F], F]:
                     cli_args = filter_arg("config_path", cli_args)
                     cfg = argtype.from_pretrained(config_path_cli, cli_args=cli_args)
                 else:
-                    cfg = draccus.parse(config_class=argtype, config_path=config_path, args=cli_args)
+                    if config_path_cli:
+                        cli_args = filter_arg("config_path", cli_args)
+                    cfg = draccus.parse(
+                        config_class=argtype,
+                        config_path=config_path_cli or config_path,
+                        args=cli_args,
+                    )
             response = fn(cfg, *args, **kwargs)
             return response
 

src/lerobot/datasets/dataset_writer.py

@@ -250,7 +250,14 @@ class DatasetWriter:
         for key, ft in self._meta.features.items():
             if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]:
                 continue
-            episode_buffer[key] = np.stack(episode_buffer[key])
+            stacked_values = np.stack(episode_buffer[key])
+
+            # `shape=(1,)` numeric features are serialized as `datasets.Value`, which expects scalars.
+            # Normalizing to `(N,)` keeps save semantics stable across dependency versions.
+            if tuple(ft["shape"]) == (1,) and ft["dtype"] != "string":
+                stacked_values = stacked_values.reshape(episode_length)
+
+            episode_buffer[key] = stacked_values
 
         # Wait for image writer to end, so that episode stats over images can be computed
         self._wait_image_writer()

src/lerobot/datasets/video_utils.py

@@ -17,11 +17,13 @@ import contextlib
 import glob
 import importlib
 import logging
+import os
 import queue
 import shutil
 import tempfile
 import threading
 import warnings
+from collections import OrderedDict
 from dataclasses import asdict, dataclass, field
 from fractions import Fraction
 from pathlib import Path
@@ -191,15 +193,70 @@ def decode_video_frames_pyav(
     return closest_frames
 
 
-class VideoDecoderCache:
-    """Thread-safe cache for video decoders to avoid expensive re-initialization."""
+DEFAULT_DECODER_CACHE_SIZE = 100
+"""Default LRU capacity for :class:`VideoDecoderCache`.
 
-    def __init__(self):
-        self._cache: dict[str, tuple[Any, Any]] = {}
+Sized to comfortably hold a small rolling window of episodes worth of decoders
+(typical recipes: 2-4 cameras per episode × tens of episodes in flight) while
+bounding host RAM. Each cached entry retains a torchcodec ``VideoDecoder`` plus
+an open ``fsspec`` file handle — on the order of a few MB per entry. Override
+via the ``LEROBOT_VIDEO_DECODER_CACHE_SIZE`` env var or by passing ``max_size``
+to the constructor (``None`` restores the legacy unbounded behaviour).
+"""
+
+
+def _default_max_cache_size() -> int | None:
+    raw = os.environ.get("LEROBOT_VIDEO_DECODER_CACHE_SIZE")
+    if raw is None:
+        return DEFAULT_DECODER_CACHE_SIZE
+    raw = raw.strip().lower()
+    if raw in ("", "none", "unbounded", "-1"):
+        return None
+    try:
+        value = int(raw)
+    except ValueError as e:
+        raise ValueError(
+            f"LEROBOT_VIDEO_DECODER_CACHE_SIZE must be an integer, 'none', or '-1'; got {raw!r}"
+        ) from e
+    if value <= 0:
+        raise ValueError(f"LEROBOT_VIDEO_DECODER_CACHE_SIZE must be positive; got {value}")
+    return value
+
+
+class VideoDecoderCache:
+    """Thread-safe LRU cache for torchcodec ``VideoDecoder`` instances.
+
+    Cached entries hold a ``VideoDecoder`` plus the open ``fsspec`` file handle
+    backing it. When the cache is full and a new path is requested, the
+    least-recently-used entry is evicted and its file handle is closed. This
+    bounds host-RAM growth when iterating over datasets with many distinct
+    video files (otherwise each ``DataLoader`` worker pins every decoder it has
+    ever opened until the process exits).
+
+    Args:
+        max_size: Maximum number of decoders to retain. ``None`` disables
+            eviction and restores legacy unbounded behaviour. Defaults to the
+            value of ``LEROBOT_VIDEO_DECODER_CACHE_SIZE`` if set, otherwise
+            :data:`DEFAULT_DECODER_CACHE_SIZE`.
+    """
+
+    _SENTINEL: ClassVar[object] = object()
+
+    def __init__(self, max_size: int | None | object = _SENTINEL):
+        if max_size is VideoDecoderCache._SENTINEL:
+            max_size = _default_max_cache_size()
+        if max_size is not None and max_size <= 0:
+            raise ValueError(f"max_size must be positive or None; got {max_size}")
+        self.max_size: int | None = max_size  # type: ignore[assignment]
+        self._cache: OrderedDict[str, tuple[Any, Any]] = OrderedDict()
         self._lock = Lock()
 
+    def __contains__(self, video_path: object) -> bool:
+        with self._lock:
+            return str(video_path) in self._cache
+
     def get_decoder(self, video_path: str):
-        """Get a cached decoder or create a new one."""
+        """Get a cached decoder or create a new one, evicting LRU if at capacity."""
         if importlib.util.find_spec("torchcodec"):
             from torchcodec.decoders import VideoDecoder
         else:
@@ -211,22 +268,36 @@ class VideoDecoderCache:
         video_path = str(video_path)
 
         with self._lock:
-            if video_path not in self._cache:
-                file_handle = fsspec.open(video_path).__enter__()
-                try:
-                    decoder = VideoDecoder(file_handle, seek_mode="approximate")
-                except Exception:
-                    file_handle.close()
-                    raise
-                self._cache[video_path] = (decoder, file_handle)
+            entry = self._cache.get(video_path)
+            if entry is not None:
+                self._cache.move_to_end(video_path)
+                return entry[0]
 
-            return self._cache[video_path][0]
+            file_handle = fsspec.open(video_path).__enter__()
+            try:
+                decoder = VideoDecoder(file_handle, seek_mode="approximate")
+            except Exception:
+                file_handle.close()
+                raise
+            self._cache[video_path] = (decoder, file_handle)
+
+            # Evict LRU entries until we are back under the cap. We close
+            # evicted file handles immediately; the associated ``VideoDecoder``
+            # is released to the GC when its last reference goes away.
+            if self.max_size is not None:
+                while len(self._cache) > self.max_size:
+                    _evicted_path, (_evicted_decoder, evicted_handle) = self._cache.popitem(last=False)
+                    with contextlib.suppress(Exception):
+                        evicted_handle.close()
+
+            return decoder
 
     def clear(self):
-        """Clear the cache and close file handles."""
+        """Clear the cache and close all file handles."""
         with self._lock:
             for _, file_handle in self._cache.values():
-                file_handle.close()
+                with contextlib.suppress(Exception):
+                    file_handle.close()
             self._cache.clear()
 
     def size(self) -> int:

src/lerobot/model/kinematics.py

@@ -18,12 +18,25 @@ from typing import TYPE_CHECKING
 
 import numpy as np
 
-from lerobot.utils.import_utils import _placo_available, require_package
+from lerobot.utils.import_utils import require_package
 
-if TYPE_CHECKING or _placo_available:
+_placo_runtime_error: ImportError | None = None
+
+if TYPE_CHECKING:
     import placo  # type: ignore[import-not-found]
 else:
-    placo = None
+    try:
+        import placo  # type: ignore[import-not-found]
+    except ImportError as _placo_import_err:
+        placo = None
+        _placo_runtime_error = _placo_import_err
+
+
+def _raise_if_placo_unusable() -> None:
+    if placo is None and _placo_runtime_error is not None:
+        raise ImportError(
+            f"placo is installed but failed to import: {_placo_runtime_error!s}"
+        ) from _placo_runtime_error
 
 
 class RobotKinematics:
@@ -44,6 +57,7 @@ class RobotKinematics:
             joint_names (list[str] | None): List of joint names to use for the kinematics solver
         """
         require_package("placo", extra="placo-dep")
+        _raise_if_placo_unusable()
 
         self.robot = placo.RobotWrapper(urdf_path)
         self.solver = placo.KinematicsSolver(self.robot)

src/lerobot/motors/robstride/robstride.py

@@ -43,6 +43,7 @@ from .tables import (
     CAN_CMD_SET_ZERO,
     DEFAULT_BAUDRATE,
     DEFAULT_TIMEOUT_MS,
+    HANDSHAKE_TIMEOUT_S,
     MODEL_RESOLUTION,
     MOTOR_LIMIT_PARAMS,
     NORMALIZED_DATA,
@@ -215,14 +216,16 @@ class RobstrideMotorsBus(MotorsBusBase):
             self._is_connected = False
             raise ConnectionError(f"Failed to connect to CAN bus: {e}") from e
 
-    def _query_status_via_clear_fault(self, motor: NameOrID) -> tuple[bool, can.Message | None]:
+    def _query_status_via_clear_fault(
+        self, motor: NameOrID, timeout: float = RUNNING_TIMEOUT
+    ) -> tuple[bool, can.Message | None]:
         motor_name = self._get_motor_name(motor)
         motor_id = self._get_motor_id(motor_name)
         recv_id = self._get_motor_recv_id(motor_name)
         data = [0xFF] * 7 + [CAN_CMD_CLEAR_FAULT]
         msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
         self._bus().send(msg)
-        return self._recv_status_via_clear_fault(expected_recv_id=recv_id)
+        return self._recv_status_via_clear_fault(expected_recv_id=recv_id, timeout=timeout)
 
     def _recv_status_via_clear_fault(
         self, expected_recv_id: int | None = None, timeout: float = RUNNING_TIMEOUT
@@ -280,7 +283,7 @@ class RobstrideMotorsBus(MotorsBusBase):
         faulted_motors = []
 
         for motor_name in self.motors:
-            has_fault, msg = self._query_status_via_clear_fault(motor_name)
+            has_fault, msg = self._query_status_via_clear_fault(motor_name, timeout=HANDSHAKE_TIMEOUT_S)
             if msg is None:
                 missing_motors.append(motor_name)
             elif has_fault:
@@ -505,6 +508,87 @@ class RobstrideMotorsBus(MotorsBusBase):
 
         return responses
 
+    def _recv_all_messages_until_quiet(
+        self,
+        *,
+        timeout: float = RUNNING_TIMEOUT,
+        max_messages: int = 4096,
+    ) -> list[can.Message]:
+        """
+        Receive frames until the bus goes quiet.
+
+        Args:
+            timeout: Poll timeout used for each recv() call. Collection stops
+                when one recv() times out (quiet gap).
+            max_messages: Safety cap to prevent unbounded loops.
+        """
+        out: list[can.Message] = []
+        max_messages = max(1, max_messages)
+        timeout = max(0.0, timeout)
+
+        try:
+            while len(out) < max_messages:
+                msg = self._bus().recv(timeout=timeout)
+                if msg is None:
+                    break
+                out.append(msg)
+        except (can.CanError, OSError) as e:
+            logger.debug(f"Error draining CAN RX queue on {self.port}: {e}")
+
+        return out
+
+    def _process_feedback_messages(self, messages: list[can.Message]) -> set[int]:
+        """
+        Decode all received feedback frames and update cached motor states.
+
+        Returns:
+            Set of payload recv_ids that were successfully mapped to motors.
+        """
+        processed_recv_ids: set[int] = set()
+        for msg in messages:
+            if len(msg.data) < 1:
+                logger.debug(
+                    f"Dropping short CAN frame on {self.port} "
+                    f"(arb=0x{int(msg.arbitration_id):02X}, data={bytes(msg.data).hex()})"
+                )
+                continue
+
+            recv_id = int(msg.data[0])
+            motor_name = self._recv_id_to_motor.get(recv_id)
+            if motor_name is None:
+                logger.debug(
+                    f"Unmapped CAN frame on {self.port} "
+                    f"(arb=0x{int(msg.arbitration_id):02X}, recv_id=0x{recv_id:02X}, data={bytes(msg.data).hex()})"
+                )
+                continue
+
+            self._process_response(motor_name, msg)
+            processed_recv_ids.add(recv_id)
+
+        return processed_recv_ids
+
+    def flush_rx_queue(self, poll_timeout_s: float = 0.0005, max_messages: int = 4096) -> int:
+        """
+        Drain pending RX frames from the CAN interface.
+
+        This is used by higher-level controllers to drop stale feedback before issuing
+        a fresh read cycle, so subsequent state reads are based on most recent replies.
+        It should also be called once when a controller instance is created/connected,
+        to clear residual frames left on the interface from previous sessions.
+        """
+        drained = 0
+        poll_timeout_s = max(0.0, poll_timeout_s)
+        max_messages = max(1, max_messages)
+        try:
+            while drained < max_messages:
+                msg = self._bus().recv(timeout=poll_timeout_s)
+                if msg is None:
+                    break
+                drained += 1
+        except (can.CanError, OSError) as e:
+            logger.debug(f"Failed to flush CAN RX queue on {self.port}: {e}")
+        return drained
+
     def _speed_control(
         self,
         motor: NameOrID,
@@ -644,11 +728,14 @@ class RobstrideMotorsBus(MotorsBusBase):
             msg = can.Message(arbitration_id=motor_id, data=data, is_extended_id=False)
             self._bus().send(msg)
             recv_id_to_motor[self._get_motor_recv_id(motor)] = motor_name
+        # Read every feedback frame until RX goes quiet, then decode all of them.
+        # This avoids dropping useful frames when responses from different motors interleave.
+        messages = self._recv_all_messages_until_quiet()
+        processed_recv_ids = self._process_feedback_messages(messages)
 
-        responses = self._recv_all_responses(list(recv_id_to_motor.keys()), timeout=RUNNING_TIMEOUT)
         for recv_id, motor_name in recv_id_to_motor.items():
-            if msg := responses.get(recv_id):
-                self._process_response(motor_name, msg)
+            if recv_id not in processed_recv_ids:
+                logger.warning(f"Packet drop: {motor_name} (ID: 0x{recv_id:02X}). Using last known state.")
 
     def _float_to_uint(self, x: float, x_min: float, x_max: float, bits: int) -> int:
         """Convert float to unsigned integer for CAN transmission."""
@@ -711,7 +798,10 @@ class RobstrideMotorsBus(MotorsBusBase):
         try:
             self._decode_motor_state(msg.data)
         except Exception as e:
-            logger.warning(f"Failed to decode response from {motor}: {e}")
+            logger.warning(
+                f"Failed to decode response from {motor} "
+                f"(arb=0x{int(msg.arbitration_id):02X}, data={bytes(msg.data).hex()}): {e}"
+            )
 
     def _get_cached_value(self, motor: str, data_name: str) -> Value:
         """Retrieve a specific value from the state cache."""
@@ -848,20 +938,12 @@ class RobstrideMotorsBus(MotorsBusBase):
             self._bus().send(msg)
             updated_motors.append(motor)
 
-        expected_recv_ids = [self._get_motor_recv_id(motor) for motor in updated_motors]
-        responses = self._recv_all_responses(expected_recv_ids, timeout=RUNNING_TIMEOUT)
-
-        for response in responses.values():
-            payload_motor_name = self._recv_id_to_motor.get(response.data[0])
-            if payload_motor_name is not None:
-                self._process_response(payload_motor_name, response)
-            else:
-                # Fallback: still attempt to decode based on payload byte0 mapping.
-                self._decode_motor_state(response.data)
+        messages = self._recv_all_messages_until_quiet()
+        processed_recv_ids = self._process_feedback_messages(messages)
 
         for motor in updated_motors:
             recv_id = self._get_motor_recv_id(motor)
-            if recv_id not in responses:
+            if recv_id not in processed_recv_ids:
                 logger.warning(f"Packet drop: {motor} (ID: 0x{recv_id:02X}). Using last known state.")
 
     def read_calibration(self) -> dict[str, MotorCalibration]:

src/lerobot/motors/robstride/tables.py

@@ -114,7 +114,8 @@ CAN_CMD_SAVE_PARAM = 0xAA
 CAN_PARAM_ID = 0x7FF
 
 
-RUNNING_TIMEOUT = 0.001
+RUNNING_TIMEOUT = 0.003
+HANDSHAKE_TIMEOUT_S = 0.05
 PARAM_TIMEOUT = 0.01
 
 STATE_CACHE_TTL_S = 0.02

src/lerobot/policies/groot/eagle2_hg_model/modeling_eagle2_5_vl.py

@@ -60,6 +60,7 @@ class Eagle25VLPreTrainedModel(PreTrainedModel):
         "SiglipEncoderLayer",
     ]
     _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn = True
     _supports_flash_attn_2 = True
     _supports_cache_class = True
     _supports_static_cache = True

src/lerobot/policies/groot/eagle2_hg_model/processing_eagle2_5_vl.py

@@ -124,7 +124,6 @@ class Eagle25VLProcessor(ProcessorMixin):
         "videos_kwargs",
         "text_kwargs",
     ]
-    image_processor_class = "AutoImageProcessor"
     tokenizer_class = "AutoTokenizer"
 
     def __init__(

src/lerobot/policies/groot/groot_n1.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
@@ -26,9 +26,14 @@ from lerobot.utils.import_utils import _transformers_available
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
+    from huggingface_hub.dataclasses import strict
     from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
     from transformers.feature_extraction_utils import BatchFeature
 else:
+
+    def strict(cls):
+        return cls
+
     AutoConfig = None
     AutoModel = None
     PretrainedConfig = object
@@ -173,19 +178,20 @@ N_COLOR_CHANNELS = 3
 
 
 # config
+@strict
 class GR00TN15Config(PretrainedConfig):
     model_type = "gr00t_n1_5"
 
-    backbone_cfg: dict
-    action_head_cfg: dict
-    action_horizon: int
-    action_dim: int
+    backbone_cfg: dict[str, Any] | None = None
+    action_head_cfg: dict[str, Any] | None = None
+    action_horizon: int = 0
+    action_dim: int = 0
     compute_dtype: str = "float32"
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        for key, value in kwargs.items():
-            setattr(self, key, value)
+    def __post_init__(self, **kwargs):
+        self.backbone_cfg = {} if self.backbone_cfg is None else self.backbone_cfg
+        self.action_head_cfg = {} if self.action_head_cfg is None else self.action_head_cfg
+        super().__post_init__(**kwargs)
 
 
 # real model

src/lerobot/policies/groot/processor_groot.py

@@ -206,7 +206,11 @@ def _build_eagle_processor(tokenizer_assets_repo: str = DEFAULT_TOKENIZER_ASSETS
             "Vendor files are copied during model creation. Create the policy/model first, "
             "or call ensure_eagle_cache_ready() before building processors."
         )
-    proc = AutoProcessor.from_pretrained(str(cache_dir), trust_remote_code=True, use_fast=True)
+    proc = AutoProcessor.from_pretrained(
+        str(cache_dir),
+        trust_remote_code=True,
+        fix_mistral_regex=False,
+    )
     proc.tokenizer.padding_side = "left"
     return proc
 

src/lerobot/policies/pi0/modeling_pi0.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import builtins
-import copy
 import logging
 import math
 from collections import deque
@@ -30,6 +29,7 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
+    from transformers.cache_utils import DynamicCache
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
 
@@ -41,6 +41,7 @@ if TYPE_CHECKING or _transformers_available:
     )
 else:
     CONFIG_MAPPING = None
+    DynamicCache = None
     modeling_gemma = None
     PiGemmaForCausalLM = None
     _gated_residual = None
@@ -141,6 +142,15 @@ def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (
     return att_2d_masks & pad_2d_masks
 
 
+def clone_past_key_values(past_key_values):
+    """Clone the DynamicCache returned by prefix prefill for compiled denoising."""
+    return DynamicCache(
+        tuple(
+            (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
+        )
+    )
+
+
 def pad_vector(vector, new_dim):
     """Pad the last dimension of a vector to new_dim with zeros.
 
@@ -227,16 +237,13 @@ def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
 
 
 # Define the complete layer computation function for gradient checkpointing
-def compute_layer_complete(
-    layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
-):
-    models = [paligemma.model.language_model, gemma_expert.model]
+def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []
     key_states = []
     value_states = []
     gates = []
     for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
+        layer = layers[i]
         hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
         gates.append(gate)
         input_shape = hidden_states.shape[:-1]
@@ -258,15 +265,16 @@ def compute_layer_complete(
         device=query_states.device,
         dtype=query_states.dtype,
     )
-    cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
+    cos, sin = rotary_emb(dummy_tensor, position_ids)
     query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
         query_states, key_states, cos, sin, unsqueeze_dim=1
     )
     batch_size = query_states.shape[0]
-    scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
+    paligemma_layer = layers[0]
+    scaling = paligemma_layer.self_attn.scaling
     # Attention computation
     att_output, _ = modeling_gemma.eager_attention_forward(
-        paligemma.model.language_model.layers[layer_idx].self_attn,
+        paligemma_layer.self_attn,
         query_states,
         key_states,
         value_states,
@@ -274,13 +282,13 @@ def compute_layer_complete(
         scaling,
     )
     # Get head_dim from the current layer, not from the model
-    head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
+    head_dim = paligemma_layer.self_attn.head_dim
     att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
     # Process layer outputs
     outputs_embeds = []
     start_pos = 0
     for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
+        layer = layers[i]
         end_pos = start_pos + hidden_states.shape[1]
         if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
             att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
@@ -488,8 +496,9 @@ class PaliGemmaWithExpertModel(
             prefix_output = None
             prefix_past_key_values = None
         else:
-            models = [self.paligemma.model.language_model, self.gemma_expert.model]
-            num_layers = self.paligemma.config.text_config.num_hidden_layers
+            paligemma_layers = self.paligemma.model.language_model.layers
+            gemma_expert_layers = self.gemma_expert.model.layers
+            rotary_emb = self.paligemma.model.language_model.rotary_emb
 
             # Check if gradient checkpointing is enabled for any of the models
             use_gradient_checkpointing = (
@@ -499,36 +508,39 @@ class PaliGemmaWithExpertModel(
             ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
 
             # Process all layers with gradient checkpointing if enabled
-            for layer_idx in range(num_layers):
+            for layers in zip(paligemma_layers, gemma_expert_layers, strict=True):
                 if use_gradient_checkpointing:
                     inputs_embeds = torch.utils.checkpoint.checkpoint(
                         compute_layer_complete,
-                        layer_idx,
                         inputs_embeds,
                         attention_mask,
                         position_ids,
                         adarms_cond,
                         use_reentrant=False,
                         preserve_rng_state=False,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
+                        layers=layers,
+                        rotary_emb=rotary_emb,
                     )
                 else:
                     inputs_embeds = compute_layer_complete(
-                        layer_idx,
                         inputs_embeds,
                         attention_mask,
                         position_ids,
                         adarms_cond,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
+                        layers=layers,
+                        rotary_emb=rotary_emb,
                     )
 
             # final norm
+            final_norms = (
+                self.paligemma.model.language_model.norm,
+                self.gemma_expert.model.norm,
+            )
+
             def compute_final_norms(inputs_embeds, adarms_cond):
                 outputs_embeds = []
                 for i, hidden_states in enumerate(inputs_embeds):
-                    out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
+                    out_emb, _ = layernorm_forward(final_norms[i], hidden_states, adarms_cond[i])
                     outputs_embeds.append(out_emb)
                 return outputs_embeds
 
@@ -907,7 +919,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
-        past_key_values = copy.deepcopy(past_key_values)
+        past_key_values = clone_past_key_values(past_key_values)
         outputs_embeds, _ = self.paligemma_with_expert.forward(
             attention_mask=full_att_2d_masks_4d,
             position_ids=position_ids,

src/lerobot/policies/pi05/modeling_pi05.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import builtins
-import copy
 import logging
 import math
 from collections import deque
@@ -30,6 +29,7 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
+    from transformers.cache_utils import DynamicCache
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
 
@@ -41,6 +41,7 @@ if TYPE_CHECKING or _transformers_available:
     )
 else:
     CONFIG_MAPPING = None
+    DynamicCache = None
     modeling_gemma = None
     PiGemmaForCausalLM = None
     _gated_residual = None
@@ -138,6 +139,15 @@ def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (
     return att_2d_masks & pad_2d_masks
 
 
+def clone_past_key_values(past_key_values):
+    """Clone the DynamicCache returned by prefix prefill for compiled denoising."""
+    return DynamicCache(
+        tuple(
+            (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
+        )
+    )
+
+
 def pad_vector(vector, new_dim):
     """Pad the last dimension of a vector to new_dim with zeros.
 
@@ -224,16 +234,13 @@ def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
 
 
 # Define the complete layer computation function for gradient checkpointing
-def compute_layer_complete(
-    layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
-):
-    models = [paligemma.model.language_model, gemma_expert.model]
+def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []
     key_states = []
     value_states = []
     gates = []
     for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
+        layer = layers[i]
         hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
         gates.append(gate)
         input_shape = hidden_states.shape[:-1]
@@ -255,15 +262,16 @@ def compute_layer_complete(
         device=query_states.device,
         dtype=query_states.dtype,
     )
-    cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
+    cos, sin = rotary_emb(dummy_tensor, position_ids)
     query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
         query_states, key_states, cos, sin, unsqueeze_dim=1
     )
     batch_size = query_states.shape[0]
-    scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
+    paligemma_layer = layers[0]
+    scaling = paligemma_layer.self_attn.scaling
     # Attention computation
     att_output, _ = modeling_gemma.eager_attention_forward(
-        paligemma.model.language_model.layers[layer_idx].self_attn,
+        paligemma_layer.self_attn,
         query_states,
         key_states,
         value_states,
@@ -271,13 +279,13 @@ def compute_layer_complete(
         scaling,
     )
     # Get head_dim from the current layer, not from the model
-    head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
+    head_dim = paligemma_layer.self_attn.head_dim
     att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
     # Process layer outputs
     outputs_embeds = []
     start_pos = 0
     for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
+        layer = layers[i]
         end_pos = start_pos + hidden_states.shape[1]
         if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
             att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
@@ -485,8 +493,9 @@ class PaliGemmaWithExpertModel(
             prefix_output = None
             prefix_past_key_values = None
         else:
-            models = [self.paligemma.model.language_model, self.gemma_expert.model]
-            num_layers = self.paligemma.config.text_config.num_hidden_layers
+            paligemma_layers = self.paligemma.model.language_model.layers
+            gemma_expert_layers = self.gemma_expert.model.layers
+            rotary_emb = self.paligemma.model.language_model.rotary_emb
 
             # Check if gradient checkpointing is enabled for any of the models
             use_gradient_checkpointing = (
@@ -496,36 +505,39 @@ class PaliGemmaWithExpertModel(
             ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
 
             # Process all layers with gradient checkpointing if enabled
-            for layer_idx in range(num_layers):
+            for layers in zip(paligemma_layers, gemma_expert_layers, strict=True):
                 if use_gradient_checkpointing:
                     inputs_embeds = torch.utils.checkpoint.checkpoint(
                         compute_layer_complete,
-                        layer_idx,
                         inputs_embeds,
                         attention_mask,
                         position_ids,
                         adarms_cond,
                         use_reentrant=False,
                         preserve_rng_state=False,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
+                        layers=layers,
+                        rotary_emb=rotary_emb,
                     )
                 else:
                     inputs_embeds = compute_layer_complete(
-                        layer_idx,
                         inputs_embeds,
                         attention_mask,
                         position_ids,
                         adarms_cond,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
+                        layers=layers,
+                        rotary_emb=rotary_emb,
                     )
 
             # final norm
+            final_norms = (
+                self.paligemma.model.language_model.norm,
+                self.gemma_expert.model.norm,
+            )
+
             def compute_final_norms(inputs_embeds, adarms_cond):
                 outputs_embeds = []
                 for i, hidden_states in enumerate(inputs_embeds):
-                    out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
+                    out_emb, _ = layernorm_forward(final_norms[i], hidden_states, adarms_cond[i])
                     outputs_embeds.append(out_emb)
                 return outputs_embeds
 
@@ -880,7 +892,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
-        past_key_values = copy.deepcopy(past_key_values)
+        past_key_values = clone_past_key_values(past_key_values)
         outputs_embeds, _ = self.paligemma_with_expert.forward(
             attention_mask=full_att_2d_masks_4d,
             position_ids=position_ids,

src/lerobot/rewards/__init__.py

@@ -21,13 +21,11 @@ from .factory import (
 )
 from .pretrained import PreTrainedRewardModel as PreTrainedRewardModel
 from .sarm.configuration_sarm import SARMConfig as SARMConfig
-from .topreward.configuration_topreward import TOPRewardConfig as TOPRewardConfig
 
 __all__ = [
     # Configuration classes
     "RewardClassifierConfig",
     "SARMConfig",
-    "TOPRewardConfig",
     # Base class
     "PreTrainedRewardModel",
     # Factory functions

src/lerobot/rewards/factory.py

@@ -26,7 +26,6 @@ from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 from .classifier.configuration_classifier import RewardClassifierConfig
 from .pretrained import PreTrainedRewardModel
 from .sarm.configuration_sarm import SARMConfig
-from .topreward.configuration_topreward import TOPRewardConfig
 
 
 def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
@@ -38,7 +37,7 @@ def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
 
     Args:
         name: The name of the reward model. Supported names are "reward_classifier",
-              "sarm", "topreward".
+              "sarm".
 
     Returns:
         The reward model class corresponding to the given name.
@@ -54,10 +53,6 @@ def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
         from lerobot.rewards.sarm.modeling_sarm import SARMRewardModel
 
         return SARMRewardModel
-    elif name == "topreward":
-        from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-        return TOPRewardModel
     else:
         try:
             return _get_reward_model_cls_from_name(name=name)
@@ -74,7 +69,7 @@ def make_reward_model_config(reward_type: str, **kwargs) -> RewardModelConfig:
 
     Args:
         reward_type: The type of the reward model. Supported types include
-                     "reward_classifier", "sarm", "topreward".
+                     "reward_classifier", "sarm".
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -87,8 +82,6 @@ def make_reward_model_config(reward_type: str, **kwargs) -> RewardModelConfig:
         return RewardClassifierConfig(**kwargs)
     elif reward_type == "sarm":
         return SARMConfig(**kwargs)
-    elif reward_type == "topreward":
-        return TOPRewardConfig(**kwargs)
     else:
         try:
             config_cls = RewardModelConfig.get_choice_class(reward_type)
@@ -169,14 +162,6 @@ def make_reward_pre_post_processors(
             dataset_meta=kwargs.get("dataset_meta"),
         )
 
-    elif isinstance(reward_cfg, TOPRewardConfig):
-        from lerobot.rewards.topreward.processor_topreward import make_topreward_pre_post_processors
-
-        return make_topreward_pre_post_processors(
-            config=reward_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
     else:
         try:
             processors = _make_processors_from_reward_model_config(

src/lerobot/rewards/topreward/__init__.py

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

src/lerobot/rewards/topreward/compute_rabc_weights.py

@@ -1,353 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 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.
-
-"""Compute per-frame TOPReward progress curves for a LeRobot dataset.
-
-For each episode, scores trajectory prefixes of increasing length using
-the TOPReward reward model, min-max normalises the raw log-prob rewards per episode,
-and writes a parquet file with one row per frame.
-
-The parquet uses the same schema as SARM's :mod:`lerobot.rewards.sarm.compute_rabc_weights`.
-
-Usage:
-    # Sparse-dense mode (15 anchors per episode, matches upstream)
-    python -m lerobot.rewards.topreward.compute_rabc_weights \\
-        --dataset-repo-id lerobot/libero_10_image \\
-        --num-samples 15
-
-    # Use a different VLM backbone
-    python -m lerobot.rewards.topreward.compute_rabc_weights \\
-        --dataset-repo-id lerobot/libero_10_image \\
-        --vlm-name Qwen/Qwen3-VL-4B-Instruct
-"""
-
-from __future__ import annotations
-
-import argparse
-import logging
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import pyarrow as pa
-import pyarrow.parquet as pq
-import torch
-from tqdm import tqdm
-
-from lerobot.datasets import LeRobotDataset
-from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig
-from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-from lerobot.rewards.topreward.processor_topreward import TOPRewardEncoderProcessorStep
-from lerobot.types import TransitionKey
-
-DEFAULT_OUTPUT_FILENAME = "topreward_progress.parquet"
-
-
-def get_reward_model_path_from_parquet(parquet_path: Path) -> str | None:
-    """Read ``reward_model_path`` from parquet metadata if available."""
-    if not parquet_path.exists():
-        return None
-    try:
-        metadata = pq.read_metadata(parquet_path).schema.to_arrow_schema().metadata
-        if metadata and b"reward_model_path" in metadata:
-            return metadata[b"reward_model_path"].decode()
-    except Exception:  # nosec B110
-        return None
-    return None
-
-
-def _resolve_task(sample: dict[str, Any], default: str) -> str:
-    """Best-effort task extraction from a dataset sample."""
-    task = sample.get("task")
-    if isinstance(task, str) and task:
-        return task
-    return default
-
-
-def normalize_rewards(rewards: list[float] | np.ndarray) -> np.ndarray:
-    """Min-max normalise raw log-prob rewards into ``[0, 1]``."""
-    rewards_arr = np.asarray(rewards, dtype=np.float64)
-    if rewards_arr.size == 0:
-        return rewards_arr.astype(np.float32)
-    if rewards_arr.size == 1:
-        return np.array([1.0], dtype=np.float32)
-    r_min, r_max = rewards_arr.min(), rewards_arr.max()
-    if r_max == r_min:
-        return np.ones_like(rewards_arr, dtype=np.float32)
-    return ((rewards_arr - r_min) / (r_max - r_min)).astype(np.float32)
-
-
-def compute_instruction_rewards_for_prefixes(
-    model: TOPRewardModel,
-    encoder: TOPRewardEncoderProcessorStep,
-    dataset: LeRobotDataset,
-    ep_start: int,
-    num_frames: int,
-    task: str,
-    image_key: str,
-    num_samples: int | None,
-    device: str,
-) -> np.ndarray:
-    """Score an episode via prefix sweep and return a per-frame normalised curve."""
-    if num_samples is None or num_samples >= num_frames:
-        prefix_lengths = np.arange(1, num_frames + 1, dtype=np.int64)
-    else:
-        prefix_lengths = np.unique(np.linspace(1, num_frames, num_samples).round().astype(np.int64))
-
-    episode_frames = torch.stack([dataset[ep_start + i][image_key] for i in range(num_frames)])
-    rewards: list[float] = []
-    for length in prefix_lengths:
-        frames = episode_frames[: int(length)].unsqueeze(0)  # (1, T, C, H, W)
-
-        transition = {
-            TransitionKey.OBSERVATION: {image_key: frames},
-            TransitionKey.COMPLEMENTARY_DATA: {"task": task},
-        }
-        encoded = encoder(transition)
-        obs = encoded[TransitionKey.OBSERVATION]
-        batch = {
-            key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in obs.items()
-        }
-
-        with torch.no_grad():
-            reward = model.compute_reward(batch)
-        rewards.append(float(reward.item()))
-
-    normalized_rewards = normalize_rewards(rewards)
-
-    if prefix_lengths.shape[0] == num_frames:
-        return normalized_rewards
-
-    return np.interp(
-        np.arange(1, num_frames + 1, dtype=np.float64),
-        prefix_lengths.astype(np.float64),
-        normalized_rewards.astype(np.float64),
-    ).astype(np.float32)
-
-
-def compute_topreward_progress(
-    dataset_repo_id: str,
-    reward_model_path: str | None = None,
-    vlm_name: str | None = None,
-    output_path: str | None = None,
-    device: str = "cuda",
-    num_samples: int | None = None,
-    fps: float | None = None,
-    episodes: list[int] | None = None,
-) -> Path:
-    """Run TOPReward over a dataset and write per-frame progress."""
-    if reward_model_path is not None:
-        logging.info(f"Loading TOPReward config from: {reward_model_path}")
-        model = TOPRewardModel.from_pretrained(reward_model_path)
-        config = model.config
-        config.device = device
-        if vlm_name is not None and vlm_name != config.vlm_name:
-            logging.info(f"Overriding vlm_name from config: {config.vlm_name} -> {vlm_name}")
-            config.vlm_name = vlm_name
-            model = TOPRewardModel(config)
-    else:
-        config_kwargs: dict[str, Any] = {"device": device}
-        if vlm_name is not None:
-            config_kwargs["vlm_name"] = vlm_name
-        if fps is not None:
-            config_kwargs["fps"] = fps
-        config = TOPRewardConfig(**config_kwargs)
-        logging.info(f"Constructing TOPReward with VLM: {config.vlm_name}")
-        model = TOPRewardModel(config)
-
-    model.to(device).eval()
-
-    encoder = TOPRewardEncoderProcessorStep(
-        vlm_name=config.vlm_name,
-        image_key=config.image_key,
-        task_key=config.task_key,
-        default_task=config.default_task,
-        max_frames=None,  # no tail-crop: we control prefix length explicitly
-        fps=config.fps,
-        prompt_prefix=config.prompt_prefix,
-        prompt_suffix_template=config.prompt_suffix_template,
-        add_chat_template=config.add_chat_template,
-        max_length=config.max_input_length,
-    )
-
-    image_key = config.image_key
-
-    logging.info(f"Loading dataset: {dataset_repo_id}")
-    dataset = LeRobotDataset(dataset_repo_id, download_videos=True)
-    logging.info(f"Dataset: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
-
-    episode_indices = list(range(dataset.num_episodes)) if episodes is None else episodes
-    logging.info(f"Processing {len(episode_indices)} episode(s)")
-
-    all_index: list[int] = []
-    all_episode: list[int] = []
-    all_frame: list[int] = []
-    all_progress: list[float] = []
-
-    for episode_idx in tqdm(episode_indices, desc="Episodes"):
-        ep = dataset.meta.episodes[episode_idx]
-        ep_start = int(ep["dataset_from_index"])
-        ep_end = int(ep["dataset_to_index"])
-        num_frames = ep_end - ep_start
-        if num_frames <= 0:
-            continue
-
-        first_sample = dataset[ep_start]
-        task = _resolve_task(first_sample, default=config.default_task or "perform the task")
-
-        per_frame = compute_instruction_rewards_for_prefixes(
-            model=model,
-            encoder=encoder,
-            dataset=dataset,
-            ep_start=ep_start,
-            num_frames=num_frames,
-            task=task,
-            image_key=image_key,
-            num_samples=num_samples,
-            device=device,
-        )
-
-        for local in range(num_frames):
-            all_index.append(ep_start + local)
-            all_episode.append(episode_idx)
-            all_frame.append(local)
-            all_progress.append(float(per_frame[local]))
-
-        if device.startswith("cuda"):
-            torch.cuda.empty_cache()
-
-    table = pa.table(
-        {
-            "index": np.asarray(all_index, dtype=np.int64),
-            "episode_index": np.asarray(all_episode, dtype=np.int64),
-            "frame_index": np.asarray(all_frame, dtype=np.int64),
-            "progress_sparse": np.asarray(all_progress, dtype=np.float32),
-        }
-    )
-
-    schema_metadata: dict[bytes, bytes] = {b"vlm_name": config.vlm_name.encode()}
-    if reward_model_path is not None:
-        schema_metadata[b"reward_model_path"] = reward_model_path.encode()
-    table = table.replace_schema_metadata(schema_metadata)
-
-    out = Path(dataset.root) / DEFAULT_OUTPUT_FILENAME if output_path is None else Path(output_path)
-    out.parent.mkdir(parents=True, exist_ok=True)
-    pq.write_table(table, out)
-    logging.info(f"Saved {len(table)} frame values to {out}")
-
-    progress_arr = np.asarray(all_progress, dtype=np.float32)
-    if progress_arr.size:
-        logging.info(
-            f"Progress: mean={float(progress_arr.mean()):.4f}, "
-            f"std={float(progress_arr.std()):.4f}, "
-            f"min={float(progress_arr.min()):.4f}, "
-            f"max={float(progress_arr.max()):.4f}"
-        )
-    return out
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Compute per-frame TOPReward progress curves for RA-BC weighting.",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog="""
-Examples:
-    # Sparse-dense mode (matches upstream TOPReward num_samples=15)
-    python -m lerobot.rewards.topreward.compute_rabc_weights \\
-        --dataset-repo-id lerobot/libero_10_image \\
-        --num-samples 15
-
-    # Use a smaller VLM
-    python -m lerobot.rewards.topreward.compute_rabc_weights \\
-        --dataset-repo-id lerobot/libero_10_image \\
-        --vlm-name Qwen/Qwen3-VL-4B-Instruct
-        """,
-    )
-    parser.add_argument(
-        "--dataset-repo-id", type=str, required=True, help="HuggingFace dataset repo id or local path."
-    )
-    parser.add_argument(
-        "--reward-model-path", type=str, default=None, help="Optional TOPReward LeRobot config."
-    )
-    parser.add_argument("--vlm-name", type=str, default=None, help="Override the VLM backbone (HF Hub id).")
-    parser.add_argument("--output-path", type=str, default=None, help="Output parquet path.")
-    parser.add_argument("--device", type=str, default="cuda", help="Device to use (default: cuda).")
-    parser.add_argument(
-        "--num-samples",
-        type=int,
-        default=None,
-        help="Anchor prefix samples per episode. None = dense. 15 matches upstream.",
-    )
-    parser.add_argument(
-        "--episodes",
-        type=int,
-        nargs="+",
-        default=None,
-        help="Process only these episode indices (e.g. --episodes 0 or --episodes 0 5 10).",
-    )
-    parser.add_argument("--fps", type=float, default=None, help="Override TOPRewardConfig.fps.")
-    parser.add_argument(
-        "--push-to-hub", action="store_true", help="Upload to the dataset repo on HuggingFace Hub."
-    )
-
-    args = parser.parse_args()
-
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-
-    output_path = compute_topreward_progress(
-        dataset_repo_id=args.dataset_repo_id,
-        reward_model_path=args.reward_model_path,
-        vlm_name=args.vlm_name,
-        output_path=args.output_path,
-        device=args.device,
-        num_samples=args.num_samples,
-        fps=args.fps,
-        episodes=args.episodes,
-    )
-
-    print(f"\nTOPReward progress saved to: {output_path}")
-
-    if args.push_to_hub:
-        from huggingface_hub import HfApi
-
-        api = HfApi()
-        hub_path = DEFAULT_OUTPUT_FILENAME
-
-        print(f"\nUploading to Hub: {args.dataset_repo_id}/{hub_path}")
-        api.upload_file(
-            path_or_fileobj=str(output_path),
-            path_in_repo=hub_path,
-            repo_id=args.dataset_repo_id,
-            repo_type="dataset",
-        )
-        print(
-            "Successfully uploaded to: "
-            f"https://huggingface.co/datasets/{args.dataset_repo_id}/blob/main/{hub_path}"
-        )
-
-        print("\nTo use in training, add to your config:")
-        print("  use_rabc: true")
-        print(f"  rabc_progress_path: hf://datasets/{args.dataset_repo_id}/{hub_path}")
-        print("  rabc_head_mode: sparse")
-    else:
-        print("\nTo use in training, add to your config:")
-        print("  use_rabc: true")
-        print(f"  rabc_progress_path: {output_path}")
-        print("  rabc_head_mode: sparse")
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/rewards/topreward/configuration_topreward.py

@@ -1,146 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
-from lerobot.configs.rewards import RewardModelConfig
-from lerobot.utils.constants import OBS_IMAGES
-
-# Default prompt scaffolding from the upstream TOPReward paper / reference
-# implementation (``QwenClient.compute_instruction_reward``). The prompt
-# scores the terminal ``True`` token in ``f"{instruction} ... True"``
-# given the video.
-DEFAULT_PROMPT_PREFIX = (
-    "The above video shows a robot manipulation trajectory that completes the following task: "
-)
-DEFAULT_PROMPT_SUFFIX_TEMPLATE = (
-    "{instruction} Decide whether the above statement is True or not. The answer is: True"
-)
-
-
-@RewardModelConfig.register_subclass("topreward")
-@dataclass
-class TOPRewardConfig(RewardModelConfig):
-    """Configuration for the TOPReward zero-shot reward model.
-
-    TOPReward is **zero-shot**: it has no learnable parameters of its own.
-    The "model" is a generic vision-language model (default
-    ``Qwen/Qwen3-VL-8B-Instruct``) used with a fixed prompt to extract
-    token log-probabilities as a reward signal. There is therefore no
-    fine-tuned checkpoint to host: ``pretrained_path`` is unused at
-    runtime — the model identity is :attr:`vlm_name` (an HF Hub id).
-
-    Args:
-        vlm_name: Hugging Face Hub id of the underlying VLM. Must be a
-            Qwen3-VL family model (the only client implemented in this
-            LeRobot port).
-        torch_dtype: Torch dtype name passed to the VLM loader
-            (``"auto"``, ``"bfloat16"``, ``"float16"``, ...).
-        attn_implementation: ``transformers`` attention implementation
-            (e.g. ``"flash_attention_2"``, ``"sdpa"``). Defaults to
-            ``None`` so the upstream picks the best available.
-        image_key: Observation key that holds the trajectory frames.
-        task_key: Complementary-data key that holds the task instruction.
-        default_task: Fallback instruction when ``task_key`` is absent.
-        max_frames: Cap on the number of frames fed to the VLM per
-            sample. ``None`` = use all frames.
-        fps: Frames-per-second metadata for the Qwen video processor.
-        prompt_prefix: Text shown to the VLM right after the video and
-            before the suffix template.
-        prompt_suffix_template: Suffix appended after ``prompt_prefix``.
-            Must contain ``{instruction}``; the VLM scores the
-            log-likelihood of the tokens that follow the prefix.
-        add_chat_template: If ``True``, wrap the full prompt with the
-            tokenizer's chat template before tokenisation (matches
-            upstream ``add_chat_template=True``).
-        success_threshold: Optional log-prob threshold. If finite,
-            :meth:`TOPRewardModel.compute_reward` returns
-            ``(reward > success_threshold).float()`` instead of the raw
-            log-prob.
-        max_input_length: Hard limit on the total tokenized input length;
-            samples that exceed it raise a ``ValueError``.
-    """
-
-    # Path to a local LeRobot dir or HF repo that holds a ``config.json``
-    # snapshot of this TOPRewardConfig. The VLM weights themselves are
-    # always identified by ``vlm_name``.
-    pretrained_path: str | None = None
-
-    vlm_name: str = "Qwen/Qwen3-VL-8B-Instruct"
-    torch_dtype: str = "auto"
-    attn_implementation: str | None = None
-
-    image_key: str = OBS_IMAGES + ".top"
-    task_key: str = "task"
-    default_task: str | None = None
-    max_frames: int | None = 16
-    fps: float = 2.0
-
-    prompt_prefix: str = DEFAULT_PROMPT_PREFIX
-    prompt_suffix_template: str = DEFAULT_PROMPT_SUFFIX_TEMPLATE
-    add_chat_template: bool = False
-
-    success_threshold: float = float("-inf")
-    max_input_length: int = 32768
-
-    license: str | None = "mit"  # matches upstream TOPReward
-    tags: list[str] | None = field(
-        default_factory=lambda: ["reward-model", "vision-language", "qwen3-vl", "zero-shot"]
-    )
-
-    input_features: dict[str, PolicyFeature] = field(default_factory=dict)
-    output_features: dict[str, PolicyFeature] = field(default_factory=dict)
-    normalization_mapping: dict[str, NormalizationMode] = field(
-        default_factory=lambda: {
-            "VISUAL": NormalizationMode.IDENTITY,
-            "REWARD": NormalizationMode.IDENTITY,
-        }
-    )
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        if self.max_frames is not None and self.max_frames < 1:
-            raise ValueError(f"max_frames must be >= 1, got {self.max_frames}")
-        if self.fps <= 0:
-            raise ValueError(f"fps must be > 0, got {self.fps}")
-        if "{instruction}" not in self.prompt_suffix_template:
-            raise ValueError(
-                "prompt_suffix_template must contain `{instruction}` so the model "
-                "scores the log-likelihood of the task suffix."
-            )
-        if self.max_input_length <= 0:
-            raise ValueError(f"max_input_length must be > 0, got {self.max_input_length}")
-
-        if self.image_key not in self.input_features:
-            self.input_features[self.image_key] = PolicyFeature(shape=(3, 224, 224), type=FeatureType.VISUAL)
-        self.output_features.setdefault("reward", PolicyFeature(shape=(1,), type=FeatureType.REWARD))
-
-    @property
-    def observation_delta_indices(self) -> list[int] | None:
-        return None
-
-    @property
-    def action_delta_indices(self) -> None:
-        return None
-
-    @property
-    def reward_delta_indices(self) -> None:
-        return None
-
-    def validate_features(self) -> None:
-        if self.image_key not in self.input_features:
-            raise ValueError(f"TOPReward requires image input feature {self.image_key!r}")

src/lerobot/rewards/topreward/modeling_topreward.py

@@ -1,238 +0,0 @@
-# Copyright 2026 Shirui Chen, Cole Harrison, Ying-Chun Lee, Angela Jin Yang,
-# Zhongzheng Ren, Lillian J. Ratliff, Jiafei Duan, Dieter Fox, Ranjay Krishna
-# and The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""TOPReward: Token Probabilities as Hidden Zero-Shot Rewards for Robotics.
-
-Paper:         https://arxiv.org/abs/2602.19313
-Project:       https://topreward.github.io/webpage/
-Original code: https://github.com/TOPReward/TOPReward
-Backbone:      https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct  (default)
-
-TOPReward is a **zero-shot** reward model: it has no fine-tuned weights of
-its own. Given a video trajectory and a task instruction, it asks an
-off-the-shelf VLM how likely the instruction is, conditioned on the video,
-and returns that log-likelihood as the reward signal.
-
-Inference recipe:
-
-1. The processor builds a chat-style prompt, tokenises it, and emits
-   ``input_ids``, ``attention_mask``, vision tensors, and ``labels``.
-   The processor label-masks everything except the terminal answer token with
-   ``-100``.
-2. Forward the full token sequence through the VLM.
-3. Read the terminal answer token log-probability from the logits as the
-   scalar reward.
-
-With the default ``prompt_suffix_template``, the only unmasked token is the
-literal ``"True"`` at the end — the reward is
-``log P("True" | video + prompt + instruction)``.
-
-This LeRobot port is **inference-only and not trainable** — :meth:`forward`
-is intentionally inherited from :class:`PreTrainedRewardModel` and raises
-``NotImplementedError``, making :attr:`PreTrainedRewardModel.is_trainable`
-return ``False``.
-
-Because the VLM weights live on the Hugging Face Hub under their canonical
-id (``Qwen/Qwen3-VL-8B-Instruct`` etc.) and TOPReward never modifies them,
-:meth:`_save_pretrained` and :meth:`from_pretrained` are overridden so a
-TOPReward LeRobot "checkpoint" is a single ``config.json`` (the VLM is
-re-fetched from the Hub at load time).
-"""
-
-from __future__ import annotations
-
-import builtins
-import logging
-import os
-from pathlib import Path
-from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, TypeVar
-
-import numpy as np
-import torch
-from huggingface_hub import HfApi, hf_hub_download
-from huggingface_hub.constants import CONFIG_NAME
-from huggingface_hub.errors import HfHubHTTPError
-from torch import Tensor
-from torch.nn.functional import cross_entropy
-
-from lerobot.configs.rewards import RewardModelConfig
-from lerobot.rewards.pretrained import PreTrainedRewardModel
-from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig
-from lerobot.rewards.topreward.processor_topreward import TOPREWARD_FEATURE_PREFIX, TOPREWARD_INPUT_KEYS
-from lerobot.utils.import_utils import _transformers_available, require_package
-
-if TYPE_CHECKING:
-    from lerobot.configs.train import TrainPipelineConfig
-
-if TYPE_CHECKING or _transformers_available:
-    from transformers import Qwen3VLForConditionalGeneration
-else:
-    Qwen3VLForConditionalGeneration = None  # type: ignore[assignment]
-
-logger = logging.getLogger(__name__)
-
-T = TypeVar("T", bound="TOPRewardModel")
-
-
-def _torch_dtype(name: str) -> torch.dtype | str:
-    """Resolve a torch dtype name; ``"auto"`` is passed through verbatim."""
-    if name == "auto":
-        return "auto"
-    dtype = getattr(torch, name, None)
-    if isinstance(dtype, torch.dtype):
-        return dtype
-    raise ValueError(f"Unknown torch dtype: {name!r}")
-
-
-class TOPRewardModel(PreTrainedRewardModel):
-    """TOPReward zero-shot reward model."""
-
-    name = "topreward"
-    config_class = TOPRewardConfig
-
-    def __init__(self, config: TOPRewardConfig) -> None:
-        require_package("transformers", extra="topreward")
-        super().__init__(config)
-        self.config = config
-
-        torch_dtype = _torch_dtype(config.torch_dtype)
-        model_kwargs: dict[str, Any] = {"dtype": torch_dtype, "trust_remote_code": True}
-        if config.attn_implementation is not None:
-            model_kwargs["attn_implementation"] = config.attn_implementation
-
-        self.model = Qwen3VLForConditionalGeneration.from_pretrained(config.vlm_name, **model_kwargs)
-
-    def compute_reward(self, batch: dict[str, Any]) -> Tensor:
-        """Return one log-prob reward per sample in the batch."""
-        inputs: dict[str, Any] = {}
-        for key in TOPREWARD_INPUT_KEYS:
-            batch_key = f"{TOPREWARD_FEATURE_PREFIX}{key}"
-            if batch_key not in batch:
-                raise KeyError(
-                    f"TOPReward batch missing `{batch_key}`. Make sure the "
-                    "TOPRewardEncoderProcessorStep ran before `compute_reward`."
-                )
-            inputs[key] = batch[batch_key]
-
-        device = next(self.model.parameters()).device
-        inputs = {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()}
-        labels = inputs.pop("labels")
-        inputs["logits_to_keep"] = 2
-
-        self.eval()
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        logits = outputs.logits
-        rewards = -cross_entropy(logits[:, -2, :].float(), labels[:, -1], reduction="none")
-        if np.isfinite(self.config.success_threshold):
-            rewards = (rewards > self.config.success_threshold).float()
-        return rewards.to(self.config.device or "cpu")
-
-    def _save_pretrained(self, save_directory: Path) -> None:
-        """Save ``config.json`` only."""
-        self.config._save_pretrained(save_directory)
-
-    @classmethod
-    def from_pretrained(
-        cls: builtins.type[T],
-        pretrained_name_or_path: str | Path,
-        *,
-        config: RewardModelConfig | None = None,
-        force_download: bool = False,
-        resume_download: bool | None = None,
-        proxies: dict | None = None,
-        token: str | bool | None = None,
-        cache_dir: str | Path | None = None,
-        local_files_only: bool = False,
-        revision: str | None = None,
-        strict: bool = False,  # noqa: ARG003 — accepted for API parity; unused (no safetensors to load)
-        **kwargs: Any,
-    ) -> T:
-        """Load a TOPReward configuration and instantiate the wrapped VLM."""
-        if config is None:
-            config = RewardModelConfig.from_pretrained(
-                pretrained_name_or_path=pretrained_name_or_path,
-                force_download=force_download,
-                resume_download=resume_download,
-                proxies=proxies,
-                token=token,
-                cache_dir=cache_dir,
-                local_files_only=local_files_only,
-                revision=revision,
-                **kwargs,
-            )
-        if not isinstance(config, TOPRewardConfig):
-            raise TypeError(
-                f"Expected a TOPRewardConfig, got {type(config).__name__}. Make sure "
-                f"`pretrained_name_or_path={pretrained_name_or_path!r}` points at a "
-                "TOPReward checkpoint."
-            )
-
-        model_id = str(pretrained_name_or_path)
-        if not os.path.isdir(model_id):
-            try:
-                hf_hub_download(
-                    repo_id=model_id,
-                    filename=CONFIG_NAME,
-                    revision=revision,
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    proxies=proxies,
-                    resume_download=resume_download,
-                    token=token,
-                    local_files_only=local_files_only,
-                )
-            except HfHubHTTPError as e:
-                raise FileNotFoundError(
-                    f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
-                ) from e
-
-        instance = cls(config, **kwargs)
-        instance.to(config.device)
-        instance.eval()
-        return instance
-
-    def push_model_to_hub(self, cfg: TrainPipelineConfig):
-        """Push the TOPReward ``config.json`` + model card to the Hub."""
-        api = HfApi()
-        repo_id = api.create_repo(
-            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
-        ).repo_id
-
-        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
-            saved_path = Path(tmp) / repo_id
-            saved_path.mkdir(parents=True, exist_ok=True)
-
-            self.config._save_pretrained(saved_path)
-
-            card = self.generate_model_card(
-                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
-            )
-            card.save(str(saved_path / "README.md"))
-
-            cfg.save_pretrained(saved_path)
-
-            commit_info = api.upload_folder(
-                repo_id=repo_id,
-                repo_type="model",
-                folder_path=saved_path,
-                commit_message="Upload TOPReward config and readme",
-                allow_patterns=["*.json", "*.yaml", "*.md"],
-                ignore_patterns=["*.tmp", "*.log", "*.safetensors"],
-            )
-
-            logger.info(f"Model pushed to {commit_info.repo_url.url}")

src/lerobot/rewards/topreward/processor_topreward.py

@@ -1,305 +0,0 @@
-# Copyright 2026 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.
-
-"""TOPReward pre/post processing pipeline."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
-
-import torch
-from torch import Tensor
-
-from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    ProcessorStep,
-    ProcessorStepRegistry,
-    policy_action_to_transition,
-)
-from lerobot.rewards.topreward.configuration_topreward import (
-    DEFAULT_PROMPT_PREFIX,
-    DEFAULT_PROMPT_SUFFIX_TEMPLATE,
-    TOPRewardConfig,
-)
-from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import (
-    OBS_IMAGES,
-    OBS_PREFIX,
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
-)
-from lerobot.utils.import_utils import _transformers_available, require_package
-
-if TYPE_CHECKING or _transformers_available:
-    from transformers import AutoProcessor
-else:
-    AutoProcessor = None
-
-TOPREWARD_FEATURE_PREFIX = f"{OBS_PREFIX}topreward."
-
-_TRUE_ANSWER = "True"
-
-TOPREWARD_VLM_INPUT_KEYS = (
-    "input_ids",
-    "attention_mask",
-    "pixel_values_videos",
-    "video_grid_thw",
-    "mm_token_type_ids",
-)
-TOPREWARD_INPUT_KEYS = TOPREWARD_VLM_INPUT_KEYS + ("labels",)
-
-
-def _prepare_video_batch(video: Tensor, *, max_frames: int | None) -> Tensor:
-    """Return videos as ``(B, T, C, H, W)`` uint8 tensors for Qwen3-VL."""
-    if video.ndim == 4:
-        video = video.unsqueeze(1)
-    elif video.ndim != 5:
-        raise ValueError(
-            f"Expected TOPReward frames with shape (B,C,H,W) or (B,T,C,H,W); got {tuple(video.shape)}"
-        )
-
-    if max_frames is not None:
-        video = video[:, -max_frames:]
-    if video.shape[-1] in (1, 3):
-        video = video.permute(0, 1, 4, 2, 3)
-    elif video.shape[2] not in (1, 3):
-        raise ValueError(f"Expected channel dim of size 1 or 3, got shape {tuple(video.shape)}")
-
-    if video.is_floating_point():
-        video = video * 255.0
-
-    return video.clamp(0, 255).to(torch.uint8).contiguous()
-
-
-def _expand_tasks(task: Any, *, batch_size: int, default: str | None) -> list[str]:
-    if task is None:
-        task = default
-    if task is None:
-        raise KeyError("TOPReward expected a task description in complementary data")
-    if isinstance(task, str):
-        return [task] * batch_size
-    if isinstance(task, tuple):
-        task = list(task)
-    if not (isinstance(task, list) and all(isinstance(item, str) for item in task)):
-        raise TypeError(f"TOPReward task must be a string or list of strings, got {type(task)}")
-    if len(task) == 1 and batch_size > 1:
-        return task * batch_size
-    if len(task) != batch_size:
-        raise ValueError(f"Expected {batch_size} tasks, got {len(task)}")
-    return task
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="topreward_encoder")
-class TOPRewardEncoderProcessorStep(ProcessorStep):
-    """Encode raw frames + task into Qwen-VL tensors for the TOPReward model.
-
-    Loads a :class:`~transformers.AutoProcessor` matching ``vlm_name`` and
-    builds the full chat prompt including the instruction suffix. The
-    resulting ``input_ids``, ``attention_mask``, vision tensors, and
-    ``labels`` are written under the ``observation.topreward.*`` namespace
-    so the model can score without re-tokenising.
-
-    At call time the step reads:
-
-    - ``observation[image_key]``: ``(B, T, C, H, W)`` or ``(B, C, H, W)`` frames.
-    - ``complementary_data[task_key]``: a string or list of strings.
-
-    and writes ``observation[f"{TOPREWARD_FEATURE_PREFIX}<name>"]`` for the
-    Qwen-VL tensors plus ``labels``.
-    """
-
-    vlm_name: str = "Qwen/Qwen3-VL-8B-Instruct"
-    image_key: str = OBS_IMAGES + ".top"
-    task_key: str = "task"
-    default_task: str | None = None
-    max_frames: int | None = 16
-    fps: float = 2.0
-    prompt_prefix: str = DEFAULT_PROMPT_PREFIX
-    prompt_suffix_template: str = DEFAULT_PROMPT_SUFFIX_TEMPLATE
-    add_chat_template: bool = False
-    max_length: int = 32768
-
-    _processor: Any = field(default=None, init=False, repr=False)
-
-    def __post_init__(self) -> None:
-        require_package("transformers", extra="topreward")
-        self._processor = AutoProcessor.from_pretrained(self.vlm_name, trust_remote_code=True)
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        observation = transition.get(TransitionKey.OBSERVATION)
-        complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
-        if self.image_key not in observation:
-            raise KeyError(f"TOPReward expected image key {self.image_key!r} in observation")
-
-        frames = observation[self.image_key]
-        videos = frames.detach().cpu() if isinstance(frames, Tensor) else torch.as_tensor(frames)
-        videos = _prepare_video_batch(videos, max_frames=self.max_frames)
-
-        batch_size = videos.shape[0]
-        tasks = _expand_tasks(
-            complementary.get(self.task_key, self.default_task),
-            batch_size=batch_size,
-            default=self.default_task,
-        )
-
-        encoded = self._encode_batch(videos, tasks, batch_size)
-
-        new_observation = dict(observation)
-        for key, value in encoded.items():
-            new_observation[f"{TOPREWARD_FEATURE_PREFIX}{key}"] = value
-
-        new_transition = transition.copy()
-        new_transition[TransitionKey.OBSERVATION] = new_observation
-        return new_transition
-
-    def _encode_batch(self, videos: Tensor, tasks: list[str], batch_size) -> dict[str, Any]:
-        """Tokenise a batch of (frames, task) pairs into Qwen-VL tensors.
-
-        The loop only builds per-sample chat strings. Tokenisation, padding,
-        video preprocessing, and label construction are batched.
-        """
-
-        texts: list[str] = []
-        video_metadata = [
-            {
-                "total_num_frames": int(videos.shape[1]),
-                "fps": float(self.fps),
-                "frames_indices": list(range(int(videos.shape[1]))),
-            }
-            for _ in range(batch_size)
-        ]
-        eos_token = self._processor.tokenizer.eos_token
-
-        for i in range(batch_size):
-            instruction_suffix = self.prompt_suffix_template.format(instruction=tasks[i])
-            if self.add_chat_template:
-                suffix_for_template = instruction_suffix.removesuffix(_TRUE_ANSWER).rstrip()
-                templated_messages = [
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "video", "video": videos[i], "fps": self.fps},
-                            {"type": "text", "text": f"{self.prompt_prefix}{suffix_for_template}"},
-                        ],
-                    }
-                ]
-                prompt_chat = self._processor.apply_chat_template(
-                    templated_messages, tokenize=False, add_generation_prompt=True
-                )
-                full_text = f"{prompt_chat}{_TRUE_ANSWER}"
-            else:
-                user_messages = [
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "video", "video": videos[i], "fps": self.fps},
-                            {"type": "text", "text": self.prompt_prefix},
-                        ],
-                    }
-                ]
-                prompt_chat = self._processor.apply_chat_template(
-                    user_messages, tokenize=False, add_generation_prompt=False
-                )
-                if eos_token is not None:
-                    prompt_chat = prompt_chat.split(eos_token)[0]
-                full_text = f"{prompt_chat}{instruction_suffix}"
-
-            texts.append(full_text)
-
-        result = self._processor(
-            text=texts,
-            videos=videos,
-            video_metadata=video_metadata,
-            do_sample_frames=False,
-            padding=True,
-            padding_side="left",
-            return_tensors="pt",
-        )
-        input_ids = result["input_ids"]
-
-        if input_ids.shape[-1] > self.max_length:
-            raise ValueError(
-                f"TOPReward input length {input_ids.shape[-1]} exceeds max_length "
-                f"{self.max_length}; lower `max_frames` or raise `max_length`."
-            )
-
-        labels = torch.full_like(input_ids, -100)
-        labels[:, -1] = input_ids[:, -1]
-        result["labels"] = labels
-        return result
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-    def get_config(self) -> dict[str, Any]:
-        return {
-            "vlm_name": self.vlm_name,
-            "image_key": self.image_key,
-            "task_key": self.task_key,
-            "default_task": self.default_task,
-            "max_frames": self.max_frames,
-            "fps": self.fps,
-            "prompt_prefix": self.prompt_prefix,
-            "prompt_suffix_template": self.prompt_suffix_template,
-            "add_chat_template": self.add_chat_template,
-            "max_length": self.max_length,
-        }
-
-
-def make_topreward_pre_post_processors(
-    config: TOPRewardConfig,
-    dataset_stats: dict[str, dict[str, Any]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """Pipeline that pre-encodes frames + task into Qwen-VL tensors.
-
-    The preprocessor adds a batch dimension if needed, runs TOPReward's
-    encoder (which tokenises the full prompt and emits ``labels``), and
-    moves everything to the configured device. The postprocessor is
-    the identity since TOPReward outputs a single reward tensor.
-    """
-    preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-        steps=[
-            AddBatchDimensionProcessorStep(),
-            TOPRewardEncoderProcessorStep(
-                vlm_name=config.vlm_name,
-                image_key=config.image_key,
-                task_key=config.task_key,
-                default_task=config.default_task,
-                max_frames=config.max_frames,
-                fps=config.fps,
-                prompt_prefix=config.prompt_prefix,
-                prompt_suffix_template=config.prompt_suffix_template,
-                add_chat_template=config.add_chat_template,
-                max_length=config.max_input_length,
-            ),
-            DeviceProcessorStep(device=config.device or "cpu"),
-        ],
-        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-    )
-    postprocessor = PolicyProcessorPipeline(
-        name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-        to_transition=policy_action_to_transition,
-    )
-    return preprocessor, postprocessor

src/lerobot/templates/lerobot_rewardmodel_modelcard_template.md

@@ -13,8 +13,6 @@
 A reward classifier is a lightweight neural network that scores observations or trajectories for task success, providing a learned reward signal or offline evaluation when explicit rewards are unavailable.
 {% elif model_name == "sarm" %}
 A Success-Aware Reward Model (SARM) predicts a dense reward signal from observations, typically used downstream for reinforcement learning or human-in-the-loop fine-tuning when task success is not directly observable.
-{% elif model_name == "topreward" %}
-TOPReward is a **zero-shot** reward model that extracts token log-probabilities from an off-the-shelf vision-language model (default Qwen3-VL) as a reward signal. Given a video trajectory and a task instruction, it returns the VLM's log-likelihood of the instruction being true, with no fine-tuning required.
 {% else %}
 _Reward model type not recognized — please update this template._
 {% endif %}

tests/datasets/test_datasets.py

@@ -24,6 +24,7 @@ import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
+import datasets
 from huggingface_hub import HfApi
 from PIL import Image
 from safetensors.torch import load_file
@@ -360,6 +361,41 @@ def test_add_frame_image_pil(image_dataset):
     assert dataset[0]["image"].shape == torch.Size(DUMMY_CHW)
 
 
+@pytest.mark.parametrize(
+    "dtype,np_dtype,values,assert_fn",
+    [
+        ("float32", np.float32, [1.0, 2.0], np.testing.assert_allclose),
+        ("int64", np.int64, [1, 2], np.testing.assert_array_equal),
+        ("bool", np.bool_, [True, False], np.testing.assert_array_equal),
+    ],
+    ids=["float32", "int64", "bool"],
+)
+def test_save_episode_shape_1_scalar_is_scalarized_before_hf_encoding(
+    tmp_path, empty_lerobot_dataset_factory, monkeypatch, dtype, np_dtype, values, assert_fn
+):
+    features = {"state": {"dtype": dtype, "shape": (1,), "names": None}}
+    dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
+    dataset.add_frame({"state": np.array([values[0]], dtype=np_dtype), "task": "Dummy task"})
+    dataset.add_frame({"state": np.array([values[1]], dtype=np_dtype), "task": "Dummy task"})
+
+    captured = {}
+    original_from_dict = datasets.Dataset.from_dict
+
+    def _from_dict_spy(cls, mapping, *args, **kwargs):
+        captured["state"] = mapping["state"]
+        return original_from_dict(mapping, *args, **kwargs)
+
+    monkeypatch.setattr(datasets.Dataset, "from_dict", classmethod(_from_dict_spy))
+
+    dataset.save_episode()
+    dataset.finalize()
+
+    assert "state" in captured
+    assert isinstance(captured["state"], np.ndarray)
+    assert captured["state"].shape == (2,)
+    assert_fn(captured["state"], np.array(values, dtype=np_dtype))
+
+
 def test_set_image_transforms_applies_transparently(image_dataset):
     dataset = image_dataset
     dataset.add_frame({"image": np.random.rand(*DUMMY_CHW), "task": "Dummy task"})

tests/datasets/test_video_decoder_cache.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python
+
+# Copyright 2026 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.
+
+"""Unit tests for ``lerobot.datasets.video_utils.VideoDecoderCache``.
+
+These cover the LRU bounding + file-handle release behaviour added to prevent
+unbounded growth when iterating over datasets with many distinct video files
+(observed: ~35 GB anon-rss per DataLoader worker on an 8 k-file dataset).
+"""
+
+import shutil
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("torchcodec", reason="torchcodec is required (install lerobot[dataset])")
+
+from lerobot.datasets.video_utils import VideoDecoderCache  # noqa: E402
+
+TEST_ARTIFACTS_DIR = Path(__file__).resolve().parent.parent / "artifacts" / "encoded_videos"
+SRC_CLIP = TEST_ARTIFACTS_DIR / "clip_4frames.mp4"
+
+
+def _make_distinct_clips(tmp_path: Path, n: int) -> list[Path]:
+    """Copy the small reference mp4 to ``n`` distinct paths.
+
+    The cache keys on absolute path, so distinct paths force distinct cache entries
+    even though the file contents are identical.
+    """
+    assert SRC_CLIP.exists(), f"missing test artifact {SRC_CLIP}"
+    paths = []
+    for i in range(n):
+        dst = tmp_path / f"clip_{i:04d}.mp4"
+        shutil.copyfile(SRC_CLIP, dst)
+        paths.append(dst)
+    return paths
+
+
+class TestVideoDecoderCacheBounded:
+    def test_default_cache_is_bounded(self):
+        """The default cache must have a finite ``max_size`` to bound RSS growth."""
+        cache = VideoDecoderCache()
+        assert cache.max_size is not None, "default cache must be bounded"
+        assert cache.max_size > 0
+
+    def test_size_capped_at_max_size(self, tmp_path):
+        """``get_decoder`` for >``max_size`` distinct paths must NOT grow without bound."""
+        paths = _make_distinct_clips(tmp_path, n=5)
+        cache = VideoDecoderCache(max_size=2)
+        for p in paths:
+            cache.get_decoder(p)
+        assert cache.size() == 2
+
+    def test_evicts_least_recently_used(self, tmp_path):
+        """Re-accessing an entry must promote it; the LRU entry is the one evicted."""
+        paths = _make_distinct_clips(tmp_path, n=3)
+        cache = VideoDecoderCache(max_size=2)
+
+        cache.get_decoder(paths[0])
+        cache.get_decoder(paths[1])
+        cache.get_decoder(paths[0])  # promote paths[0] to MRU; paths[1] is now LRU
+        cache.get_decoder(paths[2])  # should evict paths[1]
+
+        assert str(paths[0]) in cache  # MRU stays
+        assert str(paths[1]) not in cache  # LRU evicted
+        assert str(paths[2]) in cache  # newest stays
+
+    def test_eviction_closes_file_handle(self, tmp_path):
+        """Evicting an entry must close its fsspec file handle (otherwise we leak FDs)."""
+        paths = _make_distinct_clips(tmp_path, n=2)
+        cache = VideoDecoderCache(max_size=1)
+
+        cache.get_decoder(paths[0])
+        # Reach into the cache to capture the handle before it is evicted. This is
+        # the only assertion in the suite that touches a private attribute, and it
+        # is the most direct way to prove the file descriptor is actually released.
+        evicted_handle = cache._cache[str(paths[0])][1]
+        assert evicted_handle.closed is False
+
+        cache.get_decoder(paths[1])  # forces eviction of paths[0]
+
+        assert evicted_handle.closed is True
+
+    def test_clear_closes_all_file_handles(self, tmp_path):
+        """``clear()`` must close every cached file handle."""
+        paths = _make_distinct_clips(tmp_path, n=3)
+        cache = VideoDecoderCache(max_size=10)
+
+        for p in paths:
+            cache.get_decoder(p)
+        handles = [entry[1] for entry in cache._cache.values()]
+        assert all(not h.closed for h in handles)
+
+        cache.clear()
+
+        assert cache.size() == 0
+        assert all(h.closed for h in handles)
+
+    def test_hit_does_not_reopen_or_evict(self, tmp_path):
+        """A cache hit must return the same decoder instance without touching the cap."""
+        paths = _make_distinct_clips(tmp_path, n=1)
+        cache = VideoDecoderCache(max_size=2)
+
+        first = cache.get_decoder(paths[0])
+        second = cache.get_decoder(paths[0])
+
+        assert first is second
+        assert cache.size() == 1
+
+    def test_unbounded_when_max_size_none(self, tmp_path):
+        """``max_size=None`` preserves the legacy unbounded behaviour."""
+        paths = _make_distinct_clips(tmp_path, n=4)
+        cache = VideoDecoderCache(max_size=None)
+        for p in paths:
+            cache.get_decoder(p)
+        assert cache.size() == 4
+
+    def test_env_var_overrides_default(self, tmp_path, monkeypatch):
+        """``LEROBOT_VIDEO_DECODER_CACHE_SIZE`` env var sets the default ``max_size``."""
+        monkeypatch.setenv("LEROBOT_VIDEO_DECODER_CACHE_SIZE", "3")
+        cache = VideoDecoderCache()
+        assert cache.max_size == 3
+
+        paths = _make_distinct_clips(tmp_path, n=5)
+        for p in paths:
+            cache.get_decoder(p)
+        assert cache.size() == 3

tests/policies/pi0_pi05/openpi_pytorch/__init__.py

@@ -0,0 +1 @@
+"""Lightweight vendored OpenPI PyTorch modules for PI0/PI05 parity tests."""

tests/policies/pi0_pi05/openpi_pytorch/gemma.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    width: int
+    depth: int
+    mlp_dim: int
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+
+
+def get_config(variant: str) -> Config:
+    """Return the Gemma shape config needed by the OpenPI PyTorch model."""
+    if variant == "dummy":
+        return Config(width=64, depth=4, mlp_dim=128, num_heads=8, num_kv_heads=1, head_dim=16)
+    if variant == "gemma_300m":
+        return Config(width=1024, depth=18, mlp_dim=4096, num_heads=8, num_kv_heads=1, head_dim=256)
+    if variant == "gemma_2b":
+        return Config(width=2048, depth=18, mlp_dim=16_384, num_heads=8, num_kv_heads=1, head_dim=256)
+    raise ValueError(f"Unknown variant: {variant}")

tests/policies/pi0_pi05/openpi_pytorch/gemma_pytorch.py

@@ -0,0 +1,300 @@
+from typing import Literal
+
+import torch
+from torch import nn
+from transformers.models.auto import CONFIG_MAPPING
+from transformers.models.gemma import modeling_gemma
+
+from lerobot.policies.pi_gemma import (
+    PaliGemmaForConditionalGenerationWithPiGemma,
+    PiGemmaForCausalLM,
+    _gated_residual,
+    layernorm_forward,
+)
+
+
+class PaliGemmaWithExpertModel(nn.Module):
+    def __init__(
+        self,
+        vlm_config,
+        action_expert_config,
+        use_adarms=None,
+        precision: Literal["bfloat16", "float32"] = "bfloat16",
+    ):
+        if use_adarms is None:
+            use_adarms = [False, False]
+        super().__init__()
+
+        vlm_config_hf = CONFIG_MAPPING["paligemma"]()
+        vlm_config_hf._vocab_size = 257152  # noqa: SLF001
+        vlm_config_hf.image_token_index = 257152
+        vlm_config_hf.text_config.hidden_size = vlm_config.width
+        vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
+        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
+        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
+        vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
+        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
+        vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
+        vlm_config_hf.text_config.dtype = "float32"
+        vlm_config_hf.text_config.vocab_size = 257152
+        vlm_config_hf.text_config.use_adarms = use_adarms[0]
+        vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None
+        vlm_config_hf.vision_config.intermediate_size = 4304
+        vlm_config_hf.vision_config.projection_dim = 2048
+        vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
+        vlm_config_hf.vision_config.dtype = "float32"
+
+        action_expert_config_hf = CONFIG_MAPPING["gemma"](
+            head_dim=action_expert_config.head_dim,
+            hidden_size=action_expert_config.width,
+            intermediate_size=action_expert_config.mlp_dim,
+            num_attention_heads=action_expert_config.num_heads,
+            num_hidden_layers=action_expert_config.depth,
+            num_key_value_heads=action_expert_config.num_kv_heads,
+            vocab_size=257152,
+            hidden_activation="gelu_pytorch_tanh",
+            dtype="float32",
+            use_adarms=use_adarms[1],
+            adarms_cond_dim=action_expert_config.width if use_adarms[1] else None,
+        )
+
+        self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
+        self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
+        self.gemma_expert.model.embed_tokens = None
+
+        self.to_bfloat16_for_selected_params(precision)
+
+    def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"):
+        if precision == "bfloat16":
+            self.to(dtype=torch.bfloat16)
+        elif precision == "float32":
+            self.to(dtype=torch.float32)
+            return
+        else:
+            raise ValueError(f"Invalid precision: {precision}")
+
+        params_to_keep_float32 = [
+            "vision_tower",
+            "multi_modal_projector",
+            "input_layernorm",
+            "post_attention_layernorm",
+            "model.norm",
+        ]
+
+        for name, param in self.named_parameters():
+            if any(selector in name for selector in params_to_keep_float32):
+                param.data = param.data.to(dtype=torch.float32)
+
+    def embed_image(self, image: torch.Tensor):
+        # Transformers 5.4 no longer divides PaliGemma image features by sqrt(hidden_size),
+        # so the upstream helper now matches OpenPI's patched PaliGemma image-scale semantics.
+        # See https://github.com/huggingface/transformers/pull/44432/changes#diff-c916907e7e52ac85ee1a1527560eae4656cd6c76141ceb1fe3da61bd5f697d2a
+        out_dtype = image.dtype
+        if image.dtype != torch.float32:
+            image = image.to(torch.float32)
+        image_outputs = self.paligemma.model.get_image_features(image)
+        features = image_outputs.pooler_output
+        if features.dtype != out_dtype:
+            features = features.to(out_dtype)
+        return features
+
+    def embed_language_tokens(self, tokens: torch.Tensor):
+        return self.paligemma.model.language_model.get_input_embeddings()(tokens)
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | None = None,
+        inputs_embeds: list[torch.FloatTensor] | None = None,
+        use_cache: bool | None = None,
+        adarms_cond: list[torch.Tensor] | None = None,
+    ):
+        if adarms_cond is None:
+            adarms_cond = [None, None]
+        if inputs_embeds[1] is None:
+            prefix_output = self.paligemma.model.language_model.forward(
+                inputs_embeds=inputs_embeds[0],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                adarms_cond=adarms_cond[0] if adarms_cond is not None else None,
+            )
+            prefix_past_key_values = prefix_output.past_key_values
+            prefix_output = prefix_output.last_hidden_state
+            suffix_output = None
+        elif inputs_embeds[0] is None:
+            suffix_output = self.gemma_expert.model.forward(
+                inputs_embeds=inputs_embeds[1],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                adarms_cond=adarms_cond[1] if adarms_cond is not None else None,
+            )
+            suffix_output = suffix_output.last_hidden_state
+            prefix_output = None
+            prefix_past_key_values = None
+        else:
+            models = [self.paligemma.model.language_model, self.gemma_expert.model]
+            num_layers = self.paligemma.config.text_config.num_hidden_layers
+
+            # Check if gradient checkpointing is enabled for any of the models
+            use_gradient_checkpointing = (
+                hasattr(self.gemma_expert.model, "gradient_checkpointing")
+                and self.gemma_expert.model.gradient_checkpointing
+                and self.training
+            ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
+
+            # Force enable gradient checkpointing if we're in training mode and the model supports it
+            if self.training and hasattr(self.gemma_expert.model, "gradient_checkpointing"):
+                if not self.gemma_expert.model.gradient_checkpointing:
+                    print("Forcing gradient checkpointing to be enabled for Gemma expert model")
+                    self.gemma_expert.model.gradient_checkpointing = True
+                use_gradient_checkpointing = True
+
+            # Debug gradient checkpointing status
+            if hasattr(self, "_debug_gc_printed") and not self._debug_gc_printed:
+                print(f"Gemma expert model gradient checkpointing: {use_gradient_checkpointing}")
+                print(f"Model training mode: {self.training}")
+                print(
+                    f"Gemma expert model has gradient_checkpointing attr: {hasattr(self.gemma_expert.model, 'gradient_checkpointing')}"
+                )
+                if hasattr(self.gemma_expert.model, "gradient_checkpointing"):
+                    print(
+                        f"Gemma expert model gradient_checkpointing value: {self.gemma_expert.model.gradient_checkpointing}"
+                    )
+                self._debug_gc_printed = True
+
+            # Define the complete layer computation function for gradient checkpointing
+            def compute_layer_complete(layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond):
+                models = [self.paligemma.model.language_model, self.gemma_expert.model]
+
+                query_states = []
+                key_states = []
+                value_states = []
+                gates = []
+                for i, hidden_states in enumerate(inputs_embeds):
+                    layer = models[i].layers[layer_idx]
+                    hidden_states, gate = layernorm_forward(
+                        layer.input_layernorm, hidden_states, adarms_cond[i]
+                    )
+                    gates.append(gate)
+
+                    input_shape = hidden_states.shape[:-1]
+                    hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+                    query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+                    key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+                    value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+                    query_states.append(query_state)
+                    key_states.append(key_state)
+                    value_states.append(value_state)
+
+                # Concatenate and process attention
+                query_states = torch.cat(query_states, dim=2)
+                key_states = torch.cat(key_states, dim=2)
+                value_states = torch.cat(value_states, dim=2)
+
+                dummy_tensor = torch.zeros(
+                    query_states.shape[0],
+                    query_states.shape[2],
+                    query_states.shape[-1],
+                    device=query_states.device,
+                    dtype=query_states.dtype,
+                )
+                cos, sin = self.paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
+                query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
+                    query_states, key_states, cos, sin, unsqueeze_dim=1
+                )
+
+                batch_size = query_states.shape[0]
+                scaling = self.paligemma.model.language_model.layers[layer_idx].self_attn.scaling
+
+                # Attention computation
+                att_output, _ = modeling_gemma.eager_attention_forward(
+                    self.paligemma.model.language_model.layers[layer_idx].self_attn,
+                    query_states,
+                    key_states,
+                    value_states,
+                    attention_mask,
+                    scaling,
+                )
+                # Get head_dim from the current layer, not from the model
+                head_dim = self.paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
+                att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
+
+                # Process layer outputs
+                outputs_embeds = []
+                start_pos = 0
+                for i, hidden_states in enumerate(inputs_embeds):
+                    layer = models[i].layers[layer_idx]
+                    end_pos = start_pos + hidden_states.shape[1]
+
+                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
+                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
+                    out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])
+
+                    # first residual
+                    out_emb = _gated_residual(hidden_states, out_emb, gates[i])
+                    after_first_residual = out_emb.clone()
+                    out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])
+                    # Convert to bfloat16 if the next layer (mlp) uses bfloat16
+                    if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
+                        out_emb = out_emb.to(dtype=torch.bfloat16)
+
+                    out_emb = layer.mlp(out_emb)
+                    # second residual
+                    out_emb = _gated_residual(after_first_residual, out_emb, gate)
+                    outputs_embeds.append(out_emb)
+                    start_pos = end_pos
+
+                return outputs_embeds
+
+            # Process all layers with gradient checkpointing if enabled
+            for layer_idx in range(num_layers):
+                if use_gradient_checkpointing:
+                    inputs_embeds = torch.utils.checkpoint.checkpoint(
+                        compute_layer_complete,
+                        layer_idx,
+                        inputs_embeds,
+                        attention_mask,
+                        position_ids,
+                        adarms_cond,
+                        use_reentrant=False,
+                        preserve_rng_state=False,
+                    )
+                else:
+                    inputs_embeds = compute_layer_complete(
+                        layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond
+                    )
+
+                # Old code removed - now using compute_layer_complete function above
+
+            # final norm
+            # Define final norm computation function for gradient checkpointing
+            def compute_final_norms(inputs_embeds, adarms_cond):
+                outputs_embeds = []
+                for i, hidden_states in enumerate(inputs_embeds):
+                    out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
+                    outputs_embeds.append(out_emb)
+                return outputs_embeds
+
+            # Apply gradient checkpointing to final norm if enabled
+            if use_gradient_checkpointing:
+                outputs_embeds = torch.utils.checkpoint.checkpoint(
+                    compute_final_norms,
+                    inputs_embeds,
+                    adarms_cond,
+                    use_reentrant=False,
+                    preserve_rng_state=False,
+                )
+            else:
+                outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond)
+
+            prefix_output = outputs_embeds[0]
+            suffix_output = outputs_embeds[1]
+            prefix_past_key_values = None
+
+        return [prefix_output, suffix_output], prefix_past_key_values

tests/policies/pi0_pi05/openpi_pytorch/image_tools.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn.functional as F  # noqa: N812
+
+
+def resize_with_pad_torch(
+    images: torch.Tensor,
+    height: int,
+    width: int,
+    mode: str = "bilinear",
+) -> torch.Tensor:
+    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
+    by padding with black. If the image is float32, it must be in the range [-1, 1].
+
+    Args:
+        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
+        height: Target height
+        width: Target width
+        mode: Interpolation mode ('bilinear', 'nearest', etc.)
+
+    Returns:
+        Resized and padded tensor with same shape format as input
+    """
+    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
+    if images.shape[-1] <= 4:  # Assume channels-last format
+        channels_last = True
+        # Convert to channels-first for torch operations
+        if images.dim() == 3:
+            images = images.unsqueeze(0)  # Add batch dimension
+        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
+    else:
+        channels_last = False
+        if images.dim() == 3:
+            images = images.unsqueeze(0)  # Add batch dimension
+
+    batch_size, channels, cur_height, cur_width = images.shape
+
+    # Calculate resize ratio
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+
+    # Resize
+    resized_images = F.interpolate(
+        images,
+        size=(resized_height, resized_width),
+        mode=mode,
+        align_corners=False if mode == "bilinear" else None,
+    )
+
+    # Handle dtype-specific clipping
+    if images.dtype == torch.uint8:
+        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
+    elif images.dtype == torch.float32:
+        resized_images = resized_images.clamp(-1.0, 1.0)
+    else:
+        raise ValueError(f"Unsupported image dtype: {images.dtype}")
+
+    # Calculate padding
+    pad_h0, remainder_h = divmod(height - resized_height, 2)
+    pad_h1 = pad_h0 + remainder_h
+    pad_w0, remainder_w = divmod(width - resized_width, 2)
+    pad_w1 = pad_w0 + remainder_w
+
+    # Pad
+    constant_value = 0 if images.dtype == torch.uint8 else -1.0
+    padded_images = F.pad(
+        resized_images,
+        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
+        mode="constant",
+        value=constant_value,
+    )
+
+    # Convert back to original format if needed
+    if channels_last:
+        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
+        if batch_size == 1 and images.shape[0] == 1:
+            padded_images = padded_images.squeeze(0)  # Remove batch dimension if it was added
+
+    return padded_images

tests/policies/pi0_pi05/openpi_pytorch/pi0_pytorch.py

@@ -0,0 +1,471 @@
+import copy
+import logging
+import math
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import Tensor, nn
+
+import tests.policies.pi0_pi05.openpi_pytorch.gemma as _gemma
+from tests.policies.pi0_pi05.openpi_pytorch import preprocessing_pytorch as _preprocessing
+from tests.policies.pi0_pi05.openpi_pytorch.gemma_pytorch import PaliGemmaWithExpertModel
+
+
+def get_safe_dtype(target_dtype, device_type):
+    """Get a safe dtype for the given device type."""
+    if device_type == "cpu":
+        # CPU doesn't support bfloat16, use float32 instead
+        if target_dtype == torch.bfloat16:
+            return torch.float32
+        if target_dtype == torch.float64:
+            return torch.float64
+    return target_dtype
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
+) -> Tensor:
+    """Computes sine-cosine positional embedding vectors for scalar positions."""
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+
+    if time.ndim != 1:
+        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
+
+    dtype = get_safe_dtype(torch.float64, device.type)
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+
+    # Compute the outer product
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None]
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+
+
+def sample_beta(alpha, beta, bsize, device):
+    alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device)
+    beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device)
+    dist = torch.distributions.Beta(alpha_t, beta_t)
+    return dist.sample((bsize,))
+
+
+def make_att_2d_masks(pad_masks, att_masks):
+    """Copied from big_vision.
+
+    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
+    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
+    setup several types of attention, for example:
+
+      [[1 1 1 1 1 1]]: pure causal attention.
+
+      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
+          themselves and the last 3 tokens have a causal attention. The first
+          entry could also be a 1 without changing behaviour.
+
+      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
+          block can attend all previous blocks and all tokens on the same block.
+
+    Args:
+      input_mask: bool[B, N] true if its part of the input, false if padding.
+      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
+        it and 0 where it shares the same attention mask as the previous token.
+    """
+    if att_masks.ndim != 2:
+        raise ValueError(att_masks.ndim)
+    if pad_masks.ndim != 2:
+        raise ValueError(pad_masks.ndim)
+
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    return att_2d_masks & pad_2d_masks
+
+
+class PI0Pytorch(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.pi05 = config.pi05
+
+        paligemma_config = _gemma.get_config(config.paligemma_variant)
+        action_expert_config = _gemma.get_config(config.action_expert_variant)
+
+        self.paligemma_with_expert = PaliGemmaWithExpertModel(
+            paligemma_config,
+            action_expert_config,
+            use_adarms=[False, True] if self.pi05 else [False, False],
+            precision=config.dtype,
+        )
+
+        self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width)
+        self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim)
+
+        if self.pi05:
+            self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width)
+            self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width)
+        else:
+            self.state_proj = nn.Linear(config.action_dim, action_expert_config.width)
+            self.action_time_mlp_in = nn.Linear(2 * action_expert_config.width, action_expert_config.width)
+            self.action_time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width)
+
+        torch.set_float32_matmul_precision("high")
+        if config.pytorch_compile_mode is not None:
+            self.sample_actions = torch.compile(self.sample_actions, mode=config.pytorch_compile_mode)
+
+        # Initialize gradient checkpointing flag
+        self.gradient_checkpointing_enabled = False
+
+        # The upstream OpenPI module verifies a site-package Transformers patch here.
+        # This vendored test copy instead routes through LeRobot's local PiGemma compatibility layer.
+
+    def gradient_checkpointing_enable(self):
+        """Enable gradient checkpointing for memory optimization."""
+        self.gradient_checkpointing_enabled = True
+        self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
+        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
+        self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
+
+        logging.info("Enabled gradient checkpointing for PI0Pytorch model")
+
+    def gradient_checkpointing_disable(self):
+        """Disable gradient checkpointing."""
+        self.gradient_checkpointing_enabled = False
+        self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
+        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = False
+        self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
+
+        logging.info("Disabled gradient checkpointing for PI0Pytorch model")
+
+    def is_gradient_checkpointing_enabled(self):
+        """Check if gradient checkpointing is enabled."""
+        return self.gradient_checkpointing_enabled
+
+    def _apply_checkpoint(self, func, *args, **kwargs):
+        """Helper method to apply gradient checkpointing if enabled."""
+        if self.gradient_checkpointing_enabled and self.training:
+            return torch.utils.checkpoint.checkpoint(
+                func, *args, use_reentrant=False, preserve_rng_state=False, **kwargs
+            )
+        return func(*args, **kwargs)
+
+    def _prepare_attention_masks_4d(self, att_2d_masks):
+        """Helper method to prepare 4D attention masks for transformer."""
+        att_2d_masks_4d = att_2d_masks[:, None, :, :]
+        return torch.where(att_2d_masks_4d, 0.0, -2.3819763e38)
+
+    def _preprocess_observation(self, observation, *, train=True):
+        """Helper method to preprocess observation."""
+        observation = _preprocessing.preprocess_observation_pytorch(observation, train=train)
+        return (
+            list(observation.images.values()),
+            list(observation.image_masks.values()),
+            observation.tokenized_prompt,
+            observation.tokenized_prompt_mask,
+            observation.state,
+        )
+
+    def sample_noise(self, shape, device):
+        return torch.normal(
+            mean=0.0,
+            std=1.0,
+            size=shape,
+            dtype=torch.float32,
+            device=device,
+        )
+
+    def sample_time(self, bsize, device):
+        time_beta = sample_beta(1.5, 1.0, bsize, device)
+        time = time_beta * 0.999 + 0.001
+        return time.to(dtype=torch.float32, device=device)
+
+    def embed_prefix(
+        self, images, img_masks, lang_tokens, lang_masks
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Embed images with SigLIP and language tokens with embedding layer to prepare
+        for PaliGemma transformer processing.
+        """
+        embs = []
+        pad_masks = []
+        att_masks = []
+
+        # Process images
+        for img, img_mask in zip(images, img_masks, strict=True):
+
+            def image_embed_func(img):
+                return self.paligemma_with_expert.embed_image(img)
+
+            img_emb = self._apply_checkpoint(image_embed_func, img)
+
+            bsize, num_img_embs = img_emb.shape[:2]
+
+            embs.append(img_emb)
+            pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs))
+
+            # Create attention masks so that image tokens attend to each other
+            att_masks += [0] * num_img_embs
+
+        # Process language tokens
+        def lang_embed_func(lang_tokens):
+            lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens)
+            # Transformers > 5.4 scales Gemma token embeddings inside embed_tokens, matching
+            # OpenPI's former explicit sqrt(hidden_size) multiply without applying it twice.
+            # See https://github.com/huggingface/transformers/pull/44432/changes#diff-5f76eac6f18f4b491521314c318a9692318feb4d19228e9576cce7bde4240834
+            return lang_emb
+
+        lang_emb = self._apply_checkpoint(lang_embed_func, lang_tokens)
+
+        embs.append(lang_emb)
+        pad_masks.append(lang_masks)
+
+        # full attention between image and language inputs
+        num_lang_embs = lang_emb.shape[1]
+        att_masks += [0] * num_lang_embs
+
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
+
+        # Get batch size from the first dimension of the concatenated tensors
+        bsize = pad_masks.shape[0]
+        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
+
+        return embs, pad_masks, att_masks
+
+    def embed_suffix(self, state, noisy_actions, timestep):
+        """Embed state, noisy_actions, timestep to prepare for Expert Gemma processing."""
+        embs = []
+        pad_masks = []
+        att_masks = []
+
+        if not self.pi05:
+            if self.state_proj.weight.dtype == torch.float32:
+                state = state.to(torch.float32)
+
+            # Embed state
+            def state_proj_func(state):
+                return self.state_proj(state)
+
+            state_emb = self._apply_checkpoint(state_proj_func, state)
+
+            embs.append(state_emb[:, None, :])
+            bsize = state_emb.shape[0]
+            device = state_emb.device
+
+            state_mask = torch.ones(bsize, 1, dtype=torch.bool, device=device)
+            pad_masks.append(state_mask)
+
+            # Set attention masks so that image and language inputs do not attend to state or actions
+            att_masks += [1]
+
+        # Embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.action_in_proj.out_features,
+            min_period=4e-3,
+            max_period=4.0,
+            device=timestep.device,
+        )
+        time_emb = time_emb.type(dtype=timestep.dtype)
+
+        # Fuse timestep + action information using an MLP
+        def action_proj_func(noisy_actions):
+            return self.action_in_proj(noisy_actions)
+
+        action_emb = self._apply_checkpoint(action_proj_func, noisy_actions)
+
+        if not self.pi05:
+            time_emb = time_emb[:, None, :].expand_as(action_emb)
+            action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+
+            # Apply MLP layers
+            def mlp_func(action_time_emb):
+                x = self.action_time_mlp_in(action_time_emb)
+                x = F.silu(x)  # swish == silu
+                return self.action_time_mlp_out(x)
+
+            action_time_emb = self._apply_checkpoint(mlp_func, action_time_emb)
+            adarms_cond = None
+        else:
+            # time MLP (for adaRMS)
+            def time_mlp_func(time_emb):
+                x = self.time_mlp_in(time_emb)
+                x = F.silu(x)  # swish == silu
+                x = self.time_mlp_out(x)
+                return F.silu(x)
+
+            time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
+            action_time_emb = action_emb
+            adarms_cond = time_emb
+
+        # Add to input tokens
+        embs.append(action_time_emb)
+
+        bsize, action_time_dim = action_time_emb.shape[:2]
+        action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
+        pad_masks.append(action_time_mask)
+
+        # Set attention masks so that image, language and state inputs do not attend to action tokens
+        att_masks += [1] + ([0] * (self.config.action_horizon - 1))
+
+        embs = torch.cat(embs, dim=1)
+        pad_masks = torch.cat(pad_masks, dim=1)
+        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
+        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
+
+        return embs, pad_masks, att_masks, adarms_cond
+
+    def forward(self, observation, actions, noise=None, time=None) -> Tensor:
+        """Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)"""
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(
+            observation, train=True
+        )
+
+        if noise is None:
+            noise = self.sample_noise(actions.shape, actions.device)
+
+        if time is None:
+            time = self.sample_time(actions.shape[0], actions.device)
+
+        time_expanded = time[:, None, None]
+        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        u_t = noise - actions
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, img_masks, lang_tokens, lang_masks
+        )
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, time)
+        if (
+            self.paligemma_with_expert.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
+            == torch.bfloat16
+        ):
+            suffix_embs = suffix_embs.to(dtype=torch.bfloat16)
+            prefix_embs = prefix_embs.to(dtype=torch.bfloat16)
+
+        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
+        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
+
+        att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
+        position_ids = torch.cumsum(pad_masks, dim=1) - 1
+
+        # Prepare attention masks
+        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
+
+        # Apply gradient checkpointing if enabled
+        def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
+            (_, suffix_out), _ = self.paligemma_with_expert.forward(
+                attention_mask=att_2d_masks_4d,
+                position_ids=position_ids,
+                past_key_values=None,
+                inputs_embeds=[prefix_embs, suffix_embs],
+                use_cache=False,
+                adarms_cond=[None, adarms_cond],
+            )
+            return suffix_out
+
+        suffix_out = self._apply_checkpoint(
+            forward_func, prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond
+        )
+
+        suffix_out = suffix_out[:, -self.config.action_horizon :]
+        suffix_out = suffix_out.to(dtype=torch.float32)
+
+        # Apply gradient checkpointing to final action projection if enabled
+        def action_out_proj_func(suffix_out):
+            return self.action_out_proj(suffix_out)
+
+        v_t = self._apply_checkpoint(action_out_proj_func, suffix_out)
+
+        return F.mse_loss(u_t, v_t, reduction="none")
+
+    @torch.no_grad()
+    def sample_actions(self, device, observation, noise=None, num_steps=10) -> Tensor:
+        """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)"""
+        bsize = observation.state.shape[0]
+        if noise is None:
+            actions_shape = (bsize, self.config.action_horizon, self.config.action_dim)
+            noise = self.sample_noise(actions_shape, device)
+
+        images, img_masks, lang_tokens, lang_masks, state = self._preprocess_observation(
+            observation, train=False
+        )
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            images, img_masks, lang_tokens, lang_masks
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
+
+        # Compute image and language key value cache
+        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
+        self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
+
+        _, past_key_values = self.paligemma_with_expert.forward(
+            attention_mask=prefix_att_2d_masks_4d,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None],
+            use_cache=True,
+        )
+
+        dt = -1.0 / num_steps
+        dt = torch.tensor(dt, dtype=torch.float32, device=device)
+
+        x_t = noise
+        time = torch.tensor(1.0, dtype=torch.float32, device=device)
+        while time >= -dt / 2:
+            expanded_time = time.expand(bsize)
+            v_t = self.denoise_step(
+                state,
+                prefix_pad_masks,
+                past_key_values,
+                x_t,
+                expanded_time,
+            )
+
+            # Euler step - use new tensor assignment instead of in-place operation
+            x_t = x_t + dt * v_t
+            time += dt
+        return x_t
+
+    def denoise_step(
+        self,
+        state,
+        prefix_pad_masks,
+        past_key_values,
+        x_t,
+        timestep,
+    ):
+        """Apply one denoising step of the noise `x_t` at a given timestep."""
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, timestep)
+
+        suffix_len = suffix_pad_masks.shape[1]
+        batch_size = prefix_pad_masks.shape[0]
+        prefix_len = prefix_pad_masks.shape[1]
+
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+
+        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
+        position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
+
+        # Prepare attention masks
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
+        self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
+
+        past_key_values = copy.deepcopy(past_key_values)
+        outputs_embeds, _ = self.paligemma_with_expert.forward(
+            attention_mask=full_att_2d_masks_4d,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, suffix_embs],
+            use_cache=False,
+            adarms_cond=[None, adarms_cond],
+        )
+
+        suffix_out = outputs_embeds[1]
+        suffix_out = suffix_out[:, -self.config.action_horizon :]
+        suffix_out = suffix_out.to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)

tests/policies/pi0_pi05/openpi_pytorch/preprocessing_pytorch.py

@@ -0,0 +1,179 @@
+import logging
+from collections.abc import Sequence
+
+import torch
+
+from tests.policies.pi0_pi05.openpi_pytorch import image_tools
+
+logger = logging.getLogger("openpi")
+
+# Constants moved from model.py
+IMAGE_KEYS = (
+    "base_0_rgb",
+    "left_wrist_0_rgb",
+    "right_wrist_0_rgb",
+)
+
+IMAGE_RESOLUTION = (224, 224)
+
+
+def preprocess_observation_pytorch(
+    observation,
+    *,
+    train: bool = False,
+    image_keys: Sequence[str] = IMAGE_KEYS,
+    image_resolution: tuple[int, int] = IMAGE_RESOLUTION,
+):
+    """Torch.compile-compatible version of preprocess_observation_pytorch with simplified type annotations.
+
+    This function avoids complex type annotations that can cause torch.compile issues.
+    """
+    if not set(image_keys).issubset(observation.images):
+        raise ValueError(f"images dict missing keys: expected {image_keys}, got {list(observation.images)}")
+
+    batch_shape = observation.state.shape[:-1]
+
+    out_images = {}
+    for key in image_keys:
+        image = observation.images[key]
+
+        # TODO: This is a hack to handle both [B, C, H, W] and [B, H, W, C] formats
+        # Handle both [B, C, H, W] and [B, H, W, C] formats
+        is_channels_first = image.shape[1] == 3  # Check if channels are in dimension 1
+
+        if is_channels_first:
+            # Convert [B, C, H, W] to [B, H, W, C] for processing
+            image = image.permute(0, 2, 3, 1)
+
+        if image.shape[1:3] != image_resolution:
+            logger.info(f"Resizing image {key} from {image.shape[1:3]} to {image_resolution}")
+            image = image_tools.resize_with_pad_torch(image, *image_resolution)
+
+        if train:
+            # Convert from [-1, 1] to [0, 1] for PyTorch augmentations
+            image = image / 2.0 + 0.5
+
+            # Apply PyTorch-based augmentations
+            if "wrist" not in key:
+                # Geometric augmentations for non-wrist cameras
+                height, width = image.shape[1:3]
+
+                # Random crop and resize
+                crop_height = int(height * 0.95)
+                crop_width = int(width * 0.95)
+
+                # Random crop
+                max_h = height - crop_height
+                max_w = width - crop_width
+                if max_h > 0 and max_w > 0:
+                    # Use tensor operations instead of .item() for torch.compile compatibility
+                    start_h = torch.randint(0, max_h + 1, (1,), device=image.device)
+                    start_w = torch.randint(0, max_w + 1, (1,), device=image.device)
+                    image = image[:, start_h : start_h + crop_height, start_w : start_w + crop_width, :]
+
+                # Resize back to original size
+                image = torch.nn.functional.interpolate(
+                    image.permute(0, 3, 1, 2),  # [b, h, w, c] -> [b, c, h, w]
+                    size=(height, width),
+                    mode="bilinear",
+                    align_corners=False,
+                ).permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
+
+                # Random rotation (small angles)
+                # Use tensor operations instead of .item() for torch.compile compatibility
+                angle = torch.rand(1, device=image.device) * 10 - 5  # Random angle between -5 and 5 degrees
+                if torch.abs(angle) > 0.1:  # Only rotate if angle is significant
+                    # Convert to radians
+                    angle_rad = angle * torch.pi / 180.0
+
+                    # Create rotation matrix
+                    cos_a = torch.cos(angle_rad)
+                    sin_a = torch.sin(angle_rad)
+
+                    # Apply rotation using grid_sample
+                    grid_x = torch.linspace(-1, 1, width, device=image.device)
+                    grid_y = torch.linspace(-1, 1, height, device=image.device)
+
+                    # Create meshgrid
+                    grid_y, grid_x = torch.meshgrid(grid_y, grid_x, indexing="ij")
+
+                    # Expand to batch dimension
+                    grid_x = grid_x.unsqueeze(0).expand(image.shape[0], -1, -1)
+                    grid_y = grid_y.unsqueeze(0).expand(image.shape[0], -1, -1)
+
+                    # Apply rotation transformation
+                    grid_x_rot = grid_x * cos_a - grid_y * sin_a
+                    grid_y_rot = grid_x * sin_a + grid_y * cos_a
+
+                    # Stack and reshape for grid_sample
+                    grid = torch.stack([grid_x_rot, grid_y_rot], dim=-1)
+
+                    image = torch.nn.functional.grid_sample(
+                        image.permute(0, 3, 1, 2),  # [b, h, w, c] -> [b, c, h, w]
+                        grid,
+                        mode="bilinear",
+                        padding_mode="zeros",
+                        align_corners=False,
+                    ).permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
+
+            # Color augmentations for all cameras
+            # Random brightness
+            # Use tensor operations instead of .item() for torch.compile compatibility
+            brightness_factor = (
+                0.7 + torch.rand(1, device=image.device) * 0.6
+            )  # Random factor between 0.7 and 1.3
+            image = image * brightness_factor
+
+            # Random contrast
+            # Use tensor operations instead of .item() for torch.compile compatibility
+            contrast_factor = (
+                0.6 + torch.rand(1, device=image.device) * 0.8
+            )  # Random factor between 0.6 and 1.4
+            mean = image.mean(dim=[1, 2, 3], keepdim=True)
+            image = (image - mean) * contrast_factor + mean
+
+            # Random saturation (convert to HSV, modify S, convert back)
+            # For simplicity, we'll just apply a random scaling to the color channels
+            # Use tensor operations instead of .item() for torch.compile compatibility
+            saturation_factor = (
+                0.5 + torch.rand(1, device=image.device) * 1.0
+            )  # Random factor between 0.5 and 1.5
+            gray = image.mean(dim=-1, keepdim=True)
+            image = gray + (image - gray) * saturation_factor
+
+            # Clamp values to [0, 1]
+            image = torch.clamp(image, 0, 1)
+
+            # Back to [-1, 1]
+            image = image * 2.0 - 1.0
+
+        # Convert back to [B, C, H, W] format if it was originally channels-first
+        if is_channels_first:
+            image = image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]
+
+        out_images[key] = image
+
+    # obtain mask
+    out_masks = {}
+    for key in out_images:
+        if key not in observation.image_masks:
+            # do not mask by default
+            out_masks[key] = torch.ones(batch_shape, dtype=torch.bool, device=observation.state.device)
+        else:
+            out_masks[key] = observation.image_masks[key]
+
+    # Create a simple object with the required attributes instead of using the complex Observation class
+    class SimpleProcessedObservation:
+        def __init__(self, **kwargs):
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+    return SimpleProcessedObservation(
+        images=out_images,
+        image_masks=out_masks,
+        state=observation.state,
+        tokenized_prompt=observation.tokenized_prompt,
+        tokenized_prompt_mask=observation.tokenized_prompt_mask,
+        token_ar_mask=observation.token_ar_mask,
+        token_loss_mask=observation.token_loss_mask,
+    )

tests/policies/pi0_pi05/test_pi05_compile.py

@@ -0,0 +1,101 @@
+#!/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 os
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.modeling_pi05 import PI05Pytorch  # noqa: E402
+from tests.policies.pi0_pi05.utils.torch_compile import (  # noqa: E402
+    assert_cache_stability,
+    assert_compiled_output_matches_eager,
+    assert_explain_has_no_graph_breaks,
+    benchmark_runtime,
+    make_compile_config,
+    reset_compile_state,
+)
+from tests.utils import require_cuda  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="torch.compile benchmark is too slow for CI; run manually on GPU nodes",
+)
+
+
+def _make_model(*, compile_model):
+    return PI05Pytorch(make_compile_config(PI05Config, compile_model=compile_model)).cuda().eval()
+
+
+def _make_dummy_inputs(config):
+    device = torch.device("cuda")
+    common = {
+        "images": [torch.randn(1, 3, *config.image_resolution, device=device)],
+        "img_masks": [torch.ones(1, dtype=torch.bool, device=device)],
+        "tokens": torch.randint(0, 1024, (1, 5), dtype=torch.long, device=device),
+        "masks": torch.ones(1, 5, dtype=torch.bool, device=device),
+    }
+    forward_kwargs = {
+        **common,
+        "actions": torch.randn(1, config.chunk_size, config.max_action_dim, device=device),
+        "noise": torch.randn(1, config.chunk_size, config.max_action_dim, device=device),
+        "time": torch.rand(1, device=device),
+    }
+    sample_kwargs = {
+        **common,
+        "noise": torch.randn(1, config.chunk_size, config.max_action_dim, device=device),
+        "num_steps": config.num_inference_steps,
+    }
+    return forward_kwargs, sample_kwargs
+
+
+@require_cuda
+def test_pi05_torch_compile_forward_and_sample_actions():
+    if not hasattr(torch, "compile"):
+        pytest.skip("torch.compile is not available")
+    if not torch._dynamo.is_dynamo_supported():
+        pytest.skip("torch._dynamo is not supported on this platform")
+
+    torch.manual_seed(0)
+    eager_model = _make_model(compile_model=False)
+    torch.manual_seed(0)
+    compiled_model = _make_model(compile_model=True)
+    forward_kwargs, sample_kwargs = _make_dummy_inputs(compiled_model.config)
+
+    try:
+        assert_compiled_output_matches_eager(eager_model, compiled_model, forward_kwargs, sample_kwargs)
+
+        assert_explain_has_no_graph_breaks(eager_model.forward, forward_kwargs, "pi05.forward")
+        assert_explain_has_no_graph_breaks(eager_model.sample_actions, sample_kwargs, "pi05.sample_actions")
+
+        assert_cache_stability(compiled_model.forward, forward_kwargs, "pi05.forward")
+        assert_cache_stability(compiled_model.sample_actions, sample_kwargs, "pi05.sample_actions")
+
+        benchmark_runtime(eager_model.forward, compiled_model.forward, forward_kwargs, "pi05.forward")
+        benchmark_runtime(
+            eager_model.sample_actions,
+            compiled_model.sample_actions,
+            sample_kwargs,
+            "pi05.sample_actions",
+        )
+    finally:
+        reset_compile_state()
+        del eager_model
+        del compiled_model
+        torch.cuda.empty_cache()

tests/policies/pi0_pi05/test_pi05_original_vs_lerobot.py

@@ -14,52 +14,56 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Test script to verify PI0OpenPI policy integration with LeRobot vs the original implementation"""
+"""Compare LeRobot PI0.5 against the vendored OpenPI PyTorch reference."""
 
+import gc
 import os
-from copy import deepcopy
-from typing import Any
 
-import numpy as np
 import pytest
 import torch
 
-# Skip if openpi or transformers is not available
-pytest.importorskip("openpi")
 pytest.importorskip("transformers")
 
-# Skip this entire module in CI
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local OpenPI installation and is not meant for CI",
+from lerobot.configs import PreTrainedConfig  # noqa: E402
+from lerobot.policies.pi05 import PI05Policy  # noqa: E402
+from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+from tests.policies.pi0_pi05.openpi_pytorch.pi0_pytorch import PI0Pytorch  # noqa: E402
+from tests.policies.pi0_pi05.utils.openpi_parity import (  # noqa: E402
+    assert_processor_inputs_match_lerobot,
+    clone_batch,
+    deterministic_openpi_forward_preprocess,
+    fix_reference_state_dict,
+    fixed_flow_sampling,
+    load_openpi_reference_state_dict,
+    make_openpi_observation_from_raw,
+    openpi_model_actions_from_raw,
 )
 
-from openpi.models_pytorch import preprocessing_pytorch as openpi_preprocessing  # noqa: E402
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="OpenPI parity and torch.compile checks are too slow for CI; run manually on GPU nodes",
+)
 
-# NOTE: Assumes PYTHONPATH is set to include OpenPI src as per instructions.
-from openpi.models_pytorch.pi0_pytorch import PI0Pytorch  # noqa: E402
-from transformers import AutoTokenizer  # noqa: E402
-
-from lerobot.policies.pi05 import PI05Config, PI05Policy  # noqa: E402
-from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors  # noqa: E402
-from lerobot.processor import PolicyProcessorPipeline  # noqa: E402
-from lerobot.types import PolicyAction  # noqa: E402
-
-# TODO: ADDING DEFAULT IMAGES_FEATURES TO CONFIG
 DUMMY_ACTION_DIM = 32
 DUMMY_STATE_DIM = 32
 DUMMY_ACTION_HORIZON = 50
 DUMMY_MAX_TOKEN_LEN = 200
-DEVICE = "cpu"  # Use CPU to avoid memory issues for testing
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+COMPILE_MODE = "default"
+FORWARD_RTOL = 1e-4
+FORWARD_ATOL = 1e-4
+SAMPLE_RTOL = 1e-2
+SAMPLE_ATOL = 5e-3
 
 DUMMY_DATASET_STATS = {
-    "observation.state": {
+    OBS_STATE: {
         "mean": torch.zeros(DUMMY_STATE_DIM),
         "std": torch.ones(DUMMY_STATE_DIM),
         "q01": torch.zeros(DUMMY_STATE_DIM),
         "q99": torch.ones(DUMMY_STATE_DIM),
     },
-    "action": {
+    ACTION: {
         "mean": torch.zeros(DUMMY_ACTION_DIM),
         "std": torch.ones(DUMMY_ACTION_DIM),
         "q01": torch.zeros(DUMMY_ACTION_DIM),
@@ -88,6 +92,15 @@ DUMMY_DATASET_STATS = {
 }
 
 
+@pytest.fixture(autouse=True)
+def cleanup_cuda_after_test():
+    yield
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
 class PI05BaseOriginalConfig:
     action_dim: int = DUMMY_ACTION_DIM
     action_horizon: int = DUMMY_ACTION_HORIZON
@@ -96,341 +109,163 @@ class PI05BaseOriginalConfig:
     precision: str = "float32"
     pi05: bool = True
     dtype: str = "float32"
+    pytorch_compile_mode: str | None = None
 
 
-def instantiate_lerobot_pi05(
-    from_pretrained: bool = False,
-) -> tuple[
-    PI05Policy,
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    if from_pretrained:
-        # Load the policy first
-        policy = PI05Policy.from_pretrained(pretrained_name_or_path="lerobot/pi05_base", strict=True)
-    else:
-        config = PI05Config(max_action_dim=DUMMY_ACTION_DIM, max_state_dim=DUMMY_STATE_DIM, dtype="float32")
-        policy = PI05Policy(config)
+def instantiate_lerobot_pi05(*, compile_model: bool = False, gradient_checkpointing: bool = False):
+    config = PreTrainedConfig.from_pretrained("lerobot/pi05_base")
+    config.device = str(DEVICE)
+    config.dtype = "float32"
+    config.compile_model = compile_model
+    config.compile_mode = COMPILE_MODE
+    config.gradient_checkpointing = gradient_checkpointing
 
+    policy = PI05Policy.from_pretrained("lerobot/pi05_base", config=config, strict=True)
     policy.to(DEVICE)
-    policy.config.device = DEVICE
-    preprocessor, postprocessor = make_pi05_pre_post_processors(
-        config=policy.config, dataset_stats=DUMMY_DATASET_STATS
-    )
-    return (policy, preprocessor, postprocessor)
+    policy.config.device = str(DEVICE)
+    preprocessor, _ = make_pi05_pre_post_processors(config=policy.config, dataset_stats=DUMMY_DATASET_STATS)
+    return policy, preprocessor
 
 
-def instantiate_original_pi05(from_pretrained: bool = False, model_path: str | None = None):
-    config = PI05BaseOriginalConfig()
-    policy = PI0Pytorch(config)
+def instantiate_original_pi05():
+    policy = PI0Pytorch(PI05BaseOriginalConfig()).to(DEVICE)
 
-    if from_pretrained:
-        try:
-            print("Loading converted PyTorch weights from HuggingFace Hub (lerobot/pi05_base)...")
-
-            # Download the model from HuggingFace Hub
-            import safetensors.torch
-            from huggingface_hub import snapshot_download
-
-            # Download the entire repository
-            if model_path and os.path.exists(model_path):
-                cache_dir = model_path
-                print(f"Using cached model from: {cache_dir}")
-            else:
-                cache_dir = snapshot_download(repo_id="lerobot/pi05_base", repo_type="model")
-                print(f"Downloaded model to: {cache_dir}")
-
-            # Try to load safetensors format first
-            model_file = os.path.join(cache_dir, "model.safetensors")
-            if os.path.exists(model_file):
-                state_dict = safetensors.torch.load_file(model_file)
-                print(f"Loaded {len(state_dict)} parameters from safetensors")
-            else:
-                raise FileNotFoundError(f"No safetensors file found in {cache_dir}")
-
-            # Load the state dict into the model
-            missing_keys, unexpected_keys = policy.load_state_dict(state_dict, strict=False)
-
-            if missing_keys:
-                print(f"Missing keys: {len(missing_keys)}")
-                if len(missing_keys) <= 5:
-                    for key in missing_keys:
-                        print(f"    - {key}")
-                else:
-                    for key in missing_keys[:5]:
-                        print(f"    - {key}")
-                    print(f"    ... and {len(missing_keys) - 5} more")
-
-            if unexpected_keys:
-                print(f"Unexpected keys: {len(unexpected_keys)}")
-                if len(unexpected_keys) <= 5:
-                    for key in unexpected_keys:
-                        print(f"    - {key}")
-                else:
-                    for key in unexpected_keys[:5]:
-                        print(f"    - {key}")
-                    print(f"    ... and {len(unexpected_keys) - 5} more")
-
-            if not missing_keys and not unexpected_keys:
-                print("All pretrained weights loaded successfully!")
-            else:
-                print("Pretrained weights loaded with some missing/unexpected keys (this may be normal)")
-
-        except Exception as e:
-            print(f"Failed to load pretrained weights: {e}")
-            print("   Using randomly initialized weights...")
-            import traceback
-
-            traceback.print_exc()
-
-    policy.to(DEVICE)
+    # NOTE: `lerobot/pi05_base` 的 LeRobot loader 和 PI0 一样会在 strict load 前做 key
+    # 兼容转换，因此预期没有 missing_keys 或 unexpected_keys。vendored reference 则是裸
+    # `nn.Module`，需要在测试侧补齐 checkpoint 与模块命名之间的最小差异。
+    # NOTE: `lm_head.weight` 是 PaliGemma tied embedding 的保存名；LeRobot 的
+    # from_pretrained 会把它映射到内部 `embed_tokens.weight`，而 reference 模型没有这层
+    # loader，所以这里手动复用同一份 tensor，避免把权重别名差异误判成模型差异。
+    state_dict = fix_reference_state_dict(load_openpi_reference_state_dict("lerobot/pi05_base"))
+    missing_keys, unexpected_keys = policy.load_state_dict(state_dict, strict=False)
+    assert missing_keys == []
+    assert unexpected_keys == []
     return policy
 
 
 def create_dummy_data():
-    batch_size = 2  # Reduce batch size for testing
-    device = DEVICE
-
-    # Use the exact same prompt for both implementations
+    batch_size = 2
     prompt = "Pick up the red block and place it in the bin"
-
-    batch = {
-        "observation.state": torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device),
-        "action": torch.randn(
-            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=device
+    return {
+        OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=DEVICE),
+        ACTION: torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=DEVICE
         ),
-        # Create images in [0, 1] range as expected by LeRobot (will be converted to [-1, 1] internally)
         "observation.images.base_0_rgb": torch.rand(
-            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+            batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
         ),
         "observation.images.left_wrist_0_rgb": torch.rand(
-            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+            batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
         ),
         "observation.images.right_wrist_0_rgb": torch.rand(
-            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+            batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
         ),
-        # Add the task prompt for LeRobot - provide as list with single element to trigger expansion
         "task": [prompt for _ in range(batch_size)],
     }
-    return batch
 
 
-def extract_lerobot_processed_inputs(lerobot_pi0, batch):
-    """Extract the exact same processed inputs that LeRobot uses internally."""
-    # Get the tokenized language from LeRobot's internal method
-    lang_tokens, lang_masks = lerobot_pi0._tokenize_language(batch)
-
-    # Get the preprocessed images from LeRobot's internal method
-    images, img_masks = lerobot_pi0._preprocess_images(batch, train=False)
-
-    # Create dummy token_ar_mask and token_loss_mask for original implementation
-    token_ar_mask = torch.zeros_like(lang_tokens, dtype=torch.int32)
-    token_loss_mask = torch.ones_like(lang_masks, dtype=torch.bool)
-
-    return images, img_masks, lang_tokens, lang_masks, token_ar_mask, token_loss_mask
+def prepare_parity_inputs(lerobot_pi05, lerobot_preprocessor):
+    torch.manual_seed(0)
+    raw_batch = create_dummy_data()
+    lerobot_batch = lerobot_preprocessor(clone_batch(raw_batch))
+    openpi_observation = make_openpi_observation_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        max_token_len=DUMMY_MAX_TOKEN_LEN,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=True,
+    )
+    openpi_actions = openpi_model_actions_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=True,
+    )
+    assert_processor_inputs_match_lerobot(
+        lerobot_pi05,
+        lerobot_batch,
+        openpi_observation,
+        compare_state=False,
+    )
+    batch_size = raw_batch[OBS_STATE].shape[0]
+    noise = torch.randn(
+        batch_size,
+        DUMMY_ACTION_HORIZON,
+        DUMMY_ACTION_DIM,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    time = torch.linspace(0.2, 0.8, batch_size, dtype=torch.float32, device=DEVICE)
+    return lerobot_batch, openpi_observation, openpi_actions, noise, time
 
 
-class PI05Observation:
-    """Observation class that matches the original OpenPI format."""
-
-    def __init__(
-        self,
-        state,
-        images,
-        image_masks,
-        tokenized_prompt,
-        tokenized_prompt_mask,
-        token_ar_mask,
-        token_loss_mask,
-    ):
-        self.state = state
-        self.images = images
-        self.image_masks = image_masks
-        self.tokenized_prompt = tokenized_prompt
-        self.tokenized_prompt_mask = tokenized_prompt_mask
-        self.token_ar_mask = token_ar_mask
-        self.token_loss_mask = token_loss_mask
-
-
-def create_original_observation_with_openpi_preprocessing(batch):
-    """Create observation object for OpenPI using OpenPI's own preprocessing with pi05 state tokenizer."""
-    batch_size = batch["observation.state"].shape[0]
-    device = batch["observation.state"].device
-
-    # Create tokenizer for OpenPI (same as LeRobot uses)
-    tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
-
-    # Get task description (pi05 processor handles all text formatting)
-    tasks = batch.get("task", ["Pick up the object"] * batch_size)
-    if isinstance(tasks, str):
-        tasks = [tasks] * batch_size
-    elif len(tasks) == 1:
-        tasks = tasks * batch_size
-
-    # Use pi05 state and input tokenizer logic (same as Pi05PrepareStateTokenizerProcessorStep)
-    state = batch["observation.state"]
-    state = deepcopy(state)
-
-    # Prepare state (pad to max_state_dim)
-    from lerobot.policies.pi05.modeling_pi05 import pad_vector
-
-    state = pad_vector(state, DUMMY_STATE_DIM)
-
-    # Normalize state to [-1, 1] range if needed (assuming it's already normalized from normalize_inputs)
-    # Discretize into 256 bins (see openpi `PaligemmaTokenizer.tokenize()`)
-    state_np = state.cpu().numpy()
-    discretized_states = np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
-
-    # Create pi05-formatted prompts that include state information
-    full_prompts = []
-    for i, task in enumerate(tasks):
-        cleaned_text = task.strip().replace("_", " ").replace("\n", " ")
-        state_str = " ".join(map(str, discretized_states[i]))
-        full_prompt = f"Task: {cleaned_text}, State: {state_str};\nAction: "
-        full_prompts.append(full_prompt)
-
-    # Tokenize with max_length padding to match OpenPI's expected format
-    tokenized = tokenizer(
-        full_prompts,
-        padding="max_length",
-        padding_side="right",
-        truncation=True,
-        max_length=DUMMY_MAX_TOKEN_LEN,
-        return_tensors="pt",
+def assert_forward_matches(*, compile_model: bool = False, gradient_checkpointing: bool = False):
+    lerobot_pi05, lerobot_preprocessor = instantiate_lerobot_pi05(
+        compile_model=compile_model,
+        gradient_checkpointing=gradient_checkpointing,
+    )
+    original_pi05 = instantiate_original_pi05()
+    lerobot_batch, openpi_observation, openpi_actions, noise, time = prepare_parity_inputs(
+        lerobot_pi05,
+        lerobot_preprocessor,
     )
 
-    lang_tokens = tokenized["input_ids"].to(device)
-    lang_masks = tokenized["attention_mask"].to(device, dtype=torch.bool)
+    if gradient_checkpointing:
+        lerobot_pi05.train()
+    else:
+        lerobot_pi05.eval()
+    original_pi05.eval()
 
-    # Create dummy token_ar_mask and token_loss_mask for OpenPI
-    token_ar_mask = torch.zeros_like(lang_tokens, dtype=torch.int32)
-    token_loss_mask = torch.ones_like(lang_masks, dtype=torch.bool)
+    with fixed_flow_sampling(lerobot_pi05.model, noise=noise, time=time):
+        lerobot_loss, _ = lerobot_pi05(lerobot_batch, reduction="none")
+    with deterministic_openpi_forward_preprocess(original_pi05):
+        openpi_losses = original_pi05(openpi_observation, openpi_actions, noise=noise, time=time)
+    openpi_loss = openpi_losses.mean(dim=(1, 2))
 
-    # Convert LeRobot images format to OpenPI format (convert [0,1] to [-1,1] range)
-    image_dict = {
-        "base_0_rgb": batch["observation.images.base_0_rgb"] * 2.0 - 1.0,
-        "left_wrist_0_rgb": batch["observation.images.left_wrist_0_rgb"] * 2.0 - 1.0,
-        "right_wrist_0_rgb": batch["observation.images.right_wrist_0_rgb"] * 2.0 - 1.0,
-    }
+    torch.testing.assert_close(lerobot_loss, openpi_loss, rtol=FORWARD_RTOL, atol=FORWARD_ATOL)
 
-    # Create image masks (all ones for real images)
-    image_masks_dict = {}
-    for key in image_dict:
-        image_masks_dict[key] = torch.ones(batch_size, dtype=torch.bool, device=device)
 
-    # Create raw observation object (before preprocessing)
-    raw_observation = PI05Observation(
-        state=batch["observation.state"],
-        images=image_dict,
-        image_masks=image_masks_dict,
-        tokenized_prompt=lang_tokens,
-        tokenized_prompt_mask=lang_masks,
-        token_ar_mask=token_ar_mask,
-        token_loss_mask=token_loss_mask,
+def assert_sample_actions_match_openpi(*, compile_model: bool = False):
+    lerobot_pi05, lerobot_preprocessor = instantiate_lerobot_pi05(compile_model=compile_model)
+    original_pi05 = instantiate_original_pi05()
+    lerobot_batch, openpi_observation, _openpi_actions, noise, _time = prepare_parity_inputs(
+        lerobot_pi05,
+        lerobot_preprocessor,
     )
 
-    # Now use OpenPI's preprocessing
-    processed_obs = openpi_preprocessing.preprocess_observation_pytorch(raw_observation, train=False)
-
-    return processed_obs
-
-
-def create_original_observation_from_lerobot(lerobot_pi0, batch):
-    """Create observation object compatible with original OpenPI using the exact same inputs as LeRobot."""
-    _batch_size = batch["observation.state"].shape[0]
-    _device = batch["observation.state"].device
-
-    # Extract the exact same processed inputs that LeRobot uses
-    images, img_masks, lang_tokens, lang_masks, token_ar_mask, token_loss_mask = (
-        extract_lerobot_processed_inputs(lerobot_pi0, batch)
-    )
-
-    # Convert images list to dict with original OpenPI keys
-    image_dict = {
-        "base_0_rgb": images[0],
-        "left_wrist_0_rgb": images[1],
-        "right_wrist_0_rgb": images[2],
-    }
-
-    # Convert image masks list to dict with original OpenPI keys
-    image_masks_dict = {
-        "base_0_rgb": img_masks[0],
-        "left_wrist_0_rgb": img_masks[1],
-        "right_wrist_0_rgb": img_masks[2],
-    }
-
-    return PI05Observation(
-        state=batch["observation.state"],
-        images=image_dict,
-        image_masks=image_masks_dict,
-        tokenized_prompt=lang_tokens,
-        tokenized_prompt_mask=lang_masks,
-        token_ar_mask=token_ar_mask,
-        token_loss_mask=token_loss_mask,
-    )
-
-
-def test_pi05_original_vs_lerobot():
-    """Test PI05 original implementation vs LeRobot implementation."""
-    print("Initializing models...")
-    lerobot_pi05, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_pi05(
-        from_pretrained=True
-    )  # Load pretrained LeRobot model
-    original_pi0 = instantiate_original_pi05(
-        from_pretrained=True
-    )  # Load pretrained OpenPI model from HuggingFace Hub
-
-    print("Creating dummy data...")
-    batch = create_dummy_data()
-    batch_lerobot = deepcopy(batch)
-
-    # Test each model with its own preprocessing (more realistic end-to-end test)
-    print("\nTest each model with its own preprocessing")
-    print("Creating observation for OpenPI using OpenPI's own preprocessing...")
-    pi0_obs_openpi = create_original_observation_with_openpi_preprocessing(batch)
-
-    print(f"Task prompt: '{batch['task'][0]}'")
-    print(f"OpenPI tokenized prompt shape: {pi0_obs_openpi.tokenized_prompt.shape}")
-    print(f"OpenPI image shapes: {[img.shape for img in pi0_obs_openpi.images.values()]}")
-    print(f"OpenPI state shape: {pi0_obs_openpi.state.shape}")
-
-    print("Testing OpenPI with own preprocessing...")
-    original_pi0.eval()
-    torch.manual_seed(42)  # Set seed for reproducibility
-    batch_size = batch["observation.state"].shape[0]
-    noise_shape = (batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM)
-    fixed_noise = torch.randn(noise_shape, dtype=torch.float32, device=DEVICE)
-
-    with torch.no_grad():
-        openpi_actions = original_pi0.sample_actions(
-            device=DEVICE, observation=pi0_obs_openpi, noise=fixed_noise, num_steps=10
-        )
-        openpi_actions_unit = openpi_actions[:, 0, :]
-    print(f"OpenPI (own preprocessing) Actions shape: {openpi_actions.shape}")
-    print(f"OpenPI (own preprocessing) Actions unit shape: {openpi_actions_unit.shape}")
-    print(f"OpenPI (own preprocessing) Actions mean: {openpi_actions.mean().item():.6f}")
-    print(f"OpenPI (own preprocessing) Actions std: {openpi_actions.std().item():.6f}")
-
-    print("Testing LeRobot with own preprocessing...")
     lerobot_pi05.eval()
-    torch.manual_seed(42)  # Set the same seed
-
-    batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
+    original_pi05.eval()
     with torch.no_grad():
-        lerobot_actions_own = lerobot_pi05.predict_action_chunk(
-            batch_lerobot_processed
-        )  # batch_size, n_action_steps, action_dim
-        lerobot_actions_unit = lerobot_actions_own[:, 0, :]
-    print(f"LeRobot (own preprocessing) Actions shape: {lerobot_actions_own.shape}")
-    print(f"LeRobot (own preprocessing) Actions unit shape: {lerobot_actions_unit.shape}")
-    print(f"LeRobot (own preprocessing) Actions mean: {lerobot_actions_own.mean().item():.6f}")
-    print(f"LeRobot (own preprocessing) Actions std: {lerobot_actions_own.std().item():.6f}")
+        lerobot_actions = lerobot_pi05.predict_action_chunk(lerobot_batch, noise=noise, num_steps=10)
+        openpi_actions = original_pi05.sample_actions(
+            device=DEVICE,
+            observation=openpi_observation,
+            noise=noise,
+            num_steps=10,
+        )
 
-    print("\nComparing end-to-end implementations:")
-    print(f"Actions close (atol=1e-4): {torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-4)}")
-    print(f"Actions close (atol=1e-2): {torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-2)}")
-    print(f"Max absolute difference: {torch.abs(lerobot_actions_own - openpi_actions).max().item():.6f}")
+    torch.testing.assert_close(lerobot_actions, openpi_actions, rtol=SAMPLE_RTOL, atol=SAMPLE_ATOL)
 
-    assert torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-4)
-    assert torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-2)
-    assert torch.abs(lerobot_actions_own - openpi_actions).max().item() < 1e-4
+
+def test_pi05_forward_matches_openpi():
+    assert_forward_matches()
+
+
+def test_pi05_sample_actions_match_openpi():
+    assert_sample_actions_match_openpi()
+
+
+def test_pi05_gradient_checkpointing_forward_matches_openpi():
+    assert_forward_matches(gradient_checkpointing=True)
+
+
+def test_pi05_compile_forward_matches_openpi():
+    assert_forward_matches(compile_model=True)
+
+
+def test_pi05_compile_sample_actions_match_openpi():
+    assert_sample_actions_match_openpi(compile_model=True)
+
+
+def test_pi05_compile_gradient_checkpointing_forward_matches_openpi():
+    assert_forward_matches(compile_model=True, gradient_checkpointing=True)

tests/policies/pi0_pi05/test_pi0_compile.py

@@ -0,0 +1,99 @@
+#!/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 os
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi0 import PI0Config  # noqa: E402
+from lerobot.policies.pi0.modeling_pi0 import PI0Pytorch  # noqa: E402
+from tests.policies.pi0_pi05.utils.torch_compile import (  # noqa: E402
+    assert_cache_stability,
+    assert_compiled_output_matches_eager,
+    assert_explain_has_no_graph_breaks,
+    benchmark_runtime,
+    make_compile_config,
+    reset_compile_state,
+)
+from tests.utils import require_cuda  # noqa: E402
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="torch.compile benchmark is too slow for CI; run manually on GPU nodes",
+)
+
+
+def _make_model(*, compile_model):
+    return PI0Pytorch(make_compile_config(PI0Config, compile_model=compile_model)).cuda().eval()
+
+
+def _make_dummy_inputs(config):
+    device = torch.device("cuda")
+    common = {
+        "images": [torch.randn(1, 3, *config.image_resolution, device=device)],
+        "img_masks": [torch.ones(1, dtype=torch.bool, device=device)],
+        "lang_tokens": torch.randint(0, 1024, (1, 5), dtype=torch.long, device=device),
+        "lang_masks": torch.ones(1, 5, dtype=torch.bool, device=device),
+        "state": torch.randn(1, config.max_state_dim, device=device),
+    }
+    forward_kwargs = {
+        **common,
+        "actions": torch.randn(1, config.chunk_size, config.max_action_dim, device=device),
+        "noise": torch.randn(1, config.chunk_size, config.max_action_dim, device=device),
+        "time": torch.rand(1, device=device),
+    }
+    sample_kwargs = {
+        **common,
+        "noise": torch.randn(1, config.chunk_size, config.max_action_dim, device=device),
+        "num_steps": config.num_inference_steps,
+    }
+    return forward_kwargs, sample_kwargs
+
+
+@require_cuda
+def test_pi0_torch_compile_forward_and_sample_actions():
+    if not hasattr(torch, "compile"):
+        pytest.skip("torch.compile is not available")
+    if not torch._dynamo.is_dynamo_supported():
+        pytest.skip("torch._dynamo is not supported on this platform")
+
+    torch.manual_seed(0)
+    eager_model = _make_model(compile_model=False)
+    torch.manual_seed(0)
+    compiled_model = _make_model(compile_model=True)
+    forward_kwargs, sample_kwargs = _make_dummy_inputs(compiled_model.config)
+
+    try:
+        assert_compiled_output_matches_eager(eager_model, compiled_model, forward_kwargs, sample_kwargs)
+
+        assert_explain_has_no_graph_breaks(eager_model.forward, forward_kwargs, "pi0.forward")
+        assert_explain_has_no_graph_breaks(eager_model.sample_actions, sample_kwargs, "pi0.sample_actions")
+
+        assert_cache_stability(compiled_model.forward, forward_kwargs, "pi0.forward")
+        assert_cache_stability(compiled_model.sample_actions, sample_kwargs, "pi0.sample_actions")
+
+        benchmark_runtime(eager_model.forward, compiled_model.forward, forward_kwargs, "pi0.forward")
+        benchmark_runtime(
+            eager_model.sample_actions, compiled_model.sample_actions, sample_kwargs, "pi0.sample_actions"
+        )
+    finally:
+        reset_compile_state()
+        del eager_model
+        del compiled_model
+        torch.cuda.empty_cache()

tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py

@@ -14,51 +14,56 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Test script to verify PI0 policy integration with LeRobot vs the original implementation"""
+"""Compare LeRobot PI0 against the vendored OpenPI PyTorch reference."""
 
+import gc
 import os
-from copy import deepcopy
-from typing import Any
 
 import pytest
 import torch
 
-# Skip if openpi or transformers is not available
-pytest.importorskip("openpi")
 pytest.importorskip("transformers")
 
-# Skip this entire module in CI
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local OpenPI installation and is not meant for CI",
+from lerobot.configs import PreTrainedConfig  # noqa: E402
+from lerobot.policies.pi0 import PI0Policy  # noqa: E402
+from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+from tests.policies.pi0_pi05.openpi_pytorch.pi0_pytorch import PI0Pytorch  # noqa: E402
+from tests.policies.pi0_pi05.utils.openpi_parity import (  # noqa: E402
+    assert_processor_inputs_match_lerobot,
+    clone_batch,
+    deterministic_openpi_forward_preprocess,
+    fix_reference_state_dict,
+    fixed_flow_sampling,
+    load_openpi_reference_state_dict,
+    make_openpi_observation_from_raw,
+    openpi_model_actions_from_raw,
 )
 
-from openpi.models_pytorch import preprocessing_pytorch as openpi_preprocessing  # noqa: E402
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="OpenPI parity and torch.compile checks are too slow for CI; run manually on GPU nodes",
+)
 
-# NOTE: Assumes PYTHONPATH is set to include OpenPI src as per instructions.
-from openpi.models_pytorch.pi0_pytorch import PI0Pytorch  # noqa: E402
-from transformers import AutoTokenizer  # noqa: E402
-
-from lerobot.policies.pi0 import PI0Config, PI0Policy  # noqa: E402
-from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors  # noqa: E402
-from lerobot.processor import PolicyProcessorPipeline  # noqa: E402
-from lerobot.types import PolicyAction  # noqa: E402
-
-# TODO: ADDING DEFAULT IMAGES_FEATURES TO CONFIG
 DUMMY_ACTION_DIM = 32
 DUMMY_STATE_DIM = 32
 DUMMY_ACTION_HORIZON = 50
-DUMMY_MAX_TOKEN_LEN = 48  # Default for PI0 (non-pi05)
-DEVICE = "cpu"  # Use CPU to avoid memory issues for testing
+DUMMY_MAX_TOKEN_LEN = 48
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+COMPILE_MODE = "default"
+FORWARD_RTOL = 1e-4
+FORWARD_ATOL = 1e-4
+SAMPLE_RTOL = 1e-2
+SAMPLE_ATOL = 5e-3
 
 DUMMY_DATASET_STATS = {
-    "observation.state": {
+    OBS_STATE: {
         "mean": torch.zeros(DUMMY_STATE_DIM),
         "std": torch.ones(DUMMY_STATE_DIM),
         "q01": torch.zeros(DUMMY_STATE_DIM),
         "q99": torch.ones(DUMMY_STATE_DIM),
     },
-    "action": {
+    ACTION: {
         "mean": torch.zeros(DUMMY_ACTION_DIM),
         "std": torch.ones(DUMMY_ACTION_DIM),
         "q01": torch.zeros(DUMMY_ACTION_DIM),
@@ -87,6 +92,15 @@ DUMMY_DATASET_STATS = {
 }
 
 
+@pytest.fixture(autouse=True)
+def cleanup_cuda_after_test():
+    yield
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
 class PI0BaseOriginalConfig:
     action_dim: int = DUMMY_ACTION_DIM
     action_horizon: int = DUMMY_ACTION_HORIZON
@@ -95,333 +109,156 @@ class PI0BaseOriginalConfig:
     precision: str = "float32"
     pi05: bool = False
     dtype: str = "float32"
+    pytorch_compile_mode: str | None = None
 
 
-def instantiate_lerobot_pi0(
-    from_pretrained: bool = False,
-) -> tuple[
-    PI0Policy,
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    if from_pretrained:
-        # Load the policy first
-        policy = PI0Policy.from_pretrained(pretrained_name_or_path="lerobot/pi0_base", strict=True)
-    else:
-        config = PI0Config(max_action_dim=DUMMY_ACTION_DIM, max_state_dim=DUMMY_STATE_DIM, dtype="float32")
-        policy = PI0Policy(config)
+def instantiate_lerobot_pi0(*, compile_model: bool = False, gradient_checkpointing: bool = False):
+    config = PreTrainedConfig.from_pretrained("lerobot/pi0_base")
+    config.device = str(DEVICE)
+    config.dtype = "float32"
+    config.compile_model = compile_model
+    config.compile_mode = COMPILE_MODE
+    config.gradient_checkpointing = gradient_checkpointing
 
+    policy = PI0Policy.from_pretrained("lerobot/pi0_base", config=config, strict=True)
     policy.to(DEVICE)
-    policy.config.device = DEVICE
-    preprocessor, postprocessor = make_pi0_pre_post_processors(
-        config=policy.config, dataset_stats=DUMMY_DATASET_STATS
-    )
-    return (policy, preprocessor, postprocessor)
+    policy.config.device = str(DEVICE)
+    preprocessor, _ = make_pi0_pre_post_processors(config=policy.config, dataset_stats=DUMMY_DATASET_STATS)
+    return policy, preprocessor
 
 
-def instantiate_original_pi0(from_pretrained: bool = False, model_path: str = None):
-    config = PI0BaseOriginalConfig()
-    policy = PI0Pytorch(config)
-
-    if from_pretrained:
-        try:
-            print("Loading converted PyTorch weights from HuggingFace Hub (lerobot/pi0_base)...")
-
-            # Download the model from HuggingFace Hub
-            import safetensors.torch
-            from huggingface_hub import snapshot_download
-
-            # Download the entire repository
-            if model_path and os.path.exists(model_path):
-                cache_dir = model_path
-                print(f"Using cached model from: {cache_dir}")
-            else:
-                cache_dir = snapshot_download(repo_id="lerobot/pi0_base", repo_type="model")
-                print(f"Downloaded model to: {cache_dir}")
-
-            # Try to load safetensors format first
-            model_file = os.path.join(cache_dir, "model.safetensors")
-            if os.path.exists(model_file):
-                state_dict = safetensors.torch.load_file(model_file)
-                print(f"Loaded {len(state_dict)} parameters from safetensors")
-            else:
-                raise FileNotFoundError(f"No safetensors file found in {cache_dir}")
-
-            # Load the state dict into the model
-            missing_keys, unexpected_keys = policy.load_state_dict(state_dict, strict=False)
-
-            if missing_keys:
-                print(f"Missing keys: {len(missing_keys)}")
-                if len(missing_keys) <= 5:
-                    for key in missing_keys:
-                        print(f"    - {key}")
-                else:
-                    for key in missing_keys[:5]:
-                        print(f"    - {key}")
-                    print(f"    ... and {len(missing_keys) - 5} more")
-
-            if unexpected_keys:
-                print(f"Unexpected keys: {len(unexpected_keys)}")
-                if len(unexpected_keys) <= 5:
-                    for key in unexpected_keys:
-                        print(f"    - {key}")
-                else:
-                    for key in unexpected_keys[:5]:
-                        print(f"    - {key}")
-                    print(f"    ... and {len(unexpected_keys) - 5} more")
-
-            if not missing_keys and not unexpected_keys:
-                print("All pretrained weights loaded successfully!")
-            else:
-                print("Pretrained weights loaded with some missing/unexpected keys (this may be normal)")
-
-        except Exception as e:
-            print(f"Failed to load pretrained weights: {e}")
-            print("   Using randomly initialized weights...")
-            import traceback
-
-            traceback.print_exc()
-
-    policy.to(DEVICE)
+def instantiate_original_pi0():
+    policy = PI0Pytorch(PI0BaseOriginalConfig()).to(DEVICE)
+    state_dict = fix_reference_state_dict(load_openpi_reference_state_dict("lerobot/pi0_base"))
+    missing_keys, unexpected_keys = policy.load_state_dict(state_dict, strict=False)
+    assert missing_keys == []
+    assert unexpected_keys == []
     return policy
 
 
 def create_dummy_data():
-    batch_size = 2  # Reduce batch size for testing
-    device = DEVICE
-
-    # Use the exact same prompt for both implementations
+    batch_size = 2
     prompt = "Pick up the red block and place it in the bin"
-
-    batch = {
-        "observation.state": torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device),
-        "action": torch.randn(
-            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=device
+    return {
+        OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=DEVICE),
+        ACTION: torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=DEVICE
         ),
-        # Create images in [0, 1] range as expected by LeRobot (will be converted to [-1, 1] internally)
         "observation.images.base_0_rgb": torch.rand(
-            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+            batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
         ),
         "observation.images.left_wrist_0_rgb": torch.rand(
-            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+            batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
         ),
         "observation.images.right_wrist_0_rgb": torch.rand(
-            batch_size, 3, 224, 224, dtype=torch.float32, device=device
+            batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
         ),
-        # Add the task prompt for LeRobot - provide as list with single element to trigger expansion
         "task": [prompt for _ in range(batch_size)],
     }
-    return batch
 
 
-def extract_lerobot_processed_inputs(lerobot_pi0, batch):
-    """Extract the exact same processed inputs that LeRobot uses internally."""
-    # Get the tokenized language from LeRobot's internal method
-    lang_tokens, lang_masks = lerobot_pi0._tokenize_language(batch)
-
-    # Get the preprocessed images from LeRobot's internal method
-    images, img_masks = lerobot_pi0._preprocess_images(batch, train=False)
-
-    # Create dummy token_ar_mask and token_loss_mask for original implementation
-    token_ar_mask = torch.zeros_like(lang_tokens, dtype=torch.int32)
-    token_loss_mask = torch.ones_like(lang_masks, dtype=torch.bool)
-
-    return images, img_masks, lang_tokens, lang_masks, token_ar_mask, token_loss_mask
+def prepare_parity_inputs(lerobot_pi0, lerobot_preprocessor):
+    torch.manual_seed(0)
+    raw_batch = create_dummy_data()
+    lerobot_batch = lerobot_preprocessor(clone_batch(raw_batch))
+    openpi_observation = make_openpi_observation_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        max_token_len=DUMMY_MAX_TOKEN_LEN,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=False,
+    )
+    openpi_actions = openpi_model_actions_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=False,
+    )
+    assert_processor_inputs_match_lerobot(
+        lerobot_pi0,
+        lerobot_batch,
+        openpi_observation,
+        compare_state=True,
+    )
+    batch_size = raw_batch[OBS_STATE].shape[0]
+    noise = torch.randn(
+        batch_size,
+        DUMMY_ACTION_HORIZON,
+        DUMMY_ACTION_DIM,
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    time = torch.linspace(0.2, 0.8, batch_size, dtype=torch.float32, device=DEVICE)
+    return lerobot_batch, openpi_observation, openpi_actions, noise, time
 
 
-class PI0Observation:
-    """Observation class that matches the original OpenPI format."""
+def assert_forward_matches(*, compile_model: bool = False, gradient_checkpointing: bool = False):
+    lerobot_pi0, lerobot_preprocessor = instantiate_lerobot_pi0(
+        compile_model=compile_model,
+        gradient_checkpointing=gradient_checkpointing,
+    )
+    original_pi0 = instantiate_original_pi0()
+    lerobot_batch, openpi_observation, openpi_actions, noise, time = prepare_parity_inputs(
+        lerobot_pi0,
+        lerobot_preprocessor,
+    )
 
-    def __init__(
-        self,
-        state,
-        images,
-        image_masks,
-        tokenized_prompt,
-        tokenized_prompt_mask,
-        token_ar_mask,
-        token_loss_mask,
-    ):
-        self.state = state
-        self.images = images
-        self.image_masks = image_masks
-        self.tokenized_prompt = tokenized_prompt
-        self.tokenized_prompt_mask = tokenized_prompt_mask
-        self.token_ar_mask = token_ar_mask
-        self.token_loss_mask = token_loss_mask
-
-
-def create_original_observation_with_openpi_preprocessing(batch):
-    """Create observation object for OpenPI using OpenPI's own preprocessing."""
-    batch_size = batch["observation.state"].shape[0]
-    device = batch["observation.state"].device
-
-    # Create tokenizer for OpenPI (same as LeRobot uses)
-    tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
-
-    # Get task description
-    if "task" in batch:
-        tasks = batch["task"]
-        if isinstance(tasks, str):
-            # Single string: add newline if not present, then convert to list
-            if not tasks.endswith("\n"):
-                tasks = f"{tasks}\n"
-            tasks = [tasks]
-        elif isinstance(tasks, list) and all(isinstance(t, str) for t in tasks):
-            # List of strings: add newline to each if not present
-            tasks = [t if t.endswith("\n") else f"{t}\n" for t in tasks]
-            if len(tasks) == 1:
-                # Expand to batch size
-                tasks = tasks * batch_size
-                if len(tasks) != batch_size:
-                    raise ValueError(f"Expected batch size {batch_size}, got {len(tasks)}")
-        # If task is neither string nor list of strings, leave unchanged
+    if gradient_checkpointing:
+        lerobot_pi0.train()
     else:
-        # Default task if not provided
-        tasks = ["Pick up the object\n"] * batch_size
-
-    # Tokenize with max_length padding to match OpenPI's expected format
-    tokenized = tokenizer(
-        tasks,
-        padding="max_length",
-        padding_side="right",
-        truncation=True,
-        max_length=DUMMY_MAX_TOKEN_LEN,
-        return_tensors="pt",
-    )
-
-    lang_tokens = tokenized["input_ids"].to(device)
-    lang_masks = tokenized["attention_mask"].to(device, dtype=torch.bool)
-
-    # Create dummy token_ar_mask and token_loss_mask for OpenPI
-    token_ar_mask = torch.zeros_like(lang_tokens, dtype=torch.int32)
-    token_loss_mask = torch.ones_like(lang_masks, dtype=torch.bool)
-
-    # Convert LeRobot images format to OpenPI format (convert [0,1] to [-1,1] range)
-    image_dict = {
-        "base_0_rgb": batch["observation.images.base_0_rgb"] * 2.0 - 1.0,
-        "left_wrist_0_rgb": batch["observation.images.left_wrist_0_rgb"] * 2.0 - 1.0,
-        "right_wrist_0_rgb": batch["observation.images.right_wrist_0_rgb"] * 2.0 - 1.0,
-    }
-
-    # Create image masks (all ones for real images)
-    image_masks_dict = {}
-    for key in image_dict:
-        image_masks_dict[key] = torch.ones(batch_size, dtype=torch.bool, device=device)
-
-    # Create raw observation object (before preprocessing)
-    raw_observation = PI0Observation(
-        state=batch["observation.state"],
-        images=image_dict,
-        image_masks=image_masks_dict,
-        tokenized_prompt=lang_tokens,
-        tokenized_prompt_mask=lang_masks,
-        token_ar_mask=token_ar_mask,
-        token_loss_mask=token_loss_mask,
-    )
-
-    # Now use OpenPI's preprocessing
-    processed_obs = openpi_preprocessing.preprocess_observation_pytorch(raw_observation, train=False)
-
-    return processed_obs
-
-
-def create_original_observation_from_lerobot(lerobot_pi0, batch):
-    """Create observation object compatible with original OpenPI using the exact same inputs as LeRobot."""
-    _batch_size = batch["observation.state"].shape[0]
-    _device = batch["observation.state"].device
-
-    # Extract the exact same processed inputs that LeRobot uses
-    images, img_masks, lang_tokens, lang_masks, token_ar_mask, token_loss_mask = (
-        extract_lerobot_processed_inputs(lerobot_pi0, batch)
-    )
-
-    # Convert images list to dict with original OpenPI keys
-    image_dict = {
-        "base_0_rgb": images[0],
-        "left_wrist_0_rgb": images[1],
-        "right_wrist_0_rgb": images[2],
-    }
-
-    # Convert image masks list to dict with original OpenPI keys
-    image_masks_dict = {
-        "base_0_rgb": img_masks[0],
-        "left_wrist_0_rgb": img_masks[1],
-        "right_wrist_0_rgb": img_masks[2],
-    }
-
-    return PI0Observation(
-        state=batch["observation.state"],
-        images=image_dict,
-        image_masks=image_masks_dict,
-        tokenized_prompt=lang_tokens,
-        tokenized_prompt_mask=lang_masks,
-        token_ar_mask=token_ar_mask,
-        token_loss_mask=token_loss_mask,
-    )
-
-
-def test_pi0_original_vs_lerobot():
-    """Test PI0 original implementation vs LeRobot implementation."""
-    print("Initializing models...")
-    lerobot_pi0, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_pi0(
-        from_pretrained=True
-    )  # Load pretrained LeRobot model
-    original_pi0 = instantiate_original_pi0(
-        from_pretrained=True
-    )  # Load pretrained OpenPI model from HuggingFace Hub
-
-    print("Creating dummy data...")
-    batch = create_dummy_data()
-    batch_lerobot = deepcopy(batch)
-
-    # Test each model with its own preprocessing (more realistic end-to-end test)
-    print("\nTest each model with its own preprocessing")
-    print("Creating observation for OpenPI using OpenPI's own preprocessing...")
-    pi0_obs_openpi = create_original_observation_with_openpi_preprocessing(batch)
-
-    print(f"Task prompt: '{batch['task'][0]}'")
-    print(f"OpenPI tokenized prompt shape: {pi0_obs_openpi.tokenized_prompt.shape}")
-    print(f"OpenPI image shapes: {[img.shape for img in pi0_obs_openpi.images.values()]}")
-    print(f"OpenPI state shape: {pi0_obs_openpi.state.shape}")
-
-    print("Testing OpenPI with own preprocessing...")
+        lerobot_pi0.eval()
     original_pi0.eval()
-    torch.manual_seed(42)  # Set seed for reproducibility
-    batch_size = batch["observation.state"].shape[0]
-    noise_shape = (batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM)
-    fixed_noise = torch.randn(noise_shape, dtype=torch.float32, device=DEVICE)
 
-    with torch.no_grad():
-        openpi_actions = original_pi0.sample_actions(
-            device=DEVICE, observation=pi0_obs_openpi, noise=fixed_noise, num_steps=10
-        )
-        openpi_actions_unit = openpi_actions[:, 0, :]
-    print(f"OpenPI (own preprocessing) Actions shape: {openpi_actions.shape}")
-    print(f"OpenPI (own preprocessing) Actions unit shape: {openpi_actions_unit.shape}")
-    print(f"OpenPI (own preprocessing) Actions mean: {openpi_actions.mean().item():.6f}")
-    print(f"OpenPI (own preprocessing) Actions std: {openpi_actions.std().item():.6f}")
+    with fixed_flow_sampling(lerobot_pi0.model, noise=noise, time=time):
+        lerobot_loss, _ = lerobot_pi0(lerobot_batch, reduction="none")
+    with deterministic_openpi_forward_preprocess(original_pi0):
+        openpi_losses = original_pi0(openpi_observation, openpi_actions, noise=noise, time=time)
+    openpi_loss = openpi_losses.mean(dim=(1, 2))
+
+    torch.testing.assert_close(lerobot_loss, openpi_loss, rtol=FORWARD_RTOL, atol=FORWARD_ATOL)
+
+
+def assert_sample_actions_match_openpi(*, compile_model: bool = False):
+    lerobot_pi0, lerobot_preprocessor = instantiate_lerobot_pi0(compile_model=compile_model)
+    original_pi0 = instantiate_original_pi0()
+    lerobot_batch, openpi_observation, _openpi_actions, noise, _time = prepare_parity_inputs(
+        lerobot_pi0,
+        lerobot_preprocessor,
+    )
 
-    print("Testing LeRobot with own preprocessing...")
     lerobot_pi0.eval()
-    torch.manual_seed(42)  # Set the same seed
-
-    batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
+    original_pi0.eval()
     with torch.no_grad():
-        lerobot_actions_own = lerobot_pi0.predict_action_chunk(
-            batch_lerobot_processed
-        )  # batch_size, n_action_steps, action_dim
-        lerobot_actions_unit = lerobot_actions_own[:, 0, :]
-    print(f"LeRobot (own preprocessing) Actions shape: {lerobot_actions_own.shape}")
-    print(f"LeRobot (own preprocessing) Actions unit shape: {lerobot_actions_unit.shape}")
-    print(f"LeRobot (own preprocessing) Actions mean: {lerobot_actions_own.mean().item():.6f}")
-    print(f"LeRobot (own preprocessing) Actions std: {lerobot_actions_own.std().item():.6f}")
+        lerobot_actions = lerobot_pi0.predict_action_chunk(lerobot_batch, noise=noise, num_steps=10)
+        openpi_actions = original_pi0.sample_actions(
+            device=DEVICE,
+            observation=openpi_observation,
+            noise=noise,
+            num_steps=10,
+        )
 
-    print("\nComparing end-to-end implementations:")
-    print(f"Actions close (atol=1e-4): {torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-4)}")
-    print(f"Actions close (atol=1e-2): {torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-2)}")
-    print(f"Max absolute difference: {torch.abs(lerobot_actions_own - openpi_actions).max().item():.6f}")
+    torch.testing.assert_close(lerobot_actions, openpi_actions, rtol=SAMPLE_RTOL, atol=SAMPLE_ATOL)
 
-    assert torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-4)
-    assert torch.allclose(lerobot_actions_own, openpi_actions, atol=1e-2)
-    assert torch.abs(lerobot_actions_own - openpi_actions).max().item() < 1e-4
+
+def test_pi0_forward_matches_openpi():
+    assert_forward_matches()
+
+
+def test_pi0_sample_actions_match_openpi():
+    assert_sample_actions_match_openpi()
+
+
+def test_pi0_gradient_checkpointing_forward_matches_openpi():
+    assert_forward_matches(gradient_checkpointing=True)
+
+
+def test_pi0_compile_forward_matches_openpi():
+    assert_forward_matches(compile_model=True)
+
+
+def test_pi0_compile_sample_actions_match_openpi():
+    assert_sample_actions_match_openpi(compile_model=True)
+
+
+def test_pi0_compile_gradient_checkpointing_forward_matches_openpi():
+    assert_forward_matches(compile_model=True, gradient_checkpointing=True)

tests/policies/pi0_pi05/utils/__init__.py

@@ -0,0 +1 @@
+"""Utilities shared by PI0/PI05 policy tests."""

tests/policies/pi0_pi05/utils/openpi_parity.py

@@ -0,0 +1,291 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+
+import numpy as np
+import safetensors.torch
+import torch
+import torch.nn.functional as F  # noqa: N812
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer
+
+from lerobot.utils.constants import (
+    ACTION,
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+    OBS_STATE,
+)
+from tests.policies.pi0_pi05.openpi_pytorch import preprocessing_pytorch as openpi_preprocessing
+
+IMAGE_KEYS = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+TOKENIZER_NAME = "google/paligemma-3b-pt-224"
+
+
+@dataclass
+class OpenPIObservation:
+    state: torch.Tensor
+    images: dict[str, torch.Tensor]
+    image_masks: dict[str, torch.Tensor]
+    tokenized_prompt: torch.Tensor
+    tokenized_prompt_mask: torch.Tensor
+    token_ar_mask: torch.Tensor
+    token_loss_mask: torch.Tensor
+
+
+@lru_cache(maxsize=1)
+def paligemma_tokenizer():
+    return AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+
+
+def clone_batch(batch: dict) -> dict:
+    return {
+        key: value.clone() if isinstance(value, torch.Tensor) else list(value) for key, value in batch.items()
+    }
+
+
+def pad_last_dim(tensor: torch.Tensor, target_dim: int) -> torch.Tensor:
+    if tensor.shape[-1] > target_dim:
+        raise ValueError(f"Cannot pad last dimension {tensor.shape[-1]} down to {target_dim}")
+    return F.pad(tensor, (0, target_dim - tensor.shape[-1]))
+
+
+def mean_std_normalize(tensor: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    mean = stats["mean"].to(device=tensor.device, dtype=tensor.dtype)
+    std = stats["std"].to(device=tensor.device, dtype=tensor.dtype)
+    return (tensor - mean) / (std + 1e-8)
+
+
+def quantile_normalize(tensor: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    q01 = stats["q01"].to(device=tensor.device, dtype=tensor.dtype)
+    q99 = stats["q99"].to(device=tensor.device, dtype=tensor.dtype)
+    denom = torch.where(q99 == q01, torch.full_like(q99, 1e-8), q99 - q01)
+    return 2.0 * (tensor - q01) / denom - 1.0
+
+
+def openpi_model_state_from_raw(
+    batch: dict[str, torch.Tensor],
+    *,
+    action_dim: int,
+    dataset_stats: dict[str, dict[str, torch.Tensor]],
+    pi05: bool,
+) -> torch.Tensor:
+    state = batch[OBS_STATE].to(dtype=torch.float32)
+    if pi05:
+        state = quantile_normalize(state, dataset_stats[OBS_STATE])
+    else:
+        state = mean_std_normalize(state, dataset_stats[OBS_STATE])
+    return pad_last_dim(state, action_dim)
+
+
+def openpi_model_actions_from_raw(
+    batch: dict[str, torch.Tensor],
+    *,
+    action_dim: int,
+    dataset_stats: dict[str, dict[str, torch.Tensor]],
+    pi05: bool,
+) -> torch.Tensor:
+    actions = batch[ACTION].to(dtype=torch.float32)
+    if pi05:
+        actions = quantile_normalize(actions, dataset_stats[ACTION])
+    else:
+        actions = mean_std_normalize(actions, dataset_stats[ACTION])
+    return pad_last_dim(actions, action_dim)
+
+
+def _tasks_from_raw(batch: dict, batch_size: int) -> list[str]:
+    tasks = batch.get("task")
+    if tasks is None:
+        raise ValueError("The parity batch must include a task prompt.")
+    if isinstance(tasks, str):
+        return [tasks] * batch_size
+    if len(tasks) == 1:
+        return [tasks[0]] * batch_size
+    if len(tasks) != batch_size:
+        raise ValueError(f"Expected {batch_size} task prompts, got {len(tasks)}")
+    return list(tasks)
+
+
+def _format_pi0_prompts(tasks: list[str]) -> list[str]:
+    return [f"{task.strip().replace('_', ' ').replace(chr(10), ' ')}\n" for task in tasks]
+
+
+def _format_pi05_prompts(tasks: list[str], normalized_state: torch.Tensor) -> list[str]:
+    state_np = normalized_state.detach().cpu().numpy()
+    discretized_states = np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
+    prompts = []
+    for task, state in zip(tasks, discretized_states, strict=True):
+        cleaned_text = task.strip().replace("_", " ").replace("\n", " ")
+        state_str = " ".join(map(str, state))
+        prompts.append(f"Task: {cleaned_text}, State: {state_str};\nAction: ")
+    return prompts
+
+
+def _tokenize_prompts(prompts: list[str], *, max_token_len: int, device: torch.device | str):
+    tokenized = paligemma_tokenizer()(
+        prompts,
+        padding="max_length",
+        padding_side="right",
+        truncation=True,
+        max_length=max_token_len,
+        return_tensors="pt",
+    )
+    tokens = tokenized["input_ids"].to(device)
+    masks = tokenized["attention_mask"].to(device=device, dtype=torch.bool)
+    return tokens, masks
+
+
+def make_openpi_observation_from_raw(
+    batch: dict[str, torch.Tensor],
+    *,
+    action_dim: int,
+    max_token_len: int,
+    dataset_stats: dict[str, dict[str, torch.Tensor]],
+    pi05: bool,
+) -> OpenPIObservation:
+    batch_size = batch[OBS_STATE].shape[0]
+    device = batch[OBS_STATE].device
+    state = openpi_model_state_from_raw(
+        batch,
+        action_dim=action_dim,
+        dataset_stats=dataset_stats,
+        pi05=pi05,
+    )
+
+    tasks = _tasks_from_raw(batch, batch_size)
+    prompts = _format_pi05_prompts(tasks, state) if pi05 else _format_pi0_prompts(tasks)
+    tokens, masks = _tokenize_prompts(prompts, max_token_len=max_token_len, device=device)
+
+    images = {
+        key: batch[f"observation.images.{key}"].to(device=device, dtype=torch.float32) * 2.0 - 1.0
+        for key in IMAGE_KEYS
+    }
+    image_masks = {key: torch.ones(batch_size, dtype=torch.bool, device=device) for key in IMAGE_KEYS}
+
+    return OpenPIObservation(
+        state=state,
+        images=images,
+        image_masks=image_masks,
+        tokenized_prompt=tokens,
+        tokenized_prompt_mask=masks,
+        token_ar_mask=torch.zeros_like(tokens, dtype=torch.int32),
+        token_loss_mask=torch.ones_like(masks, dtype=torch.bool),
+    )
+
+
+def assert_processor_inputs_match_lerobot(
+    lerobot_policy,
+    lerobot_batch: dict[str, torch.Tensor],
+    openpi_observation: OpenPIObservation,
+    *,
+    compare_state: bool,
+):
+    openpi_processed = openpi_preprocessing.preprocess_observation_pytorch(openpi_observation, train=False)
+    lerobot_images, lerobot_image_masks = lerobot_policy._preprocess_images(lerobot_batch)
+
+    # Token IDs, token masks, images, image masks, and PI0 state are intentionally built from the same
+    # raw batch through independent LeRobot/OpenPI-style processor logic. They must be bitwise equal.
+    torch.testing.assert_close(
+        openpi_observation.tokenized_prompt, lerobot_batch[OBS_LANGUAGE_TOKENS], rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        openpi_observation.tokenized_prompt_mask,
+        lerobot_batch[OBS_LANGUAGE_ATTENTION_MASK],
+        rtol=0,
+        atol=0,
+    )
+
+    for openpi_image, lerobot_image in zip(openpi_processed.images.values(), lerobot_images, strict=True):
+        torch.testing.assert_close(openpi_image, lerobot_image, rtol=0, atol=0)
+
+    for openpi_mask, lerobot_mask in zip(
+        openpi_processed.image_masks.values(), lerobot_image_masks, strict=True
+    ):
+        torch.testing.assert_close(openpi_mask, lerobot_mask, rtol=0, atol=0)
+
+    if compare_state:
+        torch.testing.assert_close(
+            openpi_processed.state, lerobot_policy.prepare_state(lerobot_batch), rtol=0, atol=0
+        )
+
+
+def load_openpi_reference_state_dict(repo_id: str) -> dict[str, torch.Tensor]:
+    cache_dir = Path(snapshot_download(repo_id=repo_id, repo_type="model"))
+    return safetensors.torch.load_file(cache_dir / "model.safetensors")
+
+
+def fix_reference_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    fixed_state_dict = dict(state_dict)
+    lm_head_key = "paligemma_with_expert.paligemma.lm_head.weight"
+    embed_tokens_key = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
+    if lm_head_key in fixed_state_dict and embed_tokens_key not in fixed_state_dict:
+        fixed_state_dict[embed_tokens_key] = fixed_state_dict[lm_head_key].clone()
+    return fixed_state_dict
+
+
+@contextmanager
+def fixed_flow_sampling(model, *, noise: torch.Tensor, time: torch.Tensor) -> Iterator[None]:
+    original_sample_noise = model.sample_noise
+    original_sample_time = model.sample_time
+
+    def sample_noise(shape, device):
+        if tuple(shape) != tuple(noise.shape):
+            raise ValueError(f"Expected noise shape {tuple(noise.shape)}, got {tuple(shape)}")
+        return noise.to(device=device)
+
+    def sample_time(batch_size, device):
+        if batch_size != time.shape[0]:
+            raise ValueError(f"Expected time batch size {time.shape[0]}, got {batch_size}")
+        return time.to(device=device)
+
+    model.sample_noise = sample_noise
+    model.sample_time = sample_time
+    try:
+        yield
+    finally:
+        model.sample_noise = original_sample_noise
+        model.sample_time = original_sample_time
+
+
+@contextmanager
+def deterministic_openpi_forward_preprocess(openpi_policy) -> Iterator[None]:
+    """Disable OpenPI's training-time image augmentation only inside a parity forward block.
+
+    OpenPI's `forward()` calls `_preprocess_observation(..., train=True)`, which can apply stochastic
+    image augmentation. LeRobot's policy forward path does not apply that augmentation, so parity would
+    otherwise compare two different image tensors rather than two model implementations. The context manager
+    keeps the public `openpi_policy.forward(observation, ...)` call while making preprocessing deterministic.
+
+    `yield` marks the body of the caller's `with` block. The `try/finally` restores the original method even
+    if the assertion inside the block fails, so the temporary monkeypatch cannot leak into later tests.
+    """
+
+    original_preprocess_observation = openpi_policy._preprocess_observation
+
+    def preprocess_observation(observation, *, train=True):
+        return original_preprocess_observation(observation, train=False)
+
+    openpi_policy._preprocess_observation = preprocess_observation
+    try:
+        yield
+    finally:
+        openpi_policy._preprocess_observation = original_preprocess_observation

tests/policies/pi0_pi05/utils/torch_compile.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from collections.abc import Callable
+
+import torch
+from torch._dynamo.utils import counters, guard_failures
+from torch.profiler import ProfilerActivity
+
+FORWARD_RTOL = 1e-5
+FORWARD_ATOL = 5e-2
+SAMPLE_RTOL = 1e-5
+SAMPLE_ATOL = 1e-2
+COMPILE_MODE = "max-autotune"
+STEADY_STATE_WARMUPS = 3
+STEADY_STATE_REPEATS = 3
+
+
+def make_compile_config(config_cls, *, compile_model):
+    return config_cls(device="cuda", compile_model=compile_model, compile_mode=COMPILE_MODE)
+
+
+def counter_total(name):
+    return sum(counters.get(name, {}).values())
+
+
+def compile_snapshot():
+    return {
+        "graph_breaks": counter_total("graph_break"),
+        "recompiles": counter_total("recompiles"),
+        "recompile_limits": counter_total("recompile_limit"),
+        "unique_graphs": counters["stats"].get("unique_graphs", 0),
+    }
+
+
+def reset_compile_state():
+    torch._dynamo.reset()
+    counters.clear()
+    guard_failures.clear()
+
+
+def clone_cuda_graph_output(output):
+    if torch.is_tensor(output):
+        return output.clone()
+    if isinstance(output, tuple):
+        return tuple(clone_cuda_graph_output(item) for item in output)
+    if isinstance(output, list):
+        return [clone_cuda_graph_output(item) for item in output]
+    if isinstance(output, dict):
+        return {key: clone_cuda_graph_output(value) for key, value in output.items()}
+    return output
+
+
+def run_model_step(fn: Callable, kwargs: dict):
+    if hasattr(torch.compiler, "cudagraph_mark_step_begin"):
+        torch.compiler.cudagraph_mark_step_begin()
+    return fn(**kwargs)
+
+
+def assert_explain_has_no_graph_breaks(fn: Callable, kwargs: dict, label: str):
+    reset_compile_state()
+    explanation = torch._dynamo.explain(fn)(**kwargs)
+
+    assert explanation.graph_count > 0, f"{label} was not captured by Dynamo"
+    assert explanation.graph_break_count == 0, (
+        f"{label} has {explanation.graph_break_count} graph break(s): {explanation.break_reasons}"
+    )
+    assert not explanation.break_reasons, f"{label} graph break reasons: {explanation.break_reasons}"
+
+    print(
+        f"{label} capture: graphs={explanation.graph_count}, "
+        f"graph_breaks={explanation.graph_break_count}, ops={explanation.op_count}, "
+        f"guards={len(explanation.out_guards or [])}"
+    )
+    return explanation
+
+
+@torch.no_grad()
+def assert_compiled_output_matches_eager(eager_model, compiled_model, forward_kwargs, sample_kwargs):
+    eager_forward = eager_model.forward(**forward_kwargs)
+    compiled_forward = compiled_model.forward(**forward_kwargs)
+    torch.testing.assert_close(compiled_forward, eager_forward, rtol=FORWARD_RTOL, atol=FORWARD_ATOL)
+
+    eager_actions = eager_model.sample_actions(**sample_kwargs)
+    compiled_actions = compiled_model.sample_actions(**sample_kwargs)
+    torch.testing.assert_close(compiled_actions, eager_actions, rtol=SAMPLE_RTOL, atol=SAMPLE_ATOL)
+
+
+@torch.no_grad()
+def assert_cache_stability(fn: Callable, kwargs: dict, label: str):
+    reset_compile_state()
+
+    first_output = clone_cuda_graph_output(run_model_step(fn, kwargs))
+    first_snapshot = compile_snapshot()
+    second_output = clone_cuda_graph_output(run_model_step(fn, kwargs))
+    second_snapshot = compile_snapshot()
+    third_output = clone_cuda_graph_output(run_model_step(fn, kwargs))
+    third_snapshot = compile_snapshot()
+
+    torch.testing.assert_close(second_output, first_output, rtol=FORWARD_RTOL, atol=FORWARD_ATOL)
+    torch.testing.assert_close(third_output, first_output, rtol=FORWARD_RTOL, atol=FORWARD_ATOL)
+    assert first_snapshot["unique_graphs"] > 0, f"{label} did not compile any graph"
+    assert third_snapshot["graph_breaks"] == 0, f"{label} graph breaks: {third_snapshot}"
+    assert third_snapshot["recompiles"] == 0, f"{label} recompiled: {third_snapshot}"
+    assert third_snapshot["recompile_limits"] == 0, f"{label} hit recompile limit: {third_snapshot}"
+    assert second_snapshot["unique_graphs"] == first_snapshot["unique_graphs"], (
+        f"{label} compiled new graph on second call: first={first_snapshot}, second={second_snapshot}"
+    )
+    assert third_snapshot["unique_graphs"] == first_snapshot["unique_graphs"], (
+        f"{label} compiled new graph on third call: first={first_snapshot}, third={third_snapshot}"
+    )
+    assert not guard_failures, f"{label} guard failures: {dict(guard_failures)}"
+
+    print(f"{label} cache: first={first_snapshot}, third={third_snapshot}")
+
+
+@torch.no_grad()
+def benchmark_runtime(eager_fn: Callable, compiled_fn: Callable, kwargs: dict, label: str):
+    run_warmups(eager_fn, kwargs)
+    run_warmups(compiled_fn, kwargs)
+    torch.cuda.synchronize()
+
+    eager_metrics = profile_callable(eager_fn, kwargs)
+    compiled_metrics = profile_callable(compiled_fn, kwargs)
+    speedup = eager_metrics["cuda_event_ms"] / compiled_metrics["cuda_event_ms"]
+
+    print(
+        f"{label} runtime: eager_cuda={eager_metrics['cuda_event_ms']:.3f} ms, "
+        f"compiled_cuda={compiled_metrics['cuda_event_ms']:.3f} ms, speedup={speedup:.3f}x, "
+        f"host_wall_ms eager/compiled={eager_metrics['host_wall_ms']:.3f}/"
+        f"{compiled_metrics['host_wall_ms']:.3f}, "
+        f"cpu_self_time_ms eager/compiled={eager_metrics['cpu_self_time_ms']:.3f}/"
+        f"{compiled_metrics['cpu_self_time_ms']:.3f}, "
+        f"cuda_launches eager/compiled={eager_metrics['cuda_launch_count']}/"
+        f"{compiled_metrics['cuda_launch_count']}, "
+        f"profiler_events eager/compiled={eager_metrics['profiler_event_count']}/"
+        f"{compiled_metrics['profiler_event_count']}, "
+        f"peak_mem_mib eager/compiled={eager_metrics['peak_mem_mib']:.1f}/"
+        f"{compiled_metrics['peak_mem_mib']:.1f}"
+    )
+
+    assert eager_metrics["cuda_event_ms"] > 0
+    assert compiled_metrics["cuda_event_ms"] > 0
+    assert eager_metrics["profiler_event_count"] > 0
+    assert compiled_metrics["profiler_event_count"] > 0
+    return eager_metrics, compiled_metrics
+
+
+def run_warmups(fn: Callable, kwargs: dict):
+    for _ in range(STEADY_STATE_WARMUPS):
+        run_model_step(fn, kwargs)
+    torch.cuda.synchronize()
+
+
+def profile_callable(fn: Callable, kwargs: dict):
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+    host_start = time.perf_counter()
+    start_event.record()
+    for _ in range(STEADY_STATE_REPEATS):
+        run_model_step(fn, kwargs)
+    end_event.record()
+    torch.cuda.synchronize()
+    cuda_event_ms = start_event.elapsed_time(end_event) / STEADY_STATE_REPEATS
+    host_wall_ms = (time.perf_counter() - host_start) * 1000 / STEADY_STATE_REPEATS
+    peak_mem_mib = torch.cuda.max_memory_allocated() / 1024**2
+
+    with torch.profiler.profile(
+        activities=[ProfilerActivity.CPU],
+    ) as profiler:
+        run_model_step(fn, kwargs)
+        torch.cuda.synchronize()
+
+    key_averages = profiler.key_averages()
+    cpu_self_time_ms = sum(event.self_cpu_time_total for event in key_averages) / 1000
+    cuda_launch_count = sum(
+        event.count
+        for event in key_averages
+        if event.key in {"cudaLaunchKernel", "cudaGraphLaunch", "cudaLaunchKernelExC"}
+    )
+    profiler_event_count = sum(event.count for event in key_averages)
+
+    return {
+        "cuda_event_ms": cuda_event_ms,
+        "host_wall_ms": host_wall_ms,
+        "cpu_self_time_ms": cpu_self_time_ms,
+        "cuda_launch_count": cuda_launch_count,
+        "profiler_event_count": profiler_event_count,
+        "peak_mem_mib": peak_mem_mib,
+    }

tests/processor/test_pi05_processor.py

@@ -0,0 +1,155 @@
+#!/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.
+
+"""Compare the PI0.5 processor pipeline against the vendored OpenPI reference processors."""
+
+import os
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.configs import FeatureType, PolicyFeature  # noqa: E402
+from lerobot.policies.pi05 import PI05Policy  # noqa: E402
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+from tests.policies.pi0_pi05.utils.openpi_parity import (  # noqa: E402
+    IMAGE_KEYS,
+    assert_processor_inputs_match_lerobot,
+    clone_batch,
+    make_openpi_observation_from_raw,
+    openpi_model_actions_from_raw,
+)
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="OpenPI processor parity uses the PaliGemma tokenizer; run manually outside CI.",
+)
+
+DUMMY_ACTION_DIM = 32
+DUMMY_STATE_DIM = 32
+DUMMY_ACTION_HORIZON = 50
+DUMMY_MAX_TOKEN_LEN = 200
+DEVICE = torch.device("cpu")
+
+DUMMY_DATASET_STATS = {
+    OBS_STATE: {
+        "mean": torch.zeros(DUMMY_STATE_DIM),
+        "std": torch.ones(DUMMY_STATE_DIM),
+        "q01": torch.zeros(DUMMY_STATE_DIM),
+        "q99": torch.ones(DUMMY_STATE_DIM),
+    },
+    ACTION: {
+        "mean": torch.zeros(DUMMY_ACTION_DIM),
+        "std": torch.ones(DUMMY_ACTION_DIM),
+        "q01": torch.zeros(DUMMY_ACTION_DIM),
+        "q99": torch.ones(DUMMY_ACTION_DIM),
+    },
+    "images": {
+        key: {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+            "q01": torch.zeros(3, 224, 224),
+            "q99": torch.ones(3, 224, 224),
+        }
+        for key in IMAGE_KEYS
+    },
+}
+
+
+class PI05PolicyInputAdapter(torch.nn.Module):
+    """Minimal adapter exposing PI0.5 policy image preparation without loading model weights."""
+
+    _preprocess_images = PI05Policy._preprocess_images
+
+    def __init__(self, config: PI05Config) -> None:
+        super().__init__()
+        self.config = config
+        self._device_anchor = torch.nn.Parameter(torch.empty((), device=config.device), requires_grad=False)
+
+
+def create_pi05_config() -> PI05Config:
+    config = PI05Config(device=str(DEVICE))
+    config.max_state_dim = DUMMY_STATE_DIM
+    config.max_action_dim = DUMMY_ACTION_DIM
+    config.chunk_size = DUMMY_ACTION_HORIZON
+    config.n_action_steps = DUMMY_ACTION_HORIZON
+    config.tokenizer_max_length = DUMMY_MAX_TOKEN_LEN
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(DUMMY_STATE_DIM,)),
+        **{
+            f"observation.images.{key}": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224))
+            for key in IMAGE_KEYS
+        },
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(DUMMY_ACTION_DIM,)),
+    }
+    return config
+
+
+def create_dummy_data() -> dict:
+    batch_size = 2
+    prompt = "Pick up the red block and place it in the bin"
+    return {
+        OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=DEVICE),
+        ACTION: torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=DEVICE
+        ),
+        **{
+            f"observation.images.{key}": torch.rand(
+                batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
+            )
+            for key in IMAGE_KEYS
+        },
+        "task": [prompt for _ in range(batch_size)],
+    }
+
+
+def test_pi05_processor_inputs_match_openpi_reference():
+    torch.manual_seed(0)
+    config = create_pi05_config()
+    preprocessor, _ = make_pi05_pre_post_processors(config=config, dataset_stats=DUMMY_DATASET_STATS)
+
+    raw_batch = create_dummy_data()
+    lerobot_batch = preprocessor(clone_batch(raw_batch))
+    openpi_observation = make_openpi_observation_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        max_token_len=DUMMY_MAX_TOKEN_LEN,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=True,
+    )
+
+    assert_processor_inputs_match_lerobot(
+        PI05PolicyInputAdapter(config),
+        lerobot_batch,
+        openpi_observation,
+        compare_state=False,
+    )
+    torch.testing.assert_close(
+        lerobot_batch[ACTION],
+        openpi_model_actions_from_raw(
+            raw_batch,
+            action_dim=DUMMY_ACTION_DIM,
+            dataset_stats=DUMMY_DATASET_STATS,
+            pi05=True,
+        ),
+        rtol=0,
+        atol=0,
+    )

tests/processor/test_pi0_processor.py

@@ -0,0 +1,156 @@
+#!/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.
+
+"""Compare the PI0 processor pipeline against the vendored OpenPI reference processors."""
+
+import os
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.configs import FeatureType, PolicyFeature  # noqa: E402
+from lerobot.policies.pi0 import PI0Policy  # noqa: E402
+from lerobot.policies.pi0.configuration_pi0 import PI0Config  # noqa: E402
+from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+from tests.policies.pi0_pi05.utils.openpi_parity import (  # noqa: E402
+    IMAGE_KEYS,
+    assert_processor_inputs_match_lerobot,
+    clone_batch,
+    make_openpi_observation_from_raw,
+    openpi_model_actions_from_raw,
+)
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="OpenPI processor parity uses the PaliGemma tokenizer; run manually outside CI.",
+)
+
+DUMMY_ACTION_DIM = 32
+DUMMY_STATE_DIM = 32
+DUMMY_ACTION_HORIZON = 50
+DUMMY_MAX_TOKEN_LEN = 48
+DEVICE = torch.device("cpu")
+
+DUMMY_DATASET_STATS = {
+    OBS_STATE: {
+        "mean": torch.zeros(DUMMY_STATE_DIM),
+        "std": torch.ones(DUMMY_STATE_DIM),
+        "q01": torch.zeros(DUMMY_STATE_DIM),
+        "q99": torch.ones(DUMMY_STATE_DIM),
+    },
+    ACTION: {
+        "mean": torch.zeros(DUMMY_ACTION_DIM),
+        "std": torch.ones(DUMMY_ACTION_DIM),
+        "q01": torch.zeros(DUMMY_ACTION_DIM),
+        "q99": torch.ones(DUMMY_ACTION_DIM),
+    },
+    "images": {
+        key: {
+            "mean": torch.zeros(3, 224, 224),
+            "std": torch.ones(3, 224, 224),
+            "q01": torch.zeros(3, 224, 224),
+            "q99": torch.ones(3, 224, 224),
+        }
+        for key in IMAGE_KEYS
+    },
+}
+
+
+class PI0PolicyInputAdapter(torch.nn.Module):
+    """Minimal adapter exposing PI0 policy input-preparation helpers without loading model weights."""
+
+    _preprocess_images = PI0Policy._preprocess_images
+    prepare_state = PI0Policy.prepare_state
+
+    def __init__(self, config: PI0Config) -> None:
+        super().__init__()
+        self.config = config
+        self._device_anchor = torch.nn.Parameter(torch.empty((), device=config.device), requires_grad=False)
+
+
+def create_pi0_config() -> PI0Config:
+    config = PI0Config(device=str(DEVICE))
+    config.max_state_dim = DUMMY_STATE_DIM
+    config.max_action_dim = DUMMY_ACTION_DIM
+    config.chunk_size = DUMMY_ACTION_HORIZON
+    config.n_action_steps = DUMMY_ACTION_HORIZON
+    config.tokenizer_max_length = DUMMY_MAX_TOKEN_LEN
+    config.input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(DUMMY_STATE_DIM,)),
+        **{
+            f"observation.images.{key}": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224))
+            for key in IMAGE_KEYS
+        },
+    }
+    config.output_features = {
+        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(DUMMY_ACTION_DIM,)),
+    }
+    return config
+
+
+def create_dummy_data() -> dict:
+    batch_size = 2
+    prompt = "Pick up the red block and place it in the bin"
+    return {
+        OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=DEVICE),
+        ACTION: torch.randn(
+            batch_size, DUMMY_ACTION_HORIZON, DUMMY_ACTION_DIM, dtype=torch.float32, device=DEVICE
+        ),
+        **{
+            f"observation.images.{key}": torch.rand(
+                batch_size, 3, 224, 224, dtype=torch.float32, device=DEVICE
+            )
+            for key in IMAGE_KEYS
+        },
+        "task": [prompt for _ in range(batch_size)],
+    }
+
+
+def test_pi0_processor_inputs_match_openpi_reference():
+    torch.manual_seed(0)
+    config = create_pi0_config()
+    preprocessor, _ = make_pi0_pre_post_processors(config=config, dataset_stats=DUMMY_DATASET_STATS)
+
+    raw_batch = create_dummy_data()
+    lerobot_batch = preprocessor(clone_batch(raw_batch))
+    openpi_observation = make_openpi_observation_from_raw(
+        raw_batch,
+        action_dim=DUMMY_ACTION_DIM,
+        max_token_len=DUMMY_MAX_TOKEN_LEN,
+        dataset_stats=DUMMY_DATASET_STATS,
+        pi05=False,
+    )
+
+    assert_processor_inputs_match_lerobot(
+        PI0PolicyInputAdapter(config),
+        lerobot_batch,
+        openpi_observation,
+        compare_state=True,
+    )
+    torch.testing.assert_close(
+        lerobot_batch[ACTION],
+        openpi_model_actions_from_raw(
+            raw_batch,
+            action_dim=DUMMY_ACTION_DIM,
+            dataset_stats=DUMMY_DATASET_STATS,
+            pi05=False,
+        ),
+        rtol=0,
+        atol=0,
+    )

tests/rewards/test_modeling_topreward.py

@@ -1,296 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for the TOPReward reward model."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from lerobot.configs.rewards import RewardModelConfig
-from lerobot.rewards.factory import get_reward_model_class, make_reward_model_config
-from lerobot.rewards.topreward import TOPRewardConfig
-from lerobot.rewards.topreward.processor_topreward import TOPREWARD_FEATURE_PREFIX, TOPREWARD_INPUT_KEYS
-from tests.utils import skip_if_package_missing
-
-
-class _FakeQwenModel(torch.nn.Module):
-    """Stand-in for ``Qwen3VLForConditionalGeneration``.
-
-    Returns a ``SimpleNamespace`` with ``logits`` of a controlled shape so
-    the log-prob extraction path in ``compute_reward`` can be exercised
-    without downloading real VLM weights.
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-        self._param = torch.nn.Parameter(torch.zeros(1))
-        self._reward_value: float = -1.5
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):  # noqa: ARG003
-        return cls()
-
-    def forward(  # noqa: ARG002
-        self, input_ids, attention_mask=None, labels=None, logits_to_keep=0, **kwargs
-    ):
-        batch_size, seq_len = input_ids.shape
-        vocab_size = 1000
-        logits = torch.zeros(batch_size, seq_len, vocab_size)
-        # Place a controlled log-prob at the target token position so the
-        # model returns a predictable reward value.
-        # The label-masked suffix is the last token.
-        # After the causal-LM shift (logits[:, :-1], labels[:, 1:]) the scored
-        # position is logits[:, -2, :] predicting labels[:, -1].
-        # We set logits so that log_softmax at the target token ≈ _reward_value.
-        for i in range(batch_size):
-            target_idx = int(input_ids[i, -1].item())
-            logits[i, -2, target_idx] = self._reward_value * -10  # high logit -> high log-prob
-        if logits_to_keep:
-            logits = logits[:, -logits_to_keep:, :]
-        return SimpleNamespace(logits=logits)
-
-
-def _patch_build(monkeypatch) -> None:
-    """Stub out HF AutoX so TOPReward construction is cheap and offline."""
-    from lerobot.rewards.topreward import modeling_topreward
-
-    monkeypatch.setattr(modeling_topreward, "Qwen3VLForConditionalGeneration", _FakeQwenModel)
-
-
-def _make_batch(
-    input_ids: torch.Tensor,
-    attention_mask: torch.Tensor | None = None,
-    labels: torch.Tensor | None = None,
-    *,
-    omit: str | None = None,
-) -> dict[str, torch.Tensor]:
-    """Build a ``compute_reward``-ready batch using TOPReward's namespaced keys."""
-    batch_size, seq_len = input_ids.shape
-    if attention_mask is None:
-        attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)
-    batch: dict[str, torch.Tensor] = {}
-    if labels is not None:
-        batch[f"{TOPREWARD_FEATURE_PREFIX}labels"] = labels
-    batch.update(
-        {
-            f"{TOPREWARD_FEATURE_PREFIX}input_ids": input_ids,
-            f"{TOPREWARD_FEATURE_PREFIX}attention_mask": attention_mask,
-            f"{TOPREWARD_FEATURE_PREFIX}pixel_values_videos": torch.zeros(
-                batch_size, 1536, dtype=torch.float32
-            ),
-            f"{TOPREWARD_FEATURE_PREFIX}video_grid_thw": torch.ones(batch_size, 3, dtype=torch.long),
-            f"{TOPREWARD_FEATURE_PREFIX}mm_token_type_ids": torch.zeros_like(input_ids),
-        }
-    )
-    if omit is not None:
-        batch.pop(f"{TOPREWARD_FEATURE_PREFIX}{omit}", None)
-    return batch
-
-
-def _terminal_labels(input_ids: torch.Tensor) -> torch.Tensor:
-    labels = torch.full_like(input_ids, -100)
-    labels[:, -1] = input_ids[:, -1]
-    return labels
-
-
-# ---------------------------------------------------------------------------
-# Registry + factory
-# ---------------------------------------------------------------------------
-
-
-def test_topreward_config_registered():
-    assert "topreward" in RewardModelConfig.get_known_choices()
-    assert RewardModelConfig.get_choice_class("topreward") is TOPRewardConfig
-    assert isinstance(make_reward_model_config("topreward", device="cpu"), TOPRewardConfig)
-
-
-def test_topreward_factory_returns_in_tree_class():
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    assert get_reward_model_class("topreward") is TOPRewardModel
-
-
-# ---------------------------------------------------------------------------
-# Config validation
-# ---------------------------------------------------------------------------
-
-
-def test_topreward_config_rejects_zero_max_frames():
-    with pytest.raises(ValueError, match="max_frames must be >= 1"):
-        TOPRewardConfig(device="cpu", max_frames=0)
-
-
-def test_topreward_config_rejects_non_positive_fps():
-    with pytest.raises(ValueError, match="fps must be > 0"):
-        TOPRewardConfig(device="cpu", fps=0.0)
-
-
-def test_topreward_config_rejects_suffix_without_instruction_placeholder():
-    with pytest.raises(ValueError, match=r"\{instruction\}"):
-        TOPRewardConfig(device="cpu", prompt_suffix_template="no placeholder here")
-
-
-# ---------------------------------------------------------------------------
-# compute_reward
-# ---------------------------------------------------------------------------
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_compute_reward_returns_one_scalar_per_sample(monkeypatch):
-    """``compute_reward`` must return a ``(B,)`` float32 tensor with one
-    log-prob reward per sample, consuming pre-encoded Qwen-VL tensors."""
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(device="cpu")
-    model = TOPRewardModel(cfg)
-
-    input_ids = torch.randint(0, 100, (2, 10))
-    attention_mask = torch.ones(2, 10, dtype=torch.long)
-    labels = _terminal_labels(input_ids)
-
-    batch = _make_batch(input_ids, attention_mask, labels)
-    rewards = model.compute_reward(batch)
-
-    assert rewards.shape == (2,)
-    assert rewards.dtype == torch.float32
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_compute_reward_applies_success_threshold(monkeypatch):
-    """When ``success_threshold`` is finite, the model returns binary success."""
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(device="cpu", success_threshold=0.0)
-    model = TOPRewardModel(cfg)
-
-    input_ids = torch.randint(0, 100, (2, 10))
-    attention_mask = torch.ones(2, 10, dtype=torch.long)
-    labels = _terminal_labels(input_ids)
-
-    batch = _make_batch(input_ids, attention_mask, labels)
-    rewards = model.compute_reward(batch)
-
-    assert rewards.shape == (2,)
-    assert set(rewards.tolist()).issubset({0.0, 1.0})
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_compute_reward_errors_when_inputs_missing(monkeypatch):
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(device="cpu")
-    model = TOPRewardModel(cfg)
-
-    with pytest.raises(KeyError, match=r"observation\.topreward\.input_ids"):
-        model.compute_reward(_make_batch(torch.randint(0, 100, (1, 10)), omit="input_ids"))
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_compute_reward_errors_when_labels_missing(monkeypatch):
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(device="cpu")
-    model = TOPRewardModel(cfg)
-
-    input_ids = torch.randint(0, 100, (1, 10))
-    with pytest.raises(KeyError, match=r"observation\.topreward\.labels"):
-        model.compute_reward(_make_batch(input_ids, labels=None))
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_compute_reward_requires_all_encoder_keys(monkeypatch):
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(device="cpu")
-    model = TOPRewardModel(cfg)
-
-    input_ids = torch.randint(0, 100, (1, 10))
-    labels = _terminal_labels(input_ids)
-    required_encoder_keys = set(TOPREWARD_INPUT_KEYS) - {"input_ids", "labels"}
-
-    for key in required_encoder_keys:
-        with pytest.raises(KeyError, match=rf"observation\.topreward\.{key}"):
-            model.compute_reward(_make_batch(input_ids, labels=labels, omit=key))
-
-
-# ---------------------------------------------------------------------------
-# Save / load — config-only checkpoint
-# ---------------------------------------------------------------------------
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_save_pretrained_writes_only_config_json(monkeypatch, tmp_path):
-    from huggingface_hub.constants import CONFIG_NAME, SAFETENSORS_SINGLE_FILE
-
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(
-        device="cpu",
-        vlm_name="Qwen/Qwen3-VL-8B-Instruct",
-        fps=4.0,
-        image_key="observation.images.front",
-    )
-    model = TOPRewardModel(cfg)
-    model.save_pretrained(str(tmp_path))
-
-    assert (tmp_path / CONFIG_NAME).exists()
-    assert not (tmp_path / SAFETENSORS_SINGLE_FILE).exists()
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_from_pretrained_local_dir_roundtrips_config(monkeypatch, tmp_path):
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(
-        device="cpu",
-        vlm_name="Qwen/Qwen3-VL-8B-Instruct",
-        fps=4.0,
-        image_key="observation.images.front",
-        add_chat_template=True,
-        success_threshold=-1.5,
-    )
-    TOPRewardModel(cfg).save_pretrained(str(tmp_path))
-
-    reloaded = TOPRewardModel.from_pretrained(str(tmp_path))
-
-    assert isinstance(reloaded.config, TOPRewardConfig)
-    assert reloaded.config.vlm_name == "Qwen/Qwen3-VL-8B-Instruct"
-    assert reloaded.config.fps == 4.0
-    assert reloaded.config.image_key == "observation.images.front"
-    assert reloaded.config.add_chat_template is True
-    assert reloaded.config.success_threshold == -1.5
-
-
-@skip_if_package_missing("transformers")
-def test_topreward_is_not_trainable(monkeypatch):
-    from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
-
-    _patch_build(monkeypatch)
-    cfg = TOPRewardConfig(device="cpu")
-    model = TOPRewardModel(cfg)
-
-    assert model.is_trainable is False
-    with pytest.raises(NotImplementedError, match="not trainable"):
-        model.forward({"x": torch.zeros(1)})

tests/rewards/test_topreward.py

@@ -1,80 +0,0 @@
-# Copyright 2026 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.
-
-"""End-to-end TOPReward smoke test with the real Qwen3-VL model."""
-
-import os
-
-import pytest
-import torch
-
-pytest.importorskip("transformers")
-
-from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig  # noqa: E402
-from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel  # noqa: E402
-from lerobot.rewards.topreward.processor_topreward import (  # noqa: E402
-    TOPREWARD_FEATURE_PREFIX,
-    TOPREWARD_INPUT_KEYS,
-    make_topreward_pre_post_processors,
-)
-from tests.utils import require_cuda  # noqa: E402
-
-pytestmark = pytest.mark.skipif(
-    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires downloading and loading Qwen3-VL and is not meant for CI",
-)
-
-
-def _make_dummy_topreward_batch(image_key: str, task_key: str) -> dict[str, object]:
-    num_frames = 4
-    image_size = 64
-    frames = torch.zeros(1, num_frames, 3, image_size, image_size, dtype=torch.uint8)
-    for frame_idx in range(num_frames):
-        frames[0, frame_idx, 0].fill_(min(frame_idx * 48, 255))
-        frames[0, frame_idx, 1].fill_(96)
-        frames[0, frame_idx, 2].fill_(192)
-
-    return {
-        image_key: frames,
-        task_key: ["pick up the red cube"],
-    }
-
-
-@require_cuda
-def test_topreward_full_qwen3vl_preprocessor_to_compute_reward():
-    cfg = TOPRewardConfig(
-        vlm_name="Qwen/Qwen3-VL-8B-Instruct",
-        device="cuda",
-        max_frames=4,
-        fps=2.0,
-        max_input_length=4096,
-    )
-
-    preprocessor, _ = make_topreward_pre_post_processors(cfg)
-    encoded_batch = preprocessor(_make_dummy_topreward_batch(cfg.image_key, cfg.task_key))
-    for key in TOPREWARD_INPUT_KEYS:
-        assert f"{TOPREWARD_FEATURE_PREFIX}{key}" in encoded_batch
-
-    model = TOPRewardModel(cfg)
-    try:
-        model.to(cfg.device)
-        model.eval()
-        rewards = model.compute_reward(encoded_batch)
-    finally:
-        del model
-        torch.cuda.empty_cache()
-
-    assert rewards.shape == (1,)
-    assert rewards.dtype == torch.float32
-    assert torch.isfinite(rewards).all()

tests/rewards/test_topreward_processor.py

@@ -1,246 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for TOPReward's pre-processing helpers and encoder step."""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.rewards.topreward.processor_topreward import (
-    TOPREWARD_FEATURE_PREFIX,
-    TOPREWARD_INPUT_KEYS,
-    _expand_tasks,
-    _prepare_video_batch,
-)
-from lerobot.types import TransitionKey
-from tests.utils import skip_if_package_missing
-
-# ---------------------------------------------------------------------------
-# _prepare_video_batch — raw image/video batch -> (B, T, C, H, W) uint8
-# ---------------------------------------------------------------------------
-
-
-def test_prepare_video_batch_batched_chw_float_is_converted_to_uint8():
-    video = torch.rand(2, 4, 3, 8, 8)
-    tensor = _prepare_video_batch(video, max_frames=None)
-
-    assert tensor.shape == (2, 4, 3, 8, 8)
-    assert tensor.dtype == torch.uint8
-    assert tensor.min() >= 0 and tensor.max() <= 255
-
-
-def test_prepare_video_batch_batched_thwc_uint8_is_permuted_to_channel_first():
-    video = torch.randint(0, 256, (2, 3, 8, 8, 3), dtype=torch.uint8)
-    tensor = _prepare_video_batch(video, max_frames=None)
-
-    assert tensor.shape == (2, 3, 3, 8, 8)
-    assert tensor.dtype == torch.uint8
-
-
-def test_prepare_video_batch_max_frames_tail_crops_recent_frames():
-    video = torch.zeros(1, 10, 3, 4, 4)
-    for t in range(10):
-        video[:, t] = t / 9.0
-
-    tensor = _prepare_video_batch(video, max_frames=3)
-
-    assert tensor.shape == (1, 3, 3, 4, 4)
-    assert int(tensor[0, 0, 0, 0, 0]) == int(7 / 9 * 255)
-    assert int(tensor[0, -1, 0, 0, 0]) == 255
-
-
-def test_prepare_video_batch_rejects_3d_input():
-    with pytest.raises(ValueError, match="Expected TOPReward frames"):
-        _prepare_video_batch(torch.zeros(4, 8, 8), max_frames=None)
-
-
-def test_prepare_video_batch_floats_above_one_are_rescaled_and_clipped():
-    video = torch.full((1, 1, 3, 2, 2), 5.0)
-    tensor = _prepare_video_batch(video, max_frames=None)
-
-    assert tensor.shape == (1, 1, 3, 2, 2)
-    assert int(tensor.max()) == 255
-
-
-def test_prepare_video_batch_clips_very_large_floats_to_uint8_max():
-    video = torch.full((1, 1, 3, 2, 2), 300.0)
-    tensor = _prepare_video_batch(video, max_frames=None)
-
-    assert int(tensor.max()) == 255
-
-
-# ---------------------------------------------------------------------------
-# _expand_tasks — string / list / tuple broadcasting to batch size
-# ---------------------------------------------------------------------------
-
-
-def test_expand_tasks_string_is_broadcast_to_batch_size():
-    assert _expand_tasks("pick up", batch_size=3, default=None) == ["pick up", "pick up", "pick up"]
-
-
-def test_expand_tasks_list_of_matching_size_passes_through():
-    assert _expand_tasks(["a", "b", "c"], batch_size=3, default=None) == ["a", "b", "c"]
-
-
-def test_expand_tasks_tuple_is_normalised_to_list():
-    assert _expand_tasks(("a", "b"), batch_size=2, default=None) == ["a", "b"]
-
-
-def test_expand_tasks_single_element_list_is_broadcast():
-    assert _expand_tasks(["only one"], batch_size=3, default=None) == ["only one"] * 3
-
-
-def test_expand_tasks_size_mismatch_raises():
-    with pytest.raises(ValueError, match="Expected 3 tasks"):
-        _expand_tasks(["a", "b"], batch_size=3, default=None)
-
-
-def test_expand_tasks_missing_uses_default():
-    assert _expand_tasks(None, batch_size=2, default="fallback") == ["fallback", "fallback"]
-
-
-def test_expand_tasks_missing_without_default_raises():
-    with pytest.raises(KeyError, match="task description"):
-        _expand_tasks(None, batch_size=1, default=None)
-
-
-def test_expand_tasks_wrong_type_raises():
-    with pytest.raises(TypeError, match="must be a string or list"):
-        _expand_tasks(42, batch_size=1, default=None)
-
-
-# ---------------------------------------------------------------------------
-# Encoder step — stubbed AutoProcessor
-# ---------------------------------------------------------------------------
-
-
-def _skip_if_topreward_extras_missing(func):
-    func = skip_if_package_missing("transformers")(func)
-    return func
-
-
-class _FakeTokenizer:
-    eos_token = "<|endoftext|>"
-    pad_token = "<|endoftext|>"
-
-    def __call__(self, *args, **kwargs):
-        return {"input_ids": torch.zeros(1, 10, dtype=torch.long)}
-
-
-class _FakeAutoProcessor:
-    def __init__(self) -> None:
-        self.tokenizer = _FakeTokenizer()
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):  # noqa: ARG003
-        return cls()
-
-    def apply_chat_template(self, messages, **kwargs):  # noqa: ARG002
-        return "fake_prompt_text"
-
-    def __call__(self, text=None, images=None, videos=None, **kwargs):  # noqa: ARG002
-        seq_len = 10
-        batch_size = len(text) if isinstance(text, list) else 1
-        return {
-            "input_ids": torch.randint(0, 100, (batch_size, seq_len)),
-            "attention_mask": torch.ones(batch_size, seq_len, dtype=torch.long),
-            "pixel_values_videos": torch.zeros(batch_size, 1536, dtype=torch.float32),
-            "video_grid_thw": torch.ones(batch_size, 3, dtype=torch.long),
-            "mm_token_type_ids": torch.zeros(batch_size, seq_len, dtype=torch.long),
-        }
-
-
-def _build_step(monkeypatch, **overrides):
-    from lerobot.rewards.topreward import processor_topreward
-
-    monkeypatch.setattr(processor_topreward, "AutoProcessor", _FakeAutoProcessor)
-    return processor_topreward.TOPRewardEncoderProcessorStep(**overrides)
-
-
-def _make_transition(observation: dict, complementary: dict | None = None) -> dict:
-    transition: dict = {TransitionKey.OBSERVATION: observation}
-    if complementary is not None:
-        transition[TransitionKey.COMPLEMENTARY_DATA] = complementary
-    return transition
-
-
-@_skip_if_topreward_extras_missing
-def test_encoder_step_emits_input_ids_and_labels(monkeypatch):
-    """The processor must emit Qwen-VL tensors including ``input_ids`` and
-    ``labels`` under the ``observation.topreward.*`` namespace."""
-    step = _build_step(monkeypatch)
-
-    frames_batch = torch.zeros(2, 4, 3, 8, 8)
-    out = step(
-        _make_transition(
-            observation={"observation.images.top": frames_batch},
-            complementary={"task": ["pick", "place"]},
-        )
-    )
-
-    obs_out = out[TransitionKey.OBSERVATION]
-    for key in TOPREWARD_INPUT_KEYS:
-        assert f"{TOPREWARD_FEATURE_PREFIX}{key}" in obs_out
-
-    input_ids = obs_out[f"{TOPREWARD_FEATURE_PREFIX}input_ids"]
-    labels = obs_out[f"{TOPREWARD_FEATURE_PREFIX}labels"]
-    assert labels.dtype == torch.long
-    assert labels.shape == (2, 10)
-    assert labels[:, :-1].eq(-100).all()
-    assert labels[:, -1].equal(input_ids[:, -1])
-
-
-@_skip_if_topreward_extras_missing
-def test_encoder_step_get_config_roundtrips_user_fields(monkeypatch):
-    step = _build_step(
-        monkeypatch,
-        vlm_name="Qwen/Qwen3-VL-8B-Instruct",
-        image_key="observation.images.cam_top",
-        task_key="task",
-        default_task="do the thing",
-        max_frames=8,
-        fps=4.0,
-        add_chat_template=True,
-        max_length=2048,
-    )
-
-    cfg = step.get_config()
-    assert cfg["vlm_name"] == "Qwen/Qwen3-VL-8B-Instruct"
-    assert cfg["image_key"] == "observation.images.cam_top"
-    assert cfg["default_task"] == "do the thing"
-    assert cfg["max_frames"] == 8
-    assert cfg["fps"] == 4.0
-    assert cfg["add_chat_template"] is True
-    assert cfg["max_length"] == 2048
-
-
-@_skip_if_topreward_extras_missing
-def test_encoder_step_transform_features_is_identity(monkeypatch):
-    step = _build_step(monkeypatch)
-    features = {
-        PipelineFeatureType.OBSERVATION: {
-            "observation.images.top": PolicyFeature(shape=(3, 224, 224), type=FeatureType.VISUAL),
-        }
-    }
-    assert step.transform_features(features) == features
-
-
-@_skip_if_topreward_extras_missing
-def test_encoder_step_rejects_missing_image_key(monkeypatch):
-    step = _build_step(monkeypatch, image_key="observation.images.top")
-    with pytest.raises(KeyError, match="image key"):
-        step(_make_transition(observation={}, complementary={"task": "pick"}))

tests/test_yaml_policy_path.py

@@ -1,10 +1,14 @@
 """Tests for policy.path support in YAML config files (issue #2957)."""
 
 import json
+import sys
 import tempfile
+from dataclasses import dataclass, field
+from unittest.mock import patch
 
 import yaml
 
+from lerobot.configs import parser
 from lerobot.configs.parser import (
     _config_path_args,
     _config_yaml_overrides,
@@ -16,7 +20,8 @@ from lerobot.configs.parser import (
 
 
 def test_extract_path_fields_from_yaml():
-    """Test that policy.path is extracted from a YAML config and removed."""
+    """Test that policy.path is extracted from a YAML config and the policy block
+    is removed entirely (siblings are captured separately as cli_overrides)."""
     config = {
         "dataset": {"repo_id": "lerobot/pusht"},
         "policy": {"type": "smolvla", "path": "lerobot/smolvla_base", "push_to_hub": False},
@@ -26,26 +31,33 @@ def test_extract_path_fields_from_yaml():
         config_path = f.name
 
     _config_path_args.clear()
+    _config_yaml_overrides.clear()
     cleaned_path = extract_path_fields_from_config(config_path, ["policy"])
 
     # Path should be extracted and stored
     assert _config_path_args["policy"] == "lerobot/smolvla_base"
 
-    # Cleaned config should not have the path field
+    # Cleaned config should not have the policy block at all -- draccus must not
+    # try to decode it as PreTrainedConfig; the actual config comes from
+    # from_pretrained(path) with the captured overrides applied on top.
     with open(cleaned_path) as f:
         cleaned = yaml.safe_load(f)
-    assert "path" not in cleaned["policy"]
-    assert cleaned["policy"]["type"] == "smolvla"
-    assert cleaned["policy"]["push_to_hub"] is False
+    assert "policy" not in cleaned
 
     # Original dataset should be untouched
     assert cleaned["dataset"]["repo_id"] == "lerobot/pusht"
 
+    # Sibling overrides (excluding type/path) captured for from_pretrained.
+    overrides = get_yaml_overrides("policy")
+    assert any("push_to_hub=false" in o for o in overrides)
+
     _config_path_args.clear()
+    _config_yaml_overrides.clear()
 
 
 def test_extract_path_fields_from_json():
-    """Test that policy.path is extracted from a JSON config."""
+    """Test that policy.path is extracted from a JSON config and the policy
+    block is removed entirely."""
     config = {
         "policy": {"type": "act", "path": "some/local/path"},
     }
@@ -54,15 +66,17 @@ def test_extract_path_fields_from_json():
         config_path = f.name
 
     _config_path_args.clear()
+    _config_yaml_overrides.clear()
     cleaned_path = extract_path_fields_from_config(config_path, ["policy"])
 
     assert _config_path_args["policy"] == "some/local/path"
 
     with open(cleaned_path) as f:
         cleaned = json.load(f)
-    assert "path" not in cleaned["policy"]
+    assert "policy" not in cleaned
 
     _config_path_args.clear()
+    _config_yaml_overrides.clear()
 
 
 def test_extract_no_path_returns_original():
@@ -216,3 +230,91 @@ def test_flatten_nested_with_bools():
     args = _flatten_to_cli_args(d)
     assert "--optimizer.use_warmup=true" in args
     assert "--optimizer.lr=0.01" in args
+
+
+def test_extract_removes_field_with_siblings_and_no_type():
+    """Regression: when policy.path has siblings but no type:, the entire policy
+    block must still be removed from the cleaned config. Otherwise draccus tries
+    to decode the leftover dict as PreTrainedConfig and crashes on the missing
+    type discriminator.
+    """
+    config = {
+        "dataset": {"repo_id": "lerobot/pusht"},
+        "policy": {
+            "path": "lerobot/smolvla_base",
+            "n_action_steps": 10,
+            "dtype": "bfloat16",
+        },
+    }
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        yaml.dump(config, f)
+        config_path = f.name
+
+    _config_path_args.clear()
+    _config_yaml_overrides.clear()
+    cleaned_path = extract_path_fields_from_config(config_path, ["policy"])
+
+    with open(cleaned_path) as f:
+        cleaned = yaml.safe_load(f) or {}
+    assert "policy" not in cleaned, "policy block should be fully removed when path is present"
+    assert cleaned["dataset"]["repo_id"] == "lerobot/pusht"
+    assert _config_path_args["policy"] == "lerobot/smolvla_base"
+    overrides = get_yaml_overrides("policy")
+    assert any("n_action_steps=10" in o for o in overrides)
+    assert any("dtype=bfloat16" in o for o in overrides)
+
+    _config_path_args.clear()
+    _config_yaml_overrides.clear()
+
+
+@dataclass
+class _DummyNested:
+    foo: int = 0
+
+
+@dataclass
+class _DummyConfig:
+    nested: _DummyNested = field(default_factory=_DummyNested)
+    other: str = "default"
+
+    @classmethod
+    def __get_path_fields__(cls):
+        return ["nested"]
+
+
+def test_wrap_uses_cleaned_config_for_draccus_parse():
+    """Regression: wrap() updates config_path_cli to point at the cleaned temp
+    file but must propagate that to the draccus.parse fallback branch. Without
+    the fix, cli_args still contains --config_path=<original> and draccus reads
+    the original YAML with `path:` still in it, crashing on the unknown field.
+    """
+    config = {
+        "nested": {"path": "some/checkpoint", "foo": 42},
+        "other": "set-via-yaml",
+    }
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        yaml.dump(config, f)
+        config_path = f.name
+
+    _config_path_args.clear()
+    _config_yaml_overrides.clear()
+
+    captured: dict = {}
+
+    @parser.wrap()
+    def main(cfg: _DummyConfig) -> _DummyConfig:
+        captured["cfg"] = cfg
+        return cfg
+
+    with patch.object(sys, "argv", ["prog", f"--config_path={config_path}"]):
+        main()
+
+    assert captured["cfg"].other == "set-via-yaml"
+    assert _config_path_args["nested"] == "some/checkpoint"
+    # Cleaned config dropped `nested:` entirely; defaults stand for this wrapper
+    # class (a real PreTrainedConfig would now load the checkpoint and apply
+    # the captured yaml_overrides via from_pretrained()).
+    assert captured["cfg"].nested.foo == 0
+
+    _config_path_args.clear()
+    _config_yaml_overrides.clear()