Add cosmos3edge_libero_spatial_full_finetune_bs16 checkpoint

.gitattributes

@@ -49,3 +49,4 @@ gr00t_qwen3vl_0.6b_libero_spatial_full_finetune_bs64/tokenizer/tokenizer.json fi
 pi05_paligemma_robocasa_full_data_full_finetune_bs256/pi05_paligemma_robocasa_full_data_full_finetune_2026_07_21_01_22_59.jsonl filter=lfs diff=lfs merge=lfs -text
 pi05_paligemma_robocasa_full_data_full_finetune_bs256/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 fastwam_base/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+cosmos3edge_libero_spatial_full_finetune_bs16/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text

cosmos3edge_libero_spatial_full_finetune_bs16/checkpoints/step-002000-epoch-76-loss=0.7025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8765854419dd33e61b70ea40795390581835dbd0cfb4fd54eaca03eb5569e351
+size 16297490208

cosmos3edge_libero_spatial_full_finetune_bs16/config.json

@@ -0,0 +1,669 @@
+{
+  "eval": {
+    "dataset": {
+      "extra_tensor_keys": [
+        "conditioning_fps",
+        "prepend_state_to_action"
+      ],
+      "img_buffer_len": 1,
+      "transforms": [
+        {
+          "embodiment_id": 5,
+          "img_keys": [
+            "agentview_image",
+            "robot0_eye_in_hand_image"
+          ],
+          "type": "ProcessLiberoEvalInputs"
+        },
+        {
+          "conditioning_fps": 20.0,
+          "prepend_state_to_action": false,
+          "type": "SetCosmos3ActionMetadata"
+        },
+        {
+          "image_resize_strategy": "resize-naive",
+          "input_sizes": [
+            [
+              3,
+              256,
+              256
+            ],
+            [
+              3,
+              256,
+              256
+            ]
+          ],
+          "means": [
+            [
+              127.5,
+              127.5,
+              127.5
+            ],
+            [
+              127.5,
+              127.5,
+              127.5
+            ]
+          ],
+          "stds": [
+            [
+              127.5,
+              127.5,
+              127.5
+            ],
+            [
+              127.5,
+              127.5,
+              127.5
+            ]
+          ],
+          "type": "TransformImage"
+        },
+        {
+          "action_metadata": {
+            "append_viewpoint": false,
+            "conditioning_fps": 20.0,
+            "frame_window_size": 17,
+            "video_height": 256,
+            "video_width": 512
+          },
+          "cfg_dropout_rate": 0.0,
+          "max_len": 512,
+          "output_attention_mask_key": "lang_masks",
+          "output_key": "lang_tokens",
+          "tokenizer": {
+            "model_max_length": 4096,
+            "model_path": "./checkpoints/Cosmos3-Edge",
+            "padding_side": "right",
+            "trust_remote_code": true,
+            "type": "PretrainedTokenizer"
+          },
+          "type": "ProcessCosmos3Prompt"
+        },
+        {
+          "gripper_key": "robot0_gripper_qpos",
+          "norm_type": "mean_std",
+          "out_key": "states",
+          "pos_key": "robot0_eef_pos",
+          "quat_key": "robot0_eef_quat",
+          "state_dim": 64,
+          "type": "LiberoProprioFromInputs"
+        },
+        {
+          "frame_window_size": 1,
+          "num_views": 2,
+          "tile_direction": "horizontal",
+          "type": "PrepareVideo"
+        }
+      ],
+      "type": "LiberoParquetEvalDataset"
+    },
+    "denormalize_action": {
+      "action_dim": 10,
+      "norm_type": "quantile_rot",
+      "type": "DenormalizeLiberoFramewiseRot6DAction"
+    },
+    "enable_mixed_precision_training": true,
+    "eval_chunk_size": 16,
+    "inference_seed": 7,
+    "mixed_precision_dtype": "bf16",
+    "model_family": "cosmos3",
+    "norm_stats_key": "libero_spatial_no_noops",
+    "num_inference_steps": 30,
+    "num_steps_wait": 10,
+    "num_trials_per_task": 50,
+    "resize_size": 256,
+    "seed": 7,
+    "task_ids": null,
+    "task_suite_name": "libero_spatial",
+    "type": "LiberoEvalRunner"
+  },
+  "inference_model": {
+    "action_horizon": 16,
+    "action_in_proj": {
+      "input_size": 64,
+      "num_domains": 32,
+      "output_size": 2048,
+      "type": "DomainAwareLinear"
+    },
+    "action_out_proj": {
+      "input_size": 2048,
+      "num_domains": 32,
+      "output_size": 64,
+      "type": "DomainAwareLinear"
+    },
+    "base_fps": 24.0,
+    "enable_fps_modulation": true,
+    "enable_vision_loss": true,
+    "freeze_non_moe_vlm_backbone": true,
+    "freeze_vlm_backbone": false,
+    "latent_patch_size": 2,
+    "max_action_dim": 64,
+    "name_mapping": {
+      ".self_attn.k_norm_moe_gen.": ".self_attn.norm_added_k.",
+      ".self_attn.k_proj.": ".self_attn.to_k.",
+      ".self_attn.k_proj_moe_gen.": ".self_attn.add_k_proj.",
+      ".self_attn.o_proj.": ".self_attn.to_out.",
+      ".self_attn.o_proj_moe_gen.": ".self_attn.to_add_out.",
+      ".self_attn.q_norm_moe_gen.": ".self_attn.norm_added_q.",
+      ".self_attn.q_proj.": ".self_attn.to_q.",
+      ".self_attn.q_proj_moe_gen.": ".self_attn.add_q_proj.",
+      ".self_attn.v_proj.": ".self_attn.to_v.",
+      ".self_attn.v_proj_moe_gen.": ".self_attn.add_v_proj.",
+      "action_in_proj.": "action_proj_in.",
+      "action_modality_embed": "action_modality_embed",
+      "action_out_proj.": "action_proj_out.",
+      "time_embedder.mlp.0.": "time_embedder.linear_1.",
+      "time_embedder.mlp.2.": "time_embedder.linear_2.",
+      "vision_in_proj.projector.": "proj_in.",
+      "vision_out_proj.projector.": "proj_out.",
+      "vlm_backbone.lm_head.weight": "lm_head.weight",
+      "vlm_backbone.model.language_model.embed_tokens.weight": "embed_tokens.weight",
+      "vlm_backbone.model.language_model.layers.": "layers.",
+      "vlm_backbone.model.language_model.norm.weight": "norm.weight",
+      "vlm_backbone.model.language_model.norm_moe_gen.weight": "norm_moe_gen.weight"
+    },
+    "num_embodiment_domains": 32,
+    "ori_action_dim": 10,
+    "packed_attention_backend": "flash2",
+    "position_embedding_type": "unified_3d_mrope",
+    "pretrained_name_or_path": "./checkpoints/Cosmos3-Edge/transformer",
+    "rectified_flow_inference_config": {
+      "num_steps": 30,
+      "num_train_timesteps": 1000,
+      "scheduler_type": "unipc",
+      "shift": 10.0,
+      "use_dynamic_shifting": false,
+      "use_karras_sigmas": false
+    },
+    "rectified_flow_training_config": {
+      "action_loss_weight": 10.0,
+      "independent_action_schedule": false,
+      "normalize_loss_by_active": false,
+      "shift": {
+        "256": 3,
+        "480": 5,
+        "720": 10
+      },
+      "shift_action": null,
+      "train_time_action_distribution": "logitnormal",
+      "train_time_image_distribution": "logitnormal",
+      "train_time_video_distribution": "waver",
+      "train_time_weight": "uniform",
+      "use_discrete_rf": false,
+      "use_dynamic_shift": false,
+      "use_high_sigma_strategy": false,
+      "use_high_sigma_strategy_action": false,
+      "vision_loss_weight": 10.0
+    },
+    "reinitialize_action_policy": true,
+    "special_tokens": {
+      "end_of_generation": 21,
+      "eos_token_id": 11,
+      "start_of_generation": 20
+    },
+    "strict_mapping": true,
+    "timestep_scale": 0.001,
+    "type": "Cosmos3FlowMatching",
+    "unified_3d_mrope_reset_spatial_ids": true,
+    "unified_3d_mrope_temporal_modality_margin": 15000,
+    "vision_in_proj": {
+      "in_dim": 192,
+      "out_dim": 2048,
+      "type": "LinearProjector"
+    },
+    "vision_latent_dim": 48,
+    "vision_out_proj": {
+      "in_dim": 2048,
+      "out_dim": 192,
+      "type": "LinearProjector"
+    },
+    "vision_vae": {
+      "encode_exact_durations": [
+        17
+      ],
+      "pretrained_name_or_path": "./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth",
+      "type": "Cosmos3Wan22VAE"
+    },
+    "vlm_backbone": {
+      "include_visual": false,
+      "skip_init_weights": true,
+      "type": "Cosmos3MoTBackbone",
+      "vlm_config": {
+        "attention_bias": false,
+        "bos_token_id": 1,
+        "enable_mrope": true,
+        "eos_token_id": 11,
+        "head_dim": 128,
+        "hidden_size": 2048,
+        "intermediate_size": 9216,
+        "layer_norm_epsilon": 1e-05,
+        "max_position_embeddings": 131072,
+        "mlp_bias": false,
+        "mlp_hidden_act": "relu2",
+        "model_type": "nemotron_3_dense_vl_text",
+        "mrope_section": [
+          24,
+          20,
+          20
+        ],
+        "num_attention_heads": 16,
+        "num_hidden_layers": 28,
+        "num_key_value_heads": 8,
+        "pad_token_id": 11,
+        "rope_theta": 100000000.0,
+        "tie_word_embeddings": false,
+        "use_und_k_norm_for_gen": true,
+        "vocab_size": 131072
+      }
+    }
+  },
+  "model": {
+    "action_horizon": 16,
+    "action_in_proj": {
+      "input_size": 64,
+      "num_domains": 32,
+      "output_size": 2048,
+      "type": "DomainAwareLinear"
+    },
+    "action_out_proj": {
+      "input_size": 2048,
+      "num_domains": 32,
+      "output_size": 64,
+      "type": "DomainAwareLinear"
+    },
+    "base_fps": 24.0,
+    "enable_fps_modulation": true,
+    "enable_vision_loss": true,
+    "freeze_non_moe_vlm_backbone": true,
+    "freeze_vlm_backbone": false,
+    "latent_patch_size": 2,
+    "max_action_dim": 64,
+    "name_mapping": {
+      ".self_attn.k_norm_moe_gen.": ".self_attn.norm_added_k.",
+      ".self_attn.k_proj.": ".self_attn.to_k.",
+      ".self_attn.k_proj_moe_gen.": ".self_attn.add_k_proj.",
+      ".self_attn.o_proj.": ".self_attn.to_out.",
+      ".self_attn.o_proj_moe_gen.": ".self_attn.to_add_out.",
+      ".self_attn.q_norm_moe_gen.": ".self_attn.norm_added_q.",
+      ".self_attn.q_proj.": ".self_attn.to_q.",
+      ".self_attn.q_proj_moe_gen.": ".self_attn.add_q_proj.",
+      ".self_attn.v_proj.": ".self_attn.to_v.",
+      ".self_attn.v_proj_moe_gen.": ".self_attn.add_v_proj.",
+      "action_in_proj.": "action_proj_in.",
+      "action_modality_embed": "action_modality_embed",
+      "action_out_proj.": "action_proj_out.",
+      "time_embedder.mlp.0.": "time_embedder.linear_1.",
+      "time_embedder.mlp.2.": "time_embedder.linear_2.",
+      "vision_in_proj.projector.": "proj_in.",
+      "vision_out_proj.projector.": "proj_out.",
+      "vlm_backbone.lm_head.weight": "lm_head.weight",
+      "vlm_backbone.model.language_model.embed_tokens.weight": "embed_tokens.weight",
+      "vlm_backbone.model.language_model.layers.": "layers.",
+      "vlm_backbone.model.language_model.norm.weight": "norm.weight",
+      "vlm_backbone.model.language_model.norm_moe_gen.weight": "norm_moe_gen.weight"
+    },
+    "num_embodiment_domains": 32,
+    "ori_action_dim": 10,
+    "packed_attention_backend": "flash2",
+    "position_embedding_type": "unified_3d_mrope",
+    "pretrained_name_or_path": "./checkpoints/Cosmos3-Edge/transformer",
+    "rectified_flow_inference_config": {
+      "num_steps": 30,
+      "num_train_timesteps": 1000,
+      "scheduler_type": "unipc",
+      "shift": 10.0,
+      "use_dynamic_shifting": false,
+      "use_karras_sigmas": false
+    },
+    "rectified_flow_training_config": {
+      "action_loss_weight": 10.0,
+      "independent_action_schedule": false,
+      "normalize_loss_by_active": false,
+      "shift": {
+        "256": 3,
+        "480": 5,
+        "720": 10
+      },
+      "shift_action": null,
+      "train_time_action_distribution": "logitnormal",
+      "train_time_image_distribution": "logitnormal",
+      "train_time_video_distribution": "waver",
+      "train_time_weight": "uniform",
+      "use_discrete_rf": false,
+      "use_dynamic_shift": false,
+      "use_high_sigma_strategy": false,
+      "use_high_sigma_strategy_action": false,
+      "vision_loss_weight": 10.0
+    },
+    "reinitialize_action_policy": true,
+    "special_tokens": {
+      "end_of_generation": 21,
+      "eos_token_id": 11,
+      "start_of_generation": 20
+    },
+    "strict_mapping": true,
+    "timestep_scale": 0.001,
+    "type": "Cosmos3FlowMatching",
+    "unified_3d_mrope_reset_spatial_ids": true,
+    "unified_3d_mrope_temporal_modality_margin": 15000,
+    "vision_in_proj": {
+      "in_dim": 192,
+      "out_dim": 2048,
+      "type": "LinearProjector"
+    },
+    "vision_latent_dim": 48,
+    "vision_out_proj": {
+      "in_dim": 2048,
+      "out_dim": 192,
+      "type": "LinearProjector"
+    },
+    "vision_vae": {
+      "encode_exact_durations": [
+        17
+      ],
+      "pretrained_name_or_path": "./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth",
+      "type": "Cosmos3Wan22VAE"
+    },
+    "vlm_backbone": {
+      "include_visual": false,
+      "skip_init_weights": true,
+      "type": "Cosmos3MoTBackbone",
+      "vlm_config": {
+        "attention_bias": false,
+        "bos_token_id": 1,
+        "enable_mrope": true,
+        "eos_token_id": 11,
+        "head_dim": 128,
+        "hidden_size": 2048,
+        "intermediate_size": 9216,
+        "layer_norm_epsilon": 1e-05,
+        "max_position_embeddings": 131072,
+        "mlp_bias": false,
+        "mlp_hidden_act": "relu2",
+        "model_type": "nemotron_3_dense_vl_text",
+        "mrope_section": [
+          24,
+          20,
+          20
+        ],
+        "num_attention_heads": 16,
+        "num_hidden_layers": 28,
+        "num_key_value_heads": 8,
+        "pad_token_id": 11,
+        "rope_theta": 100000000.0,
+        "tie_word_embeddings": false,
+        "use_und_k_norm_for_gen": true,
+        "vocab_size": 131072
+      }
+    }
+  },
+  "runner": {
+    "change_key_name": false,
+    "collator": {
+      "list_keys": [
+        "sequence_plan"
+      ],
+      "meta_keys": [
+        "task_description",
+        "stats",
+        "info",
+        "timestamp"
+      ],
+      "pad_id": 11,
+      "sequence_keys": [
+        "text_token_ids"
+      ],
+      "tensor_keys": [
+        "images",
+        "actions",
+        "embodiment_ids",
+        "raw_action_dim",
+        "conditioning_fps",
+        "action_fps"
+      ],
+      "type": "Cosmos3Collator"
+    },
+    "enable_gradient_checkpointing": true,
+    "enable_mixed_precision_training": true,
+    "grad_accumulation_steps": 16,
+    "lr_scheduler": {
+      "cycle_length": 16000,
+      "type": "linear-warmup+linear-decay",
+      "warmup_steps": 500
+    },
+    "max_epochs": null,
+    "max_grad_norm": 1.0,
+    "max_keep_ckpts": 2,
+    "max_steps": 2000,
+    "metric": {
+      "active_trackers": [
+        "jsonl",
+        "wandb"
+      ],
+      "grad_accumulation_steps": 16,
+      "run_dir": "work_dirs",
+      "type": "VLAMetric",
+      "window_size": 1
+    },
+    "mixed_precision_dtype": "bf16",
+    "optimizer": {
+      "betas": [
+        0.9,
+        0.99
+      ],
+      "eps": 1e-08,
+      "exclude_1d_from_weight_decay": false,
+      "fused": true,
+      "lr": 5e-05,
+      "paramwise_learning_rate": {
+        "action_in_proj.": 0.00025,
+        "action_modality_embed": 0.00025,
+        "action_out_proj.": 0.00025
+      },
+      "type": "AdamW",
+      "weight_decay": 0.05
+    },
+    "sampler": null,
+    "save_iter_interval": 500,
+    "sharding_strategy": "full-shard",
+    "tokenizer": {
+      "model_max_length": 4096,
+      "model_path": "./checkpoints/Cosmos3-Edge",
+      "padding_side": "right",
+      "trust_remote_code": true,
+      "type": "PretrainedTokenizer"
+    },
+    "type": "FSDPTrainRunner"
+  },
+  "seed": 7,
+  "train_dataloader": {
+    "dataset": {
+      "datasets": {
+        "action_key": "action",
+        "action_window_size": 16,
+        "data_root_path": "/mnt/data/cpfs/mnt/data/yanis/FastWAM/data/libero_mujoco3.3.2/libero_spatial_no_noops_lerobot",
+        "frame_window_size": 17,
+        "require_full_window": true,
+        "statistic_name": "libero_spatial_no_noops",
+        "transforms": [
+          {
+            "embodiment_id": 5,
+            "name_mappings": {
+              "actions": [
+                "actions"
+              ],
+              "observation.state": [
+                "states"
+              ]
+            },
+            "parquet_keys": [
+              "observation.state",
+              "timestamp",
+              "actions",
+              "info",
+              "stats"
+            ],
+            "type": "ProcessParquetInputs",
+            "video_keys": [
+              "observation.images.image",
+              "observation.images.wrist_image"
+            ]
+          },
+          {
+            "type": "LiberoFramewiseActionToRot6D"
+          },
+          {
+            "action_metadata": {
+              "append_viewpoint": false,
+              "conditioning_fps": 20.0,
+              "frame_window_size": 17,
+              "video_height": 256,
+              "video_width": 512
+            },
+            "cfg_dropout_rate": 0.1,
+            "max_len": 512,
+            "tokenizer": {
+              "model_max_length": 4096,
+              "model_path": "./checkpoints/Cosmos3-Edge",
+              "padding_side": "right",
+              "trust_remote_code": true,
+              "type": "PretrainedTokenizer"
+            },
+            "type": "ProcessCosmos3Prompt"
+          },
+          {
+            "height": 256,
+            "type": "ResizeImages",
+            "width": 256
+          },
+          {
+            "type": "SimpleNormalizeImages"
+          },
+          {
+            "action_dim": 64,
+            "action_key": "action",
+            "action_norm_type": "quantile",
+            "state_dim": 64,
+            "state_key": "proprio",
+            "state_norm_type": "none",
+            "type": "NormalizeStatesAndActions"
+          },
+          {
+            "conditioning_fps": 20.0,
+            "frame_window_size": 17,
+            "mode": "wam",
+            "prepend_state_to_action": false,
+            "raw_action_dim": 10,
+            "type": "BuildCosmos3Sequence"
+          },
+          {
+            "frame_window_size": 17,
+            "num_views": 2,
+            "tile_direction": "horizontal",
+            "type": "PrepareVideo"
+          }
+        ],
+        "type": "ParquetDataset",
+        "use_delta": false,
+        "window_start_idx": 0
+      },
+      "name_mappings": {
+        "action": [
+          "action"
+        ],
+        "observation.state": [
+          "proprio"
+        ]
+      },
+      "statistic_keys": [
+        "observation.state",
+        "timestamp",
+        "action"
+      ],
+      "statistic_name": "libero_spatial_no_noops",
+      "statistics_overrides": {
+        "libero_spatial_no_noops": {
+          "action": {
+            "max": [
+              0.9375,
+              0.9375,
+              0.9375,
+              1.0,
+              0.368853,
+              0.341214,
+              0.356395,
+              1.0,
+              0.348251,
+              1.0
+            ],
+            "mean": [
+              0.050704,
+              0.097407,
+              -0.094833,
+              0.994873,
+              -0.004579,
+              -0.004288,
+              0.004389,
+              0.996104,
+              0.001109,
+              0.476725
+            ],
+            "min": [
+              -0.9375,
+              -0.9375,
+              -0.9375,
+              0.902028,
+              -0.356085,
+              -0.367416,
+              -0.370434,
+              0.921907,
+              -0.255,
+              0.0
+            ],
+            "q01": [
+              -0.723214,
+              -0.808929,
+              -0.9375,
+              0.934955,
+              -0.223431,
+              -0.189878,
+              -0.334735,
+              0.938516,
+              -0.107736,
+              0.0
+            ],
+            "q99": [
+              0.9375,
+              0.870536,
+              0.9375,
+              1.0,
+              0.331,
+              0.163153,
+              0.226216,
+              1.0,
+              0.127158,
+              1.0
+            ],
+            "std": [
+              0.333621,
+              0.387175,
+              0.45714,
+              0.010807,
+              0.077802,
+              0.063386,
+              0.078571,
+              0.009994,
+              0.038504,
+              0.49946
+            ]
+          }
+        }
+      },
+      "type": "DistributedRepeatingDataset"
+    },
+    "per_device_batch_size": 16,
+    "per_device_num_workers": 4
+  }
+}

