opentau.policies.cosmos3.qwen3vl_with_expert

Qwen3-VL backbone + custom flow-matching action expert (cosmos3).

This is the cosmos3 analog of paligemma_with_expert.py / gemma3_with_expert.py, adapted to the Qwen3-VL architecture and to a frozen backbone.

Design (π0.5 dual-stream, specialized for a frozen reasoner)

In the π0.5 recipe the backbone encodes images + language once (the “prefix”) and the action expert cross-attends to the backbone’s per-layer key/value cache to denoise the action chunk (the “suffix”). The two streams never run in the same forward pass — the prefix forward populates the KV cache; the suffix forward only runs the expert, reading that cache. The expert reads the backbone’s keys/values but the backbone never reads the expert.

That structure lets cosmos3 run the entire stock Qwen3VLModel.forward as a black box for the prefix (Qwen3VLWithExpertModel.run_prefix). Stock transformers handles vision encoding, image-token scatter, deepstack injection, the 3-D multimodal RoPE (MRoPE), QK-norm and the native causal mask — so the frozen Cosmos3-Super reasoning tower behaves exactly as trained, with zero reimplementation of the 32B backbone. Only the small (<1B) action expert is hand-written here.

Hard cross-attention constraints (validated at build time)

The expert’s per-layer keys/values are concatenated with the backbone’s cached KV, so the expert’s num_key_value_heads and head_dim must equal the backbone text tower’s (8 / 128 for Qwen3-VL-32B). The expert’s query head count is free (any multiple of the KV head count) because the backbone’s queries are never consumed here. The shared MRoPE (cos, sin) are produced by the backbone’s rotary_emb and reused by the expert, which is why expert_head_dim must match.

Which backbone layer the expert reads (condition_on_layer)

By default (condition_on_layer=None) the expert mirrors the backbone one-for-one: expert layer i reads the cached KV of backbone layer i, so the expert must have the same number of layers (64) as the backbone. Setting condition_on_layer=k flips this to single-layer conditioning: every expert layer cross-attends to the cached KV of backbone layer k only. Two consequences follow and are both handled here:

• The layer-count equality is dropped – the expert may be any depth >= 1 (a shallower, cheaper expert all reading the one rich layer).

• The frozen backbone is truncated to its first k + 1 text layers at build time (text_config.num_hidden_layers is lowered before the backbone is constructed / loaded), so layers k+1..N-1 are never allocated or run. The selected layer’s KV is bit-identical with or without truncation: Qwen3-VL injects deepstack vision features only into the earliest layers (after layer j for j < len(features)), so the output of layer k – and hence its KV – depends only on layers 0..k.

Classes

AdaRMSNorm(dim, cond_dim[, eps])	Adaptive RMSNorm (DiT adaLN-Zero style), conditioned on the flow-matching time.
ExpertRMSNorm(dim[, eps])	Plain RMSNorm used for the expert's per-head QK normalization.
Qwen3ActionExpert(num_hidden_layers, ...)	The trainable flow-matching action expert: a stack of AdaRMS Qwen3 decoder layers.
Qwen3ExpertAttention(hidden_size, ...)	Expert self/cross attention: expert queries over [cached_backbone_KV ; expert_KV].
Qwen3ExpertDecoderLayer(hidden_size, ...)	One expert decoder layer: AdaRMS pre-norms + gated residuals around attn / MLP.
Qwen3ExpertMLP(hidden_size, intermediate_size)	Qwen3 gated-SiLU MLP (gate/up/down), matching Qwen3VLTextMLP.
Qwen3VLWithExpertModel(qwen3vl_config, *, ...)	A frozen Qwen3-VL backbone paired with a trainable flow-matching action expert.

class opentau.policies.cosmos3.qwen3vl_with_expert.AdaRMSNorm(dim: int, cond_dim: int, eps: float = 1e-06)

Bases: Module

Adaptive RMSNorm (DiT adaLN-Zero style), conditioned on the flow-matching time.

From the conditioning vector cond (the time embedding) a single dense layer produces per-channel (scale, shift, gate); the normalized input is modulated as norm(x) * (1 + scale) + shift and the gate is returned for the gated residual x + gate * sublayer(x). The dense layer is zero-initialized so the expert starts as an exact identity/no-op residual on top of the frozen backbone — the stable adaLN-Zero initialization.

Mirrors opentau.utils.transformers_patch.PatchedGemmaRMSNorm (the AdaRMS used by the pi05/pi06 Gemma action experts) but as a standalone Qwen3 variant.

