feat(rewards): add temporal SigLIP2 and nanoVLM value functions

- add isolated configs, models, processors, and tests
- add shared distributional value utilities
- add VF overfit comparison harness
- add standalone Gemma3 VLM alignment workflow
- keep the committed RECAP baseline unchanged

docs/source/recap_value_function_experiments.mdx

@@ -0,0 +1,111 @@
+# RECAP value-function experiments
+
+All variants use the same `mc_return`, `is_terminal`, 201-bin support, Dirac/HL-Gauss
+targets, metrics, and LeRobot training pipeline. Only the representation backbone changes.
+
+## Current RECAP Gemma3 baseline
+
+```bash
+lerobot-train \
+  --reward_model.type=distributional_value_function \
+  --reward_model.target_method=dirac_delta \
+  --reward_model.device=cuda \
+  --dataset.repo_id=<dataset_repo_id> \
+  --output_dir=outputs/vf_recap_gemma3 \
+  --steps=40000 \
+  --batch_size=8
+```
+
+This initializes SigLIP2 and Gemma3-270M from unimodal checkpoints and creates a
+fresh Gemma3 multimodal connector.
+
+## Temporal SigLIP2
+
+```bash
+lerobot-train \
+  --reward_model.type=temporal_siglip_value_function \
+  --reward_model.history_steps=6 \
+  --reward_model.frame_gap=30 \
+  --reward_model.target_method=dirac_delta \
+  --reward_model.device=cuda \
+  --dataset.repo_id=<dataset_repo_id> \
+  --output_dir=outputs/vf_temporal_siglip2 \
+  --steps=40000 \
+  --batch_size=16
+```
+
+The dataset factory supplies six past-only frames for every observation key.
+The model requires `observation.state` and all configured camera streams.
+
+## nanoVLM-460M
+
+Run a frozen-backbone probe first:
+
+```bash
+lerobot-train \
+  --reward_model.type=nanovlm_value_function \
+  --reward_model.nanovlm_pretrained_path=lusxvr/nanoVLM-460M-8k \
+  --reward_model.freeze_vision_encoder=true \
+  --reward_model.freeze_multimodal_projector=true \
+  --reward_model.freeze_language_model=true \
+  --reward_model.device=cuda \
+  --dataset.repo_id=<dataset_repo_id> \
+  --output_dir=outputs/vf_nanovlm_probe \
+  --steps=5000 \
+  --batch_size=16
+```
+
+Then load the probe checkpoint and selectively fine-tune the projector/decoder
+at a lower learning rate.
+
+## Standalone Gemma3 VLM alignment
+
+The VLM trainer is intentionally separate from LeRobot:
+
+```bash
+cd third_party/nanoVLM
+
+# Projector warmup
+torchrun --standalone --nproc_per_node=4 train_recap_gemma3.py \
+  --output_dir=checkpoints/recap_vlm_warmup \
+  --steps=8000 \
+  --freeze_language_model
+
+# Full multimodal alignment
+torchrun --standalone --nproc_per_node=4 train_recap_gemma3.py \
+  --resume_from_checkpoint=checkpoints/recap_vlm_warmup/final \
+  --output_dir=checkpoints/recap_vlm_aligned \
+  --steps=50000
+```
+
+The output is a standard Hugging Face `Gemma3ForConditionalGeneration`
+checkpoint.
+
+## Aligned RECAP value function (run last)
+
+```bash
+lerobot-train \
+  --reward_model.type=distributional_value_function \
+  --reward_model.vlm_pretrained_path=third_party/nanoVLM/checkpoints/recap_vlm_aligned/final \
+  --reward_model.freeze_vision_encoder=true \
+  --reward_model.device=cuda \
+  --dataset.repo_id=<dataset_repo_id> \
+  --output_dir=outputs/vf_recap_aligned \
+  --steps=40000 \
+  --batch_size=8
+```
+
+## Small-batch verification
+
+Before each full run:
+
+```bash
+uv run python scripts/overfit_vf_variant.py \
+  --dataset_repo_id=<dataset_repo_id> \
+  --reward_type=<distributional_value_function|temporal_siglip_value_function|nanovlm_value_function> \
+  --num_samples=16 \
+  --steps=500
+```
+
+Compare runs using held-out episode NLL/MAE, per-episode return rank correlation,
+terminal success/failure separation, and the matched-versus-shuffled image loss gap.

scripts/overfit_vf_variant.py