cosmos3edge_libero_spatial_full_finetune_bs16/config.yaml

@@ -0,0 +1,549 @@
+eval:
+  dataset:
+    extra_tensor_keys:
+    - conditioning_fps
+    - prepend_state_to_action
+    img_buffer_len: 1
+    transforms:
+    - embodiment_id: 5
+      img_keys:
+      - agentview_image
+      - robot0_eye_in_hand_image
+      type: ProcessLiberoEvalInputs
+    - conditioning_fps: 20.0
+      prepend_state_to_action: false
+      type: SetCosmos3ActionMetadata
+    - image_resize_strategy: resize-naive
+      input_sizes:
+      - - 3
+        - 256
+        - 256
+      - - 3
+        - 256
+        - 256
+      means:
+      - - 127.5
+        - 127.5
+        - 127.5
+      - - 127.5
+        - 127.5
+        - 127.5
+      stds:
+      - - 127.5
+        - 127.5
+        - 127.5
+      - - 127.5
+        - 127.5
+        - 127.5
+      type: TransformImage
+    - action_metadata:
+        append_viewpoint: false
+        conditioning_fps: 20.0
+        frame_window_size: 17
+        video_height: 256
+        video_width: 512
+      cfg_dropout_rate: 0.0
+      max_len: 512
+      output_attention_mask_key: lang_masks
+      output_key: lang_tokens
+      tokenizer:
+        model_max_length: 4096
+        model_path: ./checkpoints/Cosmos3-Edge
+        padding_side: right
+        trust_remote_code: true
+        type: PretrainedTokenizer
+      type: ProcessCosmos3Prompt
+    - gripper_key: robot0_gripper_qpos
+      norm_type: mean_std
+      out_key: states
+      pos_key: robot0_eef_pos
+      quat_key: robot0_eef_quat
+      state_dim: 64
+      type: LiberoProprioFromInputs
+    - frame_window_size: 1
+      num_views: 2
+      tile_direction: horizontal
+      type: PrepareVideo
+    type: LiberoParquetEvalDataset
+  denormalize_action:
+    action_dim: 10
+    norm_type: quantile_rot
+    type: DenormalizeLiberoFramewiseRot6DAction
+  enable_mixed_precision_training: true
+  eval_chunk_size: 16
+  inference_seed: 7
+  mixed_precision_dtype: bf16
+  model_family: cosmos3
+  norm_stats_key: libero_spatial_no_noops
+  num_inference_steps: 30
+  num_steps_wait: 10
+  num_trials_per_task: 50
+  resize_size: 256
+  seed: 7
+  task_ids: null
+  task_suite_name: libero_spatial
+  type: LiberoEvalRunner
+inference_model:
+  action_horizon: 16
+  action_in_proj:
+    input_size: 64
+    num_domains: 32
+    output_size: 2048
+    type: DomainAwareLinear
+  action_out_proj:
+    input_size: 2048
+    num_domains: 32
+    output_size: 64
+    type: DomainAwareLinear
+  base_fps: 24.0
+  enable_fps_modulation: true
+  enable_vision_loss: true
+  freeze_non_moe_vlm_backbone: true
+  freeze_vlm_backbone: false
+  latent_patch_size: 2
+  max_action_dim: 64
+  name_mapping:
+    .self_attn.k_norm_moe_gen.: .self_attn.norm_added_k.
+    .self_attn.k_proj.: .self_attn.to_k.
+    .self_attn.k_proj_moe_gen.: .self_attn.add_k_proj.
+    .self_attn.o_proj.: .self_attn.to_out.
+    .self_attn.o_proj_moe_gen.: .self_attn.to_add_out.
+    .self_attn.q_norm_moe_gen.: .self_attn.norm_added_q.
+    .self_attn.q_proj.: .self_attn.to_q.
+    .self_attn.q_proj_moe_gen.: .self_attn.add_q_proj.
+    .self_attn.v_proj.: .self_attn.to_v.
+    .self_attn.v_proj_moe_gen.: .self_attn.add_v_proj.
+    action_in_proj.: action_proj_in.
+    action_modality_embed: action_modality_embed
+    action_out_proj.: action_proj_out.
+    time_embedder.mlp.0.: time_embedder.linear_1.
+    time_embedder.mlp.2.: time_embedder.linear_2.
+    vision_in_proj.projector.: proj_in.
+    vision_out_proj.projector.: proj_out.
+    vlm_backbone.lm_head.weight: lm_head.weight
+    vlm_backbone.model.language_model.embed_tokens.weight: embed_tokens.weight
+    vlm_backbone.model.language_model.layers.: layers.
+    vlm_backbone.model.language_model.norm.weight: norm.weight
+    vlm_backbone.model.language_model.norm_moe_gen.weight: norm_moe_gen.weight
+  num_embodiment_domains: 32
+  ori_action_dim: 10
+  packed_attention_backend: flash2
+  position_embedding_type: unified_3d_mrope
+  pretrained_name_or_path: ./checkpoints/Cosmos3-Edge/transformer
+  rectified_flow_inference_config:
+    num_steps: 30
+    num_train_timesteps: 1000
+    scheduler_type: unipc
+    shift: 10.0
+    use_dynamic_shifting: false
+    use_karras_sigmas: false
+  rectified_flow_training_config:
+    action_loss_weight: 10.0
+    independent_action_schedule: false
+    normalize_loss_by_active: false
+    shift:
+      '256': 3
+      '480': 5
+      '720': 10
+    shift_action: null
+    train_time_action_distribution: logitnormal
+    train_time_image_distribution: logitnormal
+    train_time_video_distribution: waver
+    train_time_weight: uniform
+    use_discrete_rf: false
+    use_dynamic_shift: false
+    use_high_sigma_strategy: false
+    use_high_sigma_strategy_action: false
+    vision_loss_weight: 10.0
+  reinitialize_action_policy: true
+  special_tokens:
+    end_of_generation: 21
+    eos_token_id: 11
+    start_of_generation: 20
+  strict_mapping: true
+  timestep_scale: 0.001
+  type: Cosmos3FlowMatching
+  unified_3d_mrope_reset_spatial_ids: true
+  unified_3d_mrope_temporal_modality_margin: 15000
+  vision_in_proj:
+    in_dim: 192
+    out_dim: 2048
+    type: LinearProjector
+  vision_latent_dim: 48
+  vision_out_proj:
+    in_dim: 2048
+    out_dim: 192
+    type: LinearProjector
+  vision_vae:
+    encode_exact_durations:
+    - 17
+    pretrained_name_or_path: ./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
+    type: Cosmos3Wan22VAE
+  vlm_backbone:
+    include_visual: false
+    skip_init_weights: true
+    type: Cosmos3MoTBackbone
+    vlm_config:
+      attention_bias: false
+      bos_token_id: 1
+      enable_mrope: true
+      eos_token_id: 11
+      head_dim: 128
+      hidden_size: 2048
+      intermediate_size: 9216
+      layer_norm_epsilon: 1.0e-05
+      max_position_embeddings: 131072
+      mlp_bias: false
+      mlp_hidden_act: relu2
+      model_type: nemotron_3_dense_vl_text
+      mrope_section:
+      - 24
+      - 20
+      - 20
+      num_attention_heads: 16
+      num_hidden_layers: 28
+      num_key_value_heads: 8
+      pad_token_id: 11
+      rope_theta: 100000000.0
+      tie_word_embeddings: false
+      use_und_k_norm_for_gen: true
+      vocab_size: 131072
+model:
+  action_horizon: 16
+  action_in_proj:
+    input_size: 64
+    num_domains: 32
+    output_size: 2048
+    type: DomainAwareLinear
+  action_out_proj:
+    input_size: 2048
+    num_domains: 32
+    output_size: 64
+    type: DomainAwareLinear
+  base_fps: 24.0
+  enable_fps_modulation: true
+  enable_vision_loss: true
+  freeze_non_moe_vlm_backbone: true
+  freeze_vlm_backbone: false
+  latent_patch_size: 2
+  max_action_dim: 64
+  name_mapping:
+    .self_attn.k_norm_moe_gen.: .self_attn.norm_added_k.
+    .self_attn.k_proj.: .self_attn.to_k.
+    .self_attn.k_proj_moe_gen.: .self_attn.add_k_proj.
+    .self_attn.o_proj.: .self_attn.to_out.
+    .self_attn.o_proj_moe_gen.: .self_attn.to_add_out.
+    .self_attn.q_norm_moe_gen.: .self_attn.norm_added_q.
+    .self_attn.q_proj.: .self_attn.to_q.
+    .self_attn.q_proj_moe_gen.: .self_attn.add_q_proj.
+    .self_attn.v_proj.: .self_attn.to_v.
+    .self_attn.v_proj_moe_gen.: .self_attn.add_v_proj.
+    action_in_proj.: action_proj_in.
+    action_modality_embed: action_modality_embed
+    action_out_proj.: action_proj_out.
+    time_embedder.mlp.0.: time_embedder.linear_1.
+    time_embedder.mlp.2.: time_embedder.linear_2.
+    vision_in_proj.projector.: proj_in.
+    vision_out_proj.projector.: proj_out.
+    vlm_backbone.lm_head.weight: lm_head.weight
+    vlm_backbone.model.language_model.embed_tokens.weight: embed_tokens.weight
+    vlm_backbone.model.language_model.layers.: layers.
+    vlm_backbone.model.language_model.norm.weight: norm.weight
+    vlm_backbone.model.language_model.norm_moe_gen.weight: norm_moe_gen.weight
+  num_embodiment_domains: 32
+  ori_action_dim: 10
+  packed_attention_backend: flash2
+  position_embedding_type: unified_3d_mrope
+  pretrained_name_or_path: ./checkpoints/Cosmos3-Edge/transformer
+  rectified_flow_inference_config:
+    num_steps: 30
+    num_train_timesteps: 1000
+    scheduler_type: unipc
+    shift: 10.0
+    use_dynamic_shifting: false
+    use_karras_sigmas: false
+  rectified_flow_training_config:
+    action_loss_weight: 10.0
+    independent_action_schedule: false
+    normalize_loss_by_active: false
+    shift:
+      '256': 3
+      '480': 5
+      '720': 10
+    shift_action: null
+    train_time_action_distribution: logitnormal
+    train_time_image_distribution: logitnormal
+    train_time_video_distribution: waver
+    train_time_weight: uniform
+    use_discrete_rf: false
+    use_dynamic_shift: false
+    use_high_sigma_strategy: false
+    use_high_sigma_strategy_action: false
+    vision_loss_weight: 10.0
+  reinitialize_action_policy: true
+  special_tokens:
+    end_of_generation: 21
+    eos_token_id: 11
+    start_of_generation: 20
+  strict_mapping: true
+  timestep_scale: 0.001
+  type: Cosmos3FlowMatching
+  unified_3d_mrope_reset_spatial_ids: true
+  unified_3d_mrope_temporal_modality_margin: 15000
+  vision_in_proj:
+    in_dim: 192
+    out_dim: 2048
+    type: LinearProjector
+  vision_latent_dim: 48
+  vision_out_proj:
+    in_dim: 2048
+    out_dim: 192
+    type: LinearProjector
+  vision_vae:
+    encode_exact_durations:
+    - 17
+    pretrained_name_or_path: ./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
+    type: Cosmos3Wan22VAE
+  vlm_backbone:
+    include_visual: false
+    skip_init_weights: true
+    type: Cosmos3MoTBackbone
+    vlm_config:
+      attention_bias: false
+      bos_token_id: 1
+      enable_mrope: true
+      eos_token_id: 11
+      head_dim: 128
+      hidden_size: 2048
+      intermediate_size: 9216
+      layer_norm_epsilon: 1.0e-05
+      max_position_embeddings: 131072
+      mlp_bias: false
+      mlp_hidden_act: relu2
+      model_type: nemotron_3_dense_vl_text
+      mrope_section:
+      - 24
+      - 20
+      - 20
+      num_attention_heads: 16
+      num_hidden_layers: 28
+      num_key_value_heads: 8
+      pad_token_id: 11
+      rope_theta: 100000000.0
+      tie_word_embeddings: false
+      use_und_k_norm_for_gen: true
+      vocab_size: 131072
+runner:
+  change_key_name: false
+  collator:
+    list_keys:
+    - sequence_plan
+    meta_keys:
+    - task_description
+    - stats
+    - info
+    - timestamp
+    pad_id: 11
+    sequence_keys:
+    - text_token_ids
+    tensor_keys:
+    - images
+    - actions
+    - embodiment_ids
+    - raw_action_dim
+    - conditioning_fps
+    - action_fps
+    type: Cosmos3Collator
+  enable_gradient_checkpointing: true
+  enable_mixed_precision_training: true
+  grad_accumulation_steps: 16
+  lr_scheduler:
+    cycle_length: 16000
+    type: linear-warmup+linear-decay
+    warmup_steps: 500
+  max_epochs: null
+  max_grad_norm: 1.0
+  max_keep_ckpts: 2
+  max_steps: 2000
+  metric:
+    active_trackers:
+    - jsonl
+    - wandb
+    grad_accumulation_steps: 16
+    run_dir: work_dirs
+    type: VLAMetric
+    window_size: 1
+  mixed_precision_dtype: bf16
+  optimizer:
+    betas:
+    - 0.9
+    - 0.99
+    eps: 1.0e-08
+    exclude_1d_from_weight_decay: false
+    fused: true
+    lr: 5.0e-05
+    paramwise_learning_rate:
+      action_in_proj.: 0.00025
+      action_modality_embed: 0.00025
+      action_out_proj.: 0.00025
+    type: AdamW
+    weight_decay: 0.05
+  sampler: null
+  save_iter_interval: 500
+  sharding_strategy: full-shard
+  tokenizer:
+    model_max_length: 4096
+    model_path: ./checkpoints/Cosmos3-Edge
+    padding_side: right
+    trust_remote_code: true
+    type: PretrainedTokenizer
+  type: FSDPTrainRunner
+seed: 7
+train_dataloader:
+  dataset:
+    datasets:
+      action_key: action
+      action_window_size: 16
+      data_root_path: /mnt/data/cpfs/mnt/data/yanis/FastWAM/data/libero_mujoco3.3.2/libero_spatial_no_noops_lerobot
+      frame_window_size: 17
+      require_full_window: true
+      statistic_name: libero_spatial_no_noops
+      transforms:
+      - embodiment_id: 5
+        name_mappings:
+          actions:
+          - actions
+          observation.state:
+          - states
+        parquet_keys:
+        - observation.state
+        - timestamp
+        - actions
+        - info
+        - stats
+        type: ProcessParquetInputs
+        video_keys:
+        - observation.images.image
+        - observation.images.wrist_image
+      - type: LiberoFramewiseActionToRot6D
+      - action_metadata:
+          append_viewpoint: false
+          conditioning_fps: 20.0
+          frame_window_size: 17
+          video_height: 256
+          video_width: 512
+        cfg_dropout_rate: 0.1
+        max_len: 512
+        tokenizer:
+          model_max_length: 4096
+          model_path: ./checkpoints/Cosmos3-Edge
+          padding_side: right
+          trust_remote_code: true
+          type: PretrainedTokenizer
+        type: ProcessCosmos3Prompt
+      - height: 256
+        type: ResizeImages
+        width: 256
+      - type: SimpleNormalizeImages
+      - action_dim: 64
+        action_key: action
+        action_norm_type: quantile
+        state_dim: 64
+        state_key: proprio
+        state_norm_type: none
+        type: NormalizeStatesAndActions
+      - conditioning_fps: 20.0
+        frame_window_size: 17
+        mode: wam
+        prepend_state_to_action: false
+        raw_action_dim: 10
+        type: BuildCosmos3Sequence
+      - frame_window_size: 17
+        num_views: 2
+        tile_direction: horizontal
+        type: PrepareVideo
+      type: ParquetDataset
+      use_delta: false
+      window_start_idx: 0
+    name_mappings:
+      action:
+      - action
+      observation.state:
+      - proprio
+    statistic_keys:
+    - observation.state
+    - timestamp
+    - action
+    statistic_name: libero_spatial_no_noops
+    statistics_overrides:
+      libero_spatial_no_noops:
+        action:
+          max:
+          - 0.9375
+          - 0.9375
+          - 0.9375
+          - 1.0
+          - 0.368853
+          - 0.341214
+          - 0.356395
+          - 1.0
+          - 0.348251
+          - 1.0
+          mean:
+          - 0.050704
+          - 0.097407
+          - -0.094833
+          - 0.994873
+          - -0.004579
+          - -0.004288
+          - 0.004389
+          - 0.996104
+          - 0.001109
+          - 0.476725
+          min:
+          - -0.9375
+          - -0.9375
+          - -0.9375
+          - 0.902028
+          - -0.356085
+          - -0.367416
+          - -0.370434
+          - 0.921907
+          - -0.255
+          - 0.0
+          q01:
+          - -0.723214
+          - -0.808929
+          - -0.9375
+          - 0.934955
+          - -0.223431
+          - -0.189878
+          - -0.334735
+          - 0.938516
+          - -0.107736
+          - 0.0
+          q99:
+          - 0.9375
+          - 0.870536
+          - 0.9375
+          - 1.0
+          - 0.331
+          - 0.163153
+          - 0.226216
+          - 1.0
+          - 0.127158
+          - 1.0
+          std:
+          - 0.333621
+          - 0.387175
+          - 0.45714
+          - 0.010807
+          - 0.077802
+          - 0.063386
+          - 0.078571
+          - 0.009994
+          - 0.038504
+          - 0.49946
+    type: DistributedRepeatingDataset
+  per_device_batch_size: 16
+  per_device_num_workers: 4

