happyhappy-jun commited on
Commit
8fe06e4
·
verified ·
1 Parent(s): ff4d9dc

Upload folder using huggingface_hub

Browse files
checkpoint-60000/config.json ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_horizon": 50,
3
+ "add_pos_embed": true,
4
+ "apply_sincos_state_encoding": true,
5
+ "architectures": [
6
+ "Gr00tN1d6"
7
+ ],
8
+ "attn_dropout": 0.2,
9
+ "attn_implementation": null,
10
+ "backbone_embedding_dim": 2048,
11
+ "backbone_model_type": "eagle",
12
+ "backbone_trainable_params_fp32": true,
13
+ "collator_overwrite_image_inputs": false,
14
+ "color_jitter_params": {
15
+ "brightness": 0.1,
16
+ "contrast": 0.1,
17
+ "hue": 0.1,
18
+ "saturation": 0.1
19
+ },
20
+ "crop_fraction": 0.95,
21
+ "diffusion_model_cfg": {
22
+ "attention_head_dim": 48,
23
+ "dropout": 0.2,
24
+ "final_dropout": true,
25
+ "interleave_self_attention": true,
26
+ "norm_type": "ada_norm",
27
+ "num_attention_heads": 32,
28
+ "num_layers": 32,
29
+ "output_dim": 1024,
30
+ "positional_embeddings": null
31
+ },
32
+ "eagle_collator": true,
33
+ "formalize_language": true,
34
+ "gemma_collator": false,
35
+ "hidden_size": 1024,
36
+ "image_crop_size": null,
37
+ "image_target_size": null,
38
+ "input_embedding_dim": 1536,
39
+ "load_bf16": true,
40
+ "max_action_dim": 128,
41
+ "max_num_embodiments": 32,
42
+ "max_seq_len": 1024,
43
+ "max_state_dim": 128,
44
+ "model_dtype": "bfloat16",
45
+ "model_name": "nvidia/Eagle-Block2A-2B-v2",
46
+ "model_type": "Gr00tN1d6",
47
+ "noise_beta_alpha": 1.5,
48
+ "noise_beta_beta": 1.0,
49
+ "noise_s": 0.999,
50
+ "num_inference_timesteps": 4,
51
+ "num_timestep_buckets": 1000,
52
+ "random_rotation_angle": null,
53
+ "reproject_vision": false,
54
+ "select_layer": 16,
55
+ "shortest_image_edge": 256,
56
+ "state_dropout_prob": 0.0,
57
+ "torch_dtype": "bfloat16",
58
+ "transformers_version": "4.51.3",
59
+ "tune_diffusion_model": true,
60
+ "tune_llm": false,
61
+ "tune_projector": true,
62
+ "tune_top_llm_layers": 4,
63
+ "tune_visual": false,
64
+ "tune_vlln": true,
65
+ "use_albumentations_transforms": true,
66
+ "use_alternate_vl_dit": true,
67
+ "use_flash_attention": true,
68
+ "use_relative_action": true,
69
+ "use_vlln": true
70
+ }
checkpoint-60000/embodiment_id.json ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "robocasa_panda_omron": 13,
3
+ "gr1": 20,
4
+ "behavior_r1_pro": 24,
5
+ "unitree_g1": 8,
6
+ "oxe_google": 0,
7
+ "oxe_widowx": 1,
8
+ "libero_panda": 2,
9
+ "oxe_droid": 16,
10
+ "new_embodiment": 10
11
+ }
checkpoint-60000/latest ADDED
@@ -0,0 +1 @@
 
 
1
+ global_step60000
checkpoint-60000/model-00001-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8fd3f824fb806d8bcd6c308b99c9fb3bc02ab201cf161343b452a60beb361255
3
+ size 4991091456
checkpoint-60000/model-00002-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9bb5cc1c97053e21fded718b9138232709488ac911d0116d4c346f299662ba3b
3
+ size 1582283096
checkpoint-60000/model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
checkpoint-60000/processor_config.json ADDED
@@ -0,0 +1,923 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "processor_class": "Gr00tN1d6Processor",
3
+ "processor_kwargs": {
4
+ "modality_configs": {
5
+ "behavior_r1_pro": {
6
+ "video": {
7
+ "delta_indices": [
8
+ 0
9
+ ],
10
+ "modality_keys": [
11
+ "observation.images.rgb.head_256_256",
12
+ "observation.images.rgb.left_wrist_256_256",
13
+ "observation.images.rgb.right_wrist_256_256"
14
+ ],
15
+ "sin_cos_embedding_keys": null,
16
+ "mean_std_embedding_keys": null,
17
+ "action_configs": null
18
+ },
19
+ "state": {
20
+ "delta_indices": [
21
+ 0
22
+ ],
23
+ "modality_keys": [
24
+ "robot_pos",
25
+ "robot_ori_cos",
26
+ "robot_ori_sin",
27
+ "robot_2d_ori",
28
+ "robot_2d_ori_cos",
29
+ "robot_2d_ori_sin",
30
+ "robot_lin_vel",
31
+ "robot_ang_vel",
32
+ "arm_left_qpos",
33
+ "arm_left_qpos_sin",
34
+ "arm_left_qpos_cos",
35
+ "eef_left_pos",
36
+ "eef_left_quat",
37
+ "gripper_left_qpos",
38
+ "arm_right_qpos",
39
+ "arm_right_qpos_sin",
40
+ "arm_right_qpos_cos",
41
+ "eef_right_pos",
42
+ "eef_right_quat",
43
+ "gripper_right_qpos",
44
+ "trunk_qpos"
45
+ ],
46
+ "sin_cos_embedding_keys": null,
47
+ "mean_std_embedding_keys": null,
48
+ "action_configs": null
49
+ },
50
+ "action": {
51
+ "delta_indices": [
52
+ 0,
53
+ 1,
54
+ 2,
55
+ 3,
56
+ 4,
57
+ 5,
58
+ 6,
59
+ 7,
60
+ 8,
61
+ 9,
62
+ 10,
63
+ 11,
64
+ 12,
65
+ 13,
66
+ 14,
67
+ 15,
68
+ 16,
69
+ 17,
70
+ 18,
71
+ 19,
72
+ 20,
73
+ 21,
74
+ 22,
75
+ 23,
76
+ 24,
77
+ 25,
78
+ 26,
79
+ 27,
80
+ 28,
81
+ 29,
82
+ 30,
83
+ 31
84
+ ],
85
+ "modality_keys": [
86
+ "base",
87
+ "torso",
88
+ "left_arm",
89
+ "left_gripper",
90
+ "right_arm",
91
+ "right_gripper"
92
+ ],
93
+ "sin_cos_embedding_keys": null,
94
+ "mean_std_embedding_keys": null,
95
+ "action_configs": [
96
+ {
97
+ "rep": "ABSOLUTE",
98
+ "type": "NON_EEF",
99
+ "format": "DEFAULT",
100
+ "state_key": null
101
+ },
102
+ {
103
+ "rep": "RELATIVE",
104
+ "type": "NON_EEF",
105
+ "format": "DEFAULT",
106
+ "state_key": "trunk_qpos"
107
+ },
108
+ {
109
+ "rep": "RELATIVE",
110
+ "type": "NON_EEF",
111
+ "format": "DEFAULT",
112
+ "state_key": "arm_left_qpos"
113
+ },
114
+ {
115
+ "rep": "ABSOLUTE",
116
+ "type": "NON_EEF",
117
+ "format": "DEFAULT",
118
+ "state_key": null
119
+ },
120
+ {
121
+ "rep": "RELATIVE",
122
+ "type": "NON_EEF",
123
+ "format": "DEFAULT",
124
+ "state_key": "arm_right_qpos"
125
+ },
126
+ {
127
+ "rep": "ABSOLUTE",
128
+ "type": "NON_EEF",
129
+ "format": "DEFAULT",
130
+ "state_key": null
131
+ }
132
+ ]
133
+ },
134
+ "language": {
135
+ "delta_indices": [
136
+ 0
137
+ ],
138
+ "modality_keys": [
139
+ "annotation.human.coarse_action"
140
+ ],
141
+ "sin_cos_embedding_keys": null,
142
+ "mean_std_embedding_keys": null,
143
+ "action_configs": null
144
+ }
145
+ },
146
+ "gr1": {
147
+ "video": {
148
+ "delta_indices": [
149
+ 0
150
+ ],
151
+ "modality_keys": [
152
+ "ego_view_bg_crop_pad_res256_freq20"
153
+ ],
154
+ "sin_cos_embedding_keys": null,
155
+ "mean_std_embedding_keys": null,
156
+ "action_configs": null
157
+ },
158
+ "state": {
159
+ "delta_indices": [
160
+ 0
161
+ ],
162
+ "modality_keys": [
163
+ "left_arm",
164
+ "right_arm",
165
+ "left_hand",
166
+ "right_hand",
167
+ "waist"
168
+ ],
169
+ "sin_cos_embedding_keys": [
170
+ "left_arm",
171
+ "right_arm",
172
+ "left_hand",
173
+ "right_hand",
174
+ "waist"
175
+ ],
176
+ "mean_std_embedding_keys": null,
177
+ "action_configs": null
178
+ },
179
+ "action": {
180
+ "delta_indices": [
181
+ 0,
182
+ 1,
183
+ 2,
184
+ 3,
185
+ 4,
186
+ 5,
187
+ 6,
188
+ 7,
189
+ 8,
190
+ 9,
191
+ 10,
192
+ 11,
193
+ 12,
194
+ 13,
195
+ 14,
196
+ 15
197
+ ],
198
+ "modality_keys": [
199
+ "left_arm",
200
+ "right_arm",
201
+ "left_hand",
202
+ "right_hand",
203
+ "waist"
204
+ ],
205
+ "sin_cos_embedding_keys": null,
206
+ "mean_std_embedding_keys": null,
207
+ "action_configs": [
208
+ {
209
+ "rep": "RELATIVE",
210
+ "type": "NON_EEF",
211
+ "format": "DEFAULT",
212
+ "state_key": null
213
+ },
214
+ {
215
+ "rep": "RELATIVE",
216
+ "type": "NON_EEF",
217
+ "format": "DEFAULT",
218
+ "state_key": null
219
+ },
220
+ {
221
+ "rep": "RELATIVE",
222
+ "type": "NON_EEF",
223
+ "format": "DEFAULT",
224
+ "state_key": null
225
+ },
226
+ {
227
+ "rep": "RELATIVE",
228
+ "type": "NON_EEF",
229
+ "format": "DEFAULT",
230
+ "state_key": null
231
+ },
232
+ {
233
+ "rep": "ABSOLUTE",
234
+ "type": "NON_EEF",
235
+ "format": "DEFAULT",
236
+ "state_key": null
237
+ }
238
+ ]
239
+ },
240
+ "language": {
241
+ "delta_indices": [
242
+ 0
243
+ ],
244
+ "modality_keys": [
245
+ "task"
246
+ ],
247
+ "sin_cos_embedding_keys": null,
248
+ "mean_std_embedding_keys": null,
249
+ "action_configs": null
250
+ }
251
+ },
252
+ "robocasa_panda_omron": {
253
+ "video": {
254
+ "delta_indices": [
255
+ 0
256
+ ],
257
+ "modality_keys": [
258
+ "robot0_eye_in_hand",
259
+ "robot0_agentview_left",
260
+ "robot0_agentview_right"
261
+ ],
262
+ "sin_cos_embedding_keys": null,
263
+ "mean_std_embedding_keys": null,
264
+ "action_configs": null
265
+ },
266
+ "state": {
267
+ "delta_indices": [
268
+ 0
269
+ ],
270
+ "modality_keys": [
271
+ "base_position",
272
+ "base_rotation",
273
+ "end_effector_position_relative",
274
+ "end_effector_rotation_relative",
275
+ "gripper_qpos"
276
+ ],
277
+ "sin_cos_embedding_keys": null,
278
+ "mean_std_embedding_keys": null,
279
+ "action_configs": null
280
+ },
281
+ "action": {
282
+ "delta_indices": [
283
+ 0,
284
+ 1,
285
+ 2,
286
+ 3,
287
+ 4,
288
+ 5,
289
+ 6,
290
+ 7,
291
+ 8,
292
+ 9,
293
+ 10,
294
+ 11,
295
+ 12,
296
+ 13,
297
+ 14,
298
+ 15
299
+ ],
300
+ "modality_keys": [
301
+ "base_motion",
302
+ "control_mode",
303
+ "end_effector_position",
304
+ "end_effector_rotation",
305
+ "gripper_close"
306
+ ],
307
+ "sin_cos_embedding_keys": null,
308
+ "mean_std_embedding_keys": null,
309
+ "action_configs": [
310
+ {
311
+ "rep": "ABSOLUTE",
312
+ "type": "NON_EEF",
313
+ "format": "DEFAULT",
314
+ "state_key": null
315
+ },
316
+ {
317
+ "rep": "ABSOLUTE",
318
+ "type": "NON_EEF",
319
+ "format": "DEFAULT",
320
+ "state_key": null
321
+ },
322
+ {
323
+ "rep": "ABSOLUTE",
324
+ "type": "EEF",
325
+ "format": "DEFAULT",
326
+ "state_key": null
327
+ },
328
+ {
329
+ "rep": "ABSOLUTE",
330
+ "type": "EEF",
331
+ "format": "DEFAULT",
332
+ "state_key": null
333
+ },
334
+ {
335
+ "rep": "ABSOLUTE",
336
+ "type": "NON_EEF",
337
+ "format": "DEFAULT",
338
+ "state_key": null
339
+ }
340
+ ]
341
+ },
342
+ "language": {
343
+ "delta_indices": [
344
+ 0
345
+ ],
346
+ "modality_keys": [
347
+ "task"
348
+ ],
349
+ "sin_cos_embedding_keys": null,
350
+ "mean_std_embedding_keys": null,
351
+ "action_configs": null
352
+ }
353
+ },
354
+ "unitree_g1": {
355
+ "video": {
356
+ "delta_indices": [
357
+ 0
358
+ ],
359
+ "modality_keys": [
360
+ "ego_view"
361
+ ],
362
+ "sin_cos_embedding_keys": null,
363
+ "mean_std_embedding_keys": null,
364
+ "action_configs": null
365
+ },
366
+ "state": {
367
+ "delta_indices": [
368
+ 0
369
+ ],
370
+ "modality_keys": [
371
+ "left_leg",
372
+ "right_leg",
373
+ "waist",
374
+ "left_arm",
375
+ "right_arm",
376
+ "left_hand",
377
+ "right_hand"
378
+ ],
379
+ "sin_cos_embedding_keys": null,
380
+ "mean_std_embedding_keys": null,
381
+ "action_configs": null
382
+ },
383
+ "action": {
384
+ "delta_indices": [
385
+ 0,
386
+ 1,
387
+ 2,
388
+ 3,
389
+ 4,
390
+ 5,
391
+ 6,
392
+ 7,
393
+ 8,
394
+ 9,
395
+ 10,
396
+ 11,
397
+ 12,
398
+ 13,
399
+ 14,
400
+ 15,
401
+ 16,
402
+ 17,
403
+ 18,
404
+ 19,
405
+ 20,
406
+ 21,
407
+ 22,
408
+ 23,
409
+ 24,
410
+ 25,
411
+ 26,
412
+ 27,
413
+ 28,
414
+ 29
415
+ ],
416
+ "modality_keys": [
417
+ "left_arm",
418
+ "right_arm",
419
+ "left_hand",
420
+ "right_hand",
421
+ "waist",
422
+ "base_height_command",
423
+ "navigate_command"
424
+ ],
425
+ "sin_cos_embedding_keys": null,
426
+ "mean_std_embedding_keys": null,
427
+ "action_configs": [
428
+ {
429
+ "rep": "RELATIVE",
430
+ "type": "NON_EEF",
431
+ "format": "DEFAULT",
432
+ "state_key": null
433
+ },
434
+ {
435
+ "rep": "RELATIVE",
436
+ "type": "NON_EEF",
437
+ "format": "DEFAULT",
438
+ "state_key": null
439
+ },
440
+ {
441
+ "rep": "ABSOLUTE",
442
+ "type": "NON_EEF",
443
+ "format": "DEFAULT",
444
+ "state_key": null
445
+ },
446
+ {
447
+ "rep": "ABSOLUTE",
448
+ "type": "NON_EEF",
449
+ "format": "DEFAULT",
450
+ "state_key": null
451
+ },
452
+ {
453
+ "rep": "ABSOLUTE",
454
+ "type": "NON_EEF",
455
+ "format": "DEFAULT",
456
+ "state_key": null
457
+ },
458
+ {
459
+ "rep": "ABSOLUTE",
460
+ "type": "NON_EEF",
461
+ "format": "DEFAULT",
462
+ "state_key": null
463
+ },
464
+ {
465
+ "rep": "ABSOLUTE",
466
+ "type": "NON_EEF",
467
+ "format": "DEFAULT",
468
+ "state_key": null
469
+ }
470
+ ]
471
+ },
472
+ "language": {
473
+ "delta_indices": [
474
+ 0
475
+ ],
476
+ "modality_keys": [
477
+ "annotation.human.task_description"
478
+ ],
479
+ "sin_cos_embedding_keys": null,
480
+ "mean_std_embedding_keys": null,
481
+ "action_configs": null
482
+ }
483
+ },
484
+ "libero_panda": {
485
+ "video": {
486
+ "delta_indices": [
487
+ 0
488
+ ],
489
+ "modality_keys": [
490
+ "image",
491
+ "wrist_image"
492
+ ],
493
+ "sin_cos_embedding_keys": null,
494
+ "mean_std_embedding_keys": null,
495
+ "action_configs": null
496
+ },
497
+ "state": {
498
+ "delta_indices": [
499
+ 0
500
+ ],
501
+ "modality_keys": [
502
+ "x",
503
+ "y",
504
+ "z",
505
+ "roll",
506
+ "pitch",
507
+ "yaw",
508
+ "gripper"
509
+ ],
510
+ "sin_cos_embedding_keys": null,
511
+ "mean_std_embedding_keys": null,
512
+ "action_configs": null
513
+ },
514
+ "action": {
515
+ "delta_indices": [
516
+ 0,
517
+ 1,
518
+ 2,
519
+ 3,
520
+ 4,
521
+ 5,
522
+ 6,
523
+ 7,
524
+ 8,
525
+ 9,
526
+ 10,
527
+ 11,
528
+ 12,
529
+ 13,
530
+ 14,
531
+ 15
532
+ ],
533
+ "modality_keys": [
534
+ "x",
535
+ "y",
536
+ "z",
537
+ "roll",
538
+ "pitch",
539
+ "yaw",
540
+ "gripper"
541
+ ],
542
+ "sin_cos_embedding_keys": null,
543
+ "mean_std_embedding_keys": null,
544
+ "action_configs": null
545
+ },
546
+ "language": {
547
+ "delta_indices": [
548
+ 0
549
+ ],
550
+ "modality_keys": [
551
+ "annotation.human.action.task_description"
552
+ ],
553
+ "sin_cos_embedding_keys": null,
554
+ "mean_std_embedding_keys": null,
555
+ "action_configs": null
556
+ }
557
+ },
558
+ "oxe_widowx": {
559
+ "video": {
560
+ "delta_indices": [
561
+ 0
562
+ ],
563
+ "modality_keys": [
564
+ "image_0"
565
+ ],
566
+ "sin_cos_embedding_keys": null,
567
+ "mean_std_embedding_keys": null,
568
+ "action_configs": null
569
+ },
570
+ "state": {
571
+ "delta_indices": [
572
+ 0
573
+ ],
574
+ "modality_keys": [
575
+ "x",
576
+ "y",
577
+ "z",
578
+ "roll",
579
+ "pitch",
580
+ "yaw",
581
+ "pad",
582
+ "gripper"
583
+ ],
584
+ "sin_cos_embedding_keys": null,
585
+ "mean_std_embedding_keys": null,
586
+ "action_configs": null
587
+ },
588
+ "action": {
589
+ "delta_indices": [
590
+ 0,
591
+ 1,
592
+ 2,
593
+ 3,
594
+ 4,
595
+ 5,
596
+ 6,
597
+ 7
598
+ ],
599
+ "modality_keys": [
600
+ "x",
601
+ "y",
602
+ "z",
603
+ "roll",
604
+ "pitch",
605
+ "yaw",
606
+ "gripper"
607
+ ],
608
+ "sin_cos_embedding_keys": null,
609
+ "mean_std_embedding_keys": [
610
+ "x",
611
+ "y",
612
+ "z",
613
+ "roll",
614
+ "pitch",
615
+ "yaw"
616
+ ],
617
+ "action_configs": null
618
+ },
619
+ "language": {
620
+ "delta_indices": [
621
+ 0
622
+ ],
623
+ "modality_keys": [
624
+ "annotation.human.action.task_description"
625
+ ],
626
+ "sin_cos_embedding_keys": null,
627
+ "mean_std_embedding_keys": null,
628
+ "action_configs": null
629
+ }
630
+ },
631
+ "oxe_google": {
632
+ "video": {
633
+ "delta_indices": [
634
+ 0
635
+ ],
636
+ "modality_keys": [
637
+ "image"
638
+ ],
639
+ "sin_cos_embedding_keys": null,
640
+ "mean_std_embedding_keys": null,
641
+ "action_configs": null
642
+ },
643
+ "state": {
644
+ "delta_indices": [
645
+ 0
646
+ ],
647
+ "modality_keys": [
648
+ "x",
649
+ "y",
650
+ "z",
651
+ "rx",
652
+ "ry",
653
+ "rz",
654
+ "rw",
655
+ "gripper"
656
+ ],
657
+ "sin_cos_embedding_keys": null,
658
+ "mean_std_embedding_keys": null,
659
+ "action_configs": null
660
+ },
661
+ "action": {
662
+ "delta_indices": [
663
+ 0,
664
+ 1,
665
+ 2,
666
+ 3,
667
+ 4,
668
+ 5,
669
+ 6,
670
+ 7
671
+ ],
672
+ "modality_keys": [
673
+ "x",
674
+ "y",
675
+ "z",
676
+ "roll",
677
+ "pitch",
678
+ "yaw",
679
+ "gripper"
680
+ ],
681
+ "sin_cos_embedding_keys": null,
682
+ "mean_std_embedding_keys": [
683
+ "x",
684
+ "y",
685
+ "z",
686
+ "roll",
687
+ "pitch",
688
+ "yaw"
689
+ ],
690
+ "action_configs": null
691
+ },
692
+ "language": {
693
+ "delta_indices": [
694
+ 0
695
+ ],
696
+ "modality_keys": [
697
+ "annotation.human.action.task_description"
698
+ ],
699
+ "sin_cos_embedding_keys": null,
700
+ "mean_std_embedding_keys": null,
701
+ "action_configs": null
702
+ }
703
+ },
704
+ "oxe_droid": {
705
+ "video": {
706
+ "delta_indices": [
707
+ 0
708
+ ],
709
+ "modality_keys": [
710
+ "exterior_image_1_left",
711
+ "wrist_image_left"
712
+ ],
713
+ "sin_cos_embedding_keys": null,
714
+ "mean_std_embedding_keys": null,
715
+ "action_configs": null
716
+ },
717
+ "state": {
718
+ "delta_indices": [
719
+ 0
720
+ ],
721
+ "modality_keys": [
722
+ "joint_position",
723
+ "gripper_position"
724
+ ],
725
+ "sin_cos_embedding_keys": null,
726
+ "mean_std_embedding_keys": null,
727
+ "action_configs": null
728
+ },
729
+ "action": {
730
+ "delta_indices": [
731
+ 0,
732
+ 1,
733
+ 2,
734
+ 3,
735
+ 4,
736
+ 5,
737
+ 6,
738
+ 7,
739
+ 8,
740
+ 9,
741
+ 10,
742
+ 11,
743
+ 12,
744
+ 13,
745
+ 14,
746
+ 15,
747
+ 16,
748
+ 17,
749
+ 18,
750
+ 19,
751
+ 20,
752
+ 21,
753
+ 22,
754
+ 23,
755
+ 24,
756
+ 25,
757
+ 26,
758
+ 27,
759
+ 28,
760
+ 29,
761
+ 30,
762
+ 31
763
+ ],
764
+ "modality_keys": [
765
+ "joint_position",
766
+ "gripper_position"
767
+ ],
768
+ "sin_cos_embedding_keys": null,
769
+ "mean_std_embedding_keys": null,
770
+ "action_configs": [
771
+ {
772
+ "rep": "RELATIVE",
773
+ "type": "NON_EEF",
774
+ "format": "DEFAULT",
775
+ "state_key": null
776
+ },
777
+ {
778
+ "rep": "ABSOLUTE",
779
+ "type": "NON_EEF",
780
+ "format": "DEFAULT",
781
+ "state_key": null
782
+ }
783
+ ]
784
+ },
785
+ "language": {
786
+ "delta_indices": [
787
+ 0
788
+ ],
789
+ "modality_keys": [
790
+ "annotation.language.language_instruction"
791
+ ],
792
+ "sin_cos_embedding_keys": null,
793
+ "mean_std_embedding_keys": null,
794
+ "action_configs": null
795
+ }
796
+ },
797
+ "new_embodiment": {
798
+ "video": {
799
+ "delta_indices": [
800
+ 0
801
+ ],
802
+ "modality_keys": [
803
+ "robot0_eye_in_hand",
804
+ "robot0_agentview_left",
805
+ "robot0_agentview_right"
806
+ ],
807
+ "sin_cos_embedding_keys": null,
808
+ "mean_std_embedding_keys": null,
809
+ "action_configs": null
810
+ },
811
+ "state": {
812
+ "delta_indices": [
813
+ 0
814
+ ],
815
+ "modality_keys": [
816
+ "base_position",
817
+ "base_rotation",
818
+ "end_effector_position_relative",
819
+ "end_effector_rotation_relative",
820
+ "gripper_qpos"
821
+ ],
822
+ "sin_cos_embedding_keys": null,
823
+ "mean_std_embedding_keys": null,
824
+ "action_configs": null
825
+ },
826
+ "action": {
827
+ "delta_indices": [
828
+ 0,
829
+ 1,
830
+ 2,
831
+ 3,
832
+ 4,
833
+ 5,
834
+ 6,
835
+ 7,
836
+ 8,
837
+ 9,
838
+ 10,
839
+ 11,
840
+ 12,
841
+ 13,
842
+ 14,
843
+ 15
844
+ ],
845
+ "modality_keys": [
846
+ "base_motion",
847
+ "control_mode",
848
+ "end_effector_position",
849
+ "end_effector_rotation",
850
+ "gripper_close"
851
+ ],
852
+ "sin_cos_embedding_keys": null,
853
+ "mean_std_embedding_keys": null,
854
+ "action_configs": [
855
+ {
856
+ "rep": "ABSOLUTE",
857
+ "type": "NON_EEF",
858
+ "format": "DEFAULT",
859
+ "state_key": null
860
+ },
861
+ {
862
+ "rep": "ABSOLUTE",
863
+ "type": "NON_EEF",
864
+ "format": "DEFAULT",
865
+ "state_key": null
866
+ },
867
+ {
868
+ "rep": "ABSOLUTE",
869
+ "type": "EEF",
870
+ "format": "DEFAULT",
871
+ "state_key": null
872
+ },
873
+ {
874
+ "rep": "ABSOLUTE",
875
+ "type": "EEF",
876
+ "format": "DEFAULT",
877
+ "state_key": null
878
+ },
879
+ {
880
+ "rep": "ABSOLUTE",
881
+ "type": "NON_EEF",
882
+ "format": "DEFAULT",
883
+ "state_key": null
884
+ }
885
+ ]
886
+ },
887
+ "language": {
888
+ "delta_indices": [
889
+ 0
890
+ ],
891
+ "modality_keys": [
892
+ "task"
893
+ ],
894
+ "sin_cos_embedding_keys": null,
895
+ "mean_std_embedding_keys": null,
896
+ "action_configs": null
897
+ }
898
+ }
899
+ },
900
+ "image_crop_size": null,
901
+ "image_target_size": null,
902
+ "use_albumentations": true,
903
+ "random_rotation_angle": null,
904
+ "color_jitter_params": {
905
+ "brightness": 0.3,
906
+ "contrast": 0.4,
907
+ "saturation": 0.5,
908
+ "hue": 0.08
909
+ },
910
+ "shortest_image_edge": 256,
911
+ "crop_fraction": 0.95,
912
+ "model_name": "nvidia/Eagle-Block2A-2B-v2",
913
+ "model_type": "eagle",
914
+ "formalize_language": true,
915
+ "max_state_dim": 128,
916
+ "max_action_dim": 128,
917
+ "max_action_horizon": 50,
918
+ "use_percentiles": false,
919
+ "clip_outliers": true,
920
+ "apply_sincos_state_encoding": true,
921
+ "use_relative_action": true
922
+ }
923
+ }
checkpoint-60000/rng_state_0.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:35dbe0a1e05f4377fe7af200bb237f2fd4fcf1f626c3ada67619a4ed5118fc5e
3
+ size 14917
checkpoint-60000/rng_state_1.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:80cdf2673e0d1e3845df5fc505dd94474498607a324d4d97974e55d06548c193
3
+ size 14917
checkpoint-60000/scheduler.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c55544f98d67f24204a10b671ebe610ad828fbb0acca03282f9d33f0f6640e8e
3
+ size 1465
checkpoint-60000/statistics.json ADDED
The diff for this file is too large to render. See raw diff
 
