opentau.policies.cosmos3.configuration_cosmos3

Configuration for the cosmos3 policy.

cosmos3 is the π0.5 flow-matching recipe with the PaliGemma backbone swapped for a frozen Qwen3-VL-32B vision-language model – the reasoning tower of NVIDIA Cosmos3-Super, extracted into a standalone Qwen3-VL-32B checkpoint by opentau.scripts.extract_cosmos3_reasoner – and a custom sub-1B Qwen3-style action expert. Continuous actions only – there is no FAST discrete-action branch and no response/subtask head, so the discrete/response fields of PI05Config are intentionally absent here.

Expert sizing note (the param budget is dominated by attention, not the MLP): the expert key/value heads (expert_num_key_value_heads) and expert_head_dim must match the Qwen3-VL text tower (8 / 128) so the expert’s keys/values concatenate with the backbone’s cached KV at every layer. The expert query head count (expert_num_attention_heads) is free (any multiple of the KV head count) because the backbone’s prefix is run through stock transformers and only its KV cache – never its queries – is consumed by the expert. With the defaults below (hidden 1024, 16 query heads, 64 layers, intermediate 2048) the trainable expert + projections total ~0.91B parameters, comfortably under 1B.

Classes

Cosmos3Config(n_obs_steps, ...)

Configuration class for the cosmos3 policy.

class opentau.policies.cosmos3.configuration_cosmos3.Cosmos3Config(n_obs_steps: int = 1, normalization_mapping: dict[str, ~opentau.configs.types.NormalizationMode] = <factory>, input_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, output_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, device: str | None = None, use_amp: bool = False, use_torch_compile: bool = False, torch_compile_mode: str = 'default', pretrained_path: str | None = None, skip_normalization_weights: bool = False, skip_input_resolution_check: bool = False, save_normalization_stats: bool = True, dataset_names: list[str] | None = None, dataset_to_norm_index: dict[str, int] | None = None, cloud_vlm_latency_mean: float = 0.0, cloud_vlm_latency_std: float = 0.0, cloud_vlm_latency_lower: float = 0.0, cloud_vlm_latency_upper: float = 0.0, action_decoder_latency_mean: float = 0.0, action_decoder_latency_std: float = 0.0, action_decoder_latency_lower: float = 0.0, action_decoder_latency_upper: float = 0.0, chunk_size: int = 50, n_action_steps: int = 50, max_state_dim: int = 32, max_action_dim: int = 32, proj_width: int = 1024, prompt_max_length: int = 256, empty_cameras: int = 0, num_steps: int = 10, max_delay: int = 0, pretrained_backbone_repo_id: str = 'TensorAuto/cosmos3-reason-32b', load_pretrained_backbone: bool = True, image_resize: int = 224, attention_implementation: str = 'sdpa', freeze_vision_encoder: bool = True, train_expert_only: bool = True, gradient_checkpointing: bool = False, dropout: float = 0.1, condition_on_layer: int | None = None, expert_hidden_size: int = 1024, expert_intermediate_size: int = 2048, expert_num_hidden_layers: int = 64, expert_num_attention_heads: int = 16, expert_num_key_value_heads: int = 8, expert_head_dim: int = 128, expert_adarms_cond_dim: int = 256, expert_rms_norm_eps: float = 1e-06, expert_rope_theta: float = 5000000.0, optimizer_lr: float = 2.5e-05, optimizer_betas: tuple[float, float] = (0.9, 0.95), optimizer_eps: float = 1e-08, optimizer_weight_decay: float = 1e-10, scheduler_warmup_steps: int = 1000, scheduler_decay_steps: int = 30000, scheduler_decay_lr: float = 2.5e-06)

Bases: PreTrainedConfig

Configuration class for the cosmos3 policy.

Parameters:

• n_obs_steps – Number of observation steps. Only 1 is supported.

• chunk_size – Trained action-chunk length (prediction horizon). Upper bound for n_action_steps. Defaults to 50.

• n_action_steps – Inference execution horizon (<= chunk_size). Defaults to 50.

• normalization_mapping – Per-feature normalization modes. Visual identity, state/action mean-std (matches π0.5).

• max_state_dim – Padded proprioceptive-state dimension. Defaults to 32.

• max_action_dim – Padded action dimension. Defaults to 32.

• proj_width – Width of the action/time projection MLPs. Defaults to 1024.

• prompt_max_length – Maximum language-prompt token length. Defaults to 256.

