Add GR00T N1.7 support

Add GR00T N1.7 policy configuration, checkpoint compatibility, processor parity, LIBERO documentation, and focused tests.

Co-authored-by: Ryan Halabi <ryhalabi@nvidia.com>

docs/source/_toctree.yml

@@ -64,7 +64,7 @@
   - local: eo1
     title: EO-1
   - local: groot
-    title: NVIDIA GR00T N1.5
+    title: NVIDIA GR00T
   - local: xvla
     title: X-VLA
   - local: multi_task_dit
@@ -75,6 +75,8 @@
 - sections:
   - local: sarm
     title: SARM
+  - local: robometer
+    title: ROBOMETER
   - local: topreward
     title: TOPReward
   title: "Reward Models"

docs/source/groot.mdx

@@ -1,16 +1,16 @@
-# GR00T N1.5 Policy
+# GR00T Policy
 
-GR00T N1.5 is an open foundation model from NVIDIA designed for generalized humanoid robot reasoning and skills. It is a cross-embodiment model that accepts multimodal input, including language and images, to perform manipulation tasks in diverse environments.
+GR00T is an NVIDIA foundation model family for generalized humanoid robot reasoning and skills. It is a cross-embodiment policy that accepts multimodal input, including language, images, and proprioception, to perform manipulation tasks in diverse environments.
 
-This document outlines the specifics of its integration and usage within the LeRobot framework.
+LeRobot integrates GR00T through the `groot` policy type. The default model family is GR00T N1.5, and GR00T N1.7 can be selected with `policy.model_version=n1.7`.
 
 ## Model Overview
 
-NVIDIA Isaac GR00T N1.5 is an upgraded version of the GR00T N1 foundation model. It is built to improve generalization and language-following abilities for humanoid robots.
+NVIDIA Isaac GR00T N1.5 is an upgraded version of the GR00T N1 foundation model. GR00T N1.7 extends the family with a Cosmos-Reason2/Qwen3-VL backbone and N1.7 checkpoints for SimplerEnv, DROID, and LIBERO.
 
-Developers and researchers can post-train GR00T N1.5 with their own real or synthetic data to adapt it for specific humanoid robots or tasks.
+Developers and researchers can post-train GR00T with their own real or synthetic data to adapt it for specific humanoid robots or tasks.
 
-GR00T N1.5 (specifically the GR00T-N1.5-3B model) is built using pre-trained vision and language encoders. It utilizes a flow matching action transformer to model a chunk of actions, conditioned on vision, language, and proprioception.
+GR00T uses pre-trained vision and language encoders with a flow matching action transformer to model a chunk of actions conditioned on vision, language, and proprioception.
 
 <img
   src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-groot-paper1%20(1).png"
@@ -28,33 +28,35 @@ This approach allows the model to be highly adaptable through post-training for
 
 ## Installation Requirements
 