__init__(dim: int, cond_dim: int, eps: float = 1e-06)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, cond: Tensor) → tuple[Tensor, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.cosmos3.qwen3vl_with_expert.ExpertRMSNorm(dim: int, eps: float = 1e-06)

Bases: Module

Plain RMSNorm used for the expert’s per-head QK normalization.

Mirrors Qwen3VLTextRMSNorm (variance in fp32, learned weight), kept as a standalone module so the frozen backbone’s norm class is never monkey-patched.

__init__(dim: int, eps: float = 1e-06)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.cosmos3.qwen3vl_with_expert.Qwen3ActionExpert(num_hidden_layers: int, hidden_size: int, intermediate_size: int, num_attention_heads: int, num_key_value_heads: int, head_dim: int, adarms_cond_dim: int, rms_norm_eps: float, dropout: float, attention_implementation: str)

Bases: Module

The trainable flow-matching action expert: a stack of AdaRMS Qwen3 decoder layers.

Operates on action-token embeddings (the suffix). Layer i cross-attends to cached_kv[i] plus the action chunk itself. The expert is mode-agnostic: the caller (Qwen3VLWithExpertModel.run_expert) hands it a cached_kv list whose length equals the expert depth – the per-layer backbone cache in the default regime, or the single selected layer’s KV broadcast to every position in single-layer mode.

__init__(num_hidden_layers: int, hidden_size: int, intermediate_size: int, num_attention_heads: int, num_key_value_heads: int, head_dim: int, adarms_cond_dim: int, rms_norm_eps: float, dropout: float, attention_implementation: str)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden: Tensor, cached_kv: list[tuple[Tensor, Tensor]], cos: Tensor, sin: Tensor, attn_mask: Tensor, adarms_cond: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.cosmos3.qwen3vl_with_expert.Qwen3ExpertAttention(hidden_size: int, num_attention_heads: int, num_key_value_heads: int, head_dim: int, rms_norm_eps: float, attention_implementation: str)

Bases: Module

Expert self/cross attention: expert queries over [cached_backbone_KV ; expert_KV].

Matches Qwen3VLTextAttention (QK-norm on the head dim before RoPE, GQA, scaling head_dim**-0.5) but the keys/values are prefixed with the frozen backbone’s cached, already-RoPE’d KV for this layer, so the expert cross-attends to the whole observation prefix as well as to the action chunk.

__init__(hidden_size: int, num_attention_heads: int, num_key_value_heads: int, head_dim: int, rms_norm_eps: float, attention_implementation: str)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden: Tensor, cached_kv: tuple[Tensor, Tensor], cos: Tensor, sin: Tensor, attn_mask: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.cosmos3.qwen3vl_with_expert.Qwen3ExpertDecoderLayer(hidden_size: int, intermediate_size: int, num_attention_heads: int, num_key_value_heads: int, head_dim: int, adarms_cond_dim: int, rms_norm_eps: float, dropout: float, attention_implementation: str)

Bases: Module

One expert decoder layer: AdaRMS pre-norms + gated residuals around attn / MLP.

__init__(hidden_size: int, intermediate_size: int, num_attention_heads: int, num_key_value_heads: int, head_dim: int, adarms_cond_dim: int, rms_norm_eps: float, dropout: float, attention_implementation: str)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden: Tensor, cached_kv: tuple[Tensor, Tensor], cos: Tensor, sin: Tensor, attn_mask: Tensor, adarms_cond: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.cosmos3.qwen3vl_with_expert.Qwen3ExpertMLP(hidden_size: int, intermediate_size: int)

Bases: Module

Qwen3 gated-SiLU MLP (gate/up/down), matching Qwen3VLTextMLP.

__init__(hidden_size: int, intermediate_size: int)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class opentau.policies.cosmos3.qwen3vl_with_expert.Qwen3VLWithExpertModel(qwen3vl_config: Qwen3VLConfig, *, expert_hidden_size: int, expert_intermediate_size: int, expert_num_hidden_layers: int, expert_num_attention_heads: int, expert_num_key_value_heads: int, expert_head_dim: int, expert_adarms_cond_dim: int, expert_rms_norm_eps: float, dropout: float, attention_implementation: str, freeze_vision_encoder: bool = True, train_expert_only: bool = True, gradient_checkpointing: bool = False, load_pretrained_backbone_repo: str | None = None, condition_on_layer: int | None = None)

Bases: Module

A frozen Qwen3-VL backbone paired with a trainable flow-matching action expert.

__init__(qwen3vl_config: Qwen3VLConfig, *, expert_hidden_size: int, expert_intermediate_size: int, expert_num_hidden_layers: int, expert_num_attention_heads: int, expert_num_key_value_heads: int, expert_head_dim: int, expert_adarms_cond_dim: int, expert_rms_norm_eps: float, dropout: float, attention_implementation: str, freeze_vision_encoder: bool = True, train_expert_only: bool = True, gradient_checkpointing: bool = False, load_pretrained_backbone_repo: str | None = None, condition_on_layer: int | None = None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

compute_rope(position_ids: Tensor, dtype: dtype, device: device) → tuple[Tensor, Tensor]

Compute MRoPE (cos, sin) for position_ids (3, B, S) using the backbone rotary.

get_rope_index(input_ids, image_grid_thw, attention_mask)

run_expert(action_embs: Tensor, cached_kv: list[tuple[Tensor, Tensor]], cos: Tensor, sin: Tensor, attn_mask: Tensor, adarms_cond: Tensor) → Tensor

run_prefix(input_ids: Tensor, attention_mask: Tensor, position_ids: Tensor, pixel_values: Tensor | None, image_grid_thw: Tensor | None) → list[tuple[Tensor, Tensor]]

Run the backbone over the observation prefix; return its per-layer (K, V) cache.

Each entry is (key, value) of shape (B, num_kv_heads, S_prefix, head_dim). The list has self.num_layers entries – the full backbone depth in the default regime, or the truncated condition_on_layer + 1 entries in single-layer mode (the selected layer being the last). run_expert maps it onto the expert layers.

When train_expert_only (the default), the backbone is fully frozen: the forward runs under no_grad and the cached KV is .detach()’d, so the expert reads it as a constant. When train_expert_only is False (partial unfreeze, e.g. a trainable text tower with a frozen vision encoder), the forward keeps its graph and the KV is not detached, so the expert’s loss backpropagates into the (unfrozen) backbone — otherwise unfreezing would be a silent no-op.

set_requires_grad() → None

property text_model

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