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

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: {

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(

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",

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