@@ -0,0 +1,206 @@
+"""Overfit any distributional VF architecture on a small real-data batch."""
+
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.datasets import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.factory import resolve_delta_timestamps
+from lerobot.rewards.distributional_value_function.configuration_distributional_value_function import (
+    DistributionalVFConfig,
+)
+from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
+    IMAGE_MASK_SUFFIX,
+)
+from lerobot.rewards.factory import make_reward_model, make_reward_pre_post_processors
+from lerobot.rewards.nanovlm_value_function.configuration_nanovlm_value_function import (
+    NanoVLMVFConfig,
+)
+from lerobot.rewards.temporal_siglip_value_function.configuration_temporal_siglip_value_function import (
+    TemporalSiglipVFConfig,
+)
+from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_repo_id", required=True)
+    parser.add_argument("--root", default=None)
+    parser.add_argument(
+        "--reward_type",
+        choices=(
+            "distributional_value_function",
+            "temporal_siglip_value_function",
+            "nanovlm_value_function",
+        ),
+        required=True,
+    )
+    parser.add_argument("--vlm_pretrained_path", default=None)
+    parser.add_argument("--nanovlm_pretrained_path", default="lusxvr/nanoVLM-460M-8k")
+    parser.add_argument("--num_samples", type=int, default=16)
+    parser.add_argument("--steps", type=int, default=500)
+    parser.add_argument("--lr_head", type=float, default=1e-3)
+    parser.add_argument("--lr_backbone", type=float, default=1e-5)
+    parser.add_argument("--history_steps", type=int, default=6)
+    parser.add_argument("--history_frame_gap", type=int, default=30)
+    parser.add_argument("--log_every", type=int, default=25)
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    metadata = LeRobotDatasetMetadata(args.dataset_repo_id, root=args.root)
+    input_features = {
+        key: PolicyFeature(type=FeatureType.VISUAL, shape=tuple(metadata.features[key]["shape"]))
+        for key in metadata.camera_keys
+    }
+    if OBS_STATE in metadata.features:
+        input_features[OBS_STATE] = PolicyFeature(
+            type=FeatureType.STATE,
+            shape=tuple(metadata.features[OBS_STATE]["shape"]),
+        )
+    common = {"input_features": input_features, "device": str(device), "target_method": "dirac_delta"}
+    if args.reward_type == "distributional_value_function":
+        config = DistributionalVFConfig(
+            **common,
+            vlm_pretrained_path=args.vlm_pretrained_path,
+            freeze_vision_encoder=True,
+        )
+    elif args.reward_type == "temporal_siglip_value_function":
+        config = TemporalSiglipVFConfig(
+            **common,
+            history_steps=args.history_steps,
+            frame_gap=args.history_frame_gap,
+        )
+        config.normalization_mapping = {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+        }
+    else:
+        config = NanoVLMVFConfig(
+            **common,
+            nanovlm_pretrained_path=args.nanovlm_pretrained_path,
+        )
+
+    delta_timestamps = resolve_delta_timestamps(config, metadata)
+    dataset = LeRobotDataset(
+        args.dataset_repo_id,
+        root=args.root,
+        delta_timestamps=delta_timestamps,
+        video_backend="pyav",
+    )
+    indices = np.linspace(0, len(dataset) - 1, args.num_samples, dtype=int).tolist()
+    preprocessor, _ = make_reward_pre_post_processors(
+        config,
+        dataset_stats=metadata.stats,
+    )
+
+    processed_samples = []
+    returns = []
+    terminals = []
+    for index in indices:
+        sample = dataset[index]
+        returns.append(torch.as_tensor(sample["mc_return"]).reshape(-1)[0])
+        terminals.append(torch.as_tensor(sample["is_terminal"]).reshape(-1)[0])
+        processed_samples.append(preprocessor(sample))
+    batch = _collate_processed(processed_samples)
+    batch["mc_return"] = torch.stack(returns).to(device)
+    batch["is_terminal"] = torch.stack(terminals).bool().to(device)
+
+    model = make_reward_model(config).to(device)
+    _run_stage(
+        model,
+        batch,
+        steps=args.steps // 2,
+        learning_rate=args.lr_head,
+        head_only=True,
+        log_every=args.log_every,
+        label="head probe",
+    )
+    _run_stage(
+        model,
+        batch,
+        steps=args.steps - args.steps // 2,
+        learning_rate=args.lr_backbone,
+        head_only=False,
+        log_every=args.log_every,
+        label="fine-tune",
+    )
+    _image_shuffle_diagnostic(model, batch, metadata.camera_keys)
+
+
+def _collate_processed(samples):
+    keys = {
+        *[key for key in samples[0] if key.startswith("observation.images.")],
+        OBS_LANGUAGE_TOKENS,
+        OBS_LANGUAGE_ATTENTION_MASK,
+    }
+    if OBS_STATE in samples[0]:
+        keys.add(OBS_STATE)
+    for key in list(keys):
+        if key.startswith("observation.images.") and not key.endswith(IMAGE_MASK_SUFFIX):
+            keys.add(key + IMAGE_MASK_SUFFIX)
+    return {key: torch.cat([sample[key] for sample in samples], dim=0) for key in keys}
+
+
+def _set_trainable(model, *, head_only: bool):
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.value_head.parameters():
+        param.requires_grad = True
+    if hasattr(model, "value_query"):
+        for param in model.value_query.parameters():
+            param.requires_grad = True
+    if head_only:
+        return
+
+    if hasattr(model, "multi_modal_projector"):
+        model.multi_modal_projector.requires_grad_(True)
+        model.language_model.requires_grad_(True)
+    elif hasattr(model, "temporal_transformer"):
+        for name, param in model.named_parameters():
+            if not name.startswith("siglip."):
+                param.requires_grad = True
+    else:
+        model.nanovlm.MP.requires_grad_(True)
+        model.nanovlm.decoder.requires_grad_(True)
+
+
+def _run_stage(model, batch, *, steps, learning_rate, head_only, log_every, label):
+    _set_trainable(model, head_only=head_only)
+    model.train()
+    params = [param for param in model.parameters() if param.requires_grad]
+    optimizer = torch.optim.AdamW(params, lr=learning_rate)
+    print(f"\n{label}: {sum(param.numel() for param in params):,} trainable parameters")
+    for step in range(steps + 1):
+        optimizer.zero_grad(set_to_none=True)
+        loss, metrics = model(batch)
+        loss.backward()
+        optimizer.step()
+        if step % log_every == 0 or step == steps:
+            print(
+                f"step={step:04d} loss={metrics['loss']:.4f} "
+                f"mae={metrics['mae']:.4f} acc={metrics['acc_neighbor']:.3f}"
+            )
+
+
+@torch.no_grad()
+def _image_shuffle_diagnostic(model, batch, camera_keys):
+    model.eval()
+    matched_loss, _ = model(batch)
+    shuffled = dict(batch)
+    permutation = torch.roll(torch.arange(batch["mc_return"].shape[0], device=matched_loss.device), 1)
+    for key in camera_keys:
+        shuffled[key] = batch[key][permutation]
+    shuffled_loss, _ = model(shuffled)
+    print(
+        f"\nvisual dependence: matched_loss={matched_loss.item():.4f} "
+        f"shuffled_loss={shuffled_loss.item():.4f} "
+        f"gap={shuffled_loss.item() - matched_loss.item():+.4f}"
+    )
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/rewards/__init__.py

@@ -22,18 +22,26 @@ from .factory import (
     make_reward_model_config as make_reward_model_config,
     make_reward_pre_post_processors as make_reward_pre_post_processors,
 )
+from .nanovlm_value_function.configuration_nanovlm_value_function import (
+    NanoVLMVFConfig as NanoVLMVFConfig,
+)
 from .pretrained import PreTrainedRewardModel as PreTrainedRewardModel
 from .robometer.configuration_robometer import RobometerConfig as RobometerConfig
 from .sarm.configuration_sarm import SARMConfig as SARMConfig