cosmos3edge_libero_spatial_full_finetune_bs16/dataset_statistics.json

@@ -0,0 +1,138 @@
+{
+  "libero_spatial_no_noops": {
+    "proprio": {
+      "mean": [
+        -0.02451290873970331,
+        0.10615150130733945,
+        1.0578400251960747,
+        3.061689349795143,
+        -0.1005965057813455,
+        0.08257375486551881,
+        0.019912415870854235,
+        -0.020113726667093317
+      ],
+      "std": [
+        0.11053132910461624,
+        0.13747598538616676,
+        0.1043268470449716,
+        0.1059396829997249,
+        0.41110779566445044,
+        0.21523187896802917,
+        0.017277071353936354,
+        0.01711604738631362
+      ],
+      "min": [
+        -0.31721994280815125,
+        -0.2926647663116455,
+        0.9094576835632324,
+        2.496100664138794,
+        -1.798353910446167,
+        -0.7207611203193665,
+        -7.665253360755742e-05,
+        -0.041329532861709595
+      ],
+      "max": [
+        0.1759040206670761,
+        0.3904264271259308,
+        1.3290090560913086,
+        3.4572718143463135,
+        1.2369194030761719,
+        1.042615294456482,
+        0.040996309369802475,
+        0.0005872611072845757
+      ],
+      "q01": null,
+      "q99": null
+    },
+    "timestamp": {
+      "mean": [
+        3.129841815551673
+      ],
+      "std": [
+        1.9217150401170382
+      ],
+      "min": [
+        0.0
+      ],
+      "max": [
+        9.6
+      ],
+      "q01": null,
+      "q99": null
+    },
+    "action": {
+      "mean": [
+        0.050704,
+        0.097407,
+        -0.094833,
+        0.994873,
+        -0.004579,
+        -0.004288,
+        0.004389,
+        0.996104,
+        0.001109,
+        0.476725
+      ],
+      "std": [
+        0.333621,
+        0.387175,
+        0.45714,
+        0.010807,
+        0.077802,
+        0.063386,
+        0.078571,
+        0.009994,
+        0.038504,
+        0.49946
+      ],
+      "min": [
+        -0.9375,
+        -0.9375,
+        -0.9375,
+        0.902028,
+        -0.356085,
+        -0.367416,
+        -0.370434,
+        0.921907,
+        -0.255,
+        0.0
+      ],
+      "max": [
+        0.9375,
+        0.9375,
+        0.9375,
+        1.0,
+        0.368853,
+        0.341214,
+        0.356395,
+        1.0,
+        0.348251,
+        1.0
+      ],
+      "q01": [
+        -0.723214,
+        -0.808929,
+        -0.9375,
+        0.934955,
+        -0.223431,
+        -0.189878,
+        -0.334735,
+        0.938516,
+        -0.107736,
+        0.0
+      ],
+      "q99": [
+        0.9375,
+        0.870536,
+        0.9375,
+        1.0,
+        0.331,
+        0.163153,
+        0.226216,
+        1.0,
+        0.127158,
+        1.0
+      ]
+    }
+  }
+}

