opentau.policies.cosmos3.modeling_cosmos3

cosmos3: a Vision-Language-Action flow-matching policy on a frozen Qwen3-VL-32B reasoner.

cosmos3 follows the π0.5 flow-matching recipe (see policies/pi05/modeling_pi05.py) but swaps the PaliGemma backbone for a frozen Qwen3-VL-32B vision-language model – the reasoning tower of NVIDIA Cosmos3-Super (extracted to a standalone Qwen3-VL-32B checkpoint by opentau.scripts.extract_cosmos3_reasoner) – and pairs it with a custom sub-1B Qwen3-style action expert (qwen3vl_with_expert.py).

Pipeline:

• The frozen reasoner encodes the camera images + language prompt once (the prefix) via the stock Qwen3VLModel.forward, producing a per-layer key/value cache.

• The trainable expert denoises a continuous action chunk (the suffix) by flow matching, cross-attending to that cache at every layer. Proprioceptive state is projected into a single token prepended to the expert’s action chunk, so actions are conditioned on state while the backbone sees only images + language.

Continuous actions only (MSE flow matching) – there is no FAST discrete-action branch and no response/subtask head, so cosmos3 always returns a zero CE term for loss-dict compatibility with scripts/train.py.

Functions

create_sinusoidal_pos_embedding(time, ...[, ...])	Sine-cosine positional embedding for scalar positions time of shape (B, N).
make_att_2d_masks(pad_masks, att_masks[, ...])	Build a 2-D attention mask from 1-D pad + block masks (π0.5 convention).

Classes

Cosmos3FlowMatching(config[, qwen3vl_config])	Flow-matching head: frozen Qwen3-VL prefix + trainable Qwen3 action expert.
Cosmos3Policy(config[, per_dataset_stats, ...])	OpenTau wrapper around Cosmos3FlowMatching (normalization, processor, action queue).

class opentau.policies.cosmos3.modeling_cosmos3.Cosmos3FlowMatching(config: Cosmos3Config, qwen3vl_config: Qwen3VLConfig | None = None)

Bases: Module

Flow-matching head: frozen Qwen3-VL prefix + trainable Qwen3 action expert.

__init__(config: Cosmos3Config, qwen3vl_config: Qwen3VLConfig | None = None)

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

embed_suffix(noisy_actions: Tensor, timestep: Tensor, state: Tensor) → tuple[Tensor, Tensor, Tensor, Tensor]

Embed the proprioceptive-state token + the noisy action chunk for the expert.

Returns (embs, pad_masks, att_masks, adarms_cond) with sequence length chunk_size + 1 (state token prepended). The state token uses a fixed time of 1.0 for its AdaRMS conditioning; the action tokens use timestep.

forward(input_ids: Tensor, attention_mask: Tensor, pixel_values: Tensor | None, image_grid_thw: Tensor | None, state: Tensor, actions: Tensor, actions_is_pad: Tensor | None = None, noise: Tensor | None = None, time: Tensor | None = None, real_action_dim: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss]

Full flow-matching training forward; returns {“MSE”, “CE”(=0)} (+ per-sample).

sample_actions(input_ids: Tensor, attention_mask: Tensor, pixel_values: Tensor | None, image_grid_thw: Tensor | None, state: Tensor, action_prefix: Tensor, delay: Tensor, noise: Tensor | None = None) → Tensor

Euler-integrate the flow from noise to an action chunk (π0.5 sampler).

sample_noise(shape: tuple[int, ...], device: device | str) → Tensor

sample_time(bsize: int, device: device | str) → Tensor

class opentau.policies.cosmos3.modeling_cosmos3.Cosmos3Policy(config: Cosmos3Config, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None, qwen3vl_config: Qwen3VLConfig | None = None)

Bases: PreTrainedPolicy

OpenTau wrapper around Cosmos3FlowMatching (normalization, processor, action queue).

__init__(config: Cosmos3Config, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None, qwen3vl_config: Qwen3VLConfig | None = None)

Initializes the PreTrainedPolicy.

Parameters:

• config – The configuration object for the policy.

• *inputs – Variable length argument list.

• **kwargs – Arbitrary keyword arguments.

Raises:

ValueError – If config is not an instance of PreTrainedConfig.

config_class

alias of Cosmos3Config

forward(batch: dict[str, Tensor], noise: Tensor | None = None, time: Tensor | None = None, return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss]

Performs a forward pass of the policy.

Parameters:

batch – A dictionary of input tensors.

Returns:

A tuple containing:

• The loss tensor.

• An optional dictionary of metrics or auxiliary outputs. Apart from the loss, items should be logging-friendly native Python types.

Return type:

tuple[Tensor, dict | None]

get_optim_params() → list[Parameter]

Returns the policy-specific parameters dict to be passed on to the optimizer.

Returns:

A dictionary of parameters to optimize.

Return type:

dict

name: None = 'cosmos3'

The name of the policy. Must be defined in subclasses.

predict_action_chunk(batch: dict[str, Tensor]) → Tensor

prepare_multimodal_inputs(batch: dict[str, Tensor]) → dict[str, Tensor]

Build Qwen3-VL input_ids/attention_mask/pixel_values/image_grid_thw.

Resizes every present camera image to image_resize and runs the Qwen3-VL chat template + processor so the language prompt is interleaved with the image tokens.

prepare_state(batch: dict[str, Tensor]) → Tensor

Return the proprioceptive state padded to max_state_dim (zeros if absent).

reset() → None

Resets the policy state.

This method should be called whenever the environment is reset. It handles tasks like clearing caches or resetting internal states for stateful policies.

sample_actions(batch: dict[str, Tensor], action_prefix: Tensor | None = None, delay: Tensor | None = None, noise: Tensor | None = None) → Tensor

select_action(batch: dict[str, Tensor], noise: Tensor | None = None) → Tensor

Selects an action based on the input batch.

This method handles action selection during inference, including caching for stateful policies (e.g. RNNs, Transformers).

Parameters:

batch – A dictionary of observation tensors.

Returns:

The selected action(s).

Return type:

Tensor

supports_torch_compile: bool = False

Whether maybe_compile_for_training() may compile this policy’s self.model. Default False: opt-in per policy, because the in-place nn.Module.compile only takes effect if the policy’s training forward invokes the submodule via self.model(...) (__call__) rather than self.model.forward(...). Subclasses whose forward has been switched to self.model(...) (currently PI05Policy and PI07LowLevelPolicy) set this True. Leaving it False makes use_torch_compile=True a loud no-op on unwired policies instead of a silent compile-but-never-dispatch.

opentau.policies.cosmos3.modeling_cosmos3.create_sinusoidal_pos_embedding(time: Tensor, dimension: int, min_period: float, max_period: float, device: device | str = 'cpu') → Tensor

Sine-cosine positional embedding for scalar positions time of shape (B, N).

Returns (B, N, dimension). Mirrors the helper in pi05/modeling_pi05.py.

opentau.policies.cosmos3.modeling_cosmos3.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

Build a 2-D attention mask from 1-D pad + block masks (π0.5 convention).

pad_masks bool (B, N): True = real token. att_masks int (B, N): 1 opens a new causal block, 0 shares the previous token’s block. Returns (B, N, N) or, when n_cross_att_tokens is given, (B, N, n_cross + N) with full cross-attention to the (valid) prefix prepended. Mirrors pi05/modeling_pi05.py::make_att_2d_masks.