refactor(pi052): remove redundant policy code

scripts/backfill_pi052_action_tokenizer.py

@@ -1,227 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Audit or backfill checkpoint-local FAST artifacts for PI052 model repositories."""
-
-from __future__ import annotations
-
-import argparse
-import hashlib
-import io
-import json
-from pathlib import Path, PurePosixPath
-from typing import Any
-
-from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
-
-DEFAULT_REPOSITORIES = (
-    "pepijn223/pi052_atomic4_01_baseline",
-    "pepijn223/pi052_atomic4_02_lr_1e5",
-    "pepijn223/pi052_atomic4_03_recipe_50_50",
-    "pepijn223/pi052_atomic4_04_flow_weight_10",
-    "pepijn223/pi052_atomic4_05_flow_repeat_1",
-    "pepijn223/pi052_atomic4_06_ki_off",
-)
-CHECKPOINT_DIRECTORIES = (
-    "",
-    "checkpoints/003000/pretrained_model",
-    "checkpoints/006000/pretrained_model",
-    "checkpoints/009000/pretrained_model",
-    "checkpoints/012000/pretrained_model",
-)
-TOKENIZER_DIRECTORY = "action_tokenizer"
-
-
-def artifact_fingerprint(files: list[tuple[str, bytes]]) -> str:
-    digest = hashlib.sha256()
-    for relative_path, content in sorted(files):
-        encoded_path = relative_path.encode()
-        digest.update(len(encoded_path).to_bytes(8, "big"))
-        digest.update(encoded_path)
-        digest.update(len(content).to_bytes(8, "big"))
-        digest.update(content)
-    return digest.hexdigest()
-
-
-def tokenizer_files(tokenizer_path: Path) -> list[tuple[str, Path]]:
-    return [
-        (path.relative_to(tokenizer_path).as_posix(), path)
-        for path in sorted(tokenizer_path.rglob("*"))
-        if path.is_file()
-    ]
-
-
-def _repo_path(directory: str, filename: str) -> str:
-    return (PurePosixPath(directory) / filename).as_posix() if directory else filename
-
-
-def _download_json(repo_id: str, path_in_repo: str, revision: str | None = None) -> dict[str, Any]:
-    path = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
-    return json.loads(Path(path).read_text())
-
-
-def make_portable_preprocessor(config: dict[str, Any]) -> dict[str, Any]:
-    config = json.loads(json.dumps(config))
-    action_steps = [
-        step for step in config["steps"] if step.get("registry_name") == "action_tokenizer_processor"
-    ]
-    if len(action_steps) != 1:
-        raise ValueError(f"Expected one action tokenizer step, found {len(action_steps)}")
-    action_step = action_steps[0]
-    action_step["config"]["action_tokenizer_name"] = TOKENIZER_DIRECTORY
-    action_step["artifacts"] = {"action_tokenizer_name": TOKENIZER_DIRECTORY}
-
-    recipe_steps = [
-        step for step in config["steps"] if step.get("registry_name") == "render_messages_processor"
-    ]
-    if len(recipe_steps) != 1 or not recipe_steps[0].get("config", {}).get("recipe"):
-        raise ValueError("PI052 preprocessor does not contain an embedded training recipe")
-    return config
-
-
-def _json_operation(path_in_repo: str, content: dict[str, Any]) -> CommitOperationAdd:
-    serialized = (json.dumps(content, indent=2) + "\n").encode()
-    return CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(serialized))
-
-
-def prepare_operations(
-    repo_id: str,
-    tokenizer_path: Path,
-    revision: str | None = None,
-) -> list[CommitOperationAdd]:
-    operations: list[CommitOperationAdd] = []
-    files = tokenizer_files(tokenizer_path)
-    for directory in CHECKPOINT_DIRECTORIES:
-        preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
-        operations.append(
-            _json_operation(
-                preprocessor_path,
-                make_portable_preprocessor(_download_json(repo_id, preprocessor_path, revision)),
-            )
-        )
-        for relative_path, local_path in files:
-            operations.append(
-                CommitOperationAdd(
-                    path_in_repo=_repo_path(
-                        directory,
-                        f"{TOKENIZER_DIRECTORY}/{relative_path}",
-                    ),
-                    path_or_fileobj=str(local_path),
-                )
-            )
-    return operations
-
-
-def audit_repository(
-    api: HfApi,
-    repo_id: str,
-    expected_tokenizer_fingerprint: str,
-    revision: str | None = None,
-) -> None:
-    info = api.model_info(repo_id, revision=revision)
-    repository_files = {sibling.rfilename for sibling in info.siblings or []}
-
-    for directory in CHECKPOINT_DIRECTORIES:
-        preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
-        policy_config_path = _repo_path(directory, "config.json")
-        postprocessor_path = _repo_path(directory, "policy_postprocessor.json")
-        for required_path in (preprocessor_path, policy_config_path, postprocessor_path):
-            if required_path not in repository_files:
-                raise FileNotFoundError(f"{repo_id}@{revision or 'main'} is missing {required_path}")
-
-        preprocessor = _download_json(repo_id, preprocessor_path, revision)
-        portable_preprocessor = make_portable_preprocessor(preprocessor)
-        if preprocessor != portable_preprocessor:
-            raise ValueError(f"{repo_id}:{preprocessor_path} is not portable")
-
-        normalizer_steps = [
-            step for step in preprocessor["steps"] if step.get("registry_name") == "normalizer_processor"
-        ]
-        if len(normalizer_steps) != 1 or "state_file" not in normalizer_steps[0]:
-            raise ValueError(f"{repo_id}:{preprocessor_path} is missing normalizer state metadata")
-        normalizer_path = _repo_path(directory, normalizer_steps[0]["state_file"])
-        if normalizer_path not in repository_files:
-            raise FileNotFoundError(f"{repo_id} is missing {normalizer_path}")
-
-        remote_tokenizer_files: list[tuple[str, bytes]] = []
-        for relative_path in _tokenizer_relative_paths(repository_files, directory):
-            path_in_repo = _repo_path(directory, f"{TOKENIZER_DIRECTORY}/{relative_path}")
-            downloaded = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
-            remote_tokenizer_files.append((relative_path, Path(downloaded).read_bytes()))
-        fingerprint = artifact_fingerprint(remote_tokenizer_files)
-        if fingerprint != expected_tokenizer_fingerprint:
-            raise ValueError(
-                f"{repo_id}:{_repo_path(directory, TOKENIZER_DIRECTORY)} fingerprint "
-                f"{fingerprint} != {expected_tokenizer_fingerprint}"
-            )
-
-
-def _tokenizer_relative_paths(repository_files: set[str], directory: str) -> list[str]:
-    prefix = _repo_path(directory, TOKENIZER_DIRECTORY).rstrip("/") + "/"
-    paths = sorted(path.removeprefix(prefix) for path in repository_files if path.startswith(prefix))
-    if not paths:
-        raise FileNotFoundError(f"Missing tokenizer artifact directory {prefix.rstrip('/')}")
-    return paths
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--tokenizer-path", type=Path, required=True)
-    parser.add_argument("--repo-id", action="append", dest="repo_ids")
-    parser.add_argument("--revision")
-    parser.add_argument("--dry-run", action="store_true")
-    parser.add_argument("--audit-only", action="store_true")
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    tokenizer_path = args.tokenizer_path.resolve()
-    if not tokenizer_path.is_dir():
-        raise FileNotFoundError(f"Tokenizer directory does not exist: {tokenizer_path}")
-
-    files = tokenizer_files(tokenizer_path)
-    fingerprint = artifact_fingerprint([(relative_path, path.read_bytes()) for relative_path, path in files])
-    api = HfApi()
-    repositories = tuple(args.repo_ids or DEFAULT_REPOSITORIES)
-    print(f"Tokenizer fingerprint: {fingerprint}")
-
-    for repo_id in repositories:
-        if args.audit_only:
-            audit_repository(api, repo_id, fingerprint, args.revision)
-            print(f"AUDIT OK {repo_id}@{args.revision or 'main'}")
-            continue
-
-        operations = prepare_operations(repo_id, tokenizer_path, args.revision)
-        if args.dry_run:
-            print(f"DRY RUN {repo_id}: {len(operations)} files")
-            for operation in operations:
-                print(f"  {operation.path_in_repo}")
-            continue
-
-        commit = api.create_commit(
-            repo_id=repo_id,
-            repo_type="model",
-            operations=operations,
-            commit_message="Embed fitted FAST tokenizer for portable PI052 checkpoints",
-            revision=args.revision,
-        )
-        audit_repository(api, repo_id, fingerprint, commit.oid)
-        print(f"BACKFILLED {repo_id}@{commit.oid}")
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/policies/pi05/modeling_pi05.py

@@ -17,7 +17,6 @@
 import builtins
 import json
 import logging
-import math
 from collections import deque
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
@@ -31,7 +30,6 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
-    from transformers.cache_utils import DynamicCache
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
     from transformers.utils import cached_file
@@ -44,7 +42,6 @@ if TYPE_CHECKING or _transformers_available:
     )
 else:
     CONFIG_MAPPING = None
-    DynamicCache = None
     modeling_gemma = None
     PiGemmaForCausalLM = None
     _gated_residual = None
@@ -60,6 +57,13 @@ from lerobot.utils.constants import (
 )
 
 from ..common.flow_matching import sample_noise, sample_time_beta
+from ..common.vla_utils import (
+    clone_past_key_values,
+    create_sinusoidal_pos_embedding,
+    make_att_2d_masks,
+    pad_vector,
+    resize_with_pad_torch,
+)
 from ..pretrained import PreTrainedPolicy, T
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
@@ -149,173 +153,6 @@ def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
     return state_dict
 
 
-def get_safe_dtype(target_dtype, device_type):
-    """Get a safe dtype for the given device type."""
-    if device_type == "mps" and target_dtype == torch.float64:
-        return torch.float32
-    if device_type == "cpu":
-        # CPU doesn't support bfloat16, use float32 instead
-        if target_dtype == torch.bfloat16:
-            return torch.float32
-        if target_dtype == torch.float64:
-            return torch.float64
-    return target_dtype
-
-
-def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
-    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
-) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    dtype = get_safe_dtype(torch.float64, device.type)
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-
-    # Compute the outer product
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-
-
-def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
-    alpha_t = torch.tensor(alpha, dtype=torch.float32)
-    beta_t = torch.tensor(beta, dtype=torch.float32)
-    dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,)).to(device)
-
-
-def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (exact copy)
-    """Copied from big_vision.
-
-    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
-    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
-    setup several types of attention, for example:
-
-      [[1 1 1 1 1 1]]: pure causal attention.
-
-      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
-          themselves and the last 3 tokens have a causal attention. The first
-          entry could also be a 1 without changing behaviour.
-
-      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
-          block can attend all previous blocks and all tokens on the same block.
-
-    Args:
-      input_mask: bool[B, N] true if its part of the input, false if padding.
-      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
-        it and 0 where it shares the same attention mask as the previous token.
-    """
-    if att_masks.ndim != 2:
-        raise ValueError(att_masks.ndim)
-    if pad_masks.ndim != 2:
-        raise ValueError(pad_masks.ndim)
-
-    cumsum = torch.cumsum(att_masks, dim=1)
-    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
-    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
-    return att_2d_masks & pad_2d_masks
-
-
-def clone_past_key_values(past_key_values):
-    """Clone the DynamicCache returned by prefix prefill for compiled denoising."""
-    return DynamicCache(
-        tuple(
-            (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
-        )
-    )
-
-
-def pad_vector(vector, new_dim):
-    """Pad the last dimension of a vector to new_dim with zeros.
-
-    Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] >= new_dim:
-        return vector
-    return F.pad(vector, (0, new_dim - vector.shape[-1]))
-
-
-def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
-    images: torch.Tensor,
-    height: int,
-    width: int,
-    mode: str = "bilinear",
-) -> torch.Tensor:
-    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
-    by padding with black. If the image is float32, it must be in the range [-1, 1].
-
-    Args:
-        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
-        height: Target height
-        width: Target width
-        mode: Interpolation mode ('bilinear', 'nearest', etc.)
-
-    Returns:
-        Resized and padded tensor with same shape format as input
-    """
-    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
-    if images.shape[-1] <= 4:  # Assume channels-last format
-        channels_last = True
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
-    else:
-        channels_last = False
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-
-    batch_size, channels, cur_height, cur_width = images.shape
-
-    # Calculate resize ratio
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-
-    # Resize
-    resized_images = F.interpolate(
-        images,
-        size=(resized_height, resized_width),
-        mode=mode,
-        align_corners=False if mode == "bilinear" else None,
-    )
-
-    # Handle dtype-specific clipping
-    if images.dtype == torch.uint8:
-        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
-    elif images.dtype == torch.float32:
-        resized_images = resized_images.clamp(0.0, 1.0)
-    else:
-        raise ValueError(f"Unsupported image dtype: {images.dtype}")
-
-    # Calculate padding
-    pad_h0, remainder_h = divmod(height - resized_height, 2)
-    pad_h1 = pad_h0 + remainder_h
-    pad_w0, remainder_w = divmod(width - resized_width, 2)
-    pad_w1 = pad_w0 + remainder_w
-
-    # Pad
-    constant_value = 0 if images.dtype == torch.uint8 else 0.0
-    padded_images = F.pad(
-        resized_images,
-        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
-        mode="constant",
-        value=constant_value,
-    )
-
-    # Convert back to original format if needed
-    if channels_last:
-        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
-
-    return padded_images
-
-
 # Define the complete layer computation function for gradient checkpointing
 def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []

src/lerobot/policies/pi052/modeling_pi052.py

@@ -970,7 +970,7 @@ class PI052Policy(PI05Policy):
             for p in norm.parameters():
                 p.requires_grad_(True)
         layers = getattr(text_model, "layers", None)
-        if isinstance(layers, (list, torch.nn.ModuleList)) and len(layers) > 0:
+        if isinstance(layers, list | torch.nn.ModuleList) and len(layers) > 0:
             for p in layers[-1].parameters():
                 p.requires_grad_(True)
 
@@ -1084,6 +1084,53 @@ class PI052Policy(PI05Policy):
         loss_dict["loss"] = total.detach().mean()
         return total, loss_dict
 
+    def _embed_supervised_prefix(
+        self,
+        batch: dict[str, Tensor],
+        text_labels: Tensor | None,
+        action_tokens: Tensor | None,
+        action_mask: Tensor | None,
+        *,
+        suppress_prefix_grads: bool = False,
+    ) -> tuple[Tensor, Tensor, Tensor, int, int]:
+        """Embed images, language, and optional FAST supervision once."""
+        images, img_masks = self._preprocess_images(batch)
+        with torch.no_grad() if suppress_prefix_grads else nullcontext():
+            prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
+                images,
+                img_masks,
+                batch[OBS_LANGUAGE_TOKENS],
+                batch[OBS_LANGUAGE_ATTENTION_MASK],
+            )
+        non_fast_prefix_len = prefix_embs.shape[1]
+
+        if text_labels is not None:
+            lang_start = non_fast_prefix_len - text_labels.shape[1]
+            if lang_start >= 0:
+                prefix_att = _mark_target_span_causal(
+                    prefix_att, text_labels, lang_start, non_fast_prefix_len
+                )
+
+        fast_len = 0
+        if action_tokens is not None and action_mask is not None:
+            fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
+            fast_len = action_tokens.shape[1]
+            prefix_embs = torch.cat([prefix_embs, fast_emb], dim=1)
+            prefix_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
+            prefix_att = torch.cat(
+                [
+                    prefix_att,
+                    torch.ones(
+                        (action_tokens.shape[0], fast_len),
+                        dtype=torch.bool,
+                        device=prefix_embs.device,
+                    ),
+                ],
+                dim=1,
+            )
+
+        return prefix_embs, prefix_pad, prefix_att, non_fast_prefix_len, fast_len
+
     def _compute_all_losses_fused(
         self,
         batch: dict[str, Tensor],
@@ -1105,37 +1152,13 @@ class PI052Policy(PI05Policy):
             and getattr(self.config, "knowledge_insulation", False)
         )
 
-        # ---- prefix: images + language + (optional FAST) -------------
+        prefix_embs, prefix_pad, prefix_att, non_fast_prefix_len, fast_len = self._embed_supervised_prefix(
-        images, img_masks = self._preprocess_images(batch)
+            batch,
-        lang_tokens = batch[OBS_LANGUAGE_TOKENS]
+            text_labels,
-        lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
+            action_tokens,
-        with torch.no_grad() if suppress_prefix_grads else nullcontext():
+            action_mask,
-            prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
+            suppress_prefix_grads=suppress_prefix_grads,
-                images, img_masks, lang_tokens, lang_masks
+        )
-            )
-        non_fast_prefix_len = prefix_embs.shape[1]  # images + language only
-
-        # Make supervised text causal rather than a bidirectional copy task.
-        if text_labels is not None:
-            lang_start = non_fast_prefix_len - text_labels.shape[1]
-            if lang_start >= 0:
-                prefix_att = _mark_target_span_causal(
-                    prefix_att, text_labels, lang_start, non_fast_prefix_len
-                )
-
-        fast_len = 0
-        if action_tokens is not None and action_mask is not None:
-            # Gemma embedding already applies its hidden-size scale.
-            fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
-            fast_len = action_tokens.shape[1]
-            ones_att = torch.ones(
-                (action_tokens.shape[0], fast_len),
-                dtype=torch.bool,
-                device=prefix_embs.device,
-            )
-            prefix_embs = torch.cat([prefix_embs, fast_emb], dim=1)
-            prefix_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
-            prefix_att = torch.cat([prefix_att, ones_att], dim=1)
 
         # Amortized flow reuses one VLM prefix across fresh denoising targets.
         num_repeats = int(getattr(self.config, "flow_num_repeats", 1))
@@ -1447,90 +1470,34 @@ class PI052Policy(PI05Policy):
         the caller doesn't want that head.
         """
 