• empty_cameras – Number of empty camera inputs to inject (sim adaptations).

• num_steps – Number of flow-matching Euler denoising steps. Defaults to 10.

• max_delay – Maximum number of frozen action-prefix steps (real-time inference).

• pretrained_backbone_repo_id – HF repo id (or local path) for the Qwen3-VL backbone weights – the Cosmos3-Super reasoning tower extracted into a standalone Qwen3-VL-32B checkpoint by opentau.scripts.extract_cosmos3_reasoner (run once on the ungated nvidia/Cosmos3-Super). Defaults to TensorAuto/cosmos3-reason-32b, the published extraction (private – the training environment needs an HF token with read access to the TensorAuto org; from_pretrained picks up the ambient HF_TOKEN / ~/.cache/huggingface/token). Re-run the extraction script if you need to reproduce or re-host it.

• load_pretrained_backbone – Whether to download/load the backbone weights on construction. Set False for CPU tests / tiny random configs.

• image_resize – Square side length (pixels) to resize every camera image to before the Qwen3-VL vision tower. Bounds the number of vision tokens to a fixed, deterministic count. Defaults to 224.

• attention_implementation – "eager" or "sdpa" for the expert attention. "flash_cuda" is unsupported (MRoPE/QK-norm). Defaults to "sdpa".

• freeze_vision_encoder – Freeze the Qwen3-VL vision tower. Defaults to True.

• train_expert_only – Freeze the entire backbone; train only the expert + projections. Defaults to True (cosmos3’s intended regime).

• gradient_checkpointing – Checkpoint the expert decoder layers to trade compute for activation memory. Defaults to False.

• dropout – Dropout probability inside the expert. Defaults to 0.1.

• expert_hidden_size – Action-expert hidden width. Defaults to 1024.

• expert_intermediate_size – Action-expert SwiGLU MLP width. Defaults to 2048.

• condition_on_layer – Which backbone (reasoner) layer the action expert cross-attends to. None (default) keeps the per-layer correspondence – expert layer i reads backbone layer i – which requires expert_num_hidden_layers == backbone depth. When set to an int k (0-indexed; Python-style negatives allowed, e.g. -1 = last layer), every expert layer cross-attends to backbone layer k instead. In this single-layer regime two things follow: (1) the expert depth is freed from the backbone depth (expert_num_hidden_layers may be anything >= 1, e.g. a shallower/cheaper expert), and (2) the frozen backbone is truncated to its first k + 1 layers at load time – the deeper layers are never allocated or run, a large VRAM + forward-compute saving on the 32B reasoner. The selected layer’s KV is bit-identical whether or not the backbone is truncated (deepstack vision features are injected only into the earliest layers, so layer k’s output depends only on layers 0..k). Defaults to None.

• expert_num_hidden_layers – Action-expert depth. With condition_on_layer=None this MUST equal the backbone text tower depth (64 for Qwen3-VL-32B) so each expert layer reads the matching backbone KV layer; with a single condition_on_layer selected the constraint is dropped and any depth >= 1 is allowed. Defaults to 64.

• expert_num_attention_heads – Action-expert query heads. Free (multiple of expert_num_key_value_heads). Defaults to 16.

• expert_num_key_value_heads – Action-expert KV heads. MUST equal the backbone text tower (8). Defaults to 8.

• expert_head_dim – Per-head dimension. MUST equal the backbone (128). Defaults to 128.

• expert_adarms_cond_dim – Width of the AdaRMS (time) conditioning vector. Defaults to 256.

• expert_rms_norm_eps – RMSNorm epsilon for the expert. Defaults to 1e-6.

• expert_rope_theta – RoPE base for the expert (matches the backbone, 5e6). Note: when the backbone’s rotary embedding is reused for the shared cos/sin this is informational; kept for standalone expert rotaries.

• optimizer_weight_decay (optimizer_lr / optimizer_betas / optimizer_eps /) – AdamW preset (π0.5 values).

• scheduler_decay_lr (scheduler_warmup_steps / scheduler_decay_steps /) – Cosine-decay-with-warmup preset (π0.5 values).

• use_torch_compile – Whether to torch.compile the model. Defaults to False (enable only after verifying bit-identical seeded runs).