+from .temporal_siglip_value_function.configuration_temporal_siglip_value_function import (
+    TemporalSiglipVFConfig as TemporalSiglipVFConfig,
+)
 from .topreward.configuration_topreward import TOPRewardConfig as TOPRewardConfig
 
 __all__ = [
     # Configuration classes
     "DistributionalVFConfig",
+    "NanoVLMVFConfig",
     "RewardClassifierConfig",
     "RobometerConfig",
     "SARMConfig",
     "TOPRewardConfig",
+    "TemporalSiglipVFConfig",
     # Base class
     "PreTrainedRewardModel",
     # Factory functions

src/lerobot/rewards/distributional_value_function/common.py

@@ -0,0 +1,120 @@
+"""Shared distributional targets, loss, and metrics for VF experiments."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import Tensor
+
+
+class DistributionalValueMixin:
+    """Mixin for models that expose ``_get_value_readout(batch)``."""
+
+    config: Any
+    value_head: Any
+    hl_gauss_sigma: float
+
+    def hl_gauss_target(self, target_value: Tensor) -> Tensor:
+        target_value = target_value.reshape(-1).clamp(
+            self.config.value_support_min, self.config.value_support_max
+        )
+        target_value = target_value.to(self.value_head.bin_centers.dtype)
+        bin_width = (self.config.value_support_max - self.config.value_support_min) / (
+            self.config.num_value_bins - 1
+        )
+        support_edges = torch.linspace(
+            self.config.value_support_min - bin_width / 2,
+            self.config.value_support_max + bin_width / 2,
+            self.config.num_value_bins + 1,
+            device=target_value.device,
+            dtype=target_value.dtype,
+        )
+        cdf = 0.5 * (
+            1.0
+            + torch.erf((support_edges[None] - target_value[:, None]) / (self.hl_gauss_sigma * math.sqrt(2)))
+        )
+        normalization = (cdf[:, -1] - cdf[:, 0]).unsqueeze(-1).clamp_min(1e-10)
+        return (cdf[:, 1:] - cdf[:, :-1]) / normalization
+
+    def dirac_delta_target(self, target_value: Tensor) -> Tensor:
+        target_value = target_value.reshape(-1).clamp(
+            self.config.value_support_min, self.config.value_support_max
+        )
+        target_value = target_value.to(self.value_head.bin_centers.dtype)
+        bin_width = self.value_head.bin_centers[1] - self.value_head.bin_centers[0]
+        position = (target_value - self.config.value_support_min) / bin_width
+        lower = position.floor().long().clamp(0, self.config.num_value_bins - 1)
+        upper = position.ceil().long().clamp(0, self.config.num_value_bins - 1)
+        weight_upper = position - lower.float()
+        weight_lower = upper.float() - position
+        same = lower == upper
+        weight_upper = torch.where(same, torch.zeros_like(weight_upper), weight_upper)
+        weight_lower = torch.where(same, torch.ones_like(weight_lower), weight_lower)
+        distribution = torch.zeros(
+            target_value.shape[0],
+            self.config.num_value_bins,
+            device=target_value.device,
+            dtype=target_value.dtype,
+        )
+        rows = torch.arange(target_value.shape[0], device=target_value.device)
+        distribution[rows, lower] += weight_lower
+        distribution[rows, upper] += weight_upper
+        return distribution
+
+    def compute_target_distribution(self, target_value: Tensor, is_terminal: Tensor) -> Tensor:
+        if self.config.target_method == "hl_gauss":
+            base = self.hl_gauss_target(target_value)
+        elif self.config.target_method == "dirac_delta":
+            base = self.dirac_delta_target(target_value)
+        else:
+            raise ValueError(f"Unknown target method: {self.config.target_method}")
+        if not self.config.use_one_hot_terminal:
+            return base
+        nearest = torch.argmin(
+            torch.abs(
+                self.value_head.bin_centers[None]
+                - target_value.reshape(-1, 1).to(self.value_head.bin_centers.dtype)
+            ),
+            dim=-1,
+        )
+        terminal = F.one_hot(nearest, num_classes=self.config.num_value_bins).to(base.dtype)
+        return torch.where(is_terminal.reshape(-1, 1).bool(), terminal, base)
+
+    def _distributional_forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
+        readout = self._get_value_readout(batch)
+        logits = self.value_head(readout)
+        probabilities = logits.softmax(-1)
+        predicted_value = (probabilities * self.value_head.bin_centers.to(probabilities.dtype)).sum(-1)
+        targets = self.compute_target_distribution(batch["mc_return"], batch["is_terminal"])
+        loss = -(targets * logits.log_softmax(-1)).sum(-1).mean()
+        target_values = (
+            batch["mc_return"].reshape(-1).clamp(self.config.value_support_min, self.config.value_support_max)
+        )
+        return loss, {
+            "loss": loss.item(),
+            "predicted_value_mean": predicted_value.mean().item(),
+            "mc_return_mean": target_values.mean().item(),
+            "mae": (predicted_value - target_values).abs().mean().item(),
+            "acc_best": (logits.argmax(-1) == targets.argmax(-1)).float().mean().item(),
+            "acc_neighbor": _neighbor_accuracy(
+                logits.argmax(-1),
+                target_values,
+                self.value_head.bin_centers,
+            ),
+        }
+
+    def compute_reward(self, batch: dict[str, Tensor]) -> Tensor:
+        logits = self.value_head(self._get_value_readout(batch))
+        probabilities = logits.softmax(-1)
+        return (probabilities * self.value_head.bin_centers.to(probabilities.dtype)).sum(-1)
+
+
+def _neighbor_accuracy(predicted_bin: Tensor, target: Tensor, centers: Tensor) -> float:
+    width = centers[1] - centers[0]
+    position = (target - centers[0]) / width
+    lower = position.floor().long().clamp(0, len(centers) - 1)
+    upper = position.ceil().long().clamp(0, len(centers) - 1)
+    return ((predicted_bin == lower) | (predicted_bin == upper)).float().mean().item()

src/lerobot/rewards/factory.py

@@ -25,9 +25,13 @@ from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 
 from .classifier.configuration_classifier import RewardClassifierConfig
 from .distributional_value_function.configuration_distributional_value_function import DistributionalVFConfig
+from .nanovlm_value_function.configuration_nanovlm_value_function import NanoVLMVFConfig
 from .pretrained import PreTrainedRewardModel
 from .robometer.configuration_robometer import RobometerConfig
 from .sarm.configuration_sarm import SARMConfig
+from .temporal_siglip_value_function.configuration_temporal_siglip_value_function import (
+    TemporalSiglipVFConfig,
+)
 from .topreward.configuration_topreward import TOPRewardConfig
 
 
@@ -70,6 +74,18 @@ def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
         )
 
         return DistributionalVFRewardModel
+    elif name == "temporal_siglip_value_function":
+        from lerobot.rewards.temporal_siglip_value_function.modeling_temporal_siglip_value_function import (
+            TemporalSiglipVFRewardModel,
+        )
+
+        return TemporalSiglipVFRewardModel
+    elif name == "nanovlm_value_function":
+        from lerobot.rewards.nanovlm_value_function.modeling_nanovlm_value_function import (
+            NanoVLMVFRewardModel,
+        )
+
+        return NanoVLMVFRewardModel
     else:
         try:
             return _get_reward_model_cls_from_name(name=name)
@@ -105,6 +121,10 @@ def make_reward_model_config(reward_type: str, **kwargs) -> RewardModelConfig:
         return TOPRewardConfig(**kwargs)
     elif reward_type == "distributional_value_function":
         return DistributionalVFConfig(**kwargs)
+    elif reward_type == "temporal_siglip_value_function":
+        return TemporalSiglipVFConfig(**kwargs)
+    elif reward_type == "nanovlm_value_function":
+        return NanoVLMVFConfig(**kwargs)
     else:
         try:
             config_cls = RewardModelConfig.get_choice_class(reward_type)
@@ -210,6 +230,24 @@ def make_reward_pre_post_processors(
             config=reward_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
         )
+    elif isinstance(reward_cfg, TemporalSiglipVFConfig):
+        from lerobot.rewards.temporal_siglip_value_function.processor_temporal_siglip_value_function import (
+            make_temporal_siglip_vf_pre_post_processors,
+        )
+
+        return make_temporal_siglip_vf_pre_post_processors(
+            config=reward_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+    elif isinstance(reward_cfg, NanoVLMVFConfig):
+        from lerobot.rewards.nanovlm_value_function.processor_nanovlm_value_function import (
+            make_nanovlm_vf_pre_post_processors,
+        )
+
+        return make_nanovlm_vf_pre_post_processors(
+            config=reward_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
 
     else:
         try:

src/lerobot/rewards/nanovlm_value_function/__init__.py

@@ -0,0 +1,9 @@
+from .configuration_nanovlm_value_function import NanoVLMVFConfig
+from .modeling_nanovlm_value_function import NanoVLMVFRewardModel
+from .processor_nanovlm_value_function import make_nanovlm_vf_pre_post_processors
+
+__all__ = [
+    "NanoVLMVFConfig",
+    "NanoVLMVFRewardModel",
+    "make_nanovlm_vf_pre_post_processors",
+]

src/lerobot/rewards/nanovlm_value_function/configuration_nanovlm_value_function.py

@@ -0,0 +1,46 @@
+"""Configuration for the pretrained nanoVLM-460M value-function experiment."""
+
+from dataclasses import dataclass, field
+
+from lerobot.configs import FeatureType, NormalizationMode
+from lerobot.configs.rewards import RewardModelConfig
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
+
+
+@RewardModelConfig.register_subclass("nanovlm_value_function")
+@dataclass
+class NanoVLMVFConfig(RewardModelConfig):
+    nanovlm_pretrained_path: str = "lusxvr/nanoVLM-460M-8k"
+    nanovlm_code_path: str = "third_party/nanoVLM"
+    tokenizer_path: str = "HuggingFaceTB/SmolLM2-360M-Instruct"
+    image_resolution: tuple[int, int] = (512, 512)
+    tokenizer_max_length: int = 256
+    num_value_bins: int = 201
+    value_support_min: float = -1.0
+    value_support_max: float = 0.0
+    hl_gauss_sigma_ratio: float = 0.75
+    target_method: str = "dirac_delta"
+    use_one_hot_terminal: bool = True
+    value_dropout: float = 0.0
+    freeze_vision_encoder: bool = True
+    freeze_multimodal_projector: bool = True
+    freeze_language_model: bool = True
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {"VISUAL": NormalizationMode.IDENTITY}
+    )
+
+    def validate_features(self) -> None:
+        if not any(feature.type == FeatureType.VISUAL for feature in self.input_features.values()):
+            raise ValueError("NanoVLMVFConfig requires visual input features")
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(lr=1e-4, weight_decay=1e-4, grad_clip_norm=1.0)
+
+    def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
+        return CosineDecayWithWarmupSchedulerConfig(
+            num_warmup_steps=500,
+            num_decay_steps=40000,
+            peak_lr=1e-4,
+            decay_lr=1e-6,
+        )

src/lerobot/rewards/nanovlm_value_function/modeling_nanovlm_value_function.py

@@ -0,0 +1,110 @@
+"""Distributional value head on the pretrained nanoVLM-460M checkpoint."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+
+from lerobot.configs.types import FeatureType
+from lerobot.rewards.distributional_value_function.common import DistributionalValueMixin
+from lerobot.rewards.distributional_value_function.modeling_distributional_value_function import ValueHead
+from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
+    IMAGE_MASK_SUFFIX,
+)
+from lerobot.rewards.pretrained import PreTrainedRewardModel
+from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+
+from .configuration_nanovlm_value_function import NanoVLMVFConfig
+
+
+class NanoVLMVFRewardModel(DistributionalValueMixin, PreTrainedRewardModel):
+    """Use nanoVLM's aligned decoder readout for RECAP return classification."""
+
+    name = "nanovlm_value_function"
+    config_class = NanoVLMVFConfig
+
+    def __init__(self, config: NanoVLMVFConfig, **kwargs):
+        super().__init__(config)
+        self.config = config
+        config.validate_features()
+        self.image_keys = [
+            key for key, feature in config.input_features.items() if feature.type == FeatureType.VISUAL
+        ]
+        code_path = Path(config.nanovlm_code_path)
+        if not code_path.is_absolute():
+            code_path = Path(__file__).resolve().parents[4] / code_path
+        if not code_path.exists():
+            raise FileNotFoundError(f"nanoVLM code not found at {code_path}")
+        if str(code_path) not in sys.path:
+            sys.path.insert(0, str(code_path))
+        from models.vision_language_model import VisionLanguageModel
+
+        self.nanovlm = VisionLanguageModel.from_pretrained(config.nanovlm_pretrained_path)
+        hidden_size = self.nanovlm.cfg.lm_hidden_dim
+        self.value_query = nn.Embedding(1, hidden_size)
+        nn.init.normal_(self.value_query.weight, std=0.02)
+        self.value_head = ValueHead(
+            hidden_size,
+            config.num_value_bins,
+            config.value_support_min,
+            config.value_support_max,
+            config.value_dropout,
+        )
+        bin_width = (config.value_support_max - config.value_support_min) / (config.num_value_bins - 1)
+        self.hl_gauss_sigma = config.hl_gauss_sigma_ratio * bin_width
+        self._set_requires_grad()
+
+    def _set_requires_grad(self):
+        if self.config.freeze_vision_encoder:
+            self.nanovlm.vision_encoder.requires_grad_(False).eval()
+        if self.config.freeze_multimodal_projector:
+            self.nanovlm.MP.requires_grad_(False).eval()
+        if self.config.freeze_language_model:
+            self.nanovlm.decoder.requires_grad_(False).eval()
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        if self.config.freeze_vision_encoder:
+            self.nanovlm.vision_encoder.eval()
+        if self.config.freeze_multimodal_projector:
+            self.nanovlm.MP.eval()
+        if self.config.freeze_language_model:
+            self.nanovlm.decoder.eval()
+        return self
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
+        return self._distributional_forward(batch)
+
+    def _get_value_readout(self, batch: dict[str, Tensor]) -> Tensor:
+        batch_size = batch[OBS_LANGUAGE_TOKENS].shape[0]
+        image_tokens = []
+        image_masks = []
+        for key in self.image_keys:
+            image = batch[key]
+            mask = batch[key + IMAGE_MASK_SUFFIX].bool()
+            features = self.nanovlm.MP(self.nanovlm.vision_encoder(image))
+            image_tokens.append(features * mask[:, None, None].to(features.dtype))
+            image_masks.append(mask[:, None].expand(batch_size, features.shape[1]))
+
+        text_tokens = self.nanovlm.decoder.token_embedding(batch[OBS_LANGUAGE_TOKENS])
+        query = self.value_query(torch.zeros(batch_size, 1, dtype=torch.long, device=text_tokens.device)).to(
+            text_tokens.dtype
+        )
+        inputs = torch.cat([*image_tokens, text_tokens, query], dim=1)
+        attention_mask = torch.cat(
+            [
+                *image_masks,
+                batch[OBS_LANGUAGE_ATTENTION_MASK].bool(),
+                torch.ones(batch_size, 1, dtype=torch.bool, device=text_tokens.device),
+            ],
+            dim=1,
+        )
+        hidden, _ = self.nanovlm.decoder(inputs, attention_mask=attention_mask)
+        return hidden[:, -1]
+
+    def get_optim_params(self):
+        return [parameter for parameter in self.parameters() if parameter.requires_grad]