-        images, img_masks = self._preprocess_images(batch)
+        prefix_embs, prefix_pad, prefix_att, _, fast_len = self._embed_supervised_prefix(
-        lang_tokens = batch[OBS_LANGUAGE_TOKENS]
+            batch,
-        lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
+            text_labels,
-
+            action_tokens,
-        prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
+            action_mask,
-            images, img_masks, lang_tokens, lang_masks
         )
-
+        att_2d = make_att_2d_masks(prefix_pad, prefix_att)
-        # Make supervised text causal before appending FAST tokens.
+        position_ids = torch.cumsum(prefix_pad, dim=1) - 1
-        if text_labels is not None:
+        att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=prefix_embs.dtype)
-            lang_start = prefix_embs.shape[1] - text_labels.shape[1]
-            if lang_start >= 0:
-                prefix_att = _mark_target_span_causal(
-                    prefix_att, text_labels, lang_start, prefix_embs.shape[1]
-                )
-
-        fast_len = 0
-        if action_tokens is not None and action_mask is not None:
-            # embed_language_tokens already applies the Gemma sqrt(hidden) scale (tf>=5.4.0);
-            # do not scale FAST action tokens again (would double-scale).
-            fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
-
-            fast_len = action_tokens.shape[1]
-            ones_att = torch.ones(
-                (action_tokens.shape[0], fast_len),
-                dtype=torch.bool,
-                device=prefix_embs.device,
-            )
-            full_embs = torch.cat([prefix_embs, fast_emb], dim=1)
-            full_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
-            full_att = torch.cat([prefix_att, ones_att], dim=1)
-        else:
-            full_embs = prefix_embs
-            full_pad = prefix_pad
-            full_att = prefix_att
-
-        att_2d = make_att_2d_masks(full_pad, full_att)
-        position_ids = torch.cumsum(full_pad, dim=1) - 1
-        att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=full_embs.dtype)
 
         (vlm_out, _), _ = self.model.paligemma_with_expert.forward(
             attention_mask=att_2d_4d,
             position_ids=position_ids,
             past_key_values=None,
-            inputs_embeds=[full_embs, None],
+            inputs_embeds=[prefix_embs, None],
             use_cache=False,
         )
         if vlm_out is None:
             raise RuntimeError("PI052 text+fast loss: VLM forward returned no hidden states.")
