#!/usr/bin/env python
# Copyright 2025 Physical Intelligence and The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. 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.
"""π05: A Vision-Language-Action Flow Model for General Robot Control
[Paper](https://www.physicalintelligence.company/download/pi05.pdf)
"""
import builtins
import logging
import math
from collections import deque
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F # noqa: N812
from einops import rearrange
from torch import Tensor, nn
from transformers import AutoProcessor, AutoTokenizer
from opentau.configs.policies import PreTrainedConfig
from opentau.configs.types import NormalizationMode
from opentau.datasets.grounding.tokenizer_utils import ensure_loc_tokens
from opentau.policies.normalize import Normalize, Unnormalize
from opentau.policies.normalize import resolve_num_datasets as _num_datasets
from opentau.policies.pi05.configuration_pi05 import PI05Config
from opentau.policies.pi05.paligemma_with_expert import (
PaliGemmaWithExpertConfig,
PaliGemmaWithExpertModel,
)
from opentau.policies.pretrained import PreTrainedPolicy, T
from opentau.policies.utils import PerSampleLoss, ce_per_sample, flow_matching_masked_mse
from opentau.utils.accelerate_utils import get_proc_accelerator
from opentau.utils.utils import get_safe_dtype
def _preferred_dtype():
return torch.float32 if torch.onnx.is_in_onnx_export() else torch.bfloat16
[docs]
def create_sinusoidal_pos_embedding(
time: Tensor, dimension: int, min_period: float, max_period: float, device: torch.device | str = "cpu"
) -> Tensor:
"""Computes sine-cosine positional embedding vectors for scalar positions.
Args:
time: A 2-D tensor of shape (batch_size, action_chunk_length).
dimension: The dimension of the embedding vectors. Must be divisible by 2.
min_period: The minimum period of the sinusoidal functions.
max_period: The maximum period of the sinusoidal functions.
device: The device to create the tensors on. Defaults to "cpu".
Returns:
A tensor of shape (batch_size, action_chunk_length, dimension) containing the positional embeddings.
Raises:
ValueError: If dimension is not divisible by 2 or if time tensor is not 2-D with shape (batch_size, action_chunk_length).
"""
if dimension % 2 != 0:
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
if time.ndim != 2:
raise ValueError("The time tensor is expected to be of shape `(batch_size, action_chunk_length)`.")
dtype = (
get_safe_dtype(torch.float64, device.type)
if isinstance(device, torch.device)
else get_safe_dtype(torch.float64, device)
)
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 = rearrange(scaling_factor, "d -> 1 1 d") * rearrange(time, "b c -> b c 1")
pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=2)
return pos_emb
[docs]
def make_att_2d_masks(
pad_masks: Tensor,
att_masks: Tensor,
n_cross_att_tokens: int | None = None,
cross_att_pad_masks: Tensor | None = None,
) -> Tensor:
"""Creates a 2-D attention mask given padding and 1-D attention masks.
Tokens can attend to valid inputs tokens which have a cumulative `att_masks`
smaller or equal to theirs. This way `att_masks` 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:
pad_masks: bool[B, N] true if its part of the input, false if padding.
att_masks: 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.
n_cross_att_tokens: Add attention mask for cross-attention tokens if
`n_cross_att_tokens` is provided.
cross_att_pad_masks: Padding masks for cross attention tokens. Required if
`n_cross_att_tokens` is provided.
Returns:
A 2D attention mask tensor of shape (B, N + n_cross_att_tokens, N + n_cross_att_tokens)
if n_cross_att_tokens is provided, else (B, N, N).
Raises:
ValueError: If att_masks or pad_masks are not 2D (including batch dimension).
AssertionError: If cross_att_pad_masks is missing when n_cross_att_tokens is set,
or if its shape is incorrect.
"""
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]
att_2d_masks = att_2d_masks & pad_2d_masks
# If `n_cross_att_tokens` is provided, we add a mask for cross-attention tokens at the end of the sequence.
if n_cross_att_tokens is not None:
assert cross_att_pad_masks is not None, (
"cross_att_pad_masks must be provided if n_cross_att_tokens is provided"
)
assert cross_att_pad_masks.shape == (att_masks.size(0), n_cross_att_tokens), (
"cross_att_pad_masks must have shape (batch_size, n_cross_att_tokens)"
)
cross_att_mask = torch.full(
(att_masks.size(0), att_masks.size(1), n_cross_att_tokens),
True,
dtype=torch.bool,
device=att_masks.device,
)
# Apply padding masks: pad_masks for rows, cross_att_pad_masks for columns
cross_att_mask = cross_att_mask & pad_masks[:, :, None] & cross_att_pad_masks[:, None, :]
# The cross_att_masks are concatenated before the att_2d_masks
att_2d_masks = torch.cat((cross_att_mask, att_2d_masks), dim=2)
return att_2d_masks
[docs]
def resize_with_pad(img: Tensor, width: int, height: int, pad_value: int = -1) -> Tensor:
"""Resizes an image to fit within the specified dimensions while maintaining aspect ratio,
and pads the remaining area with the specified value.
Args:
img: Input image tensor of shape (batch_size, channels, current_height, current_width).
width: Target width.
height: Target height.
pad_value: Value to use for padding. Defaults to -1.
Returns:
The resized and padded image tensor of shape (batch_size, channels, height, width).
Raises:
ValueError: If the input image tensor does not have 4 dimensions.
"""
if img.ndim != 4:
raise ValueError(f"(b,c,h,w) expected, but {img.shape}")
cur_height, cur_width = img.shape[2:]
# Explicit no-op when the input already matches the target — native-
# resolution inputs must pass through bit-identical, not survive a
# same-size bilinear round trip.
if (cur_height, cur_width) == (height, width):
return img
ratio = max(cur_width / width, cur_height / height)
resized_height = int(cur_height / ratio)
resized_width = int(cur_width / ratio)
resized_img = F.interpolate(
img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
)
pad_height = max(0, int(height - resized_height))
pad_width = max(0, int(width - resized_width))
# pad on left and top of image
padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
return padded_img
[docs]
def pad_discrete_tokens(tokens: list[list[int]], max_length: int) -> tuple[np.ndarray, np.ndarray]:
"""Pads or truncates a list of discrete action token sequences to a fixed length.
Args:
tokens: A list of discrete action token sequences (lists of integers).
max_length: The target length for the discrete action token sequences.
Returns:
A tuple containing:
- discrete_action_tokens: A numpy array of shape (len(tokens), max_length) containing the padded discrete action tokens.
- discrete_action_masks: A boolean numpy array of shape (len(tokens), max_length) indicating valid discrete action tokens (True) and padding (False).
"""
discrete_action_tokens = []
discrete_action_masks = []
for token in tokens:
if len(token) > max_length:
logging.warning(
f"Discrete action token length {len(token)} is greater than max_length {max_length}, truncating"
)
discrete_action_tokens.append(np.array(token[:max_length]))
discrete_action_masks.append(np.ones(max_length, dtype=bool))
else:
discrete_action_masks.append(
np.concatenate(
[np.ones(len(token), dtype=bool), np.zeros(max_length - len(token), dtype=bool)]
)
)
discrete_action_tokens.append(np.pad(token, (0, max_length - len(token)), constant_values=0))
return np.array(discrete_action_tokens), np.array(discrete_action_masks)
[docs]
class PI05Policy(PreTrainedPolicy):
"""Wrapper class around PI05FlowMatching model to train and run inference within OpenTau."""
config_class = PI05Config
name = "pi05"
# forward() dispatches self.model via __call__, so nn.Module.compile() takes effect.
supports_torch_compile = True
[docs]
def __init__(
self,
config: PI05Config,
per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None,
dataset_names: list[str] | None = None,
):
"""Initializes the PI05Policy.
Args:
config: Policy configuration class instance.
per_dataset_stats: Ordered list of per-dataset stat dicts used to
fill the stacked Normalize/Unnormalize buffers.
dataset_names: Ordered list parallel to ``per_dataset_stats``.
"""
super().__init__(config)
config.validate_features()
self.config = config
num_datasets = _num_datasets(per_dataset_stats, dataset_names, config)
self.normalize_inputs = Normalize(
config.input_features,
config.normalization_mapping,
per_dataset_stats=per_dataset_stats,
dataset_names=dataset_names,
num_datasets=num_datasets,
)
self.normalize_targets = Normalize(
config.output_features,
config.normalization_mapping,
per_dataset_stats=per_dataset_stats,
dataset_names=dataset_names,
num_datasets=num_datasets,
)
self.normalize_discrete_actions = Normalize(
config.output_features,
{"ACTION": NormalizationMode.MIN_MAX},
per_dataset_stats=per_dataset_stats,
dataset_names=dataset_names,
num_datasets=num_datasets,
)
self.unnormalize_outputs = Unnormalize(
config.output_features,
config.normalization_mapping,
per_dataset_stats=per_dataset_stats,
dataset_names=dataset_names,
num_datasets=num_datasets,
)
# The same PaliGemma tokenizer instance is shared with the inner
# `PI05FlowMatching`. The single `ensure_loc_tokens` call inside the
# inner ctor promotes the reserved <loc0000>..<loc1023> entries on
# both layers at once — no second load, no risk of revision drift.
self.language_tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
self.discrete_action_processor = AutoProcessor.from_pretrained(
config.discrete_action_tokenizer_path, trust_remote_code=True
)
# Get vocab size from processor
discrete_action_vocab_size = getattr(self.discrete_action_processor, "vocab_size", None)
self.model = PI05FlowMatching(
config,
discrete_action_vocab_size=discrete_action_vocab_size,
language_tokenizer=self.language_tokenizer,
)
self.reset()
[docs]
def reset(self) -> None:
"""This should be called whenever the environment is reset."""
self._action_queue = deque([], maxlen=self.config.n_action_steps)
[docs]
@classmethod
def from_pretrained(
cls: builtins.type[T],
pretrained_name_or_path: str | Path,
*,
config: PreTrainedConfig | None = None,
force_download: bool = False,
resume_download: bool | None = None,
proxies: dict | None = None,
token: str | bool | None = None,
cache_dir: str | Path | None = None,
local_files_only: bool = False,
revision: str | None = None,
strict: bool = True,
**kwargs,
) -> T:
"""Override the from_pretrained method to handle key remapping.
Args:
pretrained_name_or_path: Path to the pretrained model or its name on the Hub.
config: Configuration object.
force_download: Whether to force download the model weights.
resume_download: Whether to resume download.
proxies: Proxy configuration.
token: Authentication token.
cache_dir: Directory to cache downloaded files.
local_files_only: Whether to only look for files locally.
revision: Specific model revision.
strict: Whether to strictly enforce state dict matching.
**kwargs: Additional keyword arguments.
Returns:
The loaded model instance.
Raises:
ValueError: If pretrained_name_or_path is None.
"""
if pretrained_name_or_path is None:
raise ValueError("pretrained_name_or_path is required")
# Use provided config if available, otherwise create default config
if config is None:
config = PreTrainedConfig.from_pretrained(
pretrained_name_or_path=pretrained_name_or_path,
force_download=force_download,
resume_download=resume_download,
proxies=proxies,
token=token,
cache_dir=cache_dir,
local_files_only=local_files_only,
revision=revision,
**kwargs,
)
# Initialize model without loading weights
# Check if dataset_stats were provided in kwargs
model = cls(config, **kwargs)
# Now manually load and remap the state dict
acc = get_proc_accelerator()
is_main_process = acc.is_main_process if acc else True
# Populated inside the try block when skip_normalization_weights fires;
# used outside the try/except to gate the inf-buffer guard so the
# ValueError is not swallowed by the broad except below.
stripped_keys: frozenset[str] = frozenset()
try:
# Try to load the pytorch_model.bin or model.safetensors file
if is_main_process:
print(f"Loading model from: {pretrained_name_or_path}")
try:
from transformers.utils import cached_file
# Try safetensors first
resolved_file = cached_file(
pretrained_name_or_path,
"model.safetensors",
cache_dir=kwargs.get("cache_dir"),
force_download=kwargs.get("force_download", False),
resume_download=kwargs.get("resume_download"),
proxies=kwargs.get("proxies"),
use_auth_token=kwargs.get("use_auth_token"),
revision=kwargs.get("revision"),
local_files_only=kwargs.get("local_files_only", False),
)
from safetensors.torch import load_file
original_state_dict = load_file(resolved_file)
if is_main_process:
print("✓ Loaded state dict from model.safetensors")
except Exception as e:
if is_main_process:
print(f"Could not load state dict from remote files: {e}")
print("Returning model without loading pretrained weights")
return model
# First, fix any key differences # see openpi `model.py, _fix_pytorch_state_dict_keys`
fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
# Then add "model." prefix for all keys that don't already have it
remapped_state_dict = {}
remap_count = 0
for key, value in fixed_state_dict.items():
if not key.startswith("model.") and "normalize" not in key:
new_key = f"model.{key}"
remapped_state_dict[new_key] = value
remap_count += 1
if remap_count <= 10 and is_main_process: # Only print first 10 to avoid spam
print(f"Remapped: {key} -> {new_key}")
else:
remapped_state_dict[key] = value
if remap_count > 0 and is_main_process:
print(f"Remapped {remap_count} state dict keys")
# Strip saved normalize/unnormalize buffers when the user opted in
# via config.skip_normalization_weights — see PreTrainedConfig and
# PreTrainedPolicy._strip_normalization_buffers_from_state_dict.
remapped_state_dict, stripped_keys = cls._strip_normalization_buffers_from_state_dict(
remapped_state_dict, model.config, is_main_process=is_main_process
)
# Load the remapped state dict into the model
# Promote legacy single-dataset Normalize/Unnormalize buffers from
# `(*feat_shape,)` to the new `(1, *feat_shape)` stacked layout so pre-PR
# checkpoints load via `model.load_state_dict(...)`.
model._promote_legacy_norm_buffers_in_state_dict(remapped_state_dict)
missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=False)
# Hide deliberately-stripped buffer keys from the missing-keys
# warning so the noisy log does not directly contradict the INFO
# logged just above. ``stripped_keys`` is empty when the flag is
# off, so this is a no-op for default loads.
unintended_missing = [key for key in missing_keys if key not in stripped_keys]
if unintended_missing and is_main_process:
print(f"Missing keys when loading state dict: {len(unintended_missing)} keys")
if len(unintended_missing) <= 20:
for key in unintended_missing:
print(f" - {key}")
else:
for key in unintended_missing[:20]:
print(f" - {key}")
print(f" ... and {len(unintended_missing) - 20} more")
if unexpected_keys and is_main_process:
print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
if len(unexpected_keys) <= 20:
for key in unexpected_keys:
print(f" - {key}")
else:
for key in unexpected_keys[:20]:
print(f" - {key}")
print(f" ... and {len(unexpected_keys) - 20} more")
if not unintended_missing and not unexpected_keys and is_main_process:
print("All keys loaded successfully!")
except Exception as e:
if is_main_process:
print(f"Warning: Could not remap state dict keys: {e}")
# Outside the try/except so the ValueError is not swallowed by the
# broad except above. The helper no-ops when ``stripped_keys`` is
# empty (flag was off or the try block bailed before the strip ran).
cls._assert_normalize_buffers_initialized(model, stripped_keys=stripped_keys)
return model
def _fix_pytorch_state_dict_keys(
self, state_dict: dict[str, Tensor], model_config: PreTrainedConfig
) -> dict[str, Tensor]: # see openpi `BaseModelConfig, _fix_pytorch_state_dict_keys`
"""Fix state dict keys to match current model architecture.
Args:
state_dict: The state dictionary to fix.
model_config: The model configuration.
Returns:
The fixed state dictionary.
"""
import re
fixed_state_dict = {}
for key, value in state_dict.items():
new_key = key
# Handle layer norm structure changes: .weight -> .dense.weight + .dense.bias
# For gemma expert layers
if re.match(
r"paligemma_with_expert\.gemma_expert\.model\.layers\.\d+\.(input_layernorm|post_attention_layernorm)\.weight",
key,
):
# Check if the model actually has adaRMS enabled for the expert
expert_uses_adarms = getattr(
self.model.paligemma_with_expert.gemma_expert.config, "use_adarms", False
)
if expert_uses_adarms:
logging.warning(f"Skipping layer norm key (adaRMS mismatch): {key}")
continue
if re.match(r"paligemma_with_expert\.gemma_expert\.model\.norm\.weight", key):
# Check if the model actually has adaRMS enabled for the expert
expert_uses_adarms = getattr(
self.model.paligemma_with_expert.gemma_expert.config, "use_adarms", False
)
if expert_uses_adarms:
logging.warning(f"Skipping norm key (adaRMS mismatch): {key}")
continue
# Handle MLP naming changes for pi05
# pi05 model expects time_mlp_*, but checkpoint might have action_time_mlp_*
if key.startswith("action_time_mlp_in."):
new_key = key.replace("action_time_mlp_in.", "time_mlp_in.")
elif key.startswith("action_time_mlp_out."):
new_key = key.replace("action_time_mlp_out.", "time_mlp_out.")
if key.startswith("state_proj.") and model_config.state_type == "discrete":
logging.warning(f"Skipping state_proj key in discrete state mode: {key}")
continue
# Handle vision tower embedding layer potential differences
if "patch_embedding" in key:
# Some checkpoints might have this, but current model expects different structure
logging.warning(f"Vision embedding key might need handling: {key}")
fixed_state_dict[new_key] = value
return fixed_state_dict
[docs]
def get_optim_params(self) -> dict:
"""Returns the parameters to be optimized.
Returns:
A generator over the model parameters.
"""
return self.parameters()
[docs]
@torch.no_grad()
def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
"""Predict a chunk of actions given environment observations.
Args:
batch: Batch of data containing environment observations.
Returns:
The predicted action chunk.
Raises:
NotImplementedError: Always, as this method is not implemented for PI05.
"""
raise NotImplementedError("Currently not implemented for PI05")
[docs]
@torch.no_grad()
def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
"""Select a single action given environment observations.
This method uses an action queue that is replenished when it has config.max_delay or fewer actions (or is empty).
When replenishing, the current queue contents are used as action_prefix for sample_actions,
then the queue is refilled with the new chunk.
Note: This method should only be called when running a policy in simulation. For real world inference,
this method should be written in the ROS client node.
Args:
batch: Batch of data containing environment observations.
noise: Optional noise tensor to be used during sampling.
Returns:
The selected action tensor.
"""
self.eval()
if len(self._action_queue) == 0 or len(self._action_queue) <= self.config.max_delay:
# Use current queue as action prefix to replenish
action_prefix = None
delay = 0
if len(self._action_queue) > 0:
prefix_actions = list(self._action_queue)
delay = min(len(prefix_actions), self.config.max_delay)
assert delay == self.config.max_delay, f"Delay must be equal to {self.config.max_delay}"
prefix_actions = prefix_actions[-delay:]
action_prefix = torch.stack(prefix_actions, dim=1)
delay = torch.tensor(delay, dtype=torch.long, device=batch["state"].device)
actions = self.sample_actions(batch, noise=noise, action_prefix=action_prefix, delay=delay)
actions = rearrange(actions, "b c d -> c b d")
# Execute only the first n_action_steps of the predicted chunk, then
# re-query with fresh observations (receding horizon). The config guard
# (n_action_steps < chunk_size => max_delay == 0 => delay == 0) keeps this
# slice in range and exactly n_action_steps long; when n_action_steps ==
# chunk_size it clamps to actions[delay:] (unchanged behaviour).
self._action_queue.extend(actions[delay : delay + self.config.n_action_steps])
assert len(self._action_queue) == self.config.n_action_steps, (
f"Action queue must have {self.config.n_action_steps} actions"
)
action = self._action_queue.popleft()
return action
[docs]
@torch.no_grad()
def sample_actions(
self,
batch: dict[str, Tensor],
action_prefix: Tensor | None = None,
delay: Tensor | None = None,
noise: Tensor | None = None,
) -> Tensor:
"""Sample actions from the policy given environment observations.
Note: The provided action_prefix should NOT be normalized, as this method will handle normalization internally.
The action_prefix should have shape (batch_size, action_chunk_length, action_dim) where action_chunk_length is less than or equal to config.chunk_size.
Args:
batch: Batch of data containing environment observations.
action_prefix: Optional action prefix tensor of shape (batch_size, action_chunk_length, action_dim).
delay: Optional number of frozen delay actions from action_prefix.
noise: Optional noise tensor.
Returns:
The sampled actions tensor of shape (batch_size, action_chunk_length, action_dim).
"""
if not (torch.compiler.is_compiling() or torch.onnx.is_in_onnx_export()):
assert delay is None or 0 <= delay.item() <= self.config.max_delay, (
f"Delay must be None or between 0 and {self.config.max_delay}"
)
dataset_index = self._resolve_dataset_index(batch)
batch = self.normalize_inputs(batch, dataset_index)
images, img_masks = self.prepare_images(batch)
lang_tokens, lang_masks = self.prepare_language(batch)
# if delay is not provided, set it to 0
if delay is None:
delay = torch.tensor(0, dtype=torch.long, device=lang_tokens.device)
if action_prefix is None:
# create a zero filled action_prefix
bsize = lang_tokens.shape[0]
actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
action_prefix = torch.zeros(actions_shape, dtype=lang_tokens.dtype, device=lang_tokens.device)
else:
# normalize action_prefix and pad chunk dimension to config.chunk_size
action_prefix = self.normalize_targets({"actions": action_prefix}, dataset_index)["actions"]
action_prefix = F.pad( # noop if chunk_size is already the same as action_prefix.shape[1]
action_prefix,
(0, 0, 0, self.config.chunk_size - action_prefix.shape[1]),
)
state = self.prepare_state(batch) if self.config.state_type == "continuous" else None
actions = self.model.sample_actions(
images,
img_masks,
lang_tokens,
lang_masks,
action_prefix,
delay,
noise=noise,
state=state,
)
# Unpad actions
original_action_dim = self.config.action_feature.shape[0]
actions = actions[:, :, :original_action_dim]
actions = self.unnormalize_outputs({"actions": actions}, dataset_index)["actions"]
return actions
[docs]
def forward(
self,
batch: dict[str, Tensor],
noise: Tensor | None = None,
time: Tensor | None = None,
return_per_sample: bool = False,
) -> dict[str, Tensor | PerSampleLoss]:
"""Do a full training forward pass to compute the loss.
Args:
batch: Batch of data containing environment observations, actions, and targets.
noise: Optional noise tensor.
time: Optional time tensor.
return_per_sample: When True, additionally return per-sample ``MSE_per_sample``
/ ``CE_per_sample`` (:class:`PerSampleLoss`) so the validation loop can
bucket the loss by ``(dataset, control_mode)`` provenance. The scalar
``MSE``/``CE`` are unchanged, so the training path is unaffected.
Returns:
A dictionary with the loss components ("MSE" and "CE"), plus
"MSE_per_sample"/"CE_per_sample" when ``return_per_sample`` is True.
"""
dataset_index = self._resolve_dataset_index(batch)
batch = self.normalize_inputs(batch, dataset_index)
batch["discrete_actions"] = self.normalize_discrete_actions(dict(batch), dataset_index)["actions"]
batch = self.normalize_targets(batch, dataset_index)
images, img_masks = self.prepare_images(
batch
) # in img_masks we have True for real images and False for padded images
lang_tokens, lang_masks = self.prepare_language(
batch
) # in lang_masks we have True for real tokens and False for padded tokens
# response prediction is to predict the response . It will attend to image and language inputs.
response_tokens, response_masks = self.prepare_response(
batch
) # in response_masks we have True for real tokens and False for padded tokens
# discrete actions are to predict actions using autoregressive technique and not flow matching. It will attend to image, language and response inputs.
discrete_actions, discrete_action_masks = self.prepare_discrete_actions(
batch
) # in discrete_action_masks we have True for real tokens and False for padded tokens
actions = batch["actions"]
actions_is_pad = batch.get(
"action_is_pad"
) # in actions_is_pad we have False for real actions and True for padded actions
state = self.prepare_state(batch) if self.config.state_type == "continuous" else None
# Call via __call__ (not .forward) so an in-place torch.compile of
# self.model (see PreTrainedPolicy.maybe_compile_for_training) is
# actually dispatched; a direct .forward() bypasses the compiled call.
losses = self.model(
images,
img_masks,
lang_tokens,
lang_masks,
actions,
actions_is_pad,
response_tokens,
response_masks,
noise,
time,
discrete_actions,
discrete_action_masks,
state=state,
real_action_dim=batch.get("real_action_dim"),
return_per_sample=return_per_sample,
)
out: dict[str, Tensor | PerSampleLoss] = {"MSE": losses["MSE"], "CE": losses["CE"]}
if return_per_sample:
out["MSE_per_sample"] = losses["MSE_per_sample"]
out["CE_per_sample"] = losses["CE_per_sample"]
return out
[docs]
def prepare_state(self, batch: dict[str, Tensor]) -> Tensor:
"""Prepares the continuous state tensor, padding or truncating to max_state_dim.
Only used when ``state_type == "continuous"``.
Args:
batch: Batch of data containing the "state" tensor of shape (batch_size, state_dim).
Returns:
A tensor of shape (batch_size, max_state_dim).
Raises:
ValueError: If the state dimension exceeds max_state_dim.
"""
assert self.config.state_type == "continuous", "prepare_state is only used in continuous state mode"
state = batch["state"]
state_dim = state.shape[-1]
if state_dim > self.config.max_state_dim:
raise ValueError(
f"State dimension ({state_dim}) exceeds max_state_dim ({self.config.max_state_dim}). "
f"Increase max_state_dim in the config to accommodate the state vector."
)
if state_dim < self.config.max_state_dim:
state = F.pad(state, (0, self.config.max_state_dim - state_dim))
return state
[docs]
def prepare_discrete_state(self, batch: dict[str, Tensor]) -> list[str]:
"""Discretizes the state into bins and converts it to a string representation.
Each dimension of the state vector is discretized into 256 bins.
The values of each dimension of the state are expected to be in the range [-1, 1].
The discretization bins are linearly spaced between -1 and 1.
The index of the bin for each dimension is then concatenated into a space-separated string.
Args:
batch: Batch of data containing the "state" tensor.
Returns:
A list of strings, where each string is a space-separated list of discretized state values.
Raises:
ValueError: If the state values are not normalized between -1 and 1.
"""
assert self.config.state_type == "discrete", (
"prepare_discrete_state is only used in discrete state mode"
)
state = batch["state"]
state_cpu = state.to(device="cpu", dtype=torch.float32)
if torch.any(state_cpu < -1.0) or torch.any(state_cpu > 1.0):
logging.warning(
f"State values are not normalized between -1 and 1. Min: {state_cpu.min().item()}, Max: {state_cpu.max().item()}"
)
state_clipped = torch.clamp(state_cpu, -1.0, 1.0)
# replicate np.digitize with torch for torch.compile compatibility
bin_indices = ((state_clipped + 1.0) * 128.0).long().clamp(0, 255)
discretized_states = bin_indices.cpu().tolist()
return [
" ".join(map(str, row)) for row in discretized_states
] # TODO: return a tensor instead of a list of strings?
[docs]
def prepare_discrete_actions(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
"""Prepares discrete actions for the model by tokenizing and padding them.
Args:
batch: Batch of data containing the key "discrete_actions".
Returns:
A tuple containing:
- discrete_action_tokens: A tensor of shape (batch_size, max_length) containing the tokenized actions.
- discrete_action_masks: A tensor of shape (batch_size, max_length) indicating valid tokens.
"""
device = batch["discrete_actions"].device
discrete_actions = batch["discrete_actions"].to(device="cpu", dtype=torch.float32)
tokens = self.discrete_action_processor.__call__(discrete_actions)
discrete_action_tokens, discrete_action_masks = pad_discrete_tokens(
tokens, self.config.discrete_action_max_length
)
return torch.from_numpy(discrete_action_tokens).to(device=device, dtype=torch.long), torch.from_numpy(
discrete_action_masks
).to(device=device, dtype=torch.bool)
[docs]
def prepare_images(self, batch: dict[str, Tensor]) -> tuple[list[Tensor], list[Tensor]]:
"""Apply preprocessing to the images.
Resizes to 224x224 and padding to keep aspect ratio, and converts pixel range
from [0.0, 1.0] to [-1.0, 1.0] as requested by SigLIP.
Args:
batch: Batch of data containing image tensors.
Returns:
A tuple containing:
- images: A list of processed image tensors.
- img_masks: A list of image mask tensors.
Raises:
ValueError: If no image features are present in the batch.
"""
images = []
img_masks = []
present_img_keys = [key for key in self.config.image_features if key in batch]
missing_img_keys = [key for key in self.config.image_features if key not in batch]
if len(present_img_keys) == 0:
raise ValueError(
f"All image features are missing from the batch. At least one expected. (batch: {batch.keys()}) (image_features:{self.config.image_features})"
)
# Preprocess image features present in the batch
for key in present_img_keys:
img = batch[key]
if self.config.resize_imgs_with_padding is not None:
# The config tuple is (height, width); the function signature is
# (width, height) — unpack explicitly so non-square targets are not
# transposed (invisible at the square defaults).
target_h, target_w = self.config.resize_imgs_with_padding
img = resize_with_pad(img, width=target_w, height=target_h, pad_value=0)
# Normalize from range [0,1] to [-1,1] as expected by siglip
img = img * 2.0 - 1.0
bsize = img.shape[0]
device = img.device
mask = torch.ones(bsize, dtype=torch.bool, device=device)
images.append(img)
img_masks.append(mask)
# Create image features not present in the batch
# as fully 0 padded images.
for num_empty_cameras in range(len(missing_img_keys)):
if num_empty_cameras >= self.config.empty_cameras:
break
img = torch.ones_like(img) * -1
mask = torch.zeros_like(mask)
images.append(img)
img_masks.append(mask)
return images, img_masks
[docs]
def prepare_language(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
"""Tokenize the text input.
When ``state_type == "discrete"``, the state is discretized into bins and
embedded into the prompt string. When ``state_type == "continuous"``, the
state is handled separately via :meth:`prepare_state`, so the prompt only
contains the task description.
Args:
batch: Batch of data containing the key "prompt" and "state".
Returns:
A tuple containing:
- lang_tokens: Tensor of language tokens.
- lang_masks: Tensor of language attention masks.
"""
device = batch["state"].device
tasks = batch["prompt"]
if self.config.state_type == "continuous":
prompt = [f"Task: {task}, " for task in tasks]
else:
# add state to the prompt
state = self.prepare_discrete_state(batch)
# using <eos> to separate each modality
prompt = [f"Task: {task}, State: {state};\n" for task, state in zip(tasks, state, strict=False)]
tokenized_prompt = self.language_tokenizer.__call__(
prompt,
padding="max_length",
padding_side="right",
max_length=self.config.prompt_max_length,
return_tensors="pt",
truncation=True,
)
lang_tokens = tokenized_prompt["input_ids"].to(device=device)
lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
return lang_tokens, lang_masks
[docs]
def prepare_response(self, batch: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
"""Tokenize the response input.
Args:
batch: Batch of data containing the key "response".
Returns:
A tuple containing:
- response_tokens: Tensor of response language tokens.
- response_masks: Tensor of response language attention masks.
"""
if not self.config.predict_response:
return None, None
device = batch["state"].device
responses = batch["response"]
# if '' is found in response then response is not for loss calculation (used for robotic dataset with no subtask), so add pad token to the response.
response_prompt = [f"{response}<eos>;\n" for response in responses]
tokenized_response = self.language_tokenizer.__call__(
response_prompt,
padding="max_length",
padding_side="right",
max_length=self.config.response_max_length,
return_tensors="pt",
truncation=True,
)
response_tokens = tokenized_response["input_ids"].to(device=device)
response_masks = tokenized_response["attention_mask"].to(device=device, dtype=torch.bool)
return response_tokens, response_masks
[docs]
class PI05FlowMatching(nn.Module):
"""
π05: A Vision-Language-Action Flow Model for General Robot Control
[Paper](https://www.physicalintelligence.company/download/pi05.pdf)
┌──────────────────────────────────────────┐
│ actions │
│ ▲ │
│ ┌┴─────┐ │
│ kv cache │Gemma │ │
│ ┌──────────►│Expert│ │
│ │ │ │ │
│ ┌┴─────────┐ │x 10 │ │
│ │ │ └▲─────┘ │
│ │PaliGemma │ │ │
│ │ │ noise │
│ └▲──▲──▲──▲ │
│ │ │ │ └── discrete actions │
│ │ │ └───── robot state │
│ │ └──────── language tokens │
│ └─────────── image(s) │
└──────────────────────────────────────────┘
"""
# Fixed modality slots indexing the learnable modality embedding added on
# top of each token embedding when ``config.use_modality_embedding`` is set
# (see ``embed_prefix`` / ``embed_suffix``). The first five live in VLM
# hidden space (prefix tokens); the action expert is embedded separately in
# ``proj_width`` space. Slots are fixed (not a running counter over present
# blocks) so a modality gets the same embedding whether or not optional
# blocks (state, response, discrete actions) appear in a given batch.
_MODALITY_VISION = 0
_MODALITY_LANGUAGE = 1
_MODALITY_STATE = 2
_MODALITY_RESPONSE = 3
_MODALITY_DISCRETE_ACTION = 4
_NUM_PREFIX_MODALITIES = 5
[docs]
def __init__(
self,
config: PI05Config,
discrete_action_vocab_size: int | None = None,
language_tokenizer: AutoTokenizer | None = None,
):
"""Initializes the PI05FlowMatching model.
Args:
config: Model configuration.
discrete_action_vocab_size: Size of the discrete action vocabulary.
language_tokenizer: Optional pre-loaded PaliGemma tokenizer to share
with the enclosing `PI05Policy`. When ``None`` (e.g. unit tests
that construct the inner module directly) the tokenizer is
loaded here. Either way, the same instance is used by both
layers, and `ensure_loc_tokens` runs once.
"""
super().__init__()
self.config = config
paligemma_with_expert_config = PaliGemmaWithExpertConfig(
freeze_vision_encoder=self.config.freeze_vision_encoder,
train_expert_only=self.config.train_expert_only,
attention_implementation=self.config.attention_implementation,
discrete_action_vocab_size=discrete_action_vocab_size,
dropout=self.config.dropout,
gradient_checkpointing=self.config.gradient_checkpointing,
)
self.paligemma_with_expert = PaliGemmaWithExpertModel(paligemma_with_expert_config)
if self.config.state_type == "continuous":
vlm_hidden_size = self.paligemma_with_expert.config.paligemma_config.text_config.hidden_size
self.state_proj = nn.Linear(self.config.max_state_dim, vlm_hidden_size)
# Projections are float32
self.action_in_proj = nn.Linear(self.config.max_action_dim, self.config.proj_width)
self.action_out_proj = nn.Linear(self.config.proj_width, self.config.max_action_dim)
self.time_mlp_in = nn.Linear(self.config.proj_width, self.config.proj_width)
self.time_mlp_out = nn.Linear(self.config.proj_width, self.config.proj_width)
# Optional modality-specific learnable embedding added on top of the
# token embeddings (layered over RoPE — see ``config.use_modality_embedding``).
# Prefix tokens live in VLM hidden space; the action expert lives in
# ``proj_width`` space, so they get separate tables. Both are zero-init
# so turning the feature on for an existing checkpoint is a no-op at step 0.
if self.config.use_modality_embedding:
vlm_hidden_size = self.paligemma_with_expert.config.paligemma_config.text_config.hidden_size
self.modality_embedding = nn.Embedding(self._NUM_PREFIX_MODALITIES, vlm_hidden_size)
self.action_modality_embedding = nn.Embedding(1, self.config.proj_width)
nn.init.zeros_(self.modality_embedding.weight)
nn.init.zeros_(self.action_modality_embedding.weight)
if language_tokenizer is None:
language_tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
self.language_tokenizer = language_tokenizer
# PaliGemma reserves <loc0000>..<loc1023> at IDs 256000..257023, but
# the bare HF tokenizer does not register them as added/special
# tokens — a string "<loc0000>" otherwise BPE-fragments into seven
# pieces. This call promotes the reserved entries to single-token
# match mode (no new IDs, no embedding resize on PaliGemma) and
# mutates the shared tokenizer instance for `PI05Policy` too.
ensure_loc_tokens(self.language_tokenizer)
[docs]
def sample_noise(self, shape: tuple[int, ...], device: torch.device | str) -> Tensor:
"""Samples Gaussian noise.
Args:
shape: The shape of the noise tensor.
device: The device to create the tensor on.
Returns:
A tensor containing the sampled noise.
"""
noise = torch.normal(
mean=0.0,
std=1.0,
size=shape,
dtype=torch.float32,
device=device,
)
return noise
[docs]
def sample_time(self, bsize: int, device: torch.device | str) -> Tensor:
"""Samples time steps from a Beta distribution.
Args:
bsize: Batch size.
device: The device to create the tensor on.
Returns:
A tensor containing the sampled time steps.
"""
beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0)
time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=torch.float32)
time = time_beta * 0.999 + 0.001
return time
[docs]
def embed_prefix(
self,
images: list[Tensor],
img_masks: list[Tensor],
lang_tokens: Tensor,
lang_masks: Tensor,
response_tokens: Tensor | None = None,
response_masks: Tensor | None = None,
discrete_actions: Tensor | None = None,
discrete_action_masks: Tensor | None = None,
state: Tensor | None = None,
) -> tuple[Tensor, Tensor, Tensor, Tensor]:
"""Embed images with SigLIP and language tokens with embedding layer to prepare
for PaliGemma transformer processing.
When ``state_type == "continuous"`` and *state* is provided, the raw state
vector is projected into VLM embedding space and appended as a single
token (with full attention to images and language). Response tokens are
ignored in this mode.
Args:
images: List of image tensors.
img_masks: List of image mask tensors.
lang_tokens: Language token tensor.
lang_masks: Language mask tensor.
response_tokens: Optional Response language token tensor.
response_masks: Optional Response language mask tensor.
discrete_actions: Optional discrete action tensor.
discrete_action_masks: Optional discrete action mask tensor.
state: Optional continuous state tensor of shape (batch_size, max_state_dim).
indicator_tokens: Optional indicator token tensor.
indicator_masks: Optional indicator mask tensor.
Returns:
A tuple containing:
- embs: Concatenated embeddings tensor.
- pad_masks: Concatenated padding masks tensor.
- att_masks: Attention masks tensor.
- segment_ids: Per-token modality slot tensor (shape matches
``pad_masks``). Used to add the modality-specific embedding
on top of ``embs`` when ``config.use_modality_embedding`` is
set; also returned for downstream/diagnostic use.
"""
# TODO: avoid list in python and torch.cat ; prefer pre-allocation with torch.empty
embs = []
pad_masks = []
att_masks = []
# Per-token modality slot, parallel to ``att_masks``. Only consumed when
# ``config.use_modality_embedding`` is set, to add the learnable
# modality embedding on top of the token embeddings (see end of method).
segment_ids = []
# TODO: remove for loop
for (
img,
img_mask,
) in zip(images, img_masks, strict=False):
img_emb = self.paligemma_with_expert.embed_image(img)
img_emb = img_emb.to(dtype=_preferred_dtype())
# image embeddings don't need to be unnormalized because `fix/lerobot_openpi` branch of huggingface
# already removed the normalization inside PaliGemma
pass
bsize, num_img_embs = img_emb.shape[:2]
img_mask = img_mask[:, None].expand(bsize, num_img_embs)
embs.append(img_emb)
pad_masks.append(img_mask)
# Create attention masks so that image tokens attend to each other
att_masks += [0] * num_img_embs
segment_ids += [self._MODALITY_VISION] * num_img_embs
lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens)
# Normalize language embeddings
lang_emb_dim = lang_emb.shape[-1]
lang_emb = lang_emb * math.sqrt(lang_emb_dim)
embs.append(lang_emb)
pad_masks.append(lang_masks)
# full attention between image and language inputs
num_lang_embs = lang_emb.shape[1]
att_masks += [0] * num_lang_embs
segment_ids += [self._MODALITY_LANGUAGE] * num_lang_embs
if self.config.state_type == "continuous" and state is not None:
state_indicator_ids = self.language_tokenizer.encode("State: ", add_special_tokens=False)
state_indicator_tokens = torch.tensor([state_indicator_ids] * bsize, device=lang_tokens.device)
state_indicator_emb = self.paligemma_with_expert.embed_language_tokens(state_indicator_tokens)
state_indicator_emb = state_indicator_emb * math.sqrt(state_indicator_emb.shape[-1])
state_indicator_mask = torch.ones(
bsize, state_indicator_emb.shape[1], dtype=torch.bool, device=lang_tokens.device
)
embs.append(state_indicator_emb)
pad_masks.append(state_indicator_mask)
att_masks += [0] * state_indicator_emb.shape[1]
segment_ids += [self._MODALITY_STATE] * state_indicator_emb.shape[1]
state_emb = self.state_proj(state.to(dtype=_preferred_dtype()))
state_emb = rearrange(state_emb, "b d -> b 1 d")
state_mask = torch.ones(bsize, 1, dtype=torch.bool, device=state.device)
embs.append(state_emb)
pad_masks.append(state_mask)
att_masks += [0] # full attention with images and language
segment_ids += [self._MODALITY_STATE]
state_end_indicator_ids = self.language_tokenizer.encode(":\n", add_special_tokens=False)
state_end_indicator_tokens = torch.tensor(
[state_end_indicator_ids] * bsize, device=lang_tokens.device
)
state_end_indicator_emb = self.paligemma_with_expert.embed_language_tokens(
state_end_indicator_tokens
)
state_end_indicator_emb = state_end_indicator_emb * math.sqrt(state_end_indicator_emb.shape[-1])
state_end_indicator_mask = torch.ones(
bsize, state_end_indicator_emb.shape[1], dtype=torch.bool, device=lang_tokens.device
)
embs.append(state_end_indicator_emb)
pad_masks.append(state_end_indicator_mask)
att_masks += [0] * state_end_indicator_emb.shape[1]
segment_ids += [self._MODALITY_STATE] * state_end_indicator_emb.shape[1]
if self.config.predict_response:
response_indicator_ids = self.language_tokenizer.encode("Response: ", add_special_tokens=False)
response_indicator_tokens = torch.tensor(
[response_indicator_ids] * bsize, device=lang_tokens.device
)
response_indicator_emb = self.paligemma_with_expert.embed_language_tokens(
response_indicator_tokens
)
response_indicator_emb = response_indicator_emb * math.sqrt(response_indicator_emb.shape[-1])
response_indicator_mask = torch.ones(
bsize, response_indicator_emb.shape[1], dtype=torch.bool, device=lang_tokens.device
)
embs.append(response_indicator_emb)
pad_masks.append(response_indicator_mask)
att_masks += [1] * response_indicator_emb.shape[1]
segment_ids += [self._MODALITY_RESPONSE] * response_indicator_emb.shape[1]
if response_tokens is not None:
response_emb = self.paligemma_with_expert.embed_language_tokens(response_tokens)
# Normalize response language embeddings
response_emb_dim = response_emb.shape[-1]
response_emb = response_emb * math.sqrt(response_emb_dim)
embs.append(response_emb)
pad_masks.append(response_masks)
# full attention between image, language and response inputs
num_response_embs = response_emb.shape[1]
att_masks += [1] * num_response_embs
segment_ids += [self._MODALITY_RESPONSE] * num_response_embs
if discrete_actions is not None:
discrete_action_indicator_ids = self.language_tokenizer.encode(
"Action: ", add_special_tokens=False
)
discrete_action_indicator_tokens = torch.tensor(
[discrete_action_indicator_ids] * bsize, device=lang_tokens.device
)
discrete_action_indicator_emb = self.paligemma_with_expert.embed_language_tokens(
discrete_action_indicator_tokens
)
discrete_action_indicator_emb = discrete_action_indicator_emb * math.sqrt(
discrete_action_indicator_emb.shape[-1]
)
discrete_action_indicator_mask = torch.ones(
bsize, discrete_action_indicator_emb.shape[1], dtype=torch.bool, device=lang_tokens.device
)
embs.append(discrete_action_indicator_emb)
pad_masks.append(discrete_action_indicator_mask)
att_masks += [1] * discrete_action_indicator_emb.shape[1]
segment_ids += [self._MODALITY_DISCRETE_ACTION] * discrete_action_indicator_emb.shape[1]
discrete_action_emb = self.paligemma_with_expert.embed_discrete_actions(discrete_actions)
embs.append(discrete_action_emb.to(dtype=_preferred_dtype()))
pad_masks.append(discrete_action_masks)
att_masks += [1] * discrete_action_emb.shape[1]
segment_ids += [self._MODALITY_DISCRETE_ACTION] * discrete_action_emb.shape[1]
embs = torch.cat(embs, dim=1)
pad_masks = torch.cat(pad_masks, dim=1)
att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
segment_ids = torch.tensor(segment_ids, dtype=torch.long, device=pad_masks.device)
segment_ids = segment_ids[None, :].expand(bsize, segment_ids.shape[0])
# Add the learnable, modality-specific embedding on top of the token
# embeddings (layered over the existing RoPE positions, not a
# replacement). Zero-initialized, so this is a no-op until trained.
if self.config.use_modality_embedding:
modality_emb = self.modality_embedding(segment_ids)
embs = embs + modality_emb.to(dtype=embs.dtype)
return embs, pad_masks, att_masks, segment_ids
[docs]
def embed_suffix(self, noisy_actions: Tensor, timestep: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
"""Embed noisy_actions, timestep to prepare for Expert Gemma processing.
Args:
noisy_actions: Tensor containing noisy actions.
timestep: Tensor containing timesteps of shape (batch_size, action_chunk_length).
Returns:
A tuple containing:
- embs: Concatenated embeddings tensor.
- pad_masks: Concatenated padding masks tensor.
- att_masks: Attention masks tensor.
- adarms_cond: AdaRMS conditioning tensor.
"""
embs = []
pad_masks = []
att_masks = []
bsize = noisy_actions.shape[0]
dtype = _preferred_dtype()
device = noisy_actions.device
# Embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]
time_emb = create_sinusoidal_pos_embedding(
timestep, self.config.proj_width, min_period=4e-3, max_period=4.0, device=device
)
# Fuse timestep + action information using an MLP
noisy_actions = noisy_actions.to(dtype=dtype)
action_emb = self.action_in_proj(noisy_actions)
def time_mlp_func(time_emb):
x = self.time_mlp_in(time_emb)
x = F.silu(x)
x = self.time_mlp_out(x)
return F.silu(x)
time_emb = time_emb.to(dtype=dtype)
adarms_cond = time_mlp_func(time_emb)
# Add the learnable action-expert modality embedding (its own modality,
# in proj_width space), matching the prefix modality embedding in
# ``embed_prefix``. Zero-initialized, so a no-op until trained.
if self.config.use_modality_embedding:
action_emb = action_emb + self.action_modality_embedding.weight.to(dtype=action_emb.dtype)
# Add to input tokens
embs.append(action_emb)
bsize, action_dim = action_emb.shape[:2]
action_mask = torch.ones(bsize, action_dim, dtype=torch.bool, device=device)
pad_masks.append(action_mask)
# Set attention masks so that image, language and state inputs do not attend to action tokens.
# The action block spans the full chunk_size (= the noise/x_t length, here and in the
# training forward); n_action_steps is the execution horizon applied later in select_action,
# not the number of action tokens. Using n_action_steps here would mismatch the chunk_size-
# length pad mask and crash make_att_2d_masks when n_action_steps < chunk_size.
att_masks += [1] + ([0] * (self.config.chunk_size - 1))
embs = torch.cat(embs, dim=1)
pad_masks = torch.cat(pad_masks, dim=1)
att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
return embs, pad_masks, att_masks, adarms_cond
[docs]
def forward(
self,
images: list[Tensor],
img_masks: list[Tensor],
lang_tokens: Tensor,
lang_masks: Tensor,
actions: Tensor,
actions_is_pad: Tensor | None = None,
response_tokens: Tensor | None = None,
response_masks: Tensor | None = None,
noise: Tensor | None = None,
time: Tensor | None = None,
discrete_actions: Tensor | None = None,
discrete_action_masks: Tensor | None = None,
state: Tensor | None = None,
real_action_dim: Tensor | None = None,
return_per_sample: bool = False,
) -> dict[str, Tensor | PerSampleLoss]:
"""Do a full training forward pass and compute the loss.
Args:
images: List of image tensors.
img_masks: List of image mask tensors.
lang_tokens: Language token tensor.
lang_masks: Language mask tensor.
response_tokens: Response language token tensor.
response_masks: Response language mask tensor.
actions: Action tensor.
actions_is_pad: Optional action is padded mask tensor.
noise: Optional noise tensor.
time: Optional time tensor.
discrete_actions: Optional discrete action tensor.
discrete_action_masks: Optional discrete action mask tensor.
state: Optional continuous state tensor of shape (batch_size, max_state_dim).
indicator_tokens: Optional indicator token tensor.
indicator_masks: Optional indicator mask tensor.
Returns:
A dictionary containing the loss components ("MSE" and "CE").
"""
# Run VLM first to get key value cache
prefix_embs, prefix_pad_masks, prefix_att_masks, _ = self.embed_prefix(
images,
img_masks,
lang_tokens,
lang_masks,
response_tokens,
response_masks,
discrete_actions,
discrete_action_masks,
state=state,
)
vlm_2d_attention_mask = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
vlm_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
# avoids using discrete action for predicting continuous flow matching action
num_cross_att_tokens = (
prefix_embs.shape[1]
- self.config.discrete_action_indicator_max_length
- self.config.discrete_action_max_length
)
(prefix_out, _), past_key_values = self.paligemma_with_expert.forward(
attention_mask=vlm_2d_attention_mask,
position_ids=vlm_position_ids,
past_key_values=None,
inputs_embeds=[prefix_embs, None],
n_cross_att_tokens=num_cross_att_tokens,
use_cache=False,
fill_kv_cache=True,
)
# Now run action expert
batch_size = actions.shape[0]
if noise is None:
noise = self.sample_noise(actions.shape, actions.device)
if time is None:
time = self.sample_time(batch_size, actions.device)
# handle real time inference delay
delay = torch.randint(0, self.config.max_delay + 1, (batch_size,))
prefix_mask = rearrange(torch.arange(self.config.chunk_size), "c -> 1 c") < rearrange(
delay, "b -> b 1"
)
prefix_mask = prefix_mask.to(device=actions.device)
time = torch.where(
prefix_mask, 0, rearrange(time, "b -> b 1")
) # using diffusion time 0 instead of flow matching time 1
time_expanded = rearrange(time, "b c -> b c 1")
x_t = time_expanded * noise + (1 - time_expanded) * actions
u_t = noise - actions
suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, time)
action_expert_2d_attention_mask = make_att_2d_masks(
suffix_pad_masks,
suffix_att_masks,
n_cross_att_tokens=num_cross_att_tokens,
cross_att_pad_masks=prefix_pad_masks[:, :num_cross_att_tokens],
)
# We should skip the discrete action tokens as well as the discrete action indicator tokens when numbering the position ids for the action expert
prefix_offsets = torch.sum(
prefix_pad_masks[
:,
: -self.config.discrete_action_indicator_max_length - self.config.discrete_action_max_length,
],
dim=-1,
)[:, None] # action expert position ids start after prefix
action_expert_position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
# stop gradient to avoid backpropagating from action expert to VLM
if self.config.knowledge_insulation:
for layer_idx in past_key_values:
past_key_values[layer_idx]["key_states"] = past_key_values[layer_idx]["key_states"].detach()
past_key_values[layer_idx]["value_states"] = past_key_values[layer_idx][
"value_states"
].detach()
(_, suffix_out), _ = self.paligemma_with_expert.forward(
attention_mask=action_expert_2d_attention_mask,
position_ids=action_expert_position_ids,
past_key_values=past_key_values,
inputs_embeds=[None, suffix_embs],
use_cache=True,
fill_kv_cache=False,
adarms_cond=[None, adarms_cond],
)
# compute mse loss for velocity
# Supervise the whole chunk the model was trained to predict. n_action_steps
# is the inference-time execution horizon only and must not truncate the
# training target (chunk_size is the prediction horizon).
suffix_out = suffix_out[:, -self.config.chunk_size :]
# Original openpi code, upcast attention output
v_t = self.action_out_proj(suffix_out)
v_t = v_t.to(dtype=torch.float32)
# Shared masked-MSE reduction: AND-s frozen-prefix, timestep-pad, and
# dim-pad masks together and divides by the unmasked-slot count. See
# ``opentau.policies.utils.flow_matching_masked_mse`` for the full spec.
mse_result = flow_matching_masked_mse(
u_t=u_t,
v_t=v_t,
max_action_dim=self.config.max_action_dim,
prefix_mask=prefix_mask,
actions_is_pad=actions_is_pad,
real_action_dim=real_action_dim,
return_per_sample=return_per_sample,
)
mse_loss, mse_per_sample = mse_result if return_per_sample else (mse_result, None)
# compute cross entropy loss for discrete actions
batch_size, seq_len = discrete_actions.shape
discrete_token_start = -self.config.discrete_action_max_length
# The last token of response will predict the first token of discrete actions , so we need to slice from discrete_token_start -1.
# The predicted last token of discrete action is useless, so no need to include for loss calculation.
discrete_action_slice_object = slice(discrete_token_start - 1, -1)
discrete_action_out = prefix_out[:, discrete_action_slice_object]
logits = self.paligemma_with_expert.da_head(discrete_action_out)
logits = logits.to(dtype=torch.float32) # upcast to float32 for loss calculation
logits = rearrange(logits, "b s d -> (b s) d")
labels = rearrange(discrete_actions, "b s -> (b s)")
discrete_action_ce_loss = F.cross_entropy(logits, labels, reduction="none")
discrete_action_ce_loss = rearrange(discrete_action_ce_loss, "(b s) -> b s", b=batch_size, s=seq_len)
# remove pad tokens
discrete_action_is_pad = ~discrete_action_masks # convert into format where value for pad is True
discrete_action_ce_loss = discrete_action_ce_loss * ~discrete_action_is_pad
# Per-sample CE numerator/denominator over valid tokens, for the val breakdown.
discrete_action_ce_per_sample = (
ce_per_sample(discrete_action_ce_loss, ~discrete_action_is_pad) if return_per_sample else None
)
# compute mean
discrete_action_ce_loss = discrete_action_ce_loss.mean()
# compute cross entropy loss for response language only when pedict_response is set to true
if self.config.predict_response:
batch_size, seq_len = response_tokens.shape
response_token_start = (
-self.config.response_max_length
- self.config.discrete_action_max_length
- self.config.discrete_action_indicator_max_length
)
# The last token of language will predict <BOS> token of response, so no need to include for loss calculation. Hence slice starts from -self.config.discrete_action_max_length - self.config.response_max_length.
# The last token of response predicts first token of discrete actions, so no need to include for loss calculation. Hence slice ends at -self.config.discrete_action_max_length - 1.
response_token_end = (
-self.config.discrete_action_max_length - self.config.discrete_action_indicator_max_length - 1
)
response_slice_object = slice(response_token_start, response_token_end)
response_out = prefix_out[
:,
response_slice_object,
]
response_logits = self.paligemma_with_expert.paligemma.lm_head(response_out)
# response slice to exclude the <BOS> token from response while calculating loss.
response_slice = slice(1, None)
response_logits = response_logits.to(
dtype=torch.float32
) # upcast to float32 for loss calculation
response_logits = rearrange(response_logits, "b s d -> (b s) d")
response_labels = rearrange(response_tokens[:, response_slice], "b s -> (b s)")
response_ce_loss = F.cross_entropy(response_logits, response_labels, reduction="none")
response_ce_loss = rearrange(response_ce_loss, "(b s) -> b s", b=batch_size, s=seq_len - 1)
# remove pad tokens
response_is_pad = ~response_masks # convert into format where value for pad is True
# helps to control loss for response tokens in case of robotic data and VQA data
response_ce_loss = response_ce_loss * ~response_is_pad[:, response_slice]
# Per-sample response CE (valid tokens only) for the val breakdown.
response_ce_per_sample = (
ce_per_sample(response_ce_loss, ~response_is_pad[:, response_slice])
if return_per_sample
else None
)
# compute mean
response_ce_loss = response_ce_loss.mean()
else:
response_ce_loss = torch.tensor(0.0, device=mse_loss.device)
response_ce_per_sample = None
out: dict[str, Tensor | PerSampleLoss] = {
"MSE": mse_loss,
"CE": discrete_action_ce_loss + response_ce_loss,
}
if return_per_sample:
ce_ps = discrete_action_ce_per_sample
if response_ce_per_sample is not None:
ce_ps = ce_ps + response_ce_per_sample
out["MSE_per_sample"] = mse_per_sample
out["CE_per_sample"] = ce_ps
return out
[docs]
def sample_actions(
self,
images: list[Tensor],
img_masks: list[Tensor],
lang_tokens: Tensor,
lang_masks: Tensor,
action_prefix: Tensor,
delay: Tensor,
noise: Tensor | None = None,
state: Tensor | None = None,
) -> Tensor:
"""Do a full inference forward and compute the action.
Args:
images: List of image tensors.
img_masks: List of image mask tensors.
lang_tokens: Language token tensor.
lang_masks: Language mask tensor.
action_prefix: Action prefix tensor.
delay: Number of delay actions, aka number of actions frozen from the action_prefix.
noise: Optional noise tensor.
state: Optional continuous state tensor of shape (batch_size, max_state_dim).
indicator_tokens: Optional indicator token tensor.
indicator_masks: Optional indicator mask tensor.
Returns:
The sampled action tensor.
"""
bsize = lang_tokens.shape[0]
device = lang_tokens.device
if noise is None:
actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
noise = self.sample_noise(actions_shape, device)
prefix_embs, prefix_pad_masks, prefix_att_masks, _ = self.embed_prefix(
images,
img_masks,
lang_tokens,
lang_masks,
state=state,
)
prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] - 1
num_cross_att_tokens = prefix_embs.shape[1]
# Compute image and language key value cache
(prefix_out, _), past_key_values = self.paligemma_with_expert.forward(
attention_mask=prefix_att_2d_masks,
position_ids=prefix_position_ids,
past_key_values=None,
inputs_embeds=[prefix_embs, None],
n_cross_att_tokens=num_cross_att_tokens,
use_cache=False,
fill_kv_cache=True,
)
# initialize response tokens to empty tensor for storing response tokens during inference
response_tokens = torch.empty((bsize, 0), device=device, dtype=torch.long)
# if response prediction is enabled, then predict response tokens autoregressively
if self.config.predict_response:
for auto_step in range(self.config.response_max_length):
(
prefix_out,
prefix_embs,
prefix_pad_masks,
prefix_att_masks,
prefix_offsets,
response_tokens,
past_key_values,
) = self.infer_response(
prefix_out,
prefix_embs,
prefix_pad_masks,
prefix_att_masks,
past_key_values,
prefix_offsets,
response_tokens,
auto_step,
bsize,
device,
)
# perform denoising steps to get the action
dt = -1.0 / self.config.num_steps
dt = torch.tensor(dt, dtype=torch.float32, device=device)
x_t = noise
time = torch.tensor(1.0, dtype=torch.float32, device=device)
prefix_mask = rearrange(torch.arange(self.config.chunk_size, device=device), "c -> 1 c") < delay
while time >= -dt / 2:
# if delay is greater than 0, then freeze the action prefix at the beginning of action chunk
x_t = torch.where(rearrange(prefix_mask, "b c -> b c 1"), action_prefix, x_t)
masked_time = torch.where(prefix_mask, 0, time)
v_t = self.denoise_step(
prefix_pad_masks,
past_key_values,
x_t,
masked_time,
)
# Euler step
x_t += dt * v_t
time += dt
# we need to ensure the frozen actions are not modified before returning the denoised actions
x_t = torch.where(rearrange(prefix_mask, "b c -> b c 1"), action_prefix, x_t)
return x_t
[docs]
def denoise_step(
self,
prefix_pad_masks: Tensor,
past_key_values: list[dict[str, Tensor]],
x_t: Tensor,
time: Tensor,
) -> Tensor:
"""Apply one denoising step of the noise `x_t` at a given timestep.
Args:
prefix_pad_masks: Prefix padding masks.
past_key_values: Past key values from the VLM.
x_t: Current noise tensor.
time: Time tensor of shape (batch_size, action_chunk_length).
Returns:
The predicted velocity tensor (v_t).
"""
suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, time)
num_cross_att_tokens = prefix_pad_masks.shape[1]
action_expert_2d_attention_mask = make_att_2d_masks(
suffix_pad_masks,
suffix_att_masks,
n_cross_att_tokens=num_cross_att_tokens,
cross_att_pad_masks=prefix_pad_masks[:, :num_cross_att_tokens],
)
prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[
:, None
] # action expert position ids start after prefix
action_expert_position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
outputs_embeds, _ = self.paligemma_with_expert.forward(
attention_mask=action_expert_2d_attention_mask,
position_ids=action_expert_position_ids,
past_key_values=past_key_values,
inputs_embeds=[None, suffix_embs],
use_cache=True,
fill_kv_cache=False,
adarms_cond=[None, adarms_cond],
)
suffix_out = outputs_embeds[1]
# Denoise the full chunk_size chunk so v_t matches x_t in the Euler step.
# n_action_steps (execution horizon) is applied later in select_action, not
# at decode time.
suffix_out = suffix_out[:, -self.config.chunk_size :]
v_t = self.action_out_proj(suffix_out)
v_t = v_t.to(dtype=torch.float32)
return v_t
[docs]
def infer_response(
self,
prefix_out: Tensor,
prefix_embs: Tensor,
prefix_pad_masks: Tensor,
prefix_att_masks: Tensor,
past_key_values: list[dict[str, Tensor]],
prefix_offsets: Tensor,
response_tokens: Tensor,
auto_step: int,
bsize: int,
device: torch.device,
) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, list[dict[str, Tensor]], Tensor]:
"""Perform autoregressive inference for response generation.
This method generates the next token in the response sequence, updating the
various state tensors required for maintaining the generation context. It handles
the initial BOS token generation as well as subsequent tokens, and manages
padding masks to handle variable-length sequences properly.
Args:
prefix_out: Output tensor from the previous step.
prefix_embs: Embeddings for the current prefix context.
prefix_pad_masks: Boolean mask indicating valid (non-padding) tokens in the prefix.
prefix_att_masks: Attention mask for the prefix.
past_key_values: KV cache from previous transformer steps.
prefix_offsets: Position offsets for the current generation step.
response_tokens: Accumulated tokens generated so far.
auto_step: Current autoregressive step index (0 for first token).
bsize: Batch size.
device: Device to run the computation on.
Returns:
A tuple containing updated tensors for the next step:
(prefix_out, prefix_embs, prefix_pad_masks, prefix_att_masks,
prefix_offsets, response_tokens, past_key_values, response_token)
"""
EOS_TOKEN = self.language_tokenizer.convert_tokens_to_ids(self.language_tokenizer.eos_token) # noqa: N806
if auto_step == 0:
# Start the autoregressive inference with <bos> token
response_token = torch.full(
(bsize, 1),
self.language_tokenizer.bos_token_id,
device=device,
dtype=torch.long,
)
else:
# get the last predicted token from the prefix output which is predicted response
response_token = prefix_out[:, -1:]
response_token = self.paligemma_with_expert.paligemma.lm_head(response_token).argmax(dim=-1)
PAD_TOKEN = self.language_tokenizer.pad_token_id # noqa: N806
# Create pad masks: False if previous token was EOS or PAD
if response_tokens.shape[1] > 1:
prev_tokens = response_tokens
has_eos = (prev_tokens == EOS_TOKEN).any(dim=1, keepdim=True)
has_pad = (prev_tokens == PAD_TOKEN).any(dim=1, keepdim=True)
# check if the previous token was EOS or PAD. If so, then the current token should be padded, so its not attended by flow matching action expert.
response_pad_masks = ~(has_eos | has_pad)
response_token = torch.where(
response_pad_masks,
response_token,
torch.tensor(PAD_TOKEN, device=device, dtype=response_token.dtype),
)
else:
response_pad_masks = torch.ones((bsize, 1), device=device, dtype=torch.bool)
# Updating response tokens with current predicted token
response_tokens = torch.cat([response_tokens, response_token], dim=1)
# Embed the current predicted token
response_emb = self.paligemma_with_expert.embed_language_tokens(response_token)
# Normalize response language embeddings
response_emb_dim = response_emb.shape[-1]
response_emb = response_emb * math.sqrt(response_emb_dim)
response_att_masks = torch.ones((bsize, 1), device=device, dtype=response_emb.dtype)
# update the prefix embs, pad masks and att masks, so it can be used by action experts
prefix_embs = torch.cat([prefix_embs, response_emb], dim=1)
prefix_pad_masks = torch.cat([prefix_pad_masks, response_pad_masks], dim=1)
prefix_att_masks = torch.cat([prefix_att_masks, response_att_masks], dim=1)
num_cross_att_tokens = prefix_pad_masks.shape[1]
# create the attention mask for the response tokens
response_att_2d_masks = make_att_2d_masks(
response_pad_masks,
response_att_masks,
n_cross_att_tokens=num_cross_att_tokens - 1,
cross_att_pad_masks=prefix_pad_masks[:, : num_cross_att_tokens - 1],
)
prefix_offsets = prefix_offsets + response_pad_masks.long()
prefix_position_ids = prefix_offsets
# Compute image and language key value cache
(prefix_out, _), past_key_values = self.paligemma_with_expert.forward(
attention_mask=response_att_2d_masks,
position_ids=prefix_position_ids,
past_key_values=past_key_values,
inputs_embeds=[response_emb, None],
n_cross_att_tokens=num_cross_att_tokens,
use_cache=True,
fill_kv_cache=True,
)
return (
prefix_out,
prefix_embs,
prefix_pad_masks,
prefix_att_masks,
prefix_offsets,
response_tokens,
past_key_values,
)