checkpoint-60000/trainer_state.json ADDED
The diff for this file is too large to render. See raw diff
 
checkpoint-60000/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f2763d70cd1de320f4d0515e426de65befdc4afd41cd3575f080430938bc410b
3
+ size 7633
checkpoint-60000/zero_to_fp32.py ADDED
@@ -0,0 +1,760 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example:
14
+ # python zero_to_fp32.py . output_dir/
15
+ # or
16
+ # python zero_to_fp32.py . output_dir/ --safe_serialization
17
+
18
+ import argparse
19
+ import torch
20
+ import glob
21
+ import math
22
+ import os
23
+ import re
24
+ import gc
25
+ import json
26
+ import numpy as np
27
+ from tqdm import tqdm
28
+ from collections import OrderedDict
29
+ from dataclasses import dataclass
30
+
31
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
32
+ # DeepSpeed data structures it has to be available in the current python environment.
33
+ from deepspeed.utils import logger
34
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
35
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
36
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
37
+
38
+
39
+ @dataclass
40
+ class zero_model_state:
41
+ buffers: dict()
42
+ param_shapes: dict()
43
+ shared_params: list
44
+ ds_version: int
45
+ frozen_param_shapes: dict()
46
+ frozen_param_fragments: dict()
47
+
48
+
49
+ debug = 0
50
+
51
+ # load to cpu
52
+ device = torch.device('cpu')
53
+
54
+
55
+ def atoi(text):
56
+ return int(text) if text.isdigit() else text
57
+
58
+
59
+ def natural_keys(text):
60
+ '''
61
+ alist.sort(key=natural_keys) sorts in human order
62
+ http://nedbatchelder.com/blog/200712/human_sorting.html
63
+ (See Toothy's implementation in the comments)
64
+ '''
65
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
66
+
67
+
68
+ def get_model_state_file(checkpoint_dir, zero_stage):
69
+ if not os.path.isdir(checkpoint_dir):
70
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
71
+
72
+ # there should be only one file
73
+ if zero_stage <= 2:
74
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
75
+ elif zero_stage == 3:
76
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
77
+
78
+ if not os.path.exists(file):
79
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
80
+
81
+ return file
82
+
83
+
84
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
85
+ # XXX: need to test that this simple glob rule works for multi-node setup too
86
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
87
+
88
+ if len(ckpt_files) == 0:
89
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
90
+
91
+ return ckpt_files
92
+
93
+
94
+ def get_optim_files(checkpoint_dir):
95
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
96
+
97
+
98
+ def get_model_state_files(checkpoint_dir):
99
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
100
+
101
+
102
+ def parse_model_states(files):
103
+ zero_model_states = []
104
+ for file in files:
105
+ state_dict = torch.load(file, map_location=device, weights_only=False)
106
+
107
+ if BUFFER_NAMES not in state_dict:
108
+ raise ValueError(f"{file} is not a model state checkpoint")
109
+ buffer_names = state_dict[BUFFER_NAMES]
110
+ if debug:
111
+ print("Found buffers:", buffer_names)
112
+
113
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
114
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
115
+ param_shapes = state_dict[PARAM_SHAPES]
116
+
117
+ # collect parameters that are included in param_shapes
118
+ param_names = []
119
+ for s in param_shapes:
120
+ for name in s.keys():
121
+ param_names.append(name)
122
+
123
+ # update with frozen parameters
124
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
125
+ if frozen_param_shapes is not None:
126
+ if debug:
127
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
128
+ param_names += list(frozen_param_shapes.keys())
129
+
130
+ # handle shared params
131
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
132
+
133
+ ds_version = state_dict.get(DS_VERSION, None)
134
+
135
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
136
+
137
+ z_model_state = zero_model_state(buffers=buffers,
138
+ param_shapes=param_shapes,
139
+ shared_params=shared_params,
140
+ ds_version=ds_version,
141
+ frozen_param_shapes=frozen_param_shapes,
142
+ frozen_param_fragments=frozen_param_fragments)
143
+ zero_model_states.append(z_model_state)
144
+
145
+ return zero_model_states
146
+
147
+
148
+ def parse_optim_states(files, ds_checkpoint_dir):
149
+ total_files = len(files)
150
+ state_dicts = []
151
+ for f in tqdm(files, desc='Loading checkpoint shards'):
152
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
153
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
154
+ # and also handle the case where it was already removed by another helper script
155
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
156
+ state_dicts.append(state_dict)
157
+
158
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
159
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
160
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
161
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
162
+
163
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
164
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
165
+ # use the max of the partition_count to get the dp world_size.
166
+
167
+ if type(world_size) is list:
168
+ world_size = max(world_size)
169
+
170
+ if world_size != total_files:
171
+ raise ValueError(
172
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
173
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
174
+ )
175
+
176
+ # the groups are named differently in each stage
177
+ if zero_stage <= 2:
178
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
179
+ elif zero_stage == 3:
180
+ fp32_groups_key = FP32_FLAT_GROUPS
181
+ else:
182
+ raise ValueError(f"unknown zero stage {zero_stage}")
183
+
184
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
185
+ return zero_stage, world_size, fp32_flat_groups
186
+
187
+
188
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
189
+ """
190
+ Returns fp32 state_dict reconstructed from ds checkpoint
191
+
192
+ Args:
193
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
194
+
195
+ """
196
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
197
+
198
+ optim_files = get_optim_files(ds_checkpoint_dir)
199
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
200
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
201
+
202
+ model_files = get_model_state_files(ds_checkpoint_dir)
203
+
204
+ zero_model_states = parse_model_states(model_files)
205
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
206
+
207
+ if zero_stage <= 2:
208
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
209
+ exclude_frozen_parameters)
210
+ elif zero_stage == 3:
211
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
212
+ exclude_frozen_parameters)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _has_callable(obj, fn):
248
+ attr = getattr(obj, fn, None)
249
+ return callable(attr)
250
+
251
+
252
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
253
+ param_shapes = zero_model_states[0].param_shapes
254
+
255
+ # Reconstruction protocol:
256
+ #
257
+ # XXX: document this
258
+
259
+ if debug:
260
+ for i in range(world_size):
261
+ for j in range(len(fp32_flat_groups[0])):
262
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
263
+
264
+ # XXX: memory usage doubles here (zero2)
265
+ num_param_groups = len(fp32_flat_groups[0])
266
+ merged_single_partition_of_fp32_groups = []
267
+ for i in range(num_param_groups):
268
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
269
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
270
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
271
+ avail_numel = sum(
272
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
273
+
274
+ if debug:
275
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
276
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
277
+ # not asserting if there is a mismatch due to possible padding
278
+ print(f"Have {avail_numel} numels to process.")
279
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
280
+
281
+ # params
282
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
283
+ # out-of-core computing solution
284
+ total_numel = 0
285
+ total_params = 0
286
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
287
+ offset = 0
288
+ avail_numel = full_single_fp32_vector.numel()
289
+ for name, shape in shapes.items():
290
+
291
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
292
+ total_numel += unpartitioned_numel
293
+ total_params += 1
294
+
295
+ if debug:
296
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
297
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
298
+ offset += unpartitioned_numel
299
+
300
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
301
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
302
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
303
+ # live optimizer object, so we are checking that the numbers are within the right range
304
+ align_to = 2 * world_size
305
+
306
+ def zero2_align(x):
307
+ return align_to * math.ceil(x / align_to)
308
+
309
+ if debug:
310
+ print(f"original offset={offset}, avail_numel={avail_numel}")
311
+
312
+ offset = zero2_align(offset)
313
+ avail_numel = zero2_align(avail_numel)
314
+
315
+ if debug:
316
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
317
+
318
+ # Sanity check
319
+ if offset != avail_numel:
320
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
321
+
322
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
323
+
324
+
325
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
326
+ exclude_frozen_parameters):
327
+ state_dict = OrderedDict()
328
+
329
+ # buffers
330
+ buffers = zero_model_states[0].buffers
331
+ state_dict.update(buffers)
332
+ if debug:
333
+ print(f"added {len(buffers)} buffers")
334
+
335
+ if not exclude_frozen_parameters:
336
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
337
+
338
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
339
+
340
+ # recover shared parameters
341
+ for pair in zero_model_states[0].shared_params:
342
+ if pair[1] in state_dict:
343
+ state_dict[pair[0]] = state_dict[pair[1]]
344
+
345
+ return state_dict
346
+
347
+
348
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
349
+ remainder = unpartitioned_numel % world_size
350
+ padding_numel = (world_size - remainder) if remainder else 0
351
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
352
+ return partitioned_numel, padding_numel
353
+
354
+
355
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
356
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
357
+ return
358
+
359
+ if debug:
360
+ for i in range(world_size):
361
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
362
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
363
+
364
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
365
+ wanted_params = len(frozen_param_shapes)
366
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
367
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
368
+ print(f'Frozen params: Have {avail_numel} numels to process.')
369
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
370
+
371
+ total_params = 0
372
+ total_numel = 0
373
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
374
+ total_params += 1
375
+ unpartitioned_numel = shape.numel()
376
+ total_numel += unpartitioned_numel
377
+
378
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
379
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
380
+
381
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
382
+
383
+ if debug:
384
+ print(
385
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
386
+ )
387
+
388
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
389
+
390
+
391
+ class GatheredTensor:
392
+ """
393
+ A pseudo tensor that collects partitioned weights.
394
+ It is more memory efficient when there are multiple groups.
395
+ """
396
+
397
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
398
+ self.flat_groups = flat_groups
399
+ self.flat_groups_offset = flat_groups_offset
400
+ self.offset = offset
401
+ self.partitioned_numel = partitioned_numel
402
+ self.shape = shape
403
+ self.dtype = self.flat_groups[0][0].dtype
404
+
405
+ def contiguous(self):
406
+ """
407
+ Merge partitioned weights from flat_groups into a single tensor.
408
+ """
409
+ end_idx = self.offset + self.partitioned_numel
410
+ world_size = len(self.flat_groups)
411
+ pad_flat_param_chunks = []
412
+
413
+ for rank_i in range(world_size):
414
+ # for each rank, we need to collect weights from related group/groups
415
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
416
+ start_group_id = None
417
+ end_group_id = None
418
+ for group_id in range(len(self.flat_groups_offset)):
419
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
420
+ start_group_id = group_id
421
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
422
+ end_group_id = group_id
423
+ break
424
+ # collect weights from related group/groups
425
+ for group_id in range(start_group_id, end_group_id + 1):
426
+ flat_tensor = flat_groups_at_rank_i[group_id]
427
+ start_offset = self.offset - self.flat_groups_offset[group_id]
428
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
429
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
430
+
431
+ # collect weights from all ranks
432
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
433
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
434
+ return param
435
+
436
+
437
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
438
+ param_shapes = zero_model_states[0].param_shapes
439
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
440
+
441
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
442
+ # param, re-consolidating each param, while dealing with padding if any
443
+
444
+ # merge list of dicts, preserving order
445
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
446
+
447
+ if debug:
448
+ for i in range(world_size):
449
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
450
+
451
+ wanted_params = len(param_shapes)
452
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
453
+ # not asserting if there is a mismatch due to possible padding
454
+ avail_numel = fp32_flat_groups[0].numel() * world_size
455
+ print(f"Trainable params: Have {avail_numel} numels to process.")
456
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
457
+
458
+ # params
459
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
460
+ # out-of-core computing solution
461
+ offset = 0
462
+ total_numel = 0
463
+ total_params = 0
464
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
465
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
466
+ unpartitioned_numel = shape.numel()
467
+ total_numel += unpartitioned_numel
468
+ total_params += 1
469
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
470
+
471
+ if debug:
472
+ print(
473
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
474
+ )
475
+
476
+ # memory efficient tensor
477
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
478
+ state_dict[name] = tensor
479
+ offset += partitioned_numel
480
+
481
+ offset *= world_size
482
+
483
+ # Sanity check
484
+ if offset != avail_numel:
485
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
486
+
487
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
488
+
489
+
490
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
491
+ exclude_frozen_parameters):
492
+ state_dict = OrderedDict()
493
+
494
+ # buffers
495
+ buffers = zero_model_states[0].buffers
496
+ state_dict.update(buffers)
497
+ if debug:
498
+ print(f"added {len(buffers)} buffers")
499
+
500
+ if not exclude_frozen_parameters:
501
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
502
+
503
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
504
+
505
+ # recover shared parameters
506
+ for pair in zero_model_states[0].shared_params:
507
+ if pair[1] in state_dict:
508
+ state_dict[pair[0]] = state_dict[pair[1]]
509
+
510
+ return state_dict
511
+
512
+
513
+ def to_torch_tensor(state_dict, return_empty_tensor=False):
514
+ """
515
+ Convert state_dict of GatheredTensor to torch tensor
516
+ """
517
+ torch_state_dict = {}
518
+ converted_tensors = {}
519
+ for name, tensor in state_dict.items():
520
+ tensor_id = id(tensor)
521
+ if tensor_id in converted_tensors: # shared tensors
522
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
523
+ torch_state_dict[name] = shared_tensor
524
+ else:
525
+ converted_tensors[tensor_id] = name
526
+ if return_empty_tensor:
527
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
528
+ else:
529
+ torch_state_dict[name] = tensor.contiguous()
530
+ return torch_state_dict
531
+
532
+
533
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
534
+ tag=None,
535
+ exclude_frozen_parameters=False,
536
+ lazy_mode=False):
537
+ """
538
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
539
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
540
+ via a model hub.
541
+
542
+ Args:
543
+ - ``checkpoint_dir``: path to the desired checkpoint folder
544
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
545
+ - ``exclude_frozen_parameters``: exclude frozen parameters
546
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
547
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
548
+
549
+ Returns:
550
+ - pytorch ``state_dict``
551
+
552
+ A typical usage might be ::
553
+
554
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
555
+ # do the training and checkpoint saving
556
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
557
+ model = model.cpu() # move to cpu
558
+ model.load_state_dict(state_dict)
559
+ # submit to model hub or save the model to share with others
560
+
561
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
562
+ application. i.e. you will need to re-initialize the deepspeed engine, since
563
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
564
+
565
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
566
+
567
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
568
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
569
+ the checkpoint. Or you can load state_dict in lazy mode ::
570
+
571
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
572
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
573
+ for name, lazy_tensor in state_dict.item():
574
+ tensor = lazy_tensor.contiguous() # to cpu
575
+ print(name, tensor)
576
+ # del tensor to release memory if it no longer in use
577
+ """
578
+ if tag is None:
579
+ latest_path = os.path.join(checkpoint_dir, 'latest')
580
+ if os.path.isfile(latest_path):
581
+ with open(latest_path, 'r') as fd:
582
+ tag = fd.read().strip()
583
+ else:
584
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
585
+
586
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
587
+
588
+ if not os.path.isdir(ds_checkpoint_dir):
589
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
590
+
591
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
592
+ if lazy_mode:
593
+ return state_dict
594
+ else:
595
+ return to_torch_tensor(state_dict)
596
+
597
+
598
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
599
+ output_dir,
600
+ max_shard_size="5GB",
601
+ safe_serialization=False,
602
+ tag=None,
603
+ exclude_frozen_parameters=False):
604
+ """
605
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
606
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
607
+
608
+ Args:
609
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
610
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
611
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
612
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
613
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
614
+ - ``exclude_frozen_parameters``: exclude frozen parameters
615
+ """
616
+
617
+ # Dependency pre-check
618
+ if safe_serialization:
619
+ try:
620
+ from safetensors.torch import save_file
621
+ except ImportError:
622
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
623
+ raise
624
+ if max_shard_size is not None:
625
+ try:
626
+ from huggingface_hub import split_torch_state_dict_into_shards
627
+ except ImportError:
628
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
629
+ raise
630
+
631
+ # Convert zero checkpoint to state_dict
632
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
633
+ tag,
634
+ exclude_frozen_parameters,
635
+ lazy_mode=True)
636
+
637
+ # Shard the model if it is too big.
638
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
639
+ if max_shard_size is not None:
640
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
641
+ # an memory-efficient approach for sharding
642
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
643
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
644
+ filename_pattern=filename_pattern,
645
+ max_shard_size=max_shard_size)
646
+ else:
647
+ from collections import namedtuple
648
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
649
+ state_dict_split = StateDictSplit(is_sharded=False,
650
+ filename_to_tensors={weights_name: list(state_dict.keys())})
651
+
652
+ # Save the model by shard
653
+ os.makedirs(output_dir, exist_ok=True)
654
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
655
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
656
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
657
+ shard_state_dict = to_torch_tensor(shard_state_dict)
658
+ output_path = os.path.join(output_dir, shard_file)
659
+ if safe_serialization:
660
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
661
+ else:
662
+ torch.save(shard_state_dict, output_path)
663
+ # release the memory of current shard
664
+ for tensor_name in list(shard_state_dict.keys()):
665
+ del state_dict[tensor_name]
666
+ del shard_state_dict[tensor_name]
667
+ del shard_state_dict
668
+ gc.collect()
669
+
670
+ # Save index if sharded
671
+ if state_dict_split.is_sharded:
672
+ index = {
673
+ "metadata": state_dict_split.metadata,
674
+ "weight_map": state_dict_split.tensor_to_filename,
675
+ }
676
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
677
+ save_index_file = os.path.join(output_dir, save_index_file)
678
+ with open(save_index_file, "w", encoding="utf-8") as f:
679
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
680
+ f.write(content)
681
+
682
+
683
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
684
+ """
685
+ 1. Put the provided model to cpu
686
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
687
+ 3. Load it into the provided model
688
+
689
+ Args:
690
+ - ``model``: the model object to update
691
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
692
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
693
+
694
+ Returns:
695
+ - ``model`: modified model
696
+
697
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
698
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
699
+ conveniently placed for you in the checkpoint folder.
700
+
701
+ A typical usage might be ::
702
+
703
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
704
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
705
+ # submit to model hub or save the model to share with others
706
+
707
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
708
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
709
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
710
+
711
+ """
712
+ logger.info("Extracting fp32 weights")
713
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
714
+
715
+ logger.info("Overwriting model with fp32 weights")
716
+ model = model.cpu()
717
+ model.load_state_dict(state_dict, strict=False)
718
+
719
+ return model
720
+
721
+
722
+ if __name__ == "__main__":
723
+ parser = argparse.ArgumentParser()
724
+ parser.add_argument("checkpoint_dir",
725
+ type=str,
726
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
727
+ parser.add_argument("output_dir",
728
+ type=str,
729
+ help="directory to the pytorch fp32 state_dict output files"
730
+ "(e.g. path/checkpoint-12-output/)")
731
+ parser.add_argument(
732
+ "--max_shard_size",
733
+ type=str,
734
+ default="5GB",
735
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
736
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
737
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
738
+ "without CPU OOM issues.")
739
+ parser.add_argument(
740
+ "--safe_serialization",
741
+ default=False,
742
+ action='store_true',
743
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
744
+ parser.add_argument("-t",
745
+ "--tag",
746
+ type=str,
747
+ default=None,
748
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
749
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
750
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
751
+ args = parser.parse_args()
752
+
753
+ debug = args.debug
754
+
755
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
756
+ args.output_dir,
757
+ max_shard_size=args.max_shard_size,
758
+ safe_serialization=args.safe_serialization,
759
+ tag=args.tag,
760
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
config.json ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_horizon": 50,
3
+ "add_pos_embed": true,
4
+ "apply_sincos_state_encoding": true,
5
+ "architectures": [
6
+ "Gr00tN1d6"
7
+ ],
8
+ "attn_dropout": 0.2,
9
+ "attn_implementation": null,
10
+ "backbone_embedding_dim": 2048,
11
+ "backbone_model_type": "eagle",
12
+ "backbone_trainable_params_fp32": true,
13
+ "collator_overwrite_image_inputs": false,
14
+ "color_jitter_params": {
15
+ "brightness": 0.1,
16
+ "contrast": 0.1,
17
+ "hue": 0.1,
18
+ "saturation": 0.1
19
+ },
20
+ "crop_fraction": 0.95,
21
+ "diffusion_model_cfg": {
22
+ "attention_head_dim": 48,
23
+ "dropout": 0.2,
24
+ "final_dropout": true,
25
+ "interleave_self_attention": true,
26
+ "norm_type": "ada_norm",
27
+ "num_attention_heads": 32,
28
+ "num_layers": 32,
29
+ "output_dim": 1024,
30
+ "positional_embeddings": null
31
+ },
32
+ "eagle_collator": true,
33
+ "formalize_language": true,
34
+ "gemma_collator": false,
35
+ "hidden_size": 1024,
36
+ "image_crop_size": null,
37
+ "image_target_size": null,
38
+ "input_embedding_dim": 1536,
39
+ "load_bf16": true,
40
+ "max_action_dim": 128,
41
+ "max_num_embodiments": 32,
42
+ "max_seq_len": 1024,
43
+ "max_state_dim": 128,
44
+ "model_dtype": "bfloat16",
45
+ "model_name": "nvidia/Eagle-Block2A-2B-v2",
46
+ "model_type": "Gr00tN1d6",
47
+ "noise_beta_alpha": 1.5,
48
+ "noise_beta_beta": 1.0,
49
+ "noise_s": 0.999,
50
+ "num_inference_timesteps": 4,
51
+ "num_timestep_buckets": 1000,
52
+ "random_rotation_angle": null,
53
+ "reproject_vision": false,
54
+ "select_layer": 16,
55
+ "shortest_image_edge": 256,
56
+ "state_dropout_prob": 0.0,
57
+ "torch_dtype": "bfloat16",
58
+ "transformers_version": "4.51.3",
59
+ "tune_diffusion_model": true,
60
+ "tune_llm": false,
61
+ "tune_projector": true,
62
+ "tune_top_llm_layers": 4,
63
+ "tune_visual": false,
64
+ "tune_vlln": true,
65
+ "use_albumentations_transforms": true,
66
+ "use_alternate_vl_dit": true,
67
+ "use_flash_attention": true,
68
+ "use_relative_action": true,
69
+ "use_vlln": true
70
+ }
experiment_cfg/dataset_statistics.json ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "new_embodiment": {
3
+ "state": {
4
+ "base_position": {
5
+ "min": [
6
+ -4.94293737411499,
7
+ -6.890198230743408,
8
+ 0.6994727253913879
9
+ ],
10
+ "max": [
11
+ 7.778976917266846,
12
+ 0.4876587688922882,
13
+ 0.7223060727119446
14
+ ],
15
+ "mean": [
16
+ 2.6429230985097116,
17
+ -2.1040500635873367,
18
+ 0.7003864151980145
19
+ ],
20
+ "std": [
21
+ 1.6723189550183524,
22
+ 1.340646864108621,
23
+ 0.0010989839745322908
24
+ ],
25
+ "q01": [
26
+ -4.570116996765137,
27
+ -6.843641757965088,
28
+ 0.699881374835968
29
+ ],
30
+ "q99": [
31
+ 7.292758464813232,
32
+ 0.48759591579437256,
33
+ 0.714806854724884
34
+ ]
35
+ },
36
+ "base_rotation": {
37
+ "min": [
38
+ 0.0,
39
+ 0.0,
40
+ -1.0,
41
+ 0.0
42
+ ],
43
+ "max": [
44
+ 0.0,
45
+ 0.0,
46
+ 1.0,
47
+ 1.0
48
+ ],
49
+ "mean": [
50
+ 0.0,
51
+ 0.0,
52
+ 0.23872401952057823,
53
+ 0.6083901817614287
54
+ ],
55
+ "std": [
56
+ 0.0,
57
+ 0.0,
58
+ 0.6620197581758033,
59
+ 0.3668243207825513
60
+ ],
61
+ "q01": [
62
+ 0.0,
63
+ 0.0,
64
+ -1.0,
65
+ 1.9421531760599464e-07
66
+ ],
67
+ "q99": [
68
+ 0.0,
69
+ 0.0,
70
+ 1.0,
71
+ 1.0
72
+ ]
73
+ },
74
+ "end_effector_position_relative": {
75
+ "min": [
76
+ -0.5828549861907959,
77
+ -0.8486236333847046,
78
+ -0.3446122705936432
79
+ ],
80
+ "max": [
81
+ 0.9014528393745422,
82
+ 0.8569459915161133,
83
+ 1.053884744644165
84
+ ],
85
+ "mean": [
86
+ 0.25786905106953806,
87
+ -0.018494133083245046,
88
+ 0.44849795111686513
89
+ ],
90
+ "std": [
91
+ 0.16495694636764374,
92
+ 0.24473163300325904,
93
+ 0.20998219654794578
94
+ ],
95
+ "q01": [
96
+ -0.3761294186115265,
97
+ -0.721174955368042,
98
+ -0.23106376826763153
99
+ ],
100
+ "q99": [
101
+ 0.8378565311431885,
102
+ 0.7546205520629883,
103
+ 0.9249030947685242
104
+ ]
105
+ },
106
+ "end_effector_rotation_relative": {
107
+ "min": [
108
+ -0.9999930262565613,
109
+ -0.9999884366989136,
110
+ -0.9975369572639465,
111
+ 6.754255821306288e-08
112
+ ],
113
+ "max": [
114
+ 0.9999998807907104,
115
+ 0.9997354745864868,
116
+ 0.9900747537612915,
117
+ 0.9735639095306396
118
+ ],
119
+ "mean": [
120
+ -0.261215793438568,
121
+ 0.03690521291069262,
122
+ -0.07182086071331255,
123
+ 0.13733310115690522
124
+ ],
125
+ "std": [
126
+ 0.8253314117717035,
127
+ 0.297636866306729,
128
+ 0.32923964375209464,
129
+ 0.16800963654479095
130
+ ],
131
+ "q01": [
132
+ -0.9995729327201843,
133
+ -0.9585485458374023,
134
+ -0.8866488337516785,
135
+ 0.0003100793983321637
136
+ ],
137
+ "q99": [
138
+ 0.9993855357170105,
139
+ 0.9750741720199585,
140
+ 0.8878775835037231,
141
+ 0.8871033191680908
142
+ ]
143
+ },
144
+ "gripper_qpos": {
145
+ "min": [
146
+ -0.011669117026031017,
147
+ -0.06307756155729294
148
+ ],
149
+ "max": [
150
+ 0.055869169533252716,
151
+ 0.010916369967162609
152
+ ],
153
+ "mean": [
154
+ 0.03205486276150664,
155
+ -0.03223369335466349
156
+ ],
157
+ "std": [
158
+ 0.012312154775568325,
159
+ 0.012107891953724964
160
+ ],
161
+ "q01": [
162
+ -0.0028626956045627594,
163
+ -0.04449209198355675
164
+ ],
165
+ "q99": [
166
+ 0.04488873854279518,
167
+ 0.0015663582598790526
168
+ ]
169
+ }
170
+ },
171
+ "action": {
172
+ "base_motion": {
173
+ "min": [
174
+ -1.0,
175
+ -1.0,
176
+ -1.0,
177
+ 0.0
178
+ ],
179
+ "max": [
180
+ 1.0,
181
+ 1.0,
182
+ 1.0,
183
+ 0.0
184
+ ],
185
+ "mean": [
186
+ 0.02050641685926724,
187
+ 0.001128088574318351,
188
+ -0.002216507356311039,
189
+ 0.0
190
+ ],
191
+ "std": [
192
+ 0.1652318199594288,
193
+ 0.14088974424972592,
194
+ 0.1268873250272064,
195
+ 0.0
196
+ ],
197
+ "q01": [
198
+ -1.0,
199
+ -1.0,
200
+ -1.0,
201
+ 0.0
202
+ ],
203
+ "q99": [
204
+ 1.0,
205
+ 1.0,
206
+ 1.0,
207
+ 0.0
208
+ ]
209
+ },
210
+ "control_mode": {
211
+ "min": [
212
+ -1.0
213
+ ],
214
+ "max": [
215
+ 1.0
216
+ ],
217
+ "mean": [
218
+ -0.8539844503292455
219
+ ],
220
+ "std": [
221
+ 0.5202972484299755
222
+ ],
223
+ "q01": [
224
+ -1.0
225
+ ],
226
+ "q99": [
227
+ 1.0
228
+ ]
229
+ },
230
+ "end_effector_position": {
231
+ "min": [
232
+ -1.0,
233
+ -1.0,
234
+ -1.0
235
+ ],
236
+ "max": [
237
+ 1.0,
238
+ 1.0,
239
+ 1.0
240
+ ],
241
+ "mean": [
242
+ 0.002803475656952085,
243
+ -0.006818774988739351,
244
+ -0.07855993148808858
245
+ ],
246
+ "std": [
247
+ 0.44313822308380246,
248
+ 0.40503200192475836,
249
+ 0.3950365918181045
250
+ ],
251
+ "q01": [
252
+ -1.0,
253
+ -1.0,
254
+ -1.0
255
+ ],
256
+ "q99": [
257
+ 1.0,
258
+ 1.0,
259
+ 1.0
260
+ ]
261
+ },
262
+ "end_effector_rotation": {
263
+ "min": [
264
+ -1.0,
265
+ -1.0,
266
+ -1.0
267
+ ],
268
+ "max": [
269
+ 1.0,
270
+ 1.0,
271
+ 1.0
272
+ ],
273
+ "mean": [
274
+ 0.0075743831831496635,
275
+ -0.022864397394476266,
276
+ -0.0001568091574277115
277
+ ],
278
+ "std": [
279
+ 0.1104170002966275,
280
+ 0.1306405108880139,
281
+ 0.12282959716644956
282
+ ],
283
+ "q01": [
284
+ -0.6857143044471741,
285
+ -0.9028571248054504,
286
+ -1.0
287
+ ],
288
+ "q99": [
289
+ 0.9028571248054504,
290
+ 0.8342857360839844,
291
+ 0.8718582391738892
292
+ ]
293
+ },
294
+ "gripper_close": {
295
+ "min": [
296
+ -1.0
297
+ ],
298
+ "max": [
299
+ 1.0
300
+ ],
301
+ "mean": [
302
+ -0.37004847633194815
303
+ ],
304
+ "std": [
305
+ 0.9290190132822349
306
+ ],
307
+ "q01": [
308
+ -1.0
309
+ ],
310
+ "q99": [
311
+ 1.0
312
+ ]
313
+ }
314
+ },
315
+ "relative_action": {}
316
+ }
317
+ }
experiment_cfg/final_model_config.json ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_type": "Gr00tN1d6",
3
+ "model_dtype": "bfloat16",
4
+ "model_name": "nvidia/Eagle-Block2A-2B-v2",
5
+ "backbone_model_type": "eagle",
6
+ "model_revision": null,
7
+ "tune_top_llm_layers": 4,
8
+ "backbone_embedding_dim": 2048,
9
+ "tune_llm": false,
10
+ "tune_visual": false,
11
+ "select_layer": 16,
12
+ "reproject_vision": false,
13
+ "use_flash_attention": true,
14
+ "load_bf16": true,
15
+ "collator_overwrite_image_inputs": false,
16
+ "eagle_collator": true,
17
+ "backbone_trainable_params_fp32": true,
18
+ "apply_sincos_state_encoding": true,
19
+ "use_relative_action": true,
20
+ "max_state_dim": 128,
21
+ "max_action_dim": 128,
22
+ "action_horizon": 50,
23
+ "hidden_size": 1024,
24
+ "input_embedding_dim": 1536,
25
+ "add_pos_embed": true,
26
+ "attn_dropout": 0.2,
27
+ "use_vlln": true,
28
+ "max_seq_len": 1024,
29
+ "use_alternate_vl_dit": true,
30
+ "attend_text_every_n_blocks": 2,
31
+ "diffusion_model_cfg": {
32
+ "attention_head_dim": 48,
33
+ "dropout": 0.2,
34
+ "final_dropout": true,
35
+ "interleave_self_attention": true,
36
+ "norm_type": "ada_norm",
37
+ "num_attention_heads": 32,
38
+ "num_layers": 32,
39
+ "output_dim": 1024,
40
+ "positional_embeddings": null
41
+ },
42
+ "num_inference_timesteps": 4,
43
+ "noise_beta_alpha": 1.5,
44
+ "noise_beta_beta": 1.0,
45
+ "noise_s": 0.999,
46
+ "num_timestep_buckets": 1000,
47
+ "tune_projector": true,
48
+ "tune_diffusion_model": true,
49
+ "tune_vlln": true,
50
+ "state_dropout_prob": 0.0,
51
+ "state_additive_noise_scale": 0.0,
52
+ "max_num_embodiments": 32
53
+ }
experiment_cfg/final_processor_config.json ADDED
The diff for this file is too large to render. See raw diff
 