src/lerobot/rewards/nanovlm_value_function/processor_nanovlm_value_function.py

@@ -0,0 +1,69 @@
+"""Processor for the nanoVLM value-function experiment."""
+
+from typing import Any
+
+import torch
+
+from lerobot.configs import FeatureType
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    TokenizerProcessorStep,
+    batch_to_transition,
+    policy_action_to_transition,
+    transition_to_batch,
+)
+from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
+    DistributionalVFImagePreprocessorStep,
+    DistributionalVFPrepareTaskPromptStep,
+)
+from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+
+from .configuration_nanovlm_value_function import NanoVLMVFConfig
+
+
+def make_nanovlm_vf_pre_post_processors(
+    config: NanoVLMVFConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    image_keys = tuple(
+        key for key, feature in config.input_features.items() if feature.type == FeatureType.VISUAL
+    )
+    preprocessor = PolicyProcessorPipeline(
+        steps=[
+            RenameObservationsProcessorStep(rename_map={}),
+            AddBatchDimensionProcessorStep(),
+            NormalizerProcessorStep(
+                features={**config.input_features, **config.output_features},
+                norm_map=config.normalization_mapping,
+                stats=dataset_stats,
+            ),
+            DistributionalVFImagePreprocessorStep(
+                image_resolution=config.image_resolution,
+                image_keys=image_keys,
+            ),
+            DistributionalVFPrepareTaskPromptStep(),
+            TokenizerProcessorStep(
+                tokenizer_name=config.tokenizer_path,
+                max_length=config.tokenizer_max_length,
+                padding_side="right",
+                padding="max_length",
+            ),
+            DeviceProcessorStep(device=config.device or "cpu"),
+        ],
+        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        to_transition=batch_to_transition,
+        to_output=transition_to_batch,
+    )
+    postprocessor = PolicyProcessorPipeline(
+        name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+        to_transition=policy_action_to_transition,
+    )
+    return preprocessor, postprocessor

src/lerobot/rewards/temporal_siglip_value_function/__init__.py

@@ -0,0 +1,9 @@
+from .configuration_temporal_siglip_value_function import TemporalSiglipVFConfig
+from .modeling_temporal_siglip_value_function import TemporalSiglipVFRewardModel
+from .processor_temporal_siglip_value_function import make_temporal_siglip_vf_pre_post_processors
+
+__all__ = [
+    "TemporalSiglipVFConfig",
+    "TemporalSiglipVFRewardModel",
+    "make_temporal_siglip_vf_pre_post_processors",
+]

src/lerobot/rewards/temporal_siglip_value_function/configuration_temporal_siglip_value_function.py

@@ -0,0 +1,57 @@
+"""Configuration for the experimental temporal SigLIP2 value function."""
+
+from dataclasses import dataclass, field
+
+from lerobot.configs import FeatureType, NormalizationMode
+from lerobot.configs.rewards import RewardModelConfig
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
+
+
+@RewardModelConfig.register_subclass("temporal_siglip_value_function")
+@dataclass
+class TemporalSiglipVFConfig(RewardModelConfig):
+    siglip_path: str = "google/siglip2-so400m-patch14-384"
+    image_resolution: tuple[int, int] = (384, 384)
+    tokenizer_max_length: int = 64
+    history_steps: int = 6
+    frame_gap: int = 30
+    state_key: str = "observation.state"
+    state_dim: int = 32
+    hidden_size: int = 512
+    num_layers: int = 4
+    num_heads: int = 8
+    dropout: float = 0.1
+    num_value_bins: int = 201
+    value_support_min: float = -1.0
+    value_support_max: float = 0.0
+    hl_gauss_sigma_ratio: float = 0.75
+    target_method: str = "dirac_delta"
+    use_one_hot_terminal: bool = True
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+        }
+    )
+
+    @property
+    def observation_delta_indices(self) -> list[int]:
+        return [-self.frame_gap * index for index in range(self.history_steps - 1, -1, -1)]
+
+    def validate_features(self) -> None:
+        if not any(feature.type == FeatureType.VISUAL for feature in self.input_features.values()):
+            raise ValueError("TemporalSiglipVFConfig requires visual input features")
+        if self.state_key not in self.input_features:
+            raise ValueError(f"TemporalSiglipVFConfig requires {self.state_key!r}")
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        return AdamWConfig(lr=1e-4, weight_decay=1e-4, grad_clip_norm=1.0)
+
+    def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
+        return CosineDecayWithWarmupSchedulerConfig(
+            num_warmup_steps=500,
+            num_decay_steps=40000,
+            peak_lr=1e-4,
+            decay_lr=1e-6,
+        )

