opentau.policies.pi06.gemma3_with_expert

Gemma 3 backbone with Gemma-v1 action expert, for the PI06 policy.

Mirrors paligemma_with_expert.py but:

• the vision-language backbone is Gemma3ForConditionalGeneration (Gemma 3 4B, SigLIP-400m/14 + Gemma 3 text, 34 interleaved sliding-window/global layers);

• the action expert is a Gemma-v1 GemmaForCausalLM with the AdaRMS and gated-residual patches applied by opentau.utils.transformers_patch.

The per-layer attention loop below concatenates backbone and expert queries/ keys/values along the sequence dimension at every layer (the MoE-like pattern introduced in π0), so the expert can cross-attend to the backbone’s activations at every depth. Gemma 3 specifics (q_norm/k_norm, pre/post feedforward RMSNorms, per-layer local vs global RoPE, sliding-window attention) are all honored.

transformers_patch is imported at module load so the expert path picks up adaptive RMSNorm and _gated_residual. The Gemma 3 backbone remains stock — its layer-norms return a plain tensor and are used without a cond= argument.

Functions

apply_rope(x, positions[, max_wavelength])

Applies RoPE to x with the given positions and base wavelength.

Classes

Gemma3WithExpertConfig([gemma3_config, ...])	Configuration wrapper bundling a Gemma 3 VLM config and a Gemma-v1 expert config.
Gemma3WithExpertModel(config)	Gemma 3 VLM interleaved layer-wise with a Gemma-v1 action expert.

class opentau.policies.pi06.gemma3_with_expert.Gemma3WithExpertConfig(gemma3_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)

Bases: PretrainedConfig

Configuration wrapper bundling a Gemma 3 VLM config and a Gemma-v1 expert config.

__init__(gemma3_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)

Initializes the configuration.

Parameters:

• gemma3_config – Optional Gemma 3 config dict. Defaults to the google/gemma-3-4b-pt topology.

• gemma_expert_config – Optional Gemma-v1 action-expert config dict. Defaults to a ~860M-parameter Gemma with AdaRMS enabled.

• freeze_vision_encoder – Freeze the SigLIP tower during training.

• train_expert_only – Only update the expert and its heads.

• attention_implementation – “eager”, “sdpa”, or “fa2”. “fa2” is not implemented and falls back to eager with a warning. “sdpa” dispatches to torch.nn.functional.scaled_dot_product_attention; see the per-layer note about Gemma 3’s interleaved local/global pattern in forward() — π0.6 deliberately keeps the same block-causal mask at every layer, so the SDPA call sees a regular bool mask and takes the standard fused path.

• discrete_action_vocab_size – FAST tokenizer vocab size.

• dropout – Dropout probability applied in the per-layer loop.

• gradient_checkpointing – Wrap each interleaved decoder-layer body in torch.utils.checkpoint.checkpoint during training. Trades roughly one extra forward pass per step (~25-33% compute) for a large slice of activation memory per rank, enabling larger per-rank batch sizes. Only safe under plain DDP (MULTI_GPU), single-process (NO), or DeepSpeed ZeRO-1/2 — see the train.py guard. Defaults to False.

• **kwargs – Passed to PretrainedConfig.

model_type: str = 'Gemma3WithExpertModel'

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

class opentau.policies.pi06.gemma3_with_expert.Gemma3WithExpertModel(config: Gemma3WithExpertConfig)

Bases: PreTrainedModel

Gemma 3 VLM interleaved layer-wise with a Gemma-v1 action expert.

__init__(config: Gemma3WithExpertConfig)

Args: config ([PretrainedConfig]):

Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [~PreTrainedModel.from_pretrained] method to load the model weights.

config_class

alias of Gemma3WithExpertConfig

eager_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor, scaling: float | None = None) → Tensor

Standard eager scaled-dot-product attention. attention_mask is a boolean 2D mask of shape (B, Q, K) (True = attend).

embed_discrete_actions(actions: Tensor) → Tensor

