pravsels commited on
Commit
0dffada
·
verified ·
1 Parent(s): 80d9e21

Upload GR00T N1.7 checkpoint step 20000 to repo root

Browse files
config.json ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_horizon": 40,
3
+ "add_pos_embed": true,
4
+ "apply_sincos_state_encoding": false,
5
+ "architectures": [
6
+ "Gr00tN1d7"
7
+ ],
8
+ "attn_dropout": 0.2,
9
+ "attn_implementation": null,
10
+ "backbone_embedding_dim": 2048,
11
+ "backbone_trainable_params_fp32": true,
12
+ "color_jitter_params": {
13
+ "brightness": 0.3,
14
+ "contrast": 0.4,
15
+ "hue": 0.08,
16
+ "saturation": 0.5
17
+ },
18
+ "crop_fraction": 0.95,
19
+ "diffusion_model_cfg": {
20
+ "attention_head_dim": 48,
21
+ "dropout": 0.2,
22
+ "final_dropout": true,
23
+ "interleave_self_attention": true,
24
+ "norm_type": "ada_norm",
25
+ "num_attention_heads": 32,
26
+ "num_layers": 32,
27
+ "output_dim": 1024,
28
+ "positional_embeddings": null
29
+ },
30
+ "dtype": "bfloat16",
31
+ "exclude_state": false,
32
+ "formalize_language": true,
33
+ "hidden_size": 1024,
34
+ "image_crop_size": [
35
+ 230,
36
+ 230
37
+ ],
38
+ "image_target_size": [
39
+ 256,
40
+ 256
41
+ ],
42
+ "letter_box_transform": false,
43
+ "load_bf16": true,
44
+ "max_action_dim": 132,
45
+ "max_num_embodiments": 32,
46
+ "max_seq_len": 1024,
47
+ "max_state_dim": 132,
48
+ "model_dtype": "bfloat16",
49
+ "model_name": "nvidia/Cosmos-Reason2-2B",
50
+ "model_type": "Gr00tN1d7",
51
+ "noise_beta_alpha": 1.5,
52
+ "noise_beta_beta": 1.0,
53
+ "noise_s": 0.999,
54
+ "num_inference_timesteps": 4,
55
+ "num_timestep_buckets": 1000,
56
+ "random_history_crop": true,
57
+ "random_rotation_angle": 0,
58
+ "reproject_vision": false,
59
+ "rtc_ramp_rate": 6.0,
60
+ "select_layer": 16,
61
+ "shortest_image_edge": 256,
62
+ "state_dropout_prob": 0.2,
63
+ "state_gaussian_noise_std": 0.0,
64
+ "transformers_version": "4.57.3",
65
+ "tune_diffusion_model": true,
66
+ "tune_linear": true,
67
+ "tune_llm": false,
68
+ "tune_projector": true,
69
+ "tune_top_llm_layers": 0,
70
+ "tune_visual": false,
71
+ "tune_vlln": true,
72
+ "use_albumentations": true,
73
+ "use_alternate_vl_dit": true,
74
+ "use_flash_attention": true,
75
+ "use_future_tokens": false,
76
+ "use_mean_std": false,
77
+ "use_percentiles": true,
78
+ "use_vl_self_attention": true,
79
+ "use_vlln": true,
80
+ "vl_self_attention_cfg": {
81
+ "attention_head_dim": 64,
82
+ "dropout": 0.2,
83
+ "final_dropout": true,
84
+ "num_attention_heads": 32,
85
+ "num_layers": 4,
86
+ "positional_embeddings": null
87
+ }
88
+ }
embodiment_id.json ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "robocasa_panda_omron": 13,
3
+ "oxe_droid": 17,
4
+ "oxe_fractal": 18,
5
+ "oxe_language_table": 19,
6
+ "oxe_bridge": 20,
7
+ "unknown": 22,
8
+ "gr1_unified": 20,
9
+ "agibot": 26,
10
+ "sim_behavior_r1_pro": 23,
11
+ "xdof": 24,
12
+ "xdof_oss_data": 25,
13
+ "unitree_g1_full_body_with_waist_height_nav_cmd": 25,
14
+ "real_r1_pro_sharpa": 27,
15
+ "real_r1_pro_sharpa_add_view": 27,
16
+ "real_r1_pro_sharpa_relative_arm_joint": 26,
17
+ "real_r1_pro_sharpa_delta_eef": 26,
18
+ "real_r1_pro_sharpa_absolute_eef": 26,
19
+ "real_r1_pro_sharpa_meanstd": 26,
20
+ "real_r1_pro_sharpa_relative_eef": 26,
21
+ "real_r1_pro_sharpa_relative_eef_add_view": 26,
22
+ "real_r1_pro_sharpa_relative_eef_relative_hand": 26,
23
+ "real_r1_pro_sharpa_relative_eef_human": 26,
24
+ "real_r1_pro_sharpa_relative_eef_human_add_view": 26,
25
+ "real_r1_pro_sharpa_relative_eef_human_relative_hand": 26,
26
+ "real_r1_pro_sharpa_relative_eef_egodex": 26,
27
+ "real_r1_pro_sharpa_relative_eef_egodex_relative_hand": 26,
28
+ "real_r1_pro_sharpa_relative_eef_egodex_wrist_only": 26,
29
+ "real_r1_pro_sharpa_relative_eef_maxinsights": 26,
30
+ "real_r1_pro_sharpa_relative_eef_maxinsights_relative_hand": 26,
31
+ "real_r1_pro_sharpa_relative_eef_mecka": 26,
32
+ "real_r1_pro_sharpa_relative_eef_mecka_relative_hand": 26,
33
+ "real_g1_relative_eef_absolute_joints": 25,
34
+ "real_g1_relative_eef_absolute_joints_wrist_cam": 25,
35
+ "real_g1_relative_eef_relative_joints": 25,
36
+ "real_r1_pro_sharpa_relative_eef_relative_hand_relative_joint": 26,
37
+ "real_r1_pro_sharpa_relative_joint": 29,
38
+ "oxe_droid_relative_eef_relative_joint": 24,
39
+ "oxe_droid_relative_eef_relative_joint_swapped": 24,
40
+ "oxe_droid_relative_eef_relative_joint_upweight_z": 24,
41
+ "oxe_droid_relative_eef_relative_joint_upweight_z_swapped": 24,
42
+ "oxe_droid_relative_eef_relative_joint_3view": 24,
43
+ "oxe_droid_relative_eef_relative_joint_3view_swapped": 24,
44
+ "oxe_droid_relative_eef": 24,
45
+ "oxe_droid_joint_position_relative": 24,
46
+ "xdof_relative_eef_relative_joint": 27,
47
+ "xdof_relative_eef_relative_joint_subtask": 27,
48
+ "xdof_relative_eef": 27,
49
+ "xdof_relative_joint": 28,
50
+ "simpler_env_google": 0,
51
+ "simpler_env_widowx": 1,
52
+ "libero_sim": 2,
53
+ "droid_sim": 3,
54
+ "robocasa_gr1_tabletop": 10,
55
+ "new_embodiment": 10,
56
+ "unitree_g1_sonic": 11
57
+ }
experiment_cfg/conf.yaml ADDED
@@ -0,0 +1,231 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ load_config_path: null
2
+ model:
3
+ model_type: Gr00tN1d7
4
+ model_dtype: bfloat16
5
+ model_name: nvidia/Cosmos-Reason2-2B
6
+ backbone_model_type: qwen
7
+ model_revision: null
8
+ tune_top_llm_layers: 0
9
+ backbone_embedding_dim: 2048
10
+ tune_llm: false
11
+ tune_visual: false
12
+ select_layer: 12
13
+ reproject_vision: false
14
+ use_flash_attention: true
15
+ load_bf16: true
16
+ backbone_trainable_params_fp32: true
17
+ image_crop_size:
18
+ - 230
19
+ - 230
20
+ image_target_size:
21
+ - 256
22
+ - 256
23
+ shortest_image_edge: null
24
+ crop_fraction: null
25
+ random_rotation_angle: null
26
+ color_jitter_params:
27
+ brightness: 0.3
28
+ contrast: 0.4
29
+ saturation: 0.5
30
+ hue: 0.08
31
+ use_albumentations_transforms: true
32
+ letter_box_transform: false
33
+ extra_augmentation_config: null
34
+ formalize_language: true
35
+ apply_sincos_state_encoding: false
36
+ use_percentiles: true
37
+ use_relative_action: true
38
+ max_state_dim: 132
39
+ max_action_dim: 132
40
+ action_horizon: 40
41
+ hidden_size: 1024
42
+ input_embedding_dim: 1536
43
+ state_history_length: 1
44
+ add_pos_embed: true
45
+ attn_dropout: 0.2
46
+ use_vlln: true
47
+ max_seq_len: 1024
48
+ use_alternate_vl_dit: true
49
+ attend_text_every_n_blocks: 2
50
+ diffusion_model_cfg:
51
+ positional_embeddings: null
52
+ num_layers: 16
53
+ num_attention_heads: 32
54
+ attention_head_dim: 48
55
+ norm_type: ada_norm
56
+ dropout: 0.2
57
+ final_dropout: true
58
+ output_dim: 1024
59
+ interleave_self_attention: true
60
+ num_inference_timesteps: 4
61
+ noise_beta_alpha: 1.5
62
+ noise_beta_beta: 1.0
63
+ noise_s: 0.999
64
+ num_timestep_buckets: 1000
65
+ tune_projector: true
66
+ tune_diffusion_model: true
67
+ tune_vlln: true
68
+ state_dropout_prob: 0.2
69
+ exclude_state: false
70
+ use_mean_std: false
71
+ max_num_embodiments: 32
72
+ data:
73
+ datasets:
74
+ - dataset_paths:
75
+ - /workspace/vast_runs/groot/datasets/transfer_cube
76
+ embodiment_tag: new_embodiment
77
+ mix_ratio: 1.0
78
+ dataset_type: physical_embodiment
79
+ val_dataset_path: null
80
+ modality_configs:
81
+ new_embodiment:
82
+ video:
83
+ delta_indices:
84
+ - 0
85
+ modality_keys:
86
+ - left_wrist
87
+ - right_wrist
88
+ - top
89
+ sin_cos_embedding_keys: null
90
+ mean_std_embedding_keys: null
91
+ action_configs: null
92
+ state:
93
+ delta_indices:
94
+ - 0
95
+ modality_keys:
96
+ - left_arm
97
+ - left_gripper
98
+ - right_arm
99
+ - right_gripper
100
+ sin_cos_embedding_keys: null
101
+ mean_std_embedding_keys: null
102
+ action_configs: null
103
+ action:
104
+ delta_indices:
105
+ - 0
106
+ - 1
107
+ - 2
108
+ - 3
109
+ - 4
110
+ - 5
111
+ - 6
112
+ - 7
113
+ - 8
114
+ - 9
115
+ - 10
116
+ - 11
117
+ - 12
118
+ - 13
119
+ - 14
120
+ - 15
121
+ modality_keys:
122
+ - left_arm
123
+ - left_gripper
124
+ - right_arm
125
+ - right_gripper
126
+ sin_cos_embedding_keys: null
127
+ mean_std_embedding_keys: null
128
+ action_configs:
129
+ - rep: RELATIVE
130
+ type: NON_EEF
131
+ format: DEFAULT
132
+ state_key: null
133
+ - rep: ABSOLUTE
134
+ type: NON_EEF
135
+ format: DEFAULT
136
+ state_key: null
137
+ - rep: RELATIVE
138
+ type: NON_EEF
139
+ format: DEFAULT
140
+ state_key: null
141
+ - rep: ABSOLUTE
142
+ type: NON_EEF
143
+ format: DEFAULT
144
+ state_key: null
145
+ language:
146
+ delta_indices:
147
+ - 0
148
+ modality_keys:
149
+ - annotation.human.task_description
150
+ sin_cos_embedding_keys: null
151
+ mean_std_embedding_keys: null
152
+ action_configs: null
153
+ download_cache: false
154
+ shard_size: 1024
155
+ episode_sampling_rate: 0.1
156
+ num_shards_per_epoch: 100000
157
+ override_pretraining_statistics: true
158
+ mode: single_turn
159
+ random_chop: 0.0
160
+ mock_dataset_mode: false
161
+ shuffle: true
162
+ seed: 42
163
+ multiprocessing_context: fork
164
+ allow_padding: false
165
+ subsample_ratio: 1.0
166
+ image_crop_size:
167
+ - 244
168
+ - 244
169
+ image_target_size:
170
+ - 224
171
+ - 224
172
+ training:
173
+ output_dir: /workspace/vast_runs/groot/checkpoints/groot_transfer_cube
174
+ experiment_name: groot_transfer_cube
175
+ max_steps: 20000
176
+ global_batch_size: 32
177
+ per_gpu_batch_size: null
178
+ gradient_accumulation_steps: 1
179
+ learning_rate: 0.0001
180
+ lr_scheduler_type: cosine
181
+ weight_decay: 1.0e-05
182
+ warmup_ratio: 0.05
183
+ warmup_steps: 0
184
+ max_grad_norm: 1.0
185
+ optim: adamw_torch
186
+ start_from_checkpoint: nvidia/GR00T-N1.7-3B
187
+ skip_weight_loading: false
188
+ tf32: true
189
+ fp16: false
190
+ bf16: true
191
+ eval_bf16: true
192
+ logging_steps: 10
193
+ save_steps: 5000
194
+ save_total_limit: 5
195
+ save_vl_model: false
196
+ save_only_model: false
197
+ resume_from_checkpoint: false
198
+ upload_checkpoints: false
199
+ upload_every: 1000
200
+ upload_last_n_checkpoints: 5
201
+ max_concurrent_uploads: 2
202
+ eval_strategy: 'no'
203
+ eval_steps: 500
204
+ eval_set_split_ratio: 0.1
205
+ eval_batch_size: 2
206
+ save_best_eval_metric_name: ''
207
+ save_best_eval_metric_greater_is_better: true
208
+ deepspeed_stage: 2
209
+ gradient_checkpointing: false
210
+ transformers_trust_remote_code: true
211
+ transformers_local_files_only: false
212
+ transformers_cache_dir: null
213
+ transformers_access_token: null
214
+ use_ddp: false
215
+ ddp_bucket_cap_mb: 100
216
+ num_gpus: 4
217
+ dataloader_num_workers: 8
218
+ remove_unused_columns: false
219
+ use_wandb: true
220
+ wandb_project: groot_transfer_cube
221
+ enable_profiling: false
222
+ max_retries: 3
223
+ assert_loss_less_than: null
224
+ add_rl_callback: false
225
+ enable_open_loop_eval: false
226
+ open_loop_eval_traj_ids:
227
+ - 0
228
+ open_loop_eval_steps_per_traj: 100
229
+ open_loop_eval_plot_indices: null
230
+ max_steps: 20000
231
+ save_steps: 5000
experiment_cfg/config.yaml ADDED
@@ -0,0 +1,229 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ load_config_path: null
2
+ model:
3
+ model_type: Gr00tN1d7
4
+ model_dtype: bfloat16
5
+ model_name: nvidia/Cosmos-Reason2-2B
6
+ backbone_model_type: qwen
7
+ model_revision: null
8
+ tune_top_llm_layers: 0
9
+ backbone_embedding_dim: 2048
10
+ tune_llm: false
11
+ tune_visual: false
12
+ select_layer: 12
13
+ reproject_vision: false
14
+ use_flash_attention: true
15
+ load_bf16: true
16
+ backbone_trainable_params_fp32: true
17
+ image_crop_size:
18
+ - 230
19
+ - 230
20
+ image_target_size:
21
+ - 256
22
+ - 256
23
+ shortest_image_edge: null
24
+ crop_fraction: null
25
+ random_rotation_angle: null
26
+ color_jitter_params:
27
+ brightness: 0.3
28
+ contrast: 0.4
29
+ saturation: 0.5
30
+ hue: 0.08
31
+ use_albumentations_transforms: true
32
+ letter_box_transform: false
33
+ extra_augmentation_config: null
34
+ formalize_language: true
35
+ apply_sincos_state_encoding: false
36
+ use_percentiles: true
37
+ use_relative_action: true
38
+ max_state_dim: 132
39
+ max_action_dim: 132
40
+ action_horizon: 40
41
+ hidden_size: 1024
42
+ input_embedding_dim: 1536
43
+ state_history_length: 1
44
+ add_pos_embed: true
45
+ attn_dropout: 0.2
46
+ use_vlln: true
47
+ max_seq_len: 1024
48
+ use_alternate_vl_dit: true
49
+ attend_text_every_n_blocks: 2
50
+ diffusion_model_cfg:
51
+ positional_embeddings: null
52
+ num_layers: 16
53
+ num_attention_heads: 32
54
+ attention_head_dim: 48
55
+ norm_type: ada_norm
56
+ dropout: 0.2
57
+ final_dropout: true
58
+ output_dim: 1024
59
+ interleave_self_attention: true
60
+ num_inference_timesteps: 4
61
+ noise_beta_alpha: 1.5
62
+ noise_beta_beta: 1.0
63
+ noise_s: 0.999
64
+ num_timestep_buckets: 1000
65
+ tune_projector: true
66
+ tune_diffusion_model: true
67
+ tune_vlln: true
68
+ state_dropout_prob: 0.2
69
+ exclude_state: false
70
+ use_mean_std: false
71
+ max_num_embodiments: 32
72
+ data:
73
+ datasets:
74
+ - dataset_paths:
75
+ - /workspace/vast_runs/groot/datasets/transfer_cube
76
+ embodiment_tag: new_embodiment
77
+ mix_ratio: 1.0
78
+ dataset_type: physical_embodiment
79
+ val_dataset_path: null
80
+ modality_configs:
81
+ new_embodiment:
82
+ video:
83
+ delta_indices:
84
+ - 0
85
+ modality_keys:
86
+ - left_wrist
87
+ - right_wrist
88
+ - top
89
+ sin_cos_embedding_keys: null
90
+ mean_std_embedding_keys: null
91
+ action_configs: null
92
+ state:
93
+ delta_indices:
94
+ - 0
95
+ modality_keys:
96
+ - left_arm
97
+ - left_gripper
98
+ - right_arm
99
+ - right_gripper
100
+ sin_cos_embedding_keys: null
101
+ mean_std_embedding_keys: null
102
+ action_configs: null
103
+ action:
104
+ delta_indices:
105
+ - 0
106
+ - 1
107
+ - 2
108
+ - 3
109
+ - 4
110
+ - 5
111
+ - 6
112
+ - 7
113
+ - 8
114
+ - 9
115
+ - 10
116
+ - 11
117
+ - 12
118
+ - 13
119
+ - 14
120
+ - 15
121
+ modality_keys:
122
+ - left_arm
123
+ - left_gripper
124
+ - right_arm
125
+ - right_gripper
126
+ sin_cos_embedding_keys: null
127
+ mean_std_embedding_keys: null
128
+ action_configs:
129
+ - rep: relative
130
+ type: non_eef
131
+ format: default
132
+ state_key: null
133
+ - rep: absolute
134
+ type: non_eef
135
+ format: default
136
+ state_key: null
137
+ - rep: relative
138
+ type: non_eef
139
+ format: default
140
+ state_key: null
141
+ - rep: absolute
142
+ type: non_eef
143
+ format: default
144
+ state_key: null
145
+ language:
146
+ delta_indices:
147
+ - 0
148
+ modality_keys:
149
+ - annotation.human.task_description
150
+ sin_cos_embedding_keys: null
151
+ mean_std_embedding_keys: null
152
+ action_configs: null
153
+ download_cache: false
154
+ shard_size: 1024
155
+ episode_sampling_rate: 0.1
156
+ num_shards_per_epoch: 100000
157
+ override_pretraining_statistics: true
158
+ mode: single_turn
159
+ random_chop: 0.0
160
+ mock_dataset_mode: false
161
+ shuffle: true
162
+ seed: 42
163
+ multiprocessing_context: fork
164
+ allow_padding: false
165
+ subsample_ratio: 1.0
166
+ image_crop_size:
167
+ - 244
168
+ - 244
169
+ image_target_size:
170
+ - 224
171
+ - 224
172
+ training:
173
+ output_dir: /workspace/vast_runs/groot/checkpoints/groot_transfer_cube
174
+ experiment_name: groot_transfer_cube
175
+ max_steps: 20000
176
+ global_batch_size: 32
177
+ per_gpu_batch_size: null
178
+ gradient_accumulation_steps: 1
179
+ learning_rate: 0.0001
180
+ lr_scheduler_type: cosine
181
+ weight_decay: 1.0e-05
182
+ warmup_ratio: 0.05
183
+ warmup_steps: 0
184
+ max_grad_norm: 1.0
185
+ optim: adamw_torch
186
+ start_from_checkpoint: nvidia/GR00T-N1.7-3B
187
+ skip_weight_loading: false
188
+ tf32: true
189
+ fp16: false
190
+ bf16: true
191
+ eval_bf16: true
192
+ logging_steps: 10
193
+ save_steps: 5000
194
+ save_total_limit: 5
195
+ save_vl_model: false
196
+ save_only_model: false
197
+ resume_from_checkpoint: false
198
+ upload_checkpoints: false
199
+ upload_every: 1000
200
+ upload_last_n_checkpoints: 5
201
+ max_concurrent_uploads: 2
202
+ eval_strategy: 'no'
203
+ eval_steps: 500
204
+ eval_set_split_ratio: 0.1
205
+ eval_batch_size: 2
206
+ save_best_eval_metric_name: ''
207
+ save_best_eval_metric_greater_is_better: true
208
+ deepspeed_stage: 2
209
+ gradient_checkpointing: false
210
+ transformers_trust_remote_code: true
211
+ transformers_local_files_only: false
212
+ transformers_cache_dir: null
213
+ transformers_access_token: null
214
+ use_ddp: false
215
+ ddp_bucket_cap_mb: 100
216
+ num_gpus: 4
217
+ dataloader_num_workers: 8
218
+ remove_unused_columns: false
219
+ use_wandb: true
220
+ wandb_project: groot_transfer_cube
221
+ enable_profiling: false
222
+ max_retries: 3
223
+ assert_loss_less_than: null
224
+ add_rl_callback: false
225
+ enable_open_loop_eval: false
226
+ open_loop_eval_traj_ids:
227
+ - 0
228
+ open_loop_eval_steps_per_traj: 100
229
+ open_loop_eval_plot_indices: null
experiment_cfg/dataset_statistics.json ADDED
@@ -0,0 +1,1638 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "new_embodiment": {
3
+ "state": {
4
+ "left_arm": {
5
+ "min": [
6
+ -58.28571319580078,
7
+ -102.5934066772461,
8
+ -81.93406677246094,
9
+ 14.94505500793457,
10
+ -53.230770111083984
11
+ ],
12
+ "max": [
13
+ 60.30769348144531,
14
+ 69.80220031738281,
15
+ 95.82417297363281,
16
+ 100.57142639160156,
17
+ 56.30769348144531
18
+ ],
19
+ "mean": [
20
+ 19.40082359313965,
21
+ -55.4830436706543,
22
+ 42.32158279418945,
23
+ 78.39623260498047,
24
+ 6.3484392166137695
25
+ ],
26
+ "std": [
27
+ 25.444252014160156,
28
+ 46.408103942871094,
29
+ 33.32065200805664,
30
+ 16.01414680480957,
31
+ 14.616244316101074
32
+ ],
33
+ "q01": [
34
+ -48.263736724853516,
35
+ -102.5054931640625,
36
+ -53.83120796203613,
37
+ 29.77406606674195,
38
+ -39.868133544921875
39
+ ],
40
+ "q99": [
41
+ 58.28571319580078,
42
+ 48.908131332397296,
43
+ 95.56044006347656,
44
+ 100.3956069946289,
45
+ 45.31867980957031
46
+ ]
47
+ },
48
+ "left_gripper": {
49
+ "min": [
50
+ 0.4045853018760681
51
+ ],
52
+ "max": [
53
+ 52.056640625
54
+ ],
55
+ "mean": [
56
+ 6.693004608154297
57
+ ],
58
+ "std": [
59
+ 11.26343059539795
60
+ ],
61
+ "q01": [
62
+ 0.6743088364601135
63
+ ],
64
+ "q99": [
65
+ 45.7181396484375
66
+ ]
67
+ },
68
+ "right_arm": {
69
+ "min": [
70
+ -45.80219650268555,
71
+ -102.76923370361328,
72
+ -68.3956069946289,
73
+ 0.5274725556373596,
74
+ -48.131866455078125
75
+ ],
76
+ "max": [
77
+ 57.58241653442383,
78
+ 42.10988998413086,
79
+ 96.17582702636719,
80
+ 100.30769348144531,
81
+ 78.10989379882812
82
+ ],
83
+ "mean": [
84
+ -4.468321800231934,
85
+ -64.00232696533203,
86
+ 45.468753814697266,
87
+ 76.61595916748047,
88
+ 8.66189193725586
89
+ ],
90
+ "std": [
91
+ 17.140336990356445,
92
+ 41.7463493347168,
93
+ 27.779335021972656,
94
+ 17.868837356567383,
95
+ 15.68790340423584
96
+ ],
97
+ "q01": [
98
+ -37.010990142822266,
99
+ -102.76923370361328,
100
+ -31.296703338623047,
101
+ 13.274724960327148,
102
+ -38.989009857177734
103
+ ],
104
+ "q99": [
105
+ 45.27472686767578,
106
+ 31.47252655029297,
107
+ 96.0,
108
+ 98.37362670898438,
109
+ 59.03296661376953
110
+ ]
111
+ },
112
+ "right_gripper": {
113
+ "min": [
114
+ 0.664893627166748
115
+ ],
116
+ "max": [
117
+ 48.337764739990234
118
+ ],
119
+ "mean": [
120
+ 7.492151737213135
121
+ ],
122
+ "std": [
123
+ 10.930651664733887
124
+ ],
125
+ "q01": [
126
+ 0.7313829660415649
127
+ ],
128
+ "q99": [
129
+ 42.353721618652344
130
+ ]
131
+ }
132
+ },
133
+ "action": {
134
+ "left_arm": {
135
+ "min": [
136
+ -58.85714340209961,
137
+ -103.86813354492188,
138
+ -84.79120635986328,
139
+ 14.637362480163574,
140
+ -53.67033004760742
141
+ ],
142
+ "max": [
143
+ 60.70329666137695,
144
+ 69.4945068359375,
145
+ 96.92308044433594,
146
+ 100.79120635986328,
147
+ 56.659339904785156
148
+ ],
149
+ "mean": [
150
+ 19.586156845092773,
151
+ -56.023353576660156,
152
+ 41.433753967285156,
153
+ 78.04638671875,
154
+ 6.463015079498291
155
+ ],
156
+ "std": [
157
+ 25.510547637939453,
158
+ 45.92293930053711,
159
+ 33.88028335571289,
160
+ 16.092994689941406,
161
+ 14.640082359313965
162
+ ],
163
+ "q01": [
164
+ -48.483516693115234,
165
+ -102.9011001586914,
166
+ -56.131866455078125,
167
+ 28.96703338623047,
168
+ -39.809231491088866
169
+ ],
170
+ "q99": [
171
+ 58.769229888916016,
172
+ 47.9560432434082,
173
+ 96.65933990478516,
174
+ 100.61538696289062,
175
+ 45.52351703643782
176
+ ]
177
+ },
178
+ "left_gripper": {
179
+ "min": [
180
+ 0.0
181
+ ],
182
+ "max": [
183
+ 53.49794387817383
184
+ ],
185
+ "mean": [
186
+ 5.575933456420898
187
+ ],
188
+ "std": [
189
+ 11.395349502563477
190
+ ],
191
+ "q01": [
192
+ 0.4938271641731262
193
+ ],
194
+ "q99": [
195
+ 46.52921707153305
196
+ ]
197
+ },
198
+ "right_arm": {
199
+ "min": [
200
+ -46.10988998413086,
201
+ -103.38461303710938,
202
+ -71.69230651855469,
203
+ 0.5274725556373596,
204
+ -48.92307662963867
205
+ ],
206
+ "max": [
207
+ 57.978023529052734,
208
+ 41.758243560791016,
209
+ 97.45054626464844,
210
+ 100.83516693115234,
211
+ 78.72527313232422
212
+ ],
213
+ "mean": [
214
+ -4.379688262939453,
215
+ -64.219482421875,
216
+ 44.345279693603516,
217
+ 76.3670425415039,
218
+ 8.707589149475098
219
+ ],
220
+ "std": [
221
+ 17.216585159301758,
222
+ 41.30073165893555,
223
+ 28.517959594726562,
224
+ 18.001033782958984,
225
+ 15.708527565002441
226
+ ],
227
+ "q01": [
228
+ -37.05494689941406,
229
+ -102.76923370361328,
230
+ -33.626373291015625,
231
+ 12.571428298950195,
232
+ -39.42856979370117
233
+ ],
234
+ "q99": [
235
+ 45.230770111083984,
236
+ 30.534505901336505,
237
+ 97.27472686767578,
238
+ 98.57846267700178,
239
+ 59.29670333862305
240
+ ]
241
+ },
242
+ "right_gripper": {
243
+ "min": [
244
+ 0.0
245
+ ],
246
+ "max": [
247
+ 48.77049255371094
248
+ ],
249
+ "mean": [
250
+ 5.901553153991699
251
+ ],
252
+ "std": [
253
+ 11.008342742919922
254
+ ],
255
+ "q01": [
256
+ 0.4098360538482666
257
+ ],
258
+ "q99": [
259
+ 42.78688430786133
260
+ ]
261
+ }
262
+ },
263
+ "relative_action": {
264
+ "left_arm": {
265
+ "min": [
266
+ [
267
+ -17.186813354492188,
268
+ -43.38461685180664,
269
+ -19.12087631225586,
270
+ -23.252750396728516,
271
+ -18.373626708984375
272
+ ],
273
+ [
274
+ -20.08791160583496,
275
+ -48.13187026977539,
276
+ -23.16483497619629,
277
+ -23.252750396728516,
278
+ -22.857141494750977
279
+ ],
280
+ [
281
+ -22.989009857177734,
282
+ -52.70329666137695,
283
+ -27.648351669311523,
284
+ -23.252750396728516,
285
+ -26.28571319580078
286
+ ],
287
+ [
288
+ -25.802196502685547,
289
+ -57.010990142822266,
290
+ -31.956043243408203,
291
+ -23.252750396728516,
292
+ -27.956043243408203
293
+ ],
294
+ [
295
+ -28.61538314819336,
296
+ -61.758243560791016,
297
+ -36.087913513183594,
298
+ -23.60439682006836,
299
+ -28.835166931152344
300
+ ],
301
+ [
302
+ -31.25274658203125,
303
+ -66.32967376708984,
304
+ -40.57142639160156,
305
+ -27.032970428466797,
306
+ -30.593406677246094
307
+ ],
308
+ [
309
+ -33.4505500793457,
310
+ -70.54945373535156,
311
+ -44.879119873046875,
312
+ -28.967029571533203,
313
+ -34.28571319580078
314
+ ],
315
+ [
316
+ -35.29670333862305,
317
+ -74.41758728027344,
318
+ -49.098899841308594,
319
+ -30.901103973388672,
320
+ -36.30769348144531
321
+ ],
322
+ [
323
+ -35.82417297363281,
324
+ -78.10989379882812,
325
+ -52.70329666137695,
326
+ -32.83516311645508,
327
+ -37.89011001586914
328
+ ],
329
+ [
330
+ -37.93406677246094,
331
+ -81.62637329101562,
332
+ -56.30769348144531,
333
+ -34.68132400512695,
334
+ -39.56044006347656
335
+ ],
336
+ [
337
+ -41.010990142822266,
338
+ -85.4945068359375,
339
+ -59.82417297363281,
340
+ -36.967037200927734,
341
+ -41.846153259277344
342
+ ],
343
+ [
344
+ -43.736263275146484,
345
+ -89.18681335449219,
346
+ -63.164833068847656,
347
+ -39.69231033325195,
348
+ -44.0439567565918
349
+ ],
350
+ [
351
+ -45.58241653442383,
352
+ -92.6153793334961,
353
+ -66.4175796508789,
354
+ -41.714290618896484,
355
+ -45.626373291015625
356
+ ],
357
+ [
358
+ -48.0439567565918,
359
+ -95.69230651855469,
360
+ -69.67032623291016,
361
+ -43.64834976196289,
362
+ -47.12087631225586
363
+ ],
364
+ [
365
+ -50.41758346557617,
366
+ -98.24176025390625,
367
+ -72.65933990478516,
368
+ -44.967037200927734,
369
+ -47.82417297363281
370
+ ],
371
+ [
372
+ -52.52747344970703,
373
+ -101.84615325927734,
374
+ -75.64835357666016,
375
+ -46.021976470947266,
376
+ -48.351646423339844
377
+ ]
378
+ ],
379
+ "max": [
380
+ [
381
+ 21.142858505249023,
382
+ 37.758243560791016,
383
+ 43.20879364013672,
384
+ 15.60439682006836,
385
+ 11.86812973022461
386
+ ],
387
+ [
388
+ 25.538461685180664,
389
+ 42.68132400512695,
390
+ 49.5384635925293,
391
+ 19.912090301513672,
392
+ 15.296703338623047
393
+ ],
394
+ [
395
+ 29.142858505249023,
396
+ 46.90110397338867,
397
+ 55.78022003173828,
398
+ 24.043956756591797,
399
+ 18.63736343383789
400
+ ],
401
+ [
402
+ 31.516483306884766,
403
+ 50.59341049194336,
404
+ 60.79121017456055,
405
+ 28.0,
406
+ 22.065933227539062
407
+ ],
408
+ [
409
+ 35.56044006347656,
410
+ 54.90110397338867,
411
+ 65.45054626464844,
412
+ 31.604393005371094,
413
+ 24.96703338623047
414
+ ],
415
+ [
416
+ 40.21977996826172,
417
+ 59.12088394165039,
418
+ 70.02198028564453,
419
+ 35.120880126953125,
420
+ 27.78022003173828
421
+ ],
422
+ [
423
+ 44.087913513183594,
424
+ 63.164833068847656,
425
+ 74.24176025390625,
426
+ 38.81318664550781,
427
+ 30.593406677246094
428
+ ],
429
+ [
430
+ 48.30769348144531,
431
+ 69.23076629638672,
432
+ 77.84616088867188,
433
+ 42.68132019042969,
434
+ 33.05494689941406
435
+ ],
436
+ [
437
+ 52.43955993652344,
438
+ 75.38461303710938,
439
+ 81.010986328125,
440
+ 46.81318664550781,
441
+ 35.25274658203125
442
+ ],
443
+ [
444
+ 56.57142639160156,
445
+ 79.95603942871094,
446
+ 84.0,
447
+ 50.153846740722656,
448
+ 37.626373291015625
449
+ ],
450
+ [
451
+ 60.35165023803711,
452
+ 84.17582702636719,
453
+ 86.28571319580078,
454
+ 53.4945068359375,
455
+ 39.82417297363281
456
+ ],
457
+ [
458
+ 64.04396057128906,
459
+ 87.69230651855469,
460
+ 89.0988998413086,
461
+ 55.42857360839844,
462
+ 41.670326232910156
463
+ ],
464
+ [
465
+ 67.47252655029297,
466
+ 90.68132019042969,
467
+ 91.9120864868164,
468
+ 57.010990142822266,
469
+ 43.164833068847656
470
+ ],
471
+ [
472
+ 70.81318664550781,
473
+ 94.19780731201172,
474
+ 94.28571319580078,
475
+ 57.978023529052734,
476
+ 44.835166931152344
477
+ ],
478
+ [
479
+ 73.80220031738281,
480
+ 97.27472686767578,
481
+ 96.30769348144531,
482
+ 58.76923370361328,
483
+ 46.24176025390625
484
+ ],
485
+ [
486
+ 76.87911987304688,
487
+ 99.82418060302734,
488
+ 99.47252655029297,
489
+ 59.20879364013672,
490
+ 47.56044006347656
491
+ ]
492
+ ],
493
+ "mean": [
494
+ [
495
+ 0.2009088546037674,
496
+ -0.542931318283081,
497
+ -0.9684419631958008,
498
+ -0.34902822971343994,
499
+ 0.11815977096557617
500
+ ],
501
+ [
502
+ 0.20763306319713593,
503
+ -0.5415531992912292,
504
+ -0.9736434817314148,
505
+ -0.34697234630584717,
506
+ 0.11894739419221878
507
+ ],
508
+ [
509
+ 0.21414755284786224,
510
+ -0.5401749014854431,
511
+ -0.977736234664917,
512
+ -0.3446848690509796,
513
+ 0.11977844685316086
514
+ ],
515
+ [
516
+ 0.22053265571594238,
517
+ -0.5387799143791199,
518
+ -0.980843722820282,
519
+ -0.3425155282020569,
520
+ 0.12066444009542465
521
+ ],
522
+ [
523
+ 0.2267414927482605,
524
+ -0.5374149084091187,
525
+ -0.9832253456115723,
526
+ -0.34063610434532166,
527
+ 0.12155056744813919
528
+ ],
529
+ [
530
+ 0.23279529809951782,
531
+ -0.5360839366912842,
532
+ -0.9849362373352051,
533
+ -0.33914485573768616,
534
+ 0.12246031314134598
535
+ ],
536
+ [
537
+ 0.2387666255235672,
538
+ -0.5347722768783569,
539
+ -0.9861584305763245,
540
+ -0.33795133233070374,
541
+ 0.12331877648830414
542
+ ],
543
+ [
544
+ 0.24468213319778442,
545
+ -0.5334568023681641,
546
+ -0.9869793057441711,
547
+ -0.3371204733848572,
548
+ 0.12414994835853577
549
+ ],
550
+ [
551
+ 0.2505212724208832,
552
+ -0.5321456789970398,
553
+ -0.9876790046691895,
554
+ -0.33641257882118225,
555
+ 0.12503638863563538
556
+ ],
557
+ [
558
+ 0.2562320828437805,
559
+ -0.5308266878128052,
560
+ -0.9883260726928711,
561
+ -0.33572423458099365,
562
+ 0.12596991658210754
563
+ ],
564
+ [
565
+ 0.26175758242607117,
566
+ -0.5295072793960571,
567
+ -0.9889572262763977,
568
+ -0.33505547046661377,
569
+ 0.12693047523498535
570
+ ],
571
+ [
572
+ 0.2671513259410858,
573
+ -0.5281885862350464,
574
+ -0.9895422458648682,
575
+ -0.33439013361930847,
576
+ 0.12787628173828125
577
+ ],
578
+ [
579
+ 0.2723824381828308,
580
+ -0.5268692970275879,
581
+ -0.9900388121604919,
582
+ -0.333752304315567,
583
+ 0.12880182266235352
584
+ ],
585
+ [
586
+ 0.27743250131607056,
587
+ -0.5255545377731323,
588
+ -0.9903717041015625,
589
+ -0.3331150710582733,
590
+ 0.1297268271446228
591
+ ],
592
+ [
593
+ 0.28222665190696716,
594
+ -0.5242336988449097,
595
+ -0.9904801249504089,
596
+ -0.3325084447860718,
597
+ 0.13064076006412506
598
+ ],
599
+ [
600
+ 0.28677308559417725,
601
+ -0.5229145884513855,
602
+ -0.9903801679611206,
603
+ -0.33191025257110596,
604
+ 0.1316053718328476
605
+ ]
606
+ ],
607
+ "std": [
608
+ [
609
+ 2.855668067932129,
610
+ 4.9025044441223145,
611
+ 4.0851240158081055,
612
+ 2.314023494720459,
613
+ 1.810335397720337
614
+ ],
615
+ [
616
+ 3.599527597427368,
617
+ 6.334061622619629,
618
+ 5.217555046081543,
619
+ 2.934856414794922,
620
+ 2.3163375854492188
621
+ ],
622
+ [
623
+ 4.338955879211426,
624
+ 7.74597692489624,
625
+ 6.33854866027832,
626
+ 3.537440299987793,
627
+ 2.8020806312561035
628
+ ],
629
+ [
630
+ 5.064628601074219,
631
+ 9.133066177368164,
632
+ 7.442200183868408,
633
+ 4.1168084144592285,
634
+ 3.2684593200683594
635
+ ],
636
+ [
637
+ 5.774017333984375,
638
+ 10.492412567138672,
639
+ 8.525266647338867,
640
+ 4.671627521514893,
641
+ 3.717470407485962
642
+ ],
643
+ [
644
+ 6.466915130615234,
645
+ 11.822266578674316,
646
+ 9.58651065826416,
647
+ 5.202441692352295,
648
+ 4.151123046875
649
+ ],
650
+ [
651
+ 7.143940448760986,
652
+ 13.122232437133789,
653
+ 10.624748229980469,
654
+ 5.709934711456299,
655
+ 4.570662975311279
656
+ ],
657
+ [
658
+ 7.805746078491211,
659
+ 14.392159461975098,
660
+ 11.639915466308594,
661
+ 6.1947021484375,
662
+ 4.977085113525391
663
+ ],
664
+ [
665
+ 8.452713966369629,
666
+ 15.632648468017578,
667
+ 12.631338119506836,
668
+ 6.657371997833252,
669
+ 5.371180534362793
670
+ ],
671
+ [
672
+ 9.085217475891113,
673
+ 16.84394645690918,
674
+ 13.599506378173828,
675
+ 7.098426342010498,
676
+ 5.753519058227539
677
+ ],
678
+ [
679
+ 9.703451156616211,
680
+ 18.026844024658203,
681
+ 14.544489860534668,
682
+ 7.518970966339111,
683
+ 6.124891757965088
684
+ ],
685
+ [
686
+ 10.307982444763184,
687
+ 19.181936264038086,
688
+ 15.466625213623047,
689
+ 7.919764518737793,
690
+ 6.485807418823242
691
+ ],
692
+ [
693
+ 10.89881420135498,
694
+ 20.309877395629883,
695
+ 16.366405487060547,
696
+ 8.302095413208008,
697
+ 6.836681365966797
698
+ ],
699
+ [
700
+ 11.476049423217773,
701
+ 21.411109924316406,
702
+ 17.244361877441406,
703
+ 8.667145729064941,
704
+ 7.1778669357299805
705
+ ],
706
+ [
707
+ 12.039786338806152,
708
+ 22.4865665435791,
709
+ 18.10096549987793,
710
+ 9.016098022460938,
711
+ 7.509809494018555
712
+ ],
713
+ [
714
+ 12.590188980102539,
715
+ 23.53677749633789,
716
+ 18.93655776977539,
717
+ 9.349822998046875,
718
+ 7.832882404327393
719
+ ]
720
+ ],
721
+ "q01": [
722
+ [
723
+ -7.325711956024168,
724
+ -16.483516693115234,
725
+ -12.087911605834961,
726
+ -8.043952941894531,
727
+ -5.538461685180664
728
+ ],
729
+ [
730
+ -9.523515129089354,
731
+ -21.05494499206543,
732
+ -15.252747373580933,
733
+ -9.36263973236084,
734
+ -7.281758365631102
735
+ ],
736
+ [
737
+ -11.384616689682007,
738
+ -25.787251529693602,
739
+ -18.241758346557617,
740
+ -11.208786010742188,
741
+ -8.791208267211914
742
+ ],
743
+ [
744
+ -13.23076994895935,
745
+ -30.021976470947266,
746
+ -21.318681392669678,
747
+ -13.054946899414062,
748
+ -10.021978378295898
749
+ ],
750
+ [
751
+ -14.989009857177734,
752
+ -34.32966995239258,
753
+ -24.21977996826172,
754
+ -14.813182830810547,
755
+ -11.340659141540527
756
+ ],
757
+ [
758
+ -16.659339904785156,
759
+ -38.44659492492676,
760
+ -27.120880126953125,
761
+ -16.39560634613037,
762
+ -12.39560492515564
763
+ ],
764
+ [
765
+ -18.226811790466307,
766
+ -42.68132019042969,
767
+ -29.846153259277344,
768
+ -17.802196502685547,
769
+ -13.538461685180664
770
+ ],
771
+ [
772
+ -19.91208839416504,
773
+ -46.710329132080076,
774
+ -32.57142955780029,
775
+ -19.208790760040284,
776
+ -14.754285078048705
777
+ ],
778
+ [
779
+ -21.318679809570312,
780
+ -50.85714340209961,
781
+ -35.20878982543945,
782
+ -20.439559936523438,
783
+ -15.560440063476562
784
+ ],
785
+ [
786
+ -22.886153602600096,
787
+ -54.5494499206543,
788
+ -37.67033004760742,
789
+ -21.4945068359375,
790
+ -16.439559936523438
791
+ ],
792
+ [
793
+ -24.04395643234253,
794
+ -58.49054851531982,
795
+ -40.0439567565918,
796
+ -22.72527629852295,
797
+ -17.14285659790039
798
+ ],
799
+ [
800
+ -25.186813354492188,
801
+ -61.919121475219725,
802
+ -42.593406677246094,
803
+ -23.78022003173828,
804
+ -18.021976470947266
805
+ ],
806
+ [
807
+ -26.41758155822754,
808
+ -65.43560859680176,
809
+ -44.864172706604,
810
+ -24.908128814697264,
811
+ -18.813186645507812
812
+ ],
813
+ [
814
+ -27.64834976196289,
815
+ -68.52747344970703,
816
+ -47.07692337036133,
817
+ -25.890113830566406,
818
+ -19.765275058746337
819
+ ],
820
+ [
821
+ -28.879119873046875,
822
+ -71.85318824768066,
823
+ -49.18681335449219,
824
+ -26.842194938659667,
825
+ -20.571430206298828
826
+ ],
827
+ [
828
+ -29.84615484237671,
829
+ -74.84220191955566,
830
+ -51.120880126953125,
831
+ -27.809228324890135,
832
+ -21.362638473510742
833
+ ]
834
+ ],
835
+ "q99": [
836
+ [
837
+ 10.329669952392578,
838
+ 15.956047058105469,
839
+ 11.472527341842651,
840
+ 6.989013671875,
841
+ 6.417582430839539
842
+ ],
843
+ [
844
+ 13.05494499206543,
845
+ 20.967036876678467,
846
+ 15.677361869811845,
847
+ 9.186813354492188,
848
+ 8.000000791549683
849
+ ],
850
+ [
851
+ 15.516484098434448,
852
+ 25.97801971435547,
853
+ 19.457581901550128,
854
+ 11.208793640136719,
855
+ 9.655383491515947
856
+ ],
857
+ [
858
+ 17.9780216217041,
859
+ 30.798241996764972,
860
+ 23.41362514495833,
861
+ 13.054946899414062,
862
+ 11.340659141540527
863
+ ],
864
+ [
865
+ 20.61538314819336,
866
+ 35.4725297164917,
867
+ 27.38461462020874,
868
+ 14.901100158691406,
869
+ 12.923076629638672
870
+ ],
871
+ [
872
+ 23.325711956024005,
873
+ 39.956046409606934,
874
+ 30.901100158691406,
875
+ 16.571426391601562,
876
+ 14.681318283081055
877
+ ],
878
+ [
879
+ 25.89011001586914,
880
+ 44.424612121581866,
881
+ 34.68131954193115,
882
+ 17.890113830566406,
883
+ 16.33670335769637
884
+ ],
885
+ [
886
+ 28.512526893615558,
887
+ 48.30769348144531,
888
+ 37.93406677246094,
889
+ 19.457588233947586,
890
+ 17.846153259277344
891
+ ],
892
+ [
893
+ 30.813186321258545,
894
+ 52.07296756744368,
895
+ 41.259781570434406,
896
+ 20.776262359618975,
897
+ 19.340660095214844
898
+ ],
899
+ [
900
+ 33.34769191741927,
901
+ 55.76527221679671,
902
+ 44.527472801208496,
903
+ 21.934066772460938,
904
+ 20.908131656646564
905
+ ],
906
+ [
907
+ 35.736263275146484,
908
+ 59.20878917694092,
909
+ 47.516483306884766,
910
+ 22.813186645507812,
911
+ 22.31472497940047
912
+ ],
913
+ [
914
+ 38.007034988403156,
915
+ 62.28571701049805,
916
+ 50.82725028991666,
917
+ 23.85318759918196,
918
+ 23.648351669311523
919
+ ],
920
+ [
921
+ 40.131869621276856,
922
+ 65.69933517456037,
923
+ 53.23076946258545,
924
+ 24.659346237182618,
925
+ 24.9670330619812
926
+ ],
927
+ [
928
+ 42.49054851531966,
929
+ 68.23384864807096,
930
+ 56.29274818420394,
931
+ 25.538459129333496,
932
+ 26.10988998413086
933
+ ],
934
+ [
935
+ 44.263736724853516,
936
+ 70.72527313232422,
937
+ 58.8421949386595,
938
+ 26.241756439208984,
939
+ 27.149891586303546
940
+ ],
941
+ [
942
+ 46.461539573669434,
943
+ 73.36263275146484,
944
+ 61.30373451232894,
945
+ 27.120880126953125,
946
+ 28.307693481445312
947
+ ]
948
+ ]
949
+ },
950
+ "right_arm": {
951
+ "min": [
952
+ [
953
+ -13.7582426071167,
954
+ -21.9780216217041,
955
+ -27.73626708984375,
956
+ -26.901100158691406,
957
+ -14.505494117736816
958
+ ],
959
+ [
960
+ -18.065933227539062,
961
+ -27.692306518554688,
962
+ -31.340660095214844,
963
+ -26.901100158691406,
964
+ -18.549449920654297
965
+ ],
966
+ [
967
+ -21.758241653442383,
968
+ -34.021976470947266,
969
+ -35.20879364013672,
970
+ -26.901100158691406,
971
+ -22.417583465576172
972
+ ],
973
+ [
974
+ -25.362638473510742,
975
+ -40.08790969848633,
976
+ -39.25274658203125,
977
+ -26.901100158691406,
978
+ -25.846153259277344
979
+ ],
980
+ [
981
+ -28.52747344970703,
982
+ -46.06593322753906,
983
+ -45.05494689941406,
984
+ -27.78022003173828,
985
+ -28.307693481445312
986
+ ],
987
+ [
988
+ -31.164836883544922,
989
+ -51.868133544921875,
990
+ -51.03296661376953,
991
+ -29.802200317382812,
992
+ -29.97801971435547
993
+ ],
994
+ [
995
+ -34.24176025390625,
996
+ -57.142852783203125,
997
+ -57.010990142822266,
998
+ -32.0,
999
+ -30.85714340209961
1000
+ ],
1001
+ [
1002
+ -36.52747344970703,
1003
+ -62.153846740722656,
1004
+ -61.67033004760742,
1005
+ -34.54945373535156,
1006
+ -31.736265182495117
1007
+ ],
1008
+ [
1009
+ -38.10988998413086,
1010
+ -67.16484069824219,
1011
+ -66.4175796508789,
1012
+ -36.57143020629883,
1013
+ -33.31867980957031
1014
+ ],
1015
+ [
1016
+ -39.60439682006836,
1017
+ -72.26374053955078,
1018
+ -71.16483306884766,
1019
+ -38.505496978759766,
1020
+ -34.63736343383789
1021
+ ],
1022
+ [
1023
+ -40.74725341796875,
1024
+ -77.8901138305664,
1025
+ -75.29670715332031,
1026
+ -40.08790969848633,
1027
+ -36.21977996826172
1028
+ ],
1029
+ [
1030
+ -42.593406677246094,
1031
+ -83.34066009521484,
1032
+ -80.04395294189453,
1033
+ -42.10988998413086,
1034
+ -37.626373291015625
1035
+ ],
1036
+ [
1037
+ -44.0,
1038
+ -88.0,
1039
+ -84.79120635986328,
1040
+ -43.516483306884766,
1041
+ -39.384613037109375
1042
+ ],
1043
+ [
1044
+ -45.31867980957031,
1045
+ -92.30769348144531,
1046
+ -87.78022003173828,
1047
+ -45.010990142822266,
1048
+ -40.79120635986328
1049
+ ],
1050
+ [
1051
+ -46.021976470947266,
1052
+ -97.31867980957031,
1053
+ -90.06593322753906,
1054
+ -46.76923370361328,
1055
+ -42.10988998413086
1056
+ ],
1057
+ [
1058
+ -46.63736343383789,
1059
+ -101.36264038085938,
1060
+ -92.96703338623047,
1061
+ -48.5274772644043,
1062
+ -42.46154022216797
1063
+ ]
1064
+ ],
1065
+ "max": [
1066
+ [
1067
+ 16.483516693115234,
1068
+ 25.846149444580078,
1069
+ 25.010990142822266,
1070
+ 16.967037200927734,
1071
+ 21.36263656616211
1072
+ ],
1073
+ [
1074
+ 16.659339904785156,
1075
+ 29.802196502685547,
1076
+ 32.483516693115234,
1077
+ 21.45055389404297,
1078
+ 27.16483497619629
1079
+ ],
1080
+ [
1081
+ 20.61538314819336,
1082
+ 33.23076629638672,
1083
+ 39.25274658203125,
1084
+ 25.406597137451172,
1085
+ 32.96703338623047
1086
+ ],
1087
+ [
1088
+ 24.39560317993164,
1089
+ 36.21977615356445,
1090
+ 45.846153259277344,
1091
+ 27.95604705810547,
1092
+ 37.80220031738281
1093
+ ],
1094
+ [
1095
+ 28.0,
1096
+ 40.08790969848633,
1097
+ 52.263736724853516,
1098
+ 30.153846740722656,
1099
+ 42.54945373535156
1100
+ ],
1101
+ [
1102
+ 31.164833068847656,
1103
+ 44.83516311645508,
1104
+ 58.5054931640625,
1105
+ 33.05494689941406,
1106
+ 46.329673767089844
1107
+ ],
1108
+ [
1109
+ 34.241756439208984,
1110
+ 49.0549430847168,
1111
+ 64.83516693115234,
1112
+ 36.4835205078125,
1113
+ 49.406593322753906
1114
+ ],
1115
+ [
1116
+ 36.96703338623047,
1117
+ 52.83516311645508,
1118
+ 70.72527313232422,
1119
+ 39.56044006347656,
1120
+ 52.21977996826172
1121
+ ],
1122
+ [
1123
+ 39.78022003173828,
1124
+ 57.58241653442383,
1125
+ 76.61538696289062,
1126
+ 41.93406677246094,
1127
+ 54.06593322753906
1128
+ ],
1129
+ [
1130
+ 42.68131637573242,
1131
+ 61.80219650268555,
1132
+ 82.24176025390625,
1133
+ 45.010986328125,
1134
+ 56.087913513183594
1135
+ ],
1136
+ [
1137
+ 45.670326232910156,
1138
+ 66.46154022216797,
1139
+ 87.78022003173828,
1140
+ 47.82417678833008,
1141
+ 59.78022384643555
1142
+ ],
1143
+ [
1144
+ 48.395606994628906,
1145
+ 70.68132019042969,
1146
+ 93.31867980957031,
1147
+ 50.5494499206543,
1148
+ 63.120880126953125
1149
+ ],
1150
+ [
1151
+ 51.20879364013672,
1152
+ 74.81318664550781,
1153
+ 98.85714721679688,
1154
+ 53.274723052978516,
1155
+ 66.10989379882812
1156
+ ],
1157
+ [
1158
+ 53.75823974609375,
1159
+ 78.94505310058594,
1160
+ 104.04396057128906,
1161
+ 55.736263275146484,
1162
+ 69.97802734375
1163
+ ],
1164
+ [
1165
+ 56.21977996826172,
1166
+ 81.4945068359375,
1167
+ 109.14285278320312,
1168
+ 57.846153259277344,
1169
+ 71.91209411621094
1170
+ ],
1171
+ [
1172
+ 58.769229888916016,
1173
+ 84.13186645507812,
1174
+ 113.71429443359375,
1175
+ 59.69230651855469,
1176
+ 73.23077392578125
1177
+ ]
1178
+ ],
1179
+ "mean": [
1180
+ [
1181
+ 0.09307948499917984,
1182
+ -0.21033740043640137,
1183
+ -1.2144449949264526,
1184
+ -0.27963724732398987,
1185
+ 0.04348327964544296
1186
+ ],
1187
+ [
1188
+ 0.10010828077793121,
1189
+ -0.2093687504529953,
1190
+ -1.2162504196166992,
1191
+ -0.2847459018230438,
1192
+ 0.04417652264237404
1193
+ ],
1194
+ [
1195
+ 0.1072654128074646,
1196
+ -0.2083953619003296,
1197
+ -1.21742582321167,
1198
+ -0.2884536385536194,
1199
+ 0.04485388100147247
1200
+ ],
1201
+ [
1202
+ 0.11462292075157166,
1203
+ -0.20741455256938934,
1204
+ -1.2182234525680542,
1205
+ -0.2908605933189392,
1206
+ 0.0456455796957016
1207
+ ],
1208
+ [
1209
+ 0.12212661653757095,
1210
+ -0.20648866891860962,
1211
+ -1.2188026905059814,
1212
+ -0.29220396280288696,
1213
+ 0.04649243503808975
1214
+ ],
1215
+ [
1216
+ 0.12966646254062653,
1217
+ -0.2055509388446808,
1218
+ -1.2192649841308594,
1219
+ -0.2931065261363983,
1220
+ 0.04740649834275246
1221
+ ],
1222
+ [
1223
+ 0.1372615396976471,
1224
+ -0.2046326845884323,
1225
+ -1.2195936441421509,
1226
+ -0.29365071654319763,
1227
+ 0.04836377501487732
1228
+ ],
1229
+ [
1230
+ 0.1448957324028015,
1231
+ -0.203730970621109,
1232
+ -1.2196241617202759,
1233
+ -0.29413872957229614,
1234
+ 0.04934469982981682
1235
+ ],
1236
+ [
1237
+ 0.15254466235637665,
1238
+ -0.2028094083070755,
1239
+ -1.21953284740448,
1240
+ -0.29474878311157227,
1241
+ 0.05036116763949394
1242
+ ],
1243
+ [
1244
+ 0.16028816998004913,
1245
+ -0.2019113451242447,
1246
+ -1.21940016746521,
1247
+ -0.2954031229019165,
1248
+ 0.051397208124399185
1249
+ ],
1250
+ [
1251
+ 0.1685880720615387,
1252
+ -0.2010054737329483,
1253
+ -1.2190954685211182,
1254
+ -0.29616355895996094,
1255
+ 0.05248061940073967
1256
+ ],
1257
+ [
1258
+ 0.17761225998401642,
1259
+ -0.20010292530059814,
1260
+ -1.218703269958496,
1261
+ -0.29693999886512756,
1262
+ 0.053642578423023224
1263
+ ],
1264
+ [
1265
+ 0.18745994567871094,
1266
+ -0.19920113682746887,
1267
+ -1.2181789875030518,
1268
+ -0.29780125617980957,
1269
+ 0.054784782230854034
1270
+ ],
1271
+ [
1272
+ 0.19815856218338013,
1273
+ -0.19830350577831268,
1274
+ -1.2174657583236694,
1275
+ -0.2987440526485443,
1276
+ 0.0559191033244133
1277
+ ],
1278
+ [
1279
+ 0.20982207357883453,
1280
+ -0.19740089774131775,
1281
+ -1.2163711786270142,
1282
+ -0.29971298575401306,
1283
+ 0.057061586529016495
1284
+ ],
1285
+ [
1286
+ 0.22255578637123108,
1287
+ -0.19650320708751678,
1288
+ -1.2147578001022339,
1289
+ -0.30066221952438354,
1290
+ 0.05819998309016228
1291
+ ]
1292
+ ],
1293
+ "std": [
1294
+ [
1295
+ 2.368210554122925,
1296
+ 4.5004777908325195,
1297
+ 4.228534698486328,
1298
+ 2.4653029441833496,
1299
+ 2.167224407196045
1300
+ ],
1301
+ [
1302
+ 3.026092052459717,
1303
+ 5.872717380523682,
1304
+ 5.4169535636901855,
1305
+ 3.2118327617645264,
1306
+ 2.7414181232452393
1307
+ ],
1308
+ [
1309
+ 3.667999505996704,
1310
+ 7.226501941680908,
1311
+ 6.591485023498535,
1312
+ 3.9331958293914795,
1313
+ 3.2945902347564697
1314
+ ],
1315
+ [
1316
+ 4.287671089172363,
1317
+ 8.5562162399292,
1318
+ 7.745312213897705,
1319
+ 4.626955986022949,
1320
+ 3.822075366973877
1321
+ ],
1322
+ [
1323
+ 4.882786750793457,
1324
+ 9.859098434448242,
1325
+ 8.87535285949707,
1326
+ 5.29279899597168,
1327
+ 4.323612689971924
1328
+ ],
1329
+ [
1330
+ 5.453670501708984,
1331
+ 11.133562088012695,
1332
+ 9.979808807373047,
1333
+ 5.9306182861328125,
1334
+ 4.800647258758545
1335
+ ],
1336
+ [
1337
+ 6.0010576248168945,
1338
+ 12.37905216217041,
1339
+ 11.058042526245117,
1340
+ 6.5409417152404785,
1341
+ 5.255186080932617
1342
+ ],
1343
+ [
1344
+ 6.526202201843262,
1345
+ 13.595319747924805,
1346
+ 12.110023498535156,
1347
+ 7.124155044555664,
1348
+ 5.688719749450684
1349
+ ],
1350
+ [
1351
+ 7.030229091644287,
1352
+ 14.782575607299805,
1353
+ 13.13562297821045,
1354
+ 7.680479049682617,
1355
+ 6.1026129722595215
1356
+ ],
1357
+ [
1358
+ 7.514631748199463,
1359
+ 15.941143035888672,
1360
+ 14.134993553161621,
1361
+ 8.21074390411377,
1362
+ 6.498264789581299
1363
+ ],
1364
+ [
1365
+ 7.980932235717773,
1366
+ 17.071033477783203,
1367
+ 15.108572006225586,
1368
+ 8.715429306030273,
1369
+ 6.877125263214111
1370
+ ],
1371
+ [
1372
+ 8.430272102355957,
1373
+ 18.172752380371094,
1374
+ 16.056636810302734,
1375
+ 9.195815086364746,
1376
+ 7.240087985992432
1377
+ ],
1378
+ [
1379
+ 8.863301277160645,
1380
+ 19.24644660949707,
1381
+ 16.979869842529297,
1382
+ 9.652762413024902,
1383
+ 7.587943077087402
1384
+ ],
1385
+ [
1386
+ 9.28066635131836,
1387
+ 20.29257583618164,
1388
+ 17.878612518310547,
1389
+ 10.087113380432129,
1390
+ 7.9214396476745605
1391
+ ],
1392
+ [
1393
+ 9.683401107788086,
1394
+ 21.3113956451416,
1395
+ 18.75370979309082,
1396
+ 10.499829292297363,
1397
+ 8.241240501403809
1398
+ ],
1399
+ [
1400
+ 10.07201099395752,
1401
+ 22.303380966186523,
1402
+ 19.60598373413086,
1403
+ 10.892108917236328,
1404
+ 8.548293113708496
1405
+ ]
1406
+ ],
1407
+ "q01": [
1408
+ [
1409
+ -6.8131866455078125,
1410
+ -15.120880126953125,
1411
+ -12.879119873046875,
1412
+ -7.80922897338867,
1413
+ -6.578461236953734
1414
+ ],
1415
+ [
1416
+ -9.010988235473633,
1417
+ -19.60439682006836,
1418
+ -16.380655364990233,
1419
+ -10.10988998413086,
1420
+ -8.527473449707031
1421
+ ],
1422
+ [
1423
+ -10.85714340209961,
1424
+ -24.0,
1425
+ -19.545491924285887,
1426
+ -12.4835160446167,
1427
+ -10.285715103149414
1428
+ ],
1429
+ [
1430
+ -12.615385055541992,
1431
+ -28.307691249847412,
1432
+ -22.725272483825684,
1433
+ -14.681320190429688,
1434
+ -11.956044826507569
1435
+ ],
1436
+ [
1437
+ -14.373626708984375,
1438
+ -32.351649589538575,
1439
+ -25.699341506958007,
1440
+ -16.615386962890625,
1441
+ -13.61142831802368
1442
+ ],
1443
+ [
1444
+ -15.868130521774292,
1445
+ -36.48351287841797,
1446
+ -28.43956184387207,
1447
+ -18.637359619140625,
1448
+ -14.681318283081055
1449
+ ],
1450
+ [
1451
+ -17.274726543426514,
1452
+ -40.52747344970703,
1453
+ -30.9890114402771,
1454
+ -20.571426391601562,
1455
+ -16.000001583099365
1456
+ ],
1457
+ [
1458
+ -18.593406677246094,
1459
+ -44.39560317993164,
1460
+ -33.61142547607422,
1461
+ -22.065933227539062,
1462
+ -17.230769786834717
1463
+ ],
1464
+ [
1465
+ -19.91208839416504,
1466
+ -48.160881729125975,
1467
+ -35.89714183807373,
1468
+ -23.648353576660156,
1469
+ -18.270768547058104
1470
+ ],
1471
+ [
1472
+ -20.96703338623047,
1473
+ -52.0439567565918,
1474
+ -38.18285820007324,
1475
+ -25.23076629638672,
1476
+ -19.516482982635498
1477
+ ],
1478
+ [
1479
+ -22.182856616973876,
1480
+ -55.64835292816162,
1481
+ -40.0439567565918,
1482
+ -26.710329132080076,
1483
+ -20.659339904785156
1484
+ ],
1485
+ [
1486
+ -23.076923370361328,
1487
+ -58.97406520843506,
1488
+ -42.13889892578125,
1489
+ -28.131866455078125,
1490
+ -21.538461685180664
1491
+ ],
1492
+ [
1493
+ -24.043956756591797,
1494
+ -62.57846202850342,
1495
+ -44.08791286468506,
1496
+ -29.347691917419432,
1497
+ -22.417583465576172
1498
+ ],
1499
+ [
1500
+ -24.923076629638672,
1501
+ -65.84615325927734,
1502
+ -45.84615577697754,
1503
+ -30.593406352996826,
1504
+ -23.648351345062256
1505
+ ],
1506
+ [
1507
+ -26.065933227539062,
1508
+ -69.27472686767578,
1509
+ -47.58944835662842,
1510
+ -31.80922897338867,
1511
+ -24.512528476715087
1512
+ ],
1513
+ [
1514
+ -26.94505500793457,
1515
+ -72.43955993652344,
1516
+ -49.18681587219238,
1517
+ -32.87912368774414,
1518
+ -25.040001277923583
1519
+ ]
1520
+ ],
1521
+ "q99": [
1522
+ [
1523
+ 8.835165023803711,
1524
+ 12.219779968261719,
1525
+ 12.952087621688676,
1526
+ 7.648345947265625,
1527
+ 9.142857389450073
1528
+ ],
1529
+ [
1530
+ 11.120879173278809,
1531
+ 16.087913513183594,
1532
+ 17.36263656616211,
1533
+ 10.02197265625,
1534
+ 11.428570585250855
1535
+ ],
1536
+ [
1537
+ 13.230769157409668,
1538
+ 19.868133544921875,
1539
+ 21.846153259277344,
1540
+ 12.29274250030501,
1541
+ 13.538461360931397
1542
+ ],
1543
+ [
1544
+ 15.237801933288408,
1545
+ 23.736266441345215,
1546
+ 26.065933227539062,
1547
+ 14.329673767089844,
1548
+ 15.560439901351929
1549
+ ],
1550
+ [
1551
+ 17.36263656616211,
1552
+ 27.25274658203125,
1553
+ 30.25582357406583,
1554
+ 16.26373291015625,
1555
+ 17.230770111083984
1556
+ ],
1557
+ [
1558
+ 19.3696699714659,
1559
+ 31.01802196502669,
1560
+ 34.57845886230452,
1561
+ 17.846153259277344,
1562
+ 19.41362514495833
1563
+ ],
1564
+ [
1565
+ 21.230769786834717,
1566
+ 34.28571319580078,
1567
+ 38.5494499206543,
1568
+ 19.428572959899903,
1569
+ 21.186813354492188
1570
+ ],
1571
+ [
1572
+ 23.16483497619629,
1573
+ 37.626373291015625,
1574
+ 42.66637172698958,
1575
+ 20.74725658416748,
1576
+ 23.193845176696612
1577
+ ],
1578
+ [
1579
+ 24.74725341796875,
1580
+ 40.87912368774414,
1581
+ 46.27076854705794,
1582
+ 22.153846740722656,
1583
+ 25.215823230743243
1584
+ ],
1585
+ [
1586
+ 26.226811790466144,
1587
+ 43.95604705810547,
1588
+ 50.490547866821125,
1589
+ 23.47252655029297,
1590
+ 27.23780193328841
1591
+ ],
1592
+ [
1593
+ 27.54549350738509,
1594
+ 46.857142753601075,
1595
+ 53.9340661239624,
1596
+ 25.010108489989737,
1597
+ 28.878241386413077
1598
+ ],
1599
+ [
1600
+ 28.703296661376953,
1601
+ 49.846156425476075,
1602
+ 57.18681335449219,
1603
+ 25.934066123962403,
1604
+ 30.930108451843097
1605
+ ],
1606
+ [
1607
+ 30.02197805404663,
1608
+ 52.483516693115234,
1609
+ 60.864175224304034,
1610
+ 26.959119911193515,
1611
+ 32.17582321166992
1612
+ ],
1613
+ [
1614
+ 31.164836559295654,
1615
+ 55.20878917694092,
1616
+ 64.11692245483381,
1617
+ 28.116921806335284,
1618
+ 33.14285659790039
1619
+ ],
1620
+ [
1621
+ 32.39560317993164,
1622
+ 57.743294448852375,
1623
+ 67.50066017150813,
1624
+ 28.923076629638672,
1625
+ 34.622418136596515
1626
+ ],
1627
+ [
1628
+ 33.53846294403076,
1629
+ 60.029012107848956,
1630
+ 70.2707678985594,
1631
+ 29.450547847747803,
1632
+ 35.868133544921875
1633
+ ]
1634
+ ]
1635
+ }
1636
+ }
1637
+ }
1638
+ }
experiment_cfg/final_model_config.json ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_type": "Gr00tN1d7",
3
+ "model_dtype": "bfloat16",
4
+ "model_name": "nvidia/Cosmos-Reason2-2B",
5
+ "backbone_model_type": "qwen",
6
+ "model_revision": null,
7
+ "tune_top_llm_layers": 0,
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
+ "backbone_trainable_params_fp32": true,
16
+ "letter_box_transform": false,
17
+ "extra_augmentation_config": null,
18
+ "apply_sincos_state_encoding": false,
19
+ "use_percentiles": true,
20
+ "use_relative_action": false,
21
+ "max_state_dim": 132,
22
+ "max_action_dim": 132,
23
+ "action_horizon": 40,
24
+ "hidden_size": 1024,
25
+ "input_embedding_dim": 1536,
26
+ "state_history_length": 1,
27
+ "add_pos_embed": true,
28
+ "attn_dropout": 0.2,
29
+ "use_vlln": true,
30
+ "max_seq_len": 1024,
31
+ "use_alternate_vl_dit": true,
32
+ "attend_text_every_n_blocks": 2,
33
+ "diffusion_model_cfg": {
34
+ "attention_head_dim": 48,
35
+ "dropout": 0.2,
36
+ "final_dropout": true,
37
+ "interleave_self_attention": true,
38
+ "norm_type": "ada_norm",
39
+ "num_attention_heads": 32,
40
+ "num_layers": 32,
41
+ "output_dim": 1024,
42
+ "positional_embeddings": null
43
+ },
44
+ "num_inference_timesteps": 4,
45
+ "noise_beta_alpha": 1.5,
46
+ "noise_beta_beta": 1.0,
47
+ "noise_s": 0.999,
48
+ "num_timestep_buckets": 1000,
49
+ "tune_projector": true,
50
+ "tune_diffusion_model": true,
51
+ "tune_vlln": true,
52
+ "state_dropout_prob": 0.2,
53
+ "exclude_state": false,
54
+ "use_mean_std": false,
55
+ "max_num_embodiments": 32
56
+ }
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:62507a4c6cf1f37dc915b06de4ba4fa161dec720c07e48709f659b435f0c1a80
3
+ size 4990519232
model-00002-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c6f6e65091d3f01da9b9e797dfbaf25a83c5632e89797c1921a52756226d34f3
3
+ size 1919980184
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
processor_config.json ADDED
@@ -0,0 +1,1140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "processor_class": "Gr00tN1d7Processor",
3
+ "processor_kwargs": {
4
+ "modality_configs": {
5
+ "real_g1_relative_eef_relative_joints": {
6
+ "video": {
7
+ "delta_indices": [
8
+ -20,
9
+ 0
10
+ ],
11
+ "modality_keys": [
12
+ "ego_view"
13
+ ],
14
+ "sin_cos_embedding_keys": null,
15
+ "mean_std_embedding_keys": null,
16
+ "action_configs": null
17
+ },
18
+ "state": {
19
+ "delta_indices": [
20
+ 0
21
+ ],
22
+ "modality_keys": [
23
+ "left_wrist_eef_9d",
24
+ "right_wrist_eef_9d",
25
+ "left_hand",
26
+ "right_hand",
27
+ "left_arm",
28
+ "right_arm",
29
+ "waist"
30
+ ],
31
+ "sin_cos_embedding_keys": null,
32
+ "mean_std_embedding_keys": null,
33
+ "action_configs": null
34
+ },
35
+ "action": {
36
+ "delta_indices": [
37
+ 0,
38
+ 1,
39
+ 2,
40
+ 3,
41
+ 4,
42
+ 5,
43
+ 6,
44
+ 7,
45
+ 8,
46
+ 9,
47
+ 10,
48
+ 11,
49
+ 12,
50
+ 13,
51
+ 14,
52
+ 15,
53
+ 16,
54
+ 17,
55
+ 18,
56
+ 19,
57
+ 20,
58
+ 21,
59
+ 22,
60
+ 23,
61
+ 24,
62
+ 25,
63
+ 26,
64
+ 27,
65
+ 28,
66
+ 29,
67
+ 30,
68
+ 31,
69
+ 32,
70
+ 33,
71
+ 34,
72
+ 35,
73
+ 36,
74
+ 37,
75
+ 38,
76
+ 39
77
+ ],
78
+ "modality_keys": [
79
+ "left_wrist_eef_9d",
80
+ "right_wrist_eef_9d",
81
+ "left_hand",
82
+ "right_hand",
83
+ "left_arm",
84
+ "right_arm",
85
+ "waist",
86
+ "base_height_command",
87
+ "navigate_command"
88
+ ],
89
+ "sin_cos_embedding_keys": null,
90
+ "mean_std_embedding_keys": null,
91
+ "action_configs": [
92
+ {
93
+ "rep": "RELATIVE",
94
+ "type": "EEF",
95
+ "format": "XYZ_ROT6D",
96
+ "state_key": "left_wrist_eef_9d"
97
+ },
98
+ {
99
+ "rep": "RELATIVE",
100
+ "type": "EEF",
101
+ "format": "XYZ_ROT6D",
102
+ "state_key": "right_wrist_eef_9d"
103
+ },
104
+ {
105
+ "rep": "ABSOLUTE",
106
+ "type": "NON_EEF",
107
+ "format": "DEFAULT",
108
+ "state_key": "left_hand"
109
+ },
110
+ {
111
+ "rep": "ABSOLUTE",
112
+ "type": "NON_EEF",
113
+ "format": "DEFAULT",
114
+ "state_key": "right_hand"
115
+ },
116
+ {
117
+ "rep": "RELATIVE",
118
+ "type": "NON_EEF",
119
+ "format": "DEFAULT",
120
+ "state_key": "left_arm"
121
+ },
122
+ {
123
+ "rep": "RELATIVE",
124
+ "type": "NON_EEF",
125
+ "format": "DEFAULT",
126
+ "state_key": "right_arm"
127
+ },
128
+ {
129
+ "rep": "ABSOLUTE",
130
+ "type": "NON_EEF",
131
+ "format": "DEFAULT",
132
+ "state_key": "waist"
133
+ },
134
+ {
135
+ "rep": "ABSOLUTE",
136
+ "type": "NON_EEF",
137
+ "format": "DEFAULT",
138
+ "state_key": "base_height_command"
139
+ },
140
+ {
141
+ "rep": "ABSOLUTE",
142
+ "type": "NON_EEF",
143
+ "format": "DEFAULT",
144
+ "state_key": "navigate_command"
145
+ }
146
+ ]
147
+ },
148
+ "language": {
149
+ "delta_indices": [
150
+ 0
151
+ ],
152
+ "modality_keys": [
153
+ "annotation.human.task_description"
154
+ ],
155
+ "sin_cos_embedding_keys": null,
156
+ "mean_std_embedding_keys": null,
157
+ "action_configs": null
158
+ }
159
+ },
160
+ "real_r1_pro_sharpa_relative_eef_mecka": {
161
+ "video": {
162
+ "delta_indices": [
163
+ -30,
164
+ 0
165
+ ],
166
+ "modality_keys": [
167
+ "ego_view_cropratio_res320x240_freq30"
168
+ ],
169
+ "sin_cos_embedding_keys": null,
170
+ "mean_std_embedding_keys": null,
171
+ "action_configs": null
172
+ },
173
+ "state": {
174
+ "delta_indices": [
175
+ 0
176
+ ],
177
+ "modality_keys": [
178
+ "left_wrist_eef",
179
+ "right_wrist_eef",
180
+ "left_hand_joints",
181
+ "right_hand_joints"
182
+ ],
183
+ "sin_cos_embedding_keys": null,
184
+ "mean_std_embedding_keys": null,
185
+ "action_configs": null
186
+ },
187
+ "action": {
188
+ "delta_indices": [
189
+ 0,
190
+ 1,
191
+ 2,
192
+ 3,
193
+ 4,
194
+ 5,
195
+ 6,
196
+ 7,
197
+ 8,
198
+ 9,
199
+ 10,
200
+ 11,
201
+ 12,
202
+ 13,
203
+ 14,
204
+ 15,
205
+ 16,
206
+ 17,
207
+ 18,
208
+ 19,
209
+ 20,
210
+ 21,
211
+ 22,
212
+ 23,
213
+ 24,
214
+ 25,
215
+ 26,
216
+ 27,
217
+ 28,
218
+ 29,
219
+ 30,
220
+ 31,
221
+ 32,
222
+ 33,
223
+ 34,
224
+ 35,
225
+ 36,
226
+ 37,
227
+ 38,
228
+ 39
229
+ ],
230
+ "modality_keys": [
231
+ "left_wrist_eef",
232
+ "right_wrist_eef",
233
+ "left_hand_joints",
234
+ "right_hand_joints"
235
+ ],
236
+ "sin_cos_embedding_keys": null,
237
+ "mean_std_embedding_keys": null,
238
+ "action_configs": [
239
+ {
240
+ "rep": "RELATIVE",
241
+ "type": "EEF",
242
+ "format": "XYZ_ROT6D",
243
+ "state_key": "left_wrist_eef"
244
+ },
245
+ {
246
+ "rep": "RELATIVE",
247
+ "type": "EEF",
248
+ "format": "XYZ_ROT6D",
249
+ "state_key": "right_wrist_eef"
250
+ },
251
+ {
252
+ "rep": "ABSOLUTE",
253
+ "type": "NON_EEF",
254
+ "format": "DEFAULT",
255
+ "state_key": "left_hand_joints"
256
+ },
257
+ {
258
+ "rep": "ABSOLUTE",
259
+ "type": "NON_EEF",
260
+ "format": "DEFAULT",
261
+ "state_key": "right_hand_joints"
262
+ }
263
+ ]
264
+ },
265
+ "language": {
266
+ "delta_indices": [
267
+ 0
268
+ ],
269
+ "modality_keys": [
270
+ "annotation.human.coarse_action"
271
+ ],
272
+ "sin_cos_embedding_keys": null,
273
+ "mean_std_embedding_keys": null,
274
+ "action_configs": null
275
+ }
276
+ },
277
+ "real_r1_pro_sharpa_relative_eef_human": {
278
+ "video": {
279
+ "delta_indices": [
280
+ -20,
281
+ 0
282
+ ],
283
+ "modality_keys": [
284
+ "ego_view_res320x240_freq20",
285
+ "left_wrist_view_res320x240_freq20",
286
+ "right_wrist_view_res320x240_freq20"
287
+ ],
288
+ "sin_cos_embedding_keys": null,
289
+ "mean_std_embedding_keys": null,
290
+ "action_configs": null
291
+ },
292
+ "state": {
293
+ "delta_indices": [
294
+ 0
295
+ ],
296
+ "modality_keys": [
297
+ "left_wrist_eef",
298
+ "right_wrist_eef",
299
+ "left_hand_joints",
300
+ "right_hand_joints"
301
+ ],
302
+ "sin_cos_embedding_keys": null,
303
+ "mean_std_embedding_keys": null,
304
+ "action_configs": null
305
+ },
306
+ "action": {
307
+ "delta_indices": [
308
+ 0,
309
+ 1,
310
+ 2,
311
+ 3,
312
+ 4,
313
+ 5,
314
+ 6,
315
+ 7,
316
+ 8,
317
+ 9,
318
+ 10,
319
+ 11,
320
+ 12,
321
+ 13,
322
+ 14,
323
+ 15,
324
+ 16,
325
+ 17,
326
+ 18,
327
+ 19,
328
+ 20,
329
+ 21,
330
+ 22,
331
+ 23,
332
+ 24,
333
+ 25,
334
+ 26,
335
+ 27,
336
+ 28,
337
+ 29,
338
+ 30,
339
+ 31,
340
+ 32,
341
+ 33,
342
+ 34,
343
+ 35,
344
+ 36,
345
+ 37,
346
+ 38,
347
+ 39
348
+ ],
349
+ "modality_keys": [
350
+ "left_wrist_eef",
351
+ "right_wrist_eef",
352
+ "left_hand_joints",
353
+ "right_hand_joints"
354
+ ],
355
+ "sin_cos_embedding_keys": null,
356
+ "mean_std_embedding_keys": null,
357
+ "action_configs": [
358
+ {
359
+ "rep": "RELATIVE",
360
+ "type": "EEF",
361
+ "format": "XYZ_ROT6D",
362
+ "state_key": "left_wrist_eef"
363
+ },
364
+ {
365
+ "rep": "RELATIVE",
366
+ "type": "EEF",
367
+ "format": "XYZ_ROT6D",
368
+ "state_key": "right_wrist_eef"
369
+ },
370
+ {
371
+ "rep": "ABSOLUTE",
372
+ "type": "NON_EEF",
373
+ "format": "DEFAULT",
374
+ "state_key": "left_hand_joints"
375
+ },
376
+ {
377
+ "rep": "ABSOLUTE",
378
+ "type": "NON_EEF",
379
+ "format": "DEFAULT",
380
+ "state_key": "right_hand_joints"
381
+ }
382
+ ]
383
+ },
384
+ "language": {
385
+ "delta_indices": [
386
+ 0
387
+ ],
388
+ "modality_keys": [
389
+ "annotation.human.coarse_action"
390
+ ],
391
+ "sin_cos_embedding_keys": null,
392
+ "mean_std_embedding_keys": null,
393
+ "action_configs": null
394
+ }
395
+ },
396
+ "real_r1_pro_sharpa_relative_eef": {
397
+ "video": {
398
+ "delta_indices": [
399
+ -20,
400
+ 0
401
+ ],
402
+ "modality_keys": [
403
+ "ego_view_res320x240_freq20",
404
+ "left_wrist_view_res320x240_freq20",
405
+ "right_wrist_view_res320x240_freq20"
406
+ ],
407
+ "sin_cos_embedding_keys": null,
408
+ "mean_std_embedding_keys": null,
409
+ "action_configs": null
410
+ },
411
+ "state": {
412
+ "delta_indices": [
413
+ 0
414
+ ],
415
+ "modality_keys": [
416
+ "left_wrist_eef",
417
+ "right_wrist_eef",
418
+ "left_hand_joints",
419
+ "right_hand_joints"
420
+ ],
421
+ "sin_cos_embedding_keys": null,
422
+ "mean_std_embedding_keys": null,
423
+ "action_configs": null
424
+ },
425
+ "action": {
426
+ "delta_indices": [
427
+ 0,
428
+ 1,
429
+ 2,
430
+ 3,
431
+ 4,
432
+ 5,
433
+ 6,
434
+ 7,
435
+ 8,
436
+ 9,
437
+ 10,
438
+ 11,
439
+ 12,
440
+ 13,
441
+ 14,
442
+ 15,
443
+ 16,
444
+ 17,
445
+ 18,
446
+ 19,
447
+ 20,
448
+ 21,
449
+ 22,
450
+ 23,
451
+ 24,
452
+ 25,
453
+ 26,
454
+ 27,
455
+ 28,
456
+ 29,
457
+ 30,
458
+ 31,
459
+ 32,
460
+ 33,
461
+ 34,
462
+ 35,
463
+ 36,
464
+ 37,
465
+ 38,
466
+ 39
467
+ ],
468
+ "modality_keys": [
469
+ "left_wrist_eef",
470
+ "right_wrist_eef",
471
+ "left_hand_joints",
472
+ "right_hand_joints"
473
+ ],
474
+ "sin_cos_embedding_keys": null,
475
+ "mean_std_embedding_keys": null,
476
+ "action_configs": [
477
+ {
478
+ "rep": "RELATIVE",
479
+ "type": "EEF",
480
+ "format": "XYZ_ROT6D",
481
+ "state_key": "left_wrist_eef"
482
+ },
483
+ {
484
+ "rep": "RELATIVE",
485
+ "type": "EEF",
486
+ "format": "XYZ_ROT6D",
487
+ "state_key": "right_wrist_eef"
488
+ },
489
+ {
490
+ "rep": "ABSOLUTE",
491
+ "type": "NON_EEF",
492
+ "format": "DEFAULT",
493
+ "state_key": "left_hand_joints"
494
+ },
495
+ {
496
+ "rep": "ABSOLUTE",
497
+ "type": "NON_EEF",
498
+ "format": "DEFAULT",
499
+ "state_key": "right_hand_joints"
500
+ }
501
+ ]
502
+ },
503
+ "language": {
504
+ "delta_indices": [
505
+ 0
506
+ ],
507
+ "modality_keys": [
508
+ "annotation.human.coarse_action"
509
+ ],
510
+ "sin_cos_embedding_keys": null,
511
+ "mean_std_embedding_keys": null,
512
+ "action_configs": null
513
+ }
514
+ },
515
+ "xdof_relative_eef_relative_joint": {
516
+ "video": {
517
+ "delta_indices": [
518
+ -30,
519
+ 0
520
+ ],
521
+ "modality_keys": [
522
+ "top_camera-images-rgb_320_240",
523
+ "left_camera-images-rgb_320_240",
524
+ "right_camera-images-rgb_320_240"
525
+ ],
526
+ "sin_cos_embedding_keys": null,
527
+ "mean_std_embedding_keys": null,
528
+ "action_configs": null
529
+ },
530
+ "state": {
531
+ "delta_indices": [
532
+ 0
533
+ ],
534
+ "modality_keys": [
535
+ "left_wrist_eef",
536
+ "right_wrist_eef",
537
+ "left_gripper_pos",
538
+ "right_gripper_pos",
539
+ "left_joint_pos",
540
+ "right_joint_pos"
541
+ ],
542
+ "sin_cos_embedding_keys": null,
543
+ "mean_std_embedding_keys": null,
544
+ "action_configs": null
545
+ },
546
+ "action": {
547
+ "delta_indices": [
548
+ 0,
549
+ 1,
550
+ 2,
551
+ 3,
552
+ 4,
553
+ 5,
554
+ 6,
555
+ 7,
556
+ 8,
557
+ 9,
558
+ 10,
559
+ 11,
560
+ 12,
561
+ 13,
562
+ 14,
563
+ 15,
564
+ 16,
565
+ 17,
566
+ 18,
567
+ 19,
568
+ 20,
569
+ 21,
570
+ 22,
571
+ 23,
572
+ 24,
573
+ 25,
574
+ 26,
575
+ 27,
576
+ 28,
577
+ 29,
578
+ 30,
579
+ 31,
580
+ 32,
581
+ 33,
582
+ 34,
583
+ 35,
584
+ 36,
585
+ 37,
586
+ 38,
587
+ 39
588
+ ],
589
+ "modality_keys": [
590
+ "left_wrist_eef",
591
+ "right_wrist_eef",
592
+ "left_gripper_pos",
593
+ "right_gripper_pos",
594
+ "left_joint_pos",
595
+ "right_joint_pos"
596
+ ],
597
+ "sin_cos_embedding_keys": null,
598
+ "mean_std_embedding_keys": null,
599
+ "action_configs": [
600
+ {
601
+ "rep": "RELATIVE",
602
+ "type": "EEF",
603
+ "format": "XYZ_ROT6D",
604
+ "state_key": "left_wrist_eef"
605
+ },
606
+ {
607
+ "rep": "RELATIVE",
608
+ "type": "EEF",
609
+ "format": "XYZ_ROT6D",
610
+ "state_key": "right_wrist_eef"
611
+ },
612
+ {
613
+ "rep": "ABSOLUTE",
614
+ "type": "NON_EEF",
615
+ "format": "DEFAULT",
616
+ "state_key": "left_gripper_pos"
617
+ },
618
+ {
619
+ "rep": "ABSOLUTE",
620
+ "type": "NON_EEF",
621
+ "format": "DEFAULT",
622
+ "state_key": "right_gripper_pos"
623
+ },
624
+ {
625
+ "rep": "RELATIVE",
626
+ "type": "NON_EEF",
627
+ "format": "DEFAULT",
628
+ "state_key": "left_joint_pos"
629
+ },
630
+ {
631
+ "rep": "RELATIVE",
632
+ "type": "NON_EEF",
633
+ "format": "DEFAULT",
634
+ "state_key": "right_joint_pos"
635
+ }
636
+ ]
637
+ },
638
+ "language": {
639
+ "delta_indices": [
640
+ 0
641
+ ],
642
+ "modality_keys": [
643
+ "annotation.task"
644
+ ],
645
+ "sin_cos_embedding_keys": null,
646
+ "mean_std_embedding_keys": null,
647
+ "action_configs": null
648
+ }
649
+ },
650
+ "real_r1_pro_sharpa_relative_eef_maxinsights": {
651
+ "video": {
652
+ "delta_indices": [
653
+ -30,
654
+ 0
655
+ ],
656
+ "modality_keys": [
657
+ "ego_view_cropratio_res320x240_freq30"
658
+ ],
659
+ "sin_cos_embedding_keys": null,
660
+ "mean_std_embedding_keys": null,
661
+ "action_configs": null
662
+ },
663
+ "state": {
664
+ "delta_indices": [
665
+ 0
666
+ ],
667
+ "modality_keys": [
668
+ "left_wrist_eef",
669
+ "right_wrist_eef",
670
+ "left_hand_joints",
671
+ "right_hand_joints"
672
+ ],
673
+ "sin_cos_embedding_keys": null,
674
+ "mean_std_embedding_keys": null,
675
+ "action_configs": null
676
+ },
677
+ "action": {
678
+ "delta_indices": [
679
+ 0,
680
+ 1,
681
+ 2,
682
+ 3,
683
+ 4,
684
+ 5,
685
+ 6,
686
+ 7,
687
+ 8,
688
+ 9,
689
+ 10,
690
+ 11,
691
+ 12,
692
+ 13,
693
+ 14,
694
+ 15,
695
+ 16,
696
+ 17,
697
+ 18,
698
+ 19,
699
+ 20,
700
+ 21,
701
+ 22,
702
+ 23,
703
+ 24,
704
+ 25,
705
+ 26,
706
+ 27,
707
+ 28,
708
+ 29,
709
+ 30,
710
+ 31,
711
+ 32,
712
+ 33,
713
+ 34,
714
+ 35,
715
+ 36,
716
+ 37,
717
+ 38,
718
+ 39
719
+ ],
720
+ "modality_keys": [
721
+ "left_wrist_eef",
722
+ "right_wrist_eef",
723
+ "left_hand_joints",
724
+ "right_hand_joints"
725
+ ],
726
+ "sin_cos_embedding_keys": null,
727
+ "mean_std_embedding_keys": null,
728
+ "action_configs": [
729
+ {
730
+ "rep": "RELATIVE",
731
+ "type": "EEF",
732
+ "format": "XYZ_ROT6D",
733
+ "state_key": "left_wrist_eef"
734
+ },
735
+ {
736
+ "rep": "RELATIVE",
737
+ "type": "EEF",
738
+ "format": "XYZ_ROT6D",
739
+ "state_key": "right_wrist_eef"
740
+ },
741
+ {
742
+ "rep": "ABSOLUTE",
743
+ "type": "NON_EEF",
744
+ "format": "DEFAULT",
745
+ "state_key": "left_hand_joints"
746
+ },
747
+ {
748
+ "rep": "ABSOLUTE",
749
+ "type": "NON_EEF",
750
+ "format": "DEFAULT",
751
+ "state_key": "right_hand_joints"
752
+ }
753
+ ]
754
+ },
755
+ "language": {
756
+ "delta_indices": [
757
+ 0
758
+ ],
759
+ "modality_keys": [
760
+ "annotation.human.coarse_action"
761
+ ],
762
+ "sin_cos_embedding_keys": null,
763
+ "mean_std_embedding_keys": null,
764
+ "action_configs": null
765
+ }
766
+ },
767
+ "xdof_relative_eef_relative_joint_subtask": {
768
+ "video": {
769
+ "delta_indices": [
770
+ -30,
771
+ 0
772
+ ],
773
+ "modality_keys": [
774
+ "top_camera-images-rgb_320_240",
775
+ "left_camera-images-rgb_320_240",
776
+ "right_camera-images-rgb_320_240"
777
+ ],
778
+ "sin_cos_embedding_keys": null,
779
+ "mean_std_embedding_keys": null,
780
+ "action_configs": null
781
+ },
782
+ "state": {
783
+ "delta_indices": [
784
+ 0
785
+ ],
786
+ "modality_keys": [
787
+ "left_wrist_eef",
788
+ "right_wrist_eef",
789
+ "left_gripper_pos",
790
+ "right_gripper_pos",
791
+ "left_joint_pos",
792
+ "right_joint_pos"
793
+ ],
794
+ "sin_cos_embedding_keys": null,
795
+ "mean_std_embedding_keys": null,
796
+ "action_configs": null
797
+ },
798
+ "action": {
799
+ "delta_indices": [
800
+ 0,
801
+ 1,
802
+ 2,
803
+ 3,
804
+ 4,
805
+ 5,
806
+ 6,
807
+ 7,
808
+ 8,
809
+ 9,
810
+ 10,
811
+ 11,
812
+ 12,
813
+ 13,
814
+ 14,
815
+ 15,
816
+ 16,
817
+ 17,
818
+ 18,
819
+ 19,
820
+ 20,
821
+ 21,
822
+ 22,
823
+ 23,
824
+ 24,
825
+ 25,
826
+ 26,
827
+ 27,
828
+ 28,
829
+ 29,
830
+ 30,
831
+ 31,
832
+ 32,
833
+ 33,
834
+ 34,
835
+ 35,
836
+ 36,
837
+ 37,
838
+ 38,
839
+ 39
840
+ ],
841
+ "modality_keys": [
842
+ "left_wrist_eef",
843
+ "right_wrist_eef",
844
+ "left_gripper_pos",
845
+ "right_gripper_pos",
846
+ "left_joint_pos",
847
+ "right_joint_pos"
848
+ ],
849
+ "sin_cos_embedding_keys": null,
850
+ "mean_std_embedding_keys": null,
851
+ "action_configs": [
852
+ {
853
+ "rep": "RELATIVE",
854
+ "type": "EEF",
855
+ "format": "XYZ_ROT6D",
856
+ "state_key": "left_wrist_eef"
857
+ },
858
+ {
859
+ "rep": "RELATIVE",
860
+ "type": "EEF",
861
+ "format": "XYZ_ROT6D",
862
+ "state_key": "right_wrist_eef"
863
+ },
864
+ {
865
+ "rep": "ABSOLUTE",
866
+ "type": "NON_EEF",
867
+ "format": "DEFAULT",
868
+ "state_key": "left_gripper_pos"
869
+ },
870
+ {
871
+ "rep": "ABSOLUTE",
872
+ "type": "NON_EEF",
873
+ "format": "DEFAULT",
874
+ "state_key": "right_gripper_pos"
875
+ },
876
+ {
877
+ "rep": "RELATIVE",
878
+ "type": "NON_EEF",
879
+ "format": "DEFAULT",
880
+ "state_key": "left_joint_pos"
881
+ },
882
+ {
883
+ "rep": "RELATIVE",
884
+ "type": "NON_EEF",
885
+ "format": "DEFAULT",
886
+ "state_key": "right_joint_pos"
887
+ }
888
+ ]
889
+ },
890
+ "language": {
891
+ "delta_indices": [
892
+ 0
893
+ ],
894
+ "modality_keys": [
895
+ "annotation.sub_task"
896
+ ],
897
+ "sin_cos_embedding_keys": null,
898
+ "mean_std_embedding_keys": null,
899
+ "action_configs": null
900
+ }
901
+ },
902
+ "oxe_droid_relative_eef_relative_joint": {
903
+ "video": {
904
+ "delta_indices": [
905
+ -15,
906
+ 0
907
+ ],
908
+ "modality_keys": [
909
+ "exterior_image_1_left",
910
+ "wrist_image_left"
911
+ ],
912
+ "sin_cos_embedding_keys": null,
913
+ "mean_std_embedding_keys": null,
914
+ "action_configs": null
915
+ },
916
+ "state": {
917
+ "delta_indices": [
918
+ 0
919
+ ],
920
+ "modality_keys": [
921
+ "eef_9d",
922
+ "gripper_position",
923
+ "joint_position"
924
+ ],
925
+ "sin_cos_embedding_keys": null,
926
+ "mean_std_embedding_keys": null,
927
+ "action_configs": null
928
+ },
929
+ "action": {
930
+ "delta_indices": [
931
+ 0,
932
+ 1,
933
+ 2,
934
+ 3,
935
+ 4,
936
+ 5,
937
+ 6,
938
+ 7,
939
+ 8,
940
+ 9,
941
+ 10,
942
+ 11,
943
+ 12,
944
+ 13,
945
+ 14,
946
+ 15,
947
+ 16,
948
+ 17,
949
+ 18,
950
+ 19,
951
+ 20,
952
+ 21,
953
+ 22,
954
+ 23,
955
+ 24,
956
+ 25,
957
+ 26,
958
+ 27,
959
+ 28,
960
+ 29,
961
+ 30,
962
+ 31,
963
+ 32,
964
+ 33,
965
+ 34,
966
+ 35,
967
+ 36,
968
+ 37,
969
+ 38,
970
+ 39
971
+ ],
972
+ "modality_keys": [
973
+ "eef_9d",
974
+ "gripper_position",
975
+ "joint_position"
976
+ ],
977
+ "sin_cos_embedding_keys": null,
978
+ "mean_std_embedding_keys": null,
979
+ "action_configs": [
980
+ {
981
+ "rep": "RELATIVE",
982
+ "type": "EEF",
983
+ "format": "XYZ_ROT6D",
984
+ "state_key": "eef_9d"
985
+ },
986
+ {
987
+ "rep": "ABSOLUTE",
988
+ "type": "NON_EEF",
989
+ "format": "DEFAULT",
990
+ "state_key": "gripper_position"
991
+ },
992
+ {
993
+ "rep": "RELATIVE",
994
+ "type": "NON_EEF",
995
+ "format": "DEFAULT",
996
+ "state_key": "joint_position"
997
+ }
998
+ ]
999
+ },
1000
+ "language": {
1001
+ "delta_indices": [
1002
+ 0
1003
+ ],
1004
+ "modality_keys": [
1005
+ "annotation.language.language_instruction"
1006
+ ],
1007
+ "sin_cos_embedding_keys": null,
1008
+ "mean_std_embedding_keys": null,
1009
+ "action_configs": null
1010
+ }
1011
+ },
1012
+ "new_embodiment": {
1013
+ "video": {
1014
+ "delta_indices": [
1015
+ 0
1016
+ ],
1017
+ "modality_keys": [
1018
+ "left_wrist",
1019
+ "right_wrist",
1020
+ "top"
1021
+ ],
1022
+ "sin_cos_embedding_keys": null,
1023
+ "mean_std_embedding_keys": null,
1024
+ "action_configs": null
1025
+ },
1026
+ "state": {
1027
+ "delta_indices": [
1028
+ 0
1029
+ ],
1030
+ "modality_keys": [
1031
+ "left_arm",
1032
+ "left_gripper",
1033
+ "right_arm",
1034
+ "right_gripper"
1035
+ ],
1036
+ "sin_cos_embedding_keys": null,
1037
+ "mean_std_embedding_keys": null,
1038
+ "action_configs": null
1039
+ },
1040
+ "action": {
1041
+ "delta_indices": [
1042
+ 0,
1043
+ 1,
1044
+ 2,
1045
+ 3,
1046
+ 4,
1047
+ 5,
1048
+ 6,
1049
+ 7,
1050
+ 8,
1051
+ 9,
1052
+ 10,
1053
+ 11,
1054
+ 12,
1055
+ 13,
1056
+ 14,
1057
+ 15
1058
+ ],
1059
+ "modality_keys": [
1060
+ "left_arm",
1061
+ "left_gripper",
1062
+ "right_arm",
1063
+ "right_gripper"
1064
+ ],
1065
+ "sin_cos_embedding_keys": null,
1066
+ "mean_std_embedding_keys": null,
1067
+ "action_configs": [
1068
+ {
1069
+ "rep": "RELATIVE",
1070
+ "type": "NON_EEF",
1071
+ "format": "DEFAULT",
1072
+ "state_key": null
1073
+ },
1074
+ {
1075
+ "rep": "ABSOLUTE",
1076
+ "type": "NON_EEF",
1077
+ "format": "DEFAULT",
1078
+ "state_key": null
1079
+ },
1080
+ {
1081
+ "rep": "RELATIVE",
1082
+ "type": "NON_EEF",
1083
+ "format": "DEFAULT",
1084
+ "state_key": null
1085
+ },
1086
+ {
1087
+ "rep": "ABSOLUTE",
1088
+ "type": "NON_EEF",
1089
+ "format": "DEFAULT",
1090
+ "state_key": null
1091
+ }
1092
+ ]
1093
+ },
1094
+ "language": {
1095
+ "delta_indices": [
1096
+ 0
1097
+ ],
1098
+ "modality_keys": [
1099
+ "annotation.human.task_description"
1100
+ ],
1101
+ "sin_cos_embedding_keys": null,
1102
+ "mean_std_embedding_keys": null,
1103
+ "action_configs": null
1104
+ }
1105
+ }
1106
+ },
1107
+ "image_crop_size": [
1108
+ 230,
1109
+ 230
1110
+ ],
1111
+ "image_target_size": [
1112
+ 256,
1113
+ 256
1114
+ ],
1115
+ "use_albumentations": true,
1116
+ "random_rotation_angle": 0,
1117
+ "color_jitter_params": {
1118
+ "brightness": 0.3,
1119
+ "contrast": 0.4,
1120
+ "saturation": 0.5,
1121
+ "hue": 0.08
1122
+ },
1123
+ "shortest_image_edge": 256,
1124
+ "crop_fraction": 0.95,
1125
+ "letter_box_transform": false,
1126
+ "model_name": "nvidia/Cosmos-Reason2-2B",
1127
+ "model_type": "qwen",
1128
+ "formalize_language": true,
1129
+ "max_state_dim": 132,
1130
+ "max_action_dim": 132,
1131
+ "max_action_horizon": 40,
1132
+ "use_percentiles": true,
1133
+ "use_mean_std": false,
1134
+ "clip_outliers": true,
1135
+ "apply_sincos_state_encoding": false,
1136
+ "use_relative_action": true,
1137
+ "exclude_state": false,
1138
+ "state_dropout_prob": 0.2
1139
+ }
1140
+ }
statistics.json ADDED
The diff for this file is too large to render. See raw diff
 
wandb_config.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"project": "groot_transfer_cube", "run_id": "groot_transfer_cube"}
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)