masonlimx CCreater commited on
Commit
83175ee
·
1 Parent(s): 3722ac0

Add Cosmos3-Edge LIBERO checkpoints (#3)

Browse files

- Add Cosmos3-Edge LIBERO checkpoints (1a7c88e9d60e52aa820fb2e1d24f57a70c080d31)
- Add cosmos3edge_libero_spatial_full_finetune_bs16 checkpoint (47a2bbc43def50ac6166acff4f2e8c8b9fb81131)
- Add cosmos3edge_libero_object_full_finetune_bs16 checkpoint (c9ee920e863d7a64a2b85f6c082fcf26330e9c09)
- Add cosmos3edge_libero_goal_full_finetune_bs16 checkpoint (17591bd82a22033f3b1dcda473592144f9649fc6)
- Add cosmos3edge_libero_10_full_finetune_bs16 checkpoint (a469c9262304b24c15b3668d96537e0d64287495)
- Add Cosmos3-Edge mixed LIBERO checkpoint (8768ab246d27b8cbb2dd133c1df167715aa2059e)
- Remove superseded Cosmos3-Edge checkpoints (9d3c638a91fa464aa763aa0003c12ce0568ad62b)
- Resolve attributes conflict and update Cosmos3-Edge entry (e4e716e9be724992117ea0d873e11803509bc800)
- Resolve attributes conflict and update Cosmos3-Edge entry (a998cfcfb18212245970b57f6dd48eed8bc40fd4)
- Merge latest main into Cosmos3-Edge checkpoint PR (d0f118bcf6be9de3b42df18e761d202cea163ac8)
- Merge latest main and retain Edge and Nano LFS rules (49a1d4e24bcd1ce054c9e2a630894d321a9be200)


Co-authored-by: CG <CCreater@users.noreply.huggingface.co>

.gitattributes CHANGED
@@ -51,4 +51,5 @@ pi05_paligemma_robocasa_full_data_full_finetune_bs256/tokenizer/tokenizer.json f
51
  fastwam_base/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
52
  dit4dit_libero_all_full_finetune_bc071270ff2_bs256/dit4dit_libero_all_full_finetune_2026_08_12_20_51_03.jsonl filter=lfs diff=lfs merge=lfs -text
53
  dit4dit_libero_all_full_finetune_bc071270ff2_bs256/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 
54
  cosmos3_nano_libero_full_finetune_bs2048/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 
51
  fastwam_base/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
52
  dit4dit_libero_all_full_finetune_bc071270ff2_bs256/dit4dit_libero_all_full_finetune_2026_08_12_20_51_03.jsonl filter=lfs diff=lfs merge=lfs -text
53
  dit4dit_libero_all_full_finetune_bc071270ff2_bs256/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
54
+ cosmos3_edge_libero_full_finetune_bs2048/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
55
  cosmos3_nano_libero_full_finetune_bs2048/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -84,6 +84,7 @@ FluxVLA Engine supports multiple VLA model families, each optimized for differen
84
  | Model | Training Data | Performance | Download |
85
  |:---:|:---:|:---:|:---|
86
  | PI0.5 PaliGemma | LIBERO-10 | 96.0% avg | [Hugging Face](https://huggingface.co/limxdynamics/FluxVLAEngine/tree/main/pi05_paligemma_libero_10_full_finetune_bs64) |
 
87
  | GR00T Eagle 3B | LIBERO-10 | 89.4% avg | [Hugging Face](https://huggingface.co/limxdynamics/FluxVLAEngine/tree/main/gr00t_eagle_3b_libero_10_full_finetune_bs64) |
88
 
89
  ## 3. Evaluation Results
@@ -98,8 +99,9 @@ FluxVLA demonstrates state-of-the-art performance across multiple LIBERO benchma
98
  |:---:|:---:|:---:|:---:|:---:|:---:|
99
  | **FluxVLA (Pi0.5)** | 98.6 | 99.0 | 97.8 | 96.0±1.0 | 97.85 |
100
  | **FluxVLA (Qwen3VL+GR00T)** | 98.6 | 99.6 | 95.6 | 92.2±1.8 | 96.50 |
101
- | **FluxVLA (GR00T)** | 96.2 | 96.8 | 93.4 | 89.4±1.5 | 93.95 |
102
  | **FluxVLA (DreamZero)** | 96.8 | 97.4 | 90.8±1.5 | 93.6 | 94.65 |
 
 
103
 
104
  </div>
105
 
@@ -392,9 +394,9 @@ If you use FluxVLA in your research, please cite:
392
 
393
  ## 10. License
394
 
395
- This project is released under the **Apache-2.0 License**. See [LICENSE](LICENSE) for details.
396
 
397
- Both the code repository and the model weights are freely available for research and commercial use under the Apache-2.0 License.
398
 
399
  ## 11. Acknowledgements
400
 
@@ -430,4 +432,4 @@ You can also:
430
  - 🌐 GitHub: <https://github.com/limxdynamics/FluxVLA>
431
  - 🤗 Hugging Face: <https://huggingface.co/limxdynamics/FluxVLAEngine>
432
  - 📖 English Documentation: <https://fluxvla.limxdynamics.com>
433
- - 📖 中文文档: <https://fluxvla.limxdynamics.com/zh/>
 
84
  | Model | Training Data | Performance | Download |
85
  |:---:|:---:|:---:|:---|
86
  | PI0.5 PaliGemma | LIBERO-10 | 96.0% avg | [Hugging Face](https://huggingface.co/limxdynamics/FluxVLAEngine/tree/main/pi05_paligemma_libero_10_full_finetune_bs64) |
87
+ | Cosmos3-Edge | LIBERO benchmark suites (mixed training) | 94.4% avg | [Hugging Face](https://huggingface.co/limxdynamics/FluxVLAEngine/tree/main/cosmos3_edge_libero_full_finetune_bs2048) |
88
  | GR00T Eagle 3B | LIBERO-10 | 89.4% avg | [Hugging Face](https://huggingface.co/limxdynamics/FluxVLAEngine/tree/main/gr00t_eagle_3b_libero_10_full_finetune_bs64) |
89
 
90
  ## 3. Evaluation Results
 
99
  |:---:|:---:|:---:|:---:|:---:|:---:|
100
  | **FluxVLA (Pi0.5)** | 98.6 | 99.0 | 97.8 | 96.0±1.0 | 97.85 |
101
  | **FluxVLA (Qwen3VL+GR00T)** | 98.6 | 99.6 | 95.6 | 92.2±1.8 | 96.50 |
 
102
  | **FluxVLA (DreamZero)** | 96.8 | 97.4 | 90.8±1.5 | 93.6 | 94.65 |
103
+ | **FluxVLA (Cosmos3-Edge)** | 95.6 | 95.6 | 91.6 | 94.8 | 94.4 |
104
+ | **FluxVLA (GR00T)** | 96.2 | 96.8 | 93.4 | 89.4±1.5 | 93.95 |
105
 
106
  </div>
107
 
 
394
 
395
  ## 10. License
396
 
397
+ The FluxVLA code is released under the **Apache-2.0 License**. See [LICENSE](LICENSE) for details.
398
 
399
+ Model checkpoints derived from third-party base models remain subject to their respective licenses. The Cosmos3-Edge checkpoints are derived from NVIDIA Cosmos3-Edge and are subject to the [Open Model Development and Weight License 1.1 (OpenMDW 1.1)](https://openmdw.ai/license/1-1/).
400
 
401
  ## 11. Acknowledgements
402
 
 
432
  - 🌐 GitHub: <https://github.com/limxdynamics/FluxVLA>
433
  - 🤗 Hugging Face: <https://huggingface.co/limxdynamics/FluxVLAEngine>
434
  - 📖 English Documentation: <https://fluxvla.limxdynamics.com>
435
+ - 📖 中文文档: <https://fluxvla.limxdynamics.com/zh/>
cosmos3_edge_libero_full_finetune_bs2048/checkpoints/step-005000-epoch-22-loss=0.5737.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d5b028fd50784f5070a390de63f9df06187c0fc8a4ab8e415113a68a124e8a0e
3
+ size 16297490208
cosmos3_edge_libero_full_finetune_bs2048/config.json ADDED
@@ -0,0 +1,695 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "eval": {
3
+ "dataset": {
4
+ "extra_tensor_keys": [
5
+ "conditioning_fps",
6
+ "prepend_state_to_action"
7
+ ],
8
+ "img_buffer_len": 1,
9
+ "transforms": [
10
+ {
11
+ "embodiment_id": 5,
12
+ "img_keys": [
13
+ "agentview_image",
14
+ "robot0_eye_in_hand_image"
15
+ ],
16
+ "type": "ProcessLiberoEvalInputs"
17
+ },
18
+ {
19
+ "conditioning_fps": 20.0,
20
+ "prepend_state_to_action": false,
21
+ "type": "SetCosmos3ActionMetadata"
22
+ },
23
+ {
24
+ "image_resize_strategy": "resize-crop",
25
+ "input_sizes": [
26
+ [
27
+ 3,
28
+ 256,
29
+ 256
30
+ ],
31
+ [
32
+ 3,
33
+ 256,
34
+ 256
35
+ ]
36
+ ],
37
+ "means": [
38
+ [
39
+ 127.5,
40
+ 127.5,
41
+ 127.5
42
+ ],
43
+ [
44
+ 127.5,
45
+ 127.5,
46
+ 127.5
47
+ ]
48
+ ],
49
+ "stds": [
50
+ [
51
+ 127.5,
52
+ 127.5,
53
+ 127.5
54
+ ],
55
+ [
56
+ 127.5,
57
+ 127.5,
58
+ 127.5
59
+ ]
60
+ ],
61
+ "type": "TransformImage"
62
+ },
63
+ {
64
+ "action_metadata": {
65
+ "append_viewpoint": false,
66
+ "conditioning_fps": 20.0,
67
+ "frame_window_size": 17,
68
+ "video_height": 192,
69
+ "video_width": 320
70
+ },
71
+ "cfg_dropout_rate": 0.0,
72
+ "format_prompt_as_json": true,
73
+ "max_len": 512,
74
+ "output_attention_mask_key": "lang_masks",
75
+ "output_key": "lang_tokens",
76
+ "tokenizer": {
77
+ "model_max_length": 4096,
78
+ "model_path": "./checkpoints/Cosmos3-Edge",
79
+ "padding_side": "right",
80
+ "trust_remote_code": true,
81
+ "type": "PretrainedTokenizer"
82
+ },
83
+ "type": "ProcessCosmos3Prompt"
84
+ },
85
+ {
86
+ "gripper_key": "robot0_gripper_qpos",
87
+ "norm_type": "mean_std",
88
+ "out_key": "states",
89
+ "pos_key": "robot0_eef_pos",
90
+ "quat_key": "robot0_eef_quat",
91
+ "state_dim": 64,
92
+ "type": "LiberoProprioFromInputs"
93
+ },
94
+ {
95
+ "frame_window_size": 1,
96
+ "num_views": 2,
97
+ "tile_direction": "horizontal",
98
+ "type": "PrepareVideo"
99
+ },
100
+ {
101
+ "height": 192,
102
+ "type": "ResizeAndReflectPad",
103
+ "width": 320
104
+ }
105
+ ],
106
+ "type": "LiberoParquetEvalDataset"
107
+ },
108
+ "denormalize_action": {
109
+ "action_dim": 10,
110
+ "norm_type": "quantile_rot",
111
+ "type": "DenormalizeLiberoFramewiseRot6DAction"
112
+ },
113
+ "enable_mixed_precision_training": true,
114
+ "eval_chunk_size": 16,
115
+ "inference_seed": 7,
116
+ "mixed_precision_dtype": "bf16",
117
+ "model_family": "cosmos3",
118
+ "norm_stats_key": "all_libero_no_noops",
119
+ "num_inference_steps": 30,
120
+ "num_steps_wait": 10,
121
+ "num_trials_per_task": 50,
122
+ "resize_size": 256,
123
+ "run_id_suffix": "mixed-v21-shift10-after-train",
124
+ "seed": 7,
125
+ "task_ids": null,
126
+ "task_suite_name": [
127
+ "libero_spatial",
128
+ "libero_object",
129
+ "libero_goal",
130
+ "libero_10"
131
+ ],
132
+ "type": "LiberoEvalRunner"
133
+ },
134
+ "inference_model": {
135
+ "action_horizon": 16,
136
+ "action_in_proj": {
137
+ "input_size": 64,
138
+ "num_domains": 32,
139
+ "output_size": 2048,
140
+ "type": "DomainAwareLinear"
141
+ },
142
+ "action_out_proj": {
143
+ "input_size": 2048,
144
+ "num_domains": 32,
145
+ "output_size": 64,
146
+ "type": "DomainAwareLinear"
147
+ },
148
+ "base_fps": 24.0,
149
+ "enable_fps_modulation": true,
150
+ "enable_vision_loss": true,
151
+ "freeze_non_moe_vlm_backbone": true,
152
+ "freeze_vlm_backbone": false,
153
+ "latent_patch_size": 2,
154
+ "max_action_dim": 64,
155
+ "name_mapping": {
156
+ ".self_attn.k_norm_moe_gen.": ".self_attn.norm_added_k.",
157
+ ".self_attn.k_proj.": ".self_attn.to_k.",
158
+ ".self_attn.k_proj_moe_gen.": ".self_attn.add_k_proj.",
159
+ ".self_attn.o_proj.": ".self_attn.to_out.",
160
+ ".self_attn.o_proj_moe_gen.": ".self_attn.to_add_out.",
161
+ ".self_attn.q_norm_moe_gen.": ".self_attn.norm_added_q.",
162
+ ".self_attn.q_proj.": ".self_attn.to_q.",
163
+ ".self_attn.q_proj_moe_gen.": ".self_attn.add_q_proj.",
164
+ ".self_attn.v_proj.": ".self_attn.to_v.",
165
+ ".self_attn.v_proj_moe_gen.": ".self_attn.add_v_proj.",
166
+ "action_in_proj.": "action_proj_in.",
167
+ "action_modality_embed": "action_modality_embed",
168
+ "action_out_proj.": "action_proj_out.",
169
+ "time_embedder.mlp.0.": "time_embedder.linear_1.",
170
+ "time_embedder.mlp.2.": "time_embedder.linear_2.",
171
+ "vision_in_proj.projector.": "proj_in.",
172
+ "vision_out_proj.projector.": "proj_out.",
173
+ "vlm_backbone.lm_head.weight": "lm_head.weight",
174
+ "vlm_backbone.model.language_model.embed_tokens.weight": "embed_tokens.weight",
175
+ "vlm_backbone.model.language_model.layers.": "layers.",
176
+ "vlm_backbone.model.language_model.norm.weight": "norm.weight",
177
+ "vlm_backbone.model.language_model.norm_moe_gen.weight": "norm_moe_gen.weight"
178
+ },
179
+ "num_embodiment_domains": 32,
180
+ "ori_action_dim": 10,
181
+ "packed_attention_backend": "flash2",
182
+ "position_embedding_type": "unified_3d_mrope",
183
+ "pretrained_name_or_path": "./checkpoints/Cosmos3-Edge/transformer",
184
+ "rectified_flow_inference_config": {
185
+ "num_steps": 30,
186
+ "num_train_timesteps": 1000,
187
+ "scheduler_type": "unipc",
188
+ "shift": 10.0,
189
+ "use_dynamic_shifting": false,
190
+ "use_karras_sigmas": false
191
+ },
192
+ "rectified_flow_training_config": {
193
+ "action_loss_weight": 10.0,
194
+ "independent_action_schedule": false,
195
+ "normalize_loss_by_active": false,
196
+ "shift": {
197
+ "256": 3,
198
+ "480": 5,
199
+ "720": 10
200
+ },
201
+ "shift_action": null,
202
+ "train_time_action_distribution": "logitnormal",
203
+ "train_time_image_distribution": "logitnormal",
204
+ "train_time_video_distribution": "waver",
205
+ "train_time_weight": "uniform",
206
+ "use_discrete_rf": false,
207
+ "use_dynamic_shift": false,
208
+ "use_high_sigma_strategy": false,
209
+ "use_high_sigma_strategy_action": false,
210
+ "vision_loss_weight": 10.0
211
+ },
212
+ "reinitialize_action_policy": true,
213
+ "special_tokens": {
214
+ "end_of_generation": 21,
215
+ "eos_token_id": 11,
216
+ "start_of_generation": 20
217
+ },
218
+ "strict_mapping": true,
219
+ "timestep_scale": 0.001,
220
+ "type": "Cosmos3FlowMatching",
221
+ "unified_3d_mrope_reset_spatial_ids": true,
222
+ "unified_3d_mrope_temporal_modality_margin": 15000,
223
+ "vision_in_proj": {
224
+ "in_dim": 192,
225
+ "out_dim": 2048,
226
+ "type": "LinearProjector"
227
+ },
228
+ "vision_latent_dim": 48,
229
+ "vision_out_proj": {
230
+ "in_dim": 2048,
231
+ "out_dim": 192,
232
+ "type": "LinearProjector"
233
+ },
234
+ "vision_vae": {
235
+ "encode_exact_durations": [
236
+ 17
237
+ ],
238
+ "pretrained_name_or_path": "./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth",
239
+ "type": "Cosmos3Wan22VAE"
240
+ },
241
+ "vlm_backbone": {
242
+ "include_visual": false,
243
+ "skip_init_weights": true,
244
+ "type": "Cosmos3MoTBackbone",
245
+ "vlm_config": {
246
+ "attention_bias": false,
247
+ "bos_token_id": 1,
248
+ "enable_mrope": true,
249
+ "eos_token_id": 11,
250
+ "head_dim": 128,
251
+ "hidden_size": 2048,
252
+ "intermediate_size": 9216,
253
+ "layer_norm_epsilon": 1e-05,
254
+ "max_position_embeddings": 131072,
255
+ "mlp_bias": false,
256
+ "mlp_hidden_act": "relu2",
257
+ "model_type": "nemotron_3_dense_vl_text",
258
+ "mrope_section": [
259
+ 24,
260
+ 20,
261
+ 20
262
+ ],
263
+ "num_attention_heads": 16,
264
+ "num_hidden_layers": 28,
265
+ "num_key_value_heads": 8,
266
+ "pad_token_id": 11,
267
+ "rope_theta": 100000000.0,
268
+ "tie_word_embeddings": false,
269
+ "use_und_k_norm_for_gen": true,
270
+ "vocab_size": 131072
271
+ }
272
+ }
273
+ },
274
+ "model": {
275
+ "action_horizon": 16,
276
+ "action_in_proj": {
277
+ "input_size": 64,
278
+ "num_domains": 32,
279
+ "output_size": 2048,
280
+ "type": "DomainAwareLinear"
281
+ },
282
+ "action_out_proj": {
283
+ "input_size": 2048,
284
+ "num_domains": 32,
285
+ "output_size": 64,
286
+ "type": "DomainAwareLinear"
287
+ },
288
+ "base_fps": 24.0,
289
+ "enable_fps_modulation": true,
290
+ "enable_vision_loss": true,
291
+ "freeze_non_moe_vlm_backbone": true,
292
+ "freeze_vlm_backbone": false,
293
+ "latent_patch_size": 2,
294
+ "max_action_dim": 64,
295
+ "name_mapping": {
296
+ ".self_attn.k_norm_moe_gen.": ".self_attn.norm_added_k.",
297
+ ".self_attn.k_proj.": ".self_attn.to_k.",
298
+ ".self_attn.k_proj_moe_gen.": ".self_attn.add_k_proj.",
299
+ ".self_attn.o_proj.": ".self_attn.to_out.",
300
+ ".self_attn.o_proj_moe_gen.": ".self_attn.to_add_out.",
301
+ ".self_attn.q_norm_moe_gen.": ".self_attn.norm_added_q.",
302
+ ".self_attn.q_proj.": ".self_attn.to_q.",
303
+ ".self_attn.q_proj_moe_gen.": ".self_attn.add_q_proj.",
304
+ ".self_attn.v_proj.": ".self_attn.to_v.",
305
+ ".self_attn.v_proj_moe_gen.": ".self_attn.add_v_proj.",
306
+ "action_in_proj.": "action_proj_in.",
307
+ "action_modality_embed": "action_modality_embed",
308
+ "action_out_proj.": "action_proj_out.",
309
+ "time_embedder.mlp.0.": "time_embedder.linear_1.",
310
+ "time_embedder.mlp.2.": "time_embedder.linear_2.",
311
+ "vision_in_proj.projector.": "proj_in.",
312
+ "vision_out_proj.projector.": "proj_out.",
313
+ "vlm_backbone.lm_head.weight": "lm_head.weight",
314
+ "vlm_backbone.model.language_model.embed_tokens.weight": "embed_tokens.weight",
315
+ "vlm_backbone.model.language_model.layers.": "layers.",
316
+ "vlm_backbone.model.language_model.norm.weight": "norm.weight",
317
+ "vlm_backbone.model.language_model.norm_moe_gen.weight": "norm_moe_gen.weight"
318
+ },
319
+ "num_embodiment_domains": 32,
320
+ "ori_action_dim": 10,
321
+ "packed_attention_backend": "flash2",
322
+ "position_embedding_type": "unified_3d_mrope",
323
+ "pretrained_name_or_path": "./checkpoints/Cosmos3-Edge/transformer",
324
+ "rectified_flow_inference_config": {
325
+ "num_steps": 30,
326
+ "num_train_timesteps": 1000,
327
+ "scheduler_type": "unipc",
328
+ "shift": 10.0,
329
+ "use_dynamic_shifting": false,
330
+ "use_karras_sigmas": false
331
+ },
332
+ "rectified_flow_training_config": {
333
+ "action_loss_weight": 10.0,
334
+ "independent_action_schedule": false,
335
+ "normalize_loss_by_active": false,
336
+ "shift": {
337
+ "256": 3,
338
+ "480": 5,
339
+ "720": 10
340
+ },
341
+ "shift_action": null,
342
+ "train_time_action_distribution": "logitnormal",
343
+ "train_time_image_distribution": "logitnormal",
344
+ "train_time_video_distribution": "waver",
345
+ "train_time_weight": "uniform",
346
+ "use_discrete_rf": false,
347
+ "use_dynamic_shift": false,
348
+ "use_high_sigma_strategy": false,
349
+ "use_high_sigma_strategy_action": false,
350
+ "vision_loss_weight": 10.0
351
+ },
352
+ "reinitialize_action_policy": true,
353
+ "special_tokens": {
354
+ "end_of_generation": 21,
355
+ "eos_token_id": 11,
356
+ "start_of_generation": 20
357
+ },
358
+ "strict_mapping": true,
359
+ "timestep_scale": 0.001,
360
+ "type": "Cosmos3FlowMatching",
361
+ "unified_3d_mrope_reset_spatial_ids": true,
362
+ "unified_3d_mrope_temporal_modality_margin": 15000,
363
+ "vision_in_proj": {
364
+ "in_dim": 192,
365
+ "out_dim": 2048,
366
+ "type": "LinearProjector"
367
+ },
368
+ "vision_latent_dim": 48,
369
+ "vision_out_proj": {
370
+ "in_dim": 2048,
371
+ "out_dim": 192,
372
+ "type": "LinearProjector"
373
+ },
374
+ "vision_vae": {
375
+ "encode_exact_durations": [
376
+ 17
377
+ ],
378
+ "pretrained_name_or_path": "./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth",
379
+ "type": "Cosmos3Wan22VAE"
380
+ },
381
+ "vlm_backbone": {
382
+ "include_visual": false,
383
+ "skip_init_weights": true,
384
+ "type": "Cosmos3MoTBackbone",
385
+ "vlm_config": {
386
+ "attention_bias": false,
387
+ "bos_token_id": 1,
388
+ "enable_mrope": true,
389
+ "eos_token_id": 11,
390
+ "head_dim": 128,
391
+ "hidden_size": 2048,
392
+ "intermediate_size": 9216,
393
+ "layer_norm_epsilon": 1e-05,
394
+ "max_position_embeddings": 131072,
395
+ "mlp_bias": false,
396
+ "mlp_hidden_act": "relu2",
397
+ "model_type": "nemotron_3_dense_vl_text",
398
+ "mrope_section": [
399
+ 24,
400
+ 20,
401
+ 20
402
+ ],
403
+ "num_attention_heads": 16,
404
+ "num_hidden_layers": 28,
405
+ "num_key_value_heads": 8,
406
+ "pad_token_id": 11,
407
+ "rope_theta": 100000000.0,
408
+ "tie_word_embeddings": false,
409
+ "use_und_k_norm_for_gen": true,
410
+ "vocab_size": 131072
411
+ }
412
+ }
413
+ },
414
+ "runner": {
415
+ "change_key_name": false,
416
+ "collator": {
417
+ "list_keys": [
418
+ "sequence_plan"
419
+ ],
420
+ "meta_keys": [
421
+ "task_description",
422
+ "stats",
423
+ "info",
424
+ "timestamp"
425
+ ],
426
+ "pad_id": 11,
427
+ "sequence_keys": [
428
+ "text_token_ids"
429
+ ],
430
+ "tensor_keys": [
431
+ "images",
432
+ "actions",
433
+ "embodiment_ids",
434
+ "raw_action_dim",
435
+ "conditioning_fps",
436
+ "action_fps"
437
+ ],
438
+ "type": "Cosmos3Collator"
439
+ },
440
+ "enable_gradient_checkpointing": true,
441
+ "enable_mixed_precision_training": true,
442
+ "grad_accumulation_steps": 4,
443
+ "lr_scheduler": {
444
+ "cycle_length": 16000,
445
+ "type": "linear-warmup+linear-decay",
446
+ "warmup_steps": 500
447
+ },
448
+ "max_epochs": null,
449
+ "max_grad_norm": 1.0,
450
+ "max_keep_ckpts": 2,
451
+ "max_steps": 5000,
452
+ "metric": {
453
+ "active_trackers": [
454
+ "jsonl",
455
+ "wandb"
456
+ ],
457
+ "grad_accumulation_steps": 4,
458
+ "run_dir": "work_dirs",
459
+ "type": "VLAMetric",
460
+ "window_size": 1
461
+ },
462
+ "mixed_precision_dtype": "bf16",
463
+ "optimizer": {
464
+ "betas": [
465
+ 0.9,
466
+ 0.99
467
+ ],
468
+ "eps": 1e-08,
469
+ "exclude_1d_from_weight_decay": false,
470
+ "fused": true,
471
+ "lr": 5e-05,
472
+ "paramwise_learning_rate": {
473
+ "action_in_proj.": 0.00025,
474
+ "action_modality_embed": 0.00025,
475
+ "action_out_proj.": 0.00025
476
+ },
477
+ "type": "AdamW",
478
+ "weight_decay": 0.05
479
+ },
480
+ "sampler": null,
481
+ "save_iter_interval": 500,
482
+ "sharding_strategy": "full-shard",
483
+ "tokenizer": {
484
+ "model_max_length": 4096,
485
+ "model_path": "./checkpoints/Cosmos3-Edge",
486
+ "padding_side": "right",
487
+ "trust_remote_code": true,
488
+ "type": "PretrainedTokenizer"
489
+ },
490
+ "type": "FSDPTrainRunner"
491
+ },
492
+ "seed": 7,
493
+ "train_dataloader": {
494
+ "dataset": {
495
+ "datasets": {
496
+ "action_key": "action",
497
+ "action_window_size": 16,
498
+ "data_root_path": [
499
+ "./datasets/libero_spatial_no_noops_lerobotv2.1",
500
+ "./datasets/libero_spatial_no_noops_lerobotv2.1",
501
+ "./datasets/libero_object_no_noops_lerobotv2.1",
502
+ "./datasets/libero_object_no_noops_lerobotv2.1",
503
+ "./datasets/libero_goal_no_noops_lerobotv2.1",
504
+ "./datasets/libero_goal_no_noops_lerobotv2.1",
505
+ "./datasets/libero_10_no_noops_lerobotv2.1"
506
+ ],
507
+ "frame_window_size": 17,
508
+ "require_full_window": true,
509
+ "statistic_name": "all_libero_no_noops",
510
+ "transforms": [
511
+ {
512
+ "embodiment_id": 5,
513
+ "name_mappings": {
514
+ "actions": [
515
+ "actions"
516
+ ],
517
+ "observation.state": [
518
+ "states"
519
+ ]
520
+ },
521
+ "parquet_keys": [
522
+ "observation.state",
523
+ "timestamp",
524
+ "actions",
525
+ "info",
526
+ "stats"
527
+ ],
528
+ "type": "ProcessParquetInputs",
529
+ "video_keys": [
530
+ "observation.images.image",
531
+ "observation.images.wrist_image"
532
+ ]
533
+ },
534
+ {
535
+ "type": "LiberoFramewiseActionToRot6D"
536
+ },
537
+ {
538
+ "action_metadata": {
539
+ "append_viewpoint": false,
540
+ "conditioning_fps": 20.0,
541
+ "frame_window_size": 17,
542
+ "video_height": 192,
543
+ "video_width": 320
544
+ },
545
+ "cfg_dropout_rate": 0.1,
546
+ "format_prompt_as_json": true,
547
+ "max_len": 512,
548
+ "tokenizer": {
549
+ "model_max_length": 4096,
550
+ "model_path": "./checkpoints/Cosmos3-Edge",
551
+ "padding_side": "right",
552
+ "trust_remote_code": true,
553
+ "type": "PretrainedTokenizer"
554
+ },
555
+ "type": "ProcessCosmos3Prompt"
556
+ },
557
+ {
558
+ "height": 256,
559
+ "type": "ResizeImages",
560
+ "width": 256
561
+ },
562
+ {
563
+ "type": "SimpleNormalizeImages"
564
+ },
565
+ {
566
+ "action_dim": 64,
567
+ "action_key": "action",
568
+ "action_norm_type": "quantile",
569
+ "state_dim": 64,
570
+ "state_key": "proprio",
571
+ "state_norm_type": "none",
572
+ "type": "NormalizeStatesAndActions"
573
+ },
574
+ {
575
+ "conditioning_fps": 20.0,
576
+ "frame_window_size": 17,
577
+ "mode": "wam",
578
+ "prepend_state_to_action": false,
579
+ "raw_action_dim": 10,
580
+ "type": "BuildCosmos3Sequence"
581
+ },
582
+ {
583
+ "frame_window_size": 17,
584
+ "num_views": 2,
585
+ "tile_direction": "horizontal",
586
+ "type": "PrepareVideo"
587
+ },
588
+ {
589
+ "height": 192,
590
+ "type": "ResizeAndReflectPad",
591
+ "width": 320
592
+ }
593
+ ],
594
+ "type": "ParquetDataset",
595
+ "use_delta": false,
596
+ "window_start_idx": 0
597
+ },
598
+ "name_mappings": {
599
+ "action": [
600
+ "action"
601
+ ],
602
+ "observation.state": [
603
+ "proprio"
604
+ ]
605
+ },
606
+ "statistic_keys": [
607
+ "observation.state",
608
+ "timestamp",
609
+ "action"
610
+ ],
611
+ "statistic_name": "all_libero_no_noops",
612
+ "statistics_overrides": {
613
+ "all_libero_no_noops": {
614
+ "action": {
615
+ "max": [
616
+ 0.9375,
617
+ 0.9375,
618
+ 0.9375,
619
+ 1.0,
620
+ 0.368853,
621
+ 0.341214,
622
+ 0.356395,
623
+ 1.0,
624
+ 0.348251,
625
+ 1.0
626
+ ],
627
+ "mean": [
628
+ 0.050704,
629
+ 0.097407,
630
+ -0.094833,
631
+ 0.994873,
632
+ -0.004579,
633
+ -0.004288,
634
+ 0.004389,
635
+ 0.996104,
636
+ 0.001109,
637
+ 0.476725
638
+ ],
639
+ "min": [
640
+ -0.9375,
641
+ -0.9375,
642
+ -0.9375,
643
+ 0.902028,
644
+ -0.356085,
645
+ -0.367416,
646
+ -0.370434,
647
+ 0.921907,
648
+ -0.255,
649
+ 0.0
650
+ ],
651
+ "q01": [
652
+ -0.723214,
653
+ -0.808929,
654
+ -0.9375,
655
+ 0.934955,
656
+ -0.223431,
657
+ -0.189878,
658
+ -0.334735,
659
+ 0.938516,
660
+ -0.107736,
661
+ 0.0
662
+ ],
663
+ "q99": [
664
+ 0.9375,
665
+ 0.870536,
666
+ 0.9375,
667
+ 1.0,
668
+ 0.331,
669
+ 0.163153,
670
+ 0.226216,
671
+ 1.0,
672
+ 0.127158,
673
+ 1.0
674
+ ],
675
+ "std": [
676
+ 0.333621,
677
+ 0.387175,
678
+ 0.45714,
679
+ 0.010807,
680
+ 0.077802,
681
+ 0.063386,
682
+ 0.078571,
683
+ 0.009994,
684
+ 0.038504,
685
+ 0.49946
686
+ ]
687
+ }
688
+ }
689
+ },
690
+ "type": "DistributedRepeatingDataset"
691
+ },
692
+ "per_device_batch_size": 32,
693
+ "per_device_num_workers": 4
694
+ }
695
+ }
cosmos3_edge_libero_full_finetune_bs2048/config.yaml ADDED
@@ -0,0 +1,569 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ eval:
2
+ dataset:
3
+ extra_tensor_keys:
4
+ - conditioning_fps
5
+ - prepend_state_to_action
6
+ img_buffer_len: 1
7
+ transforms:
8
+ - embodiment_id: 5
9
+ img_keys:
10
+ - agentview_image
11
+ - robot0_eye_in_hand_image
12
+ type: ProcessLiberoEvalInputs
13
+ - conditioning_fps: 20.0
14
+ prepend_state_to_action: false
15
+ type: SetCosmos3ActionMetadata
16
+ - image_resize_strategy: resize-crop
17
+ input_sizes:
18
+ - - 3
19
+ - 256
20
+ - 256
21
+ - - 3
22
+ - 256
23
+ - 256
24
+ means:
25
+ - - 127.5
26
+ - 127.5
27
+ - 127.5
28
+ - - 127.5
29
+ - 127.5
30
+ - 127.5
31
+ stds:
32
+ - - 127.5
33
+ - 127.5
34
+ - 127.5
35
+ - - 127.5
36
+ - 127.5
37
+ - 127.5
38
+ type: TransformImage
39
+ - action_metadata:
40
+ append_viewpoint: false
41
+ conditioning_fps: 20.0
42
+ frame_window_size: 17
43
+ video_height: 192
44
+ video_width: 320
45
+ cfg_dropout_rate: 0.0
46
+ format_prompt_as_json: true
47
+ max_len: 512
48
+ output_attention_mask_key: lang_masks
49
+ output_key: lang_tokens
50
+ tokenizer:
51
+ model_max_length: 4096
52
+ model_path: ./checkpoints/Cosmos3-Edge
53
+ padding_side: right
54
+ trust_remote_code: true
55
+ type: PretrainedTokenizer
56
+ type: ProcessCosmos3Prompt
57
+ - gripper_key: robot0_gripper_qpos
58
+ norm_type: mean_std
59
+ out_key: states
60
+ pos_key: robot0_eef_pos
61
+ quat_key: robot0_eef_quat
62
+ state_dim: 64
63
+ type: LiberoProprioFromInputs
64
+ - frame_window_size: 1
65
+ num_views: 2
66
+ tile_direction: horizontal
67
+ type: PrepareVideo
68
+ - height: 192
69
+ type: ResizeAndReflectPad
70
+ width: 320
71
+ type: LiberoParquetEvalDataset
72
+ denormalize_action:
73
+ action_dim: 10
74
+ norm_type: quantile_rot
75
+ type: DenormalizeLiberoFramewiseRot6DAction
76
+ enable_mixed_precision_training: true
77
+ eval_chunk_size: 16
78
+ inference_seed: 7
79
+ mixed_precision_dtype: bf16
80
+ model_family: cosmos3
81
+ norm_stats_key: all_libero_no_noops
82
+ num_inference_steps: 30
83
+ num_steps_wait: 10
84
+ num_trials_per_task: 50
85
+ resize_size: 256
86
+ run_id_suffix: mixed-v21-shift10-after-train
87
+ seed: 7
88
+ task_ids: null
89
+ task_suite_name:
90
+ - libero_spatial
91
+ - libero_object
92
+ - libero_goal
93
+ - libero_10
94
+ type: LiberoEvalRunner
95
+ inference_model:
96
+ action_horizon: 16
97
+ action_in_proj:
98
+ input_size: 64
99
+ num_domains: 32
100
+ output_size: 2048
101
+ type: DomainAwareLinear
102
+ action_out_proj:
103
+ input_size: 2048
104
+ num_domains: 32
105
+ output_size: 64
106
+ type: DomainAwareLinear
107
+ base_fps: 24.0
108
+ enable_fps_modulation: true
109
+ enable_vision_loss: true
110
+ freeze_non_moe_vlm_backbone: true
111
+ freeze_vlm_backbone: false
112
+ latent_patch_size: 2
113
+ max_action_dim: 64
114
+ name_mapping:
115
+ .self_attn.k_norm_moe_gen.: .self_attn.norm_added_k.
116
+ .self_attn.k_proj.: .self_attn.to_k.
117
+ .self_attn.k_proj_moe_gen.: .self_attn.add_k_proj.
118
+ .self_attn.o_proj.: .self_attn.to_out.
119
+ .self_attn.o_proj_moe_gen.: .self_attn.to_add_out.
120
+ .self_attn.q_norm_moe_gen.: .self_attn.norm_added_q.
121
+ .self_attn.q_proj.: .self_attn.to_q.
122
+ .self_attn.q_proj_moe_gen.: .self_attn.add_q_proj.
123
+ .self_attn.v_proj.: .self_attn.to_v.
124
+ .self_attn.v_proj_moe_gen.: .self_attn.add_v_proj.
125
+ action_in_proj.: action_proj_in.
126
+ action_modality_embed: action_modality_embed
127
+ action_out_proj.: action_proj_out.
128
+ time_embedder.mlp.0.: time_embedder.linear_1.
129
+ time_embedder.mlp.2.: time_embedder.linear_2.
130
+ vision_in_proj.projector.: proj_in.
131
+ vision_out_proj.projector.: proj_out.
132
+ vlm_backbone.lm_head.weight: lm_head.weight
133
+ vlm_backbone.model.language_model.embed_tokens.weight: embed_tokens.weight
134
+ vlm_backbone.model.language_model.layers.: layers.
135
+ vlm_backbone.model.language_model.norm.weight: norm.weight
136
+ vlm_backbone.model.language_model.norm_moe_gen.weight: norm_moe_gen.weight
137
+ num_embodiment_domains: 32
138
+ ori_action_dim: 10
139
+ packed_attention_backend: flash2
140
+ position_embedding_type: unified_3d_mrope
141
+ pretrained_name_or_path: ./checkpoints/Cosmos3-Edge/transformer
142
+ rectified_flow_inference_config:
143
+ num_steps: 30
144
+ num_train_timesteps: 1000
145
+ scheduler_type: unipc
146
+ shift: 10.0
147
+ use_dynamic_shifting: false
148
+ use_karras_sigmas: false
149
+ rectified_flow_training_config:
150
+ action_loss_weight: 10.0
151
+ independent_action_schedule: false
152
+ normalize_loss_by_active: false
153
+ shift:
154
+ '256': 3
155
+ '480': 5
156
+ '720': 10
157
+ shift_action: null
158
+ train_time_action_distribution: logitnormal
159
+ train_time_image_distribution: logitnormal
160
+ train_time_video_distribution: waver
161
+ train_time_weight: uniform
162
+ use_discrete_rf: false
163
+ use_dynamic_shift: false
164
+ use_high_sigma_strategy: false
165
+ use_high_sigma_strategy_action: false
166
+ vision_loss_weight: 10.0
167
+ reinitialize_action_policy: true
168
+ special_tokens:
169
+ end_of_generation: 21
170
+ eos_token_id: 11
171
+ start_of_generation: 20
172
+ strict_mapping: true
173
+ timestep_scale: 0.001
174
+ type: Cosmos3FlowMatching
175
+ unified_3d_mrope_reset_spatial_ids: true
176
+ unified_3d_mrope_temporal_modality_margin: 15000
177
+ vision_in_proj:
178
+ in_dim: 192
179
+ out_dim: 2048
180
+ type: LinearProjector
181
+ vision_latent_dim: 48
182
+ vision_out_proj:
183
+ in_dim: 2048
184
+ out_dim: 192
185
+ type: LinearProjector
186
+ vision_vae:
187
+ encode_exact_durations:
188
+ - 17
189
+ pretrained_name_or_path: ./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
190
+ type: Cosmos3Wan22VAE
191
+ vlm_backbone:
192
+ include_visual: false
193
+ skip_init_weights: true
194
+ type: Cosmos3MoTBackbone
195
+ vlm_config:
196
+ attention_bias: false
197
+ bos_token_id: 1
198
+ enable_mrope: true
199
+ eos_token_id: 11
200
+ head_dim: 128
201
+ hidden_size: 2048
202
+ intermediate_size: 9216
203
+ layer_norm_epsilon: 1.0e-05
204
+ max_position_embeddings: 131072
205
+ mlp_bias: false
206
+ mlp_hidden_act: relu2
207
+ model_type: nemotron_3_dense_vl_text
208
+ mrope_section:
209
+ - 24
210
+ - 20
211
+ - 20
212
+ num_attention_heads: 16
213
+ num_hidden_layers: 28
214
+ num_key_value_heads: 8
215
+ pad_token_id: 11
216
+ rope_theta: 100000000.0
217
+ tie_word_embeddings: false
218
+ use_und_k_norm_for_gen: true
219
+ vocab_size: 131072
220
+ model:
221
+ action_horizon: 16
222
+ action_in_proj:
223
+ input_size: 64
224
+ num_domains: 32
225
+ output_size: 2048
226
+ type: DomainAwareLinear
227
+ action_out_proj:
228
+ input_size: 2048
229
+ num_domains: 32
230
+ output_size: 64
231
+ type: DomainAwareLinear
232
+ base_fps: 24.0
233
+ enable_fps_modulation: true
234
+ enable_vision_loss: true
235
+ freeze_non_moe_vlm_backbone: true
236
+ freeze_vlm_backbone: false
237
+ latent_patch_size: 2
238
+ max_action_dim: 64
239
+ name_mapping:
240
+ .self_attn.k_norm_moe_gen.: .self_attn.norm_added_k.
241
+ .self_attn.k_proj.: .self_attn.to_k.
242
+ .self_attn.k_proj_moe_gen.: .self_attn.add_k_proj.
243
+ .self_attn.o_proj.: .self_attn.to_out.
244
+ .self_attn.o_proj_moe_gen.: .self_attn.to_add_out.
245
+ .self_attn.q_norm_moe_gen.: .self_attn.norm_added_q.
246
+ .self_attn.q_proj.: .self_attn.to_q.
247
+ .self_attn.q_proj_moe_gen.: .self_attn.add_q_proj.
248
+ .self_attn.v_proj.: .self_attn.to_v.
249
+ .self_attn.v_proj_moe_gen.: .self_attn.add_v_proj.
250
+ action_in_proj.: action_proj_in.
251
+ action_modality_embed: action_modality_embed
252
+ action_out_proj.: action_proj_out.
253
+ time_embedder.mlp.0.: time_embedder.linear_1.
254
+ time_embedder.mlp.2.: time_embedder.linear_2.
255
+ vision_in_proj.projector.: proj_in.
256
+ vision_out_proj.projector.: proj_out.
257
+ vlm_backbone.lm_head.weight: lm_head.weight
258
+ vlm_backbone.model.language_model.embed_tokens.weight: embed_tokens.weight
259
+ vlm_backbone.model.language_model.layers.: layers.
260
+ vlm_backbone.model.language_model.norm.weight: norm.weight
261
+ vlm_backbone.model.language_model.norm_moe_gen.weight: norm_moe_gen.weight
262
+ num_embodiment_domains: 32
263
+ ori_action_dim: 10
264
+ packed_attention_backend: flash2
265
+ position_embedding_type: unified_3d_mrope
266
+ pretrained_name_or_path: ./checkpoints/Cosmos3-Edge/transformer
267
+ rectified_flow_inference_config:
268
+ num_steps: 30
269
+ num_train_timesteps: 1000
270
+ scheduler_type: unipc
271
+ shift: 10.0
272
+ use_dynamic_shifting: false
273
+ use_karras_sigmas: false
274
+ rectified_flow_training_config:
275
+ action_loss_weight: 10.0
276
+ independent_action_schedule: false
277
+ normalize_loss_by_active: false
278
+ shift:
279
+ '256': 3
280
+ '480': 5
281
+ '720': 10
282
+ shift_action: null
283
+ train_time_action_distribution: logitnormal
284
+ train_time_image_distribution: logitnormal
285
+ train_time_video_distribution: waver
286
+ train_time_weight: uniform
287
+ use_discrete_rf: false
288
+ use_dynamic_shift: false
289
+ use_high_sigma_strategy: false
290
+ use_high_sigma_strategy_action: false
291
+ vision_loss_weight: 10.0
292
+ reinitialize_action_policy: true
293
+ special_tokens:
294
+ end_of_generation: 21
295
+ eos_token_id: 11
296
+ start_of_generation: 20
297
+ strict_mapping: true
298
+ timestep_scale: 0.001
299
+ type: Cosmos3FlowMatching
300
+ unified_3d_mrope_reset_spatial_ids: true
301
+ unified_3d_mrope_temporal_modality_margin: 15000
302
+ vision_in_proj:
303
+ in_dim: 192
304
+ out_dim: 2048
305
+ type: LinearProjector
306
+ vision_latent_dim: 48
307
+ vision_out_proj:
308
+ in_dim: 2048
309
+ out_dim: 192
310
+ type: LinearProjector
311
+ vision_vae:
312
+ encode_exact_durations:
313
+ - 17
314
+ pretrained_name_or_path: ./checkpoints/Wan2.2-TI2V-5B/Wan2.2_VAE.pth
315
+ type: Cosmos3Wan22VAE
316
+ vlm_backbone:
317
+ include_visual: false
318
+ skip_init_weights: true
319
+ type: Cosmos3MoTBackbone
320
+ vlm_config:
321
+ attention_bias: false
322
+ bos_token_id: 1
323
+ enable_mrope: true
324
+ eos_token_id: 11
325
+ head_dim: 128
326
+ hidden_size: 2048
327
+ intermediate_size: 9216
328
+ layer_norm_epsilon: 1.0e-05
329
+ max_position_embeddings: 131072
330
+ mlp_bias: false
331
+ mlp_hidden_act: relu2
332
+ model_type: nemotron_3_dense_vl_text
333
+ mrope_section:
334
+ - 24
335
+ - 20
336
+ - 20
337
+ num_attention_heads: 16
338
+ num_hidden_layers: 28
339
+ num_key_value_heads: 8
340
+ pad_token_id: 11
341
+ rope_theta: 100000000.0
342
+ tie_word_embeddings: false
343
+ use_und_k_norm_for_gen: true
344
+ vocab_size: 131072
345
+ runner:
346
+ change_key_name: false
347
+ collator:
348
+ list_keys:
349
+ - sequence_plan
350
+ meta_keys:
351
+ - task_description
352
+ - stats
353
+ - info
354
+ - timestamp
355
+ pad_id: 11
356
+ sequence_keys:
357
+ - text_token_ids
358
+ tensor_keys:
359
+ - images
360
+ - actions
361
+ - embodiment_ids
362
+ - raw_action_dim
363
+ - conditioning_fps
364
+ - action_fps
365
+ type: Cosmos3Collator
366
+ enable_gradient_checkpointing: true
367
+ enable_mixed_precision_training: true
368
+ grad_accumulation_steps: 4
369
+ lr_scheduler:
370
+ cycle_length: 16000
371
+ type: linear-warmup+linear-decay
372
+ warmup_steps: 500
373
+ max_epochs: null
374
+ max_grad_norm: 1.0
375
+ max_keep_ckpts: 2
376
+ max_steps: 5000
377
+ metric:
378
+ active_trackers:
379
+ - jsonl
380
+ - wandb
381
+ grad_accumulation_steps: 4
382
+ run_dir: work_dirs
383
+ type: VLAMetric
384
+ window_size: 1
385
+ mixed_precision_dtype: bf16
386
+ optimizer:
387
+ betas:
388
+ - 0.9
389
+ - 0.99
390
+ eps: 1.0e-08
391
+ exclude_1d_from_weight_decay: false
392
+ fused: true
393
+ lr: 5.0e-05
394
+ paramwise_learning_rate:
395
+ action_in_proj.: 0.00025
396
+ action_modality_embed: 0.00025
397
+ action_out_proj.: 0.00025
398
+ type: AdamW
399
+ weight_decay: 0.05
400
+ sampler: null
401
+ save_iter_interval: 500
402
+ sharding_strategy: full-shard
403
+ tokenizer:
404
+ model_max_length: 4096
405
+ model_path: ./checkpoints/Cosmos3-Edge
406
+ padding_side: right
407
+ trust_remote_code: true
408
+ type: PretrainedTokenizer
409
+ type: FSDPTrainRunner
410
+ seed: 7
411
+ train_dataloader:
412
+ dataset:
413
+ datasets:
414
+ action_key: action
415
+ action_window_size: 16
416
+ data_root_path:
417
+ - ./datasets/libero_spatial_no_noops_lerobotv2.1
418
+ - ./datasets/libero_spatial_no_noops_lerobotv2.1
419
+ - ./datasets/libero_object_no_noops_lerobotv2.1
420
+ - ./datasets/libero_object_no_noops_lerobotv2.1
421
+ - ./datasets/libero_goal_no_noops_lerobotv2.1
422
+ - ./datasets/libero_goal_no_noops_lerobotv2.1
423
+ - ./datasets/libero_10_no_noops_lerobotv2.1
424
+ frame_window_size: 17
425
+ require_full_window: true
426
+ statistic_name: all_libero_no_noops
427
+ transforms:
428
+ - embodiment_id: 5
429
+ name_mappings:
430
+ actions:
431
+ - actions
432
+ observation.state:
433
+ - states
434
+ parquet_keys:
435
+ - observation.state
436
+ - timestamp
437
+ - actions
438
+ - info
439
+ - stats
440
+ type: ProcessParquetInputs
441
+ video_keys:
442
+ - observation.images.image
443
+ - observation.images.wrist_image
444
+ - type: LiberoFramewiseActionToRot6D
445
+ - action_metadata:
446
+ append_viewpoint: false
447
+ conditioning_fps: 20.0
448
+ frame_window_size: 17
449
+ video_height: 192
450
+ video_width: 320
451
+ cfg_dropout_rate: 0.1
452
+ format_prompt_as_json: true
453
+ max_len: 512
454
+ tokenizer:
455
+ model_max_length: 4096
456
+ model_path: ./checkpoints/Cosmos3-Edge
457
+ padding_side: right
458
+ trust_remote_code: true
459
+ type: PretrainedTokenizer
460
+ type: ProcessCosmos3Prompt
461
+ - height: 256
462
+ type: ResizeImages
463
+ width: 256
464
+ - type: SimpleNormalizeImages
465
+ - action_dim: 64
466
+ action_key: action
467
+ action_norm_type: quantile
468
+ state_dim: 64
469
+ state_key: proprio
470
+ state_norm_type: none
471
+ type: NormalizeStatesAndActions
472
+ - conditioning_fps: 20.0
473
+ frame_window_size: 17
474
+ mode: wam
475
+ prepend_state_to_action: false
476
+ raw_action_dim: 10
477
+ type: BuildCosmos3Sequence
478
+ - frame_window_size: 17
479
+ num_views: 2
480
+ tile_direction: horizontal
481
+ type: PrepareVideo
482
+ - height: 192
483
+ type: ResizeAndReflectPad
484
+ width: 320
485
+ type: ParquetDataset
486
+ use_delta: false
487
+ window_start_idx: 0
488
+ name_mappings:
489
+ action:
490
+ - action
491
+ observation.state:
492
+ - proprio
493
+ statistic_keys:
494
+ - observation.state
495
+ - timestamp
496
+ - action
497
+ statistic_name: all_libero_no_noops
498
+ statistics_overrides:
499
+ all_libero_no_noops:
500
+ action:
501
+ max:
502
+ - 0.9375
503
+ - 0.9375
504
+ - 0.9375
505
+ - 1.0
506
+ - 0.368853
507
+ - 0.341214
508
+ - 0.356395
509
+ - 1.0
510
+ - 0.348251
511
+ - 1.0
512
+ mean:
513
+ - 0.050704
514
+ - 0.097407
515
+ - -0.094833
516
+ - 0.994873
517
+ - -0.004579
518
+ - -0.004288
519
+ - 0.004389
520
+ - 0.996104
521
+ - 0.001109
522
+ - 0.476725
523
+ min:
524
+ - -0.9375
525
+ - -0.9375
526
+ - -0.9375
527
+ - 0.902028
528
+ - -0.356085
529
+ - -0.367416
530
+ - -0.370434
531
+ - 0.921907
532
+ - -0.255
533
+ - 0.0
534
+ q01:
535
+ - -0.723214
536
+ - -0.808929
537
+ - -0.9375
538
+ - 0.934955
539
+ - -0.223431
540
+ - -0.189878
541
+ - -0.334735
542
+ - 0.938516
543
+ - -0.107736
544
+ - 0.0
545
+ q99:
546
+ - 0.9375
547
+ - 0.870536
548
+ - 0.9375
549
+ - 1.0
550
+ - 0.331
551
+ - 0.163153
552
+ - 0.226216
553
+ - 1.0
554
+ - 0.127158
555
+ - 1.0
556
+ std:
557
+ - 0.333621
558
+ - 0.387175
559
+ - 0.45714
560
+ - 0.010807
561
+ - 0.077802
562
+ - 0.063386
563
+ - 0.078571
564
+ - 0.009994
565
+ - 0.038504
566
+ - 0.49946
567
+ type: DistributedRepeatingDataset
568
+ per_device_batch_size: 32
569
+ per_device_num_workers: 4
cosmos3_edge_libero_full_finetune_bs2048/cosmos3edge_libero_full_finetune_2026_08_17_03_08_18.jsonl ADDED
The diff for this file is too large to render. See raw diff
 
cosmos3_edge_libero_full_finetune_bs2048/dataset_statistics.json ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "all_libero_no_noops": {
3
+ "proprio": {
4
+ "mean": [
5
+ -0.047496414237850684,
6
+ 0.03435983471608021,
7
+ 0.7514258672095461,
8
+ 2.983973905120257,
9
+ -0.1419395165888572,
10
+ -0.118218935315545,
11
+ 0.026695325676034158,
12
+ -0.026933706126891162
13
+ ],
14
+ "std": [
15
+ 0.10441071071075952,
16
+ 0.15324089039592545,
17
+ 0.39886824990682457,
18
+ 0.34338260700499407,
19
+ 0.8068781715283699,
20
+ 0.31379284611306424,
21
+ 0.014319374319652562,
22
+ 0.014208265894281714
23
+ ],
24
+ "min": [
25
+ -0.48278069496154785,
26
+ -0.3309336006641388,
27
+ 0.008128181099891663,
28
+ 1.002794623374939,
29
+ -3.6312508583068848,
30
+ -1.842738389968872,
31
+ -0.005453015677630901,
32
+ -0.042015016078948975
33
+ ],
34
+ "max": [
35
+ 0.2103137969970703,
36
+ 0.3904264271259308,
37
+ 1.472778081893921,
38
+ 3.7248642444610596,
39
+ 3.5618896484375,
40
+ 1.3863215446472168,
41
+ 0.04232141748070717,
42
+ 0.0013126095291227102
43
+ ],
44
+ "q01": null,
45
+ "q99": null
46
+ },
47
+ "timestamp": {
48
+ "mean": [
49
+ 4.252653686389771
50
+ ],
51
+ "std": [
52
+ 3.2577109230601815
53
+ ],
54
+ "min": [
55
+ 0.0
56
+ ],
57
+ "max": [
58
+ 25.2
59
+ ],
60
+ "q01": null,
61
+ "q99": null
62
+ },
63
+ "action": {
64
+ "mean": [
65
+ 0.050704,
66
+ 0.097407,
67
+ -0.094833,
68
+ 0.994873,
69
+ -0.004579,
70
+ -0.004288,
71
+ 0.004389,
72
+ 0.996104,
73
+ 0.001109,
74
+ 0.476725
75
+ ],
76
+ "std": [
77
+ 0.333621,
78
+ 0.387175,
79
+ 0.45714,
80
+ 0.010807,
81
+ 0.077802,
82
+ 0.063386,
83
+ 0.078571,
84
+ 0.009994,
85
+ 0.038504,
86
+ 0.49946
87
+ ],
88
+ "min": [
89
+ -0.9375,
90
+ -0.9375,
91
+ -0.9375,
92
+ 0.902028,
93
+ -0.356085,
94
+ -0.367416,
95
+ -0.370434,
96
+ 0.921907,
97
+ -0.255,
98
+ 0.0
99
+ ],
100
+ "max": [
101
+ 0.9375,
102
+ 0.9375,
103
+ 0.9375,
104
+ 1.0,
105
+ 0.368853,
106
+ 0.341214,
107
+ 0.356395,
108
+ 1.0,
109
+ 0.348251,
110
+ 1.0
111
+ ],
112
+ "q01": [
113
+ -0.723214,
114
+ -0.808929,
115
+ -0.9375,
116
+ 0.934955,
117
+ -0.223431,
118
+ -0.189878,
119
+ -0.334735,
120
+ 0.938516,
121
+ -0.107736,
122
+ 0.0
123
+ ],
124
+ "q99": [
125
+ 0.9375,
126
+ 0.870536,
127
+ 0.9375,
128
+ 1.0,
129
+ 0.331,
130
+ 0.163153,
131
+ 0.226216,
132
+ 1.0,
133
+ 0.127158,
134
+ 1.0
135
+ ]
136
+ }
137
+ }
138
+ }
cosmos3_edge_libero_full_finetune_bs2048/run-metrics.jsonl ADDED
@@ -0,0 +1 @@
 
 
1
+ {"hparams": "{'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}}}, '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}}}, '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': 4, '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': 5000, 'metric': {'active_trackers': ('jsonl', 'wandb'), 'grad_accumulation_steps': 4, 'run_dir': '/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_mixed_edge_v21_192x320_gb2048_2x8_h100_5k', '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_full_finetune.py): {'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}}}, '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}}}, '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': 4, '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': 5000, 'metric': {'active_trackers': ('jsonl', 'wandb'), 'grad_accumulation_steps': 4, 'run_dir': '/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_mixed_edge_v21_192x320_gb2048_2x8_h100_5k', '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_full_finetune.py): {...}, 'args': Namespace(config='configs/cosmos3/cosmos3edge_libero_full_finetune.py', work_dir='/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_mixed_edge_v21_192x320_gb2048_2x8_h100_5k', cfg_options={'train_dataloader.per_device_batch_size': 32, 'runner.grad_accumulation_steps': 4, 'runner.metric.grad_accumulation_steps': 4, 'inference_model.rectified_flow_inference_config.shift': 10.0, 'eval.task_suite_name': ['libero_spatial', 'libero_object', 'libero_goal', 'libero_10'], 'eval.num_trials_per_task': 50, 'eval.task_ids': None, 'eval.run_id_suffix': 'mixed-v21-shift10-after-train'}, eval_after_train=True, resume_from=None)}, 'seed': 7, 'train_dataloader': {'dataset': {'datasets': {'action_key': 'action', 'action_window_size': 16, 'data_root_path': ['./datasets/libero_spatial_no_noops_lerobotv2.1', './datasets/libero_spatial_no_noops_lerobotv2.1', './datasets/libero_object_no_noops_lerobotv2.1', './datasets/libero_object_no_noops_lerobotv2.1', './datasets/libero_goal_no_noops_lerobotv2.1', './datasets/libero_goal_no_noops_lerobotv2.1', './datasets/libero_10_no_noops_lerobotv2.1'], 'frame_window_size': 17, 'require_full_window': True, 'statistic_name': 'all_libero_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': 192, 'video_width': 320}, 'cfg_dropout_rate': 0.1, 'format_prompt_as_json': True, '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'}, {'height': 192, 'type': 'ResizeAndReflectPad', 'width': 320}], 'type': 'ParquetDataset', 'use_delta': False, 'window_start_idx': 0}, 'name_mappings': {'action': ['action'], 'observation.state': ['proprio']}, 'statistic_keys': ['observation.state', 'timestamp', 'action'], 'statistic_name': 'all_libero_no_noops', 'statistics_overrides': {'all_libero_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': 32, 'per_device_num_workers': 4}, '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-crop', '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': 192, 'video_width': 320}, 'cfg_dropout_rate': 0.0, 'format_prompt_as_json': True, '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'}, {'height': 192, 'type': 'ResizeAndReflectPad', 'width': 320}], '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': 'all_libero_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', 'libero_object', 'libero_goal', 'libero_10'], 'type': 'LiberoEvalRunner', 'run_id_suffix': 'mixed-v21-shift10-after-train'}}, 'args': Namespace(config='configs/cosmos3/cosmos3edge_libero_full_finetune.py', work_dir='/mnt/data/cpfs/users/leo/wk_dir/cosmos3/libero_mixed_edge_v21_192x320_gb2048_2x8_h100_5k', cfg_options={'train_dataloader.per_device_batch_size': 32, 'runner.grad_accumulation_steps': 4, 'runner.metric.grad_accumulation_steps': 4, 'inference_model.rectified_flow_inference_config.shift': 10.0, 'eval.task_suite_name': ['libero_spatial', 'libero_object', 'libero_goal', 'libero_10'], 'eval.num_trials_per_task': 50, 'eval.task_ids': None, 'eval.run_id_suffix': 'mixed-v21-shift10-after-train'}, eval_after_train=True, resume_from=None)}, 'seed': 7, 'train_dataloader': {'dataset': {'datasets': {'action_key': 'action', 'action_window_size': 16, 'data_root_path': ['./datasets/libero_spatial_no_noops_lerobotv2.1', './datasets/libero_spatial_no_noops_lerobotv2.1', './datasets/libero_object_no_noops_lerobotv2.1', './datasets/libero_object_no_noops_lerobotv2.1', './datasets/libero_goal_no_noops_lerobotv2.1', './datasets/libero_goal_no_noops_lerobotv2.1', './datasets/libero_10_no_noops_lerobotv2.1'], 'frame_window_size': 17, 'require_full_window': True, 'statistic_name': 'all_libero_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': 192, 'video_width': 320}, 'cfg_dropout_rate': 0.1, 'format_prompt_as_json': True, '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'}, {'height': 192, 'type': 'ResizeAndReflectPad', 'width': 320}], 'type': 'ParquetDataset', 'use_delta': False, 'window_start_idx': 0}, 'name_mappings': {'action': ['action'], 'observation.state': ['proprio']}, 'statistic_keys': ['observation.state', 'timestamp', 'action'], 'statistic_name': 'all_libero_no_noops', 'statistics_overrides': {'all_libero_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': 32, 'per_device_num_workers': 4}, '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-crop', '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': 192, 'video_width': 320}, 'cfg_dropout_rate': 0.0, 'format_prompt_as_json': True, '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'}, {'height': 192, 'type': 'ResizeAndReflectPad', 'width': 320}], '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': 'all_libero_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', 'libero_object', 'libero_goal', 'libero_10'], 'type': 'LiberoEvalRunner', 'run_id_suffix': 'mixed-v21-shift10-after-train'}}", "run_id": "cosmos3edge_libero_full_finetune_2026_08_17_03_08_18"}
cosmos3_edge_libero_full_finetune_bs2048/tokenizer/chat_template.jinja ADDED
@@ -0,0 +1,324 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ {% macro render_extra_keys(json_dict, handled_keys) %}
3
+ {%- if json_dict is mapping %}
4
+ {%- for json_key in json_dict if json_key not in handled_keys %}
5
+ {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
6
+ {{- '
7
+ <' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
8
+ {%- else %}
9
+ {{-'
10
+ <' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
11
+ {%- endif %}
12
+ {%- endfor %}
13
+ {%- endif %}
14
+ {% endmacro %}
15
+ {%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
16
+ {%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
17
+
18
+ {%- set ns = namespace(last_user_idx = -1) %}
19
+ {%- set loop_messages = messages %}
20
+ {%- for m in loop_messages %}
21
+ {%- if m["role"] == "user" %}
22
+ {%- set ns.last_user_idx = loop.index0 %}
23
+ {%- endif %}
24
+ {%- endfor %}
25
+
26
+ {%- if messages[0]["role"] == "system" %}
27
+ {%- if messages[0]["content"] is string %}
28
+ {%- set system_message = messages[0]["content"] %}
29
+ {%- else %}
30
+ {%- set ns = namespace(system_text="") %}
31
+ {%- for part in messages[0]["content"] %}
32
+ {%- if 'text' in part %}
33
+ {%- set ns.system_text = ns.system_text + part.text %}
34
+ {%- endif %}
35
+ {%- endfor %}
36
+ {%- set system_message = ns.system_text %}
37
+ {%- endif %}
38
+
39
+ {%- set loop_messages = messages[1:] %}
40
+ {%- else %}
41
+ {%- set system_message = "" %}
42
+ {%- set loop_messages = messages %}
43
+ {%- endif %}
44
+ {%- if not tools is defined %}
45
+ {%- set tools = [] %}
46
+ {%- endif %}
47
+ {# Recompute last_user_idx relative to loop_messages after handling system #}
48
+ {%- set ns = namespace(last_user_idx = -1) %}
49
+ {%- for m in loop_messages %}
50
+ {%- if m["role"] == "user" %}
51
+ {%- set ns.last_user_idx = loop.index0 %}
52
+ {%- endif %}
53
+ {%- endfor %}
54
+ {%- if system_message is defined %}
55
+ {{- "<|im_start|>system
56
+ " + system_message }}
57
+ {%- else %}
58
+ {%- if tools is iterable and tools | length > 0 %}
59
+ {{- "<|im_start|>system
60
+ " }}
61
+ {%- endif %}
62
+ {%- endif %}
63
+ {%- if tools is iterable and tools | length > 0 %}
64
+ {%- if system_message is defined and system_message | length > 0 %}
65
+ {{- "
66
+
67
+ " }}
68
+ {%- endif %}
69
+ {{- "# Tools
70
+
71
+ You have access to the following functions:
72
+
73
+ " }}
74
+ {{- "<tools>" }}
75
+ {%- for tool in tools %}
76
+ {%- if tool.function is defined %}
77
+ {%- set tool = tool.function %}
78
+ {%- endif %}
79
+ {{- "
80
+ <function>
81
+ <name>" ~ tool.name ~ "</name>" }}
82
+ {%- if tool.description is defined %}
83
+ {{- '
84
+ <description>' ~ (tool.description | trim) ~ '</description>' }}
85
+ {%- endif %}
86
+ {{- '
87
+ <parameters>' }}
88
+ {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
89
+ {%- for param_name, param_fields in tool.parameters.properties|items %}
90
+ {{- '
91
+ <parameter>' }}
92
+ {{- '
93
+ <name>' ~ param_name ~ '</name>' }}
94
+ {%- if param_fields.type is defined %}
95
+ {{- '
96
+ <type>' ~ (param_fields.type | string) ~ '</type>' }}
97
+ {%- endif %}
98
+ {%- if param_fields.description is defined %}
99
+ {{- '
100
+ <description>' ~ (param_fields.description | trim) ~ '</description>' }}
101
+ {%- endif %}
102
+ {%- if param_fields.enum is defined %}
103
+ {{- '
104
+ <enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
105
+ {%- endif %}
106
+ {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
107
+ {{- render_extra_keys(param_fields, handled_keys) }}
108
+ {{- '
109
+ </parameter>' }}
110
+ {%- endfor %}
111
+ {%- endif %}
112
+ {% set handled_keys = ['type', 'properties', 'required'] %}
113
+ {{- render_extra_keys(tool.parameters, handled_keys) }}
114
+ {%- if tool.parameters is defined and tool.parameters.required is defined %}
115
+ {{- '
116
+ <required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
117
+ {%- endif %}
118
+ {{- '
119
+ </parameters>' }}
120
+ {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
121
+ {{- render_extra_keys(tool, handled_keys) }}
122
+ {{- '
123
+ </function>' }}
124
+ {%- endfor %}
125
+ {{- "
126
+ </tools>" }}
127
+
128
+ {{- '
129
+
130
+ If you choose to call a function ONLY reply in the following format with NO suffix:
131
+
132
+ <tool_call>
133
+ <function=example_function_name>
134
+ <parameter=example_parameter_1>
135
+ value_1
136
+ </parameter>
137
+ <parameter=example_parameter_2>
138
+ This is the value for the second parameter
139
+ that can span
140
+ multiple lines
141
+ </parameter>
142
+ </function>
143
+ </tool_call>
144
+
145
+ <IMPORTANT>
146
+ Reminder:
147
+ - Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
148
+ - Required parameters MUST be specified
149
+ - You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
150
+ - 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
151
+ </IMPORTANT>' }}
152
+ {%- endif %}
153
+
154
+
155
+ {%- if system_message is defined %}
156
+ {{- '<|im_end|>
157
+ ' }}
158
+ {%- else %}
159
+ {%- if tools is iterable and tools | length > 0 %}
160
+ {{- '<|im_end|>
161
+ ' }}
162
+ {%- endif %}
163
+ {%- endif %}
164
+
165
+ {%- set image_count = namespace(value=0) %}
166
+ {%- set video_count = namespace(value=0) %}
167
+ {%- for message in loop_messages %}
168
+ {%- if message.role == "assistant" %}
169
+ {# Add reasoning content in to content field for unified processing below. #}
170
+ {%- set text_content = namespace(text_content="") -%}
171
+ {%- if message.content is string %}
172
+ {%- set text_content.text_content = message.content %}
173
+ {%- else %}
174
+ {%- for content in message.content %}
175
+ {%- if 'text' in content %}
176
+ {%- set text_content.text_content = content.text %}
177
+ {%- break %}
178
+ {%- endif %}
179
+ {%- endfor %}
180
+ {%- endif %}
181
+ {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
182
+ {%- set content = "<think>
183
+ " ~ message.reasoning_content ~ "
184
+ </think>
185
+ " ~ (text_content.text_content | default('', true)) %}
186
+ {%- else %}
187
+ {%- set content = text_content.text_content | default('', true) %}
188
+ {# Allow downstream logic to to take care of broken thought, only handle coherent reasoning here. #}
189
+ {%- if '<think>' not in content and '</think>' not in content %}
190
+ {%- set content = "<think></think>" ~ content %}
191
+ {%- else %}
192
+ {%- set content = content %}
193
+ {%- endif -%}
194
+ {%- endif %}
195
+ {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
196
+ {# Assistant message has tool calls. #}
197
+ {{- '<|im_start|>assistant
198
+ ' }}
199
+ {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
200
+ {%- if content is string and content | trim | length > 0 %}
201
+ {%- if include_content %}
202
+ {{- (content | trim) ~ '
203
+ ' -}}
204
+ {%- else %}
205
+ {%- set c = (content | string) %}
206
+ {%- if '</think>' in c %}
207
+ {# Keep only content after the last closing think. Also generation prompt causes this. #}
208
+ {%- set c = c.split('</think>')[-1] %}
209
+ {%- elif '<think>' in c %}
210
+ {# If <think> was opened but never closed, drop the trailing think segment #}
211
+ {%- set c = c.split('<think>')[0] %}
212
+ {%- endif %}
213
+ {%- set c = "<think></think>" ~ c | trim %}
214
+ {%- if c | length > 0 %}
215
+ {{- c ~ '
216
+ ' -}}
217
+ {%- endif %}
218
+ {%- endif %}
219
+ {%- else %}
220
+ {{- "<think></think>" -}}
221
+ {%- endif %}
222
+ {%- for tool_call in message.tool_calls %}
223
+ {%- if tool_call.function is defined %}
224
+ {%- set tool_call = tool_call.function %}
225
+ {%- endif %}
226
+ {{- '<tool_call>
227
+ <function=' ~ tool_call.name ~ '>
228
+ ' -}}
229
+ {%- if tool_call.arguments is defined %}
230
+ {%- for args_name, args_value in tool_call.arguments|items %}
231
+ {{- '<parameter=' ~ args_name ~ '>
232
+ ' -}}
233
+ {%- 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 %}
234
+ {{- args_value ~ '
235
+ </parameter>
236
+ ' -}}
237
+ {%- endfor %}
238
+ {%- endif %}
239
+ {{- '</function>
240
+ </tool_call>
241
+ ' -}}
242
+ {%- endfor %}
243
+ {{- '<|im_end|>
244
+ ' }}
245
+ {%- else %}
246
+ {# Assistant message doesn't have tool calls. #}
247
+ {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
248
+ {{- '<|im_start|>assistant
249
+ ' ~ (content | default('', true) | string | trim) ~ '<|im_end|>
250
+ ' }}
251
+ {%- else %}
252
+ {%- set c = (content | default('', true) | string) %}
253
+ {%- if '<think>' in c and '</think>' in c %}
254
+ {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
255
+ {%- endif %}
256
+ {%- set c = c | trim %}
257
+ {%- if c | length > 0 %}
258
+ {{- '<|im_start|>assistant
259
+ ' ~ c ~ '<|im_end|>
260
+ ' }}
261
+ {%- else %}
262
+ {{- '<|im_start|>assistant
263
+ <|im_end|>
264
+ ' }}
265
+ {%- endif %}
266
+ {%- endif %}
267
+ {%- endif %}
268
+ {%- elif message.role == "user" or message.role == "system" %}
269
+ {{- '<|im_start|>' + message.role + '
270
+ ' }}
271
+ {%- if message.content is string %}
272
+ {{- message.content }}
273
+ {%- else %}
274
+ {%- for content in message.content %}
275
+ {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}
276
+ {%- set image_count.value = image_count.value + 1 %}
277
+ {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
278
+ <|vision_start|><|image_pad|><|vision_end|>
279
+ {%- elif content.type == 'video' or 'video' in content %}
280
+ {%- set video_count.value = video_count.value + 1 %}
281
+ {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
282
+ <|vision_start|><|video_pad|><|vision_end|>
283
+ {%- elif 'text' in content %}
284
+ {{- content.text }}
285
+ {%- endif %}
286
+ {%- endfor %}
287
+ {%- endif %}
288
+ {{- '<|im_end|>
289
+ ' }}
290
+ {%- elif message.role == "tool" %}
291
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
292
+ {{- '<|im_start|>user
293
+ ' }}
294
+ {%- endif %}
295
+ {{- '<tool_response>
296
+ ' }}
297
+ {{- message.content }}
298
+ {{- '
299
+ </tool_response>
300
+ ' }}
301
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
302
+ {{- '<|im_end|>
303
+ ' }}
304
+ {%- elif loop.last %}
305
+ {{- '<|im_end|>
306
+ ' }}
307
+ {%- endif %}
308
+ {%- else %}
309
+ {{- '<|im_start|>' + message.role + '
310
+ ' + message.content + '<|im_end|>
311
+ ' }}
312
+ {%- endif %}
313
+ {%- endfor %}
314
+
315
+ {%- if add_generation_prompt %}
316
+ {%- if enable_thinking %}
317
+ {{- '<|im_start|>assistant
318
+ <think>
319
+ ' }}
320
+ {%- else %}
321
+ {{- '<|im_start|>assistant
322
+ <think></think>' }}
323
+ {%- endif %}
324
+ {%- endif %}
cosmos3_edge_libero_full_finetune_bs2048/tokenizer/special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "<|im_end|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "<|im_end|>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<unk>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
cosmos3_edge_libero_full_finetune_bs2048/tokenizer/tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4dc692a99dca6d92a44e9529ffb53541eb79bb260eec7281bc51268c05d7311c
3
+ size 17077500
cosmos3_edge_libero_full_finetune_bs2048/tokenizer/tokenizer_config.json ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "backend": "tokenizers",
4
+ "bos_token": "<s>",
5
+ "clean_up_tokenization_spaces": false,
6
+ "eos_token": "<|im_end|>",
7
+ "is_local": true,
8
+ "local_files_only": false,
9
+ "model_input_names": [
10
+ "input_ids",
11
+ "attention_mask"
12
+ ],
13
+ "model_max_length": 262144,
14
+ "pad_token": "<|im_end|>",
15
+ "tokenizer_class": "TokenizersBackend",
16
+ "unk_token": "<unk>"
17
+ }
cosmos3_edge_libero_full_finetune_bs2048/vlm_backbone_config.json ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "attention_bias": false,
3
+ "attention_dropout": 0.0,
4
+ "bos_token_id": 1,
5
+ "enable_mrope": true,
6
+ "enable_rope": true,
7
+ "eos_token_id": 11,
8
+ "head_dim": 128,
9
+ "hidden_dropout": 0.0,
10
+ "hidden_size": 2048,
11
+ "initializer_range": 0.02,
12
+ "intermediate_size": 9216,
13
+ "layer_norm_epsilon": 1e-05,
14
+ "max_position_embeddings": 131072,
15
+ "mlp_bias": false,
16
+ "mlp_hidden_act": "relu2",
17
+ "model_type": "nemotron_3_dense_vl_text",
18
+ "mrope_section": [
19
+ 24,
20
+ 20,
21
+ 20
22
+ ],
23
+ "num_attention_heads": 16,
24
+ "num_hidden_layers": 28,
25
+ "num_key_value_heads": 8,
26
+ "num_logits_to_keep": 1,
27
+ "packed_attention_backend": "flash2",
28
+ "pad_token_id": 11,
29
+ "residual_in_fp32": false,
30
+ "rope_parameters": {
31
+ "rope_theta": 100000000.0,
32
+ "rope_type": "default"
33
+ },
34
+ "rope_theta": 100000000.0,
35
+ "sliding_window": null,
36
+ "tie_word_embeddings": false,
37
+ "transformers_version": "5.3.0",
38
+ "use_cache": true,
39
+ "use_und_k_norm_for_gen": true,
40
+ "vocab_size": 131072
41
+ }