-
+        return self._prefix_ce_losses(
-        lm_head = self.model.paligemma_with_expert.paligemma.lm_head
+            vlm_out,
-
+            text_labels,
-        text_loss: Tensor | None = None
+            action_tokens,
-        if text_labels is not None:
+            action_code_mask,
-            lang_len = text_labels.shape[1]
+            fast_len,
-            # embed_prefix lays out as [images, language]; with FAST
+            predict_actions_t,
-            # appended the full sequence is [images, language, FAST].
+            reduction,
-            if fast_len > 0:
+        )
-                text_hidden = vlm_out[:, -(fast_len + lang_len) : -fast_len, :]
-            else:
-                text_hidden = vlm_out[:, -lang_len:, :]
-            text_loss = _shifted_lin_ce(
-                text_hidden,
-                lm_head.weight,
-                text_labels,
-                z_loss_weight=getattr(self.config, "text_ce_z_loss_weight", 0.0),
-                compiled=self.config.use_compiled_text_ce,
-                reduction=reduction,
-            )
-
-        fast_loss: Tensor | None = None
-        if action_tokens is not None and action_code_mask is not None and fast_len > 0:
-            fast_hidden = vlm_out[:, -fast_len:, :]
-            fast_loss = _fast_lin_ce(
-                fast_hidden,
-                lm_head.weight,
-                action_tokens,
-                action_code_mask,
-                predict_actions_t,
-                compiled=self.config.use_compiled_text_ce,
-                reduction=reduction,
-            )
-
-        return text_loss, fast_loss
 
     def select_message(
         self,

src/lerobot/policies/pi_gemma.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
 import torch
 from torch import nn
@@ -22,19 +22,11 @@ from torch.nn import functional as F  # noqa: N812
 
 from lerobot.utils.import_utils import _transformers_available
 
-# Default PaliGemma SigLIP input resolution. Mirrors
-# ``pi05.configuration_pi05.DEFAULT_IMAGE_SIZE``; duplicated as a plain constant
-# to avoid importing the pi05 package here (which would create an import cycle:
-# pi_gemma -> pi05.__init__ -> modeling_pi05 -> pi_gemma).
-DEFAULT_IMAGE_SIZE = 224
-
 if TYPE_CHECKING or _transformers_available:
     from transformers.cache_utils import DynamicCache
     from transformers.masking_utils import create_causal_mask
     from transformers.modeling_layers import GradientCheckpointingLayer
     from transformers.modeling_outputs import BaseModelOutputWithPast
-    from transformers.models.auto import CONFIG_MAPPING
-    from transformers.models.gemma import modeling_gemma
     from transformers.models.gemma.modeling_gemma import (
         GemmaAttention,
         GemmaConfig,
@@ -58,8 +50,6 @@ else:
     GradientCheckpointingLayer = None
     BaseModelOutputWithPast = None
     create_causal_mask = None
-    CONFIG_MAPPING = None
-    modeling_gemma = None
 
 
 def _gated_residual(
@@ -425,332 +415,3 @@ def sdpa_attention_forward(
         scale=scaling,
     )
     return attn_output.transpose(1, 2).contiguous(), None
-
-
-# Define the complete layer computation function for gradient checkpointing
-def compute_layer_complete(
-    layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
-):
-    models = [paligemma.model.language_model, gemma_expert.model]
-    query_states = []
-    key_states = []
-    value_states = []
-    gates = []
-    for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
-        hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
-        gates.append(gate)
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-        query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-        key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-        value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-        query_states.append(query_state)
-        key_states.append(key_state)
-        value_states.append(value_state)
-    # Concatenate and process attention
-    query_states = torch.cat(query_states, dim=2)
-    key_states = torch.cat(key_states, dim=2)
-    value_states = torch.cat(value_states, dim=2)
-    dummy_tensor = torch.zeros(
-        query_states.shape[0],
-        query_states.shape[2],
-        query_states.shape[-1],
-        device=query_states.device,
-        dtype=query_states.dtype,
-    )
-    cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
-    query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, unsqueeze_dim=1
-    )
-    batch_size = query_states.shape[0]
-    scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
-    att_output, _ = sdpa_attention_forward(
-        paligemma.model.language_model.layers[layer_idx].self_attn,
-        query_states,
-        key_states,
-        value_states,
-        attention_mask,
-        scaling,
-    )
-    # Get head_dim from the current layer, not from the model
-    head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
-    att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
-    # Process layer outputs
-    outputs_embeds = []
-    start_pos = 0
-    for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
-        end_pos = start_pos + hidden_states.shape[1]
-        if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
-            att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
-        out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])
-        # first residual
-        out_emb = _gated_residual(hidden_states, out_emb, gates[i])
-        after_first_residual = out_emb.clone()
-        out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])
-        # Convert to bfloat16 if the next layer (mlp) uses bfloat16
-        if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
-            out_emb = out_emb.to(dtype=torch.bfloat16)
-        out_emb = layer.mlp(out_emb)
-        # second residual
-        out_emb = _gated_residual(after_first_residual, out_emb, gate)
-        outputs_embeds.append(out_emb)
-        start_pos = end_pos
-    return outputs_embeds
-
-
-class GemmaVariantConfig:  # see openpi `gemma.py: Config`
-    """Configuration for Gemma model variants."""
-
-    def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
-        self.width = width
-        self.depth = depth
-        self.mlp_dim = mlp_dim
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_dim = head_dim
-
-
-def get_gemma_config(variant: str) -> GemmaVariantConfig:  # see openpi `gemma.py: get_config`
-    """Returns config for specified gemma variant."""
-    if variant == "gemma_300m":
-        return GemmaVariantConfig(
-            width=1024,
-            depth=18,
-            mlp_dim=4096,
-            num_heads=8,
-            num_kv_heads=1,
-            head_dim=256,
-        )
-    elif variant == "gemma_2b":
-        return GemmaVariantConfig(
-            width=2048,
-            depth=18,
-            mlp_dim=16_384,
-            num_heads=8,
-            num_kv_heads=1,
-            head_dim=256,
-        )
-    else:
-        raise ValueError(f"Unknown variant: {variant}")
-
-
-class PaliGemmaWithExpertModel(
-    nn.Module
-):  # see openpi `gemma_pytorch.py: PaliGemmaWithExpertModel` this class is almost a exact copy of PaliGemmaWithExpertModel in openpi
-    """PaliGemma model with action expert for PI05."""
-
-    def __init__(
-        self,
-        vlm_config,
-        action_expert_config,
-        use_adarms=None,
-        precision: Literal["bfloat16", "float32"] = "bfloat16",
-        image_size: int = DEFAULT_IMAGE_SIZE,
-        freeze_vision_encoder: bool = False,
-        train_expert_only: bool = False,
-    ):
-        if use_adarms is None:
-            use_adarms = [False, False]
-        super().__init__()
-        self.freeze_vision_encoder = freeze_vision_encoder
-        self.train_expert_only = train_expert_only
-
-        vlm_config_hf = CONFIG_MAPPING["paligemma"]()
-        vlm_config_hf._vocab_size = 257152  # noqa: SLF001
-        vlm_config_hf.image_token_index = 257152
-        vlm_config_hf.text_config.hidden_size = vlm_config.width
-        vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
-        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
-        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
-        vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
-        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
-        vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
-        vlm_config_hf.text_config.dtype = "float32"
-        vlm_config_hf.text_config.vocab_size = 257152
-        vlm_config_hf.text_config.use_adarms = use_adarms[0]
-        vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None
-        vlm_config_hf.vision_config.image_size = image_size
-        vlm_config_hf.vision_config.intermediate_size = 4304
-        vlm_config_hf.vision_config.projection_dim = 2048
-        vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
-        vlm_config_hf.vision_config.dtype = "float32"
-
-        action_expert_config_hf = CONFIG_MAPPING["gemma"](
-            head_dim=action_expert_config.head_dim,
-            hidden_size=action_expert_config.width,
-            intermediate_size=action_expert_config.mlp_dim,
-            num_attention_heads=action_expert_config.num_heads,
-            num_hidden_layers=action_expert_config.depth,
-            num_key_value_heads=action_expert_config.num_kv_heads,
-            vocab_size=257152,
-            hidden_activation="gelu_pytorch_tanh",
-            dtype="float32",
-            use_adarms=use_adarms[1],
-            adarms_cond_dim=action_expert_config.width if use_adarms[1] else None,
-        )
-
-        self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
-        self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
-        self.gemma_expert.model.embed_tokens = None
-
-        self.to_bfloat16_for_selected_params(precision)
-        self._set_requires_grad()
-
-    def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"):
-        if precision == "bfloat16":
-            self.to(dtype=torch.bfloat16)
-        elif precision == "float32":
-            self.to(dtype=torch.float32)
-            return
-        else:
-            raise ValueError(f"Invalid precision: {precision}")
-
-        # Keep full vision path in float32 so we never toggle (toggle causes optimizer
-        # "same dtype" error). Saves memory vs full float32; more memory than only 3 params.
-        params_to_keep_float32 = [
-            "vision_tower",
-            "multi_modal_projector",
-            "lm_head",
-            "input_layernorm",
-            "post_attention_layernorm",
-            "model.norm",
-        ]
-
-        for name, param in self.named_parameters():
-            if any(selector in name for selector in params_to_keep_float32):
-                param.data = param.data.to(dtype=torch.float32)
-
-    def _set_requires_grad(self):
-        if self.freeze_vision_encoder:
-            self.paligemma.model.vision_tower.eval()
-            for param in self.paligemma.model.vision_tower.parameters():
-                param.requires_grad = False
-        if self.train_expert_only:
-            self.paligemma.eval()
-            for param in self.paligemma.parameters():
-                param.requires_grad = False
-
-    def train(self, mode: bool = True):
-        super().train(mode)
-        if self.freeze_vision_encoder:
-            self.paligemma.model.vision_tower.eval()
-        if self.train_expert_only:
-            self.paligemma.eval()
-
-    def embed_image(self, image: torch.Tensor):
-        # Vision tower and multi_modal_projector are kept in float32 (params_to_keep_float32).
-        out_dtype = image.dtype
-        if image.dtype != torch.float32:
-            image = image.to(torch.float32)
-        image_outputs = self.paligemma.model.get_image_features(image)
-        # OpenPI / big_vision convention: image (soft) tokens are NOT scaled by the
-        # Gemma embedder normalizer (sqrt(hidden_size)) — only text tokens are. lerobot/pi05_base
-        # was trained in this regime, so scaling image features here over-scales them ~45x and
-        # breaks the pretrained vision-language alignment. Keep image features un-normalized.
-        features = image_outputs.pooler_output
-        if features.dtype != out_dtype:
-            features = features.to(out_dtype)
-        return features
-
-    def embed_language_tokens(self, tokens: torch.Tensor):
-        return self.paligemma.model.language_model.embed_tokens(tokens)
-
-    def forward(
-        self,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: list[torch.FloatTensor] | None = None,
-        inputs_embeds: list[torch.FloatTensor] | None = None,
-        use_cache: bool | None = None,
-        adarms_cond: list[torch.Tensor] | None = None,
-    ):
-        if adarms_cond is None:
-            adarms_cond = [None, None]
-        if inputs_embeds[1] is None:
-            prefix_output = self.paligemma.model.language_model.forward(
-                inputs_embeds=inputs_embeds[0],
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                use_cache=use_cache,
-                adarms_cond=adarms_cond[0] if adarms_cond is not None else None,
-            )
-            prefix_past_key_values = prefix_output.past_key_values
-            prefix_output = prefix_output.last_hidden_state
-            suffix_output = None
-        elif inputs_embeds[0] is None:
-            suffix_output = self.gemma_expert.model.forward(
-                inputs_embeds=inputs_embeds[1],
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                use_cache=use_cache,
-                adarms_cond=adarms_cond[1] if adarms_cond is not None else None,
-            )
-            suffix_output = suffix_output.last_hidden_state
-            prefix_output = None
-            prefix_past_key_values = None
-        else:
-            models = [self.paligemma.model.language_model, self.gemma_expert.model]
-            num_layers = self.paligemma.config.text_config.num_hidden_layers
-
-            # Check if gradient checkpointing is enabled for any of the models
-            use_gradient_checkpointing = (
-                hasattr(self.gemma_expert.model, "gradient_checkpointing")
-                and self.gemma_expert.model.gradient_checkpointing
-                and self.training
-            ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
-
-            # Process all layers with gradient checkpointing if enabled
-            for layer_idx in range(num_layers):
-                if use_gradient_checkpointing:
-                    inputs_embeds = torch.utils.checkpoint.checkpoint(
-                        compute_layer_complete,
-                        layer_idx,
-                        inputs_embeds,
-                        attention_mask,
-                        position_ids,
-                        adarms_cond,
-                        use_reentrant=False,
-                        preserve_rng_state=False,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
-                    )
-                else:
-                    inputs_embeds = compute_layer_complete(
-                        layer_idx,
-                        inputs_embeds,
-                        attention_mask,
-                        position_ids,
-                        adarms_cond,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
-                    )
-
-            # final norm
-            def compute_final_norms(inputs_embeds, adarms_cond):
-                outputs_embeds = []
-                for i, hidden_states in enumerate(inputs_embeds):
-                    out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
-                    outputs_embeds.append(out_emb)
-                return outputs_embeds
-
-            # Apply gradient checkpointing to final norm if enabled
-            if use_gradient_checkpointing:
-                outputs_embeds = torch.utils.checkpoint.checkpoint(
-                    compute_final_norms,
-                    inputs_embeds,
-                    adarms_cond,
-                    use_reentrant=False,
-                    preserve_rng_state=False,
-                )
-            else:
-                outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond)
-
-            prefix_output = outputs_embeds[0]
-            suffix_output = outputs_embeds[1]
-            prefix_past_key_values = None
-
-        return [prefix_output, suffix_output], prefix_past_key_values

