opentau.policies.cosmos3_nano.configuration_cosmos3_nano
Configuration for the cosmos3_nano policy.
cosmos3_nano is the cosmos3 recipe (see configuration_cosmos3.py) with the frozen
Qwen3-VL-32B backbone (the Cosmos3-Super reasoning tower) swapped for a frozen
Qwen3-VL-8B — the reasoning tower of NVIDIA Cosmos3-Nano, extracted into a
standalone checkpoint by opentau.scripts.extract_cosmos3_reasoner and published as
TensorAuto/cosmos3-reason-8b (private — like the 32B extraction, 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).
The Nano tower is byte-for-byte the Qwen3-VL-8B geometry: hidden 4096, 36 layers,
32 attention heads, and — decisive for the action expert’s cross-attention — the same
8 key/value heads x head_dim 128 as the 32B tower. The expert’s hard KV-concat
constraints therefore carry over unchanged, and the only default that must move is
expert_num_hidden_layers (36, matching the backbone depth for the default per-layer
correspondence). With the inherited expert widths (hidden 1024, 16 query heads,
intermediate 2048) the trainable expert + projections total ~0.51B parameters.
condition_on_layer works exactly as in cosmos3 (valid range [-36, 35]); every other
field, constraint, and preset is inherited from Cosmos3Config.
Classes
|
Configuration class for the cosmos3_nano policy. |
- class opentau.policies.cosmos3_nano.configuration_cosmos3_nano.Cosmos3NanoConfig(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-8b', 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 = 36, 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)[source]
Bases:
Cosmos3ConfigConfiguration class for the cosmos3_nano policy.
Identical to
Cosmos3Configexcept for the two nano defaults below; see the cosmos3 docstrings for the meaning of every field and the hard expert/backbone geometry constraints (re-validated against the loaded backbone at model-build time).- Parameters:
pretrained_backbone_repo_id – HF repo id (or local path) for the Qwen3-VL-8B backbone weights — the Cosmos3-Nano reasoning tower extracted by
opentau.scripts.extract_cosmos3_reasoner(run once on the ungatednvidia/Cosmos3-Nano). Defaults toTensorAuto/cosmos3-reason-8b, the published extraction (private — needs an HF token with TensorAuto org read access). Re-run the extraction script to reproduce or re-host it.expert_num_hidden_layers – Action-expert depth. With
condition_on_layer=Nonethis MUST equal the backbone text tower depth (36 for Qwen3-VL-8B); with a singlecondition_on_layerselected, any depth >= 1 is allowed. Defaults to 36.
- __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-8b', 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 = 36, 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
- pretrained_backbone_repo_id: str = 'TensorAuto/cosmos3-reason-8b'