opentau.policies.pi05.paligemma_with_expert

PaliGemma with Expert Module.

This module implements the PaliGemma model with an additional expert module, specifically designed for the Pi05 policy. It combines a pre-trained PaliGemma Vision-Language Model (VLM) with a Gemma-based expert model to handle action generation and conditioning.

Functions

apply_rope(x, positions[, max_wavelength])

Applies (multimodal) RoPE positions to the input tensor.

Classes

`PaliGemmaWithExpertConfig`([...])	Configuration class for PaliGemmaWithExpertModel.
`PaliGemmaWithExpertModel`(config)	PaliGemma model with an additional expert module for action generation.

class opentau.policies.pi05.paligemma_with_expert.PaliGemmaWithExpertConfig(paligemma_config: dict | None = None, gemma_expert_config: dict | None = None, freeze_vision_encoder: bool = True, train_expert_only: bool = True, attention_implementation: str = 'eager', discrete_action_vocab_size: int | None = None, dropout: float = 0.1, gradient_checkpointing: bool = False, **kwargs)[source]

Bases: PretrainedConfig

Configuration class for PaliGemmaWithExpertModel.

__init__(paligemma_config: dict | None = None, gemma_expert_config: dict | None = None, freeze_vision_encoder: bool = True, train_expert_only: bool = True, attention_implementation: str = 'eager', discrete_action_vocab_size: int | None = None, dropout: float = 0.1, gradient_checkpointing: bool = False, **kwargs)[source]

Initializes the configuration.

Parameters:

paligemma_config – Configuration dictionary for the PaliGemma model.
gemma_expert_config – Configuration dictionary for the Gemma expert model.
freeze_vision_encoder – Whether to freeze the vision encoder. Defaults to True.
train_expert_only – Whether to train only the expert model. Defaults to True.
attention_implementation – Attention implementation to use (“eager”, “sdpa”, or “fa2”). Defaults to “eager”.
discrete_action_vocab_size – Vocabulary size for discrete actions.
dropout – Dropout probability. Defaults to 0.1.
gradient_checkpointing – Wrap each decoder-layer body in torch.utils.checkpoint.checkpoint during training. Trades roughly one extra forward pass per step (~25-33% compute) for ~30-40 GB of activation memory per rank, which enables larger per-rank batch sizes and amortizes per-step fixed cost. Only safe under plain DDP (MULTI_GPU), single-process (NO), or DeepSpeed ZeRO-1/2 — see the train.py guard. Defaults to False.
**kwargs – Additional keyword arguments passed to PretrainedConfig.

model_type: str = 'PaliGemmaWithExpertModel'

sub_configs: dict[str, type['PretrainedConfig']] = {'gemma_expert_config': <class 'transformers.models.auto.configuration_auto.AutoConfig'>, 'paligemma_config': <class 'transformers.models.auto.configuration_auto.AutoConfig'>}

class opentau.policies.pi05.paligemma_with_expert.PaliGemmaWithExpertModel(config: PaliGemmaWithExpertConfig)[source]

Bases: PreTrainedModel

PaliGemma model with an additional expert module for action generation.

__init__(config: PaliGemmaWithExpertConfig)[source]

Initializes the PaliGemmaWithExpertModel.

Parameters:: config – Configuration object for the model.

config_class: alias of PaliGemmaWithExpertConfig

eager_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor) → Tensor[source]

Eager attention forward pass using standard matrix multiplications.

Parameters:

attention_mask – Attention mask tensor.
batch_size – Batch size.
head_dim – Head dimension.
query_states – Query states tensor.
key_states – Key states tensor.
value_states – Value states tensor.

Returns:

Attention output.

Return type:

torch.Tensor

embed_discrete_actions(actions: Tensor) → Tensor[source]

Embeds discrete action tokens.

Parameters:: actions – Input discrete action indices.
Returns:: Action embeddings.
Return type:: torch.Tensor

embed_image(image: Tensor) → Tensor[source]

Computes image embeddings.

Parameters:: image – Input image tensor.
Returns:: Image embeddings.
Return type:: torch.Tensor

embed_language_tokens(tokens: Tensor) → Tensor[source]

Embeds language tokens.

Parameters:: tokens – Input token indices.
Returns:: Token embeddings.
Return type:: torch.Tensor

flash_attention_forward(attention_block_ids: tuple[Tensor, ...], batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor) → Tensor[source]

Custom block-causal CUDA flash attention forward.

Unlike the eager/sdpa interfaces this consumes the compact block-id representation (q_blk, k_blk, q_valid, k_valid) and reconstructs the block-causal mask inside the kernel, so the dense (B, Sq, Sk) mask is never materialized. GQA/MQA is handled natively (no K/V expansion), and the kernel runs both matmuls on Tensor Cores (fp16/bf16) with fp32 accumulation, matching eager/sdpa outputs within fp/bf16 noise.

Parameters:

attention_block_ids – (q_blk, k_blk, q_valid, k_valid) from flash_attn_cuda.make_att_block_ids.
batch_size – Batch size.
head_dim – Per-head dimension.
query_states – (B, Sq, num_attention_heads, head_dim).
key_states – (B, Sk, num_key_value_heads, head_dim).
value_states – (B, Sk, num_key_value_heads, head_dim).

Returns:

(B, Sq, num_attention_heads * head_dim) output, matching the eager/sdpa interfaces.

Return type:

torch.Tensor

Forward pass of the model.

Parameters:

attention_mask – Dense bool attention mask (B, Sq, Sk), True=attend. Used by the eager/sdpa backends. None when the flash_cuda backend is active (the mask is reconstructed in-kernel instead).
attention_block_ids – Compact (q_blk, k_blk, q_valid, k_valid) block-id representation for the flash_cuda backend (see flash_attn_cuda.make_att_block_ids). None for eager/sdpa.
position_ids – Position IDs tensor.
past_key_values – Past key values for caching.
inputs_embeds – List of input embeddings for the different model parts.
n_cross_att_tokens – Number of cross-attention tokens.
use_cache – Whether to use KV cache.
fill_kv_cache – Whether to fill the KV cache.
adarms_cond – List of AdaRMS conditioning tensors.

Returns:

A tuple containing:

outputs_embeds: List of output embeddings.
past_key_values: Updated past key values.

Return type:

tuple

Raises:

ValueError – If n_cross_att_tokens is not provided when fill_kv_cache is True.

get_attention_interface()[source]

Returns the attention implementation function based on config.

Dispatches on self.config.attention_implementation:

"eager": naive matmul-softmax-matmul in fp32 (historical default; see eager_attention_forward).
"sdpa": torch.nn.functional.scaled_dot_product_attention which on A100 dispatches to FlashAttention-2 or the memory-efficient backend — 2-3x faster than eager at pi05’s sequence length.
"fa2": accepted for backward compatibility; falls back to eager with a warning emitted at config validation time.
"flash_cuda": custom block-causal CUDA flash kernel, dispatched in _run_layer via flash_attention_forward whenever attention_block_ids is supplied. There is no silent fallback: if this method is ever reached under flash_cuda it means a caller failed to thread the block-ids, which is a bug, so it raises.

Returns:: The attention function to use.
Return type:: callable
Raises:: RuntimeError – If attention_implementation == "flash_cuda" (the flash path must go through flash_attention_forward with block-ids, not this interface).

sdpa_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor) → Tensor[source]

SDPA attention forward pass using F.scaled_dot_product_attention.

Produces the same output shape and semantic as eager_attention_forward but delegates the scores-softmax-matmul chain to PyTorch’s fused SDPA kernel. On A100 + bf16 PyTorch typically dispatches to FlashAttention-2, which keeps the S×S attention matrix in on-chip SRAM and runs the matmul in bf16 with fp32 accumulation in the softmax. There is a deliberate numerical difference vs. the eager kernel: we do not upcast Q/K to float32 before the matmul, because modern attention kernels do fp32 accumulation internally — cleaner and faster. Training dynamics are equivalent within bf16 reassociation noise.

Parameters:

attention_mask – Boolean mask of shape (B, S, S); True = attend.
batch_size – Batch size.
head_dim – Per-head dimension.
query_states – Query states of shape (B, S, num_attention_heads, head_dim).
key_states – Key states of shape (B, S, num_key_value_heads, head_dim).
value_states – Value states of shape (B, S, num_key_value_heads, head_dim).

Returns:

Attention output of shape (B, S, num_attention_heads * head_dim).

Return type:

torch.Tensor

set_requires_grad() → None[source]: Sets the requires_grad attribute for model parameters based on configuration.

to_bfloat16_like_physical_intelligence() → None[source]: Casts specific model components to bfloat16 dtype.

train(mode: bool = True) → None[source]

Sets the module in training mode.

Parameters:: mode – whether to set training mode (True) or evaluation mode (False). Defaults to True.

opentau.policies.pi05.paligemma_with_expert.apply_rope(x: Tensor, positions: Tensor, max_wavelength: int = 10000) → Tensor[source]

Applies (multimodal) RoPE positions to the input tensor.

Two position layouts are supported, selected by positions.ndim:

1-D RoPE — positions of shape [B, L]: a single scalar position per token (the historical behaviour). Every rotary frequency band is driven by that one position.
Interleaved MRoPE — positions of shape [A, B, L] (here A == 3 for the temporal / height / width axes): frequency band i is driven by axis i % A, so the three axes are interleaved across the full frequency spectrum rather than each owning a contiguous block (the Qwen3-VL “interleaved MRoPE” scheme). When all A axes carry the same positions this is bit-identical to the 1-D path, so text-style tokens (t == h == w) rotate exactly as plain RoPE.

Parameters:

x – Input tensor of shape [B, L, H, D].
positions – Position tensor of shape [B, L] (1-D RoPE) or [A, B, L] (interleaved MRoPE).
max_wavelength – Maximum wavelength for RoPE. Defaults to 10_000.

Returns:

The input tensor with RoPE applied, of shape [B, L, H, D].

Return type:

Tensor