src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py

@@ -0,0 +1,154 @@
+"""Past-only temporal SigLIP2 distributional value function."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import Tensor, nn
+
+from lerobot.configs.types import FeatureType
+from lerobot.rewards.distributional_value_function.common import DistributionalValueMixin
+from lerobot.rewards.distributional_value_function.modeling_distributional_value_function import ValueHead
+from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
+    IMAGE_MASK_SUFFIX,
+)
+from lerobot.rewards.pretrained import PreTrainedRewardModel
+from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+from .configuration_temporal_siglip_value_function import TemporalSiglipVFConfig
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoModel
+else:
+    AutoModel = None  # type: ignore[assignment]
+
+
+class TemporalSiglipVFRewardModel(DistributionalValueMixin, PreTrainedRewardModel):
+    """Fuse three-camera history, task, and state before causal temporal attention."""
+
+    name = "temporal_siglip_value_function"
+    config_class = TemporalSiglipVFConfig
+
+    def __init__(self, config: TemporalSiglipVFConfig, **kwargs):
+        require_package("transformers", extra="recap")
+        super().__init__(config)
+        self.config = config
+        config.validate_features()
+        self.image_keys = [
+            key for key, feature in config.input_features.items() if feature.type == FeatureType.VISUAL
+        ]
+        self.siglip = AutoModel.from_pretrained(config.siglip_path)
+        self.siglip.requires_grad_(False).eval()
+
+        vision_dim = self.siglip.config.vision_config.hidden_size
+        text_dim = self.siglip.config.text_config.hidden_size
+        hidden_size = config.hidden_size
+        self.camera_proj = nn.Linear(vision_dim, hidden_size)
+        self.camera_embedding = nn.Embedding(len(self.image_keys), hidden_size)
+        self.task_proj = nn.Linear(text_dim, hidden_size)
+        self.state_proj = nn.Linear(config.state_dim, hidden_size)
+        self.frame_fusion = nn.Sequential(
+            nn.LayerNorm((len(self.image_keys) + 2) * hidden_size),
+            nn.Linear((len(self.image_keys) + 2) * hidden_size, hidden_size),
+            nn.GELU(),
+        )
+        layer = nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            nhead=config.num_heads,
+            dim_feedforward=4 * hidden_size,
+            dropout=config.dropout,
+            batch_first=True,
+            activation="gelu",
+        )
+        self.temporal_transformer = nn.TransformerEncoder(
+            layer,
+            num_layers=config.num_layers,
+            norm=nn.LayerNorm(hidden_size),
+        )
+        self.time_embedding = nn.Embedding(config.history_steps, hidden_size)
+        self.value_head = ValueHead(
+            hidden_size,
+            config.num_value_bins,
+            config.value_support_min,
+            config.value_support_max,
+            config.dropout,
+        )
+        bin_width = (config.value_support_max - config.value_support_min) / (config.num_value_bins - 1)
+        self.hl_gauss_sigma = config.hl_gauss_sigma_ratio * bin_width
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        self.siglip.eval()
+        return self
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
+        return self._distributional_forward(batch)
+
+    def _get_value_readout(self, batch: dict[str, Tensor]) -> Tensor:
+        images = [batch[key] for key in self.image_keys]
+        masks = [batch[key + IMAGE_MASK_SUFFIX].bool() for key in self.image_keys]
+        state = batch[self.config.state_key]
+        if state.ndim == 2:
+            state = state[:, None]
+        batch_size, history_steps = images[0].shape[:2]
+        if history_steps != self.config.history_steps:
+            raise ValueError(f"Expected {self.config.history_steps} frames, got {history_steps}")
+
+        camera_tokens = []
+        vision_dtype = next(self.siglip.vision_model.parameters()).dtype
+        for camera_index, (image, mask) in enumerate(zip(images, masks, strict=True)):
+            with torch.no_grad():
+                features = self.siglip.vision_model(
+                    pixel_values=image.flatten(0, 1).to(vision_dtype),
+                    interpolate_pos_encoding=True,
+                    return_dict=True,
+                ).pooler_output
+            features = self.camera_proj(features).unflatten(0, (batch_size, history_steps))
+            camera_ids = torch.full(
+                (batch_size, history_steps),
+                camera_index,
+                dtype=torch.long,
+                device=features.device,
+            )
+            camera_tokens.append(
+                (features + self.camera_embedding(camera_ids)) * mask[..., None].to(features.dtype)
+            )
+
+        with torch.no_grad():
+            task_features = self.siglip.text_model(
+                input_ids=batch[OBS_LANGUAGE_TOKENS],
+                attention_mask=batch[OBS_LANGUAGE_ATTENTION_MASK],
+                return_dict=True,
+            ).pooler_output
+        task_token = self.task_proj(task_features)[:, None].expand(-1, history_steps, -1)
+        state = self._fit_state_dim(state).to(task_token.dtype)
+        state_token = self.state_proj(state)
+        frame_tokens = self.frame_fusion(torch.cat([*camera_tokens, task_token, state_token], -1))
+        frame_tokens = (
+            frame_tokens + self.time_embedding(torch.arange(history_steps, device=frame_tokens.device))[None]
+        )
+
+        causal_mask = torch.triu(
+            torch.ones(history_steps, history_steps, dtype=torch.bool, device=frame_tokens.device),
+            diagonal=1,
+        )
+        frame_valid = torch.stack(masks).any(0)
+        hidden = self.temporal_transformer(
+            frame_tokens,
+            mask=causal_mask,
+            src_key_padding_mask=~frame_valid,
+            is_causal=True,
+        )
+        last_valid = frame_valid.long().sum(-1).sub(1).clamp_min(0)
+        return hidden[torch.arange(batch_size, device=hidden.device), last_valid]
+
+    def _fit_state_dim(self, state: Tensor) -> Tensor:
+        if state.shape[-1] > self.config.state_dim:
+            return state[..., : self.config.state_dim]
+        return F.pad(state, (0, self.config.state_dim - state.shape[-1]))
+
+    def get_optim_params(self):
+        return [parameter for parameter in self.parameters() if parameter.requires_grad]

src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py

@@ -0,0 +1,131 @@
+"""Processor for past-only temporal SigLIP2 value inputs."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor
+
+from lerobot.configs import FeatureType
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TokenizerProcessorStep,
+    batch_to_transition,
+    policy_action_to_transition,
+    transition_to_batch,
+)
+from lerobot.rewards.distributional_value_function.processor_distributional_value_function import (
+    IMAGE_MASK_SUFFIX,
+    DistributionalVFPrepareTaskPromptStep,
+    resize_with_pad_torch,
+)
+from lerobot.types import EnvTransition, TransitionKey
+from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+
+from .configuration_temporal_siglip_value_function import TemporalSiglipVFConfig
+
+
+@ProcessorStepRegistry.register(name="temporal_siglip_vf_image_processor")
+@dataclass
+class TemporalSiglipImageProcessorStep(ProcessorStep):
+    image_resolution: tuple[int, int]
+    image_keys: tuple[str, ...]
+    history_steps: int
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        transition = transition.copy()
+        observation = dict(transition.get(TransitionKey.OBSERVATION, {}))
+        batch_size = self._batch_size(observation)
+        for key in self.image_keys:
+            if key not in observation:
+                height, width = self.image_resolution
+                observation[key] = torch.full((batch_size, self.history_steps, 3, height, width), -1.0)
+                observation[key + IMAGE_MASK_SUFFIX] = torch.zeros(
+                    batch_size, self.history_steps, dtype=torch.bool
+                )
+                continue
+
+            image = observation[key].float()
+            if image.ndim == 4:
+                image = image[:, None]
+            if image.ndim != 5 or image.shape[2] != 3:
+                raise ValueError(f"Expected {key} as [B,T,3,H,W], got {tuple(image.shape)}")
+            batch_size, history = image.shape[:2]
+            image = image.flatten(0, 1).permute(0, 2, 3, 1)
+            image = image * 2.0 - 1.0
+            if image.shape[1:3] != self.image_resolution:
+                image = resize_with_pad_torch(image, *self.image_resolution)
+            observation[key] = image.permute(0, 3, 1, 2).unflatten(0, (batch_size, history))
+            observation[key + IMAGE_MASK_SUFFIX] = torch.ones(
+                batch_size, history, dtype=torch.bool, device=image.device
+            )
+        transition[TransitionKey.OBSERVATION] = observation
+        return transition
+
+    @staticmethod
+    def _batch_size(observation: dict[str, Any]) -> int:
+        for value in observation.values():
+            if isinstance(value, Tensor) and value.ndim >= 2:
+                return value.shape[0]
+        return 1
+
+    def transform_features(self, features):
+        return features
+
+    def get_config(self):
+        return {
+            "image_resolution": self.image_resolution,
+            "image_keys": self.image_keys,
+            "history_steps": self.history_steps,
+        }
+
+
+def make_temporal_siglip_vf_pre_post_processors(
+    config: TemporalSiglipVFConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    image_keys = tuple(
+        key for key, feature in config.input_features.items() if feature.type == FeatureType.VISUAL
+    )
+    preprocessor = PolicyProcessorPipeline(
+        steps=[
+            RenameObservationsProcessorStep(rename_map={}),
+            AddBatchDimensionProcessorStep(),
+            NormalizerProcessorStep(
+                features={**config.input_features, **config.output_features},
+                norm_map=config.normalization_mapping,
+                stats=dataset_stats,
+            ),
+            TemporalSiglipImageProcessorStep(
+                image_resolution=config.image_resolution,
+                image_keys=image_keys,
+                history_steps=config.history_steps,
+            ),
+            DistributionalVFPrepareTaskPromptStep(),
+            TokenizerProcessorStep(
+                tokenizer_name=config.siglip_path,
+                max_length=config.tokenizer_max_length,
+                padding_side="right",
+                padding="max_length",
+            ),
+            DeviceProcessorStep(device=config.device or "cpu"),
+        ],
+        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        to_transition=batch_to_transition,
+        to_output=transition_to_batch,
+    )
+    postprocessor = PolicyProcessorPipeline(
+        name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+        to_transition=policy_action_to_transition,
+    )
+    return preprocessor, postprocessor

tests/rewards/test_nanovlm_value_function.py

@@ -0,0 +1,84 @@
+import sys
+from types import ModuleType, SimpleNamespace
+
+import torch
+from torch import nn
+
+from lerobot.configs import FeatureType, PolicyFeature
+from lerobot.rewards.factory import get_reward_model_class, make_reward_model_config
+from lerobot.rewards.nanovlm_value_function.configuration_nanovlm_value_function import (
+    NanoVLMVFConfig,
+)
+from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+
+CAMERA = "observation.images.top"
+
+
+def test_config_and_factory_registration():
+    config = make_reward_model_config("nanovlm_value_function")
+    assert isinstance(config, NanoVLMVFConfig)
+    assert get_reward_model_class("nanovlm_value_function").__name__ == "NanoVLMVFRewardModel"
+
+
+def test_nanovlm_model_forward(monkeypatch):
+    from lerobot.rewards.nanovlm_value_function.modeling_nanovlm_value_function import (
+        NanoVLMVFRewardModel,
+    )
+
+    class FakeVision(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, image):
+            return torch.ones(image.shape[0], 4, 6)
+
+    class FakeProjector(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.proj = nn.Linear(6, 8)
+
+        def forward(self, features):
+            return self.proj(features)
+
+    class FakeDecoder(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.token_embedding = nn.Embedding(100, 8)
+
+        def forward(self, inputs, attention_mask=None):
+            return inputs, None
+
+    class FakeNano(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.cfg = SimpleNamespace(lm_hidden_dim=8)
+            self.vision_encoder = FakeVision()
+            self.MP = FakeProjector()
+            self.decoder = FakeDecoder()
+
+        @classmethod
+        def from_pretrained(cls, path):
+            return cls()
+
+    fake_module = ModuleType("models.vision_language_model")
+    fake_module.VisionLanguageModel = FakeNano
+    monkeypatch.setitem(sys.modules, "models.vision_language_model", fake_module)
+
+    config = NanoVLMVFConfig(
+        device="cpu",
+        nanovlm_code_path="third_party/nanoVLM",
+    )
+    config.input_features = {CAMERA: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16))}
+    model = NanoVLMVFRewardModel(config)
+    batch = {
+        CAMERA: torch.rand(1, 3, 16, 16),
+        CAMERA + ".mask": torch.ones(1, dtype=torch.bool),
+        OBS_LANGUAGE_TOKENS: torch.ones(1, 4, dtype=torch.long),
+        OBS_LANGUAGE_ATTENTION_MASK: torch.ones(1, 4, dtype=torch.bool),
+        "mc_return": torch.tensor([-0.5]),
+        "is_terminal": torch.tensor([False]),
+    }
+    loss, metrics = model(batch)
+    assert torch.isfinite(loss)
+    assert -1.0 <= metrics["predicted_value_mean"] <= 0.0

tests/rewards/test_temporal_siglip_value_function.py

@@ -0,0 +1,101 @@
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from lerobot.configs import FeatureType, PolicyFeature
+from lerobot.rewards.factory import get_reward_model_class, make_reward_model_config
+from lerobot.rewards.temporal_siglip_value_function.configuration_temporal_siglip_value_function import (
+    TemporalSiglipVFConfig,
+)
+from lerobot.rewards.temporal_siglip_value_function.processor_temporal_siglip_value_function import (
+    TemporalSiglipImageProcessorStep,
+)
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
+
+CAMERAS = ("observation.images.top", "observation.images.left", "observation.images.right")
+
+
+def _config(**kwargs):
+    config = TemporalSiglipVFConfig(
+        device="cpu",
+        hidden_size=8,
+        num_layers=1,
+        num_heads=2,
+        history_steps=2,
+        state_dim=4,
+        **kwargs,
+    )
+    config.input_features = {
+        **{key: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)) for key in CAMERAS},
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(4,)),
+    }
+    return config
+
+
+def test_config_and_factory_registration():
+    config = make_reward_model_config("temporal_siglip_value_function")
+    assert isinstance(config, TemporalSiglipVFConfig)
+    assert get_reward_model_class("temporal_siglip_value_function").__name__ == "TemporalSiglipVFRewardModel"
+
+
+def test_history_offsets_are_past_only():
+    config = TemporalSiglipVFConfig(history_steps=4, frame_gap=10)
+    assert config.observation_delta_indices == [-30, -20, -10, 0]
+
+
+def test_temporal_image_processor():
+    step = TemporalSiglipImageProcessorStep(
+        image_resolution=(32, 32),
+        image_keys=(CAMERAS[0],),
+        history_steps=2,
+    )
+    transition = {
+        TransitionKey.OBSERVATION: {
+            CAMERAS[0]: torch.full((1, 2, 3, 20, 16), 0.5),
+        }
+    }
+    observation = step(transition)[TransitionKey.OBSERVATION]
+    assert observation[CAMERAS[0]].shape == (1, 2, 3, 32, 32)
+    assert observation[CAMERAS[0] + ".mask"].shape == (1, 2)
+
+
+def test_temporal_model_forward(monkeypatch):
+    from lerobot.rewards.temporal_siglip_value_function import (
+        modeling_temporal_siglip_value_function as modeling,
+    )
+
+    class FakeEncoder(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, pixel_values=None, input_ids=None, **kwargs):
+            batch = pixel_values.shape[0] if pixel_values is not None else input_ids.shape[0]
+            return SimpleNamespace(pooler_output=torch.ones(batch, 8))
+
+    class FakeSiglip(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.vision_model = FakeEncoder()
+            self.text_model = FakeEncoder()
+            self.config = SimpleNamespace(
+                vision_config=SimpleNamespace(hidden_size=8),
+                text_config=SimpleNamespace(hidden_size=8),
+            )
+
+    monkeypatch.setattr(modeling.AutoModel, "from_pretrained", lambda *args, **kwargs: FakeSiglip())
+    model = modeling.TemporalSiglipVFRewardModel(_config())
+    batch = {
+        **{key: torch.rand(1, 2, 3, 16, 16) for key in CAMERAS},
+        **{key + ".mask": torch.ones(1, 2, dtype=torch.bool) for key in CAMERAS},
+        OBS_STATE: torch.rand(1, 2, 4),
+        OBS_LANGUAGE_TOKENS: torch.ones(1, 4, dtype=torch.long),
+        OBS_LANGUAGE_ATTENTION_MASK: torch.ones(1, 4, dtype=torch.bool),
+        "mc_return": torch.tensor([-0.5]),
+        "is_terminal": torch.tensor([False]),
+    }
+    loss, metrics = model(batch)
+    assert torch.isfinite(loss)
+    assert -1.0 <= metrics["predicted_value_mean"] <= 0.0