tests/scripts/test_backfill_pi052_action_tokenizer.py

@@ -1,67 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from scripts.backfill_pi052_action_tokenizer import (
-    CHECKPOINT_DIRECTORIES,
-    DEFAULT_REPOSITORIES,
-    artifact_fingerprint,
-    make_portable_preprocessor,
-)
-
-
-def test_atomic4_backfill_covers_every_repository_and_checkpoint():
-    assert len(DEFAULT_REPOSITORIES) == 6
-    assert CHECKPOINT_DIRECTORIES == (
-        "",
-        "checkpoints/003000/pretrained_model",
-        "checkpoints/006000/pretrained_model",
-        "checkpoints/009000/pretrained_model",
-        "checkpoints/012000/pretrained_model",
-    )
-
-
-def test_backfill_embeds_recipe_and_declares_relative_tokenizer():
-    recipe = {"messages": [{"role": "user", "content": "${task}", "stream": "low_level"}]}
-    preprocessor = {
-        "name": "policy_preprocessor",
-        "steps": [
-            {
-                "registry_name": "normalizer_processor",
-                "config": {},
-                "state_file": "normalizer.safetensors",
-            },
-            {"registry_name": "render_messages_processor", "config": {"recipe": recipe}},
-            {
-                "registry_name": "action_tokenizer_processor",
-                "config": {"action_tokenizer_name": "/fsx/original/tokenizer"},
-            },
-        ],
-    }
-
-    portable = make_portable_preprocessor(preprocessor)
-
-    assert portable["steps"][1]["config"]["recipe"] == recipe
-    assert portable["steps"][2]["config"]["action_tokenizer_name"] == "action_tokenizer"
-    assert portable["steps"][2]["artifacts"] == {"action_tokenizer_name": "action_tokenizer"}
-    assert preprocessor["steps"][2]["config"]["action_tokenizer_name"].startswith("/fsx/")
-
-
-def test_artifact_fingerprint_includes_paths_and_contents():
-    first = artifact_fingerprint([("a/file", b"same"), ("b/file", b"content")])
-
-    assert first == artifact_fingerprint([("b/file", b"content"), ("a/file", b"same")])
-    assert first != artifact_fingerprint([("a/renamed", b"same"), ("b/file", b"content")])
-    assert first != artifact_fingerprint([("a/file", b"changed"), ("b/file", b"content")])