cosmos3edge_libero_spatial_full_finetune_bs16/run-metrics.jsonl

@@ -0,0 +1 @@
+{"hparams": "{'eval': {'dataset': {'extra_tensor_keys': ['conditioning_fps', 'prepend_state_to_action'], 'img_buffer_len': 1, 'transforms': [{'embodiment_id': 5, 'img_keys': ['agentview_image', 'robot0_eye_in_hand_image'], 'type': 'ProcessLiberoEvalInputs'}, {'conditioning_fps': 20.0, 'prepend_state_to_action': False, 'type': 'SetCosmos3ActionMetadata'}, {'image_resize_strategy': 'resize-naive', 'input_sizes': [[3, 256, 256], [3, 256, 256]], 'means': [[127.5, 127.5, 127.5], [127.5, 127.5, 127.5]], 'stds': [[127.5, 127.5, 127.5], [127.5, 127.5, 127.5]], 'type': 'TransformImage'}, {'action_metadata': {'append_viewpoint': False, 'conditioning_fps': 20.0, 'frame_window_size': 17, 'video_height': 256, 'video_width': 512}, 'cfg_dropout_rate': 0.0, 'max_len': 512, 'output_attention_mask_key': 'lang_masks', 'output_key': 'lang_tokens', 'tokenizer': {'model_max_length': 4096, 'model_path': './checkpoints/Cosmos3-Edge', 'padding_side': 'right', 'trust_remote_code': True, 'type': 'PretrainedTokenizer'}, 'type': 'ProcessCosmos3Prompt'}, {'gripper_key': 'robot0_gripper_qpos', 'norm_type': 'mean_std', 'out_key': 'states', 'pos_key': 'robot0_eef_pos', 'quat_key': 'robot0_eef_quat', 'state_dim': 64, 'type': 'LiberoProprioFromInputs'}, {'frame_window_size': 1, 'num_views': 2, 'tile_direction': 'horizontal', 'type': 'PrepareVideo'}], 'type': 'LiberoParquetEvalDataset'}, 'denormalize_action': {'action_dim': 10, 'norm_type': 'quantile_rot', 'type': 'DenormalizeLiberoFramewiseRot6DAction'}, 'enable_mixed_precision_training': True, 'eval_chunk_size': 16, 'inference_seed': 7, 'mixed_precision_dtype': 'bf16', 'model_family': 'cosmos3', 'norm_stats_key': 'libero_spatial_no_noops', 'num_inference_steps': 30, 'num_steps_wait': 10, 'num_trials_per_task': 50, 'resize_size': 256, 'seed': 7, 'task_ids': None, 'task_suite_name': 'libero_spatial', 'type': 'LiberoEvalRunner'}, 'inference_model': {'action_horizon': 16, 'action_in_proj': {'input_size': 64, 'num_domains': 32, 'output_size': 2048, 'type': 'DomainAwareLinear'}, 'action_out_proj': {'input_size': 2048, 'num_domains': 32, 'output_size': 64, 'type': 'DomainAwareLinear'}, 'base_fps': 24.0, 'enable_fps_modulation': True, 'enable_vision_loss': True, 'freeze_non_moe_vlm_backbone': True, 'freeze_vlm_backbone': False, 'latent_patch_size': 2, 'max_action_dim': 64, 'name_mapping': {'.self_attn.k_norm_moe_gen.': '.self_attn.norm_added_k.', '.self_attn.k_proj.': '.self_attn.to_k.', '.self_attn.k_proj_moe_gen.': '.self_attn.add_k_proj.', '.self_attn.o_proj.': '.self_attn.to_out.', '.self_attn.o_proj_moe_gen.': '.self_attn.to_add_out.', '.self_attn.q_norm_moe_gen.': '.self_attn.norm_added_q.', '.self_attn.q_proj.': '.self_attn.to_q.', '.self_attn.q_proj_moe_gen.': '.self_attn.add_q_proj.', '.self_attn.v_proj.': '.self_attn.to_v.', '.self_attn.v_proj_moe_gen.': '.self_attn.add_v_proj.', 'action_in_proj.': 'action_proj_in.', 'action_modality_embed': 'action_modality_embed', 'action_out_proj.': 'action_proj_out.', 'time_embedder.mlp.0.': 'time_embedder.linear_1.', 'time_embedder.mlp.2.': 'time_embedder.linear_2.', 'vision_in_proj.projector.': 'proj_in.', 'vision_out_proj.projector.': 'proj_out.', 'vlm_backbone.lm_head.weight': 'lm_head.weight', 'vlm_backbone.model.language_model.embed_tokens.weight': 'embed_tokens.weight', 'vlm_backbone.model.language_model.layers.': 'layers.', 'vlm_backbone.model.language_model.norm.weight': 'norm.weight', 'vlm_backbone.model.language_model.norm_moe_gen.weight': 'norm_moe_gen.weight'}, 'num_embodiment_domains': 32, 'ori_action_dim': 10, 'packed_attention_backend': 'flash2', 'position_embedding_type': 'unified_3d_mrope', 'pretrained_name_or_path': './checkpoints/Cosmos3-Edge/transformer', 'reinitialize_action_policy': True, 'rectified_flow_inference_config': {'num_steps': 30, 'num_train_timesteps': 1000, 'scheduler_type': 'unipc', 'shift': 10.0, 'use_dynamic_shifting': False, 'use_karras_sigmas': False}, 'rectified_flow_training_config': {'action_loss_weight': 10.0, 'independent_action_schedule': False, 'normalize_loss_by_active': False, 'shift': {'256': 3, '480': 5, '720': 10}, 'shift_action': None, 'train_time_action_distribution': 'logitnormal', 'train_time_image_distribution': 'logitnormal', 'train_time_video_distribution': 'waver', 'train_time_weight': 'uniform', 'use_discrete_rf': False, 'use_dynamic_shift': False, 'use_high_sigma_strategy': False, 'use_high_sigma_strategy_action': False, 'vision_loss_weight': 10.0}, 'special_tokens': {'end_of_generation': 21, 'eos_token_id': 11, 'start_of_generation': 20}, 'strict_mapping': True, 'timestep_scale': 0.001, 'type': 'Cosmos3FlowMatching', 'unified_3d_mrope_reset_spatial_ids': True, 'unified_3d_mrope_temporal_modality_margin': 15000, 'vision_in_proj': {'in_dim': 192, 'out_dim': 2048, 'type': 'LinearProjector'}, 'vision_latent_dim': 48, 'vision_out_proj': {'in_dim': 2048, 'out_dim': 192, 'type': 'LinearProjector'}, 'vision_vae': {'encode_exact_durations': [17], 'pretrained_name_or_path': './checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth', 'type': 'Cosmos3Wan22VAE'}, 'vlm_backbone': {'include_visual': False, 'skip_init_weights': True, 'type': 'Cosmos3MoTBackbone', 'vlm_config': {'attention_bias': False, 'bos_token_id': 1, 'enable_mrope': True, 'eos_token_id': 11, 'head_dim': 128, 'hidden_size': 2048, 'intermediate_size': 9216, 'layer_norm_epsilon': 1e-05, 'max_position_embeddings': 131072, 'mlp_bias': False, 'mlp_hidden_act': 'relu2', 'model_type': 'nemotron_3_dense_vl_text', 'mrope_section': [24, 20, 20], 'num_attention_heads': 16, 'num_hidden_layers': 28, 'num_key_value_heads': 8, 'pad_token_id': 11, 'rope_theta': 100000000.0, 'tie_word_embeddings': False, 'use_und_k_norm_for_gen': True, 'vocab_size': 131072}}}, 'model': {'action_horizon': 16, 'action_in_proj': {'input_size': 64, 'num_domains': 32, 'output_size': 2048, 'type': 'DomainAwareLinear'}, 'action_out_proj': {'input_size': 2048, 'num_domains': 32, 'output_size': 64, 'type': 'DomainAwareLinear'}, 'base_fps': 24.0, 'enable_fps_modulation': True, 'enable_vision_loss': True, 'freeze_non_moe_vlm_backbone': True, 'freeze_vlm_backbone': False, 'latent_patch_size': 2, 'max_action_dim': 64, 'name_mapping': {'.self_attn.k_norm_moe_gen.': '.self_attn.norm_added_k.', '.self_attn.k_proj.': '.self_attn.to_k.', '.self_attn.k_proj_moe_gen.': '.self_attn.add_k_proj.', '.self_attn.o_proj.': '.self_attn.to_out.', '.self_attn.o_proj_moe_gen.': '.self_attn.to_add_out.', '.self_attn.q_norm_moe_gen.': '.self_attn.norm_added_q.', '.self_attn.q_proj.': '.self_attn.to_q.', '.self_attn.q_proj_moe_gen.': '.self_attn.add_q_proj.', '.self_attn.v_proj.': '.self_attn.to_v.', '.self_attn.v_proj_moe_gen.': '.self_attn.add_v_proj.', 'action_in_proj.': 'action_proj_in.', 'action_modality_embed': 'action_modality_embed', 'action_out_proj.': 'action_proj_out.', 'time_embedder.mlp.0.': 'time_embedder.linear_1.', 'time_embedder.mlp.2.': 'time_embedder.linear_2.', 'vision_in_proj.projector.': 'proj_in.', 'vision_out_proj.projector.': 'proj_out.', 'vlm_backbone.lm_head.weight': 'lm_head.weight', 'vlm_backbone.model.language_model.embed_tokens.weight': 'embed_tokens.weight', 'vlm_backbone.model.language_model.layers.': 'layers.', 'vlm_backbone.model.language_model.norm.weight': 'norm.weight', 'vlm_backbone.model.language_model.norm_moe_gen.weight': 'norm_moe_gen.weight'}, 'num_embodiment_domains': 32, 'ori_action_dim': 10, 'packed_attention_backend': 'flash2', 'position_embedding_type': 'unified_3d_mrope', 'pretrained_name_or_path': './checkpoints/Cosmos3-Edge/transformer', 'reinitialize_action_policy': True, 'rectified_flow_inference_config': {'num_steps': 30, 'num_train_timesteps': 1000, 'scheduler_type': 'unipc', 'shift': 10.0, 'use_dynamic_shifting': False, 'use_karras_sigmas': False}, 'rectified_flow_training_config': {'action_loss_weight': 10.0, 'independent_action_schedule': False, 'normalize_loss_by_active': False, 'shift': {'256': 3, '480': 5, '720': 10}, 'shift_action': None, 'train_time_action_distribution': 'logitnormal', 'train_time_image_distribution': 'logitnormal', 'train_time_video_distribution': 'waver', 'train_time_weight': 'uniform', 'use_discrete_rf': False, 'use_dynamic_shift': False, 'use_high_sigma_strategy': False, 'use_high_sigma_strategy_action': False, 'vision_loss_weight': 10.0}, 'special_tokens': {'end_of_generation': 21, 'eos_token_id': 11, 'start_of_generation': 20}, 'strict_mapping': True, 'timestep_scale': 0.001, 'type': 'Cosmos3FlowMatching', 'unified_3d_mrope_reset_spatial_ids': True, 'unified_3d_mrope_temporal_modality_margin': 15000, 'vision_in_proj': {'in_dim': 192, 'out_dim': 2048, 'type': 'LinearProjector'}, 'vision_latent_dim': 48, 'vision_out_proj': {'in_dim': 2048, 'out_dim': 192, 'type': 'LinearProjector'}, 'vision_vae': {'encode_exact_durations': [17], 'pretrained_name_or_path': './checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth', 'type': 'Cosmos3Wan22VAE'}, 'vlm_backbone': {'include_visual': False, 'skip_init_weights': True, 'type': 'Cosmos3MoTBackbone', 'vlm_config': {'attention_bias': False, 'bos_token_id': 1, 'enable_mrope': True, 'eos_token_id': 11, 'head_dim': 128, 'hidden_size': 2048, 'intermediate_size': 9216, 'layer_norm_epsilon': 1e-05, 'max_position_embeddings': 131072, 'mlp_bias': False, 'mlp_hidden_act': 'relu2', 'model_type': 'nemotron_3_dense_vl_text', 'mrope_section': [24, 20, 20], 'num_attention_heads': 16, 'num_hidden_layers': 28, 'num_key_value_heads': 8, 'pad_token_id': 11, 'rope_theta': 100000000.0, 'tie_word_embeddings': False, 'use_und_k_norm_for_gen': True, 'vocab_size': 131072}}}, 'runner': {'change_key_name': False, 'collator': {'list_keys': ['sequence_plan'], 'meta_keys': ['task_description', 'stats', 'info', 'timestamp'], 'pad_id': 11, 'sequence_keys': ['text_token_ids'], 'tensor_keys': ['images', 'actions', 'embodiment_ids', 'raw_action_dim', 'conditioning_fps', 'action_fps'], 'type': 'Cosmos3Collator'}, 'enable_gradient_checkpointing': True, 'enable_mixed_precision_training': True, 'grad_accumulation_steps': 16, 'lr_scheduler': {'cycle_length': 16000, 'type': 'linear-warmup+linear-decay', 'warmup_steps': 500}, 'max_epochs': None, 'max_grad_norm': 1.0, 'max_keep_ckpts': 2, 'max_steps': 2000, 'metric': {'active_trackers': ('jsonl', 'wandb'), 'grad_accumulation_steps': 16, 'run_dir': '/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_spatial_edge_pr51_latest_bs16', 'type': 'VLAMetric', 'window_size': 1}, 'mixed_precision_dtype': 'bf16', 'optimizer': {'betas': (0.9, 0.99), 'eps': 1e-08, 'exclude_1d_from_weight_decay': False, 'fused': True, 'lr': 5e-05, 'paramwise_learning_rate': {'action_in_proj.': 0.00025, 'action_modality_embed': 0.00025, 'action_out_proj.': 0.00025}, 'type': 'AdamW', 'weight_decay': 0.05}, 'sampler': None, 'save_iter_interval': 500, 'sharding_strategy': 'full-shard', 'tokenizer': {'model_max_length': 4096, 'model_path': './checkpoints/Cosmos3-Edge', 'padding_side': 'right', 'trust_remote_code': True, 'type': 'PretrainedTokenizer'}, 'type': 'FSDPTrainRunner', 'cfg': Config (path: configs/cosmos3/cosmos3edge_libero_spatial_full_finetune.py): {'eval': {'dataset': {'extra_tensor_keys': ['conditioning_fps', 'prepend_state_to_action'], 'img_buffer_len': 1, 'transforms': [{'embodiment_id': 5, 'img_keys': ['agentview_image', 'robot0_eye_in_hand_image'], 'type': 'ProcessLiberoEvalInputs'}, {'conditioning_fps': 20.0, 'prepend_state_to_action': False, 'type': 'SetCosmos3ActionMetadata'}, {'image_resize_strategy': 'resize-naive', 'input_sizes': [[3, 256, 256], [3, 256, 256]], 'means': [[127.5, 127.5, 127.5], [127.5, 127.5, 127.5]], 'stds': [[127.5, 127.5, 127.5], [127.5, 127.5, 127.5]], 'type': 'TransformImage'}, {'action_metadata': {'append_viewpoint': False, 'conditioning_fps': 20.0, 'frame_window_size': 17, 'video_height': 256, 'video_width': 512}, 'cfg_dropout_rate': 0.0, 'max_len': 512, 'output_attention_mask_key': 'lang_masks', 'output_key': 'lang_tokens', 'tokenizer': {'model_max_length': 4096, 'model_path': './checkpoints/Cosmos3-Edge', 'padding_side': 'right', 'trust_remote_code': True, 'type': 'PretrainedTokenizer'}, 'type': 'ProcessCosmos3Prompt'}, {'gripper_key': 'robot0_gripper_qpos', 'norm_type': 'mean_std', 'out_key': 'states', 'pos_key': 'robot0_eef_pos', 'quat_key': 'robot0_eef_quat', 'state_dim': 64, 'type': 'LiberoProprioFromInputs'}, {'frame_window_size': 1, 'num_views': 2, 'tile_direction': 'horizontal', 'type': 'PrepareVideo'}], 'type': 'LiberoParquetEvalDataset'}, 'denormalize_action': {'action_dim': 10, 'norm_type': 'quantile_rot', 'type': 'DenormalizeLiberoFramewiseRot6DAction'}, 'enable_mixed_precision_training': True, 'eval_chunk_size': 16, 'inference_seed': 7, 'mixed_precision_dtype': 'bf16', 'model_family': 'cosmos3', 'norm_stats_key': 'libero_spatial_no_noops', 'num_inference_steps': 30, 'num_steps_wait': 10, 'num_trials_per_task': 50, 'resize_size': 256, 'seed': 7, 'task_ids': None, 'task_suite_name': 'libero_spatial', 'type': 'LiberoEvalRunner'}, 'inference_model': {'action_horizon': 16, 'action_in_proj': {'input_size': 64, 'num_domains': 32, 'output_size': 2048, 'type': 'DomainAwareLinear'}, 'action_out_proj': {'input_size': 2048, 'num_domains': 32, 'output_size': 64, 'type': 'DomainAwareLinear'}, 'base_fps': 24.0, 'enable_fps_modulation': True, 'enable_vision_loss': True, 'freeze_non_moe_vlm_backbone': True, 'freeze_vlm_backbone': False, 'latent_patch_size': 2, 'max_action_dim': 64, 'name_mapping': {'.self_attn.k_norm_moe_gen.': '.self_attn.norm_added_k.', '.self_attn.k_proj.': '.self_attn.to_k.', '.self_attn.k_proj_moe_gen.': '.self_attn.add_k_proj.', '.self_attn.o_proj.': '.self_attn.to_out.', '.self_attn.o_proj_moe_gen.': '.self_attn.to_add_out.', '.self_attn.q_norm_moe_gen.': '.self_attn.norm_added_q.', '.self_attn.q_proj.': '.self_attn.to_q.', '.self_attn.q_proj_moe_gen.': '.self_attn.add_q_proj.', '.self_attn.v_proj.': '.self_attn.to_v.', '.self_attn.v_proj_moe_gen.': '.self_attn.add_v_proj.', 'action_in_proj.': 'action_proj_in.', 'action_modality_embed': 'action_modality_embed', 'action_out_proj.': 'action_proj_out.', 'time_embedder.mlp.0.': 'time_embedder.linear_1.', 'time_embedder.mlp.2.': 'time_embedder.linear_2.', 'vision_in_proj.projector.': 'proj_in.', 'vision_out_proj.projector.': 'proj_out.', 'vlm_backbone.lm_head.weight': 'lm_head.weight', 'vlm_backbone.model.language_model.embed_tokens.weight': 'embed_tokens.weight', 'vlm_backbone.model.language_model.layers.': 'layers.', 'vlm_backbone.model.language_model.norm.weight': 'norm.weight', 'vlm_backbone.model.language_model.norm_moe_gen.weight': 'norm_moe_gen.weight'}, 'num_embodiment_domains': 32, 'ori_action_dim': 10, 'packed_attention_backend': 'flash2', 'position_embedding_type': 'unified_3d_mrope', 'pretrained_name_or_path': './checkpoints/Cosmos3-Edge/transformer', 'reinitialize_action_policy': True, 'rectified_flow_inference_config': {'num_steps': 30, 'num_train_timesteps': 1000, 'scheduler_type': 'unipc', 'shift': 10.0, 'use_dynamic_shifting': False, 'use_karras_sigmas': False}, 'rectified_flow_training_config': {'action_loss_weight': 10.0, 'independent_action_schedule': False, 'normalize_loss_by_active': False, 'shift': {'256': 3, '480': 5, '720': 10}, 'shift_action': None, 'train_time_action_distribution': 'logitnormal', 'train_time_image_distribution': 'logitnormal', 'train_time_video_distribution': 'waver', 'train_time_weight': 'uniform', 'use_discrete_rf': False, 'use_dynamic_shift': False, 'use_high_sigma_strategy': False, 'use_high_sigma_strategy_action': False, 'vision_loss_weight': 10.0}, 'special_tokens': {'end_of_generation': 21, 'eos_token_id': 11, 'start_of_generation': 20}, 'strict_mapping': True, 'timestep_scale': 0.001, 'type': 'Cosmos3FlowMatching', 'unified_3d_mrope_reset_spatial_ids': True, 'unified_3d_mrope_temporal_modality_margin': 15000, 'vision_in_proj': {'in_dim': 192, 'out_dim': 2048, 'type': 'LinearProjector'}, 'vision_latent_dim': 48, 'vision_out_proj': {'in_dim': 2048, 'out_dim': 192, 'type': 'LinearProjector'}, 'vision_vae': {'encode_exact_durations': [17], 'pretrained_name_or_path': './checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth', 'type': 'Cosmos3Wan22VAE'}, 'vlm_backbone': {'include_visual': False, 'skip_init_weights': True, 'type': 'Cosmos3MoTBackbone', 'vlm_config': {'attention_bias': False, 'bos_token_id': 1, 'enable_mrope': True, 'eos_token_id': 11, 'head_dim': 128, 'hidden_size': 2048, 'intermediate_size': 9216, 'layer_norm_epsilon': 1e-05, 'max_position_embeddings': 131072, 'mlp_bias': False, 'mlp_hidden_act': 'relu2', 'model_type': 'nemotron_3_dense_vl_text', 'mrope_section': [24, 20, 20], 'num_attention_heads': 16, 'num_hidden_layers': 28, 'num_key_value_heads': 8, 'pad_token_id': 11, 'rope_theta': 100000000.0, 'tie_word_embeddings': False, 'use_und_k_norm_for_gen': True, 'vocab_size': 131072}}}, 'model': {'action_horizon': 16, 'action_in_proj': {'input_size': 64, 'num_domains': 32, 'output_size': 2048, 'type': 'DomainAwareLinear'}, 'action_out_proj': {'input_size': 2048, 'num_domains': 32, 'output_size': 64, 'type': 'DomainAwareLinear'}, 'base_fps': 24.0, 'enable_fps_modulation': True, 'enable_vision_loss': True, 'freeze_non_moe_vlm_backbone': True, 'freeze_vlm_backbone': False, 'latent_patch_size': 2, 'max_action_dim': 64, 'name_mapping': {'.self_attn.k_norm_moe_gen.': '.self_attn.norm_added_k.', '.self_attn.k_proj.': '.self_attn.to_k.', '.self_attn.k_proj_moe_gen.': '.self_attn.add_k_proj.', '.self_attn.o_proj.': '.self_attn.to_out.', '.self_attn.o_proj_moe_gen.': '.self_attn.to_add_out.', '.self_attn.q_norm_moe_gen.': '.self_attn.norm_added_q.', '.self_attn.q_proj.': '.self_attn.to_q.', '.self_attn.q_proj_moe_gen.': '.self_attn.add_q_proj.', '.self_attn.v_proj.': '.self_attn.to_v.', '.self_attn.v_proj_moe_gen.': '.self_attn.add_v_proj.', 'action_in_proj.': 'action_proj_in.', 'action_modality_embed': 'action_modality_embed', 'action_out_proj.': 'action_proj_out.', 'time_embedder.mlp.0.': 'time_embedder.linear_1.', 'time_embedder.mlp.2.': 'time_embedder.linear_2.', 'vision_in_proj.projector.': 'proj_in.', 'vision_out_proj.projector.': 'proj_out.', 'vlm_backbone.lm_head.weight': 'lm_head.weight', 'vlm_backbone.model.language_model.embed_tokens.weight': 'embed_tokens.weight', 'vlm_backbone.model.language_model.layers.': 'layers.', 'vlm_backbone.model.language_model.norm.weight': 'norm.weight', 'vlm_backbone.model.language_model.norm_moe_gen.weight': 'norm_moe_gen.weight'}, 'num_embodiment_domains': 32, 'ori_action_dim': 10, 'packed_attention_backend': 'flash2', 'position_embedding_type': 'unified_3d_mrope', 'pretrained_name_or_path': './checkpoints/Cosmos3-Edge/transformer', 'reinitialize_action_policy': True, 'rectified_flow_inference_config': {'num_steps': 30, 'num_train_timesteps': 1000, 'scheduler_type': 'unipc', 'shift': 10.0, 'use_dynamic_shifting': False, 'use_karras_sigmas': False}, 'rectified_flow_training_config': {'action_loss_weight': 10.0, 'independent_action_schedule': False, 'normalize_loss_by_active': False, 'shift': {'256': 3, '480': 5, '720': 10}, 'shift_action': None, 'train_time_action_distribution': 'logitnormal', 'train_time_image_distribution': 'logitnormal', 'train_time_video_distribution': 'waver', 'train_time_weight': 'uniform', 'use_discrete_rf': False, 'use_dynamic_shift': False, 'use_high_sigma_strategy': False, 'use_high_sigma_strategy_action': False, 'vision_loss_weight': 10.0}, 'special_tokens': {'end_of_generation': 21, 'eos_token_id': 11, 'start_of_generation': 20}, 'strict_mapping': True, 'timestep_scale': 0.001, 'type': 'Cosmos3FlowMatching', 'unified_3d_mrope_reset_spatial_ids': True, 'unified_3d_mrope_temporal_modality_margin': 15000, 'vision_in_proj': {'in_dim': 192, 'out_dim': 2048, 'type': 'LinearProjector'}, 'vision_latent_dim': 48, 'vision_out_proj': {'in_dim': 2048, 'out_dim': 192, 'type': 'LinearProjector'}, 'vision_vae': {'encode_exact_durations': [17], 'pretrained_name_or_path': './checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth', 'type': 'Cosmos3Wan22VAE'}, 'vlm_backbone': {'include_visual': False, 'skip_init_weights': True, 'type': 'Cosmos3MoTBackbone', 'vlm_config': {'attention_bias': False, 'bos_token_id': 1, 'enable_mrope': True, 'eos_token_id': 11, 'head_dim': 128, 'hidden_size': 2048, 'intermediate_size': 9216, 'layer_norm_epsilon': 1e-05, 'max_position_embeddings': 131072, 'mlp_bias': False, 'mlp_hidden_act': 'relu2', 'model_type': 'nemotron_3_dense_vl_text', 'mrope_section': [24, 20, 20], 'num_attention_heads': 16, 'num_hidden_layers': 28, 'num_key_value_heads': 8, 'pad_token_id': 11, 'rope_theta': 100000000.0, 'tie_word_embeddings': False, 'use_und_k_norm_for_gen': True, 'vocab_size': 131072}}}, 'runner': {'change_key_name': False, 'collator': {'list_keys': ['sequence_plan'], 'meta_keys': ['task_description', 'stats', 'info', 'timestamp'], 'pad_id': 11, 'sequence_keys': ['text_token_ids'], 'tensor_keys': ['images', 'actions', 'embodiment_ids', 'raw_action_dim', 'conditioning_fps', 'action_fps'], 'type': 'Cosmos3Collator'}, 'enable_gradient_checkpointing': True, 'enable_mixed_precision_training': True, 'grad_accumulation_steps': 16, 'lr_scheduler': {'cycle_length': 16000, 'type': 'linear-warmup+linear-decay', 'warmup_steps': 500}, 'max_epochs': None, 'max_grad_norm': 1.0, 'max_keep_ckpts': 2, 'max_steps': 2000, 'metric': {'active_trackers': ('jsonl', 'wandb'), 'grad_accumulation_steps': 16, 'run_dir': '/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_spatial_edge_pr51_latest_bs16', 'type': 'VLAMetric', 'window_size': 1}, 'mixed_precision_dtype': 'bf16', 'optimizer': {'betas': (0.9, 0.99), 'eps': 1e-08, 'exclude_1d_from_weight_decay': False, 'fused': True, 'lr': 5e-05, 'paramwise_learning_rate': {'action_in_proj.': 0.00025, 'action_modality_embed': 0.00025, 'action_out_proj.': 0.00025}, 'type': 'AdamW', 'weight_decay': 0.05}, 'sampler': None, 'save_iter_interval': 500, 'sharding_strategy': 'full-shard', 'tokenizer': {'model_max_length': 4096, 'model_path': './checkpoints/Cosmos3-Edge', 'padding_side': 'right', 'trust_remote_code': True, 'type': 'PretrainedTokenizer'}, 'type': 'FSDPTrainRunner', 'cfg': Config (path: configs/cosmos3/cosmos3edge_libero_spatial_full_finetune.py): {...}, 'args': Namespace(config='configs/cosmos3/cosmos3edge_libero_spatial_full_finetune.py', work_dir='/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_spatial_edge_pr51_latest_bs16', cfg_options={'train_dataloader.per_device_batch_size': 16, 'runner.grad_accumulation_steps': 16, 'runner.metric.grad_accumulation_steps': 16}, eval_after_train=True, resume_from=None)}, 'seed': 7, 'train_dataloader': {'dataset': {'datasets': {'action_key': 'action', 'action_window_size': 16, 'data_root_path': '/mnt/data/cpfs/mnt/data/yanis/FastWAM/data/libero_mujoco3.3.2/libero_spatial_no_noops_lerobot', 'frame_window_size': 17, 'require_full_window': True, 'statistic_name': 'libero_spatial_no_noops', 'transforms': [{'embodiment_id': 5, 'name_mappings': {'actions': ['actions'], 'observation.state': ['states']}, 'parquet_keys': ['observation.state', 'timestamp', 'actions', 'info', 'stats'], 'type': 'ProcessParquetInputs', 'video_keys': ['observation.images.image', 'observation.images.wrist_image']}, {'type': 'LiberoFramewiseActionToRot6D'}, {'action_metadata': {'append_viewpoint': False, 'conditioning_fps': 20.0, 'frame_window_size': 17, 'video_height': 256, 'video_width': 512}, 'cfg_dropout_rate': 0.1, 'max_len': 512, 'tokenizer': {'model_max_length': 4096, 'model_path': './checkpoints/Cosmos3-Edge', 'padding_side': 'right', 'trust_remote_code': True, 'type': 'PretrainedTokenizer'}, 'type': 'ProcessCosmos3Prompt'}, {'height': 256, 'type': 'ResizeImages', 'width': 256}, {'type': 'SimpleNormalizeImages'}, {'action_dim': 64, 'action_key': 'action', 'action_norm_type': 'quantile', 'state_norm_type': 'none', 'state_dim': 64, 'state_key': 'proprio', 'type': 'NormalizeStatesAndActions'}, {'conditioning_fps': 20.0, 'frame_window_size': 17, 'mode': 'wam', 'prepend_state_to_action': False, 'raw_action_dim': 10, 'type': 'BuildCosmos3Sequence'}, {'frame_window_size': 17, 'num_views': 2, 'tile_direction': 'horizontal', 'type': 'PrepareVideo'}], 'type': 'ParquetDataset', 'use_delta': False, 'window_start_idx': 0}, 'name_mappings': {'action': ['action'], 'observation.state': ['proprio']}, 'statistic_keys': ['observation.state', 'timestamp', 'action'], 'statistic_name': 'libero_spatial_no_noops', 'statistics_overrides': {'libero_spatial_no_noops': {'action': {'mean': [0.050704, 0.097407, -0.094833, 0.994873, -0.004579, -0.004288, 0.004389, 0.996104, 0.001109, 0.476725], 'std': [0.333621, 0.387175, 0.45714, 0.010807, 0.077802, 0.063386, 0.078571, 0.009994, 0.038504, 0.49946], 'min': [-0.9375, -0.9375, -0.9375, 0.902028, -0.356085, -0.367416, -0.370434, 0.921907, -0.255, 0.0], 'max': [0.9375, 0.9375, 0.9375, 1.0, 0.368853, 0.341214, 0.356395, 1.0, 0.348251, 1.0], 'q01': [-0.723214, -0.808929, -0.9375, 0.934955, -0.223431, -0.189878, -0.334735, 0.938516, -0.107736, 0.0], 'q99': [0.9375, 0.870536, 0.9375, 1.0, 0.331, 0.163153, 0.226216, 1.0, 0.127158, 1.0]}}}, 'type': 'DistributedRepeatingDataset'}, 'per_device_batch_size': 16, 'per_device_num_workers': 4}}, 'args': Namespace(config='configs/cosmos3/cosmos3edge_libero_spatial_full_finetune.py', work_dir='/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_spatial_edge_pr51_latest_bs16', cfg_options={'train_dataloader.per_device_batch_size': 16, 'runner.grad_accumulation_steps': 16, 'runner.metric.grad_accumulation_steps': 16}, eval_after_train=True, resume_from=None)}, 'seed': 7, 'train_dataloader': {'dataset': {'datasets': {'action_key': 'action', 'action_window_size': 16, 'data_root_path': '/mnt/data/cpfs/mnt/data/yanis/FastWAM/data/libero_mujoco3.3.2/libero_spatial_no_noops_lerobot', 'frame_window_size': 17, 'require_full_window': True, 'statistic_name': 'libero_spatial_no_noops', 'transforms': [{'embodiment_id': 5, 'name_mappings': {'actions': ['actions'], 'observation.state': ['states']}, 'parquet_keys': ['observation.state', 'timestamp', 'actions', 'info', 'stats'], 'type': 'ProcessParquetInputs', 'video_keys': ['observation.images.image', 'observation.images.wrist_image']}, {'type': 'LiberoFramewiseActionToRot6D'}, {'action_metadata': {'append_viewpoint': False, 'conditioning_fps': 20.0, 'frame_window_size': 17, 'video_height': 256, 'video_width': 512}, 'cfg_dropout_rate': 0.1, 'max_len': 512, 'tokenizer': {'model_max_length': 4096, 'model_path': './checkpoints/Cosmos3-Edge', 'padding_side': 'right', 'trust_remote_code': True, 'type': 'PretrainedTokenizer'}, 'type': 'ProcessCosmos3Prompt'}, {'height': 256, 'type': 'ResizeImages', 'width': 256}, {'type': 'SimpleNormalizeImages'}, {'action_dim': 64, 'action_key': 'action', 'action_norm_type': 'quantile', 'state_norm_type': 'none', 'state_dim': 64, 'state_key': 'proprio', 'type': 'NormalizeStatesAndActions'}, {'conditioning_fps': 20.0, 'frame_window_size': 17, 'mode': 'wam', 'prepend_state_to_action': False, 'raw_action_dim': 10, 'type': 'BuildCosmos3Sequence'}, {'frame_window_size': 17, 'num_views': 2, 'tile_direction': 'horizontal', 'type': 'PrepareVideo'}], 'type': 'ParquetDataset', 'use_delta': False, 'window_start_idx': 0}, 'name_mappings': {'action': ['action'], 'observation.state': ['proprio']}, 'statistic_keys': ['observation.state', 'timestamp', 'action'], 'statistic_name': 'libero_spatial_no_noops', 'statistics_overrides': {'libero_spatial_no_noops': {'action': {'mean': [0.050704, 0.097407, -0.094833, 0.994873, -0.004579, -0.004288, 0.004389, 0.996104, 0.001109, 0.476725], 'std': [0.333621, 0.387175, 0.45714, 0.010807, 0.077802, 0.063386, 0.078571, 0.009994, 0.038504, 0.49946], 'min': [-0.9375, -0.9375, -0.9375, 0.902028, -0.356085, -0.367416, -0.370434, 0.921907, -0.255, 0.0], 'max': [0.9375, 0.9375, 0.9375, 1.0, 0.368853, 0.341214, 0.356395, 1.0, 0.348251, 1.0], 'q01': [-0.723214, -0.808929, -0.9375, 0.934955, -0.223431, -0.189878, -0.334735, 0.938516, -0.107736, 0.0], 'q99': [0.9375, 0.870536, 0.9375, 1.0, 0.331, 0.163153, 0.226216, 1.0, 0.127158, 1.0]}}}, 'type': 'DistributedRepeatingDataset'}, 'per_device_batch_size': 16, 'per_device_num_workers': 4}}", "run_id": "cosmos3edge_libero_spatial_full_finetune_2026_08_04_01_21_37"}