__init__(n_obs_steps: int = 1, normalization_mapping: dict[str, ~opentau.configs.types.NormalizationMode] = <factory>, input_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, output_features: dict[str, ~opentau.configs.types.PolicyFeature] = <factory>, device: str | None = None, use_amp: bool = False, use_torch_compile: bool = False, torch_compile_mode: str = 'default', pretrained_path: str | None = None, skip_normalization_weights: bool = False, skip_input_resolution_check: bool = False, save_normalization_stats: bool = True, dataset_names: list[str] | None = None, dataset_to_norm_index: dict[str, int] | None = None, cloud_vlm_latency_mean: float = 0.0, cloud_vlm_latency_std: float = 0.0, cloud_vlm_latency_lower: float = 0.0, cloud_vlm_latency_upper: float = 0.0, action_decoder_latency_mean: float = 0.0, action_decoder_latency_std: float = 0.0, action_decoder_latency_lower: float = 0.0, action_decoder_latency_upper: float = 0.0, chunk_size: int = 50, n_action_steps: int = 50, max_state_dim: int = 32, max_action_dim: int = 32, proj_width: int = 1024, prompt_max_length: int = 256, empty_cameras: int = 0, num_steps: int = 10, max_delay: int = 0, pretrained_backbone_repo_id: str = 'TensorAuto/cosmos3-reason-32b', load_pretrained_backbone: bool = True, image_resize: int = 224, attention_implementation: str = 'sdpa', freeze_vision_encoder: bool = True, train_expert_only: bool = True, gradient_checkpointing: bool = False, dropout: float = 0.1, condition_on_layer: int | None = None, expert_hidden_size: int = 1024, expert_intermediate_size: int = 2048, expert_num_hidden_layers: int = 64, expert_num_attention_heads: int = 16, expert_num_key_value_heads: int = 8, expert_head_dim: int = 128, expert_adarms_cond_dim: int = 256, expert_rms_norm_eps: float = 1e-06, expert_rope_theta: float = 5000000.0, optimizer_lr: float = 2.5e-05, optimizer_betas: tuple[float, float] = (0.9, 0.95), optimizer_eps: float = 1e-08, optimizer_weight_decay: float = 1e-10, scheduler_warmup_steps: int = 1000, scheduler_decay_steps: int = 30000, scheduler_decay_lr: float = 2.5e-06) → None

property action_delta_indices: list[int]

Get indices for action delta features.

Returns:

List of indices indicating which action features should be treated as deltas, or None if no delta features are used.

attention_implementation: str = 'sdpa'

chunk_size: int = 50

condition_on_layer: int | None = None

dropout: float = 0.1

empty_cameras: int = 0

expert_adarms_cond_dim: int = 256

expert_head_dim: int = 128

expert_hidden_size: int = 1024

expert_intermediate_size: int = 2048

expert_num_attention_heads: int = 16

expert_num_hidden_layers: int = 64

expert_num_key_value_heads: int = 8

expert_rms_norm_eps: float = 1e-06

expert_rope_theta: float = 5000000.0

freeze_vision_encoder: bool = True

get_optimizer_preset() → AdamWConfig

Return the default AdamW optimizer configuration.

get_scheduler_preset() → LRSchedulerConfig

Return the default cosine-decay-with-warmup scheduler configuration.

gradient_checkpointing: bool = False

image_resize: int = 224

load_pretrained_backbone: bool = True

max_action_dim: int = 32

max_delay: int = 0

max_state_dim: int = 32

n_action_steps: int = 50

n_obs_steps: int = 1

normalization_mapping: dict[str, NormalizationMode]

num_steps: int = 10

property observation_delta_indices: None

Get indices for observation delta features.

Returns:

List of indices indicating which observation features should be treated as deltas, or None if no delta features are used.

optimizer_betas: tuple[float, float] = (0.9, 0.95)

optimizer_eps: float = 1e-08

optimizer_lr: float = 2.5e-05

optimizer_weight_decay: float = 1e-10

pretrained_backbone_repo_id: str = 'TensorAuto/cosmos3-reason-32b'

proj_width: int = 1024

prompt_max_length: int = 256

property reward_delta_indices: None

Get indices for reward delta features.

Returns:

List of indices indicating which reward features should be treated as deltas, or None if no delta features are used.

scheduler_decay_lr: float = 2.5e-06

scheduler_decay_steps: int = 30000

scheduler_warmup_steps: int = 1000

train_expert_only: bool = True

use_torch_compile: bool = False

validate_features() → None

Add empty cameras to input_features if configured.