From 2aa7f601cd450e16e11c8e79991db950f94aae81 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Thu, 23 Jul 2026 13:38:07 +0200 Subject: [PATCH] feat(recap): align value function with pi0.6 Gemma3 VLM - use 448px SigLIP2 inputs with 2x2 pooling - adopt native Gemma3 projector and attention masking - support aligned VLM checkpoints - update preprocessing, targets, and tests --- ...iguration_distributional_value_function.py | 37 +-- .../modeling_distributional_value_function.py | 228 ++++++++++++------ ...processor_distributional_value_function.py | 15 +- .../test_distributional_value_function.py | 110 ++++----- 4 files changed, 241 insertions(+), 149 deletions(-) diff --git a/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py index f2fed31ae..812275b53 100644 --- a/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py +++ b/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py @@ -16,17 +16,19 @@ Paper: "π*0.6: a VLA That Learns From Experience" (Physical Intelligence, 2025) https://pi.website/blog/pistar06 +Architecture source of truth: "π0.6 Model Card", Section 2 (Model Design) + https://website.pi-asset.com/pi06star/PI06_model_card.pdf Distributional value function V^{pi_ref}(o_t, l) (Section IV-A). Architecture (~670M params): - Vision: SigLIP2-so400m — 27 layers, 1152-dim, 256 patches/image + Vision: SigLIP2-so400m — 27 layers, 1152-dim, 1024 patches/image at 448px LM: Gemma3-270M — 18 layers, 640-dim - Proj: Linear(1152, 640) fresh init - Head: [CLS] → Linear(640→320) → LN → GELU → Dropout → Linear(320→201) + Proj: 2x2 pool → RMSNorm → Linear(1152, 640), 256 soft tokens/image + Readout: one-way learned value query → 2-layer MLP → 201 bins -Inputs: multi-camera images (3 x 256 patches) + ``"Task: {task}."`` prompt -Targets: MC returns in [-1, 0], cross-entropy on HL-Gauss (default) or Dirac delta +Inputs: multi-camera images (3 x 256 soft tokens) + ``"Task: {task}."`` prompt +Targets: MC returns in [-1, 0], cross-entropy on Dirac delta (default) or HL-Gauss Init: SigLIP2 + Gemma3 from pretrained HF checkpoints; head normal_(std=0.02) """ @@ -43,16 +45,23 @@ class DistributionalVFConfig(RewardModelConfig): """Configuration for RECAP's distributional value function. Predicts V^{pi_ref}(o_t, l) as a categorical distribution over B=201 bins in [-1, 0]. - Trained with cross-entropy on HL-Gauss soft targets (default) or Dirac delta (C51), + Trained with cross-entropy on Dirac delta (C51, default) or HL-Gauss soft targets, with optional one-hot targets for terminal states. - Architecture: monolithic VLM — SigLIP2-so400m (vision) + Gemma3-270M (language), - bidirectional prefix attention, one-way [CLS] readout, 2-layer MLP value head. + Architecture: adapted from the native Gemma3 multimodal VLM design and + scaled to π0.6's ~670M value backbone: + 448px SigLIP2-so400m images are pooled from 1024 patches to 256 soft + tokens, RMS-normalized, projected into Gemma3-270M, and followed by a + one-way learned value-query token. Image tokens attend bidirectionally; + text and the value query remain causal. """ # Backbone pretrained paths - siglip_path: str = "google/siglip2-so400m-patch14-224" + siglip_path: str = "google/siglip2-so400m-patch14-384" gemma3_path: str = "google/gemma-3-270m" + # Optional standard Gemma3ForConditionalGeneration checkpoint produced by + # standalone VLM alignment. When set, it supplies vision, connector, and LM. + vlm_pretrained_path: str | None = None # Distributional head num_value_bins: int = 201 @@ -60,14 +69,15 @@ class DistributionalVFConfig(RewardModelConfig): value_support_max: float = 0.0 hl_gauss_sigma_ratio: float = 5.0 - # Target distribution method: "hl_gauss" (default, soft) or "dirac_delta" (C51, hard) - target_method: str = "hl_gauss" + # Target distribution method: "dirac_delta" (paper-faithful C51) or "hl_gauss" (soft) + target_method: str = "dirac_delta" # Whether to use one-hot targets for terminal states (exact return, no smoothing). use_one_hot_terminal: bool = True # Image - image_resolution: tuple[int, int] = (224, 224) + image_resolution: tuple[int, int] = (448, 448) + num_image_tokens: int = 256 # Tokenizer (uses Gemma3's tokenizer) tokenizer_max_length: int = 200 @@ -79,9 +89,6 @@ class DistributionalVFConfig(RewardModelConfig): stop_gradient_to_vlm: bool = False vision_encoder_lr_multiplier: float = 0.5 - # Readout: "mean_pool" (average all tokens) or "last_token" (causal LM last position) - readout: str = "mean_pool" - # Normalization normalization_mapping: dict[str, NormalizationMode] = field( default_factory=lambda: { diff --git a/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py index e423ad112..53427b95d 100644 --- a/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py +++ b/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py @@ -16,18 +16,19 @@ Paper: "π*0.6: a VLA That Learns From Experience" (Physical Intelligence, 2025) https://pi.website/blog/pistar06 +Architecture source of truth: "π0.6 Model Card", Section 2 (Model Design) + https://website.pi-asset.com/pi06star/PI06_model_card.pdf Implements the distributional value function V^{pi_ref}(o_t, l) from Section IV-A. -Architecture: the paper uses a 670M-parameter Gemma 3 VLM (Figure 3) — -SigLIP2-so400m (27 layers, 1152-dim) + Gemma3-270M (18 layers, 640-dim), -with a [CLS] token readout predicting a categorical distribution over -B=201 discrete value bins in [-1, 0]. This implementation uses a 2-layer -MLP value head (Linear→LN→GELU→Dropout→Linear) inspired by Robometer -(Chen et al., 2025). +It adapts the native Gemma3 multimodal VLM design to π0.6's smaller ~670M scale: +448px SigLIP images are pooled to 256 soft tokens, RMS-normalized, projected +into Gemma3-270M, and processed with bidirectional image / causal text +attention. A final learned value-query token supplies the 201-bin readout. """ from __future__ import annotations +import copy import math from typing import TYPE_CHECKING, Any @@ -40,7 +41,6 @@ from lerobot.rewards.pretrained import PreTrainedRewardModel from lerobot.utils.constants import ( OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, - OPENPI_ATTENTION_MASK_VALUE as _ATTENTION_MASK_VALUE, ) from lerobot.utils.import_utils import _transformers_available, require_package @@ -48,16 +48,33 @@ from .configuration_distributional_value_function import DistributionalVFConfig from .processor_distributional_value_function import IMAGE_MASK_SUFFIX if TYPE_CHECKING or _transformers_available: - from transformers import Gemma3ForCausalLM, SiglipVisionModel + from transformers import ( + Gemma3Config, + Gemma3ForCausalLM, + Gemma3ForConditionalGeneration, + SiglipVisionModel, + ) + from transformers.models.gemma3.modeling_gemma3 import ( + Gemma3MultiModalProjector, + create_causal_mask_mapping, + ) else: + Gemma3Config = None # type: ignore[assignment] Gemma3ForCausalLM = None # type: ignore[assignment] + Gemma3ForConditionalGeneration = None # type: ignore[assignment] + Gemma3MultiModalProjector = None # type: ignore[assignment] SiglipVisionModel = None # type: ignore[assignment] + create_causal_mask_mapping = None # type: ignore[assignment] class ValueHead(nn.Module): """Categorical value projection: hidden state → bin logits. - 2-layer MLP: Linear → LayerNorm → GELU → Dropout → Linear. + The 2-layer MLP topology is adapted from Robometer's prediction head: + Linear → LayerNorm → GELU → Dropout → Linear. Unlike Robometer's progress + and success heads, this head predicts RECAP's 201-bin MC-return + distribution over [-1, 0] from the final value-query representation. + Also holds the ``bin_centers`` buffer used to compute E[V] = Σ p_i · c_i. """ @@ -96,12 +113,11 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): Trained with cross-entropy on HL-Gauss or Dirac delta targets centered on per-task normalized Monte Carlo returns. - Architecture: SigLIP2-so400m + Linear(1152→640) + Gemma3-270M. - Multi-camera images are encoded by SigLIP2 (256 patches each), projected to - Gemma3's hidden dim, concatenated with tokenized language, and processed by - all 18 Gemma3 transformer layers. - - Mean-pooled last-layer hidden states are read out through a 2-layer MLP value head. + Architecture: adapted from the native Gemma3 multimodal VLM using a + 448px SigLIP2-so400m vision tower, Gemma3 multimodal projector, and + Gemma3-270M language backbone. Each camera is represented by 256 soft + image tokens. Image tokens are bidirectional, text is causal, and a final + one-way value query supplies the hidden state consumed by the value head. """ name = "distributional_value_function" @@ -112,18 +128,42 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): super().__init__(config) self.config = config - self.vision_encoder = SiglipVisionModel.from_pretrained(config.siglip_path) - siglip_hidden = self.vision_encoder.config.hidden_size # 1152 + if config.vlm_pretrained_path: + aligned_vlm = Gemma3ForConditionalGeneration.from_pretrained(config.vlm_pretrained_path) + self._vlm_config = aligned_vlm.config + self.vision_encoder = aligned_vlm.model.vision_tower + self.multi_modal_projector = aligned_vlm.model.multi_modal_projector + self.language_model = aligned_vlm.model.language_model + else: + self.vision_encoder = SiglipVisionModel.from_pretrained(config.siglip_path) + gemma3 = Gemma3ForCausalLM.from_pretrained(config.gemma3_path) + self.language_model = gemma3.model - self.gemma3 = Gemma3ForCausalLM.from_pretrained(config.gemma3_path) - self.gemma3_hidden = self.gemma3.config.hidden_size # 640 + # Adapt Gemma3's native multimodal connector to the 270M text + # backbone and π0.6's 448px input layout. + vision_config = copy.deepcopy(self.vision_encoder.config) + vision_config.image_size = config.image_resolution[0] + text_config = copy.deepcopy(gemma3.config) + self._vlm_config = Gemma3Config( + vision_config=vision_config, + text_config=text_config, + mm_tokens_per_image=config.num_image_tokens, + ) + self.multi_modal_projector = Gemma3MultiModalProjector(self._vlm_config) + nn.init.normal_( + self.multi_modal_projector.mm_input_projection_weight, + mean=0.0, + std=self._vlm_config.initializer_range, + ) - # Fresh image projection: SigLIP2 1152-dim → Gemma3 640-dim - self.image_proj = nn.Linear(siglip_hidden, self.gemma3_hidden, bias=True) - nn.init.normal_(self.image_proj.weight, std=0.02) - nn.init.zeros_(self.image_proj.bias) + self._validate_vlm_config() + self.gemma3_hidden = self._vlm_config.text_config.hidden_size # 640 - # Value head: last-token hidden state → MLP → num_bins logits + # One-way suffix query, analogous to PI05's suffix/action tokens. + self.value_query = nn.Embedding(1, self.gemma3_hidden) + nn.init.normal_(self.value_query.weight, std=0.02) + + # Value head: value-query hidden state → MLP → num_bins logits self.value_head = ValueHead( hidden_size=self.gemma3_hidden, num_bins=config.num_value_bins, @@ -139,6 +179,35 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): # Apply freezing self._set_requires_grad() + @property + def gemma3(self) -> nn.Module: + """Backward-compatible access to the Gemma3 text backbone.""" + return self.language_model + + def _validate_vlm_config(self) -> None: + """Validate the π0.6 448px → 256-token multimodal layout.""" + vision_config = self._vlm_config.vision_config + image_size = self.config.image_resolution[0] + if vision_config.image_size != image_size: + raise ValueError( + f"VLM vision image_size ({vision_config.image_size}) does not match " + f"DistributionalVFConfig.image_resolution ({image_size})" + ) + + patches_per_side = image_size // vision_config.patch_size + tokens_per_side = int(self.config.num_image_tokens**0.5) + if tokens_per_side**2 != self.config.num_image_tokens: + raise ValueError("num_image_tokens must be a perfect square") + if patches_per_side % tokens_per_side: + raise ValueError( + f"{patches_per_side} patches/side cannot be evenly pooled to {tokens_per_side} tokens/side" + ) + if self._vlm_config.mm_tokens_per_image != self.config.num_image_tokens: + raise ValueError( + f"VLM emits {self._vlm_config.mm_tokens_per_image} image tokens, " + f"expected {self.config.num_image_tokens}" + ) + def _set_requires_grad(self) -> None: if self.config.freeze_vision_encoder: for param in self.vision_encoder.parameters(): @@ -146,16 +215,16 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): self.vision_encoder.eval() if self.config.freeze_language_model: - for param in self.gemma3.parameters(): + for param in self.language_model.parameters(): param.requires_grad = False - self.gemma3.eval() + self.language_model.eval() def train(self, mode: bool = True): super().train(mode) if self.config.freeze_vision_encoder: self.vision_encoder.eval() if self.config.freeze_language_model: - self.gemma3.eval() + self.language_model.eval() return self def get_optim_params(self) -> list[dict]: @@ -178,18 +247,25 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): ] def embed_image(self, image: Tensor) -> Tensor: - """Embed images: SigLIP2 → projection → [B, num_patches, gemma3_hidden]. + """Embed images with π0.6's Gemma3 visual connector. Args: image: [batch_size, channels, height, width] preprocessed image in [-1, 1]. Returns: - [B, 256, gemma3_hidden] projected image features. + [B, 256, gemma3_hidden] pooled, normalized, projected features. """ - if image.dtype != torch.float32: - image = image.to(torch.float32) - feats = self.vision_encoder(pixel_values=image).last_hidden_state - return self.image_proj(feats) + vision_dtype = next(self.vision_encoder.parameters()).dtype + image_features = self.vision_encoder( + pixel_values=image.to(dtype=vision_dtype), + interpolate_pos_encoding=True, + ).last_hidden_state + projected_features = self.multi_modal_projector(image_features) + if projected_features.shape[1] != self.config.num_image_tokens: + raise RuntimeError( + f"Expected {self.config.num_image_tokens} image tokens, got {projected_features.shape[1]}" + ) + return projected_features def embed_text(self, token_ids: Tensor) -> Tensor: """Embed text using Gemma3's embedding table (includes sqrt(d) scaling). @@ -200,7 +276,7 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): Returns: [B, seq_len, gemma3_hidden] text embeddings. """ - return self.gemma3.model.embed_tokens(token_ids) + return self.language_model.embed_tokens(token_ids) def embed_prefix( self, @@ -208,28 +284,30 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): img_masks: list[Tensor], text_embeddings: Tensor, text_padding_mask: Tensor, - ) -> tuple[Tensor, Tensor]: - """Build prefix: [img1_patches, img2_patches, ..., lang_tokens]. - - All prefix tokens use bidirectional attention (att_mask=0). + ) -> tuple[Tensor, Tensor, Tensor]: + """Build [image soft tokens..., text] plus masks. Returns: embs: [B, total_prefix_len, hidden_dim] pad_masks: [B, total_prefix_len] boolean + token_type_ids: [B, total_prefix_len], 1=image and 0=text """ embs: list[Tensor] = [] pad_masks: list[Tensor] = [] + token_types: list[Tensor] = [] for img, img_mask in zip(images, img_masks, strict=True): img_emb = self.embed_image(img) - bsize, num_patches = img_emb.shape[:2] + bsize, num_image_tokens = img_emb.shape[:2] embs.append(img_emb) - pad_masks.append(img_mask[:, None].expand(bsize, num_patches)) + pad_masks.append(img_mask[:, None].expand(bsize, num_image_tokens)) + token_types.append(torch.ones(bsize, num_image_tokens, dtype=torch.long, device=img_emb.device)) embs.append(text_embeddings) pad_masks.append(text_padding_mask) + token_types.append(torch.zeros_like(text_padding_mask, dtype=torch.long)) - return torch.cat(embs, dim=1), torch.cat(pad_masks, dim=1) + return torch.cat(embs, dim=1), torch.cat(pad_masks, dim=1), torch.cat(token_types, dim=1) def hl_gauss_target(self, target_value: Tensor) -> Tensor: """HL-Gauss soft target distribution. @@ -371,50 +449,64 @@ class DistributionalVFRewardModel(PreTrainedRewardModel): terminal_distribution = self.one_hot_target(target_value) return torch.where(is_terminal[:, None].bool(), terminal_distribution, base_distribution) - def _vlm_forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]: - """Shared VLM forward: images + text → Gemma3 → last-token hidden → logits. - - Returns: - (value_logits [B, num_bins], predicted_value [B, 1]) - """ + def _get_vlm_readout(self, batch: dict[str, Tensor]) -> Tensor: + """Run Gemma3 image-bidirectional/text-causal attention plus value query.""" images, img_masks, token_ids, text_pad_mask = self._get_model_inputs(batch) text_embs = self.embed_text(token_ids) - prefix_embs, prefix_pad_masks = self.embed_prefix(images, img_masks, text_embs, text_pad_mask) + prefix_embs, prefix_pad_masks, prefix_token_types = self.embed_prefix( + images, img_masks, text_embs, text_pad_mask + ) if self.config.stop_gradient_to_vlm: prefix_embs = prefix_embs.detach() + batch_size, prefix_len = prefix_pad_masks.shape device = prefix_embs.device - model_dtype = next(self.gemma3.parameters()).dtype + model_dtype = next(self.language_model.parameters()).dtype - # Bidirectional attention: every valid token attends to every valid token - att_2d = prefix_pad_masks[:, None, :] * prefix_pad_masks[:, :, None] - att_4d = torch.where( - att_2d[:, None, :, :], - torch.tensor(0.0, dtype=model_dtype, device=device), - torch.tensor(_ATTENTION_MASK_VALUE, dtype=model_dtype, device=device), + query_ids = torch.zeros(batch_size, 1, dtype=torch.long, device=device) + query_emb = self.value_query(query_ids) + hidden_states = torch.cat([prefix_embs, query_emb], dim=1) + + query_pad_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=device) + pad_masks = torch.cat([prefix_pad_masks, query_pad_mask], dim=1) + token_type_ids = torch.cat( + [prefix_token_types, torch.zeros(batch_size, 1, dtype=torch.long, device=device)], + dim=1, ) - position_ids = torch.cumsum(prefix_pad_masks.long(), dim=1) - 1 + prefix_position_ids = torch.cumsum(prefix_pad_masks.long(), dim=1) - 1 + query_position_ids = prefix_pad_masks.sum(dim=1, keepdim=True).long() + position_ids = torch.cat([prefix_position_ids, query_position_ids], dim=1).clamp_min(0) - if prefix_embs.dtype != model_dtype: - prefix_embs = prefix_embs.to(model_dtype) + if hidden_states.dtype != model_dtype: + hidden_states = hidden_states.to(model_dtype) - outputs = self.gemma3.model( - inputs_embeds=prefix_embs, - attention_mask=att_4d, + attention_masks = create_causal_mask_mapping( + self._vlm_config, + inputs_embeds=hidden_states, + attention_mask=pad_masks, + past_key_values=None, position_ids=position_ids, + token_type_ids=token_type_ids, + is_training=self.training, ) + outputs = self.language_model( + inputs_embeds=hidden_states, + attention_mask=attention_masks, + position_ids=position_ids, + use_cache=False, + ) + return outputs.last_hidden_state[:, -1, :] - # Readout from last hidden layer - if self.config.readout == "mean_pool": - hidden = outputs.last_hidden_state - mask = prefix_pad_masks.unsqueeze(-1).to(dtype=hidden.dtype) - readout = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1) - else: - readout = outputs.last_hidden_state[:, -1, :] + def _vlm_forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]: + """Run the VLM and value head. + Returns: + (value_logits [B, num_bins], predicted_value [B, 1]) + """ + readout = self._get_vlm_readout(batch) value_logits = self.value_head(readout) value_probs = F.softmax(value_logits, dim=-1) predicted_value = (value_probs * self.value_head.bin_centers.to(dtype=value_probs.dtype)).sum( diff --git a/src/lerobot/rewards/distributional_value_function/processor_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/processor_distributional_value_function.py index c6f938da2..b9a86aae1 100644 --- a/src/lerobot/rewards/distributional_value_function/processor_distributional_value_function.py +++ b/src/lerobot/rewards/distributional_value_function/processor_distributional_value_function.py @@ -18,7 +18,7 @@ Paper: "π*0.6: a VLA That Learns From Experience" (Physical Intelligence, 2025) https://pi.website/blog/pistar06 Prepares inputs for V^{pi_ref}(o_t, l): -1. Resize multi-camera images to 224x224 (with aspect-preserving padding) +1. Resize multi-camera images to 448x448 (with aspect-preserving padding) 2. Normalize images from [0,1] → [-1,1] (SigLIP standard) 3. Handle missing cameras (placeholder + mask) 4. Format task prompt: ``"Task: {task}."`` @@ -132,12 +132,13 @@ def resize_with_pad_torch( class DistributionalVFImagePreprocessorStep(ProcessorStep): """Resize and normalize multi-camera images for the VF. + Expects LeRobot's standard float image range [0, 1]. Produces [B, 3, H, W] tensors in [-1, 1] for each camera, plus boolean masks indicating which cameras are present. Missing cameras get a black placeholder image and mask=False. """ - image_resolution: tuple[int, int] = (224, 224) + image_resolution: tuple[int, int] = (448, 448) image_keys: tuple[str, ...] = () def __call__(self, transition: EnvTransition) -> EnvTransition: @@ -154,12 +155,12 @@ class DistributionalVFImagePreprocessorStep(ProcessorStep): if is_channels_first: img = img.permute(0, 2, 3, 1) # BCHW → BHWC + # Gemma3's SigLIP vision tower expects [-1, 1]. + img = img * 2.0 - 1.0 + if img.shape[1:3] != self.image_resolution: img = resize_with_pad_torch(img, *self.image_resolution) - if img.min() >= 0.0 and img.max() <= 1.0: - img = img * 2.0 - 1.0 - observation[key] = img.permute(0, 3, 1, 2) # BHWC → BCHW observation[key + IMAGE_MASK_SUFFIX] = torch.ones( img.shape[0], dtype=torch.bool, device=img.device @@ -240,7 +241,7 @@ def make_distributional_vf_pre_post_processors( 1. Rename observations (no-op by default) 2. Add a batch dimension 3. Normalize features (identity for images) - 4. Resize + normalize images → [B, 3, 224, 224] in [-1, 1] + 4. Resize + normalize images → [B, 3, 448, 448] in [-1, 1] 5. Format task prompt: ``"Task: {task}."`` 6. Tokenize with Gemma3 tokenizer 7. Move tensors to the configured device @@ -265,7 +266,7 @@ def make_distributional_vf_pre_post_processors( ), DistributionalVFPrepareTaskPromptStep(), TokenizerProcessorStep( - tokenizer_name=config.gemma3_path, + tokenizer_name=config.vlm_pretrained_path or config.gemma3_path, max_length=config.tokenizer_max_length, padding_side="right", padding="max_length", diff --git a/tests/rewards/test_distributional_value_function.py b/tests/rewards/test_distributional_value_function.py index 2519fe7e2..77a383ecf 100644 --- a/tests/rewards/test_distributional_value_function.py +++ b/tests/rewards/test_distributional_value_function.py @@ -28,8 +28,9 @@ from lerobot.types import TransitionKey from lerobot.utils.constants import OBS_IMAGES from tests.utils import skip_if_package_missing -BATCH_SIZE = 4 +BATCH_SIZE = 1 NUM_BINS = 201 +IMAGE_SIZE = 448 IMAGE_KEY = f"{OBS_IMAGES}.top" IMAGE_KEY_WRIST_LEFT = f"{OBS_IMAGES}.wrist_left" IMAGE_KEY_WRIST_RIGHT = f"{OBS_IMAGES}.wrist_right" @@ -38,14 +39,14 @@ IMAGE_KEY_WRIST_RIGHT = f"{OBS_IMAGES}.wrist_right" def _make_config(**overrides) -> DistributionalVFConfig: defaults = { "device": "cpu", - "image_resolution": (224, 224), + "image_resolution": (IMAGE_SIZE, IMAGE_SIZE), } defaults.update(overrides) config = DistributionalVFConfig(**defaults) config.input_features = { - IMAGE_KEY: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)), - IMAGE_KEY_WRIST_LEFT: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)), - IMAGE_KEY_WRIST_RIGHT: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)), + IMAGE_KEY: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_SIZE, IMAGE_SIZE)), + IMAGE_KEY_WRIST_LEFT: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_SIZE, IMAGE_SIZE)), + IMAGE_KEY_WRIST_RIGHT: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_SIZE, IMAGE_SIZE)), } config.output_features = {} config.normalization_mapping = { @@ -69,11 +70,11 @@ def _make_batch(batch_size: int = BATCH_SIZE, device: str = "cpu") -> dict[str, from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS return { - IMAGE_KEY: torch.rand(batch_size, 3, 224, 224, device=device) * 2 - 1, + IMAGE_KEY: torch.rand(batch_size, 3, IMAGE_SIZE, IMAGE_SIZE, device=device) * 2 - 1, IMAGE_KEY + IMAGE_MASK_SUFFIX: torch.ones(batch_size, dtype=torch.bool, device=device), - IMAGE_KEY_WRIST_LEFT: torch.rand(batch_size, 3, 224, 224, device=device) * 2 - 1, + IMAGE_KEY_WRIST_LEFT: torch.rand(batch_size, 3, IMAGE_SIZE, IMAGE_SIZE, device=device) * 2 - 1, IMAGE_KEY_WRIST_LEFT + IMAGE_MASK_SUFFIX: torch.ones(batch_size, dtype=torch.bool, device=device), - IMAGE_KEY_WRIST_RIGHT: torch.rand(batch_size, 3, 224, 224, device=device) * 2 - 1, + IMAGE_KEY_WRIST_RIGHT: torch.rand(batch_size, 3, IMAGE_SIZE, IMAGE_SIZE, device=device) * 2 - 1, IMAGE_KEY_WRIST_RIGHT + IMAGE_MASK_SUFFIX: torch.ones(batch_size, dtype=torch.bool, device=device), OBS_LANGUAGE_TOKENS: torch.randint(0, 1000, (batch_size, 200), device=device), OBS_LANGUAGE_ATTENTION_MASK: torch.ones(batch_size, 200, dtype=torch.bool, device=device), @@ -114,6 +115,13 @@ def test_make_reward_model_config_factory(): assert config.num_value_bins == 101 +def test_config_defaults_match_pi06_gemma3_layout(): + config = DistributionalVFConfig() + assert config.image_resolution == (448, 448) + assert config.num_image_tokens == 256 + assert config.target_method == "dirac_delta" + + # ------------------------------------------------------------------ # Target distribution tests (HL-Gauss, Dirac delta, one-hot) # ------------------------------------------------------------------ @@ -276,9 +284,9 @@ def test_model_has_expected_components(): assert hasattr(model, "vision_encoder") assert hasattr(model, "gemma3") - assert hasattr(model, "image_proj") + assert hasattr(model, "multi_modal_projector") assert hasattr(model, "value_head") - assert hasattr(model, "cls_embedding") + assert hasattr(model, "value_query") assert hasattr(model.value_head, "mlp") assert hasattr(model.value_head, "bin_centers") @@ -298,23 +306,26 @@ def test_value_head_output_dim(): @skip_if_package_missing("transformers") -def test_cls_embedding_is_nn_embedding(): - """CLS is nn.Embedding (FSDP-safe) with correct shape.""" +def test_value_query_is_nn_embedding(): + """Value query is nn.Embedding (FSDP-safe) with correct shape.""" model = _make_model() from torch import nn - assert isinstance(model.cls_embedding, nn.Embedding) - assert model.cls_embedding.num_embeddings == 1 - assert model.cls_embedding.embedding_dim == model.gemma3_hidden + assert isinstance(model.value_query, nn.Embedding) + assert model.value_query.num_embeddings == 1 + assert model.value_query.embedding_dim == model.gemma3_hidden @skip_if_package_missing("transformers") -def test_image_proj_dimensions(): - """Image projection maps SigLIP2 hidden to Gemma3 hidden.""" +def test_multimodal_projector_dimensions_and_pooling(): + """Gemma3 connector pools 448px patches to 256 tokens and projects to LM width.""" model = _make_model() siglip_hidden = model.vision_encoder.config.hidden_size - assert model.image_proj.in_features == siglip_hidden - assert model.image_proj.out_features == model.gemma3_hidden + projector = model.multi_modal_projector + assert projector.mm_input_projection_weight.shape == (siglip_hidden, model.gemma3_hidden) + assert projector.patches_per_image == 32 + assert projector.tokens_per_side == 16 + assert projector.kernel_size == 2 # ------------------------------------------------------------------ @@ -356,12 +367,12 @@ def test_compute_reward_returns_correct_shape(): """compute_reward returns [batch_size] tensor of finite float32 values.""" model = _make_model() model.eval() - batch = _make_batch(batch_size=3) + batch = _make_batch(batch_size=1) with torch.no_grad(): values = model.compute_reward(batch) - assert values.shape == (3,) + assert values.shape == (1,) assert values.dtype == torch.float32 assert torch.isfinite(values).all() @@ -371,7 +382,7 @@ def test_compute_reward_values_in_support_range(): """Predicted values lie within [value_support_min, value_support_max].""" model = _make_model() model.eval() - batch = _make_batch(batch_size=8) + batch = _make_batch(batch_size=1) with torch.no_grad(): values = model.compute_reward(batch) @@ -400,8 +411,8 @@ def test_gradient_flows_through_value_head(): @skip_if_package_missing("transformers") -def test_gradient_flows_through_cls_embedding(): - """Backprop produces non-zero gradients on the learned [CLS] embedding.""" +def test_gradient_flows_through_value_query(): + """Backprop produces non-zero gradients on the learned value query.""" model = _make_model() model.train() batch = _make_batch() @@ -409,13 +420,13 @@ def test_gradient_flows_through_cls_embedding(): loss, _ = model.forward(batch) loss.backward() - assert model.cls_embedding.weight.grad is not None - assert not torch.all(model.cls_embedding.weight.grad == 0) + assert model.value_query.weight.grad is not None + assert not torch.all(model.value_query.weight.grad == 0) @skip_if_package_missing("transformers") -def test_gradient_flows_through_image_proj(): - """Backprop produces non-zero gradients on the image projection.""" +def test_gradient_flows_through_multimodal_projector(): + """Backprop produces non-zero gradients on the Gemma3 multimodal projection.""" model = _make_model() model.train() batch = _make_batch() @@ -423,8 +434,9 @@ def test_gradient_flows_through_image_proj(): loss, _ = model.forward(batch) loss.backward() - assert model.image_proj.weight.grad is not None - assert not torch.all(model.image_proj.weight.grad == 0) + projector_weight = model.multi_modal_projector.mm_input_projection_weight + assert projector_weight.grad is not None + assert not torch.all(projector_weight.grad == 0) # ------------------------------------------------------------------ @@ -459,8 +471,8 @@ def test_freeze_language_model(): @skip_if_package_missing("transformers") -def test_stop_gradient_to_vlm_preserves_cls_grad(): - """With stop_gradient_to_vlm, CLS embedding still gets gradients.""" +def test_stop_gradient_to_vlm_preserves_value_query_grad(): + """With stop_gradient_to_vlm, the value query still gets gradients.""" config = _make_config(stop_gradient_to_vlm=True) from lerobot.rewards.distributional_value_function.modeling_distributional_value_function import ( DistributionalVFRewardModel, @@ -473,8 +485,8 @@ def test_stop_gradient_to_vlm_preserves_cls_grad(): loss, _ = model.forward(batch) loss.backward() - assert model.cls_embedding.weight.grad is not None - assert not torch.all(model.cls_embedding.weight.grad == 0) + assert model.value_query.weight.grad is not None + assert not torch.all(model.value_query.weight.grad == 0) # ------------------------------------------------------------------ @@ -533,7 +545,7 @@ def test_processor_pipeline_produces_expected_keys(): assert IMAGE_KEY_WRIST_RIGHT + IMAGE_MASK_SUFFIX in processed img = processed[IMAGE_KEY] - assert img.shape == (1, 3, 224, 224) + assert img.shape == (1, 3, IMAGE_SIZE, IMAGE_SIZE) assert img.min() >= -1.0 - 1e-5 assert img.max() <= 1.0 + 1e-5 @@ -604,7 +616,7 @@ def test_image_preprocessor_resize_and_normalize(): transition = { TransitionKey.OBSERVATION: { - IMAGE_KEY: torch.rand(2, 3, 320, 240), # non-square, [0, 1] + IMAGE_KEY: torch.full((2, 3, 320, 240), 0.5), # non-square, [0, 1] } } @@ -614,6 +626,9 @@ def test_image_preprocessor_resize_and_normalize(): assert obs[IMAGE_KEY].shape == (2, 3, 224, 224) assert obs[IMAGE_KEY].min() >= -1.0 - 1e-5 assert obs[IMAGE_KEY].max() <= 1.0 + 1e-5 + # Content value 0.5 must map to 0.0. This also verifies normalization + # happens before resize padding introduces -1 values. + assert torch.allclose(obs[IMAGE_KEY][:, :, 112, 112], torch.zeros(2, 3), atol=1e-5) assert obs[IMAGE_KEY + IMAGE_MASK_SUFFIX].all() @@ -668,29 +683,6 @@ def test_save_load_pretrained_roundtrip(tmp_path): torch.testing.assert_close(orig_sd[key], loaded_sd[key], msg=f"Mismatch in {key}") -# ------------------------------------------------------------------ -# Attention mask utility test -# ------------------------------------------------------------------ - - -@skip_if_package_missing("transformers") -def test_make_att_2d_masks(): - """Verify attention mask construction for prefix + CLS.""" - from lerobot.rewards.distributional_value_function.modeling_distributional_value_function import ( - make_att_2d_masks, - ) - - pad = torch.ones(1, 4, dtype=torch.bool) - att = torch.tensor([[0, 0, 0, 1]]) - mask = make_att_2d_masks(pad, att)[0] - - assert mask[0, 0] # prefix sees prefix - assert mask[1, 2] # prefix sees prefix - assert not mask[0, 3] # prefix does NOT see CLS - assert mask[3, 0] # CLS sees prefix - assert mask[3, 3] # CLS sees itself - - # ------------------------------------------------------------------ # Categorical metrics test # ------------------------------------------------------------------