cosmos3edge_libero_spatial_full_finetune_bs16/tokenizer/chat_template.jinja

@@ -0,0 +1,324 @@
+
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '
+<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'
+<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{% endmacro %}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+
+{%- if messages[0]["role"] == "system" %}
+    {%- if messages[0]["content"] is string %}
+        {%- set system_message = messages[0]["content"] %}
+    {%- else %}
+        {%- set ns = namespace(system_text="") %}
+        {%- for part in messages[0]["content"] %}
+            {%- if 'text' in part %}
+                {%- set ns.system_text = ns.system_text + part.text %}
+            {%- endif %}
+        {%- endfor %}
+        {%- set system_message = ns.system_text %}
+    {%- endif %}
+
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{# Recompute last_user_idx relative to loop_messages after handling system #}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system
+" + system_message }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system
+" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "
+
+" }}
+    {%- endif %}
+    {{- "# Tools
+
+You have access to the following functions:
+
+" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "
+<function>
+<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '
+<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '
+<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '
+<parameter>' }}
+                {{- '
+<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '
+<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '
+<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '
+<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '
+</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {% set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '
+<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '
+</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '
+</function>' }}
+    {%- endfor %}
+    {{- "
+</tools>" }}
+
+    {{- '
+
+If you choose to call a function ONLY reply in the following format with NO suffix:
+
+<tool_call>
+<function=example_function_name>
+<parameter=example_parameter_1>
+value_1
+</parameter>
+<parameter=example_parameter_2>
+This is the value for the second parameter
+that can span
+multiple lines
+</parameter>
+</function>
+</tool_call>
+
+<IMPORTANT>
+Reminder:
+- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
+- Required parameters MUST be specified
+- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
+- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
+</IMPORTANT>' }}
+{%- endif %}
+
+
+{%- if system_message is defined %}
+    {{- '<|im_end|>
+' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>
+' }}
+    {%- endif %}
+{%- endif %}
+
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {# Add reasoning content in to content field for unified processing below. #}
+        {%- set text_content = namespace(text_content="") -%}
+        {%- if message.content is string %}
+            {%- set text_content.text_content = message.content %}
+        {%- else %}
+            {%- for content in message.content %}
+                {%- if 'text' in content %}
+                    {%- set text_content.text_content = content.text %}
+                    {%- break %}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>
+" ~ message.reasoning_content ~ "
+</think>
+" ~ (text_content.text_content | default('', true)) %}
+        {%- else %}
+            {%- set content = text_content.text_content | default('', true) %}
+            {# Allow downstream logic to to take care of broken thought, only handle coherent reasoning here. #}
+            {%- if '<think>' not in content and '</think>' not in content %}
+                {%- set content = "<think></think>" ~ content %}
+            {%- else %}
+                {%- set content = content %}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {# Assistant message has tool calls. #}
+            {{- '<|im_start|>assistant
+' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '
+' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {# Keep only content after the last closing think. Also generation prompt causes this. #}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {# If <think> was opened but never closed, drop the trailing think segment #}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c | trim %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '
+' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>
+<function=' ~ tool_call.name ~ '>
+' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>
+' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '
+</parameter>
+' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>
+</tool_call>
+' -}}
+                {%- endfor %}
+                {{- '<|im_end|>
+' }}
+        {%- else %}
+            {# Assistant message doesn't have tool calls. #}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant
+' ~ (content | default('', true) | string | trim) ~ '<|im_end|>
+' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant
+' ~ c ~ '<|im_end|>
+' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant
+<|im_end|>
+' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '
+' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- else %}
+            {%- for content in message.content %}
+                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
+                    <|vision_start|><|image_pad|><|vision_end|>
+                {%- elif content.type == 'video' or 'video' in content %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
+                    <|vision_start|><|video_pad|><|vision_end|>
+                {%- elif 'text' in content %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>
+' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user
+' }}
+        {%- endif %}
+        {{- '<tool_response>
+' }}
+        {{- message.content }}
+        {{- '
+</tool_response>
+' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>
+' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>
+' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '
+' + message.content + '<|im_end|>
+' }}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {%- if enable_thinking %}
+        {{- '<|im_start|>assistant
+<think>
+' }}
+    {%- else %}
+        {{- '<|im_start|>assistant
+<think></think>' }}
+    {%- endif %}
+{%- endif %}

cosmos3edge_libero_spatial_full_finetune_bs16/tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

cosmos3edge_libero_spatial_full_finetune_bs16/tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4dc692a99dca6d92a44e9529ffb53541eb79bb260eec7281bc51268c05d7311c
+size 17077500

cosmos3edge_libero_spatial_full_finetune_bs16/tokenizer/tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 262144,
+  "pad_token": "<|im_end|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

cosmos3edge_libero_spatial_full_finetune_bs16/vlm_backbone_config.json

@@ -0,0 +1,41 @@
+{
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "enable_mrope": true,
+  "enable_rope": true,
+  "eos_token_id": 11,
+  "head_dim": 128,
+  "hidden_dropout": 0.0,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 9216,
+  "layer_norm_epsilon": 1e-05,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "mlp_hidden_act": "relu2",
+  "model_type": "nemotron_3_dense_vl_text",
+  "mrope_section": [
+    24,
+    20,
+    20
+  ],
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "num_logits_to_keep": 1,
+  "packed_attention_backend": "flash2",
+  "pad_token_id": 11,
+  "residual_in_fp32": false,
+  "rope_parameters": {
+    "rope_theta": 100000000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 100000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "use_cache": true,
+  "use_und_k_norm_for_gen": true,
+  "vocab_size": 131072
+}