diff --git a/scripts/backfill_pi052_action_tokenizer.py b/scripts/backfill_pi052_action_tokenizer.py deleted file mode 100644 index d97f62d37..000000000 --- a/scripts/backfill_pi052_action_tokenizer.py +++ /dev/null @@ -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() diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index b232c3906..9ce2d80f1 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/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 = [] diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index b22423f46..47d7af249 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/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) ------------- - images, img_masks = self._preprocess_images(batch) - lang_tokens = batch[OBS_LANGUAGE_TOKENS] - lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK] - with torch.no_grad() if suppress_prefix_grads else nullcontext(): - prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix( - 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) + prefix_embs, prefix_pad, prefix_att, non_fast_prefix_len, fast_len = self._embed_supervised_prefix( + batch, + text_labels, + action_tokens, + action_mask, + suppress_prefix_grads=suppress_prefix_grads, + ) # 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) - lang_tokens = batch[OBS_LANGUAGE_TOKENS] - lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK] - - prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix( - images, img_masks, lang_tokens, lang_masks + prefix_embs, prefix_pad, prefix_att, _, fast_len = self._embed_supervised_prefix( + batch, + text_labels, + action_tokens, + action_mask, ) - - # Make supervised text causal before appending FAST tokens. - if text_labels is not None: - 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) + att_2d = make_att_2d_masks(prefix_pad, prefix_att) + position_ids = torch.cumsum(prefix_pad, dim=1) - 1 + att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=prefix_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.") - - lm_head = self.model.paligemma_with_expert.paligemma.lm_head - - text_loss: Tensor | None = None - if text_labels is not None: - lang_len = text_labels.shape[1] - # embed_prefix lays out as [images, language]; with FAST - # appended the full sequence is [images, language, FAST]. - 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 + return self._prefix_ce_losses( + vlm_out, + text_labels, + action_tokens, + action_code_mask, + fast_len, + predict_actions_t, + reduction, + ) def select_message( self, diff --git a/src/lerobot/policies/pi_gemma.py b/src/lerobot/policies/pi_gemma.py index dcaf145dd..4e80db78d 100644 --- a/src/lerobot/policies/pi_gemma.py +++ b/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 diff --git a/tests/scripts/test_backfill_pi052_action_tokenizer.py b/tests/scripts/test_backfill_pi052_action_tokenizer.py deleted file mode 100644 index d3d664785..000000000 --- a/tests/scripts/test_backfill_pi052_action_tokenizer.py +++ /dev/null @@ -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")])