diff --git a/docs/source/recap_value_function_experiments.mdx b/docs/source/recap_value_function_experiments.mdx new file mode 100644 index 000000000..0554cc545 --- /dev/null +++ b/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= \ + --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= \ + --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= \ + --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= \ + --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= \ + --reward_type= \ + --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. diff --git a/scripts/overfit_vf_variant.py b/scripts/overfit_vf_variant.py new file mode 100644 index 000000000..236f3e0f1 --- /dev/null +++ b/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() diff --git a/src/lerobot/rewards/__init__.py b/src/lerobot/rewards/__init__.py index 91e973afa..88af781b2 100644 --- a/src/lerobot/rewards/__init__.py +++ b/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 diff --git a/src/lerobot/rewards/distributional_value_function/common.py b/src/lerobot/rewards/distributional_value_function/common.py new file mode 100644 index 000000000..0ffaedede --- /dev/null +++ b/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() diff --git a/src/lerobot/rewards/factory.py b/src/lerobot/rewards/factory.py index c791555f8..9e376638d 100644 --- a/src/lerobot/rewards/factory.py +++ b/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: diff --git a/src/lerobot/rewards/nanovlm_value_function/__init__.py b/src/lerobot/rewards/nanovlm_value_function/__init__.py new file mode 100644 index 000000000..bc05b9d77 --- /dev/null +++ b/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", +] diff --git a/src/lerobot/rewards/nanovlm_value_function/configuration_nanovlm_value_function.py b/src/lerobot/rewards/nanovlm_value_function/configuration_nanovlm_value_function.py new file mode 100644 index 000000000..df2853315 --- /dev/null +++ b/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, + ) diff --git a/src/lerobot/rewards/nanovlm_value_function/modeling_nanovlm_value_function.py b/src/lerobot/rewards/nanovlm_value_function/modeling_nanovlm_value_function.py new file mode 100644 index 000000000..a61c4b41b --- /dev/null +++ b/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] diff --git a/src/lerobot/rewards/nanovlm_value_function/processor_nanovlm_value_function.py b/src/lerobot/rewards/nanovlm_value_function/processor_nanovlm_value_function.py new file mode 100644 index 000000000..f456b146e --- /dev/null +++ b/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 diff --git a/src/lerobot/rewards/temporal_siglip_value_function/__init__.py b/src/lerobot/rewards/temporal_siglip_value_function/__init__.py new file mode 100644 index 000000000..5c9d99af1 --- /dev/null +++ b/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", +] diff --git a/src/lerobot/rewards/temporal_siglip_value_function/configuration_temporal_siglip_value_function.py b/src/lerobot/rewards/temporal_siglip_value_function/configuration_temporal_siglip_value_function.py new file mode 100644 index 000000000..16effff8d --- /dev/null +++ b/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, + ) diff --git a/src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py b/src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py new file mode 100644 index 000000000..9c6cfa13c --- /dev/null +++ b/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] diff --git a/src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py b/src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py new file mode 100644 index 000000000..263e3ba43 --- /dev/null +++ b/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 diff --git a/tests/rewards/test_nanovlm_value_function.py b/tests/rewards/test_nanovlm_value_function.py new file mode 100644 index 000000000..6d76ef518 --- /dev/null +++ b/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 diff --git a/tests/rewards/test_temporal_siglip_value_function.py b/tests/rewards/test_temporal_siglip_value_function.py new file mode 100644 index 000000000..a0e40c589 --- /dev/null +++ b/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