embed_image(image: Tensor) → Tensor

Runs the SigLIP tower + multimodal projector to obtain image tokens.

embed_language_tokens(tokens: Tensor) → Tensor

Embed token ids through Gemma 3’s shared text embedding table.

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, n_cross_att_tokens: int | None = None, use_cache: bool | None = None, fill_kv_cache: bool | None = None, adarms_cond: list[Tensor] | None = None) → tuple[list[FloatTensor | None], list[FloatTensor] | Cache | None]

Interleaved per-layer forward for the Gemma 3 backbone and Gemma-v1 expert.

The two streams (index 0 = backbone, index 1 = expert) share each layer’s attention — queries and KVs are concatenated along the sequence axis. When one stream’s embeddings are None the other runs alone, pulling KVs for the missing stream from past_key_values when use_cache=True.

Parameters:

• attention_mask – 2D boolean mask of shape (B, Q, K_total). See opentau.policies.pi05.modeling_pi05.make_att_2d_masks.

• position_ids – (B, L_total) token positions, used for RoPE.

• past_key_values – Per-layer KV cache populated on a previous call.

• inputs_embeds – [backbone_embeds, expert_embeds]. Either may be None.

• n_cross_att_tokens – Number of prefix tokens to retain in the cache (must be provided when fill_kv_cache=True).

• use_cache – Read KVs from past_key_values (prefix cross-attention).

• fill_kv_cache – Write this call’s KVs into past_key_values.

• adarms_cond – Per-stream AdaRMS conditioning tensors [None, cond].

Returns:

A pair (outputs_embeds, past_key_values) where outputs_embeds is a two-element list mirroring the inputs_embeds layout.

get_attention_interface()

Returns the attention implementation function based on config.

Dispatches on self.config.attention_implementation:

• "eager": per-layer matmul-softmax-matmul in fp32 (historical default; see eager_attention_forward).

• "sdpa": torch.nn.functional.scaled_dot_product_attention which on A100 + bf16 dispatches to FlashAttention-2 / mem-efficient backends. Note π0.6 keeps the same block-causal mask at every layer (sliding window deliberately not enforced — see the comment near apply_rope), so SDPA sees a regular bool mask and does not need a per-layer mask shape branch.

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

sdpa_attention_forward(attention_mask: Tensor, batch_size: int, head_dim: int, query_states: Tensor, key_states: Tensor, value_states: Tensor, scaling: float | None = None) → Tensor

SDPA attention forward pass using F.scaled_dot_product_attention.

Same output shape and semantics 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. Q/K are not upcast to float32 (modern attention kernels accumulate the softmax in fp32 internally); training dynamics match eager within bf16 reassociation noise.

Parameters:

• attention_mask – Boolean mask of shape (B, Q, K_total); True = attend. π0.6 keeps the same block-causal mask at every layer.

• batch_size – Batch size.

• head_dim – Per-head dimension.

• query_states – (B, Q, num_attention_heads, head_dim).

• key_states – (B, K_total, num_key_value_heads, head_dim).

• value_states – (B, K_total, num_key_value_heads, head_dim).

• scaling – Override for the QK scaling factor. Gemma 3 uses query_pre_attn_scalar ** -0.5 rather than the default head_dim ** -0.5; the caller passes this through.

Returns:

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

Return type:

torch.Tensor

set_requires_grad() → None

to_bfloat16_like_physical_intelligence() → None

train(mode: bool = True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

opentau.policies.pi06.gemma3_with_expert.apply_rope(x: Tensor, positions: Tensor, max_wavelength: float = 10000.0) → Tensor

Applies RoPE to x with the given positions and base wavelength.

Parameters:

• x – Tensor of shape (B, L, H, D).

• positions – Tensor of shape (B, L).

• max_wavelength – RoPE base frequency. Gemma 3 uses 10_000 for sliding (local) layers and 1_000_000 for full (global) layers; Gemma-v1 expert uses 10_000.

Returns:

RoPE-transformed tensor, same shape / dtype as the input.