src/lerobot/rewards/robometer/modeling_robometer.py

# 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)