config.json

{
  "type": "cosmos3",
  "n_obs_steps": 1,
  "input_features": {
    "observation.images.wrist": {
      "type": "VISUAL",
      "shape": [
        3,
        360,
        640
      ]
    },
    "observation.images.top": {
      "type": "VISUAL",
      "shape": [
        3,
        360,
        640
      ]
    },
    "observation.state": {
      "type": "STATE",
      "shape": [
        6
      ]
    }
  },
  "output_features": {
    "action": {
      "type": "ACTION",
      "shape": [
        6
      ]
    }
  },
  "device": "cpu",
  "use_amp": false,
  "use_peft": false,
  "push_to_hub": true,
  "repo_id": null,
  "private": null,
  "tags": null,
  "license": null,
  "pretrained_path": null,
  "pretrained_revision": null,
  "text_processor_name_or_path": "geonmin-kim/Cosmos3-Edge-Policy-SO101-init",
  "transformer_config": {
    "action_dim": 64,
    "action_gen": true,
    "attention_bias": false,
    "attention_dropout": 0.0,
    "base_fps": 24,
    "enable_fps_modulation": true,
    "head_dim": 128,
    "hidden_act": "relu2",
    "hidden_size": 2048,
    "intermediate_size": 9216,
    "latent_channel": 48,
    "latent_patch_size": 2,
    "num_attention_heads": 16,
    "num_embodiment_domains": 32,
    "num_hidden_layers": 28,
    "num_key_value_heads": 8,
    "patch_latent_dim": 192,
    "qk_norm_for_text": false,
    "rms_norm_eps": 1e-05,
    "rope_axes_dim": [
      24,
      20,
      20
    ],
    "rope_scaling": {
      "mrope_section": [
        24,
        20,
        20
      ]
    },
    "rope_theta": 100000000.0,
    "sound_dim": null,
    "sound_gen": false,
    "sound_latent_fps": 25.0,
    "timestep_scale": 0.001,
    "unified_3d_mrope_reset_spatial_ids": true,
    "unified_3d_mrope_temporal_modality_margin": 15000,
    "use_und_k_norm_for_gen": true,
    "vocab_size": 131072
  },
  "vae_config": {
    "attn_scales": [],
    "base_dim": 160,
    "decoder_base_dim": 256,
    "dim_mult": [
      1,
      2,
      4,
      4
    ],
    "dropout": 0.0,
    "in_channels": 12,
    "is_residual": true,
    "latents_mean": [
      -0.2289,
      -0.0052,
      -0.1323,
      -0.2339,
      -0.2799,
      0.0174,
      0.1838,
      0.1557,
      -0.1382,
      0.0542,
      0.2813,
      0.0891,
      0.157,
      -0.0098,
      0.0375,
      -0.1825,
      -0.2246,
      -0.1207,
      -0.0698,
      0.5109,
      0.2665,
      -0.2108,
      -0.2158,
      0.2502,
      -0.2055,
      -0.0322,
      0.1109,
      0.1567,
      -0.0729,
      0.0899,
      -0.2799,
      -0.123,
      -0.0313,
      -0.1649,
      0.0117,
      0.0723,
      -0.2839,
      -0.2083,
      -0.052,
      0.3748,
      0.0152,
      0.1957,
      0.1433,
      -0.2944,
      0.3573,
      -0.0548,
      -0.1681,
      -0.0667
    ],
    "latents_std": [
      0.4765,
      1.0364,
      0.4514,
      1.1677,
      0.5313,
      0.499,
      0.4818,
      0.5013,
      0.8158,
      1.0344,
      0.5894,
      1.0901,
      0.6885,
      0.6165,
      0.8454,
      0.4978,
      0.5759,
      0.3523,
      0.7135,
      0.6804,
      0.5833,
      1.4146,
      0.8986,
      0.5659,
      0.7069,
      0.5338,
      0.4889,
      0.4917,
      0.4069,
      0.4999,
      0.6866,
      0.4093,
      0.5709,
      0.6065,
      0.6415,
      0.4944,
      0.5726,
      1.2042,
      0.5458,
      1.6887,
      0.3971,
      1.06,
      0.3943,
      0.5537,
      0.5444,
      0.4089,
      0.7468,
      0.7744
    ],
    "num_res_blocks": 2,
    "out_channels": 12,
    "patch_size": 2,
    "scale_factor_spatial": 16,
    "scale_factor_temporal": 4,
    "temperal_downsample": [
      false,
      true,
      true
    ],
    "z_dim": 48
  },
  "scheduler_config": {
    "_class_name": "UniPCMultistepScheduler",
    "_diffusers_version": "0.40.0.dev0",
    "beta_end": 0.02,
    "beta_schedule": "linear",
    "beta_start": 0.0001,
    "disable_corrector": [],
    "dynamic_thresholding_ratio": 0.995,
    "final_sigmas_type": "zero",
    "flow_shift": 5.0,
    "lower_order_final": true,
    "num_train_timesteps": 1000,
    "predict_x0": true,
    "prediction_type": "flow_prediction",
    "rescale_betas_zero_snr": false,
    "sample_max_value": 1.0,
    "shift_terminal": null,
    "sigma_max": null,
    "sigma_min": null,
    "solver_order": 2,
    "solver_p": null,
    "solver_type": "bh2",
    "steps_offset": 0,
    "thresholding": false,
    "time_shift_type": "exponential",
    "timestep_spacing": "linspace",
    "trained_betas": null,
    "use_beta_sigmas": false,
    "use_dynamic_shifting": false,
    "use_exponential_sigmas": false,
    "use_flow_sigmas": true,
    "use_karras_sigmas": false
  },
  "freeze_vae": true,
  "dtype": "bfloat16",
  "local_files_only": false,
  "chunk_size": 32,
  "n_action_steps": 32,
  "raw_action_dim": 6,
  "max_action_dim": 64,
  "max_state_dim": 6,
  "joint_position_dim": 5,
  "gripper_position_dim": 1,
  "use_state": true,
  "history_length": 1,
  "action_space": "joint_pos",
  "invert_gripper": false,
  "domain_name": "droid_lerobot",
  "domain_id": 8,
  "eos_token_id": 11,
  "start_of_generation_token_id": 20,
  "mode": "policy",
  "viewpoint": "concat_view",
  "additional_view_description": "The top row is from the wrist-mounted camera. The bottom row contains a third-person top-down view of the scene with the robot visible, and a blank padded view.",
  "conditioning_fps": 30.0,
  "resolution_tier": 480,
  "guidance_scale": 3.0,
  "num_inference_steps": 4,
  "shift": 5.0,
  "seed": 0,
  "deterministic_seed": false,
  "generate_video": false,
  "output_type": "latent",
  "train_time_video_distribution": "waver",
  "video_loss_weight": 10.0,
  "action_loss_weight": 10.0,
  "normalize_loss_by_active": false,
  "image_keys": [
    "observation.images.wrist",
    "observation.images.top"
  ],
  "num_views": 3,
  "image_height": 360,
  "image_width": 640,
  "composed_image_height": 540,
  "composed_image_width": 640,
  "prompt_key": "task",
  "normalization_mapping": {
    "VISUAL": "IDENTITY",
    "STATE": "QUANTILES",
    "ACTION": "QUANTILES"
  },
  "optimizer_lr": 1e-05,
  "optimizer_betas": [
    0.9,
    0.95
  ],
  "optimizer_eps": 1e-08,
  "optimizer_weight_decay": 0.01,
  "optimizer_grad_clip_norm": 1.0,
  "scheduler_warmup_steps": 1000,
  "scheduler_decay_steps": 30000,
  "scheduler_decay_lr": 0.0
}