model-00001-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8fd3f824fb806d8bcd6c308b99c9fb3bc02ab201cf161343b452a60beb361255
3
+ size 4991091456
model-00002-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9bb5cc1c97053e21fded718b9138232709488ac911d0116d4c346f299662ba3b
3
+ size 1582283096
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
processor/embodiment_id.json ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "robocasa_panda_omron": 13,
3
+ "gr1": 20,
4
+ "behavior_r1_pro": 24,
5
+ "unitree_g1": 8,
6
+ "oxe_google": 0,
7
+ "oxe_widowx": 1,
8
+ "libero_panda": 2,
9
+ "oxe_droid": 16,
10
+ "new_embodiment": 10
11
+ }
processor/processor_config.json ADDED
@@ -0,0 +1,923 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "processor_class": "Gr00tN1d6Processor",
3
+ "processor_kwargs": {
4
+ "modality_configs": {
5
+ "behavior_r1_pro": {
6
+ "video": {
7
+ "delta_indices": [
8
+ 0
9
+ ],
10
+ "modality_keys": [
11
+ "observation.images.rgb.head_256_256",
12
+ "observation.images.rgb.left_wrist_256_256",
13
+ "observation.images.rgb.right_wrist_256_256"
14
+ ],
15
+ "sin_cos_embedding_keys": null,
16
+ "mean_std_embedding_keys": null,
17
+ "action_configs": null
18
+ },
19
+ "state": {
20
+ "delta_indices": [
21
+ 0
22
+ ],
23
+ "modality_keys": [
24
+ "robot_pos",
25
+ "robot_ori_cos",
26
+ "robot_ori_sin",
27
+ "robot_2d_ori",
28
+ "robot_2d_ori_cos",
29
+ "robot_2d_ori_sin",
30
+ "robot_lin_vel",
31
+ "robot_ang_vel",
32
+ "arm_left_qpos",
33
+ "arm_left_qpos_sin",
34
+ "arm_left_qpos_cos",
35
+ "eef_left_pos",
36
+ "eef_left_quat",
37
+ "gripper_left_qpos",
38
+ "arm_right_qpos",
39
+ "arm_right_qpos_sin",
40
+ "arm_right_qpos_cos",
41
+ "eef_right_pos",
42
+ "eef_right_quat",
43
+ "gripper_right_qpos",
44
+ "trunk_qpos"
45
+ ],
46
+ "sin_cos_embedding_keys": null,
47
+ "mean_std_embedding_keys": null,
48
+ "action_configs": null
49
+ },
50
+ "action": {
51
+ "delta_indices": [
52
+ 0,
53
+ 1,
54
+ 2,
55
+ 3,
56
+ 4,
57
+ 5,
58
+ 6,
59
+ 7,
60
+ 8,
61
+ 9,
62
+ 10,
63
+ 11,
64
+ 12,
65
+ 13,
66
+ 14,
67
+ 15,
68
+ 16,
69
+ 17,
70
+ 18,
71
+ 19,
72
+ 20,
73
+ 21,
74
+ 22,
75
+ 23,
76
+ 24,
77
+ 25,
78
+ 26,
79
+ 27,
80
+ 28,
81
+ 29,
82
+ 30,
83
+ 31
84
+ ],
85
+ "modality_keys": [
86
+ "base",
87
+ "torso",
88
+ "left_arm",
89
+ "left_gripper",
90
+ "right_arm",
91
+ "right_gripper"
92
+ ],
93
+ "sin_cos_embedding_keys": null,
94
+ "mean_std_embedding_keys": null,
95
+ "action_configs": [
96
+ {
97
+ "rep": "ABSOLUTE",
98
+ "type": "NON_EEF",
99
+ "format": "DEFAULT",
100
+ "state_key": null
101
+ },
102
+ {
103
+ "rep": "RELATIVE",
104
+ "type": "NON_EEF",
105
+ "format": "DEFAULT",
106
+ "state_key": "trunk_qpos"
107
+ },
108
+ {
109
+ "rep": "RELATIVE",
110
+ "type": "NON_EEF",
111
+ "format": "DEFAULT",
112
+ "state_key": "arm_left_qpos"
113
+ },
114
+ {
115
+ "rep": "ABSOLUTE",
116
+ "type": "NON_EEF",
117
+ "format": "DEFAULT",
118
+ "state_key": null
119
+ },
120
+ {
121
+ "rep": "RELATIVE",
122
+ "type": "NON_EEF",
123
+ "format": "DEFAULT",
124
+ "state_key": "arm_right_qpos"
125
+ },
126
+ {
127
+ "rep": "ABSOLUTE",
128
+ "type": "NON_EEF",
129
+ "format": "DEFAULT",
130
+ "state_key": null
131
+ }
132
+ ]
133
+ },
134
+ "language": {
135
+ "delta_indices": [
136
+ 0
137
+ ],
138
+ "modality_keys": [
139
+ "annotation.human.coarse_action"
140
+ ],
141
+ "sin_cos_embedding_keys": null,
142
+ "mean_std_embedding_keys": null,
143
+ "action_configs": null
144
+ }
145
+ },
146
+ "gr1": {
147
+ "video": {
148
+ "delta_indices": [
149
+ 0
150
+ ],
151
+ "modality_keys": [
152
+ "ego_view_bg_crop_pad_res256_freq20"
153
+ ],
154
+ "sin_cos_embedding_keys": null,
155
+ "mean_std_embedding_keys": null,
156
+ "action_configs": null
157
+ },
158
+ "state": {
159
+ "delta_indices": [
160
+ 0
161
+ ],
162
+ "modality_keys": [
163
+ "left_arm",
164
+ "right_arm",
165
+ "left_hand",
166
+ "right_hand",
167
+ "waist"
168
+ ],
169
+ "sin_cos_embedding_keys": [
170
+ "left_arm",
171
+ "right_arm",
172
+ "left_hand",
173
+ "right_hand",
174
+ "waist"
175
+ ],
176
+ "mean_std_embedding_keys": null,
177
+ "action_configs": null
178
+ },
179
+ "action": {
180
+ "delta_indices": [
181
+ 0,
182
+ 1,
183
+ 2,
184
+ 3,
185
+ 4,
186
+ 5,
187
+ 6,
188
+ 7,
189
+ 8,
190
+ 9,
191
+ 10,
192
+ 11,
193
+ 12,
194
+ 13,
195
+ 14,
196
+ 15
197
+ ],
198
+ "modality_keys": [
199
+ "left_arm",
200
+ "right_arm",
201
+ "left_hand",
202
+ "right_hand",
203
+ "waist"
204
+ ],
205
+ "sin_cos_embedding_keys": null,
206
+ "mean_std_embedding_keys": null,
207
+ "action_configs": [
208
+ {
209
+ "rep": "RELATIVE",
210
+ "type": "NON_EEF",
211
+ "format": "DEFAULT",
212
+ "state_key": null
213
+ },
214
+ {
215
+ "rep": "RELATIVE",
216
+ "type": "NON_EEF",
217
+ "format": "DEFAULT",
218
+ "state_key": null
219
+ },
220
+ {
221
+ "rep": "RELATIVE",
222
+ "type": "NON_EEF",
223
+ "format": "DEFAULT",
224
+ "state_key": null
225
+ },
226
+ {
227
+ "rep": "RELATIVE",
228
+ "type": "NON_EEF",
229
+ "format": "DEFAULT",
230
+ "state_key": null
231
+ },
232
+ {
233
+ "rep": "ABSOLUTE",
234
+ "type": "NON_EEF",
235
+ "format": "DEFAULT",
236
+ "state_key": null
237
+ }
238
+ ]
239
+ },
240
+ "language": {
241
+ "delta_indices": [
242
+ 0
243
+ ],
244
+ "modality_keys": [
245
+ "task"
246
+ ],
247
+ "sin_cos_embedding_keys": null,
248
+ "mean_std_embedding_keys": null,
249
+ "action_configs": null
250
+ }
251
+ },
252
+ "robocasa_panda_omron": {
253
+ "video": {
254
+ "delta_indices": [
255
+ 0
256
+ ],
257
+ "modality_keys": [
258
+ "robot0_eye_in_hand",
259
+ "robot0_agentview_left",
260
+ "robot0_agentview_right"
261
+ ],
262
+ "sin_cos_embedding_keys": null,
263
+ "mean_std_embedding_keys": null,
264
+ "action_configs": null
265
+ },
266
+ "state": {
267
+ "delta_indices": [
268
+ 0
269
+ ],
270
+ "modality_keys": [
271
+ "base_position",
272
+ "base_rotation",
273
+ "end_effector_position_relative",
274
+ "end_effector_rotation_relative",
275
+ "gripper_qpos"
276
+ ],
277
+ "sin_cos_embedding_keys": null,
278
+ "mean_std_embedding_keys": null,
279
+ "action_configs": null
280
+ },
281
+ "action": {
282
+ "delta_indices": [
283
+ 0,
284
+ 1,
285
+ 2,
286
+ 3,
287
+ 4,
288
+ 5,
289
+ 6,
290
+ 7,
291
+ 8,
292
+ 9,
293
+ 10,
294
+ 11,
295
+ 12,
296
+ 13,
297
+ 14,
298
+ 15
299
+ ],
300
+ "modality_keys": [
301
+ "base_motion",
302
+ "control_mode",
303
+ "end_effector_position",
304
+ "end_effector_rotation",
305
+ "gripper_close"
306
+ ],
307
+ "sin_cos_embedding_keys": null,
308
+ "mean_std_embedding_keys": null,
309
+ "action_configs": [
310
+ {
311
+ "rep": "ABSOLUTE",
312
+ "type": "NON_EEF",
313
+ "format": "DEFAULT",
314
+ "state_key": null
315
+ },
316
+ {
317
+ "rep": "ABSOLUTE",
318
+ "type": "NON_EEF",
319
+ "format": "DEFAULT",
320
+ "state_key": null
321
+ },
322
+ {
323
+ "rep": "ABSOLUTE",
324
+ "type": "EEF",
325
+ "format": "DEFAULT",
326
+ "state_key": null
327
+ },
328
+ {
329
+ "rep": "ABSOLUTE",
330
+ "type": "EEF",
331
+ "format": "DEFAULT",
332
+ "state_key": null
333
+ },
334
+ {
335
+ "rep": "ABSOLUTE",
336
+ "type": "NON_EEF",
337
+ "format": "DEFAULT",
338
+ "state_key": null
339
+ }
340
+ ]
341
+ },
342
+ "language": {
343
+ "delta_indices": [
344
+ 0
345
+ ],
346
+ "modality_keys": [
347
+ "task"
348
+ ],
349
+ "sin_cos_embedding_keys": null,
350
+ "mean_std_embedding_keys": null,
351
+ "action_configs": null
352
+ }
353
+ },
354
+ "unitree_g1": {
355
+ "video": {
356
+ "delta_indices": [
357
+ 0
358
+ ],
359
+ "modality_keys": [
360
+ "ego_view"
361
+ ],
362
+ "sin_cos_embedding_keys": null,
363
+ "mean_std_embedding_keys": null,
364
+ "action_configs": null
365
+ },
366
+ "state": {
367
+ "delta_indices": [
368
+ 0
369
+ ],
370
+ "modality_keys": [
371
+ "left_leg",
372
+ "right_leg",
373
+ "waist",
374
+ "left_arm",
375
+ "right_arm",
376
+ "left_hand",
377
+ "right_hand"
378
+ ],
379
+ "sin_cos_embedding_keys": null,
380
+ "mean_std_embedding_keys": null,
381
+ "action_configs": null
382
+ },
383
+ "action": {
384
+ "delta_indices": [
385
+ 0,
386
+ 1,
387
+ 2,
388
+ 3,
389
+ 4,
390
+ 5,
391
+ 6,
392
+ 7,
393
+ 8,
394
+ 9,
395
+ 10,
396
+ 11,
397
+ 12,
398
+ 13,
399
+ 14,
400
+ 15,
401
+ 16,
402
+ 17,
403
+ 18,
404
+ 19,
405
+ 20,
406
+ 21,
407
+ 22,
408
+ 23,
409
+ 24,
410
+ 25,
411
+ 26,
412
+ 27,
413
+ 28,
414
+ 29
415
+ ],
416
+ "modality_keys": [
417
+ "left_arm",
418
+ "right_arm",
419
+ "left_hand",
420
+ "right_hand",
421
+ "waist",
422
+ "base_height_command",
423
+ "navigate_command"
424
+ ],
425
+ "sin_cos_embedding_keys": null,
426
+ "mean_std_embedding_keys": null,
427
+ "action_configs": [
428
+ {
429
+ "rep": "RELATIVE",
430
+ "type": "NON_EEF",
431
+ "format": "DEFAULT",
432
+ "state_key": null
433
+ },
434
+ {
435
+ "rep": "RELATIVE",
436
+ "type": "NON_EEF",
437
+ "format": "DEFAULT",
438
+ "state_key": null
439
+ },
440
+ {
441
+ "rep": "ABSOLUTE",
442
+ "type": "NON_EEF",
443
+ "format": "DEFAULT",
444
+ "state_key": null
445
+ },
446
+ {
447
+ "rep": "ABSOLUTE",
448
+ "type": "NON_EEF",
449
+ "format": "DEFAULT",
450
+ "state_key": null
451
+ },
452
+ {
453
+ "rep": "ABSOLUTE",
454
+ "type": "NON_EEF",
455
+ "format": "DEFAULT",
456
+ "state_key": null
457
+ },
458
+ {
459
+ "rep": "ABSOLUTE",
460
+ "type": "NON_EEF",
461
+ "format": "DEFAULT",
462
+ "state_key": null
463
+ },
464
+ {
465
+ "rep": "ABSOLUTE",
466
+ "type": "NON_EEF",
467
+ "format": "DEFAULT",
468
+ "state_key": null
469
+ }
470
+ ]
471
+ },
472
+ "language": {
473
+ "delta_indices": [
474
+ 0
475
+ ],
476
+ "modality_keys": [
477
+ "annotation.human.task_description"
478
+ ],
479
+ "sin_cos_embedding_keys": null,
480
+ "mean_std_embedding_keys": null,
481
+ "action_configs": null
482
+ }
483
+ },
484
+ "libero_panda": {
485
+ "video": {
486
+ "delta_indices": [
487
+ 0
488
+ ],
489
+ "modality_keys": [
490
+ "image",
491
+ "wrist_image"
492
+ ],
493
+ "sin_cos_embedding_keys": null,
494
+ "mean_std_embedding_keys": null,
495
+ "action_configs": null
496
+ },
497
+ "state": {
498
+ "delta_indices": [
499
+ 0
500
+ ],
501
+ "modality_keys": [
502
+ "x",
503
+ "y",
504
+ "z",
505
+ "roll",
506
+ "pitch",
507
+ "yaw",
508
+ "gripper"
509
+ ],
510
+ "sin_cos_embedding_keys": null,
511
+ "mean_std_embedding_keys": null,
512
+ "action_configs": null
513
+ },
514
+ "action": {
515
+ "delta_indices": [
516
+ 0,
517
+ 1,
518
+ 2,
519
+ 3,
520
+ 4,
521
+ 5,
522
+ 6,
523
+ 7,
524
+ 8,
525
+ 9,
526
+ 10,
527
+ 11,
528
+ 12,
529
+ 13,
530
+ 14,
531
+ 15
532
+ ],
533
+ "modality_keys": [
534
+ "x",
535
+ "y",
536
+ "z",
537
+ "roll",
538
+ "pitch",
539
+ "yaw",
540
+ "gripper"
541
+ ],
542
+ "sin_cos_embedding_keys": null,
543
+ "mean_std_embedding_keys": null,
544
+ "action_configs": null
545
+ },
546
+ "language": {
547
+ "delta_indices": [
548
+ 0
549
+ ],
550
+ "modality_keys": [
551
+ "annotation.human.action.task_description"
552
+ ],
553
+ "sin_cos_embedding_keys": null,
554
+ "mean_std_embedding_keys": null,
555
+ "action_configs": null
556
+ }
557
+ },
558
+ "oxe_widowx": {
559
+ "video": {
560
+ "delta_indices": [
561
+ 0
562
+ ],
563
+ "modality_keys": [
564
+ "image_0"
565
+ ],
566
+ "sin_cos_embedding_keys": null,
567
+ "mean_std_embedding_keys": null,
568
+ "action_configs": null
569
+ },
570
+ "state": {
571
+ "delta_indices": [
572
+ 0
573
+ ],
574
+ "modality_keys": [
575
+ "x",
576
+ "y",
577
+ "z",
578
+ "roll",
579
+ "pitch",
580
+ "yaw",
581
+ "pad",
582
+ "gripper"
583
+ ],
584
+ "sin_cos_embedding_keys": null,
585
+ "mean_std_embedding_keys": null,
586
+ "action_configs": null
587
+ },
588
+ "action": {
589
+ "delta_indices": [
590
+ 0,
591
+ 1,
592
+ 2,
593
+ 3,
594
+ 4,
595
+ 5,
596
+ 6,
597
+ 7
598
+ ],
599
+ "modality_keys": [
600
+ "x",
601
+ "y",
602
+ "z",
603
+ "roll",
604
+ "pitch",
605
+ "yaw",
606
+ "gripper"
607
+ ],
608
+ "sin_cos_embedding_keys": null,
609
+ "mean_std_embedding_keys": [
610
+ "x",
611
+ "y",
612
+ "z",
613
+ "roll",
614
+ "pitch",
615
+ "yaw"
616
+ ],
617
+ "action_configs": null
618
+ },
619
+ "language": {
620
+ "delta_indices": [
621
+ 0
622
+ ],
623
+ "modality_keys": [
624
+ "annotation.human.action.task_description"
625
+ ],
626
+ "sin_cos_embedding_keys": null,
627
+ "mean_std_embedding_keys": null,
628
+ "action_configs": null
629
+ }
630
+ },
631
+ "oxe_google": {
632
+ "video": {
633
+ "delta_indices": [
634
+ 0
635
+ ],
636
+ "modality_keys": [
637
+ "image"
638
+ ],
639
+ "sin_cos_embedding_keys": null,
640
+ "mean_std_embedding_keys": null,
641
+ "action_configs": null
642
+ },
643
+ "state": {
644
+ "delta_indices": [
645
+ 0
646
+ ],
647
+ "modality_keys": [
648
+ "x",
649
+ "y",
650
+ "z",
651
+ "rx",
652
+ "ry",
653
+ "rz",
654
+ "rw",
655
+ "gripper"
656
+ ],
657
+ "sin_cos_embedding_keys": null,
658
+ "mean_std_embedding_keys": null,
659
+ "action_configs": null
660
+ },
661
+ "action": {
662
+ "delta_indices": [
663
+ 0,
664
+ 1,
665
+ 2,
666
+ 3,
667
+ 4,
668
+ 5,
669
+ 6,
670
+ 7
671
+ ],
672
+ "modality_keys": [
673
+ "x",
674
+ "y",
675
+ "z",
676
+ "roll",
677
+ "pitch",
678
+ "yaw",
679
+ "gripper"
680
+ ],
681
+ "sin_cos_embedding_keys": null,
682
+ "mean_std_embedding_keys": [
683
+ "x",
684
+ "y",
685
+ "z",
686
+ "roll",
687
+ "pitch",
688
+ "yaw"
689
+ ],
690
+ "action_configs": null
691
+ },
692
+ "language": {
693
+ "delta_indices": [
694
+ 0
695
+ ],
696
+ "modality_keys": [
697
+ "annotation.human.action.task_description"
698
+ ],
699
+ "sin_cos_embedding_keys": null,
700
+ "mean_std_embedding_keys": null,
701
+ "action_configs": null
702
+ }
703
+ },
704
+ "oxe_droid": {
705
+ "video": {
706
+ "delta_indices": [
707
+ 0
708
+ ],
709
+ "modality_keys": [
710
+ "exterior_image_1_left",
711
+ "wrist_image_left"
712
+ ],
713
+ "sin_cos_embedding_keys": null,
714
+ "mean_std_embedding_keys": null,
715
+ "action_configs": null
716
+ },
717
+ "state": {
718
+ "delta_indices": [
719
+ 0
720
+ ],
721
+ "modality_keys": [
722
+ "joint_position",
723
+ "gripper_position"
724
+ ],
725
+ "sin_cos_embedding_keys": null,
726
+ "mean_std_embedding_keys": null,
727
+ "action_configs": null
728
+ },
729
+ "action": {
730
+ "delta_indices": [
731
+ 0,
732
+ 1,
733
+ 2,
734
+ 3,
735
+ 4,
736
+ 5,
737
+ 6,
738
+ 7,
739
+ 8,
740
+ 9,
741
+ 10,
742
+ 11,
743
+ 12,
744
+ 13,
745
+ 14,
746
+ 15,
747
+ 16,
748
+ 17,
749
+ 18,
750
+ 19,
751
+ 20,
752
+ 21,
753
+ 22,
754
+ 23,
755
+ 24,
756
+ 25,
757
+ 26,
758
+ 27,
759
+ 28,
760
+ 29,
761
+ 30,
762
+ 31
763
+ ],
764
+ "modality_keys": [
765
+ "joint_position",
766
+ "gripper_position"
767
+ ],
768
+ "sin_cos_embedding_keys": null,
769
+ "mean_std_embedding_keys": null,
770
+ "action_configs": [
771
+ {
772
+ "rep": "RELATIVE",
773
+ "type": "NON_EEF",
774
+ "format": "DEFAULT",
775
+ "state_key": null
776
+ },
777
+ {
778
+ "rep": "ABSOLUTE",
779
+ "type": "NON_EEF",
780
+ "format": "DEFAULT",
781
+ "state_key": null
782
+ }
783
+ ]
784
+ },
785
+ "language": {
786
+ "delta_indices": [
787
+ 0
788
+ ],
789
+ "modality_keys": [
790
+ "annotation.language.language_instruction"
791
+ ],
792
+ "sin_cos_embedding_keys": null,
793
+ "mean_std_embedding_keys": null,
794
+ "action_configs": null
795
+ }
796
+ },
797
+ "new_embodiment": {
798
+ "video": {
799
+ "delta_indices": [
800
+ 0
801
+ ],
802
+ "modality_keys": [
803
+ "robot0_eye_in_hand",
804
+ "robot0_agentview_left",
805
+ "robot0_agentview_right"
806
+ ],
807
+ "sin_cos_embedding_keys": null,
808
+ "mean_std_embedding_keys": null,
809
+ "action_configs": null
810
+ },
811
+ "state": {
812
+ "delta_indices": [
813
+ 0
814
+ ],
815
+ "modality_keys": [
816
+ "base_position",
817
+ "base_rotation",
818
+ "end_effector_position_relative",
819
+ "end_effector_rotation_relative",
820
+ "gripper_qpos"
821
+ ],
822
+ "sin_cos_embedding_keys": null,
823
+ "mean_std_embedding_keys": null,
824
+ "action_configs": null
825
+ },
826
+ "action": {
827
+ "delta_indices": [
828
+ 0,
829
+ 1,
830
+ 2,
831
+ 3,
832
+ 4,
833
+ 5,
834
+ 6,
835
+ 7,
836
+ 8,
837
+ 9,
838
+ 10,
839
+ 11,
840
+ 12,
841
+ 13,
842
+ 14,
843
+ 15
844
+ ],
845
+ "modality_keys": [
846
+ "base_motion",
847
+ "control_mode",
848
+ "end_effector_position",
849
+ "end_effector_rotation",
850
+ "gripper_close"
851
+ ],
852
+ "sin_cos_embedding_keys": null,
853
+ "mean_std_embedding_keys": null,
854
+ "action_configs": [
855
+ {
856
+ "rep": "ABSOLUTE",
857
+ "type": "NON_EEF",
858
+ "format": "DEFAULT",
859
+ "state_key": null
860
+ },
861
+ {
862
+ "rep": "ABSOLUTE",
863
+ "type": "NON_EEF",
864
+ "format": "DEFAULT",
865
+ "state_key": null
866
+ },
867
+ {
868
+ "rep": "ABSOLUTE",
869
+ "type": "EEF",
870
+ "format": "DEFAULT",
871
+ "state_key": null
872
+ },
873
+ {
874
+ "rep": "ABSOLUTE",
875
+ "type": "EEF",
876
+ "format": "DEFAULT",
877
+ "state_key": null
878
+ },
879
+ {
880
+ "rep": "ABSOLUTE",
881
+ "type": "NON_EEF",
882
+ "format": "DEFAULT",
883
+ "state_key": null
884
+ }
885
+ ]
886
+ },
887
+ "language": {
888
+ "delta_indices": [
889
+ 0
890
+ ],
891
+ "modality_keys": [
892
+ "task"
893
+ ],
894
+ "sin_cos_embedding_keys": null,
895
+ "mean_std_embedding_keys": null,
896
+ "action_configs": null
897
+ }
898
+ }
899
+ },
900
+ "image_crop_size": null,
901
+ "image_target_size": null,
902
+ "use_albumentations": true,
903
+ "random_rotation_angle": null,
904
+ "color_jitter_params": {
905
+ "brightness": 0.3,
906
+ "contrast": 0.4,
907
+ "saturation": 0.5,
908
+ "hue": 0.08
909
+ },
910
+ "shortest_image_edge": 256,
911
+ "crop_fraction": 0.95,
912
+ "model_name": "nvidia/Eagle-Block2A-2B-v2",
913
+ "model_type": "eagle",
914
+ "formalize_language": true,
915
+ "max_state_dim": 128,
916
+ "max_action_dim": 128,
917
+ "max_action_horizon": 50,
918
+ "use_percentiles": false,
919
+ "clip_outliers": true,
920
+ "apply_sincos_state_encoding": true,
921
+ "use_relative_action": true
922
+ }
923
+ }
processor/statistics.json ADDED
The diff for this file is too large to render. See raw diff
 
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f2763d70cd1de320f4d0515e426de65befdc4afd41cd3575f080430938bc410b
3
+ size 7633