-As of today, GR00T N1.5 requires flash attention for it's internal working.
-
-We are working on making this optional, but in the meantime that means that we require an extra installation step and it can only be used in CUDA enabled devices.
-
-1. Following the Environment Setup of our [Installation Guide](./installation). **Attention** don't install `lerobot` in this step.
-2. Install [Flash Attention](https://github.com/Dao-AILab/flash-attention) by running:
+Install LeRobot with the GR00T extra:
 
 ```bash
-# Check https://pytorch.org/get-started/locally/ for your system
-pip install "torch>=2.2.1,<2.8.0" "torchvision>=0.21.0,<0.23.0" # --index-url https://download.pytorch.org/whl/cu1XX
-pip install ninja "packaging>=24.2,<26.0" # flash attention dependencies
-pip install "flash-attn>=2.5.9,<3.0.0" --no-build-isolation
-python -c "import flash_attn; print(f'Flash Attention {flash_attn.__version__} imported successfully')"
+pip install "lerobot[groot]"
 ```
 
-3. Install LeRobot by running:
+GR00T is intended for NVIDIA GPU-accelerated systems. The `groot` extra installs the policy dependencies, including `transformers`, `diffusers`, `peft`, `dm-tree`, and Flash Attention where available. If Flash Attention is unavailable or incompatible, LeRobot falls back to SDPA attention in supported GR00T paths, with lower expected throughput.
+
+For a source checkout, follow the Environment Setup in the [Installation Guide](./installation), then install the extra:
 
 ```bash
-pip install lerobot[groot]
+uv sync --locked --extra groot
 ```
 
+If you need to install Flash Attention manually for your CUDA/PyTorch build, use the wheel or source build recommended by the [Flash Attention project](https://github.com/Dao-AILab/flash-attention).
+
 ## Usage
 
-To use GR00T in your LeRobot configuration, specify the policy type as:
+To use GR00T N1.5 in your LeRobot configuration, specify the policy type:
 
-```python
-policy.type=groot
+```bash
+--policy.type=groot
+```
+
+To use GR00T N1.7:
+
+```bash
+--policy.type=groot \
+--policy.model_version=n1.7
 ```
 
 ## Training
@@ -85,14 +87,20 @@ accelerate launch \
   --job_name=$JOB_NAME
 ```
 
+For N1.7, add:
+
+```bash
+--policy.model_version=n1.7
+```
+
 ## Performance Results
 
-### Libero Benchmark Results
+### LIBERO Benchmark Results
 
 > [!NOTE]
-> Follow our instructions for Libero usage: [Libero](./libero)
+> Follow the [LIBERO](./libero) setup instructions before running `lerobot-eval`.
 
-GR00T has demonstrated strong performance on the Libero benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the Libero dataset and compared the results to the GR00T reference results.
+GR00T has demonstrated strong performance on the LIBERO benchmark suite. To compare and test its LeRobot implementation, we finetuned the GR00T N1.5 model for 30k steps on the LIBERO dataset and compared the results to the GR00T reference results.
 
 | Benchmark          | LeRobot Implementation | GR00T Reference |
 | ------------------ | ---------------------- | --------------- |
@@ -101,7 +109,49 @@ GR00T has demonstrated strong performance on the Libero benchmark suite. To comp
 | **Libero Long**    | 82.0%                  | 76.0%           |
 | **Average**        | 87.0%                  | 87.0%           |
 
-These results demonstrate GR00T's strong generalization capabilities across diverse robotic manipulation tasks. To reproduce these results, you can follow the instructions in the [Libero](https://huggingface.co/docs/lerobot/libero) section.
+These results demonstrate GR00T's strong generalization capabilities across diverse robotic manipulation tasks. To reproduce these results, follow the instructions in the [LIBERO](./libero) section.
+
+### GR00T N1.7 LIBERO Checkpoints
+
+NVIDIA publishes GR00T N1.7 LIBERO checkpoints at [`nvidia/GR00T-N1.7-LIBERO`](https://huggingface.co/nvidia/GR00T-N1.7-LIBERO), with one subdirectory per LIBERO suite:
+
+| Suite          | Checkpoint subdirectory |
+| -------------- | ----------------------- |
+| LIBERO Spatial | `libero_spatial`        |
+| LIBERO Object  | `libero_object`         |
+| LIBERO Goal    | `libero_goal`           |
+| LIBERO 10      | `libero_10`             |
+
+Preliminary LeRobot integration results:
+
+| Suite          | Status | Success rate | n_episodes |
+| -------------- | ------ | -----------: | ---------: |
+| LIBERO Spatial | ✓      |         ~95% |         XX |
+| LIBERO Object  | ✓      |          XX% |         XX |
+| LIBERO Goal    | ✓      |          XX% |         XX |
+| LIBERO 10      | ✓      |          XX% |         XX |
+| **Average**    | ✓      |      **XX%** |     **XX** |
+
+Replace the `XX` placeholders with final eval artifacts before merge.
+
+Download the suite checkpoint locally, then point `--policy.base_model_path` at the downloaded subdirectory. `--policy.path` is reserved for LeRobot checkpoints that contain a LeRobot `config.json` with a `type` field.
+
+```bash
+huggingface-cli download nvidia/GR00T-N1.7-LIBERO \
+  --include "libero_spatial/*" \
+  --local-dir ./GR00T-N1.7-LIBERO
+
+lerobot-eval \
+  --policy.type=groot \
+  --policy.model_version=n1.7 \
+  --policy.base_model_path=./GR00T-N1.7-LIBERO/libero_spatial \
+  --policy.embodiment_tag=libero_sim \
+  --env.type=libero \
+  --env.task=libero_spatial \
+  --eval.n_episodes=50
+```
+
+Use `eval.n_episodes >= 50` per suite when reporting success rates.
 
 ### Evaluate in your hardware setup
 
@@ -131,4 +181,4 @@ lerobot-rollout\
 
 ## License
 
-This model follows NVIDIA's proprietary license, consistent with the original [GR00T repository](https://github.com/NVIDIA/Isaac-GR00T). Future versions (starting from N1.7) will follow **Apache 2.0 License**.
+GR00T N1.5 follows NVIDIA's license terms, consistent with the original [GR00T repository](https://github.com/NVIDIA/Isaac-GR00T). GR00T N1.7 is released under the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/).

docs/source/policy_groot_README.md

@@ -24,4 +24,8 @@ Code: https://github.com/NVIDIA/Isaac-GR00T
 
 Blog: https://developer.nvidia.com/isaac/gr00t
 
-Hugging Face Model: https://huggingface.co/nvidia/GR00T-N1.5-3B
+Hugging Face Models:
+
+- GR00T N1.5: https://huggingface.co/nvidia/GR00T-N1.5-3B
+- GR00T N1.7: https://huggingface.co/nvidia/GR00T-N1.7-3B
+- GR00T N1.7 LIBERO checkpoints: https://huggingface.co/nvidia/GR00T-N1.7-LIBERO

docs/source/robometer.mdx

@@ -0,0 +1,185 @@
+# ROBOMETER
+
+ROBOMETER is a **general-purpose video-language robotic reward model**. It predicts dense, frame-level task progress and frame-level success from a trajectory video and a task description.
+
+**Paper**: [ROBOMETER: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons](https://arxiv.org/abs/2603.02115)
+**Project**: [robometer.github.io](https://robometer.github.io/)
+**Original code**: [github.com/robometer/robometer](https://github.com/robometer/robometer)
+**Checkpoint**: [lerobot/Robometer-4B](https://huggingface.co/lerobot/Robometer-4B)
+
+## Overview
+
+ROBOMETER builds on `Qwen/Qwen3-VL-4B-Instruct` and adds three lightweight prediction heads:
+
+- **Progress head**: predicts per-frame task progress in `[0, 1]`.
+- **Success head**: predicts per-frame task success probability.
+- **Preference head**: predicts which of two trajectories better completes the task during training.
+
+The paper trains ROBOMETER with a composite objective:
+
+```text
+L = L_pref + L_prog + L_succ
+```
+
+The LeRobot integration is currently **inference-only**. It preserves the preference head so that the published `Robometer-4B` checkpoint loads without remapping, but `compute_reward()` queries the progress or success head only.
+
+## What the LeRobot Integration Covers
+
+- Standard `reward_model.type=robometer` configuration through LeRobot.
+- Qwen3-VL image and text preprocessing through `RobometerEncoderProcessorStep`.
+- LeRobot reward-model save/load APIs through `PreTrainedRewardModel`.
+- Dense, frame-level progress and success predictions internally.
+- A scalar reward through `compute_reward()` for downstream LeRobot reward-model usage.
+
+This page focuses on using the published ROBOMETER checkpoint as a zero-shot reward model. Training ROBOMETER from scratch is outside the current LeRobot integration.
+
+## Installation Requirements
+
+1. Install LeRobot by following the [Installation Guide](./installation).
+2. Install the ROBOMETER dependencies:
+
+```bash
+pip install -e ".[robometer]"
+```
+
+If you use `uv` directly from a source checkout:
+
+```bash
+uv sync --extra robometer
+```
+
+ROBOMETER uses a Qwen3-VL-4B backbone, so GPU inference is strongly recommended.
+
+## Model Inputs and Outputs
+
+ROBOMETER 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/video observation used by ROBOMETER            |
+| `reward_model.task_key`   | `task`                   | Key in complementary data that stores the task string |
+| `reward_model.max_frames` | `8`                      | Maximum number of frames passed to ROBOMETER          |
+
+The model predicts per-frame progress and success internally. The LeRobot reward API returns a scalar per sample:
+
+- `reward_output="progress"` (default): return the last-frame progress, clamped to `[0, 1]`.
+- `reward_output="success"`: return `1.0` if the last-frame success probability is above `success_threshold`, otherwise `0.0`.
+
+## Usage
+
+### Load the Reward Model Directly
+
+```python
+from lerobot.rewards.robometer import RobometerConfig, RobometerRewardModel
+
+cfg = RobometerConfig(
+    pretrained_path="lerobot/Robometer-4B",
+    device="cuda",
+    reward_output="progress",
+)
+reward_model = RobometerRewardModel.from_pretrained(cfg.pretrained_path, config=cfg)
+```
+
+### Encode Frames and Compute a Reward
+
+For a direct Python call, provide frames as `uint8` arrays with shape `(T, H, W, C)` and a task string:
+
+```python
+from lerobot.rewards.robometer.modeling_robometer import ROBOMETER_FEATURE_PREFIX
+from lerobot.rewards.robometer.processor_robometer import RobometerEncoderProcessorStep
+
+# frames: np.ndarray, shape (T, H, W, C), dtype uint8
+# task: str
+encoder = RobometerEncoderProcessorStep(
+    base_model_id=cfg.base_model_id,
+    use_multi_image=cfg.use_multi_image,
+    use_per_frame_progress_token=cfg.use_per_frame_progress_token,
+    max_frames=cfg.max_frames,
+)
+
+encoded = encoder.encode_samples([(frames, task)])
+batch = {f"{ROBOMETER_FEATURE_PREFIX}{key}": value for key, value in encoded.items()}
+
+reward = reward_model.compute_reward(batch)
+```
+
+`reward` is a tensor of shape `(batch_size,)`.
+
+### Use the Reward Factory
+
+You can also instantiate ROBOMETER through 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(
+    "robometer",
+    pretrained_path="lerobot/Robometer-4B",
+    device="cuda",
+    image_key="observation.images.top",
+)
+reward_model = make_reward_model(cfg)
+preprocessor, postprocessor = make_reward_pre_post_processors(cfg)
+```
+
+The preprocessor writes Qwen-VL tensors under the `observation.robometer.*` namespace, and `compute_reward()` reads those encoded tensors.
+
+## Configuration Notes
+
+### Backbone and Vocabulary
+
+The published checkpoint uses a Qwen3-VL-4B backbone. ROBOMETER adds five special tokens to the tokenizer in a fixed order:
+
+```text
+<|split_token|>
+<|reward_token|>
+<|pref_token|>
+<|sim_token|>
+<|prog_token|>
+```
+
+`<|prog_token|>` is inserted after each frame and is the hidden-state position used for per-frame progress and success prediction. `<|split_token|>` and `<|pref_token|>` are used by the paper's pairwise trajectory preference objective. `<|reward_token|>` and `<|sim_token|>` are preserved for checkpoint compatibility.
+
+The LeRobot config stores a serialized `vlm_config` with the post-resize vocabulary so the model can reload from `config.json` without downloading the base Qwen weights first. For `Qwen/Qwen3-VL-4B-Instruct`, the tokenizer length is `151669`, and the five ROBOMETER tokens produce the checkpoint vocabulary size `151674`.
+
+### Progress Prediction
+
+In the published checkpoint, progress is discrete. The progress head outputs logits over `progress_discrete_bins=10` uniformly spaced bin centers in `[0, 1]`. LeRobot converts these logits into a continuous value by applying a softmax and taking the expectation over bin centers, matching the upstream ROBOMETER implementation.
+
+### Success Prediction
+
+The success head outputs raw logits per frame. LeRobot converts them to probabilities with `sigmoid`. When `reward_output="success"`, `compute_reward()` thresholds the last-frame success probability using `success_threshold`.
+
+## Limitations
+
+- The current LeRobot integration is inference-only; it does not implement ROBOMETER training or preference-pair training.
+- `compute_reward()` returns a scalar per sample for the LeRobot reward-model API, even though ROBOMETER predicts per-frame progress and success internally.
+- ROBOMETER is video-language based; it does not use privileged robot state such as contact forces or object poses.
+
+## References
+
+- [ROBOMETER project](https://robometer.github.io/)
+- [ROBOMETER paper](https://arxiv.org/abs/2603.02115)
+- [Original ROBOMETER code](https://github.com/robometer/robometer)
+- [Published ROBOMETER-4B checkpoint](https://huggingface.co/lerobot/Robometer-4B)
+- [Qwen3-VL-4B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct)
+
+## Citation
+
+```bibtex
+@inproceedings{liang2026robometer,
+title = {Robometer: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons},
+author={Anthony Liang and Yigit Korkmaz and Jiahui Zhang and Minyoung Hwang and Abrar Anwar and Sidhant Kaushik and Aditya Shah and Alex S. Huang and Luke Zettlemoyer and Dieter Fox and Yu Xiang and Anqi Li and Andreea Bobu and Abhishek Gupta and Stephen Tu and Erdem Biyik and Jesse Zhang},
+year={2026},
+booktitle={Robotics: Science and Systems 2026},
+}
+```
+
+## License
+
+This LeRobot integration follows the **Apache 2.0 License** used by LeRobot. Check the upstream ROBOMETER code and model pages for the licenses of the original implementation and released checkpoints.

pyproject.toml

@@ -212,6 +212,7 @@ 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]"]
+robometer = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]", "lerobot[peft-dep]"]
 topreward = ["lerobot[transformers-dep]"]
 xvla = ["lerobot[transformers-dep]"]
 eo1 = ["lerobot[transformers-dep]", "lerobot[qwen-vl-utils-dep]"]
@@ -291,6 +292,7 @@ all = [
     "lerobot[libero]; sys_platform == 'linux'",
     "lerobot[metaworld]",
     "lerobot[sarm]",
+    "lerobot[robometer]",
     "lerobot[topreward]",
     "lerobot[peft]",
     # "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2

src/lerobot/policies/factory.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 
 import importlib
 import logging
+from copy import copy
 from typing import TYPE_CHECKING, Any, TypedDict, Unpack
 
 import torch
@@ -48,7 +49,7 @@ from .act.configuration_act import ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig
 from .eo1.configuration_eo1 import EO1Config
 from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
-from .groot.configuration_groot import GrootConfig
+from .groot.configuration_groot import GROOT_N1_7, GrootConfig
 from .molmoact2.configuration_molmoact2 import MolmoAct2Config
 from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
 from .pi0.configuration_pi0 import PI0Config
@@ -273,24 +274,47 @@ def make_pre_post_processors(
             policy configuration type.
     """
     if pretrained_path:
+        if isinstance(policy_cfg, GrootConfig):
+            from .groot.configuration_groot import is_raw_groot_n1_7_checkpoint
+
+            if is_raw_groot_n1_7_checkpoint(pretrained_path):
+                from .groot.processor_groot import make_groot_pre_post_processors
+
+                processor_cfg = copy(policy_cfg)
+                processor_cfg.base_model_path = str(pretrained_path)
+                return make_groot_pre_post_processors(
+                    config=processor_cfg,
+                    dataset_stats=kwargs.get("dataset_stats"),
+                )
+
         # TODO(Steven): Temporary patch, implement correctly the processors for Gr00t
         if isinstance(policy_cfg, GrootConfig):
-            # GROOT handles normalization in groot_pack_inputs_v3 step
+            # GROOT handles normalization in its pack-inputs step
             # Need to override both stats AND normalize_min_max since saved config might be empty
-            preprocessor_overrides = {}
-            postprocessor_overrides = {}
-            preprocessor_overrides["groot_pack_inputs_v3"] = {
-                "stats": kwargs.get("dataset_stats"),
-                "normalize_min_max": True,
-            }
+            dataset_stats = kwargs.get("dataset_stats")
+            preprocessor_overrides = dict(kwargs.get("preprocessor_overrides", {}))
+            postprocessor_overrides = dict(kwargs.get("postprocessor_overrides", {}))
+            pack_inputs_key = (
+                "groot_n1_7_pack_inputs_v1"
+                if policy_cfg.model_version == GROOT_N1_7
+                else "groot_pack_inputs_v3"
+            )
+            pack_input_overrides = dict(preprocessor_overrides.get(pack_inputs_key, {}))
+            pack_input_overrides["normalize_min_max"] = True
+            if dataset_stats is not None and policy_cfg.model_version != GROOT_N1_7:
+                pack_input_overrides["stats"] = dataset_stats
+            preprocessor_overrides[pack_inputs_key] = pack_input_overrides
 
             # Also ensure postprocessing slices to env action dim and unnormalizes with dataset stats
             env_action_dim = policy_cfg.output_features[ACTION].shape[0]
-            postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = {
-                "stats": kwargs.get("dataset_stats"),
-                "normalize_min_max": True,
-                "env_action_dim": env_action_dim,
-            }
+            action_unpack_overrides = dict(
+                postprocessor_overrides.get("groot_action_unpack_unnormalize_v1", {})
+            )
+            action_unpack_overrides["normalize_min_max"] = True
+            action_unpack_overrides["env_action_dim"] = env_action_dim
+            if dataset_stats is not None and policy_cfg.model_version != GROOT_N1_7:
+                action_unpack_overrides["stats"] = dataset_stats
+            postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = action_unpack_overrides
             kwargs["preprocessor_overrides"] = preprocessor_overrides
             kwargs["postprocessor_overrides"] = postprocessor_overrides
 

src/lerobot/policies/groot/__init__.py

@@ -18,4 +18,12 @@ from .configuration_groot import GrootConfig
 from .modeling_groot import GrootPolicy
 from .processor_groot import make_groot_pre_post_processors
 
-__all__ = ["GrootConfig", "GrootPolicy", "make_groot_pre_post_processors"]
+__all__ = ["GR00TN17", "GR00TN17Config", "GrootConfig", "GrootPolicy", "make_groot_pre_post_processors"]
+
+
+def __getattr__(name: str):
+    if name in {"GR00TN17", "GR00TN17Config"}:
+        from .groot_n1_7 import GR00TN17, GR00TN17Config
+
+        return {"GR00TN17": GR00TN17, "GR00TN17Config": GR00TN17Config}[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/lerobot/policies/groot/action_head/cross_attention_dit.py

@@ -181,8 +181,7 @@ class BasicTransformerBlock(nn.Module):
         attn_output = self.attn1(
             norm_hidden_states,
             encoder_hidden_states=encoder_hidden_states,
-            attention_mask=attention_mask,
-            # encoder_attention_mask=encoder_attention_mask,
+            attention_mask=encoder_attention_mask if encoder_hidden_states is not None else attention_mask,
         )
         if self.final_dropout:
             attn_output = self.final_dropout(attn_output)
@@ -318,6 +317,71 @@ class DiT(ModelMixin, ConfigMixin):
             return self.proj_out_2(hidden_states)
 
 
+class AlternateVLDiT(DiT):
+    """N1.7 DiT variant that alternates cross-attention over image and text tokens."""
+
+    def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.attend_text_every_n_blocks = attend_text_every_n_blocks
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        return_all_hidden_states: bool = False,
+        image_mask: torch.Tensor | None = None,
+        backbone_attention_mask: torch.Tensor | None = None,
+    ):
+        if image_mask is None:
+            raise ValueError("image_mask is required for AlternateVLDiT.")
+        if backbone_attention_mask is None:
+            raise ValueError("backbone_attention_mask is required for AlternateVLDiT.")
+
+        temb = self.timestep_encoder(timestep)
+        hidden_states = hidden_states.contiguous()
+        encoder_hidden_states = encoder_hidden_states.contiguous()
+
+        image_attention_mask = image_mask & backbone_attention_mask
+        non_image_attention_mask = (~image_mask) & backbone_attention_mask
+
+        all_hidden_states = [hidden_states]
+        if not self.config.interleave_self_attention:
+            raise ValueError("AlternateVLDiT requires interleave_self_attention=True.")
+
+        for idx, block in enumerate(self.transformer_blocks):
+            if idx % 2 == 1:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=None,
+                    encoder_hidden_states=None,
+                    encoder_attention_mask=None,
+                    temb=temb,
+                )
+            else:
+                curr_encoder_attention_mask = (
+                    non_image_attention_mask
+                    if idx % (2 * self.attend_text_every_n_blocks) == 0
+                    else image_attention_mask
+                )
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=None,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=curr_encoder_attention_mask,
+                    temb=temb,
+                )
+            all_hidden_states.append(hidden_states)
+
+        conditioning = temb
+        shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+        if return_all_hidden_states:
+            return self.proj_out_2(hidden_states), all_hidden_states
+        return self.proj_out_2(hidden_states)
+
+
 class SelfAttentionTransformer(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
 

src/lerobot/policies/groot/configuration_groot.py

@@ -14,12 +14,295 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+import os
 from dataclasses import dataclass, field
+from pathlib import Path
 
 from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
 from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
 
+GROOT_N1_5 = "n1.5"
+GROOT_N1_7 = "n1.7"
+GROOT_N1_5_BASE_MODEL = "nvidia/GR00T-N1.5-3B"
+GROOT_N1_7_BASE_MODEL = "nvidia/GR00T-N1.7-3B"
+GROOT_N1_7_BACKBONE_MODEL = "nvidia/Cosmos-Reason2-2B"
+GROOT_ACTION_DECODE_TRANSFORM_LIBERO = "libero"
+
+_GROOT_MODEL_VERSION_ALIASES = {
+    "n1.5": GROOT_N1_5,
+    "n1_5": GROOT_N1_5,
+    "n15": GROOT_N1_5,
+    "1.5": GROOT_N1_5,
+    "n1.7": GROOT_N1_7,
+    "n1_7": GROOT_N1_7,
+    "n1d7": GROOT_N1_7,
+    "n17": GROOT_N1_7,
+    "1.7": GROOT_N1_7,
+}
+
+_GROOT_ACTION_DECODE_TRANSFORM_ALIASES = {
+    "none": None,
+    "": None,
+    GROOT_ACTION_DECODE_TRANSFORM_LIBERO: GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
+}
+
+
+def normalize_groot_model_version(model_version: str) -> str:
+    normalized = _GROOT_MODEL_VERSION_ALIASES.get(model_version.lower())
+    if normalized is None:
+        supported = ", ".join(sorted({GROOT_N1_5, GROOT_N1_7}))
+        raise ValueError(
+            f"Unsupported GR00T model_version '{model_version}'. Supported versions: {supported}."
+        )
+    return normalized
+
+
+def normalize_groot_action_decode_transform(transform: str | None) -> str | None:
+    if transform is None:
+        return None
+    normalized = _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.get(transform.lower())
+    if normalized is None and transform.lower() not in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES:
+        supported = ", ".join(
+            sorted(key for key, value in _GROOT_ACTION_DECODE_TRANSFORM_ALIASES.items() if value is not None)
+        )
+        raise ValueError(
+            f"Unsupported GR00T N1.7 action decode transform '{transform}'. "
+            f"Supported transforms: none, {supported}."
+        )
+    return normalized
+
+
+def infer_groot_model_version(model_path: str | None) -> str | None:
+    if not model_path:
+        return None
+    model_path_lower = model_path.lower()
+    if "gr00t-n1.7" in model_path_lower or "gr00t_n1.7" in model_path_lower:
+        return GROOT_N1_7
+    if "gr00t-n1.5" in model_path_lower or "gr00t_n1.5" in model_path_lower:
+        return GROOT_N1_5
+    config_version = _infer_groot_model_version_from_local_config(model_path)
+    if config_version is not None:
+        return config_version
+    return None
+
+
+def is_raw_groot_n1_7_checkpoint(model_path: str | Path | None) -> bool:
+    if model_path is None:
+        return False
+
+    path = Path(model_path).expanduser()
+    if path.is_dir():
+        config_path = path / "config.json"
+    elif path.name == "config.json":
+        config_path = path
+    else:
+        return False
+
+    try:
+        with config_path.open() as f:
+            config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return False
+
+    return "type" not in config and _infer_groot_model_version_from_config(config) == GROOT_N1_7
+
+
+def infer_groot_n1_7_embodiment_tag(model_path: str | Path | None) -> str | None:
+    if model_path is None:
+        return None
+
+    processor_config_path = Path(model_path).expanduser() / "processor_config.json"
+    try:
+        with processor_config_path.open() as f:
+            processor_config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+    modality_configs = processor_config.get("processor_kwargs", {}).get("modality_configs", {})
+    if not isinstance(modality_configs, dict):
+        return None
+    if "libero_sim" in modality_configs:
+        return "libero_sim"
+    if len(modality_configs) == 1:
+        return next(iter(modality_configs))
+    return None
+
+
+def infer_groot_n1_7_action_horizon(
+    model_path: str | Path | None, embodiment_tag: str | None = None
+) -> int | None:
+    if model_path is None:
+        return None
+
+    processor_config_path = Path(model_path).expanduser() / "processor_config.json"
+    try:
+        with processor_config_path.open() as f:
+            processor_config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+    processor_kwargs = processor_config.get("processor_kwargs", {})
+    if not isinstance(processor_kwargs, dict):
+        return None
+    modality_configs = processor_kwargs.get("modality_configs", {})
+    if not isinstance(modality_configs, dict):
+        return None
+
+    if embodiment_tag is None:
+        embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path)
+    if embodiment_tag is None:
+        return None
+
+    embodiment_config = modality_configs.get(embodiment_tag, {})
+    if not isinstance(embodiment_config, dict):
+        return None
+    action_config = embodiment_config.get("action", {})
+    if not isinstance(action_config, dict):
+        return None
+    delta_indices = action_config.get("delta_indices", [])
+    if not isinstance(delta_indices, list):
+        return None
+    return len(delta_indices) or None
+
+
+def infer_groot_n1_7_action_execution_horizon(
+    model_path: str | Path | None, embodiment_tag: str | None = None
+) -> int | None:
+    action_horizon = infer_groot_n1_7_action_horizon(model_path, embodiment_tag)
+    if action_horizon is None:
+        return None
+
+    if embodiment_tag is None:
+        embodiment_tag = infer_groot_n1_7_embodiment_tag(model_path)
+    if embodiment_tag == "libero_sim":
+        # NVIDIA's N1.7 LIBERO rollout wrapper replans after 8 of the 16 decoded
+        # actions. Keeping that execution cadence avoids stale open-loop chunks.
+        return min(action_horizon, 8)
+    return action_horizon
+
+
+def resolve_groot_n1_7_backbone_model(model_name: str, cache_dir: str | Path | None = None) -> str:
+    model_path = Path(model_name).expanduser()
+    if model_path.exists():
+        return str(model_path)
+
+    cached_snapshot = _find_cached_hf_snapshot(model_name, cache_dir=cache_dir)
+    return str(cached_snapshot) if cached_snapshot is not None else model_name
+
+
+def _find_cached_hf_snapshot(repo_id: str, cache_dir: str | Path | None = None) -> Path | None:
+    repo_cache_name = f"models--{repo_id.replace('/', '--')}"
+    required_files = (
+        "config.json",
+        "tokenizer_config.json",
+        "preprocessor_config.json",
+        "video_preprocessor_config.json",
+    )
+
+    for hub_cache in _candidate_hf_hub_caches(cache_dir):
+        repo_cache = hub_cache / repo_cache_name
+        snapshots_dir = repo_cache / "snapshots"
+        if not snapshots_dir.is_dir():
+            continue
+
+        candidates: list[Path] = []
+        ref_path = repo_cache / "refs" / "main"
+        try:
+            ref = ref_path.read_text().strip()
+        except OSError:
+            ref = ""
+        if ref:
+            candidates.append(snapshots_dir / ref)
+        candidates.extend(
+            sorted(
+                (path for path in snapshots_dir.iterdir() if path.is_dir()),
+                key=lambda path: path.stat().st_mtime,
+                reverse=True,
+            )
+        )
+
+        seen: set[Path] = set()
+        for snapshot in candidates:
+            if snapshot in seen:
+                continue
+            seen.add(snapshot)
+            if all((snapshot / filename).exists() for filename in required_files):
+                return snapshot
+    return None
+
+
+def _candidate_hf_hub_caches(cache_dir: str | Path | None) -> list[Path]:
+    candidates: list[Path] = []
+    if cache_dir is not None:
+        cache_path = Path(cache_dir).expanduser()
+        candidates.append(cache_path)
+        candidates.append(cache_path / "hub")
+
+    hub_cache = os.environ.get("HUGGINGFACE_HUB_CACHE")
+    if hub_cache:
+        candidates.append(Path(hub_cache).expanduser())
+
+    hf_home = os.environ.get("HF_HOME")
+    if hf_home:
+        candidates.append(Path(hf_home).expanduser() / "hub")
+
+    candidates.append(Path.home() / ".cache" / "huggingface" / "hub")
+
+    deduped: list[Path] = []
+    seen: set[Path] = set()
+    for candidate in candidates:
+        resolved = candidate.resolve() if candidate.exists() else candidate
+        if resolved not in seen:
+            seen.add(resolved)
+            deduped.append(candidate)
+    return deduped
+
+
+def _infer_groot_model_version_from_local_config(model_path: str) -> str | None:
+    path = Path(model_path).expanduser()
+    if path.is_dir():
+        config_path = path / "config.json"
+    elif path.name == "config.json":
+        config_path = path
+    else:
+        return None
+
+    if not config_path.exists():
+        return None
+
+    try:
+        with config_path.open() as f:
+            config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+    return _infer_groot_model_version_from_config(config)
+
+
+def _infer_groot_model_version_from_config(config: dict) -> str | None:
+    model_version = config.get("model_version")
+    if isinstance(model_version, str):
+        try:
+            return normalize_groot_model_version(model_version)
+        except ValueError:
+            return None
+
+    candidates = [config.get("model_type"), *(config.get("architectures") or [])]
+    for candidate in candidates:
+        if not isinstance(candidate, str):
+            continue
+        normalized = candidate.lower().replace("-", "_")
+        if normalized in {"gr00tn1d7", "gr00t_n1d7", "gr00t_n1_7"}:
+            return GROOT_N1_7
+        if normalized in {"gr00t_n1_5", "gr00tn15", "gr00t_n1d5"}:
+            return GROOT_N1_5
+
+    if config.get("model_name") == GROOT_N1_7_BACKBONE_MODEL:
+        return GROOT_N1_7
+    return None
+
 
 @PreTrainedConfig.register_subclass("groot")
 @dataclass
@@ -52,12 +335,21 @@ class GrootConfig(PreTrainedConfig):
 
     # Groot-specific model parameters (from groot_finetune_script.py)
 
+    # Explicit GR00T model family selection. Defaults to N1.5 to preserve existing behavior.
+    model_version: str = GROOT_N1_5
+
     # Path or HuggingFace model ID for the base Groot model
-    base_model_path: str = "nvidia/GR00T-N1.5-3B"
+    base_model_path: str | None = None
 
     # HF repo ID (or local path) that hosts vocab.json and merges.txt for Eagle tokenizer.
     tokenizer_assets_repo: str = "lerobot/eagle2hg-processor-groot-n1p5"
 
+    # HF repo ID (or local path) for the GR00T N1.7 Cosmos/Qwen3-VL backbone processor.
+    n1_7_backbone_model: str = GROOT_N1_7_BACKBONE_MODEL
+
+    # Optional named action transform applied after raw N1.7 checkpoint decoding and before env.step().
+    action_decode_transform: str | None = None
+
     # Embodiment tag to use for training (e.g. 'new_embodiment', 'gr1')
     embodiment_tag: str = "new_embodiment"
 
@@ -117,6 +409,35 @@ class GrootConfig(PreTrainedConfig):
     resume: bool = False
 
     def __post_init__(self):
+        self.model_version = normalize_groot_model_version(self.model_version)
+        self.action_decode_transform = normalize_groot_action_decode_transform(self.action_decode_transform)
+        if self.base_model_path is None:
+            self.base_model_path = (
+                GROOT_N1_7_BASE_MODEL if self.model_version == GROOT_N1_7 else GROOT_N1_5_BASE_MODEL
+            )
+
+        if self.action_decode_transform is not None and self.model_version != GROOT_N1_7:
+            raise ValueError("action_decode_transform can only be used with model_version='n1.7'.")
+
+        if self.model_version == GROOT_N1_7:
+            if self.max_state_dim == 64:
+                self.max_state_dim = 132
+            if self.max_action_dim == 32:
+                self.max_action_dim = 132
+            if self.chunk_size == 50:
+                self.chunk_size = 40
+            if self.n_action_steps == 50:
+                self.n_action_steps = 40
+            if tuple(self.image_size) == (224, 224):
+                self.image_size = (256, 256)
+
+        inferred_version = infer_groot_model_version(self.base_model_path)
+        if inferred_version is not None and inferred_version != self.model_version:
+            raise ValueError(
+                f"GR00T model_version '{self.model_version}' does not match base_model_path "
+                f"'{self.base_model_path}', which looks like '{inferred_version}'."
+            )
+
         super().__post_init__()
 
         if self.n_action_steps > self.chunk_size:
@@ -192,7 +513,12 @@ class GrootConfig(PreTrainedConfig):
     @property
     def action_delta_indices(self) -> list[int]:
         """Return indices for delta actions."""
-        return list(range(min(self.chunk_size, 16)))
+        model_action_horizon = 16
+        if self.model_version == GROOT_N1_7:
+            model_action_horizon = (
+                infer_groot_n1_7_action_horizon(self.base_model_path, self.embodiment_tag) or 40
+            )
+        return list(range(min(self.chunk_size, model_action_horizon)))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/groot/groot_n1_7.py

@@ -0,0 +1,962 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import importlib
+import json
+import logging
+from contextlib import suppress
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from huggingface_hub import snapshot_download
+from huggingface_hub.errors import HFValidationError, RepositoryNotFoundError
+from torch import nn
+from torch.distributions import Beta
+
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+from .action_head.cross_attention_dit import AlternateVLDiT, DiT, SelfAttentionTransformer
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+    from transformers.feature_extraction_utils import BatchFeature
+else:
+    AutoConfig = None
+    AutoModel = None
+    PretrainedConfig = object
+    PreTrainedModel = object
+    BatchFeature = None
+
+try:
+    import tree
+except ImportError:
+    tree = None
+
+try:
+    from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration
+except ImportError:
+    Qwen3VLConfig = None
+    Qwen3VLForConditionalGeneration = None
+
+logger = logging.getLogger(__name__)
+
+
+def _copy_default(value: Any) -> Any:
+    return deepcopy(value)
+
+
+GR00T_N1_7_DEFAULTS: dict[str, Any] = {
+    "model_dtype": "bfloat16",
+    "dtype": "bfloat16",
+    "model_name": "nvidia/Cosmos-Reason2-2B",
+    "backbone_model_type": "qwen",
+    "model_revision": None,
+    "tune_top_llm_layers": 0,
+    "backbone_embedding_dim": 2048,
+    "tune_llm": False,
+    "tune_visual": False,
+    "select_layer": 12,
+    "reproject_vision": False,
+    "use_flash_attention": True,
+    "load_bf16": False,
+    "backbone_trainable_params_fp32": True,
+    "image_crop_size": (230, 230),
+    "image_target_size": (256, 256),
+    "shortest_image_edge": None,
+    "crop_fraction": None,
+    "random_rotation_angle": None,
+    "color_jitter_params": None,
+    "use_albumentations_transforms": True,
+    "extra_augmentation_config": None,
+    "formalize_language": True,
+    "apply_sincos_state_encoding": False,
+    "use_percentiles": True,
+    "use_relative_action": False,
+    "max_state_dim": 132,
+    "max_action_dim": 132,
+    "action_horizon": 40,
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "state_history_length": 1,
+    "add_pos_embed": True,
+    "attn_dropout": 0.2,
+    "use_vlln": True,
+    "max_seq_len": 1024,
+    "use_alternate_vl_dit": True,
+    "attend_text_every_n_blocks": 2,
+    "diffusion_model_cfg": {
+        "positional_embeddings": None,
+        "num_layers": 32,
+        "num_attention_heads": 32,
+        "attention_head_dim": 48,
+        "norm_type": "ada_norm",
+        "dropout": 0.2,
+        "final_dropout": True,
+        "output_dim": 1024,
+        "interleave_self_attention": True,
+    },
+    "vl_self_attention_cfg": {
+        "positional_embeddings": None,
+        "num_layers": 4,
+        "num_attention_heads": 32,
+        "attention_head_dim": 64,
+        "dropout": 0.2,
+        "final_dropout": True,
+    },
+    "num_inference_timesteps": 4,
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_timestep_buckets": 1000,
+    "tune_projector": True,
+    "tune_diffusion_model": True,
+    "tune_vlln": True,
+    "state_dropout_prob": 0.2,
+    "exclude_state": False,
+    "use_mean_std": False,
+    "max_num_embodiments": 32,
+    "rtc_ramp_rate": 6.0,
+}
+
+
+class GR00TN17Config(PretrainedConfig):
+    """Configuration for NVIDIA GR00T N1.7.
+
+    N1.7 uses the Cosmos-Reason2-2B / Qwen3-VL backbone and a multi-embodiment
+    flow-matching action head. This mirrors the public N1.7 checkpoint config
+    while keeping it local to LeRobot and independent from the external
+    Isaac-GR00T ``gr00t`` Python package.
+    """
+
+    model_type = "Gr00tN1d7"
+
+    _defaults = GR00T_N1_7_DEFAULTS
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        for key, value in GR00T_N1_7_DEFAULTS.items():
+            setattr(self, key, _copy_default(kwargs.pop(key, value)))
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def to_filtered_dict(self, exclude_augment: bool = True) -> dict[str, Any]:
+        cfg = self.to_dict()
+        if not exclude_augment:
+            return cfg
+        exclude_keys = {
+            "random_rotation_angle",
+            "color_jitter_params",
+            "use_albumentations_transforms",
+            "formalize_language",
+            "image_crop_size",
+            "image_target_size",
+            "shortest_image_edge",
+            "crop_fraction",
+        }
+        return {k: v for k, v in cfg.items() if k not in exclude_keys}
+
+    def to_filtered_json(self, exclude_augment: bool = True, **kwargs) -> str:
+        return json.dumps(self.to_filtered_dict(exclude_augment), indent=2, default=str, **kwargs)
+
+
+class CategorySpecificLinear(nn.Module):
+    """Linear layer with category-specific weights for multi-embodiment support."""
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int):
+        super().__init__()
+        self.num_categories = num_categories
+        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
+        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        selected_w = self.W[cat_ids]
+        selected_b = self.b[cat_ids]
+        return torch.bmm(x, selected_w) + selected_b.unsqueeze(1)
+
+
+class CategorySpecificMLP(nn.Module):
+    """Two-layer MLP with category-specific weights."""
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
+        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        hidden = F.relu(self.layer1(x, cat_ids))
+        return self.layer2(hidden, cat_ids)
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """Sinusoidal encoding of shape ``(B, T, D)`` for timestep tensors ``(B, T)``.
+
+    The frequency scalar is intentionally created on CPU and then broadcast with
+    the device-local arange result. That mirrors Isaac-GR00T's N1.7 timestep
+    embedding and avoids tiny dtype/device construction differences in parity
+    tests.
+    """
+
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        timesteps = timesteps.float()
+        half_dim = self.embedding_dim // 2
+        exponent = -torch.arange(half_dim, dtype=torch.float, device=timesteps.device) * (
+            torch.log(torch.tensor(10000.0)) / half_dim
+        )
+        freqs = timesteps.unsqueeze(-1) * exponent.exp()
+        return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)
+
+
+def swish(x: torch.Tensor) -> torch.Tensor:
+    return x * torch.sigmoid(x)
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    """Action encoder with category-specific projections and sinusoidal time encoding."""
+
+    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int):
+        super().__init__()
+        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
+        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
+        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
+
+    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        batch_size, horizon, _ = actions.shape
+        if timesteps.dim() != 1 or timesteps.shape[0] != batch_size:
+            raise ValueError("Expected `timesteps` to have shape (B,).")
+        timesteps = timesteps.unsqueeze(1).expand(-1, horizon)
+        action_emb = self.W1(actions, cat_ids)
+        time_emb = self.pos_encoding(timesteps).to(dtype=action_emb.dtype)
+        x = swish(self.W2(torch.cat([action_emb, time_emb], dim=-1), cat_ids))
+        return self.W3(x, cat_ids)
+
+
+class Qwen3Backbone(nn.Module):
+    """Cosmos-Reason2/Qwen3-VL backbone used by GR00T N1.7.
+
+    The public checkpoint stores the action head in the GR00T checkpoint but
+    uses a Hugging Face Qwen3-VL-compatible backbone interface. This wrapper
+    keeps the nested HF module layout compatible across transformer versions
+    and exposes the hidden states consumed by the action head.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "nvidia/Cosmos-Reason2-2B",
+        tune_llm: bool = False,
+        tune_visual: bool = False,
+        select_layer: int = -1,
+        reproject_vision: bool = False,
+        use_flash_attention: bool = False,
+        load_bf16: bool = False,
+        tune_top_llm_layers: int = 0,
+        trainable_params_fp32: bool = False,
+        transformers_loading_kwargs: dict[str, Any] | None = None,
+        load_pretrained_weights: bool = True,
+    ):
+        if Qwen3VLForConditionalGeneration is None:
+            raise ImportError(
+                "Qwen3VLForConditionalGeneration is required for GR00T N1.7. "
+                "Install the GR00T optional dependencies with `pip install 'lerobot[groot]'` "
+                "or use a transformers version that provides Qwen3-VL support."
+            )
+
+        super().__init__()
+        transformers_loading_kwargs = transformers_loading_kwargs or {"trust_remote_code": True}
+
+        extra_kwargs: dict[str, Any] = {}
+        if use_flash_attention:
+            try:
+                import flash_attn  # noqa: F401
+
+                extra_kwargs["attn_implementation"] = "flash_attention_2"
+            except ImportError:
+                logger.warning("flash_attn is not installed. Falling back to SDPA attention.")
+                extra_kwargs["attn_implementation"] = "sdpa"
+        if load_bf16:
+            extra_kwargs["torch_dtype"] = torch.bfloat16
+
+        if load_pretrained_weights:
+            self.model = Qwen3VLForConditionalGeneration.from_pretrained(
+                model_name,
+                **extra_kwargs,
+                **transformers_loading_kwargs,
+            ).eval()
+        else:
+            self.model = self._from_backbone_config(
+                model_name=model_name,
+                model_kwargs=extra_kwargs,
+                config_kwargs=transformers_loading_kwargs,
+            ).eval()
+
+        while len(self.language_model.layers) > select_layer:
+            self.language_model.layers.pop(-1)
+
+        self.select_layer = select_layer
+        self.set_trainable_parameters(tune_llm, tune_visual, tune_top_llm_layers)
+        if load_bf16 and trainable_params_fp32:
+            for parameter in self.parameters():
+                if parameter.requires_grad:
+                    parameter.data = parameter.data.to(torch.float32)
+
+    def set_trainable_parameters(
+        self, tune_llm: bool, tune_visual: bool, tune_top_llm_layers: int = 0
+    ) -> None:
+        self.tune_llm = tune_llm
+        self.tune_visual = tune_visual
+        for parameter in self.parameters():
+            parameter.requires_grad = True
+        if not tune_llm:
+            self.language_model.requires_grad_(False)
+        if not tune_visual:
+            self.visual.requires_grad_(False)
+        if tune_top_llm_layers > 0:
+            for layer in self.language_model.layers[-tune_top_llm_layers:]:
+                for parameter in layer.parameters():
+                    parameter.requires_grad = True
+
+    def set_frozen_modules_to_eval_mode(self) -> None:
+        if self.training:
+            if self.language_model and not self.tune_llm:
+                self.language_model.eval()
+            if self.visual and not self.tune_visual:
+                self.visual.eval()
+
+    @property
+    def language_model(self) -> nn.Module:
+        return getattr(self.model, "model", self.model).language_model
+
+    @property
+    def visual(self) -> nn.Module:
+        return getattr(self.model, "model", self.model).visual
+
+    def _from_backbone_config(
+        self,
+        *,
+        model_name: str,
+        model_kwargs: dict[str, Any],
+        config_kwargs: dict[str, Any],
+    ) -> nn.Module:
+        if _is_cosmos_reason2_backbone(model_name):
+            backbone_config = _cosmos_reason2_qwen3_vl_config()
+        else:
+            if AutoConfig is None:
+                raise ImportError(
+                    "AutoConfig is required to initialize a GR00T N1.7 backbone from config. "
+                    "Install the GR00T optional dependencies with `pip install 'lerobot[groot]'`."
+                )
+            backbone_config = AutoConfig.from_pretrained(model_name, **config_kwargs)
+        return Qwen3VLForConditionalGeneration._from_config(backbone_config, **model_kwargs)
+
+    def prepare_input(self, batch: dict[str, Any]) -> BatchFeature:
+        return BatchFeature(data=batch)
+
+    def _ensure_mm_token_type_ids(self, model_input: dict[str, torch.Tensor]) -> None:
+        if "mm_token_type_ids" in model_input:
+            return
+        if "image_grid_thw" not in model_input and "video_grid_thw" not in model_input:
+            return
+
+        input_ids = model_input.get("input_ids")
+        if input_ids is None:
+            return
+
+        mm_token_type_ids = torch.zeros(input_ids.shape, dtype=torch.int32, device=input_ids.device)
+        image_token_id = getattr(self.model.config, "image_token_id", None)
+        video_token_id = getattr(self.model.config, "video_token_id", None)
+        if image_token_id is not None:
+            mm_token_type_ids[input_ids == image_token_id] = 1
+        if video_token_id is not None:
+            mm_token_type_ids[input_ids == video_token_id] = 2
+
+        model_input["mm_token_type_ids"] = mm_token_type_ids
+
+    def _ensure_legacy_qwen3_position_ids(self, model_input: dict[str, torch.Tensor]) -> None:
+        """Restore the Qwen3-VL text position ids used by older Transformers releases.
+
+        Transformers 5.x computes 3-row multimodal RoPE ids for Qwen3-VL and then
+        drops text position ids before calling text-layer flash attention. GR00T
+        N1.7 was aligned against the older Transformers path, where a fourth text
+        position row is forwarded alongside the temporal/height/width rows. Adding
+        the row here preserves the newer multimodal position computation while
+        keeping flash attention on the legacy code path.
+        """
+
+        if "position_ids" in model_input:
+            return
+
+        qwen3_model = getattr(self.model, "model", self.model)
+        compute_3d_position_ids = getattr(qwen3_model, "compute_3d_position_ids", None)
+        if compute_3d_position_ids is None:
+            return
+
+        position_ids = compute_3d_position_ids(
+            input_ids=model_input.get("input_ids"),
+            image_grid_thw=model_input.get("image_grid_thw"),
+            video_grid_thw=model_input.get("video_grid_thw"),
+            inputs_embeds=None,
+            attention_mask=model_input.get("attention_mask"),
+            past_key_values=None,
+            mm_token_type_ids=model_input.get("mm_token_type_ids"),
+        )
+        if position_ids.ndim == 3 and position_ids.shape[0] == 3:
+            position_ids = torch.cat([position_ids[:1], position_ids], dim=0)
+
+        model_input["position_ids"] = position_ids
+
+    def _last_decoder_layer_output(self, model_input: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Return the pre-final-norm decoder output consumed by the N1.7 action head.
+
+        Older Transformers releases exposed this tensor as ``hidden_states[-1]``.
+        Newer releases expose the post-final-norm tensor there instead. Capturing
+        the last decoder layer output directly keeps the N1.7 action head input
+        stable across Transformers versions.
+        """
+
+        captured: dict[str, torch.Tensor] = {}
+
+        def capture_output(_module: nn.Module, _inputs: tuple[Any, ...], output: Any) -> None:
+            if isinstance(output, torch.Tensor):
+                captured["features"] = output
+            elif isinstance(output, (tuple, list)) and output:
+                captured["features"] = output[0]
+            elif hasattr(output, "last_hidden_state"):
+                captured["features"] = output.last_hidden_state
+
+        hook = self.language_model.layers[-1].register_forward_hook(capture_output)
+        try:
+            outputs = self.model(**model_input, output_hidden_states=True)
+        finally:
+            hook.remove()
+
+        return captured.get("features", outputs.hidden_states[-1])
+
+    def forward(self, vl_input: BatchFeature) -> BatchFeature:
+        self.set_frozen_modules_to_eval_mode()
+        keys_to_use = ["input_ids", "attention_mask", "pixel_values", "image_grid_thw"]
+        optional_keys = ["mm_token_type_ids", "pixel_values_videos", "video_grid_thw"]
+        model_input = {key: vl_input[key] for key in keys_to_use}
+        model_input.update({key: vl_input[key] for key in optional_keys if key in vl_input})
+        self._ensure_mm_token_type_ids(model_input)
+        self._ensure_legacy_qwen3_position_ids(model_input)
+        features = self._last_decoder_layer_output(model_input)
+        image_mask = model_input["input_ids"] == self.model.config.image_token_id
+        attention_mask = model_input["attention_mask"] == 1
+        return BatchFeature(
+            data={
+                "backbone_features": features,
+                "backbone_attention_mask": attention_mask,
+                "image_mask": image_mask,
+            }
+        )
+
+
+class GR00TN17ActionHead(nn.Module):
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: GR00TN17Config):
+        require_package("diffusers", extra="groot")
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.input_embedding_dim = config.input_embedding_dim
+
+        if config.use_alternate_vl_dit:
+            self.model = AlternateVLDiT(
+                **config.diffusion_model_cfg,
+                cross_attention_dim=config.backbone_embedding_dim,
+                attend_text_every_n_blocks=config.attend_text_every_n_blocks,
+            )
+        else:
+            self.model = DiT(
+                **config.diffusion_model_cfg,
+                cross_attention_dim=config.backbone_embedding_dim,
+            )
+
+        self.action_dim = config.max_action_dim
+        self.action_horizon = config.action_horizon
+        self.num_inference_timesteps = config.num_inference_timesteps
+        self.state_encoder = CategorySpecificMLP(
+            num_categories=config.max_num_embodiments,
+            input_dim=config.max_state_dim * config.state_history_length,
+            hidden_dim=self.hidden_size,
+            output_dim=self.input_embedding_dim,
+        )
+        self.action_encoder = MultiEmbodimentActionEncoder(
+            action_dim=self.action_dim,
+            hidden_size=self.input_embedding_dim,
+            num_embodiments=config.max_num_embodiments,
+        )
+        self.action_decoder = CategorySpecificMLP(
+            num_categories=config.max_num_embodiments,
+            input_dim=self.hidden_size,
+            hidden_dim=self.hidden_size,
+            output_dim=self.action_dim,
+        )
+        self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity()
+        vl_self_attention_cfg = getattr(config, "vl_self_attention_cfg", None)
+        if vl_self_attention_cfg and vl_self_attention_cfg.get("num_layers", 0) > 0:
+            self.vl_self_attention = SelfAttentionTransformer(**vl_self_attention_cfg)
+        else:
+            self.vl_self_attention = nn.Identity()
+        if config.add_pos_embed:
+            self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim)
+            nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
+        self.state_dropout_prob = config.state_dropout_prob
+        self._noise_beta_alpha = config.noise_beta_alpha
+        self._noise_beta_beta = config.noise_beta_beta
+        self._beta_dist = None
+        self.num_timestep_buckets = config.num_timestep_buckets
+        self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model, config.tune_vlln)
+
+    def set_trainable_parameters(
+        self, tune_projector: bool, tune_diffusion_model: bool, tune_vlln: bool
+    ) -> None:
+        self.tune_projector = tune_projector
+        self.tune_diffusion_model = tune_diffusion_model
+        self.tune_vlln = tune_vlln
+        for parameter in self.parameters():
+            parameter.requires_grad = True
+        if not tune_projector:
+            self.state_encoder.requires_grad_(False)
+            self.action_encoder.requires_grad_(False)
+            self.action_decoder.requires_grad_(False)
+            if self.config.add_pos_embed:
+                self.position_embedding.requires_grad_(False)
+        if not tune_diffusion_model:
+            self.model.requires_grad_(False)
+        if not tune_vlln:
+            self.vlln.requires_grad_(False)
+            self.vl_self_attention.requires_grad_(False)
+
+    def set_frozen_modules_to_eval_mode(self) -> None:
+        if self.training:
+            if not self.tune_projector:
+                self.state_encoder.eval()
+                self.action_encoder.eval()
+                self.action_decoder.eval()
+                if self.config.add_pos_embed:
+                    self.position_embedding.eval()
+            if not self.tune_diffusion_model:
+                self.model.eval()
+            if not self.tune_vlln:
+                self.vlln.eval()
+                self.vl_self_attention.eval()
+
+    def sample_time(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        if self._beta_dist is None:
+            beta_alpha = torch.tensor(self._noise_beta_alpha, device="cpu", dtype=torch.float32)
+            beta_beta = torch.tensor(self._noise_beta_beta, device="cpu", dtype=torch.float32)
+            self._beta_dist = Beta(beta_alpha, beta_beta, validate_args=False)
+        sample = self._beta_dist.sample([batch_size]).to(device, dtype=dtype)
+        return (1 - sample) * self.config.noise_s
+
+    def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature:
+        backbone_features = self.vlln(backbone_output["backbone_features"])
+        backbone_output["backbone_features"] = self.vl_self_attention(backbone_features)
+        return backbone_output
+
+    def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
+        self.set_frozen_modules_to_eval_mode()
+        backbone_output = self.process_backbone_output(backbone_output)
+        vl_embeds = backbone_output.backbone_features
+        device = vl_embeds.device
+        embodiment_id = action_input.embodiment_id
+
+        if action_input.state.shape[1] != self.config.state_history_length:
+            raise ValueError("state history length does not match GR00T N1.7 config.")
+        state = action_input.state.view(action_input.state.shape[0], 1, -1)
+        state_features = self.state_encoder(state, embodiment_id)
+
+        if self.training and self.state_dropout_prob > 0:
+            do_dropout = (
+                torch.rand(state_features.shape[0], device=state_features.device) < self.state_dropout_prob
+            )
+            state_features = state_features * (1 - do_dropout[:, None, None].to(dtype=state_features.dtype))
+
+        actions = action_input.action
+        noise = torch.randn(actions.shape, device=actions.device, dtype=actions.dtype)
+        t = self.sample_time(actions.shape[0], device=actions.device, dtype=actions.dtype)
+        t = t[:, None, None]
+        noisy_trajectory = (1 - t) * noise + t * actions
+        velocity = actions - noise
+        t_discretized = (t[:, 0, 0] * self.num_timestep_buckets).long()
+        action_features = self.action_encoder(noisy_trajectory, t_discretized, embodiment_id)
+
+        if self.config.add_pos_embed:
+            pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
+            action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0)
+
+        sa_embs = torch.cat((state_features, action_features), dim=1)
+        if self.config.use_alternate_vl_dit:
+            model_output, _ = self.model(
+                hidden_states=sa_embs,
+                encoder_hidden_states=vl_embeds,
+                encoder_attention_mask=backbone_output.backbone_attention_mask,
+                timestep=t_discretized,
+                return_all_hidden_states=True,
+                image_mask=backbone_output.image_mask,
+                backbone_attention_mask=backbone_output.backbone_attention_mask,
+            )
+        else:
+            model_output, _ = self.model(
+                hidden_states=sa_embs,
+                encoder_hidden_states=vl_embeds,
+                encoder_attention_mask=backbone_output.backbone_attention_mask,
+                timestep=t_discretized,
+                return_all_hidden_states=True,
+            )
+
+        pred = self.action_decoder(model_output, embodiment_id)
+        pred_actions = pred[:, -actions.shape[1] :]
+        action_mask = action_input.action_mask.to(dtype=pred_actions.dtype)
+        action_loss = F.mse_loss(pred_actions, velocity, reduction="none") * action_mask
+        loss = action_loss.sum() / (action_mask.sum() + 1e-6)
+        return BatchFeature(
+            data={
+                "loss": loss,
+                "action_loss": action_loss,
+                "action_mask": action_mask,
+                "backbone_features": vl_embeds,
+                "state_features": state_features,
+            }
+        )
+
+    def _encode_features(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature:
+        backbone_output = self.process_backbone_output(backbone_output)
+        state = action_input.state
+        if state.shape[1] != self.config.state_history_length:
+            raise ValueError("state history length does not match GR00T N1.7 config.")
+        state = state.view(state.shape[0], 1, -1)
+        state_features = self.state_encoder(state, action_input.embodiment_id)
+        return BatchFeature(
+            data={"backbone_features": backbone_output.backbone_features, "state_features": state_features}
+        )
+
+    @torch.no_grad()
+    def get_action_with_features(
+        self,
+        backbone_features: torch.Tensor,
+        state_features: torch.Tensor,
+        embodiment_id: torch.Tensor,
+        backbone_output: BatchFeature,
+        action_input: BatchFeature,
+        options: dict[str, Any] | None = None,
+    ) -> BatchFeature:
+        vl_embeds = backbone_features
+        batch_size = vl_embeds.shape[0]
+        device = vl_embeds.device
+        actions = torch.randn(
+            size=(batch_size, self.config.action_horizon, self.action_dim),
+            dtype=vl_embeds.dtype,
+            device=device,
+        )
+        dt = 1.0 / self.num_inference_timesteps
+        vel_strength = torch.ones_like(actions)
+
+        if "action" in action_input:
+            if options is None:
+                raise ValueError("RTC options are required when action is provided to get_action.")
+            action_horizon_before_padding = options["action_horizon"]
+            actions[:, : options["rtc_overlap_steps"], :] = action_input["action"][
+                :,
+                action_horizon_before_padding - options["rtc_overlap_steps"] : action_horizon_before_padding,
+                :,
+            ]
+            vel_strength[:, : options["rtc_frozen_steps"], :] = 0.0
+            intermediate_steps = options["rtc_overlap_steps"] - options["rtc_frozen_steps"]
+            t = torch.linspace(0.0, 1.0, intermediate_steps + 2, device=device)
+            ramp = 1 - torch.exp(-options["rtc_ramp_rate"] * t)
+            ramp = ramp / ramp[-1].clamp_min(1e-8)
+            vel_strength[:, options["rtc_frozen_steps"] : options["rtc_overlap_steps"], :] = ramp[1:-1][
+                None, :, None
+            ].to(device)
+
+        for t_step in range(self.num_inference_timesteps):
+            t_cont = t_step / float(self.num_inference_timesteps)
+            t_discretized = int(t_cont * self.num_timestep_buckets)
+            timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device)
+            action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id)
+            if self.config.add_pos_embed:
+                pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device)
+                action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0)
+            sa_embs = torch.cat((state_features, action_features), dim=1)
+
+            if self.config.use_alternate_vl_dit:
+                model_output = self.model(
+                    hidden_states=sa_embs,
+                    encoder_hidden_states=vl_embeds,
+                    timestep=timesteps_tensor,
+                    image_mask=backbone_output.image_mask,
+                    backbone_attention_mask=backbone_output.backbone_attention_mask,
+                )
+            else:
+                model_output = self.model(
+                    hidden_states=sa_embs,
+                    encoder_hidden_states=vl_embeds,
+                    timestep=timesteps_tensor,
+                )
+            pred = self.action_decoder(model_output, embodiment_id)
+            actions = actions + dt * pred[:, -self.action_horizon :] * vel_strength
+
+        return BatchFeature(
+            data={
+                "action_pred": actions,
+                "backbone_features": vl_embeds,
+                "state_features": state_features,
+            }
+        )
+
+    @torch.no_grad()
+    def get_action(
+        self,
+        backbone_output: BatchFeature,
+        action_input: BatchFeature,
+        options: dict[str, Any] | None = None,
+    ) -> BatchFeature:
+        features = self._encode_features(backbone_output, action_input)
+        return self.get_action_with_features(
+            backbone_features=features.backbone_features,
+            state_features=features.state_features,
+            embodiment_id=action_input.embodiment_id,
+            backbone_output=backbone_output,
+            action_input=action_input,
+            options=options,
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(iter(self.parameters())).dtype
+
+    def prepare_input(self, batch: dict[str, Any]) -> BatchFeature:
+        return BatchFeature(data=batch)
+
+
+def _is_cosmos_reason2_backbone(model_name: str) -> bool:
+    return str(model_name).rstrip("/") == "nvidia/Cosmos-Reason2-2B"
+
+
+def _cosmos_reason2_qwen3_vl_config() -> PretrainedConfig:
+    if Qwen3VLConfig is None:
+        raise ImportError(
+            "Qwen3VLConfig is required for GR00T N1.7. "
+            "Install the GR00T optional dependencies with `pip install 'lerobot[groot]'`."
+        )
+    return Qwen3VLConfig(
+        image_token_id=151655,
+        video_token_id=151656,
+        vision_start_token_id=151652,
+        vision_end_token_id=151653,
+        tie_word_embeddings=True,
+        text_config={
+            "attention_bias": False,
+            "attention_dropout": 0.0,
+            "bos_token_id": 151643,
+            "dtype": "bfloat16",
+            "eos_token_id": 151645,
+            "head_dim": 128,
+            "hidden_act": "silu",
+            "hidden_size": 2048,
+            "initializer_range": 0.02,
+            "intermediate_size": 6144,
+            "max_position_embeddings": 262144,
+            "model_type": "qwen3_vl_text",
+            "num_attention_heads": 16,
+            "num_hidden_layers": 28,
+            "num_key_value_heads": 8,
+            "rms_norm_eps": 1e-6,
+            "rope_scaling": {
+                "mrope_interleaved": True,
+                "mrope_section": [24, 20, 20],
+                "rope_type": "default",
+            },
+            "rope_theta": 5000000,
+            "tie_word_embeddings": True,
+            "use_cache": True,
+            "vocab_size": 151936,
+        },
+        vision_config={
+            "deepstack_visual_indexes": [5, 11, 17],
+            "depth": 24,
+            "hidden_act": "gelu_pytorch_tanh",
+            "hidden_size": 1024,
+            "in_channels": 3,
+            "initializer_range": 0.02,
+            "intermediate_size": 4096,
+            "model_type": "qwen3_vl",
+            "num_heads": 16,
+            "num_position_embeddings": 2304,
+            "out_hidden_size": 2048,
+            "patch_size": 16,
+            "spatial_merge_size": 2,
+            "temporal_patch_size": 2,
+        },
+    )
+
+
+def get_backbone_cls(config: GR00TN17Config):
+    if (
+        config.backbone_model_type == "qwen"
+        or "nvidia/Cosmos-Reason2" in config.model_name
+        or "Qwen/Qwen3-VL" in config.model_name
+    ):
+        return Qwen3Backbone
+    raise ValueError(f"Unsupported GR00T N1.7 backbone model: {config.model_name}")
+
+
+class GR00TN17(PreTrainedModel):
+    """GR00T N1.7 model with a Cosmos-Reason2/Qwen3-VL backbone."""
+
+    config_class = GR00TN17Config
+    supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        config: GR00TN17Config,
+        transformers_loading_kwargs: dict[str, Any] | None = None,
+        load_backbone_weights: bool = True,
+    ):
+        super().__init__(config)
+        transformers_loading_kwargs = transformers_loading_kwargs or {"trust_remote_code": True}
+        self.config = config
+        backbone_cls = get_backbone_cls(config)
+        self.backbone = backbone_cls(
+            model_name=config.model_name,
+            tune_llm=config.tune_llm,
+            tune_visual=config.tune_visual,
+            select_layer=config.select_layer,
+            reproject_vision=config.reproject_vision,
+            use_flash_attention=config.use_flash_attention,
+            load_bf16=config.load_bf16,
+            tune_top_llm_layers=config.tune_top_llm_layers,
+            trainable_params_fp32=config.backbone_trainable_params_fp32,
+            transformers_loading_kwargs=transformers_loading_kwargs,
+            load_pretrained_weights=load_backbone_weights,
+        )
+        self.action_head = GR00TN17ActionHead(config)
+        self.post_init()
+
+    def prepare_input(self, inputs: dict[str, Any]) -> tuple[BatchFeature, BatchFeature]:
+        global tree
+        if tree is None:
+            require_package("dm-tree", extra="groot", import_name="tree")
+            tree = importlib.import_module("tree")
+        backbone_inputs = self.backbone.prepare_input(inputs)
+        action_inputs = self.action_head.prepare_input(inputs)
+
+        def to_device_with_dtype(x):
+            if not isinstance(x, torch.Tensor):
+                return x
+            if torch.is_floating_point(x):
+                return x.to(self.device, dtype=self.dtype)
+            return x.to(self.device)
+
+        return (
+            tree.map_structure(to_device_with_dtype, backbone_inputs),
+            tree.map_structure(to_device_with_dtype, action_inputs),
+        )
+
+    def forward(self, inputs: dict[str, Any]) -> BatchFeature:
+        backbone_inputs, action_inputs = self.prepare_input(inputs)
+        backbone_outputs = self.backbone(backbone_inputs)
+        return self.action_head(backbone_outputs, action_inputs)
+
+    def get_action(self, inputs: dict[str, Any], options: dict[str, Any] | None = None) -> BatchFeature:
+        backbone_inputs, action_inputs = self.prepare_input(inputs)
+        backbone_outputs = self.backbone(backbone_inputs)
+        return self.action_head.get_action(backbone_outputs, action_inputs, options)
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.parameters())).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(iter(self.parameters())).dtype
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        tune_visual = kwargs.pop("tune_visual", True)
+        tune_llm = kwargs.pop("tune_llm", False)
+        tune_projector = kwargs.pop("tune_projector", True)
+        tune_diffusion_model = kwargs.pop("tune_diffusion_model", True)
+        tune_vlln = kwargs.pop("tune_vlln", True)
+        transformers_loading_kwargs = kwargs.pop("transformers_loading_kwargs", None) or {
+            "trust_remote_code": True
+        }
+        load_backbone_weights = kwargs.pop("load_backbone_weights", False)
+        for key in ("revision", "cache_dir", "local_files_only", "token"):
+            if key in kwargs:
+                transformers_loading_kwargs.setdefault(key, kwargs[key])
+
+        try:
+            local_model_path = snapshot_download(
+                pretrained_model_name_or_path,
+                repo_type="model",
+                revision=kwargs.get("revision"),
+                cache_dir=kwargs.get("cache_dir"),
+                local_files_only=kwargs.get("local_files_only", False),
+                token=kwargs.get("token"),
+            )
+        except (HFValidationError, RepositoryNotFoundError):
+            local_model_path = pretrained_model_name_or_path
+
+        pretrained_model = super().from_pretrained(
+            local_model_path,
+            transformers_loading_kwargs=transformers_loading_kwargs,
+            load_backbone_weights=load_backbone_weights,
+            **kwargs,
+        )
+        pretrained_model.backbone.set_trainable_parameters(
+            tune_visual=tune_visual,
+            tune_llm=tune_llm,
+            tune_top_llm_layers=pretrained_model.config.tune_top_llm_layers,
+        )
+        pretrained_model.action_head.set_trainable_parameters(
+            tune_projector=tune_projector,
+            tune_diffusion_model=tune_diffusion_model,
+            tune_vlln=tune_vlln,
+        )
+        return pretrained_model
+
+
+def _register_with_transformers() -> None:
+    if AutoConfig is None or AutoModel is None:
+        return
+    try:
+        AutoConfig.register(GR00TN17Config.model_type, GR00TN17Config, exist_ok=True)
+    except TypeError:
+        with suppress(ValueError):
+            AutoConfig.register(GR00TN17Config.model_type, GR00TN17Config)
+    try:
+        AutoModel.register(GR00TN17Config, GR00TN17, exist_ok=True)
+    except TypeError:
+        with suppress(ValueError):
+            AutoModel.register(GR00TN17Config, GR00TN17)
+
+
+_register_with_transformers()

src/lerobot/policies/groot/modeling_groot.py

@@ -46,7 +46,15 @@ from lerobot.utils.constants import ACTION, OBS_IMAGES
 from lerobot.utils.import_utils import require_package
 
 from ..pretrained import PreTrainedPolicy
-from .configuration_groot import GrootConfig
+from .configuration_groot import (
+    GROOT_N1_5,
+    GROOT_N1_7,
+    GrootConfig,
+    infer_groot_model_version,
+    infer_groot_n1_7_action_execution_horizon,
+    infer_groot_n1_7_action_horizon,
+    normalize_groot_model_version,
+)
 from .groot_n1 import GR00TN15
 
 T = TypeVar("T", bound="GrootPolicy")
@@ -67,6 +75,7 @@ class GrootPolicy(PreTrainedPolicy):
 
         # Initialize GR00T model using ported components
         self._groot_model = self._create_groot_model()
+        self._action_queue_steps = self._resolve_action_queue_steps()
 
         self.reset()
 
@@ -82,13 +91,23 @@ class GrootPolicy(PreTrainedPolicy):
         # Handle Flash Attention compatibility issues
         self._handle_flash_attention_compatibility()
 
-        model = GR00TN15.from_pretrained(
-            pretrained_model_name_or_path=self.config.base_model_path,
-            tune_llm=self.config.tune_llm,
-            tune_visual=self.config.tune_visual,
-            tune_projector=self.config.tune_projector,
-            tune_diffusion_model=self.config.tune_diffusion_model,
-        )
+        model_kwargs = {
+            "pretrained_model_name_or_path": self.config.base_model_path,
+            "tune_llm": self.config.tune_llm,
+            "tune_visual": self.config.tune_visual,
+            "tune_projector": self.config.tune_projector,
+            "tune_diffusion_model": self.config.tune_diffusion_model,
+        }
+        if self.config.model_version == GROOT_N1_7:
+            from .groot_n1_7 import GR00TN17
+
+            model = GR00TN17.from_pretrained(
+                **model_kwargs,
+                tune_vlln=True,
+                transformers_loading_kwargs={"trust_remote_code": True},
+            )
+        else:
+            model = GR00TN15.from_pretrained(**model_kwargs)
 
         model.compute_dtype = "bfloat16" if self.config.use_bf16 else model.compute_dtype
         model.config.compute_dtype = model.compute_dtype
@@ -97,7 +116,7 @@ class GrootPolicy(PreTrainedPolicy):
 
     def reset(self):
         """Reset policy state when environment resets."""
-        self._action_queue = deque([], maxlen=self.config.n_action_steps)
+        self._action_queue = deque([], maxlen=self._action_queue_steps)
 
     @classmethod
     def from_pretrained(
@@ -141,8 +160,13 @@ class GrootPolicy(PreTrainedPolicy):
         from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
         from huggingface_hub.errors import HfHubHTTPError
 
+        requested_version = (
+            normalize_groot_model_version(config.model_version)
+            if config is not None
+            else infer_groot_model_version(str(pretrained_name_or_path)) or GROOT_N1_5
+        )
         print(
-            "The Groot policy is a wrapper around Nvidia's GR00T N1.5 model.\n"
+            f"The Groot policy is a wrapper around Nvidia's GR00T {requested_version} model.\n"
             f"Loading pretrained model from: {pretrained_name_or_path}"
         )
 
@@ -193,8 +217,12 @@ class GrootPolicy(PreTrainedPolicy):
         print("Detected base GR00T model, loading from HuggingFace...")
 
         if config is None:
+            model_version = infer_groot_model_version(str(pretrained_name_or_path)) or GROOT_N1_5
             # Create default config with the pretrained path
-            config = GrootConfig(base_model_path=str(pretrained_name_or_path))
+            config = GrootConfig(
+                model_version=model_version,
+                base_model_path=str(pretrained_name_or_path),
+            )
 
             # Add minimal visual feature required for validation
             # validate_features() will automatically add state and action features
@@ -215,6 +243,25 @@ class GrootPolicy(PreTrainedPolicy):
             if hasattr(config, key):
                 setattr(config, key, value)
 
+        config.model_version = normalize_groot_model_version(config.model_version)
+        inferred_version = infer_groot_model_version(config.base_model_path)
+        if inferred_version is not None and inferred_version != config.model_version:
+            raise ValueError(
+                f"GR00T model_version '{config.model_version}' does not match base_model_path "
+                f"'{config.base_model_path}', which looks like '{inferred_version}'."
+            )
+        if config.model_version == GROOT_N1_7:
+            if config.max_state_dim == 64:
+                config.max_state_dim = 132
+            if config.max_action_dim == 32:
+                config.max_action_dim = 132
+            if config.chunk_size == 50:
+                config.chunk_size = 40
+            if config.n_action_steps == 50:
+                config.n_action_steps = 40
+            if tuple(config.image_size) == (224, 224):
+                config.image_size = (256, 256)
+
         # Create a fresh policy instance - this will automatically load the GR00T model
         # in __init__ via _create_groot_model()
         policy = cls(config)
@@ -225,18 +272,59 @@ class GrootPolicy(PreTrainedPolicy):
     def get_optim_params(self) -> dict:
         return self.parameters()
 
+    def _resolve_action_queue_steps(self) -> int:
+        n_action_steps = int(self.config.n_action_steps)
+        if self.config.model_version != GROOT_N1_7:
+            return n_action_steps
+
+        checkpoint_action_horizon = infer_groot_n1_7_action_horizon(
+            self.config.base_model_path,
+            self.config.embodiment_tag,
+        )
+        execution_horizon = infer_groot_n1_7_action_execution_horizon(
+            self.config.base_model_path,
+            self.config.embodiment_tag,
+        )
+        horizons = [n_action_steps]
+        if checkpoint_action_horizon is not None:
+            horizons.append(checkpoint_action_horizon)
+        if execution_horizon is not None:
+            horizons.append(execution_horizon)
+        return min(horizons)
+
+    def _filter_groot_inputs(self, batch: dict[str, Tensor], *, include_action: bool) -> dict[str, Tensor]:
+        allowed_base = {"state", "state_mask", "embodiment_id"}
+        if include_action:
+            allowed_base.update({"action", "action_mask"})
+
+        if self.config.model_version == GROOT_N1_7:
+            allowed_base.update(
+                {
+                    "input_ids",
+                    "attention_mask",
+                    "pixel_values",
+                    "image_grid_thw",
+                    "mm_token_type_ids",
+                    "pixel_values_videos",
+                    "video_grid_thw",
+                }
+            )
+            allowed_base.add("action_mask")
+        else:
+            allowed_base.update({"action_mask"} if include_action else set())
+
+        return {
+            k: v
+            for k, v in batch.items()
+            if (k in allowed_base or k.startswith("eagle_")) and not (k.startswith("next.") or k == "info")
+        }
+
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Training forward pass.
 
         Delegates to Isaac-GR00T model.forward when inputs are compatible.
         """
-        # Build a clean input dict for GR00T: keep only tensors GR00T consumes
-        allowed_base = {"state", "state_mask", "action", "action_mask", "embodiment_id"}
-        groot_inputs = {
-            k: v
-            for k, v in batch.items()
-            if (k in allowed_base or k.startswith("eagle_")) and not (k.startswith("next.") or k == "info")
-        }
+        groot_inputs = self._filter_groot_inputs(batch, include_action=True)
 
         # Get device from model parameters
         device = next(self.parameters()).device
@@ -261,15 +349,10 @@ class GrootPolicy(PreTrainedPolicy):
         """
         self.eval()
 
-        # Build a clean input dict for GR00T: keep only tensors GR00T consumes
-        # Preprocessing is handled by the processor pipeline, so we just filter the batch
-        # NOTE: During inference, we should NOT pass action/action_mask (that's what we're predicting)
-        allowed_base = {"state", "state_mask", "embodiment_id"}
-        groot_inputs = {
-            k: v
-            for k, v in batch.items()
-            if (k in allowed_base or k.startswith("eagle_")) and not (k.startswith("next.") or k == "info")
-        }
+        # Preprocessing is handled by the processor pipeline, so we just filter the batch.
+        # During inference, we do not pass action because it is predicted.
+        # N1.7 still carries a 2-D action horizon mask from its checkpoint processor.
+        groot_inputs = self._filter_groot_inputs(batch, include_action=False)
 
         # Get device from model parameters
         device = next(self.parameters()).device
@@ -292,7 +375,7 @@ class GrootPolicy(PreTrainedPolicy):
 
         if len(self._action_queue) == 0:
             actions = self.predict_action_chunk(batch)
-            self._action_queue.extend(actions.transpose(0, 1))
+            self._action_queue.extend(actions[:, : self._action_queue_steps].transpose(0, 1))
         return self._action_queue.popleft()
 
     # -------------------------

src/lerobot/rewards/__init__.py

@@ -20,12 +20,14 @@ from .factory import (
     make_reward_pre_post_processors as make_reward_pre_post_processors,
 )
 from .pretrained import PreTrainedRewardModel as PreTrainedRewardModel
+from .robometer.configuration_robometer import RobometerConfig as RobometerConfig
 from .sarm.configuration_sarm import SARMConfig as SARMConfig
 from .topreward.configuration_topreward import TOPRewardConfig as TOPRewardConfig
 
 __all__ = [
     # Configuration classes
     "RewardClassifierConfig",
+    "RobometerConfig",
     "SARMConfig",
     "TOPRewardConfig",
     # Base class

src/lerobot/rewards/factory.py

@@ -25,6 +25,7 @@ from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 
 from .classifier.configuration_classifier import RewardClassifierConfig
 from .pretrained import PreTrainedRewardModel
+from .robometer.configuration_robometer import RobometerConfig
 from .sarm.configuration_sarm import SARMConfig
 from .topreward.configuration_topreward import TOPRewardConfig
 
@@ -38,7 +39,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", "robometer", "topreward".
 
     Returns:
         The reward model class corresponding to the given name.
@@ -54,6 +55,10 @@ def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
         from lerobot.rewards.sarm.modeling_sarm import SARMRewardModel
 
         return SARMRewardModel
+    elif name == "robometer":
+        from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+
+        return RobometerRewardModel
     elif name == "topreward":
         from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel
 
@@ -74,7 +79,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", "robometer", "topreward".
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -87,6 +92,8 @@ def make_reward_model_config(reward_type: str, **kwargs) -> RewardModelConfig:
         return RewardClassifierConfig(**kwargs)
     elif reward_type == "sarm":
         return SARMConfig(**kwargs)
+    elif reward_type == "robometer":
+        return RobometerConfig(**kwargs)
     elif reward_type == "topreward":
         return TOPRewardConfig(**kwargs)
     else:
@@ -168,6 +175,13 @@ def make_reward_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
             dataset_meta=kwargs.get("dataset_meta"),
         )
+    elif isinstance(reward_cfg, RobometerConfig):
+        from lerobot.rewards.robometer.processor_robometer import make_robometer_pre_post_processors
+
+        return make_robometer_pre_post_processors(
+            config=reward_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
 
     elif isinstance(reward_cfg, TOPRewardConfig):
         from lerobot.rewards.topreward.processor_topreward import make_topreward_pre_post_processors

src/lerobot/rewards/robometer/__init__.py

@@ -0,0 +1,19 @@
+# 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_robometer import RobometerConfig
+from .modeling_robometer import RobometerRewardModel
+from .processor_robometer import make_robometer_pre_post_processors
+
+__all__ = ["RobometerConfig", "RobometerRewardModel", "make_robometer_pre_post_processors"]

src/lerobot/rewards/robometer/compute_rabc_weights.py

@@ -0,0 +1,320 @@
+#!/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 Robometer progress and success curves for a LeRobot dataset.
+
+For each episode, builds per-frame sub-samples using the frame-steps
+strategy from the Robometer eval server: for each original frame ``t``,
+linspace-subsample ``[0, t]`` into ``K`` frames (default 4, matching
+``NUM_SUBSAMPLED_FRAMES`` in the eval server), run one forward through
+the Robometer processor + model, and keep the last-frame progress value.
+All sub-samples are the same size ``K`` so they batch cleanly.
+
+The parquet uses the same schema as SARM's
+:mod:`lerobot.rewards.sarm.compute_rabc_weights` so existing consumers —
+:class:`lerobot.rewards.sarm.rabc.RABCWeights` (which reads
+``progress_sparse``) and the progress-overlay script in
+``examples/dataset/create_progress_videos.py`` — work without modification.
+
+Usage:
+    # Dense per-frame progress for one episode
+    python -m lerobot.rewards.robometer.compute_rabc_weights \\
+        --dataset-repo-id lerobot/libero_10_image \\
+        --reward-model-path lerobot/Robometer-4B \\
+        --episodes 0
+
+    # All episodes with batching
+    python -m lerobot.rewards.robometer.compute_rabc_weights \\
+        --dataset-repo-id lerobot/libero_10_image \\
+        --reward-model-path lerobot/Robometer-4B \\
+        --batch-size 16
+"""
+
+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.robometer.configuration_robometer import RobometerConfig
+from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+from lerobot.rewards.robometer.processor_robometer import RobometerEncoderProcessorStep
+from lerobot.types import TransitionKey
+
+DEFAULT_OUTPUT_FILENAME = "robometer_progress.parquet"
+
+# Upstream Robometer eval server uses K=4 for frame-steps sub-samples.
+DEFAULT_NUM_SUBSAMPLED_FRAMES = 4
+
+
+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 _build_subsample_indices(num_frames: int, num_subsampled_frames: int) -> list[np.ndarray]:
+    """Frame-steps linspace expansion.
+
+    For each ``t in [0, num_frames - 1]`` returns ``num_subsampled_frames``
+    indices from ``np.linspace(0, t, num_subsampled_frames)`` — the first
+    and last frames are always included. Each entry is a fixed-size array
+    so the model can batch them.
+    """
+    return [np.linspace(0, t, num_subsampled_frames).round().astype(np.int64) for t in range(num_frames)]
+
+
+def compute_robometer_progress(
+    dataset_repo_id: str,
+    reward_model_path: str,
+    output_path: str | None = None,
+    device: str = "cuda",
+    batch_size: int = 32,
+    num_subsampled_frames: int = DEFAULT_NUM_SUBSAMPLED_FRAMES,
+    episodes: list[int] | None = None,
+    image_key: str | None = None,
+) -> Path:
+    """Run Robometer over a dataset and write per-frame progress + success."""
+    logging.info(f"Loading Robometer: {reward_model_path}")
+    config = RobometerConfig(pretrained_path=reward_model_path, device=device)
+    if image_key is not None:
+        config.image_key = image_key
+    model = RobometerRewardModel.from_pretrained(reward_model_path, config=config)
+    model.to(device).eval()
+
+    encoder = RobometerEncoderProcessorStep(
+        base_model_id=config.base_model_id,
+        image_key=config.image_key,
+        task_key=config.task_key,
+        default_task=config.default_task,
+        max_frames=num_subsampled_frames,
+        use_multi_image=config.use_multi_image,
+        use_per_frame_progress_token=config.use_per_frame_progress_token,
+    )
+
+    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")
+
+        ep_frames = torch.stack([dataset[ep_start + i][image_key] for i in range(num_frames)])
+
+        sub_indices = _build_subsample_indices(num_frames, num_subsampled_frames)
+
+        progress_per_frame = np.zeros(num_frames, dtype=np.float32)
+
+        for start in tqdm(range(0, num_frames, batch_size), desc=f"  Ep {episode_idx}", leave=False):
+            end = min(start + batch_size, num_frames)
+            frames_batch = torch.stack([ep_frames[sub_indices[i]] for i in range(start, end)])
+
+            transition = {
+                TransitionKey.OBSERVATION: {image_key: frames_batch},
+                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():
+                rewards = model.compute_reward(batch)
+            progress_per_frame[start:end] = rewards.cpu().numpy()
+
+        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(progress_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),
+        }
+    ).replace_schema_metadata({b"reward_model_path": reward_model_path.encode()})
+
+    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 Robometer progress curves for RA-BC weighting.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Dense per-frame progress for one episode
+    python -m lerobot.rewards.robometer.compute_rabc_weights \\
+        --dataset-repo-id lerobot/libero_10_image \\
+        --reward-model-path lerobot/Robometer-4B \\
+        --episodes 0
+
+    # All episodes, smaller batches for memory-constrained GPUs
+    python -m lerobot.rewards.robometer.compute_rabc_weights \\
+        --dataset-repo-id lerobot/libero_10_image \\
+        --reward-model-path lerobot/Robometer-4B \\
+        --batch-size 16
+        """,
+    )
+    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="Robometer checkpoint repo id or local path."
+    )
+    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(
+        "--batch-size", type=int, default=32, help="Sub-samples per Qwen forward (default: 32)."
+    )
+    parser.add_argument(
+        "--num-subsampled-frames",
+        type=int,
+        default=DEFAULT_NUM_SUBSAMPLED_FRAMES,
+        help=f"Frames per sub-sample (default: {DEFAULT_NUM_SUBSAMPLED_FRAMES}, matches eval server).",
+    )
+    parser.add_argument(
+        "--episodes", type=int, nargs="+", default=None, help="Process only these episode indices."
+    )
+    parser.add_argument(
+        "--image-key", type=str, default=None, help="Image observation key (default: from config)."
+    )
+    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")
+
+    reward_model_path = args.reward_model_path
+    if reward_model_path is None:
+        temp_dataset = LeRobotDataset(args.dataset_repo_id, download_videos=False)
+        parquet_path = Path(temp_dataset.root) / DEFAULT_OUTPUT_FILENAME
+        reward_model_path = get_reward_model_path_from_parquet(parquet_path)
+        if reward_model_path:
+            logging.info(f"Using reward model from parquet metadata: {reward_model_path}")
+        else:
+            raise ValueError(
+                "--reward-model-path is required (no existing parquet with model metadata found)."
+            )
+
+    output_path = compute_robometer_progress(
+        dataset_repo_id=args.dataset_repo_id,
+        reward_model_path=reward_model_path,
+        output_path=args.output_path,
+        device=args.device,
+        batch_size=args.batch_size,
+        num_subsampled_frames=args.num_subsampled_frames,
+        episodes=args.episodes,
+        image_key=args.image_key,
+    )
+
+    print(f"\nRobometer 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/robometer/configuration_robometer.py

@@ -0,0 +1,158 @@
+# 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 copy import deepcopy
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.rewards import RewardModelConfig
+from lerobot.utils.constants import OBS_IMAGES
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoConfig, AutoTokenizer
+else:
+    AutoConfig = None  # type: ignore[assignment]
+    AutoTokenizer = None  # type: ignore[assignment]
+
+
+# Special tokens Robometer adds to the Qwen-VL tokenizer at construction time.
+# The order is part of the data contract: upstream resized ``embed_tokens``
+# after adding these tokens in this exact order, so changing the set or order
+# would silently misalign the saved embedding rows with their token ids.
+# ``<|reward_token|>`` and ``<|sim_token|>`` are leftover from earlier upstream
+# heads (never read at inference) but still occupy rows the checkpoint expects.
+ROBOMETER_SPECIAL_TOKENS = (
+    "<|split_token|>",
+    "<|reward_token|>",
+    "<|pref_token|>",
+    "<|sim_token|>",
+    "<|prog_token|>",
+)
+
+
+@RewardModelConfig.register_subclass("robometer")
+@dataclass
+class RobometerConfig(RewardModelConfig):
+    """Configuration for the Robometer reward model."""
+
+    pretrained_path: str | None = "lerobot/Robometer-4B"
+    image_key: str = OBS_IMAGES + ".top"
+    task_key: str = "task"
+    default_task: str | None = None
+
+    max_frames: int | None = 8
+    reward_output: str = "progress"  # "progress" or "success"
+    success_threshold: float = 0.5
+
+    license: str | None = "apache-2.0"
+    tags: list[str] | None = field(
+        default_factory=lambda: ["reward-model", "vision-language", "qwen3-vl", "zero-shot"]
+    )
+
+    base_model_id: str = "Qwen/Qwen3-VL-4B-Instruct"
+    torch_dtype: str = "bfloat16"
+    use_multi_image: bool = True
+    use_per_frame_progress_token: bool = True
+    average_temporal_patches: bool = True
+    frame_pooling: str = "mean"  # "mean" | "boundary" | "attention"
+    frame_pooling_attn_temperature: float = 1.0
+    progress_loss_type: str = "discrete"  # "l1" | "l2" | "discrete"
+    progress_discrete_bins: int = 10
+
+    # Serialised Qwen backbone config (post-resize). Always populated by
+    # ``__post_init__`` from ``base_model_id`` + ``len(tokenizer) + 5``, so it
+    # is non-empty after construction. Saved into ``config.json`` automatically
+    # by the base ``_save_pretrained``.
+    vlm_config: dict[str, Any] = field(default_factory=dict)
+
+    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.reward_output not in {"progress", "success"}:
+            raise ValueError(f"reward_output must be 'progress' or 'success', got {self.reward_output!r}")
+        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.frame_pooling not in {"mean", "boundary", "attention"}:
+            raise ValueError(f"frame_pooling must be mean/boundary/attention; got {self.frame_pooling!r}")
+        if self.frame_pooling_attn_temperature <= 0:
+            raise ValueError("frame_pooling_attn_temperature must be > 0")
+        if self.progress_loss_type not in {"l1", "l2", "discrete"}:
+            raise ValueError(f"progress_loss_type must be l1/l2/discrete; got {self.progress_loss_type!r}")
+        if self.use_per_frame_progress_token and not self.use_multi_image:
+            raise ValueError("use_per_frame_progress_token=True requires use_multi_image=True")
+
+        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("progress", PolicyFeature(shape=(1,), type=FeatureType.REWARD))
+        self.output_features.setdefault("success", PolicyFeature(shape=(1,), type=FeatureType.REWARD))
+
+        # Deterministically populate ``vlm_config`` so it is non-empty after
+        # construction. For ``Qwen/Qwen3-VL-4B-Instruct`` this gives
+        # ``len(tokenizer) + 5 = 151,669 + 5 = 151,674`` — the exact post-resize
+        # vocab the published ``Robometer-4B`` checkpoint was saved with.
+        if not self.vlm_config:
+            require_package("transformers", extra="robometer")
+            vlm = AutoConfig.from_pretrained(self.base_model_id).to_dict()
+            tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
+            text_config = vlm.get("text_config")
+            if not isinstance(text_config, dict):
+                raise ValueError(
+                    f"Backbone config for {self.base_model_id!r} has no nested `text_config`; "
+                    "Robometer expects a Qwen-VL-style config."
+                )
+            text_config["vocab_size"] = len(tokenizer) + len(ROBOMETER_SPECIAL_TOKENS)
+            self.vlm_config = vlm
+
+    @property
+    def use_discrete_progress(self) -> bool:
+        """Whether the progress head outputs distribution logits over bins."""
+        return self.progress_loss_type.lower() == "discrete"
+
+    @property
+    def vlm_backbone_config(self):
+        """Reconstruct the Qwen backbone config from :attr:`vlm_config`."""
+        require_package("transformers", extra="robometer")
+        config_dict = deepcopy(self.vlm_config)
+        model_type = config_dict.pop("model_type", None)
+        if model_type is None:
+            raise ValueError("vlm_config must include `model_type` to reconstruct the backbone config")
+        return AutoConfig.for_model(model_type, **config_dict)
+
+    @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"Robometer requires image input feature {self.image_key!r}")

src/lerobot/rewards/robometer/modeling_robometer.py

@@ -0,0 +1,481 @@
+# Copyright 2026 Anthony Liang, Yigit Korkmaz, Stephen Tu, Erdem Bıyık, Jesse Zhang
+# 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.
+
+"""ROBOMETER: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons.
+
+Paper:         https://arxiv.org/abs/2603.02115
+Project:       https://robometer.github.io
+Original code: https://github.com/aliang8/robometer
+Model:         https://huggingface.co/robometer/Robometer-4B
+
+Robometer is a general-purpose, video-language-input reward model built on
+``Qwen/Qwen3-VL-4B-Instruct``. It is trained with a dual reward-prediction
+objective:
+
+- A frame-level progress loss anchoring reward magnitude on expert data.
+- A trajectory-comparison preference loss imposing global ordering constraints
+  across trajectories sharing the same instruction.
+
+To support downstream RL it also predicts a frame-level binary success. The
+training prompt inserts three learnable tokens:
+
+- ``<|prog_token|>`` after each frame to read per-frame progress and success.
+- ``<|pref_token|>`` at the end to read pairwise preference (training-only).
+- ``<|split_token|>`` between two trajectories in preference samples
+  (training-only).
+
+Progress is modeled as a categorical distribution over ``progress_discrete_bins``
+uniformly-spaced centers in ``[0, 1]`` (C51-style), and the continuous estimate
+is recovered as the softmax-weighted mean of those centers — see
+:func:`convert_bins_to_continuous`.
+
+This LeRobot port is **inference-only**: the preference head is preserved in
+the state dict for byte-equivalence with the published ``Robometer-4B``
+checkpoint but is not queried by :meth:`RobometerRewardModel.compute_reward`,
+which returns the last-frame progress (clamped to ``[0, 1]``) or sigmoid'd
+success probability depending on :attr:`RobometerConfig.reward_output`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import torch
+from torch import Tensor, nn
+
+from lerobot.rewards.pretrained import PreTrainedRewardModel
+from lerobot.rewards.robometer.configuration_robometer import RobometerConfig
+from lerobot.utils.constants import OBS_PREFIX
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoModelForImageTextToText
+else:
+    AutoModelForImageTextToText = None  # type: ignore[assignment]
+
+logger = logging.getLogger(__name__)
+
+# Namespace for Robometer's pre-encoded Qwen-VL observation tensors.
+ROBOMETER_FEATURE_PREFIX = f"{OBS_PREFIX}robometer."
+ROBOMETER_QWEN_INPUT_KEYS = (
+    "input_ids",
+    "attention_mask",
+    "pixel_values",
+    "pixel_values_videos",
+    "image_grid_thw",
+    "video_grid_thw",
+    "second_per_grid_ts",
+    "mm_token_type_ids",
+)
+ROBOMETER_METADATA_KEYS = (
+    "prog_token_id",
+    "vision_start_token_id",
+    "vision_end_token_id",
+    "video_merge_size",
+)
+ROBOMETER_INPUT_KEYS = ROBOMETER_QWEN_INPUT_KEYS + ROBOMETER_METADATA_KEYS
+
+
+def convert_bins_to_continuous(bin_logits: Tensor) -> Tensor:
+    """Collapse per-bin logits into a single value in ``[0, 1]``.
+
+    The discrete progress head outputs ``num_bins`` logits per frame. Bins are
+    evenly spaced centers in ``[0, 1]``; the continuous prediction is the
+    softmax-weighted mean of those centers.
+    """
+    bin_probs = torch.softmax(bin_logits, dim=-1)
+    num_bins = bin_logits.shape[-1]
+    bin_centers = torch.linspace(0.0, 1.0, num_bins, device=bin_logits.device, dtype=bin_logits.dtype)
+    return (bin_probs * bin_centers).sum(dim=-1)
+
+
+def _squeeze_last_safe(x: Tensor) -> Tensor:
+    """Drop a trailing singleton dim only when present."""
+    return x.squeeze(-1) if x.ndim > 1 and x.shape[-1] == 1 else x
+
+
+def _torch_dtype(name: str) -> torch.dtype:
+    dtype = getattr(torch, name, None)
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    raise ValueError(f"Unknown torch dtype: {name!r}")
+
+
+class RobometerPredictionHead(nn.Sequential):
+    """Small MLP head used for Robometer's progress / success / preference outputs."""
+
+    def __init__(self, hidden_dim: int, output_size: int, *, dropout: float, with_sigmoid: bool) -> None:
+        layers: list[nn.Module] = [
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.LayerNorm(hidden_dim // 2),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim // 2, output_size),
+        ]
+        if with_sigmoid:
+            layers.append(nn.Sigmoid())
+        super().__init__(*layers)
+
+
+def decode_progress_outputs(
+    progress_logits: Tensor | None,
+    success_logits: Tensor | None,
+    *,
+    is_discrete_mode: bool,
+) -> dict[str, list[list[float]]]:
+    """Decode RBM head outputs into per-frame floats.
+
+    Args:
+        progress_logits: ``(B, T)`` (continuous) or ``(B, T, num_bins)`` (discrete).
+        success_logits: ``(B, T)`` raw logits, ``sigmoid``-ed to probabilities.
+        is_discrete_mode: if True the progress logits get a softmax over bins
+            and are projected onto bin centers via :func:`convert_bins_to_continuous`.
+
+    Returns:
+        Dict with ``progress_pred`` and ``success_probs``, each a list of
+        length ``B`` of per-frame float lists.
+    """
+    progress_pred: list[list[float]] = []
+    success_probs: list[list[float]] = []
+
+    if progress_logits is not None:
+        for sample_logits in progress_logits:
+            if is_discrete_mode:
+                continuous = convert_bins_to_continuous(sample_logits.detach().float().cpu())
+                progress_pred.append(continuous.flatten().tolist())
+            else:
+                progress_pred.append(sample_logits.detach().float().cpu().flatten().tolist())
+
+    if success_logits is not None:
+        for sample_logits in success_logits:
+            success_probs.append(torch.sigmoid(sample_logits.detach().float().cpu()).flatten().tolist())
+
+    return {"progress_pred": progress_pred, "success_probs": success_probs}
+
+
+class RobometerRewardModel(PreTrainedRewardModel):
+    """Robometer (RBM) reward model — inference-only LeRobot port.
+
+    Wraps a Qwen-VL backbone (default: ``Qwen/Qwen3-VL-4B-Instruct``) with three
+    prediction heads from the paper (progress, success, preference). At
+    inference time only the progress and success heads are queried; the
+    preference head is kept on the module so the published ``Robometer-4B``
+    safetensors load unchanged.
+    """
+
+    name = "robometer"
+    config_class = RobometerConfig
+
+    def __init__(self, config: RobometerConfig, *, dropout: float = 0.1) -> None:
+        require_package("transformers", extra="robometer")
+        super().__init__(config)
+        self.config = config
+
+        # Two backbone-build paths (EO-1 style, branched on ``pretrained_path``):
+        #
+        #   - Fresh training (``pretrained_path is None``): download the base
+        #     Qwen weights and resize the embed table to match
+        #     ``vlm_config.text_config.vocab_size`` — populated deterministically
+        #     in ``RobometerConfig.__post_init__`` as
+        #     ``len(tokenizer) + len(ROBOMETER_SPECIAL_TOKENS)``
+        #
+        #   - Loading a saved checkpoint (``pretrained_path`` is set): rebuild
+        #     the empty architecture from ``vlm_config`` via
+        #     ``AutoModelForImageTextToText.from_config`` so the subsequent
+        #     ``model.safetensors`` load is a direct fill of the right shape —
+        #     no redundant Qwen weight download.
+        torch_dtype = _torch_dtype(config.torch_dtype)
+        if config.pretrained_path is None:
+            self.model = AutoModelForImageTextToText.from_pretrained(
+                config.base_model_id,
+                dtype=torch_dtype,
+                trust_remote_code=True,
+            )
+            target_vocab = config.vlm_config["text_config"]["vocab_size"]
+            self.model.resize_token_embeddings(target_vocab)
+        else:
+            self.model = AutoModelForImageTextToText.from_config(
+                config.vlm_backbone_config,
+                dtype=torch_dtype,
+                trust_remote_code=True,
+            )
+
+        # All Qwen-VL backbones Robometer supports expose `text_config.hidden_size`.
+        # Falls back to the top-level `hidden_size` so future non-multimodal
+        # variants would still resolve.
+        backbone_config = self.model.config
+        text_config = getattr(backbone_config, "text_config", None)
+        hidden_size = getattr(text_config, "hidden_size", None) if text_config is not None else None
+        if hidden_size is None:
+            hidden_size = getattr(backbone_config, "hidden_size", None)
+        if hidden_size is None:
+            raise AttributeError(
+                f"Could not infer hidden_size from backbone config of {config.base_model_id}"
+            )
+        hidden_dim = int(hidden_size)
+
+        # Robometer's three prediction heads + frame-pool attention.
+        progress_output = config.progress_discrete_bins if config.use_discrete_progress else 1
+        self.progress_head = RobometerPredictionHead(
+            hidden_dim,
+            progress_output,
+            dropout=dropout,
+            with_sigmoid=not config.use_discrete_progress,
+        )
+        self.preference_head = RobometerPredictionHead(hidden_dim, 1, dropout=dropout, with_sigmoid=False)
+        self.success_head = RobometerPredictionHead(hidden_dim, 1, dropout=dropout, with_sigmoid=False)
+        self.frame_pool_attn = nn.Linear(hidden_dim, 1, bias=False)
+
+        # Match the dtype of the loaded base model so weight loading is a no-op cast.
+        model_dtype = next(self.model.parameters()).dtype
+        self.progress_head.to(dtype=model_dtype)
+        self.preference_head.to(dtype=model_dtype)
+        self.success_head.to(dtype=model_dtype)
+        self.frame_pool_attn.to(dtype=model_dtype)
+
+    def compute_reward(self, batch: dict[str, Tensor]) -> Tensor:
+        inputs = {
+            key: batch[f"{ROBOMETER_FEATURE_PREFIX}{key}"]
+            for key in ROBOMETER_INPUT_KEYS
+            if f"{ROBOMETER_FEATURE_PREFIX}{key}" in batch
+        }
+        if "input_ids" not in inputs:
+            raise KeyError(
+                f"Robometer batch missing pre-encoded inputs (expected "
+                f"`{ROBOMETER_FEATURE_PREFIX}input_ids`). Make sure the "
+                "RobometerEncoderProcessorStep ran before `compute_reward`."
+            )
+
+        device = next(self.model.parameters()).device
+        inputs = {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()}
+
+        self.eval()
+        with torch.no_grad():
+            progress_logits, success_logits = self._compute_rbm_logits(inputs)
+
+        decoded = decode_progress_outputs(
+            progress_logits,
+            success_logits,
+            is_discrete_mode=self.config.use_discrete_progress,
+        )
+        values = (
+            decoded["success_probs"] if self.config.reward_output == "success" else decoded["progress_pred"]
+        )
+
+        rewards = torch.stack([torch.as_tensor(seq, dtype=torch.float32)[-1] for seq in values])
+        if self.config.reward_output == "success":
+            rewards = (rewards > self.config.success_threshold).float()
+        else:
+            # Match upstream Robometer's ``extract_rewards_from_output``: per-frame
+            # progress predictions are clamped to ``[0, 1]`` before being returned.
+            rewards = rewards.clamp(0.0, 1.0)
+        return rewards.to(self.config.device or "cpu")
+
+    def _compute_rbm_logits(
+        self,
+        inputs: dict[str, Any],
+    ) -> tuple[Tensor, Tensor]:
+        """Run the Qwen3-VL backbone and apply Robometer's heads.
+
+        ``inputs`` is the encoded batch produced by
+        :class:`RobometerEncoderProcessorStep`. It carries Qwen tensors as well
+        as Robometer-specific metadata (``prog_token_id``,
+        ``vision_start_token_id``, ``vision_end_token_id``, ``video_merge_size``)
+        — the metadata is popped here so the rest can be forwarded straight to
+        the Qwen model.
+
+        Returns ``(progress_logits, success_logits)``. Shapes:
+
+        - ``progress_logits``: ``(B, T)`` (continuous) or ``(B, T, num_bins)`` (discrete).
+        - ``success_logits``: ``(B, T)`` raw logits (sigmoid happens at decode time).
+        """
+        prog_token_id = inputs.pop("prog_token_id", None)
+        vision_start_token_id = inputs.pop("vision_start_token_id", None)
+        vision_end_token_id = inputs.pop("vision_end_token_id", None)
+        video_merge_size = inputs.pop("video_merge_size", 14)
+
+        # Qwen3-VL doesn't reliably populate `last_hidden_state`; ask for the
+        # full hidden-state tuple and take the last layer. This matches the
+        # `is_qwen3` path in upstream Robometer's `RBM.forward_qwen` (main).
+        outputs = self.model(**inputs, output_hidden_states=True, return_dict=True)
+        hidden_state = (
+            outputs.hidden_states[-1]
+            if getattr(outputs, "hidden_states", None)
+            else outputs.last_hidden_state
+        )
+
+        input_ids = inputs["input_ids"]
+        if self.config.use_per_frame_progress_token:
+            if prog_token_id is None:
+                raise KeyError("`prog_token_id` missing in batch (run RobometerEncoderProcessorStep first)")
+            return self._process_token_extraction(hidden_state, input_ids, prog_token_id=prog_token_id)
+        if self.config.use_multi_image:
+            if vision_start_token_id is None or vision_end_token_id is None:
+                raise KeyError(
+                    "`vision_start_token_id` / `vision_end_token_id` missing in batch "
+                    "(run RobometerEncoderProcessorStep first)"
+                )
+            return self._process_multi_image_frames(
+                hidden_state,
+                input_ids,
+                start_id=vision_start_token_id,
+                end_id=vision_end_token_id,
+            )
+        video_grid_thw = inputs.get("video_grid_thw")
+        if video_grid_thw is None:
+            raise ValueError("video_grid_thw is required for video-mode Robometer inference")
+        if vision_start_token_id is None:
+            raise KeyError("`vision_start_token_id` missing in batch")
+        return self._process_video_frames(
+            hidden_state,
+            input_ids,
+            video_grid_thw,
+            start_id=vision_start_token_id,
+            merge_size=video_merge_size,
+        )
+
+    def _apply_heads_to_hidden_states(self, frame_embeddings: Tensor) -> tuple[Tensor, Tensor]:
+        """Apply progress + success heads to a tensor of frame embeddings."""
+        progress_out = self.progress_head(frame_embeddings)
+        progress = progress_out if self.config.use_discrete_progress else _squeeze_last_safe(progress_out)
+        success = _squeeze_last_safe(self.success_head(frame_embeddings))
+        return progress, success
+
+    def _process_token_extraction(
+        self,
+        hidden_state: Tensor,
+        input_ids: Tensor,
+        *,
+        prog_token_id: int,
+    ) -> tuple[Tensor, Tensor]:
+        """Per-frame progress/success from ``<|prog_token|>`` positions."""
+        token_mask = input_ids == prog_token_id
+        batch_indices, positions = token_mask.nonzero(as_tuple=True)
+        if positions.numel() == 0:
+            raise ValueError("`<|prog_token|>` not found in any sequence")
+
+        per_sample_hidden = [
+            hidden_state[i, positions[batch_indices == i]] for i in range(input_ids.shape[0])
+        ]
+        progress_list, success_list = [], []
+        for embeddings in per_sample_hidden:
+            if embeddings.shape[0] == 0:
+                raise ValueError("`<|prog_token|>` missing in a sequence")
+            progress, success = self._apply_heads_to_hidden_states(embeddings)
+            progress_list.append(progress)
+            success_list.append(success)
+
+        return torch.stack(progress_list), torch.stack(success_list)
+
+    def _process_multi_image_frames(
+        self,
+        hidden_state: Tensor,
+        input_ids: Tensor,
+        *,
+        start_id: int,
+        end_id: int,
+    ) -> tuple[Tensor, Tensor]:
+        """Per-frame progress/success in multi-image mode (Qwen-VL)."""
+        progress_list, success_list = [], []
+        for batch_idx in range(input_ids.shape[0]):
+            seq_ids = input_ids[batch_idx]
+            seq_hidden = hidden_state[batch_idx]
+            frame_embeddings = self._extract_hidden_states_from_token_pairs(
+                seq_hidden, seq_ids, start_id, end_id
+            )
+            progress, success = self._apply_heads_to_hidden_states(frame_embeddings)
+            progress_list.append(progress)
+            success_list.append(success)
+
+        return torch.stack(progress_list), torch.stack(success_list)
+
+    def _extract_hidden_states_from_token_pairs(
+        self,
+        hidden_state: Tensor,
+        input_ids: Tensor,
+        start_id: int,
+        end_id: int,
+    ) -> Tensor:
+        start_positions = (input_ids == start_id).nonzero(as_tuple=True)[0]
+        end_positions = (input_ids == end_id).nonzero(as_tuple=True)[0]
+        if start_positions.numel() == 0:
+            raise ValueError("`<|vision_start|>` not found in sequence")
+        if start_positions.numel() != end_positions.numel():
+            raise ValueError(
+                f"Mismatched vision token counts: {start_positions.numel()} start vs "
+                f"{end_positions.numel()} end"
+            )
+
+        frames: list[Tensor] = []
+        for start, end in zip(start_positions.tolist(), end_positions.tolist(), strict=True):
+            if start >= end:
+                raise ValueError(f"Invalid vision token pair: start={start} end={end}")
+            patch_tokens = hidden_state[start + 1 : end]
+            if patch_tokens.shape[0] == 0:
+                frames.append((hidden_state[start] + hidden_state[end]) / 2.0)
+                continue
+
+            pooling = self.config.frame_pooling
+            if pooling == "mean":
+                frames.append(patch_tokens.mean(dim=0))
+            elif pooling == "boundary":
+                frames.append(patch_tokens[-1])
+            else:  # attention
+                scores = (
+                    self.frame_pool_attn(patch_tokens).squeeze(-1)
+                    / self.config.frame_pooling_attn_temperature
+                )
+                weights = torch.softmax(scores, dim=0).unsqueeze(-1)
+                frames.append((weights * patch_tokens).sum(dim=0))
+
+        return torch.stack(frames)
+
+    def _process_video_frames(
+        self,
+        hidden_state: Tensor,
+        input_ids: Tensor,
+        video_grid_thw: Tensor,
+        *,
+        start_id: int,
+        merge_size: int,
+    ) -> tuple[Tensor, Tensor]:
+        """Per-frame progress/success in video mode (Qwen-VL)."""
+        progress_list, success_list = [], []
+        for batch_idx in range(input_ids.shape[0]):
+            seq_ids = input_ids[batch_idx]
+            seq_hidden = hidden_state[batch_idx]
+            start_positions = (seq_ids == start_id).nonzero(as_tuple=True)[0]
+            if start_positions.numel() == 0:
+                raise ValueError("`<|vision_start|>` not found in sequence")
+            t_dim, h_dim, w_dim = (int(x) for x in video_grid_thw[batch_idx].tolist())
+            tokens_per_frame = (h_dim * w_dim) // (merge_size**2)
+
+            cursor = start_positions[0].item()
+            frame_embeddings: list[Tensor] = []
+            for _ in range(t_dim):
+                if self.config.average_temporal_patches:
+                    patch = seq_hidden[cursor : cursor + tokens_per_frame]
+                    frame_embeddings.append(patch.mean(dim=0))
+                else:
+                    frame_embeddings.append(seq_hidden[cursor + tokens_per_frame])
+                cursor += tokens_per_frame
+
+            stacked = torch.stack(frame_embeddings)
+            progress, success = self._apply_heads_to_hidden_states(stacked)
+            progress_list.append(progress)
+            success_list.append(success)
+
+        return torch.stack(progress_list), torch.stack(success_list)

src/lerobot/rewards/robometer/processor_robometer.py

@@ -0,0 +1,338 @@
+# 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.
+
+"""Robometer pre/post processing pipelines."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from PIL import Image
+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.robometer.configuration_robometer import (
+    ROBOMETER_SPECIAL_TOKENS,
+    RobometerConfig,
+)
+from lerobot.rewards.robometer.modeling_robometer import ROBOMETER_FEATURE_PREFIX
+from lerobot.types import EnvTransition, TransitionKey
+from lerobot.utils.constants import (
+    OBS_IMAGES,
+    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
+
+PROGRESS_PROMPT = (
+    "The task for the robot is '{task}'. Given the trajectory video, predict "
+    "the task progress at each frame, how far along the robot is towards "
+    "completing the task, a float between 0 and 1, where 0 is the starting "
+    "state and 1 is when the task is completed. If the robot is not "
+    "performing the same task, predict 0 progress."
+)
+
+
+def _frames_to_pil(frames: np.ndarray) -> list[Image.Image]:
+    """Convert ``(T, H, W, C)`` uint8 frames to a list of PIL images."""
+    if frames.ndim != 4:
+        raise ValueError(f"Expected (T,H,W,C) frames; got shape {frames.shape}")
+    if frames.dtype != np.uint8:
+        frames = np.clip(frames, 0, 255).astype(np.uint8)
+    return [Image.fromarray(frames[i]) for i in range(frames.shape[0])]
+
+
+def _video_to_numpy(video: Tensor, *, max_frames: int | None) -> np.ndarray:
+    """Convert one trajectory tensor to a ``(T, H, W, C) uint8`` numpy array."""
+    if max_frames is not None:
+        video = video[-max_frames:]
+    if video.shape[1] in (1, 3):
+        video = video.permute(0, 2, 3, 1)
+    elif video.shape[-1] not in (1, 3):
+        raise ValueError(f"Expected channel dim of size 1 or 3, got shape {tuple(video.shape)}")
+
+    array = video.detach().cpu().numpy()
+    if np.issubdtype(array.dtype, np.floating) and array.size > 0 and array.max() <= 1.0:
+        array = array * 255.0
+    return np.clip(array, 0, 255).astype(np.uint8)
+
+
+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("Robometer 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"Robometer 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="robometer_encoder")
+class RobometerEncoderProcessorStep(ProcessorStep):
+    """Encode raw frames + task into Qwen-VL tensors for the Robometer model.
+
+    Loads a :class:`~transformers.AutoProcessor` matching ``base_model_id`` and
+    registers Robometer's special tokens on the tokenizer. The matching
+    embedding resize happens model-side in
+    :meth:`RobometerRewardModel.__init__`.
+
+    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"{ROBOMETER_FEATURE_PREFIX}<name>"]`` for:
+
+    - the Qwen-VL processor outputs: ``input_ids``, ``attention_mask``,
+      ``pixel_values``, ``image_grid_thw``, ``video_grid_thw``, ...
+    - Robometer-specific token ids consumed by the model heads:
+      ``prog_token_id``, ``vision_start_token_id``, ``vision_end_token_id``,
+      ``video_merge_size``.
+    """
+
+    base_model_id: str = "Qwen/Qwen3-VL-4B-Instruct"
+    image_key: str = OBS_IMAGES + ".top"
+    task_key: str = "task"
+    default_task: str | None = None
+    max_frames: int | None = 8
+    use_multi_image: bool = True
+    use_per_frame_progress_token: bool = True
+    max_length: int = 1024
+
+    _processor: Any = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        require_package("transformers", extra="robometer")
+        require_package("qwen-vl-utils", extra="robometer", import_name="qwen_vl_utils")
+
+        self._processor = AutoProcessor.from_pretrained(
+            self.base_model_id,
+            trust_remote_code=True,
+            do_sample_frames=False,
+            padding_side="right",
+        )
+
+        # Register Robometer's special tokens on the tokenizer. The matching
+        # embedding resize happens model-side in `RobometerRewardModel.__init__`.
+        tokenizer = self._processor.tokenizer
+        # Qwen tokenizers may not define a pad token, but batched prompts/videos
+        # require padding, so reuse EOS as the padding token.
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        for token in ROBOMETER_SPECIAL_TOKENS:
+            if token not in tokenizer.get_vocab():
+                tokenizer.add_special_tokens({"additional_special_tokens": [token]})
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        observation = transition.get(TransitionKey.OBSERVATION)
+        complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
+        if not isinstance(observation, dict):
+            raise ValueError("RobometerEncoderProcessorStep requires an observation dict")
+
+        if self.image_key not in observation:
+            raise KeyError(f"Robometer expected image key {self.image_key!r} in observation")
+
+        frames = observation[self.image_key]
+        tensor = frames.detach().cpu() if isinstance(frames, Tensor) else torch.as_tensor(frames)
+        if tensor.ndim == 4:
+            tensor = tensor.unsqueeze(1)
+        elif tensor.ndim != 5:
+            raise ValueError(
+                f"Expected Robometer frames with shape (B,C,H,W) or (B,T,C,H,W); got {tuple(tensor.shape)}"
+            )
+
+        batch_size = tensor.shape[0]
+        tasks = _expand_tasks(
+            complementary.get(self.task_key, self.default_task),
+            batch_size=batch_size,
+            default=self.default_task,
+        )
+
+        samples = [
+            (_video_to_numpy(tensor[i], max_frames=self.max_frames), tasks[i]) for i in range(batch_size)
+        ]
+        encoded = self.encode_samples(samples)
+
+        new_observation = dict(observation)
+        for key, value in encoded.items():
+            new_observation[f"{ROBOMETER_FEATURE_PREFIX}{key}"] = value
+
+        new_transition = transition.copy()
+        new_transition[TransitionKey.OBSERVATION] = new_observation
+        return new_transition
+
+    def encode_samples(self, samples: list[tuple[np.ndarray, str]]) -> dict[str, Tensor]:
+        """Run the Qwen-VL processor on a list of ``(frames, task)`` samples."""
+        from qwen_vl_utils import process_vision_info
+
+        conversations = [self._build_conversation(frames, task) for frames, task in samples]
+
+        texts = [
+            self._processor.apply_chat_template(
+                msg,
+                tokenize=False,
+                add_generation_prompt=False,
+                add_vision_id=True,
+                enable_thinking=False,
+                fps=1,
+            )
+            for msg in conversations
+        ]
+
+        process_kwargs: dict[str, Any] = {
+            "return_video_kwargs": True,
+            "return_video_metadata": True,
+        }
+        image_processor = getattr(self._processor, "image_processor", None)
+        if image_processor is not None and hasattr(image_processor, "patch_size"):
+            process_kwargs["image_patch_size"] = image_processor.patch_size
+
+        image_inputs, video_inputs, video_kwargs = process_vision_info(conversations, **process_kwargs)
+
+        videos: list[Any] | None = None
+        video_metadatas: list[Any] | None = None
+        if video_inputs:
+            if isinstance(video_inputs[0], tuple) and len(video_inputs[0]) == 2:
+                videos_seq, metadatas_seq = zip(*video_inputs, strict=False)
+                videos = list(videos_seq)
+                video_metadatas = list(metadatas_seq)
+            else:
+                videos = list(video_inputs)
+
+        processor_kwargs: dict[str, Any] = {
+            "text": texts,
+            "images": image_inputs,
+            "padding": True,
+            "truncation": False,
+            "max_length": self.max_length,
+            "return_tensors": "pt",
+            "do_resize": False,
+        }
+        if videos is not None:
+            processor_kwargs["videos"] = videos
+        if video_metadatas is not None:
+            processor_kwargs["video_metadata"] = video_metadatas
+        if video_kwargs:
+            processor_kwargs.update(video_kwargs)
+
+        encoded = self._processor(**processor_kwargs)
+
+        # Write Robometer-specific token ids and the video patch merge size into
+        # the encoded batch so `RobometerRewardModel` doesn't need its own
+        # tokenizer at inference (EO1-style separation: the processor owns the
+        # tokenizer, the model owns the backbone and heads).
+        tokenizer = self._processor.tokenizer
+        encoded["prog_token_id"] = tokenizer.convert_tokens_to_ids("<|prog_token|>")
+        encoded["vision_start_token_id"] = tokenizer.convert_tokens_to_ids("<|vision_start|>")
+        encoded["vision_end_token_id"] = tokenizer.convert_tokens_to_ids("<|vision_end|>")
+        video_processor = getattr(self._processor, "video_processor", None)
+        encoded["video_merge_size"] = int(getattr(video_processor, "merge_size", 14))
+        return encoded
+
+    def _build_conversation(self, frames: np.ndarray, task: str) -> list[dict[str, Any]]:
+        pil_frames = _frames_to_pil(frames)
+        prompt = PROGRESS_PROMPT.format(task=task)
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+
+        if self.use_multi_image:
+            for image in pil_frames:
+                content.append({"type": "image", "image": image})
+                if self.use_per_frame_progress_token:
+                    content.append({"type": "text", "text": "<|prog_token|>"})
+        else:
+            content.append({"type": "video", "video": pil_frames, "sample_fps": 1.0})
+
+        return [{"role": "user", "content": content}]
+
+    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 {
+            "base_model_id": self.base_model_id,
+            "image_key": self.image_key,
+            "task_key": self.task_key,
+            "default_task": self.default_task,
+            "max_frames": self.max_frames,
+            "use_multi_image": self.use_multi_image,
+            "use_per_frame_progress_token": self.use_per_frame_progress_token,
+            "max_length": self.max_length,
+        }
+
+
+def make_robometer_pre_post_processors(
+    config: RobometerConfig,
+    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 Robometer's
+    encoder, and moves everything to the configured device. The
+    postprocessor is the identity since Robometer outputs a single reward
+    tensor.
+    """
+    del dataset_stats  # Robometer has its own normalisation inside the Qwen-VL processor.
+
+    preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+        steps=[
+            AddBatchDimensionProcessorStep(),
+            RobometerEncoderProcessorStep(
+                base_model_id=config.base_model_id,
+                image_key=config.image_key,
+                task_key=config.task_key,
+                default_task=config.default_task,
+                max_frames=config.max_frames,
+                use_multi_image=config.use_multi_image,
+                use_per_frame_progress_token=config.use_per_frame_progress_token,
+            ),
+            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,6 +13,8 @@
 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 == "robometer" %}
+ROBOMETER is a general-purpose video-language robotic reward model built on a fine-tuned Qwen3-VL-4B backbone with progress, preference, and success heads. Given a trajectory video and a task description, it predicts dense, frame-level task progress in [0, 1] and frame-level success probabilities for downstream robot learning, including offline RL, online RL, data filtering and retrieval, and automated failure detection.
 {% 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 %}

tests/rewards/test_modeling_robometer.py

@@ -0,0 +1,340 @@
+# 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 Robometer 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.robometer import RobometerConfig
+from lerobot.rewards.robometer.configuration_robometer import ROBOMETER_SPECIAL_TOKENS
+from lerobot.rewards.robometer.modeling_robometer import (
+    ROBOMETER_FEATURE_PREFIX,
+    convert_bins_to_continuous,
+    decode_progress_outputs,
+)
+from tests.utils import skip_if_package_missing
+
+# Length of the fake tokenizer used in `_patch_build`. The deterministic
+# resize target derived in ``RobometerConfig.__post_init__`` is therefore
+# ``_FAKE_TOKENIZER_LEN + len(ROBOMETER_SPECIAL_TOKENS)``.
+_FAKE_TOKENIZER_LEN = 100
+_EXPECTED_RESIZED_VOCAB = _FAKE_TOKENIZER_LEN + len(ROBOMETER_SPECIAL_TOKENS)
+
+
+class _FakeQwenConfig:
+    """Stand-in for a Qwen3-VL config (the `model.config` attribute).
+
+    ``to_dict`` matches HF's ``PretrainedConfig.to_dict`` closely enough for
+    ``RobometerConfig.__post_init__`` to snapshot a meaningful ``vlm_config``
+    into the saved ``config.json`` and for the reload path to round-trip
+    through ``AutoConfig.for_model``.
+    """
+
+    def __init__(self, hidden_dim: int = 8, vocab_size: int = _FAKE_TOKENIZER_LEN) -> None:
+        # `vocab_size` here is the *pre-resize* value the fake backbone advertises.
+        # `__post_init__` is expected to overwrite it with `len(tokenizer) + 5`.
+        self.text_config = SimpleNamespace(hidden_size=hidden_dim, vocab_size=vocab_size)
+        self._hidden_dim = hidden_dim
+        self._vocab_size = vocab_size
+
+    def to_dict(self) -> dict:
+        return {
+            "model_type": "fake_qwen",
+            "text_config": {
+                "hidden_size": self._hidden_dim,
+                "vocab_size": self._vocab_size,
+            },
+        }
+
+
+class _FakeEmbeddings(torch.nn.Module):
+    def __init__(self, num_embeddings: int = _FAKE_TOKENIZER_LEN) -> None:
+        super().__init__()
+        self.num_embeddings = num_embeddings
+
+
+class _FakeBaseModel(torch.nn.Module):
+    """Stand-in for the Qwen3-VL backbone during tests.
+
+    Provides the minimum surface `RobometerRewardModel.__init__` and
+    `_compute_rbm_logits` rely on: a `parameters()` iterator (for dtype +
+    device), a `config.text_config.hidden_size`, a `config.to_dict()` so
+    `_save_pretrained` can snapshot `vlm_config`,
+    `get_input_embeddings()` / `resize_token_embeddings()` so the fresh-init
+    embed resize is a no-op, and a forward that returns a `SimpleNamespace`
+    with a `hidden_states` tuple.
+    """
+
+    def __init__(self, hidden_dim: int = 8) -> None:
+        super().__init__()
+        self._param = torch.nn.Parameter(torch.zeros(1))
+        self.hidden_dim = hidden_dim
+        self.config = _FakeQwenConfig(hidden_dim)
+        self._embeddings = _FakeEmbeddings()
+
+    def get_input_embeddings(self) -> _FakeEmbeddings:
+        return self._embeddings
+
+    def resize_token_embeddings(self, new_size: int) -> None:
+        self._embeddings.num_embeddings = new_size
+
+    def forward(self, **kwargs):  # noqa: ARG002 - intentional kwargs sink
+        input_ids = kwargs["input_ids"]
+        return SimpleNamespace(
+            hidden_states=(torch.zeros(input_ids.shape[0], input_ids.shape[1], self.hidden_dim),),
+            last_hidden_state=torch.zeros(input_ids.shape[0], input_ids.shape[1], self.hidden_dim),
+        )
+
+
+class _FakeTokenizer:
+    """Minimal stand-in for an HF tokenizer.
+
+    ``RobometerConfig.__post_init__`` uses ``len(tokenizer)`` to compute the
+    deterministic resize target ``len(tokenizer) + len(ROBOMETER_SPECIAL_TOKENS)``,
+    so a working ``__len__`` is all we need.
+    """
+
+    def __init__(self, length: int = _FAKE_TOKENIZER_LEN) -> None:
+        self._length = length
+
+    def __len__(self) -> int:
+        return self._length
+
+
+def _patch_build(monkeypatch) -> None:
+    """Stub out the HF AutoX calls so Robometer construction stays cheap in tests.
+
+    Covers (EO-1 style — no model-side override hooks):
+    * ``AutoConfig.from_pretrained`` (config side) — used by
+      ``RobometerConfig.__post_init__`` to snapshot the backbone config.
+    * ``AutoTokenizer.from_pretrained`` (config side) — used by
+      ``__post_init__`` to compute ``len(tokenizer) + 5``.
+    * ``AutoConfig.for_model``                       — used by
+      ``RobometerConfig.vlm_backbone_config`` when rebuilding for ``from_config``.
+    * ``AutoModelForImageTextToText.from_pretrained`` — fresh-training path
+      (``pretrained_path is None``).
+    * ``AutoModelForImageTextToText.from_config``    — checkpoint-reload path
+      (``pretrained_path`` is set).
+    """
+    from lerobot.rewards.robometer import configuration_robometer, modeling_robometer
+
+    monkeypatch.setattr(
+        modeling_robometer.AutoModelForImageTextToText,
+        "from_pretrained",
+        lambda *args, **kwargs: _FakeBaseModel(hidden_dim=8),
+    )
+    monkeypatch.setattr(
+        modeling_robometer.AutoModelForImageTextToText,
+        "from_config",
+        lambda *args, **kwargs: _FakeBaseModel(hidden_dim=8),
+    )
+    monkeypatch.setattr(
+        configuration_robometer.AutoConfig,
+        "for_model",
+        lambda *args, **kwargs: _FakeQwenConfig(hidden_dim=8),
+    )
+    monkeypatch.setattr(
+        configuration_robometer.AutoConfig,
+        "from_pretrained",
+        lambda *args, **kwargs: _FakeQwenConfig(hidden_dim=8),
+    )
+    monkeypatch.setattr(
+        configuration_robometer.AutoTokenizer,
+        "from_pretrained",
+        lambda *args, **kwargs: _FakeTokenizer(length=_FAKE_TOKENIZER_LEN),
+    )
+
+
+def _make_batch(features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Build a `compute_reward`-ready batch using Robometer's namespaced keys."""
+    return {f"{ROBOMETER_FEATURE_PREFIX}{key}": value for key, value in features.items()}
+
+
+@skip_if_package_missing("transformers")
+def test_robometer_config_registered(monkeypatch):
+    _patch_build(monkeypatch)
+    assert "robometer" in RewardModelConfig.get_known_choices()
+    assert RewardModelConfig.get_choice_class("robometer") is RobometerConfig
+    assert isinstance(make_reward_model_config("robometer", device="cpu"), RobometerConfig)
+
+
+def test_robometer_factory_returns_in_tree_class():
+    from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+
+    assert get_reward_model_class("robometer") is RobometerRewardModel
+
+
+def test_convert_bins_to_continuous_returns_expected_values():
+    # Two frames: first peaks at bin 0 (center 0.0), second peaks at bin 9 (center 1.0).
+    bin_logits = torch.full((2, 10), -10.0)
+    bin_logits[0, 0] = 10.0
+    bin_logits[1, -1] = 10.0
+    values = convert_bins_to_continuous(bin_logits)
+    assert values.shape == (2,)
+    assert torch.allclose(values, torch.tensor([0.0, 1.0]), atol=1e-3)
+
+
+def test_decode_progress_outputs_returns_last_frame_values():
+    progress = torch.tensor([[0.1, 0.9], [0.4, 0.6]])
+    success_logits = torch.tensor([[0.0, 5.0], [0.0, -5.0]])
+
+    outputs = decode_progress_outputs(progress, success_logits, is_discrete_mode=False)
+
+    assert outputs["progress_pred"] == [pytest.approx([0.1, 0.9]), pytest.approx([0.4, 0.6])]
+    assert outputs["success_probs"][0][-1] == pytest.approx(torch.sigmoid(torch.tensor(5.0)).item(), abs=1e-3)
+    assert outputs["success_probs"][1][-1] == pytest.approx(
+        torch.sigmoid(torch.tensor(-5.0)).item(), abs=1e-3
+    )
+
+
+def test_decode_progress_outputs_discrete_mode_softmaxes_over_bins():
+    # 2 frames, peaks at bin 0 and bin 9 → continuous predictions 0.0 and 1.0
+    bin_logits = torch.full((1, 2, 10), -10.0)
+    bin_logits[0, 0, 0] = 10.0
+    bin_logits[0, 1, -1] = 10.0
+
+    outputs = decode_progress_outputs(bin_logits, success_logits=None, is_discrete_mode=True)
+
+    assert outputs["success_probs"] == []
+    assert outputs["progress_pred"][0] == pytest.approx([0.0, 1.0], abs=1e-3)
+
+
+@skip_if_package_missing("transformers")
+def test_robometer_post_init_overwrites_vocab_size_with_tokenizer_length(monkeypatch):
+    """``RobometerConfig.__post_init__`` must overwrite the backbone's stale
+    ``text_config.vocab_size`` (which on the real Qwen3-VL config is the
+    padded embedding size, ``151,936``) with ``len(tokenizer) + 5``. This is
+    the contract that makes the published ``Robometer-4B`` checkpoint load
+    byte-equivalently."""
+    _patch_build(monkeypatch)
+
+    cfg = RobometerConfig(device="cpu", progress_loss_type="l2")
+
+    assert cfg.vlm_config["text_config"]["vocab_size"] == _EXPECTED_RESIZED_VOCAB
+
+
+@skip_if_package_missing("transformers")
+def test_robometer_compute_reward_reads_pre_encoded_inputs(monkeypatch):
+    from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+
+    progress = torch.tensor([[0.1, 0.9], [0.4, 0.6]])
+    success_logits = torch.tensor([[0.0, 5.0], [0.0, -5.0]])
+    _patch_build(monkeypatch)
+
+    cfg = RobometerConfig(device="cpu", reward_output="progress", progress_loss_type="l2")
+    model = RobometerRewardModel(cfg)
+    # Bypass the Qwen3-VL forward + head extraction with deterministic logits.
+    monkeypatch.setattr(model, "_compute_rbm_logits", lambda _inputs: (progress, success_logits))
+
+    batch = _make_batch({"input_ids": torch.zeros(2, 2, dtype=torch.long)})
+    rewards = model.compute_reward(batch)
+
+    assert torch.allclose(rewards, torch.tensor([0.9, 0.6]))
+
+
+@skip_if_package_missing("transformers")
+def test_robometer_compute_reward_can_return_binary_success(monkeypatch):
+    from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+
+    progress = torch.tensor([[0.1, 0.9], [0.4, 0.6]])
+    success_logits = torch.tensor([[0.0, 5.0], [0.0, -5.0]])  # sigmoid(5) > 0.5; sigmoid(-5) < 0.5
+    _patch_build(monkeypatch)
+
+    cfg = RobometerConfig(
+        device="cpu",
+        reward_output="success",
+        success_threshold=0.5,
+        progress_loss_type="l2",
+    )
+    model = RobometerRewardModel(cfg)
+    monkeypatch.setattr(model, "_compute_rbm_logits", lambda _inputs: (progress, success_logits))
+
+    batch = _make_batch({"input_ids": torch.zeros(2, 2, dtype=torch.long)})
+    rewards = model.compute_reward(batch)
+
+    assert torch.equal(rewards, torch.tensor([1.0, 0.0]))
+
+
+@skip_if_package_missing("transformers")
+def test_robometer_compute_reward_errors_when_inputs_missing(monkeypatch):
+    from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+
+    _patch_build(monkeypatch)
+
+    cfg = RobometerConfig(device="cpu", progress_loss_type="l2")
+    model = RobometerRewardModel(cfg)
+
+    with pytest.raises(KeyError, match=r"observation\.robometer\.input_ids"):
+        model.compute_reward({})
+
+
+@skip_if_package_missing("transformers")
+def test_robometer_save_pretrained_roundtrips(monkeypatch, tmp_path):
+    """Saving and reloading a Robometer model in LeRobot HF format must produce
+    a single ``model.safetensors`` + ``config.json`` (no Hydra ``config.yaml``),
+    must round-trip user-tunable config fields, and must persist all three
+    prediction heads (``progress_head``, ``success_head``, ``preference_head``)
+    so the published ``Robometer-4B`` checkpoint loads byte-equivalently.
+    """
+    from huggingface_hub.constants import CONFIG_NAME, SAFETENSORS_SINGLE_FILE
+    from safetensors.torch import load_file
+
+    from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel
+
+    _patch_build(monkeypatch)
+    cfg = RobometerConfig(
+        device="cpu",
+        pretrained_path="robometer/Robometer-4B",
+        # Knobs the user might tweak — must survive the round-trip.
+        image_key="observation.images.cam_top",
+        task_key="task",
+        reward_output="success",
+        success_threshold=0.7,
+        progress_loss_type="l2",
+    )
+    model = RobometerRewardModel(cfg)
+    model.save_pretrained(str(tmp_path))
+
+    # Exactly the files LeRobot's HubMixin promises.
+    assert (tmp_path / CONFIG_NAME).exists()
+    assert (tmp_path / SAFETENSORS_SINGLE_FILE).exists()
+    assert not (tmp_path / "config.yaml").exists()  # we want HF-style, not Hydra
+
+    # All three heads must be present in the saved safetensors. The preference
+    # head is unused at inference but the published checkpoint expects its
+    # rows — losing it would silently break weight loading.
+    state = load_file(str(tmp_path / SAFETENSORS_SINGLE_FILE))
+    assert any(k.startswith("progress_head.") for k in state), "progress_head weights missing"
+    assert any(k.startswith("success_head.") for k in state), "success_head weights missing"
+    assert any(k.startswith("preference_head.") for k in state), "preference_head weights missing"
+
+    # Reload from the local directory: no Hub fetch, no YAML overlay. The
+    # base class drives subclass dispatch via the `type` field in config.json.
+    reloaded_cfg = RewardModelConfig.from_pretrained(str(tmp_path))
+    assert isinstance(reloaded_cfg, RobometerConfig)
+    reloaded_cfg.pretrained_path = str(tmp_path)  # mimic lerobot-train's `validate()`
+    reloaded = RobometerRewardModel.from_pretrained(str(tmp_path), config=reloaded_cfg)
+
+    assert reloaded.config.image_key == "observation.images.cam_top"
+    assert reloaded.config.task_key == "task"
+    assert reloaded.config.reward_output == "success"
+    assert reloaded.config.success_threshold == 0.7
+    assert reloaded.config.progress_loss_type == "l2"  # came back from config.json

tests/rewards/test_robometer_processor.py

@@ -0,0 +1,354 @@
+# 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 Robometer's pre-processing helpers and encoder step.
+
+Covers the pure helpers (``_video_to_numpy`` and ``_expand_tasks``) directly,
+and exercises :class:`RobometerEncoderProcessorStep` with a stubbed
+``AutoProcessor`` so we don't need to download Qwen-VL just to test the
+dataclass plumbing (``transform_features`` / ``get_config``).
+
+The full ``__call__`` path that runs ``process_vision_info`` + the Qwen
+processor is intentionally *not* covered here — it is essentially HF glue
+that's exercised by the integration / parity scripts.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.rewards.robometer.processor_robometer import (
+    PROGRESS_PROMPT,
+    _expand_tasks,
+    _frames_to_pil,
+    _video_to_numpy,
+)
+from tests.utils import skip_if_package_missing
+
+
+def _skip_if_robometer_extras_missing(func):
+    """Apply both optional-dependency guards in one shot.
+
+    ``RobometerEncoderProcessorStep.__post_init__`` calls
+    ``require_package("transformers", ...)`` *and*
+    ``require_package("qwen-vl-utils", ...)``, so both need to be present
+    before we can instantiate the step.
+    """
+    func = skip_if_package_missing("qwen-vl-utils", import_name="qwen_vl_utils")(func)
+    func = skip_if_package_missing("transformers")(func)
+    return func
+
+
+# ---------------------------------------------------------------------------
+# _video_to_numpy — pure tensor → uint8 (T, H, W, C) conversion
+# ---------------------------------------------------------------------------
+
+
+def test_video_to_numpy_chw_float_is_converted_to_thwc_uint8():
+    video = torch.rand(4, 3, 8, 8)  # (T, C, H, W) floats in [0, 1]
+    array = _video_to_numpy(video, max_frames=None)
+
+    assert array.shape == (4, 8, 8, 3)
+    assert array.dtype == np.uint8
+    assert array.min() >= 0 and array.max() <= 255
+
+
+def test_video_to_numpy_already_thwc_uint8_passes_through():
+    video = torch.randint(0, 256, (3, 8, 8, 3), dtype=torch.uint8)  # (T, H, W, C)
+    array = _video_to_numpy(video, max_frames=None)
+
+    assert array.shape == (3, 8, 8, 3)
+    assert array.dtype == np.uint8
+
+
+def test_video_to_numpy_max_frames_tail_crops_recent_frames():
+    """``max_frames`` should keep the **last** K frames (most recent)."""
+    video = torch.zeros(10, 3, 4, 4)
+    for t in range(10):
+        video[t] = t / 9.0  # marker: 0 at t=0, ≈1 at t=9
+
+    array = _video_to_numpy(video, max_frames=3)
+
+    assert array.shape == (3, 4, 4, 3)
+    # The first kept frame is t=7 → marker ≈ 7/9 → uint8 ≈ 198
+    assert int(array[0, 0, 0, 0]) == int(round(7 / 9 * 255))
+    # The last kept frame is t=9 → marker = 1.0 → uint8 = 255
+    assert int(array[-1, 0, 0, 0]) == 255
+
+
+def test_video_to_numpy_rejects_3d_input():
+    with pytest.raises(ValueError, match="Expected channel dim"):
+        _video_to_numpy(torch.zeros(4, 8, 8), max_frames=None)
+
+
+def test_video_to_numpy_floats_above_one_pass_through_without_rescaling():
+    """If ``array.max() > 1`` the helper assumes the tensor is already in the
+    [0, 255] range (uint8-as-float), so values pass through unchanged."""
+    video = torch.full((1, 3, 2, 2), 5.0)
+    array = _video_to_numpy(video, max_frames=None)
+
+    assert array.shape == (1, 2, 2, 3)
+    assert int(array.max()) == 5
+
+
+def test_video_to_numpy_clips_very_large_floats_to_uint8_max():
+    """Out-of-uint8-range floats are clipped at 255 before the cast."""
+    video = torch.full((1, 3, 2, 2), 300.0)
+    array = _video_to_numpy(video, max_frames=None)
+
+    assert int(array.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)
+
+
+# ---------------------------------------------------------------------------
+# _frames_to_pil — uint8 (T, H, W, C) → list[PIL.Image]
+# ---------------------------------------------------------------------------
+
+
+def test_frames_to_pil_returns_one_image_per_frame():
+    frames = np.zeros((4, 8, 8, 3), dtype=np.uint8)
+    images = _frames_to_pil(frames)
+
+    assert len(images) == 4
+    assert all(img.size == (8, 8) for img in images)
+
+
+def test_frames_to_pil_casts_floats_to_uint8():
+    frames = np.full((2, 4, 4, 3), 200.0, dtype=np.float32)
+    images = _frames_to_pil(frames)
+
+    assert len(images) == 2
+    # PIL converted from clipped uint8 - sanity check pixel values come through.
+    assert np.asarray(images[0]).dtype == np.uint8
+
+
+def test_frames_to_pil_rejects_non_4d_input():
+    with pytest.raises(ValueError, match=r"\(T,H,W,C\)"):
+        _frames_to_pil(np.zeros((4, 8, 8), dtype=np.uint8))
+
+
+# ---------------------------------------------------------------------------
+# Encoder step plumbing — exercise dataclass surface with a stubbed AutoProcessor
+# ---------------------------------------------------------------------------
+
+
+class _FakeTokenizer:
+    """Tokenizer surface the encoder step touches in ``__post_init__``."""
+
+    def __init__(self) -> None:
+        self.pad_token: str | None = None
+        self.eos_token = "<|endoftext|>"
+        self._vocab: dict[str, int] = {"<|endoftext|>": 0}
+        self.added: list[str] = []
+
+    def get_vocab(self) -> dict[str, int]:
+        return self._vocab
+
+    def add_special_tokens(self, payload: dict[str, Any]) -> int:
+        for token in payload.get("additional_special_tokens", []):
+            if token not in self._vocab:
+                self._vocab[token] = len(self._vocab)
+                self.added.append(token)
+        return len(self.added)
+
+
+class _FakeAutoProcessor:
+    """Stand-in returned by ``AutoProcessor.from_pretrained`` during tests."""
+
+    def __init__(self) -> None:
+        self.tokenizer = _FakeTokenizer()
+        self.image_processor = None
+        self.video_processor = None
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):  # noqa: ARG003
+        return cls()
+
+
+def _build_step(monkeypatch, **overrides):
+    from lerobot.rewards.robometer import processor_robometer
+
+    monkeypatch.setattr(processor_robometer, "AutoProcessor", _FakeAutoProcessor)
+
+    return processor_robometer.RobometerEncoderProcessorStep(**overrides)
+
+
+@_skip_if_robometer_extras_missing
+def test_encoder_step_registers_special_tokens_on_tokenizer(monkeypatch):
+    """``__post_init__`` must register Robometer's five special tokens on the
+    tokenizer that ships with the chosen Qwen-VL checkpoint."""
+    from lerobot.rewards.robometer.configuration_robometer import ROBOMETER_SPECIAL_TOKENS
+
+    step = _build_step(monkeypatch)
+
+    vocab = step._processor.tokenizer.get_vocab()
+    for token in ROBOMETER_SPECIAL_TOKENS:
+        assert token in vocab, f"{token} not registered on the tokenizer"
+
+
+@_skip_if_robometer_extras_missing
+def test_encoder_step_sets_pad_token_to_eos_when_missing(monkeypatch):
+    """Qwen tokenizers ship without a pad token; the step must reuse EOS so
+    batched processing doesn't crash on padding."""
+    step = _build_step(monkeypatch)
+
+    assert step._processor.tokenizer.pad_token == "<|endoftext|>"
+
+
+@_skip_if_robometer_extras_missing
+def test_encoder_step_get_config_roundtrips_user_fields(monkeypatch):
+    """``get_config`` must serialise every user-tunable field — these are what
+    the processor pipeline saves under ``preprocessor_config.json``."""
+    step = _build_step(
+        monkeypatch,
+        base_model_id="Qwen/Qwen3-VL-4B-Instruct",
+        image_key="observation.images.cam_top",
+        task_key="task",
+        default_task="do the thing",
+        max_frames=12,
+        use_multi_image=True,
+        use_per_frame_progress_token=True,
+        max_length=2048,
+    )
+
+    cfg = step.get_config()
+    assert cfg == {
+        "base_model_id": "Qwen/Qwen3-VL-4B-Instruct",
+        "image_key": "observation.images.cam_top",
+        "task_key": "task",
+        "default_task": "do the thing",
+        "max_frames": 12,
+        "use_multi_image": True,
+        "use_per_frame_progress_token": True,
+        "max_length": 2048,
+    }
+
+
+@_skip_if_robometer_extras_missing
+def test_encoder_step_transform_features_is_identity(monkeypatch):
+    """The encoder step writes Qwen tensors into ``observation`` at call time,
+    but it does **not** advertise new typed features at pipeline-build time —
+    the downstream model consumes them via the ``ROBOMETER_FEATURE_PREFIX``
+    namespace, not via the typed feature map.
+    """
+    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_robometer_extras_missing
+def test_encoder_step_build_conversation_inserts_prog_token_per_frame(monkeypatch):
+    """In multi-image mode with per-frame progress tokens, the conversation
+    must alternate ``image`` and ``<|prog_token|>`` text entries, one pair
+    per frame, after the task prompt."""
+    step = _build_step(
+        monkeypatch,
+        use_multi_image=True,
+        use_per_frame_progress_token=True,
+    )
+
+    frames = np.zeros((3, 8, 8, 3), dtype=np.uint8)
+    conversation = step._build_conversation(frames, task="pick up the cube")
+
+    assert len(conversation) == 1 and conversation[0]["role"] == "user"
+    content = conversation[0]["content"]
+
+    # First entry is the task prompt.
+    assert content[0] == {"type": "text", "text": PROGRESS_PROMPT.format(task="pick up the cube")}
+
+    # Then 3 (image, <|prog_token|>) pairs.
+    expected_tail = [
+        item
+        for _ in range(3)
+        for item in (
+            {"type": "image"},  # value asserted below
+            {"type": "text", "text": "<|prog_token|>"},
+        )
+    ]
+    assert len(content) == 1 + len(expected_tail)
+    for got, exp in zip(content[1:], expected_tail, strict=True):
+        assert got["type"] == exp["type"]
+        if exp["type"] == "text":
+            assert got["text"] == exp["text"]
+
+
+@_skip_if_robometer_extras_missing
+def test_encoder_step_build_conversation_video_mode_uses_single_video_entry(monkeypatch):
+    """When ``use_multi_image=False``, frames are bundled into a single
+    ``video`` content entry instead of individual ``image`` entries."""
+    step = _build_step(
+        monkeypatch,
+        use_multi_image=False,
+        use_per_frame_progress_token=False,
+    )
+
+    frames = np.zeros((4, 8, 8, 3), dtype=np.uint8)
+    conversation = step._build_conversation(frames, task="pour the water")
+
+    content = conversation[0]["content"]
+    # Exactly two entries: the prompt and one video entry.
+    assert len(content) == 2
+    assert content[0]["type"] == "text"
+    assert content[1]["type"] == "video"
+    # The video entry carries all four frames.
+    assert len(content[1]["video"]) == 4

uv.lock

@@ -2989,6 +2989,11 @@ rebot = [
     { name = "motorbridge" },
     { name = "motorbridge-smart-servo" },
 ]
+robometer = [
+    { name = "peft" },
+    { name = "qwen-vl-utils" },
+    { name = "transformers" },
+]
 robstride = [
     { name = "python-can" },
 ]
@@ -3146,6 +3151,7 @@ requires-dist = [
     { name = "lerobot", extras = ["peft-dep"], marker = "extra == 'groot'" },
     { name = "lerobot", extras = ["peft-dep"], marker = "extra == 'molmoact2'" },
     { name = "lerobot", extras = ["peft-dep"], marker = "extra == 'peft'" },
+    { name = "lerobot", extras = ["peft-dep"], marker = "extra == 'robometer'" },
     { name = "lerobot", extras = ["peft-dep"], marker = "extra == 'wallx'" },
     { name = "lerobot", extras = ["phone"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["pi"], marker = "extra == 'all'" },
@@ -3163,10 +3169,12 @@ requires-dist = [
     { name = "lerobot", extras = ["pyzmq-dep"], marker = "extra == 'lekiwi'" },
     { name = "lerobot", extras = ["pyzmq-dep"], marker = "extra == 'unitree-g1'" },
     { name = "lerobot", extras = ["qwen-vl-utils-dep"], marker = "extra == 'eo1'" },
+    { name = "lerobot", extras = ["qwen-vl-utils-dep"], marker = "extra == 'robometer'" },
     { name = "lerobot", extras = ["qwen-vl-utils-dep"], marker = "extra == 'sarm'" },
     { name = "lerobot", extras = ["qwen-vl-utils-dep"], marker = "extra == 'wallx'" },
     { name = "lerobot", extras = ["reachy2"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["rebot"], marker = "extra == 'all'" },
+    { name = "lerobot", extras = ["robometer"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["robstride"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["sarm"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'aloha'" },
@@ -3188,6 +3196,7 @@ requires-dist = [
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'multi-task-dit'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'peft'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'pi'" },
+    { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'robometer'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'sarm'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'smolvla'" },
     { name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'topreward'" },
@@ -3258,7 +3267,7 @@ requires-dist = [
     { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
     { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.25.0" },
 ]
-provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "topreward", "xvla", "eo1", "hilserl", "async", "peft", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
+provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "async", "peft", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
 
 [[package]]
 name = "librt"