opentau.policies.pi0.paligemma_with_expert

PaliGemma with Expert Module for PI0.

This module implements the PaliGemma model with an additional expert module, specifically designed for the Pi0 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 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.pi0.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', dropout: float = 0.1, gradient_checkpointing: bool = False, **kwargs)

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', dropout: float = 0.1, gradient_checkpointing: bool = False, **kwargs)

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”.

• 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.pi0.paligemma_with_expert.PaliGemmaWithExpertModel(config: PaliGemmaWithExpertConfig)

Bases: PreTrainedModel

PaliGemma model with an additional expert module for action generation.

__init__(config: PaliGemmaWithExpertConfig)

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

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_image(image: Tensor) → Tensor

Computes image embeddings.

Parameters:

image – Input image tensor.

Returns:

Image embeddings.

Return type:

torch.Tensor

embed_language_tokens(tokens: Tensor) → Tensor

Embeds language tokens.

Parameters:

tokens – Input token indices.

Returns:

Token embeddings.

Return type:

torch.Tensor

forward(attention_mask: Tensor | None = None, position_ids: LongTensor | None = None, past_key_values: list[FloatTensor] | Cache | None = None, inputs_embeds: list[FloatTensor] | None = None, use_cache: bool | None = None, fill_kv_cache: bool | None = None) → tuple[list[FloatTensor | None], list[FloatTensor] | Cache | None]

Forward pass of the model.

Parameters:

• attention_mask – Attention mask tensor.

• position_ids – Position IDs tensor.

• past_key_values – Past key values for caching.

• inputs_embeds – List of input embeddings for the different model parts.

• use_cache – Whether to use KV cache.

• fill_kv_cache – Whether to fill the KV cache.

Returns:

A tuple containing:

• outputs_embeds: List of output embeddings.

• past_key_values: Updated past key values.

Return type:

tuple

get_attention_interface()

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 — typically 2-3x faster than eager.

• "fa2": accepted for backward compatibility; falls back to eager with a warning emitted at config validation time.

Returns:

The attention function to use.

Return type:

callable

sdpa_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor) → Tensor

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

Sets the requires_grad attribute for model parameters based on configuration.

to_bfloat16_like_physical_intelligence() → None

Casts specific model components to bfloat16 dtype.

train(mode: bool = True) → None

Sets the module in training mode.

Parameters:

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

opentau.policies.pi0.paligemma_with_expert.apply_rope(x: Tensor, positions: Tensor, max_wavelength: int = 10000) → Tensor

Applies RoPE positions to the input tensor.

Parameters:

• x – Input tensor of shape [B, L, H, D].

• positions – Position tensor of shape [B, L].

• 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