Add files using upload-large-folder tool

checkpoint-12000/config.json

@@ -0,0 +1,70 @@
+{
+  "action_horizon": 50,
+  "add_pos_embed": true,
+  "apply_sincos_state_encoding": true,
+  "architectures": [
+    "Gr00tN1d6"
+  ],
+  "attn_dropout": 0.2,
+  "attn_implementation": null,
+  "backbone_embedding_dim": 2048,
+  "backbone_model_type": "eagle",
+  "backbone_trainable_params_fp32": true,
+  "collator_overwrite_image_inputs": false,
+  "color_jitter_params": {
+    "brightness": 0.1,
+    "contrast": 0.1,
+    "hue": 0.1,
+    "saturation": 0.1
+  },
+  "crop_fraction": 0.95,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "eagle_collator": true,
+  "formalize_language": true,
+  "gemma_collator": false,
+  "hidden_size": 1024,
+  "image_crop_size": null,
+  "image_target_size": null,
+  "input_embedding_dim": 1536,
+  "load_bf16": true,
+  "max_action_dim": 128,
+  "max_num_embodiments": 32,
+  "max_seq_len": 1024,
+  "max_state_dim": 128,
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "model_type": "Gr00tN1d6",
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_inference_timesteps": 4,
+  "num_timestep_buckets": 1000,
+  "random_rotation_angle": null,
+  "reproject_vision": false,
+  "select_layer": 16,
+  "shortest_image_edge": 256,
+  "state_dropout_prob": 0.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "tune_diffusion_model": true,
+  "tune_llm": false,
+  "tune_projector": true,
+  "tune_top_llm_layers": 4,
+  "tune_visual": false,
+  "tune_vlln": true,
+  "use_albumentations_transforms": true,
+  "use_alternate_vl_dit": true,
+  "use_flash_attention": true,
+  "use_relative_action": true,
+  "use_vlln": true
+}

checkpoint-12000/experiment_cfg/conf.yaml

@@ -0,0 +1,203 @@
+load_config_path: null
+model:
+  model_type: Gr00tN1d6
+  model_dtype: bfloat16
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  backbone_model_type: eagle
+  model_revision: null
+  tune_top_llm_layers: 4
+  backbone_embedding_dim: 2048
+  tune_llm: false
+  tune_visual: false
+  select_layer: 16
+  reproject_vision: false
+  use_flash_attention: true
+  load_bf16: false
+  collator_overwrite_image_inputs: false
+  eagle_collator: true
+  backbone_trainable_params_fp32: true
+  image_crop_size: null
+  image_target_size: null
+  shortest_image_edge: 256
+  crop_fraction: 0.95
+  random_rotation_angle: null
+  color_jitter_params: null
+  use_albumentations_transforms: true
+  formalize_language: true
+  apply_sincos_state_encoding: false
+  use_relative_action: true
+  max_state_dim: 29
+  max_action_dim: 29
+  action_horizon: 16
+  hidden_size: 1024
+  input_embedding_dim: 1536
+  add_pos_embed: true
+  attn_dropout: 0.2
+  use_vlln: true
+  max_seq_len: 1024
+  use_alternate_vl_dit: true
+  attend_text_every_n_blocks: 2
+  diffusion_model_cfg:
+    positional_embeddings: null
+    num_layers: 32
+    num_attention_heads: 32
+    attention_head_dim: 48
+    norm_type: ada_norm
+    dropout: 0.2
+    final_dropout: true
+    output_dim: 1024
+    interleave_self_attention: true
+  num_inference_timesteps: 4
+  noise_beta_alpha: 1.5
+  noise_beta_beta: 1.0
+  noise_s: 0.999
+  num_timestep_buckets: 1000
+  tune_projector: true
+  tune_diffusion_model: true
+  tune_vlln: true
+  state_dropout_prob: 0.0
+  state_additive_noise_scale: 0.0
+  max_num_embodiments: 32
+data:
+  datasets:
+  - dataset_paths:
+    - ./datasets/reachy2_100
+    embodiment_tag: reachy2
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    reachy2:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front_cam
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - arm_joints
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      action:
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        modality_keys:
+        - arm_joints
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - rep: RELATIVE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: arm_joints
+        - rep: ABSOLUTE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: null
+      language:
+        delta_indices:
+        - 0
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+  download_cache: false
+  shard_size: 1024
+  episode_sampling_rate: 0.1
+  num_shards_per_epoch: 10000
+  override_pretraining_statistics: false
+  mode: single_turn
+  random_chop: 0.0
+  mock_dataset_mode: false
+  shuffle: true
+  seed: 42
+  multiprocessing_context: fork
+  allow_padding: false
+  subsample_ratio: 1.0
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  video_backend: decord
+training:
+  output_dir: /tmp/groot_output/reachy2
+  experiment_name: null
+  max_steps: 30000
+  global_batch_size: 64
+  batch_size: null
+  gradient_accumulation_steps: 1
+  learning_rate: 0.0001
+  lr_scheduler_type: cosine
+  weight_decay: 1.0e-05
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  max_grad_norm: 1.0
+  optim: adamw_torch
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  fp16: false
+  bf16: true
+  eval_bf16: true
+  logging_steps: 10
+  save_steps: 3000
+  save_total_limit: 5
+  save_vl_model: false
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  max_concurrent_uploads: 2
+  eval_strategy: 'no'
+  eval_steps: 500
+  eval_set_split_ratio: 0.1
+  eval_batch_size: 2
+  save_best_eval_metric_name: ''
+  save_best_eval_metric_greater_is_better: true
+  deepspeed_stage: 2
+  gradient_checkpointing: false
+  transformers_trust_remote_code: true
+  transformers_local_files_only: false
+  transformers_cache_dir: null
+  transformers_access_token: null
+  use_ddp: false
+  ddp_bucket_cap_mb: 100
+  num_gpus: 1
+  dataloader_num_workers: 8
+  remove_unused_columns: false
+  use_wandb: false
+  wandb_project: finetune-gr00t-n1d6
+  enable_profiling: false
+  max_retries: 3
+  assert_loss_less_than: null
+  add_rl_callback: false
+  enable_open_loop_eval: false
+  open_loop_eval_traj_ids:
+  - 0
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_plot_indices: null
+max_steps: 30000
+save_steps: 3000

checkpoint-12000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,1020 @@
+{
+  "reachy2": {
+    "state": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.316539764404297,
+          -0.3203781247138977,
+          2.417543411254883,
+          -109.398681640625,
+          -1.9847908020019531,
+          -3.484698534011841,
+          -0.3555510342121124
+        ],
+        "std": [
+          7.066186428070068,
+          3.507993459701538,
+          22.217893600463867,
+          14.851909637451172,
+          13.270956039428711,
+          11.239269256591797,
+          11.768560409545898
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723615646362305,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      }
+    },
+    "action": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.172462463378906,
+          -0.3404896557331085,
+          2.3248910903930664,
+          -109.23794555664062,
+          -1.812583327293396,
+          -3.5105223655700684,
+          -0.2696441113948822
+        ],
+        "std": [
+          7.037580966949463,
+          3.620150566101074,
+          22.205087661743164,
+          14.961851119995117,
+          13.164973258972168,
+          11.261707305908203,
+          11.734737396240234
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723625183105469,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.5416929721832275
+        ],
+        "max": [
+          0.9385063648223877
+        ],
+        "mean": [
+          0.8345981240272522
+        ],
+        "std": [
+          0.12303706258535363
+        ],
+        "q01": [
+          0.6467909216880798
+        ],
+        "q99": [
+          0.9350630044937134
+        ]
+      }
+    },
+    "relative_action": {
+      "arm_joints": {
+        "min": [
+          [
+            -20.440982818603516,
+            -11.665962219238281,
+            -14.02474594116211,
+            -20.817272186279297,
+            -10.09469223022461,
+            -6.401888370513916,
+            -8.767365455627441
+          ],
+          [
+            -20.449968338012695,
+            -11.677082061767578,
+            -25.868675231933594,
+            -30.582874298095703,
+            -17.461694717407227,
+            -11.940333366394043,
+            -16.13103485107422
+          ],
+          [
+            -19.697128295898438,
+            -10.979728698730469,
+            -35.53736114501953,
+            -32.2390022277832,
+            -22.624008178710938,
+            -17.247066497802734,
+            -22.225337982177734
+          ],
+          [
+            -20.74054527282715,
+            -10.334735870361328,
+            -44.72303009033203,
+            -33.42095947265625,
+            -25.056644439697266,
+            -22.322742462158203,
+            -27.965492248535156
+          ],
+          [
+            -24.075529098510742,
+            -10.338696479797363,
+            -50.480987548828125,
+            -34.134925842285156,
+            -26.33944320678711,
+            -25.707611083984375,
+            -31.17104721069336
+          ],
+          [
+            -26.567087173461914,
+            -10.34083366394043,
+            -54.44721603393555,
+            -34.8769645690918,
+            -26.375856399536133,
+            -29.134244918823242,
+            -33.48938751220703
+          ],
+          [
+            -27.659536361694336,
+            -10.340897560119629,
+            -57.30182647705078,
+            -36.380828857421875,
+            -26.375410079956055,
+            -32.329833984375,
+            -35.47602081298828
+          ],
+          [
+            -27.924762725830078,
+            -11.4292573928833,
+            -59.24040603637695,
+            -36.7144775390625,
+            -26.372446060180664,
+            -33.91874313354492,
+            -35.908573150634766
+          ],
+          [
+            -27.95497703552246,
+            -11.528078079223633,
+            -59.352848052978516,
+            -36.841575622558594,
+            -26.24493980407715,
+            -35.34848403930664,
+            -36.83350372314453
+          ],
+          [
+            -28.4979248046875,
+            -11.443194389343262,
+            -59.4117431640625,
+            -36.853206634521484,
+            -26.16107940673828,
+            -36.28956604003906,
+            -36.86201858520508
+          ],
+          [
+            -28.523822784423828,
+            -11.443158149719238,
+            -59.41852951049805,
+            -36.85321807861328,
+            -26.160951614379883,
+            -36.43997573852539,
+            -36.88083267211914
+          ],
+          [
+            -28.522306442260742,
+            -11.443062782287598,
+            -59.424468994140625,
+            -37.14503860473633,
+            -26.16082763671875,
+            -36.443668365478516,
+            -36.880828857421875
+          ],
+          [
+            -28.468467712402344,
+            -11.44290828704834,
+            -59.42174530029297,
+            -37.880104064941406,
+            -26.160717010498047,
+            -36.44268035888672,
+            -36.880828857421875
+          ],
+          [
+            -28.41590118408203,
+            -11.442895889282227,
+            -59.4072265625,
+            -38.47637939453125,
+            -26.160614013671875,
+            -36.4410400390625,
+            -36.880828857421875
+          ],
+          [
+            -28.362062454223633,
+            -11.442879676818848,
+            -59.38927459716797,
+            -41.26643371582031,
+            -26.160587310791016,
+            -36.43976974487305,
+            -36.88082504272461
+          ],
+          [
+            -28.244752883911133,
+            -11.442869186401367,
+            -59.369014739990234,
+            -44.968875885009766,
+            -26.160572052001953,
+            -36.433982849121094,
+            -36.880794525146484
+          ]
+        ],
+        "max": [
+          [
+            5.384329319000244,
+            11.368682861328125,
+            12.580580711364746,
+            21.664548873901367,
+            22.868919372558594,
+            13.739058494567871,
+            11.726442337036133
+          ],
+          [
+            9.6729097366333,
+            11.36864185333252,
+            23.997724533081055,
+            21.659494400024414,
+            22.868900299072266,
+            23.967060089111328,
+            14.994425773620605
+          ],
+          [
+            13.331961631774902,
+            11.190152168273926,
+            33.41366958618164,
+            21.635488510131836,
+            25.30813217163086,
+            26.21414566040039,
+            20.79837417602539
+          ],
+          [
+            16.202699661254883,
+            10.40443229675293,
+            40.09684753417969,
+            23.701637268066406,
+            28.603330612182617,
+            27.53230857849121,
+            25.446346282958984
+          ],
+          [
+            19.05048370361328,
+            10.40436840057373,
+            46.37767791748047,
+            28.047138214111328,
+            31.07068634033203,
+            28.03820037841797,
+            29.473896026611328
+          ],
+          [
+            20.245254516601562,
+            10.402230262756348,
+            51.10308837890625,
+            31.863624572753906,
+            31.10450553894043,
+            29.278892517089844,
+            32.1208381652832
+          ],
+          [
+            21.428510665893555,
+            10.398269653320312,
+            53.42000961303711,
+            34.682308197021484,
+            31.103679656982422,
+            30.250988006591797,
+            33.45231246948242
+          ],
+          [
+            21.98176383972168,
+            10.395852088928223,
+            55.17178726196289,
+            36.64023208618164,
+            31.101837158203125,
+            30.935718536376953,
+            34.229305267333984
+          ],
+          [
+            22.397876739501953,
+            11.4427490234375,
+            55.23532485961914,
+            37.58878707885742,
+            31.096826553344727,
+            31.260026931762695,
+            34.44834899902344
+          ],
+          [
+            22.6441650390625,
+            11.780808448791504,
+            55.34986114501953,
+            38.26850128173828,
+            31.096656799316406,
+            31.333267211914062,
+            34.469146728515625
+          ],
+          [
+            22.692224502563477,
+            11.812055587768555,
+            55.36343765258789,
+            38.49715805053711,
+            31.61040687561035,
+            31.37093162536621,
+            34.469146728515625
+          ],
+          [
+            22.72762107849121,
+            11.678536415100098,
+            55.36940383911133,
+            38.59236526489258,
+            35.083839416503906,
+            31.3709659576416,
+            34.46915054321289
+          ],
+          [
+            24.14029312133789,
+            11.678519248962402,
+            55.369537353515625,
+            38.63944625854492,
+            38.55842208862305,
+            31.37123680114746,
+            34.46915054321289
+          ],
+          [
+            24.973907470703125,
+            11.678510665893555,
+            55.36623001098633,
+            38.656288146972656,
+            39.96430587768555,
+            31.371545791625977,
+            34.46915054321289
+          ],
+          [
+            25.06302261352539,
+            11.678531646728516,
+            55.3525390625,
+            38.694881439208984,
+            40.0225944519043,
+            31.37181854248047,
+            34.468910217285156
+          ],
+          [
+            25.270801544189453,
+            11.678507804870605,
+            55.302616119384766,
+            38.694969177246094,
+            40.022586822509766,
+            31.37185287475586,
+            34.385284423828125
+          ]
+        ],
+        "mean": [
+          [
+            -0.09851349890232086,
+            -1.3889506362829707e-06,
+            -0.021199805662035942,
+            -0.0799250528216362,
+            0.1505303829908371,
+            0.11832442134618759,
+            0.13495467603206635
+          ],
+          [
+            -0.17057012021541595,
+            0.004252429585903883,
+            -0.008995133452117443,
+            -0.22457502782344818,
+            0.261105477809906,
+            0.27207958698272705,
+            0.24137930572032928
+          ],
+          [
+            -0.10028649866580963,
+            0.03141691908240318,
+            0.09048548340797424,
+            -0.5493719577789307,
+            0.19902515411376953,
+            0.47134673595428467,
+            0.2771521210670471
+          ],
+          [
+            -0.05444306880235672,
+            0.05378008633852005,
+            0.1796490103006363,
+            -0.879206120967865,
+            0.1520315706729889,
+            0.6866292357444763,
+            0.3248023986816406
+          ],
+          [
+            -0.013876532204449177,
+            0.07615811377763748,
+            0.2686113715171814,
+            -1.2192035913467407,
+            0.10656605660915375,
+            0.9128072261810303,
+            0.3744547367095947
+          ],
+          [
+            0.025331225246191025,
+            0.09861738234758377,
+            0.3578471839427948,
+            -1.5637919902801514,
+            0.06106000766158104,
+            1.1433591842651367,
+            0.4248007833957672
+          ],
+          [
+            0.0637386217713356,
+            0.12045339494943619,
+            0.44647839665412903,
+            -1.908605933189392,
+            0.016334885731339455,
+            1.3740676641464233,
+            0.4745999872684479
+          ],
+          [
+            0.10053692013025284,
+            0.13872262835502625,
+            0.5284987092018127,
+            -2.2462470531463623,
+            -0.02209978550672531,
+            1.5981824398040771,
+            0.518027663230896
+          ],
+          [
+            0.1247839704155922,
+            0.14948606491088867,
+            0.583483874797821,
+            -2.525705337524414,
+            -0.04458107799291611,
+            1.7731598615646362,
+            0.5445129871368408
+          ],
+          [
+            0.12165514379739761,
+            0.14827895164489746,
+            0.602064847946167,
+            -2.593496561050415,
+            -0.047096796333789825,
+            1.797343134880066,
+            0.5500790476799011
+          ],
+          [
+            0.1124163493514061,
+            0.14204099774360657,
+            0.6022222638130188,
+            -2.55938458442688,
+            -0.04651280865073204,
+            1.7794476747512817,
+            0.5512804985046387
+          ],
+          [
+            0.09131923317909241,
+            0.12857742607593536,
+            0.6059610843658447,
+            -2.4963152408599854,
+            -0.04961543530225754,
+            1.740905523300171,
+            0.5537145137786865
+          ],
+          [
+            0.054126936942338943,
+            0.10883964598178864,
+            0.6135755181312561,
+            -2.381096601486206,
+            -0.056390050798654556,
+            1.6682276725769043,
+            0.5572347044944763
+          ],
+          [
+            -0.001200970378704369,
+            0.08611582219600677,
+            0.6268700361251831,
+            -2.207465648651123,
+            -0.06850837916135788,
+            1.5567927360534668,
+            0.5623834133148193
+          ],
+          [
+            -0.07379738241434097,
+            0.06273127347230911,
+            0.6453807950019836,
+            -1.9786254167556763,
+            -0.08601187914609909,
+            1.4089199304580688,
+            0.5690953135490417
+          ],
+          [
+            -0.15925677120685577,
+            0.039619192481040955,
+            0.6695932149887085,
+            -1.7069525718688965,
+            -0.10912654548883438,
+            1.2322944402694702,
+            0.5775030851364136
+          ]
+        ],
+        "std": [
+          [
+            1.9917564392089844,
+            1.1211100816726685,
+            2.417202949523926,
+            2.94158935546875,
+            2.39365291595459,
+            1.874914526939392,
+            1.6537472009658813
+          ],
+          [
+            3.056582450866699,
+            1.6457527875900269,
+            4.458590507507324,
+            5.178750514984131,
+            3.765504837036133,
+            3.6351206302642822,
+            2.9263687133789062
+          ],
+          [
+            3.596614122390747,
+            1.8142179250717163,
+            6.294006824493408,
+            7.028163433074951,
+            4.514294147491455,
+            5.296512126922607,
+            4.0286865234375
+          ],
+          [
+            4.277676582336426,
+            2.040328025817871,
+            8.043697357177734,
+            8.868635177612305,
+            5.327646255493164,
+            6.84896993637085,
+            5.116058349609375
+          ],
+          [
+            4.939493179321289,
+            2.248694896697998,
+            9.646915435791016,
+            10.583718299865723,
+            6.0862040519714355,
+            8.271647453308105,
+            6.122095584869385
+          ],
+          [
+            5.552292823791504,
+            2.439715623855591,
+            11.097370147705078,
+            12.150360107421875,
+            6.7813005447387695,
+            9.56061840057373,
+            7.034397125244141
+          ],
+          [
+            6.109565258026123,
+            2.6163902282714844,
+            12.404134750366211,
+            13.574675559997559,
+            7.419105052947998,
+            10.72553539276123,
+            7.855382919311523
+          ],
+          [
+            6.608177661895752,
+            2.7779171466827393,
+            13.585582733154297,
+            14.875726699829102,
+            8.011663436889648,
+            11.787652015686035,
+            8.598040580749512
+          ],
+          [
+            7.046725273132324,
+            2.93743634223938,
+            14.688329696655273,
+            16.078651428222656,
+            8.60019302368164,
+            12.769647598266602,
+            9.277371406555176
+          ],
+          [
+            7.43002986907959,
+            3.1344540119171143,
+            15.875958442687988,
+            17.27137565612793,
+            9.292807579040527,
+            13.705930709838867,
+            9.92088794708252
+          ],
+          [
+            7.734869003295898,
+            3.2678585052490234,
+            17.0732421875,
+            18.39610481262207,
+            9.9412841796875,
+            14.559659957885742,
+            10.473697662353516
+          ],
+          [
+            7.963422775268555,
+            3.390157699584961,
+            18.116336822509766,
+            19.400726318359375,
+            10.5236177444458,
+            15.339380264282227,
+            10.94632339477539
+          ],
+          [
+            8.121000289916992,
+            3.5208446979522705,
+            19.046070098876953,
+            20.327335357666016,
+            11.053685188293457,
+            16.05907440185547,
+            11.358160972595215
+          ],
+          [
+            8.218527793884277,
+            3.651193857192993,
+            19.873519897460938,
+            21.182926177978516,
+            11.531027793884277,
+            16.716957092285156,
+            11.718016624450684
+          ],
+          [
+            8.26684284210205,
+            3.7774546146392822,
+            20.612743377685547,
+            21.96902847290039,
+            11.957940101623535,
+            17.311002731323242,
+            12.034592628479004
+          ],
+          [
+            8.276911735534668,
+            3.8994669914245605,
+            21.281835556030273,
+            22.691150665283203,
+            12.342350959777832,
+            17.84532928466797,
+            12.317776679992676
+          ]
+        ],
+        "q01": [
+          [
+            -9.697267532348633,
+            -3.6728341579437256,
+            -10.083383560180664,
+            -7.500851631164551,
+            -5.057114124298096,
+            -5.191580295562744,
+            -5.047628879547119
+          ],
+          [
+            -17.174776077270508,
+            -7.670596122741699,
+            -15.547335624694824,
+            -14.150874137878418,
+            -9.842418670654297,
+            -9.951410293579102,
+            -9.844037055969238
+          ],
+          [
+            -16.49785804748535,
+            -9.060977935791016,
+            -23.215665817260742,
+            -19.967872619628906,
+            -13.348301887512207,
+            -14.247873306274414,
+            -14.610755920410156
+          ],
+          [
+            -16.7535400390625,
+            -9.395726203918457,
+            -29.43296241760254,
+            -25.006587982177734,
+            -16.37109375,
+            -18.130178451538086,
+            -18.5999755859375
+          ],
+          [
+            -17.191621780395508,
+            -9.43246078491211,
+            -36.27531051635742,
+            -28.893749237060547,
+            -18.89956283569336,
+            -21.859458923339844,
+            -22.91756248474121
+          ],
+          [
+            -19.174758911132812,
+            -9.506808280944824,
+            -40.94317626953125,
+            -31.8222713470459,
+            -20.72899055480957,
+            -24.93976402282715,
+            -25.59789276123047
+          ],
+          [
+            -20.843280792236328,
+            -9.511780738830566,
+            -43.895530700683594,
+            -33.7496337890625,
+            -21.83380699157715,
+            -27.3231201171875,
+            -27.804710388183594
+          ],
+          [
+            -22.112207412719727,
+            -9.513916015625,
+            -47.63551330566406,
+            -34.7276611328125,
+            -22.691164016723633,
+            -29.399824142456055,
+            -30.01004981994629
+          ],
+          [
+            -23.385971069335938,
+            -9.525580406188965,
+            -49.76005172729492,
+            -35.31100082397461,
+            -23.10930824279785,
+            -30.991064071655273,
+            -31.236299514770508
+          ],
+          [
+            -23.984119415283203,
+            -9.732311248779297,
+            -51.624576568603516,
+            -35.76605224609375,
+            -23.42118263244629,
+            -31.96640396118164,
+            -32.20992660522461
+          ],
+          [
+            -24.42165184020996,
+            -10.062808990478516,
+            -52.54084777832031,
+            -36.21389389038086,
+            -23.58216094970703,
+            -32.878292083740234,
+            -33.002262115478516
+          ],
+          [
+            -24.806774139404297,
+            -10.06833267211914,
+            -53.43137741088867,
+            -36.29438781738281,
+            -23.675121307373047,
+            -33.4473876953125,
+            -33.2683219909668
+          ],
+          [
+            -24.799230575561523,
+            -10.070316314697266,
+            -54.24421691894531,
+            -36.37746810913086,
+            -23.685352325439453,
+            -33.9517822265625,
+            -33.42109298706055
+          ],
+          [
+            -24.732126235961914,
+            -10.07029914855957,
+            -54.64339828491211,
+            -36.4773063659668,
+            -23.682329177856445,
+            -34.212223052978516,
+            -33.55178451538086
+          ],
+          [
+            -24.3509578704834,
+            -10.070287704467773,
+            -54.56159210205078,
+            -36.60988998413086,
+            -23.682201385498047,
+            -34.409423828125,
+            -33.26326370239258
+          ],
+          [
+            -23.759092330932617,
+            -10.070272445678711,
+            -54.44050216674805,
+            -37.09710693359375,
+            -23.682085037231445,
+            -34.68149948120117,
+            -32.9384765625
+          ]
+        ],
+        "q99": [
+          [
+            3.203047752380371,
+            3.7505993843078613,
+            8.72464370727539,
+            9.948126792907715,
+            11.195253372192383,
+            5.477096080780029,
+            7.006125450134277
+          ],
+          [
+            5.7390570640563965,
+            6.635564804077148,
+            16.971813201904297,
+            18.642690658569336,
+            19.90561294555664,
+            10.401065826416016,
+            11.342881202697754
+          ],
+          [
+            8.532023429870605,
+            8.251216888427734,
+            24.75363540649414,
+            18.472192764282227,
+            19.774227142333984,
+            14.655866622924805,
+            15.70905590057373
+          ],
+          [
+            10.869905471801758,
+            9.230382919311523,
+            31.436893463134766,
+            19.715234756469727,
+            19.81361961364746,
+            18.759492874145508,
+            20.138126373291016
+          ],
+          [
+            12.716119766235352,
+            9.484061241149902,
+            37.04032516479492,
+            23.830785751342773,
+            19.81861114501953,
+            21.909305572509766,
+            23.74039649963379
+          ],
+          [
+            14.470975875854492,
+            9.510160446166992,
+            41.770057678222656,
+            27.173765182495117,
+            20.0618839263916,
+            24.20587921142578,
+            26.801916122436523
+          ],
+          [
+            16.024580001831055,
+            9.510558128356934,
+            45.34690475463867,
+            30.040668487548828,
+            21.872987747192383,
+            25.943317413330078,
+            28.912029266357422
+          ],
+          [
+            16.895124435424805,
+            9.513190269470215,
+            47.93039321899414,
+            32.24497985839844,
+            22.715166091918945,
+            27.396934509277344,
+            30.489883422851562
+          ],
+          [
+            17.761186599731445,
+            9.570911407470703,
+            50.055416107177734,
+            34.065670013427734,
+            23.788578033447266,
+            28.185653686523438,
+            31.546728134155273
+          ],
+          [
+            18.391948699951172,
+            9.712384223937988,
+            51.48924255371094,
+            35.40724182128906,
+            24.940738677978516,
+            28.75865364074707,
+            32.44330978393555
+          ],
+          [
+            19.106931686401367,
+            9.978495597839355,
+            52.564918518066406,
+            36.375064849853516,
+            25.661256790161133,
+            29.218664169311523,
+            33.03133010864258
+          ],
+          [
+            19.62432861328125,
+            10.027887344360352,
+            53.07897186279297,
+            36.976646423339844,
+            26.609840393066406,
+            29.49842071533203,
+            33.33045959472656
+          ],
+          [
+            19.768310546875,
+            10.027881622314453,
+            53.27336883544922,
+            37.44525909423828,
+            28.005691528320312,
+            29.92152214050293,
+            33.3731803894043
+          ],
+          [
+            19.806724548339844,
+            10.027848243713379,
+            53.31312942504883,
+            37.71003341674805,
+            29.430599212646484,
+            29.987842559814453,
+            33.3889274597168
+          ],
+          [
+            20.140544891357422,
+            10.027847290039062,
+            53.163578033447266,
+            37.966617584228516,
+            30.082942962646484,
+            30.24797248840332,
+            33.202571868896484
+          ],
+          [
+            20.164806365966797,
+            10.02784538269043,
+            52.79570770263672,
+            38.0666389465332,
+            30.232667922973633,
+            30.37464714050293,
+            32.94834899902344
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-12000/experiment_cfg/final_model_config.json

@@ -0,0 +1,53 @@
+{
+  "model_type": "Gr00tN1d6",
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "backbone_model_type": "eagle",
+  "model_revision": null,
+  "tune_top_llm_layers": 4,
+  "backbone_embedding_dim": 2048,
+  "tune_llm": false,
+  "tune_visual": false,
+  "select_layer": 16,
+  "reproject_vision": false,
+  "use_flash_attention": true,
+  "load_bf16": true,
+  "collator_overwrite_image_inputs": false,
+  "eagle_collator": true,
+  "backbone_trainable_params_fp32": true,
+  "apply_sincos_state_encoding": true,
+  "use_relative_action": true,
+  "max_state_dim": 128,
+  "max_action_dim": 128,
+  "action_horizon": 50,
+  "hidden_size": 1024,
+  "input_embedding_dim": 1536,
+  "add_pos_embed": true,
+  "attn_dropout": 0.2,
+  "use_vlln": true,
+  "max_seq_len": 1024,
+  "use_alternate_vl_dit": true,
+  "attend_text_every_n_blocks": 2,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "num_inference_timesteps": 4,
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_timestep_buckets": 1000,
+  "tune_projector": true,
+  "tune_diffusion_model": true,
+  "tune_vlln": true,
+  "state_dropout_prob": 0.0,
+  "state_additive_noise_scale": 0.0,
+  "max_num_embodiments": 32
+}

checkpoint-12000/processor_config.json

@@ -0,0 +1,453 @@
+{
+  "processor_class": "Gr00tN1d6Processor",
+  "processor_kwargs": {
+    "modality_configs": {
+      "behavior_r1_pro": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "observation.images.rgb.head_256_256",
+            "observation.images.rgb.left_wrist_256_256",
+            "observation.images.rgb.right_wrist_256_256"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "robot_pos",
+            "robot_ori_cos",
+            "robot_ori_sin",
+            "robot_2d_ori",
+            "robot_2d_ori_cos",
+            "robot_2d_ori_sin",
+            "robot_lin_vel",
+            "robot_ang_vel",
+            "arm_left_qpos",
+            "arm_left_qpos_sin",
+            "arm_left_qpos_cos",
+            "eef_left_pos",
+            "eef_left_quat",
+            "gripper_left_qpos",
+            "arm_right_qpos",
+            "arm_right_qpos_sin",
+            "arm_right_qpos_cos",
+            "eef_right_pos",
+            "eef_right_quat",
+            "gripper_right_qpos",
+            "trunk_qpos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31
+          ],
+          "modality_keys": [
+            "base",
+            "torso",
+            "left_arm",
+            "left_gripper",
+            "right_arm",
+            "right_gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "trunk_qpos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_left_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_right_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "gr1": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "ego_view_bg_crop_pad_res256_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "robocasa_panda_omron": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "res256_image_side_0",
+            "res256_image_side_1",
+            "res256_image_wrist_0"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "end_effector_position_relative",
+            "end_effector_rotation_relative",
+            "gripper_qpos",
+            "base_position",
+            "base_rotation"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "end_effector_position",
+            "end_effector_rotation",
+            "gripper_close",
+            "base_motion",
+            "control_mode"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.action.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "reachy2": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front_cam"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "arm_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "arm_joints",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      }
+    },
+    "image_crop_size": null,
+    "image_target_size": null,
+    "use_albumentations": true,
+    "random_rotation_angle": null,
+    "color_jitter_params": {
+      "brightness": 0.3,
+      "contrast": 0.4,
+      "saturation": 0.5,
+      "hue": 0.08
+    },
+    "shortest_image_edge": 256,
+    "crop_fraction": 0.95,
+    "model_name": "nvidia/Eagle-Block2A-2B-v2",
+    "model_type": "eagle",
+    "formalize_language": true,
+    "max_state_dim": 128,
+    "max_action_dim": 128,
+    "max_action_horizon": 50,
+    "use_percentiles": false,
+    "clip_outliers": true,
+    "apply_sincos_state_encoding": true,
+    "use_relative_action": true
+  }
+}

checkpoint-12000/rng_state.pth

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

checkpoint-12000/wandb_config.json

@@ -0,0 +1 @@
+{"project": "finetune-gr00t-n1d6", "run_id": "reachy2"}

checkpoint-3000/config.json

@@ -0,0 +1,70 @@
+{
+  "action_horizon": 50,
+  "add_pos_embed": true,
+  "apply_sincos_state_encoding": true,
+  "architectures": [
+    "Gr00tN1d6"
+  ],
+  "attn_dropout": 0.2,
+  "attn_implementation": null,
+  "backbone_embedding_dim": 2048,
+  "backbone_model_type": "eagle",
+  "backbone_trainable_params_fp32": true,
+  "collator_overwrite_image_inputs": false,
+  "color_jitter_params": {
+    "brightness": 0.1,
+    "contrast": 0.1,
+    "hue": 0.1,
+    "saturation": 0.1
+  },
+  "crop_fraction": 0.95,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "eagle_collator": true,
+  "formalize_language": true,
+  "gemma_collator": false,
+  "hidden_size": 1024,
+  "image_crop_size": null,
+  "image_target_size": null,
+  "input_embedding_dim": 1536,
+  "load_bf16": true,
+  "max_action_dim": 128,
+  "max_num_embodiments": 32,
+  "max_seq_len": 1024,
+  "max_state_dim": 128,
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "model_type": "Gr00tN1d6",
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_inference_timesteps": 4,
+  "num_timestep_buckets": 1000,
+  "random_rotation_angle": null,
+  "reproject_vision": false,
+  "select_layer": 16,
+  "shortest_image_edge": 256,
+  "state_dropout_prob": 0.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "tune_diffusion_model": true,
+  "tune_llm": false,
+  "tune_projector": true,
+  "tune_top_llm_layers": 4,
+  "tune_visual": false,
+  "tune_vlln": true,
+  "use_albumentations_transforms": true,
+  "use_alternate_vl_dit": true,
+  "use_flash_attention": true,
+  "use_relative_action": true,
+  "use_vlln": true
+}

checkpoint-3000/embodiment_id.json

@@ -0,0 +1,12 @@
+{
+  "robocasa_panda_omron": 13,
+  "gr1": 20,
+  "behavior_r1_pro": 24,
+  "unitree_g1": 8,
+  "oxe_google": 0,
+  "oxe_widowx": 1,
+  "libero_panda": 2,
+  "oxe_droid": 16,
+  "new_embodiment": 10,
+  "reachy2": 11
+}

checkpoint-3000/experiment_cfg/conf.yaml

@@ -0,0 +1,203 @@
+load_config_path: null
+model:
+  model_type: Gr00tN1d6
+  model_dtype: bfloat16
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  backbone_model_type: eagle
+  model_revision: null
+  tune_top_llm_layers: 4
+  backbone_embedding_dim: 2048
+  tune_llm: false
+  tune_visual: false
+  select_layer: 16
+  reproject_vision: false
+  use_flash_attention: true
+  load_bf16: false
+  collator_overwrite_image_inputs: false
+  eagle_collator: true
+  backbone_trainable_params_fp32: true
+  image_crop_size: null
+  image_target_size: null
+  shortest_image_edge: 256
+  crop_fraction: 0.95
+  random_rotation_angle: null
+  color_jitter_params: null
+  use_albumentations_transforms: true
+  formalize_language: true
+  apply_sincos_state_encoding: false
+  use_relative_action: true
+  max_state_dim: 29
+  max_action_dim: 29
+  action_horizon: 16
+  hidden_size: 1024
+  input_embedding_dim: 1536
+  add_pos_embed: true
+  attn_dropout: 0.2
+  use_vlln: true
+  max_seq_len: 1024
+  use_alternate_vl_dit: true
+  attend_text_every_n_blocks: 2
+  diffusion_model_cfg:
+    positional_embeddings: null
+    num_layers: 32
+    num_attention_heads: 32
+    attention_head_dim: 48
+    norm_type: ada_norm
+    dropout: 0.2
+    final_dropout: true
+    output_dim: 1024
+    interleave_self_attention: true
+  num_inference_timesteps: 4
+  noise_beta_alpha: 1.5
+  noise_beta_beta: 1.0
+  noise_s: 0.999
+  num_timestep_buckets: 1000
+  tune_projector: true
+  tune_diffusion_model: true
+  tune_vlln: true
+  state_dropout_prob: 0.0
+  state_additive_noise_scale: 0.0
+  max_num_embodiments: 32
+data:
+  datasets:
+  - dataset_paths:
+    - ./datasets/reachy2_100
+    embodiment_tag: reachy2
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    reachy2:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front_cam
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - arm_joints
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      action:
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        modality_keys:
+        - arm_joints
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - rep: RELATIVE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: arm_joints
+        - rep: ABSOLUTE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: null
+      language:
+        delta_indices:
+        - 0
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+  download_cache: false
+  shard_size: 1024
+  episode_sampling_rate: 0.1
+  num_shards_per_epoch: 10000
+  override_pretraining_statistics: false
+  mode: single_turn
+  random_chop: 0.0
+  mock_dataset_mode: false
+  shuffle: true
+  seed: 42
+  multiprocessing_context: fork
+  allow_padding: false
+  subsample_ratio: 1.0
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  video_backend: decord
+training:
+  output_dir: /tmp/groot_output/reachy2
+  experiment_name: null
+  max_steps: 30000
+  global_batch_size: 64
+  batch_size: null
+  gradient_accumulation_steps: 1
+  learning_rate: 0.0001
+  lr_scheduler_type: cosine
+  weight_decay: 1.0e-05
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  max_grad_norm: 1.0
+  optim: adamw_torch
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  fp16: false
+  bf16: true
+  eval_bf16: true
+  logging_steps: 10
+  save_steps: 3000
+  save_total_limit: 5
+  save_vl_model: false
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  max_concurrent_uploads: 2
+  eval_strategy: 'no'
+  eval_steps: 500
+  eval_set_split_ratio: 0.1
+  eval_batch_size: 2
+  save_best_eval_metric_name: ''
+  save_best_eval_metric_greater_is_better: true
+  deepspeed_stage: 2
+  gradient_checkpointing: false
+  transformers_trust_remote_code: true
+  transformers_local_files_only: false
+  transformers_cache_dir: null
+  transformers_access_token: null
+  use_ddp: false
+  ddp_bucket_cap_mb: 100
+  num_gpus: 1
+  dataloader_num_workers: 8
+  remove_unused_columns: false
+  use_wandb: false
+  wandb_project: finetune-gr00t-n1d6
+  enable_profiling: false
+  max_retries: 3
+  assert_loss_less_than: null
+  add_rl_callback: false
+  enable_open_loop_eval: false
+  open_loop_eval_traj_ids:
+  - 0
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_plot_indices: null
+max_steps: 30000
+save_steps: 3000

checkpoint-3000/experiment_cfg/config.yaml

@@ -0,0 +1,237 @@
+!!python/object:gr00t.configs.base_config.Config
+data: !!python/object:gr00t.configs.data.data_config.DataConfig
+  allow_padding: false
+  datasets:
+  - !!python/object:gr00t.configs.data.data_config.SingleDatasetConfig
+    dataset_paths:
+    - ./datasets/reachy2_100
+    dataset_type: physical_embodiment
+    embodiment_tag: reachy2
+    mix_ratio: 1.0
+    val_dataset_path: null
+  download_cache: false
+  episode_sampling_rate: 0.1
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  mock_dataset_mode: false
+  modality_configs:
+    reachy2:
+      action: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs:
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: &id001 !!python/object/apply:gr00t.data.types.ActionFormat
+          - default
+          rep: !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - relative
+          state_key: arm_joints
+          type: &id002 !!python/object/apply:gr00t.data.types.ActionType
+          - non_eef
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: *id001
+          rep: !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - absolute
+          state_key: null
+          type: *id002
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        mean_std_embedding_keys: null
+        modality_keys:
+        - arm_joints
+        - gripper
+        sin_cos_embedding_keys: null
+      language: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+      state: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - arm_joints
+        sin_cos_embedding_keys: null
+      video: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - front_cam
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 10000
+  override_pretraining_statistics: false
+  random_chop: 0.0
+  seed: 42
+  shard_size: 1024
+  shuffle: true
+  subsample_ratio: 1.0
+  video_backend: decord
+load_config_path: null
+model: !!python/object:gr00t.configs.model.gr00t_n1d6.Gr00tN1d6Config
+  _attn_implementation_autoset: false
+  _attn_implementation_internal: null
+  _commit_hash: null
+  _name_or_path: ''
+  add_cross_attention: false
+  architectures: null
+  backbone_model_type: eagle
+  backbone_trainable_params_fp32: true
+  bad_words_ids: null
+  begin_suppress_tokens: null
+  bos_token_id: null
+  chunk_size_feed_forward: 0
+  color_jitter_params: null
+  cross_attention_hidden_size: null
+  decoder_start_token_id: null
+  diffusion_model_cfg:
+    attention_head_dim: 48
+    dropout: 0.2
+    final_dropout: true
+    interleave_self_attention: true
+    norm_type: ada_norm
+    num_attention_heads: 32
+    num_layers: 32
+    output_dim: 1024
+    positional_embeddings: null
+  diversity_penalty: 0.0
+  do_sample: false
+  eagle_collator: true
+  early_stopping: false
+  encoder_no_repeat_ngram_size: 0
+  eos_token_id: null
+  exponential_decay_length_penalty: null
+  finetuning_task: null
+  forced_bos_token_id: null
+  forced_eos_token_id: null
+  id2label:
+    0: LABEL_0
+    1: LABEL_1
+  is_decoder: false
+  is_encoder_decoder: false
+  label2id:
+    LABEL_0: 0
+    LABEL_1: 1
+  length_penalty: 1.0
+  load_bf16: false
+  max_length: 20
+  min_length: 0
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  no_repeat_ngram_size: 0
+  num_beam_groups: 1
+  num_beams: 1
+  num_return_sequences: 1
+  output_attentions: false
+  output_hidden_states: false
+  output_scores: false
+  pad_token_id: null
+  prefix: null
+  problem_type: null
+  pruned_heads: {}
+  random_rotation_angle: null
+  remove_invalid_values: false
+  repetition_penalty: 1.0
+  reproject_vision: false
+  return_dict: true
+  return_dict_in_generate: false
+  sep_token_id: null
+  state_dropout_prob: 0.0
+  suppress_tokens: null
+  task_specific_params: null
+  temperature: 1.0
+  tf_legacy_loss: false
+  tie_encoder_decoder: false
+  tie_word_embeddings: true
+  tokenizer_class: null
+  top_k: 50
+  top_p: 1.0
+  torch_dtype: null
+  torchscript: false
+  transformers_version: null
+  tune_diffusion_model: true
+  tune_llm: false
+  tune_projector: true
+  tune_visual: false
+  typical_p: 1.0
+  use_bfloat16: false
+  use_relative_action: true
+training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
+  add_rl_callback: false
+  assert_loss_less_than: null
+  batch_size: null
+  bf16: true
+  dataloader_num_workers: 8
+  ddp_bucket_cap_mb: 100
+  deepspeed_stage: 2
+  enable_open_loop_eval: false
+  enable_profiling: false
+  eval_batch_size: 2
+  eval_bf16: true
+  eval_set_split_ratio: 0.1
+  eval_steps: 500
+  eval_strategy: 'no'
+  experiment_name: null
+  fp16: false
+  global_batch_size: 64
+  gradient_accumulation_steps: 1
+  gradient_checkpointing: false
+  learning_rate: 0.0001
+  logging_steps: 10
+  lr_scheduler_type: cosine
+  max_concurrent_uploads: 2
+  max_grad_norm: 1.0
+  max_retries: 3
+  max_steps: 30000
+  num_gpus: 1
+  open_loop_eval_plot_indices: null
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_traj_ids:
+  - 0
+  optim: adamw_torch
+  output_dir: /tmp/groot_output/reachy2
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_steps: 3000
+  save_total_limit: 5
+  save_vl_model: false
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  transformers_access_token: null
+  transformers_cache_dir: null
+  transformers_local_files_only: false
+  transformers_trust_remote_code: true
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  use_ddp: false
+  use_wandb: false
+  wandb_project: finetune-gr00t-n1d6
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

checkpoint-3000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,1020 @@
+{
+  "reachy2": {
+    "state": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.316539764404297,
+          -0.3203781247138977,
+          2.417543411254883,
+          -109.398681640625,
+          -1.9847908020019531,
+          -3.484698534011841,
+          -0.3555510342121124
+        ],
+        "std": [
+          7.066186428070068,
+          3.507993459701538,
+          22.217893600463867,
+          14.851909637451172,
+          13.270956039428711,
+          11.239269256591797,
+          11.768560409545898
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723615646362305,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      }
+    },
+    "action": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.172462463378906,
+          -0.3404896557331085,
+          2.3248910903930664,
+          -109.23794555664062,
+          -1.812583327293396,
+          -3.5105223655700684,
+          -0.2696441113948822
+        ],
+        "std": [
+          7.037580966949463,
+          3.620150566101074,
+          22.205087661743164,
+          14.961851119995117,
+          13.164973258972168,
+          11.261707305908203,
+          11.734737396240234
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723625183105469,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.5416929721832275
+        ],
+        "max": [
+          0.9385063648223877
+        ],
+        "mean": [
+          0.8345981240272522
+        ],
+        "std": [
+          0.12303706258535363
+        ],
+        "q01": [
+          0.6467909216880798
+        ],
+        "q99": [
+          0.9350630044937134
+        ]
+      }
+    },
+    "relative_action": {
+      "arm_joints": {
+        "min": [
+          [
+            -20.440982818603516,
+            -11.665962219238281,
+            -14.02474594116211,
+            -20.817272186279297,
+            -10.09469223022461,
+            -6.401888370513916,
+            -8.767365455627441
+          ],
+          [
+            -20.449968338012695,
+            -11.677082061767578,
+            -25.868675231933594,
+            -30.582874298095703,
+            -17.461694717407227,
+            -11.940333366394043,
+            -16.13103485107422
+          ],
+          [
+            -19.697128295898438,
+            -10.979728698730469,
+            -35.53736114501953,
+            -32.2390022277832,
+            -22.624008178710938,
+            -17.247066497802734,
+            -22.225337982177734
+          ],
+          [
+            -20.74054527282715,
+            -10.334735870361328,
+            -44.72303009033203,
+            -33.42095947265625,
+            -25.056644439697266,
+            -22.322742462158203,
+            -27.965492248535156
+          ],
+          [
+            -24.075529098510742,
+            -10.338696479797363,
+            -50.480987548828125,
+            -34.134925842285156,
+            -26.33944320678711,
+            -25.707611083984375,
+            -31.17104721069336
+          ],
+          [
+            -26.567087173461914,
+            -10.34083366394043,
+            -54.44721603393555,
+            -34.8769645690918,
+            -26.375856399536133,
+            -29.134244918823242,
+            -33.48938751220703
+          ],
+          [
+            -27.659536361694336,
+            -10.340897560119629,
+            -57.30182647705078,
+            -36.380828857421875,
+            -26.375410079956055,
+            -32.329833984375,
+            -35.47602081298828
+          ],
+          [
+            -27.924762725830078,
+            -11.4292573928833,
+            -59.24040603637695,
+            -36.7144775390625,
+            -26.372446060180664,
+            -33.91874313354492,
+            -35.908573150634766
+          ],
+          [
+            -27.95497703552246,
+            -11.528078079223633,
+            -59.352848052978516,
+            -36.841575622558594,
+            -26.24493980407715,
+            -35.34848403930664,
+            -36.83350372314453
+          ],
+          [
+            -28.4979248046875,
+            -11.443194389343262,
+            -59.4117431640625,
+            -36.853206634521484,
+            -26.16107940673828,
+            -36.28956604003906,
+            -36.86201858520508
+          ],
+          [
+            -28.523822784423828,
+            -11.443158149719238,
+            -59.41852951049805,
+            -36.85321807861328,
+            -26.160951614379883,
+            -36.43997573852539,
+            -36.88083267211914
+          ],
+          [
+            -28.522306442260742,
+            -11.443062782287598,
+            -59.424468994140625,
+            -37.14503860473633,
+            -26.16082763671875,
+            -36.443668365478516,
+            -36.880828857421875
+          ],
+          [
+            -28.468467712402344,
+            -11.44290828704834,
+            -59.42174530029297,
+            -37.880104064941406,
+            -26.160717010498047,
+            -36.44268035888672,
+            -36.880828857421875
+          ],
+          [
+            -28.41590118408203,
+            -11.442895889282227,
+            -59.4072265625,
+            -38.47637939453125,
+            -26.160614013671875,
+            -36.4410400390625,
+            -36.880828857421875
+          ],
+          [
+            -28.362062454223633,
+            -11.442879676818848,
+            -59.38927459716797,
+            -41.26643371582031,
+            -26.160587310791016,
+            -36.43976974487305,
+            -36.88082504272461
+          ],
+          [
+            -28.244752883911133,
+            -11.442869186401367,
+            -59.369014739990234,
+            -44.968875885009766,
+            -26.160572052001953,
+            -36.433982849121094,
+            -36.880794525146484
+          ]
+        ],
+        "max": [
+          [
+            5.384329319000244,
+            11.368682861328125,
+            12.580580711364746,
+            21.664548873901367,
+            22.868919372558594,
+            13.739058494567871,
+            11.726442337036133
+          ],
+          [
+            9.6729097366333,
+            11.36864185333252,
+            23.997724533081055,
+            21.659494400024414,
+            22.868900299072266,
+            23.967060089111328,
+            14.994425773620605
+          ],
+          [
+            13.331961631774902,
+            11.190152168273926,
+            33.41366958618164,
+            21.635488510131836,
+            25.30813217163086,
+            26.21414566040039,
+            20.79837417602539
+          ],
+          [
+            16.202699661254883,
+            10.40443229675293,
+            40.09684753417969,
+            23.701637268066406,
+            28.603330612182617,
+            27.53230857849121,
+            25.446346282958984
+          ],
+          [
+            19.05048370361328,
+            10.40436840057373,
+            46.37767791748047,
+            28.047138214111328,
+            31.07068634033203,
+            28.03820037841797,
+            29.473896026611328
+          ],
+          [
+            20.245254516601562,
+            10.402230262756348,
+            51.10308837890625,
+            31.863624572753906,
+            31.10450553894043,
+            29.278892517089844,
+            32.1208381652832
+          ],
+          [
+            21.428510665893555,
+            10.398269653320312,
+            53.42000961303711,
+            34.682308197021484,
+            31.103679656982422,
+            30.250988006591797,
+            33.45231246948242
+          ],
+          [
+            21.98176383972168,
+            10.395852088928223,
+            55.17178726196289,
+            36.64023208618164,
+            31.101837158203125,
+            30.935718536376953,
+            34.229305267333984
+          ],
+          [
+            22.397876739501953,
+            11.4427490234375,
+            55.23532485961914,
+            37.58878707885742,
+            31.096826553344727,
+            31.260026931762695,
+            34.44834899902344
+          ],
+          [
+            22.6441650390625,
+            11.780808448791504,
+            55.34986114501953,
+            38.26850128173828,
+            31.096656799316406,
+            31.333267211914062,
+            34.469146728515625
+          ],
+          [
+            22.692224502563477,
+            11.812055587768555,
+            55.36343765258789,
+            38.49715805053711,
+            31.61040687561035,
+            31.37093162536621,
+            34.469146728515625
+          ],
+          [
+            22.72762107849121,
+            11.678536415100098,
+            55.36940383911133,
+            38.59236526489258,
+            35.083839416503906,
+            31.3709659576416,
+            34.46915054321289
+          ],
+          [
+            24.14029312133789,
+            11.678519248962402,
+            55.369537353515625,
+            38.63944625854492,
+            38.55842208862305,
+            31.37123680114746,
+            34.46915054321289
+          ],
+          [
+            24.973907470703125,
+            11.678510665893555,
+            55.36623001098633,
+            38.656288146972656,
+            39.96430587768555,
+            31.371545791625977,
+            34.46915054321289
+          ],
+          [
+            25.06302261352539,
+            11.678531646728516,
+            55.3525390625,
+            38.694881439208984,
+            40.0225944519043,
+            31.37181854248047,
+            34.468910217285156
+          ],
+          [
+            25.270801544189453,
+            11.678507804870605,
+            55.302616119384766,
+            38.694969177246094,
+            40.022586822509766,
+            31.37185287475586,
+            34.385284423828125
+          ]
+        ],
+        "mean": [
+          [
+            -0.09851349890232086,
+            -1.3889506362829707e-06,
+            -0.021199805662035942,
+            -0.0799250528216362,
+            0.1505303829908371,
+            0.11832442134618759,
+            0.13495467603206635
+          ],
+          [
+            -0.17057012021541595,
+            0.004252429585903883,
+            -0.008995133452117443,
+            -0.22457502782344818,
+            0.261105477809906,
+            0.27207958698272705,
+            0.24137930572032928
+          ],
+          [
+            -0.10028649866580963,
+            0.03141691908240318,
+            0.09048548340797424,
+            -0.5493719577789307,
+            0.19902515411376953,
+            0.47134673595428467,
+            0.2771521210670471
+          ],
+          [
+            -0.05444306880235672,
+            0.05378008633852005,
+            0.1796490103006363,
+            -0.879206120967865,
+            0.1520315706729889,
+            0.6866292357444763,
+            0.3248023986816406
+          ],
+          [
+            -0.013876532204449177,
+            0.07615811377763748,
+            0.2686113715171814,
+            -1.2192035913467407,
+            0.10656605660915375,
+            0.9128072261810303,
+            0.3744547367095947
+          ],
+          [
+            0.025331225246191025,
+            0.09861738234758377,
+            0.3578471839427948,
+            -1.5637919902801514,
+            0.06106000766158104,
+            1.1433591842651367,
+            0.4248007833957672
+          ],
+          [
+            0.0637386217713356,
+            0.12045339494943619,
+            0.44647839665412903,
+            -1.908605933189392,
+            0.016334885731339455,
+            1.3740676641464233,
+            0.4745999872684479
+          ],
+          [
+            0.10053692013025284,
+            0.13872262835502625,
+            0.5284987092018127,
+            -2.2462470531463623,
+            -0.02209978550672531,
+            1.5981824398040771,
+            0.518027663230896
+          ],
+          [
+            0.1247839704155922,
+            0.14948606491088867,
+            0.583483874797821,
+            -2.525705337524414,
+            -0.04458107799291611,
+            1.7731598615646362,
+            0.5445129871368408
+          ],
+          [
+            0.12165514379739761,
+            0.14827895164489746,
+            0.602064847946167,
+            -2.593496561050415,
+            -0.047096796333789825,
+            1.797343134880066,
+            0.5500790476799011
+          ],
+          [
+            0.1124163493514061,
+            0.14204099774360657,
+            0.6022222638130188,
+            -2.55938458442688,
+            -0.04651280865073204,
+            1.7794476747512817,
+            0.5512804985046387
+          ],
+          [
+            0.09131923317909241,
+            0.12857742607593536,
+            0.6059610843658447,
+            -2.4963152408599854,
+            -0.04961543530225754,
+            1.740905523300171,
+            0.5537145137786865
+          ],
+          [
+            0.054126936942338943,
+            0.10883964598178864,
+            0.6135755181312561,
+            -2.381096601486206,
+            -0.056390050798654556,
+            1.6682276725769043,
+            0.5572347044944763
+          ],
+          [
+            -0.001200970378704369,
+            0.08611582219600677,
+            0.6268700361251831,
+            -2.207465648651123,
+            -0.06850837916135788,
+            1.5567927360534668,
+            0.5623834133148193
+          ],
+          [
+            -0.07379738241434097,
+            0.06273127347230911,
+            0.6453807950019836,
+            -1.9786254167556763,
+            -0.08601187914609909,
+            1.4089199304580688,
+            0.5690953135490417
+          ],
+          [
+            -0.15925677120685577,
+            0.039619192481040955,
+            0.6695932149887085,
+            -1.7069525718688965,
+            -0.10912654548883438,
+            1.2322944402694702,
+            0.5775030851364136
+          ]
+        ],
+        "std": [
+          [
+            1.9917564392089844,
+            1.1211100816726685,
+            2.417202949523926,
+            2.94158935546875,
+            2.39365291595459,
+            1.874914526939392,
+            1.6537472009658813
+          ],
+          [
+            3.056582450866699,
+            1.6457527875900269,
+            4.458590507507324,
+            5.178750514984131,
+            3.765504837036133,
+            3.6351206302642822,
+            2.9263687133789062
+          ],
+          [
+            3.596614122390747,
+            1.8142179250717163,
+            6.294006824493408,
+            7.028163433074951,
+            4.514294147491455,
+            5.296512126922607,
+            4.0286865234375
+          ],
+          [
+            4.277676582336426,
+            2.040328025817871,
+            8.043697357177734,
+            8.868635177612305,
+            5.327646255493164,
+            6.84896993637085,
+            5.116058349609375
+          ],
+          [
+            4.939493179321289,
+            2.248694896697998,
+            9.646915435791016,
+            10.583718299865723,
+            6.0862040519714355,
+            8.271647453308105,
+            6.122095584869385
+          ],
+          [
+            5.552292823791504,
+            2.439715623855591,
+            11.097370147705078,
+            12.150360107421875,
+            6.7813005447387695,
+            9.56061840057373,
+            7.034397125244141
+          ],
+          [
+            6.109565258026123,
+            2.6163902282714844,
+            12.404134750366211,
+            13.574675559997559,
+            7.419105052947998,
+            10.72553539276123,
+            7.855382919311523
+          ],
+          [
+            6.608177661895752,
+            2.7779171466827393,
+            13.585582733154297,
+            14.875726699829102,
+            8.011663436889648,
+            11.787652015686035,
+            8.598040580749512
+          ],
+          [
+            7.046725273132324,
+            2.93743634223938,
+            14.688329696655273,
+            16.078651428222656,
+            8.60019302368164,
+            12.769647598266602,
+            9.277371406555176
+          ],
+          [
+            7.43002986907959,
+            3.1344540119171143,
+            15.875958442687988,
+            17.27137565612793,
+            9.292807579040527,
+            13.705930709838867,
+            9.92088794708252
+          ],
+          [
+            7.734869003295898,
+            3.2678585052490234,
+            17.0732421875,
+            18.39610481262207,
+            9.9412841796875,
+            14.559659957885742,
+            10.473697662353516
+          ],
+          [
+            7.963422775268555,
+            3.390157699584961,
+            18.116336822509766,
+            19.400726318359375,
+            10.5236177444458,
+            15.339380264282227,
+            10.94632339477539
+          ],
+          [
+            8.121000289916992,
+            3.5208446979522705,
+            19.046070098876953,
+            20.327335357666016,
+            11.053685188293457,
+            16.05907440185547,
+            11.358160972595215
+          ],
+          [
+            8.218527793884277,
+            3.651193857192993,
+            19.873519897460938,
+            21.182926177978516,
+            11.531027793884277,
+            16.716957092285156,
+            11.718016624450684
+          ],
+          [
+            8.26684284210205,
+            3.7774546146392822,
+            20.612743377685547,
+            21.96902847290039,
+            11.957940101623535,
+            17.311002731323242,
+            12.034592628479004
+          ],
+          [
+            8.276911735534668,
+            3.8994669914245605,
+            21.281835556030273,
+            22.691150665283203,
+            12.342350959777832,
+            17.84532928466797,
+            12.317776679992676
+          ]
+        ],
+        "q01": [
+          [
+            -9.697267532348633,
+            -3.6728341579437256,
+            -10.083383560180664,
+            -7.500851631164551,
+            -5.057114124298096,
+            -5.191580295562744,
+            -5.047628879547119
+          ],
+          [
+            -17.174776077270508,
+            -7.670596122741699,
+            -15.547335624694824,
+            -14.150874137878418,
+            -9.842418670654297,
+            -9.951410293579102,
+            -9.844037055969238
+          ],
+          [
+            -16.49785804748535,
+            -9.060977935791016,
+            -23.215665817260742,
+            -19.967872619628906,
+            -13.348301887512207,
+            -14.247873306274414,
+            -14.610755920410156
+          ],
+          [
+            -16.7535400390625,
+            -9.395726203918457,
+            -29.43296241760254,
+            -25.006587982177734,
+            -16.37109375,
+            -18.130178451538086,
+            -18.5999755859375
+          ],
+          [
+            -17.191621780395508,
+            -9.43246078491211,
+            -36.27531051635742,
+            -28.893749237060547,
+            -18.89956283569336,
+            -21.859458923339844,
+            -22.91756248474121
+          ],
+          [
+            -19.174758911132812,
+            -9.506808280944824,
+            -40.94317626953125,
+            -31.8222713470459,
+            -20.72899055480957,
+            -24.93976402282715,
+            -25.59789276123047
+          ],
+          [
+            -20.843280792236328,
+            -9.511780738830566,
+            -43.895530700683594,
+            -33.7496337890625,
+            -21.83380699157715,
+            -27.3231201171875,
+            -27.804710388183594
+          ],
+          [
+            -22.112207412719727,
+            -9.513916015625,
+            -47.63551330566406,
+            -34.7276611328125,
+            -22.691164016723633,
+            -29.399824142456055,
+            -30.01004981994629
+          ],
+          [
+            -23.385971069335938,
+            -9.525580406188965,
+            -49.76005172729492,
+            -35.31100082397461,
+            -23.10930824279785,
+            -30.991064071655273,
+            -31.236299514770508
+          ],
+          [
+            -23.984119415283203,
+            -9.732311248779297,
+            -51.624576568603516,
+            -35.76605224609375,
+            -23.42118263244629,
+            -31.96640396118164,
+            -32.20992660522461
+          ],
+          [
+            -24.42165184020996,
+            -10.062808990478516,
+            -52.54084777832031,
+            -36.21389389038086,
+            -23.58216094970703,
+            -32.878292083740234,
+            -33.002262115478516
+          ],
+          [
+            -24.806774139404297,
+            -10.06833267211914,
+            -53.43137741088867,
+            -36.29438781738281,
+            -23.675121307373047,
+            -33.4473876953125,
+            -33.2683219909668
+          ],
+          [
+            -24.799230575561523,
+            -10.070316314697266,
+            -54.24421691894531,
+            -36.37746810913086,
+            -23.685352325439453,
+            -33.9517822265625,
+            -33.42109298706055
+          ],
+          [
+            -24.732126235961914,
+            -10.07029914855957,
+            -54.64339828491211,
+            -36.4773063659668,
+            -23.682329177856445,
+            -34.212223052978516,
+            -33.55178451538086
+          ],
+          [
+            -24.3509578704834,
+            -10.070287704467773,
+            -54.56159210205078,
+            -36.60988998413086,
+            -23.682201385498047,
+            -34.409423828125,
+            -33.26326370239258
+          ],
+          [
+            -23.759092330932617,
+            -10.070272445678711,
+            -54.44050216674805,
+            -37.09710693359375,
+            -23.682085037231445,
+            -34.68149948120117,
+            -32.9384765625
+          ]
+        ],
+        "q99": [
+          [
+            3.203047752380371,
+            3.7505993843078613,
+            8.72464370727539,
+            9.948126792907715,
+            11.195253372192383,
+            5.477096080780029,
+            7.006125450134277
+          ],
+          [
+            5.7390570640563965,
+            6.635564804077148,
+            16.971813201904297,
+            18.642690658569336,
+            19.90561294555664,
+            10.401065826416016,
+            11.342881202697754
+          ],
+          [
+            8.532023429870605,
+            8.251216888427734,
+            24.75363540649414,
+            18.472192764282227,
+            19.774227142333984,
+            14.655866622924805,
+            15.70905590057373
+          ],
+          [
+            10.869905471801758,
+            9.230382919311523,
+            31.436893463134766,
+            19.715234756469727,
+            19.81361961364746,
+            18.759492874145508,
+            20.138126373291016
+          ],
+          [
+            12.716119766235352,
+            9.484061241149902,
+            37.04032516479492,
+            23.830785751342773,
+            19.81861114501953,
+            21.909305572509766,
+            23.74039649963379
+          ],
+          [
+            14.470975875854492,
+            9.510160446166992,
+            41.770057678222656,
+            27.173765182495117,
+            20.0618839263916,
+            24.20587921142578,
+            26.801916122436523
+          ],
+          [
+            16.024580001831055,
+            9.510558128356934,
+            45.34690475463867,
+            30.040668487548828,
+            21.872987747192383,
+            25.943317413330078,
+            28.912029266357422
+          ],
+          [
+            16.895124435424805,
+            9.513190269470215,
+            47.93039321899414,
+            32.24497985839844,
+            22.715166091918945,
+            27.396934509277344,
+            30.489883422851562
+          ],
+          [
+            17.761186599731445,
+            9.570911407470703,
+            50.055416107177734,
+            34.065670013427734,
+            23.788578033447266,
+            28.185653686523438,
+            31.546728134155273
+          ],
+          [
+            18.391948699951172,
+            9.712384223937988,
+            51.48924255371094,
+            35.40724182128906,
+            24.940738677978516,
+            28.75865364074707,
+            32.44330978393555
+          ],
+          [
+            19.106931686401367,
+            9.978495597839355,
+            52.564918518066406,
+            36.375064849853516,
+            25.661256790161133,
+            29.218664169311523,
+            33.03133010864258
+          ],
+          [
+            19.62432861328125,
+            10.027887344360352,
+            53.07897186279297,
+            36.976646423339844,
+            26.609840393066406,
+            29.49842071533203,
+            33.33045959472656
+          ],
+          [
+            19.768310546875,
+            10.027881622314453,
+            53.27336883544922,
+            37.44525909423828,
+            28.005691528320312,
+            29.92152214050293,
+            33.3731803894043
+          ],
+          [
+            19.806724548339844,
+            10.027848243713379,
+            53.31312942504883,
+            37.71003341674805,
+            29.430599212646484,
+            29.987842559814453,
+            33.3889274597168
+          ],
+          [
+            20.140544891357422,
+            10.027847290039062,
+            53.163578033447266,
+            37.966617584228516,
+            30.082942962646484,
+            30.24797248840332,
+            33.202571868896484
+          ],
+          [
+            20.164806365966797,
+            10.02784538269043,
+            52.79570770263672,
+            38.0666389465332,
+            30.232667922973633,
+            30.37464714050293,
+            32.94834899902344
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-3000/experiment_cfg/final_model_config.json

@@ -0,0 +1,53 @@
+{
+  "model_type": "Gr00tN1d6",
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "backbone_model_type": "eagle",
+  "model_revision": null,
+  "tune_top_llm_layers": 4,
+  "backbone_embedding_dim": 2048,
+  "tune_llm": false,
+  "tune_visual": false,
+  "select_layer": 16,
+  "reproject_vision": false,
+  "use_flash_attention": true,
+  "load_bf16": true,
+  "collator_overwrite_image_inputs": false,
+  "eagle_collator": true,
+  "backbone_trainable_params_fp32": true,
+  "apply_sincos_state_encoding": true,
+  "use_relative_action": true,
+  "max_state_dim": 128,
+  "max_action_dim": 128,
+  "action_horizon": 50,
+  "hidden_size": 1024,
+  "input_embedding_dim": 1536,
+  "add_pos_embed": true,
+  "attn_dropout": 0.2,
+  "use_vlln": true,
+  "max_seq_len": 1024,
+  "use_alternate_vl_dit": true,
+  "attend_text_every_n_blocks": 2,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "num_inference_timesteps": 4,
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_timestep_buckets": 1000,
+  "tune_projector": true,
+  "tune_diffusion_model": true,
+  "tune_vlln": true,
+  "state_dropout_prob": 0.0,
+  "state_additive_noise_scale": 0.0,
+  "max_num_embodiments": 32
+}

checkpoint-3000/processor_config.json

@@ -0,0 +1,453 @@
+{
+  "processor_class": "Gr00tN1d6Processor",
+  "processor_kwargs": {
+    "modality_configs": {
+      "behavior_r1_pro": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "observation.images.rgb.head_256_256",
+            "observation.images.rgb.left_wrist_256_256",
+            "observation.images.rgb.right_wrist_256_256"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "robot_pos",
+            "robot_ori_cos",
+            "robot_ori_sin",
+            "robot_2d_ori",
+            "robot_2d_ori_cos",
+            "robot_2d_ori_sin",
+            "robot_lin_vel",
+            "robot_ang_vel",
+            "arm_left_qpos",
+            "arm_left_qpos_sin",
+            "arm_left_qpos_cos",
+            "eef_left_pos",
+            "eef_left_quat",
+            "gripper_left_qpos",
+            "arm_right_qpos",
+            "arm_right_qpos_sin",
+            "arm_right_qpos_cos",
+            "eef_right_pos",
+            "eef_right_quat",
+            "gripper_right_qpos",
+            "trunk_qpos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31
+          ],
+          "modality_keys": [
+            "base",
+            "torso",
+            "left_arm",
+            "left_gripper",
+            "right_arm",
+            "right_gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "trunk_qpos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_left_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_right_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "gr1": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "ego_view_bg_crop_pad_res256_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "robocasa_panda_omron": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "res256_image_side_0",
+            "res256_image_side_1",
+            "res256_image_wrist_0"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "end_effector_position_relative",
+            "end_effector_rotation_relative",
+            "gripper_qpos",
+            "base_position",
+            "base_rotation"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "end_effector_position",
+            "end_effector_rotation",
+            "gripper_close",
+            "base_motion",
+            "control_mode"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.action.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "reachy2": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front_cam"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "arm_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "arm_joints",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      }
+    },
+    "image_crop_size": null,
+    "image_target_size": null,
+    "use_albumentations": true,
+    "random_rotation_angle": null,
+    "color_jitter_params": {
+      "brightness": 0.3,
+      "contrast": 0.4,
+      "saturation": 0.5,
+      "hue": 0.08
+    },
+    "shortest_image_edge": 256,
+    "crop_fraction": 0.95,
+    "model_name": "nvidia/Eagle-Block2A-2B-v2",
+    "model_type": "eagle",
+    "formalize_language": true,
+    "max_state_dim": 128,
+    "max_action_dim": 128,
+    "max_action_horizon": 50,
+    "use_percentiles": false,
+    "clip_outliers": true,
+    "apply_sincos_state_encoding": true,
+    "use_relative_action": true
+  }
+}

checkpoint-3000/trainer_state.json

@@ -0,0 +1,1834 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.1,
+  "eval_steps": 500,
+  "global_step": 3000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "grad_norm": 1.3009417057037354,
+      "learning_rate": 6.000000000000001e-07,
+      "loss": 1.1749,
+      "step": 10
+    },
+    {
+      "grad_norm": 1.2199763059616089,
+      "learning_rate": 1.2666666666666667e-06,
+      "loss": 1.1792,
+      "step": 20
+    },
+    {
+      "grad_norm": 1.2334861755371094,
+      "learning_rate": 1.9333333333333336e-06,
+      "loss": 1.1634,
+      "step": 30
+    },
+    {
+      "grad_norm": 0.49021822214126587,
+      "learning_rate": 2.6e-06,
+      "loss": 1.1531,
+      "step": 40
+    },
+    {
+      "grad_norm": 0.5607965588569641,
+      "learning_rate": 3.2666666666666666e-06,
+      "loss": 1.1421,
+      "step": 50
+    },
+    {
+      "grad_norm": 0.5553359389305115,
+      "learning_rate": 3.9333333333333335e-06,
+      "loss": 1.1477,
+      "step": 60
+    },
+    {
+      "grad_norm": 0.46306923031806946,
+      "learning_rate": 4.6e-06,
+      "loss": 1.1444,
+      "step": 70
+    },
+    {
+      "grad_norm": 0.37802454829216003,
+      "learning_rate": 5.266666666666667e-06,
+      "loss": 1.141,
+      "step": 80
+    },
+    {
+      "grad_norm": 0.30536743998527527,
+      "learning_rate": 5.933333333333334e-06,
+      "loss": 1.1446,
+      "step": 90
+    },
+    {
+      "grad_norm": 0.30887874960899353,
+      "learning_rate": 6.6e-06,
+      "loss": 1.1291,
+      "step": 100
+    },
+    {
+      "grad_norm": 0.3307710289955139,
+      "learning_rate": 7.266666666666668e-06,
+      "loss": 1.1313,
+      "step": 110
+    },
+    {
+      "grad_norm": 0.41364726424217224,
+      "learning_rate": 7.933333333333334e-06,
+      "loss": 1.1059,
+      "step": 120
+    },
+    {
+      "grad_norm": 0.6503350734710693,
+      "learning_rate": 8.599999999999999e-06,
+      "loss": 1.092,
+      "step": 130
+    },
+    {
+      "grad_norm": 0.4717787802219391,
+      "learning_rate": 9.266666666666667e-06,
+      "loss": 1.0759,
+      "step": 140
+    },
+    {
+      "grad_norm": 0.4680929481983185,
+      "learning_rate": 9.933333333333334e-06,
+      "loss": 1.0662,
+      "step": 150
+    },
+    {
+      "grad_norm": 0.5255718231201172,
+      "learning_rate": 1.06e-05,
+      "loss": 1.0681,
+      "step": 160
+    },
+    {
+      "grad_norm": 0.5515910983085632,
+      "learning_rate": 1.1266666666666667e-05,
+      "loss": 1.0722,
+      "step": 170
+    },
+    {
+      "grad_norm": 0.48872169852256775,
+      "learning_rate": 1.1933333333333333e-05,
+      "loss": 1.0508,
+      "step": 180
+    },
+    {
+      "grad_norm": 0.684299111366272,
+      "learning_rate": 1.2600000000000001e-05,
+      "loss": 1.0376,
+      "step": 190
+    },
+    {
+      "grad_norm": 1.1723419427871704,
+      "learning_rate": 1.3266666666666666e-05,
+      "loss": 1.0092,
+      "step": 200
+    },
+    {
+      "grad_norm": 1.824255347251892,
+      "learning_rate": 1.3933333333333334e-05,
+      "loss": 0.9294,
+      "step": 210
+    },
+    {
+      "grad_norm": 1.4841794967651367,
+      "learning_rate": 1.4599999999999999e-05,
+      "loss": 0.8456,
+      "step": 220
+    },
+    {
+      "grad_norm": 2.235757350921631,
+      "learning_rate": 1.5266666666666667e-05,
+      "loss": 0.7766,
+      "step": 230
+    },
+    {
+      "grad_norm": 2.0225579738616943,
+      "learning_rate": 1.5933333333333332e-05,
+      "loss": 0.7053,
+      "step": 240
+    },
+    {
+      "grad_norm": 2.08072829246521,
+      "learning_rate": 1.66e-05,
+      "loss": 0.6041,
+      "step": 250
+    },
+    {
+      "grad_norm": 1.6011135578155518,
+      "learning_rate": 1.726666666666667e-05,
+      "loss": 0.5229,
+      "step": 260
+    },
+    {
+      "grad_norm": 2.2321128845214844,
+      "learning_rate": 1.7933333333333337e-05,
+      "loss": 0.4281,
+      "step": 270
+    },
+    {
+      "grad_norm": 2.5934298038482666,
+      "learning_rate": 1.86e-05,
+      "loss": 0.3515,
+      "step": 280
+    },
+    {
+      "grad_norm": 1.9225094318389893,
+      "learning_rate": 1.926666666666667e-05,
+      "loss": 0.2963,
+      "step": 290
+    },
+    {
+      "grad_norm": 2.2619729042053223,
+      "learning_rate": 1.9933333333333334e-05,
+      "loss": 0.2743,
+      "step": 300
+    },
+    {
+      "grad_norm": 1.548316478729248,
+      "learning_rate": 2.06e-05,
+      "loss": 0.253,
+      "step": 310
+    },
+    {
+      "grad_norm": 2.221480369567871,
+      "learning_rate": 2.1266666666666667e-05,
+      "loss": 0.2073,
+      "step": 320
+    },
+    {
+      "grad_norm": 1.8750739097595215,
+      "learning_rate": 2.1933333333333332e-05,
+      "loss": 0.2117,
+      "step": 330
+    },
+    {
+      "grad_norm": 2.543257236480713,
+      "learning_rate": 2.26e-05,
+      "loss": 0.2149,
+      "step": 340
+    },
+    {
+      "grad_norm": 2.042208433151245,
+      "learning_rate": 2.326666666666667e-05,
+      "loss": 0.1861,
+      "step": 350
+    },
+    {
+      "grad_norm": 2.3000669479370117,
+      "learning_rate": 2.3933333333333337e-05,
+      "loss": 0.1862,
+      "step": 360
+    },
+    {
+      "grad_norm": 1.547419786453247,
+      "learning_rate": 2.46e-05,
+      "loss": 0.1628,
+      "step": 370
+    },
+    {
+      "grad_norm": 2.104421615600586,
+      "learning_rate": 2.5266666666666666e-05,
+      "loss": 0.161,
+      "step": 380
+    },
+    {
+      "grad_norm": 1.269059181213379,
+      "learning_rate": 2.5933333333333338e-05,
+      "loss": 0.1528,
+      "step": 390
+    },
+    {
+      "grad_norm": 1.5469577312469482,
+      "learning_rate": 2.6600000000000003e-05,
+      "loss": 0.1441,
+      "step": 400
+    },
+    {
+      "grad_norm": 1.4905346632003784,
+      "learning_rate": 2.7266666666666668e-05,
+      "loss": 0.1473,
+      "step": 410
+    },
+    {
+      "grad_norm": 1.8448561429977417,
+      "learning_rate": 2.7933333333333332e-05,
+      "loss": 0.1371,
+      "step": 420
+    },
+    {
+      "grad_norm": 1.6076915264129639,
+      "learning_rate": 2.86e-05,
+      "loss": 0.1409,
+      "step": 430
+    },
+    {
+      "grad_norm": 1.9364864826202393,
+      "learning_rate": 2.926666666666667e-05,
+      "loss": 0.1318,
+      "step": 440
+    },
+    {
+      "grad_norm": 1.8688479661941528,
+      "learning_rate": 2.9933333333333337e-05,
+      "loss": 0.1294,
+      "step": 450
+    },
+    {
+      "grad_norm": 1.5966501235961914,
+      "learning_rate": 3.06e-05,
+      "loss": 0.1204,
+      "step": 460
+    },
+    {
+      "grad_norm": 1.4227409362792969,
+      "learning_rate": 3.126666666666666e-05,
+      "loss": 0.1104,
+      "step": 470
+    },
+    {
+      "grad_norm": 1.5022822618484497,
+      "learning_rate": 3.1933333333333335e-05,
+      "loss": 0.1155,
+      "step": 480
+    },
+    {
+      "grad_norm": 1.629330039024353,
+      "learning_rate": 3.26e-05,
+      "loss": 0.1186,
+      "step": 490
+    },
+    {
+      "grad_norm": 1.404849648475647,
+      "learning_rate": 3.326666666666667e-05,
+      "loss": 0.1172,
+      "step": 500
+    },
+    {
+      "grad_norm": 1.5502809286117554,
+      "learning_rate": 3.3933333333333336e-05,
+      "loss": 0.1093,
+      "step": 510
+    },
+    {
+      "grad_norm": 1.4023650884628296,
+      "learning_rate": 3.46e-05,
+      "loss": 0.111,
+      "step": 520
+    },
+    {
+      "grad_norm": 1.4860507249832153,
+      "learning_rate": 3.526666666666667e-05,
+      "loss": 0.1104,
+      "step": 530
+    },
+    {
+      "grad_norm": 1.537412405014038,
+      "learning_rate": 3.593333333333334e-05,
+      "loss": 0.1084,
+      "step": 540
+    },
+    {
+      "grad_norm": 1.4830931425094604,
+      "learning_rate": 3.66e-05,
+      "loss": 0.0916,
+      "step": 550
+    },
+    {
+      "grad_norm": 1.4352984428405762,
+      "learning_rate": 3.726666666666667e-05,
+      "loss": 0.104,
+      "step": 560
+    },
+    {
+      "grad_norm": 1.3771204948425293,
+      "learning_rate": 3.793333333333334e-05,
+      "loss": 0.1079,
+      "step": 570
+    },
+    {
+      "grad_norm": 1.2277265787124634,
+      "learning_rate": 3.86e-05,
+      "loss": 0.1036,
+      "step": 580
+    },
+    {
+      "grad_norm": 1.1596946716308594,
+      "learning_rate": 3.926666666666667e-05,
+      "loss": 0.0972,
+      "step": 590
+    },
+    {
+      "grad_norm": 1.66548752784729,
+      "learning_rate": 3.993333333333333e-05,
+      "loss": 0.1048,
+      "step": 600
+    },
+    {
+      "grad_norm": 1.2037672996520996,
+      "learning_rate": 4.0600000000000004e-05,
+      "loss": 0.0939,
+      "step": 610
+    },
+    {
+      "grad_norm": 1.3335858583450317,
+      "learning_rate": 4.126666666666667e-05,
+      "loss": 0.0933,
+      "step": 620
+    },
+    {
+      "grad_norm": 1.2477325201034546,
+      "learning_rate": 4.1933333333333334e-05,
+      "loss": 0.0955,
+      "step": 630
+    },
+    {
+      "grad_norm": 1.316272497177124,
+      "learning_rate": 4.26e-05,
+      "loss": 0.0999,
+      "step": 640
+    },
+    {
+      "grad_norm": 0.7914609313011169,
+      "learning_rate": 4.3266666666666664e-05,
+      "loss": 0.1001,
+      "step": 650
+    },
+    {
+      "grad_norm": 0.918965756893158,
+      "learning_rate": 4.3933333333333335e-05,
+      "loss": 0.0919,
+      "step": 660
+    },
+    {
+      "grad_norm": 1.0262644290924072,
+      "learning_rate": 4.46e-05,
+      "loss": 0.084,
+      "step": 670
+    },
+    {
+      "grad_norm": 1.2336540222167969,
+      "learning_rate": 4.526666666666667e-05,
+      "loss": 0.0854,
+      "step": 680
+    },
+    {
+      "grad_norm": 0.9136370420455933,
+      "learning_rate": 4.5933333333333336e-05,
+      "loss": 0.0812,
+      "step": 690
+    },
+    {
+      "grad_norm": 1.1595962047576904,
+      "learning_rate": 4.660000000000001e-05,
+      "loss": 0.083,
+      "step": 700
+    },
+    {
+      "grad_norm": 1.0120071172714233,
+      "learning_rate": 4.726666666666667e-05,
+      "loss": 0.0828,
+      "step": 710
+    },
+    {
+      "grad_norm": 0.8206706643104553,
+      "learning_rate": 4.793333333333334e-05,
+      "loss": 0.0829,
+      "step": 720
+    },
+    {
+      "grad_norm": 0.9440912008285522,
+      "learning_rate": 4.86e-05,
+      "loss": 0.0848,
+      "step": 730
+    },
+    {
+      "grad_norm": 0.8843382596969604,
+      "learning_rate": 4.926666666666667e-05,
+      "loss": 0.0893,
+      "step": 740
+    },
+    {
+      "grad_norm": 1.190615177154541,
+      "learning_rate": 4.993333333333334e-05,
+      "loss": 0.0866,
+      "step": 750
+    },
+    {
+      "grad_norm": 0.9545953869819641,
+      "learning_rate": 5.0600000000000003e-05,
+      "loss": 0.0851,
+      "step": 760
+    },
+    {
+      "grad_norm": 1.1350778341293335,
+      "learning_rate": 5.1266666666666675e-05,
+      "loss": 0.0906,
+      "step": 770
+    },
+    {
+      "grad_norm": 0.9359301924705505,
+      "learning_rate": 5.193333333333333e-05,
+      "loss": 0.0807,
+      "step": 780
+    },
+    {
+      "grad_norm": 0.884219765663147,
+      "learning_rate": 5.2600000000000005e-05,
+      "loss": 0.0775,
+      "step": 790
+    },
+    {
+      "grad_norm": 1.1784193515777588,
+      "learning_rate": 5.326666666666666e-05,
+      "loss": 0.081,
+      "step": 800
+    },
+    {
+      "grad_norm": 1.332310438156128,
+      "learning_rate": 5.3933333333333334e-05,
+      "loss": 0.0772,
+      "step": 810
+    },
+    {
+      "grad_norm": 1.1053208112716675,
+      "learning_rate": 5.4600000000000006e-05,
+      "loss": 0.0773,
+      "step": 820
+    },
+    {
+      "grad_norm": 0.779946506023407,
+      "learning_rate": 5.5266666666666664e-05,
+      "loss": 0.0757,
+      "step": 830
+    },
+    {
+      "grad_norm": 0.8009672164916992,
+      "learning_rate": 5.5933333333333335e-05,
+      "loss": 0.0784,
+      "step": 840
+    },
+    {
+      "grad_norm": 1.1611248254776,
+      "learning_rate": 5.66e-05,
+      "loss": 0.0756,
+      "step": 850
+    },
+    {
+      "grad_norm": 0.9510448575019836,
+      "learning_rate": 5.726666666666667e-05,
+      "loss": 0.0793,
+      "step": 860
+    },
+    {
+      "grad_norm": 0.7078857421875,
+      "learning_rate": 5.7933333333333337e-05,
+      "loss": 0.0724,
+      "step": 870
+    },
+    {
+      "grad_norm": 1.2287640571594238,
+      "learning_rate": 5.86e-05,
+      "loss": 0.0745,
+      "step": 880
+    },
+    {
+      "grad_norm": 0.8060562014579773,
+      "learning_rate": 5.926666666666667e-05,
+      "loss": 0.0723,
+      "step": 890
+    },
+    {
+      "grad_norm": 0.6258071660995483,
+      "learning_rate": 5.9933333333333345e-05,
+      "loss": 0.0752,
+      "step": 900
+    },
+    {
+      "grad_norm": 0.5880326628684998,
+      "learning_rate": 6.06e-05,
+      "loss": 0.0688,
+      "step": 910
+    },
+    {
+      "grad_norm": 0.82889723777771,
+      "learning_rate": 6.126666666666667e-05,
+      "loss": 0.0704,
+      "step": 920
+    },
+    {
+      "grad_norm": 0.964141845703125,
+      "learning_rate": 6.193333333333333e-05,
+      "loss": 0.0734,
+      "step": 930
+    },
+    {
+      "grad_norm": 0.6693416237831116,
+      "learning_rate": 6.26e-05,
+      "loss": 0.0703,
+      "step": 940
+    },
+    {
+      "grad_norm": 0.7100448608398438,
+      "learning_rate": 6.326666666666667e-05,
+      "loss": 0.0731,
+      "step": 950
+    },
+    {
+      "grad_norm": 0.6688232421875,
+      "learning_rate": 6.393333333333333e-05,
+      "loss": 0.0689,
+      "step": 960
+    },
+    {
+      "grad_norm": 1.0207332372665405,
+      "learning_rate": 6.460000000000001e-05,
+      "loss": 0.0701,
+      "step": 970
+    },
+    {
+      "grad_norm": 0.649103045463562,
+      "learning_rate": 6.526666666666666e-05,
+      "loss": 0.0679,
+      "step": 980
+    },
+    {
+      "grad_norm": 0.8316885232925415,
+      "learning_rate": 6.593333333333334e-05,
+      "loss": 0.0755,
+      "step": 990
+    },
+    {
+      "grad_norm": 0.5805196762084961,
+      "learning_rate": 6.66e-05,
+      "loss": 0.0623,
+      "step": 1000
+    },
+    {
+      "grad_norm": 0.6480134129524231,
+      "learning_rate": 6.726666666666667e-05,
+      "loss": 0.0626,
+      "step": 1010
+    },
+    {
+      "grad_norm": 0.8561599850654602,
+      "learning_rate": 6.793333333333334e-05,
+      "loss": 0.0704,
+      "step": 1020
+    },
+    {
+      "grad_norm": 0.5639697313308716,
+      "learning_rate": 6.860000000000001e-05,
+      "loss": 0.0622,
+      "step": 1030
+    },
+    {
+      "grad_norm": 0.5264688730239868,
+      "learning_rate": 6.926666666666667e-05,
+      "loss": 0.0576,
+      "step": 1040
+    },
+    {
+      "grad_norm": 0.6329230666160583,
+      "learning_rate": 6.993333333333334e-05,
+      "loss": 0.0591,
+      "step": 1050
+    },
+    {
+      "grad_norm": 0.5251981019973755,
+      "learning_rate": 7.06e-05,
+      "loss": 0.0569,
+      "step": 1060
+    },
+    {
+      "grad_norm": 0.80583655834198,
+      "learning_rate": 7.126666666666667e-05,
+      "loss": 0.0619,
+      "step": 1070
+    },
+    {
+      "grad_norm": 0.7077835202217102,
+      "learning_rate": 7.193333333333334e-05,
+      "loss": 0.0638,
+      "step": 1080
+    },
+    {
+      "grad_norm": 0.6568991541862488,
+      "learning_rate": 7.26e-05,
+      "loss": 0.0635,
+      "step": 1090
+    },
+    {
+      "grad_norm": 0.6936503648757935,
+      "learning_rate": 7.326666666666667e-05,
+      "loss": 0.0586,
+      "step": 1100
+    },
+    {
+      "grad_norm": 0.8385834693908691,
+      "learning_rate": 7.393333333333333e-05,
+      "loss": 0.0563,
+      "step": 1110
+    },
+    {
+      "grad_norm": 0.6665258407592773,
+      "learning_rate": 7.46e-05,
+      "loss": 0.0598,
+      "step": 1120
+    },
+    {
+      "grad_norm": 0.7633928060531616,
+      "learning_rate": 7.526666666666668e-05,
+      "loss": 0.0682,
+      "step": 1130
+    },
+    {
+      "grad_norm": 0.6557093858718872,
+      "learning_rate": 7.593333333333334e-05,
+      "loss": 0.0622,
+      "step": 1140
+    },
+    {
+      "grad_norm": 0.7676181793212891,
+      "learning_rate": 7.66e-05,
+      "loss": 0.0554,
+      "step": 1150
+    },
+    {
+      "grad_norm": 0.49698543548583984,
+      "learning_rate": 7.726666666666667e-05,
+      "loss": 0.0619,
+      "step": 1160
+    },
+    {
+      "grad_norm": 0.658534824848175,
+      "learning_rate": 7.793333333333333e-05,
+      "loss": 0.0628,
+      "step": 1170
+    },
+    {
+      "grad_norm": 0.7802417874336243,
+      "learning_rate": 7.860000000000001e-05,
+      "loss": 0.0606,
+      "step": 1180
+    },
+    {
+      "grad_norm": 0.7362853288650513,
+      "learning_rate": 7.926666666666666e-05,
+      "loss": 0.057,
+      "step": 1190
+    },
+    {
+      "grad_norm": 0.7263256311416626,
+      "learning_rate": 7.993333333333334e-05,
+      "loss": 0.0658,
+      "step": 1200
+    },
+    {
+      "grad_norm": 0.605475902557373,
+      "learning_rate": 8.060000000000001e-05,
+      "loss": 0.0585,
+      "step": 1210
+    },
+    {
+      "grad_norm": 0.6713262796401978,
+      "learning_rate": 8.126666666666667e-05,
+      "loss": 0.0613,
+      "step": 1220
+    },
+    {
+      "grad_norm": 0.9524453282356262,
+      "learning_rate": 8.193333333333334e-05,
+      "loss": 0.0623,
+      "step": 1230
+    },
+    {
+      "grad_norm": 0.7799460887908936,
+      "learning_rate": 8.26e-05,
+      "loss": 0.0637,
+      "step": 1240
+    },
+    {
+      "grad_norm": 0.5529918670654297,
+      "learning_rate": 8.326666666666667e-05,
+      "loss": 0.0582,
+      "step": 1250
+    },
+    {
+      "grad_norm": 0.48169708251953125,
+      "learning_rate": 8.393333333333335e-05,
+      "loss": 0.0613,
+      "step": 1260
+    },
+    {
+      "grad_norm": 0.4488230049610138,
+      "learning_rate": 8.46e-05,
+      "loss": 0.0533,
+      "step": 1270
+    },
+    {
+      "grad_norm": 0.43133044242858887,
+      "learning_rate": 8.526666666666667e-05,
+      "loss": 0.0619,
+      "step": 1280
+    },
+    {
+      "grad_norm": 0.6242606043815613,
+      "learning_rate": 8.593333333333333e-05,
+      "loss": 0.056,
+      "step": 1290
+    },
+    {
+      "grad_norm": 0.48737120628356934,
+      "learning_rate": 8.66e-05,
+      "loss": 0.056,
+      "step": 1300
+    },
+    {
+      "grad_norm": 0.47289788722991943,
+      "learning_rate": 8.726666666666667e-05,
+      "loss": 0.0547,
+      "step": 1310
+    },
+    {
+      "grad_norm": 0.8228473663330078,
+      "learning_rate": 8.793333333333333e-05,
+      "loss": 0.0539,
+      "step": 1320
+    },
+    {
+      "grad_norm": 0.6020798087120056,
+      "learning_rate": 8.86e-05,
+      "loss": 0.0639,
+      "step": 1330
+    },
+    {
+      "grad_norm": 0.6205971837043762,
+      "learning_rate": 8.926666666666668e-05,
+      "loss": 0.0644,
+      "step": 1340
+    },
+    {
+      "grad_norm": 0.5030979514122009,
+      "learning_rate": 8.993333333333334e-05,
+      "loss": 0.0584,
+      "step": 1350
+    },
+    {
+      "grad_norm": 0.7016464471817017,
+      "learning_rate": 9.06e-05,
+      "loss": 0.0559,
+      "step": 1360
+    },
+    {
+      "grad_norm": 0.4915011525154114,
+      "learning_rate": 9.126666666666667e-05,
+      "loss": 0.0565,
+      "step": 1370
+    },
+    {
+      "grad_norm": 0.3934551775455475,
+      "learning_rate": 9.193333333333334e-05,
+      "loss": 0.052,
+      "step": 1380
+    },
+    {
+      "grad_norm": 0.5540144443511963,
+      "learning_rate": 9.260000000000001e-05,
+      "loss": 0.0523,
+      "step": 1390
+    },
+    {
+      "grad_norm": 0.5887146592140198,
+      "learning_rate": 9.326666666666667e-05,
+      "loss": 0.0528,
+      "step": 1400
+    },
+    {
+      "grad_norm": 0.6534768342971802,
+      "learning_rate": 9.393333333333334e-05,
+      "loss": 0.0518,
+      "step": 1410
+    },
+    {
+      "grad_norm": 0.5691227316856384,
+      "learning_rate": 9.46e-05,
+      "loss": 0.0536,
+      "step": 1420
+    },
+    {
+      "grad_norm": 0.3767082691192627,
+      "learning_rate": 9.526666666666667e-05,
+      "loss": 0.0568,
+      "step": 1430
+    },
+    {
+      "grad_norm": 0.47469252347946167,
+      "learning_rate": 9.593333333333334e-05,
+      "loss": 0.0569,
+      "step": 1440
+    },
+    {
+      "grad_norm": 0.5082659125328064,
+      "learning_rate": 9.66e-05,
+      "loss": 0.0549,
+      "step": 1450
+    },
+    {
+      "grad_norm": 0.6236725449562073,
+      "learning_rate": 9.726666666666667e-05,
+      "loss": 0.0636,
+      "step": 1460
+    },
+    {
+      "grad_norm": 0.5344048738479614,
+      "learning_rate": 9.793333333333333e-05,
+      "loss": 0.0559,
+      "step": 1470
+    },
+    {
+      "grad_norm": 0.5537348985671997,
+      "learning_rate": 9.86e-05,
+      "loss": 0.0573,
+      "step": 1480
+    },
+    {
+      "grad_norm": 0.6604304313659668,
+      "learning_rate": 9.926666666666668e-05,
+      "loss": 0.0559,
+      "step": 1490
+    },
+    {
+      "grad_norm": 0.6182398796081543,
+      "learning_rate": 9.993333333333334e-05,
+      "loss": 0.0526,
+      "step": 1500
+    },
+    {
+      "grad_norm": 0.4472421109676361,
+      "learning_rate": 9.999997539434007e-05,
+      "loss": 0.0554,
+      "step": 1510
+    },
+    {
+      "grad_norm": 0.5673554539680481,
+      "learning_rate": 9.999989033776898e-05,
+      "loss": 0.0498,
+      "step": 1520
+    },
+    {
+      "grad_norm": 0.631686806678772,
+      "learning_rate": 9.999974452661641e-05,
+      "loss": 0.0558,
+      "step": 1530
+    },
+    {
+      "grad_norm": 0.5880892276763916,
+      "learning_rate": 9.999953796105959e-05,
+      "loss": 0.0565,
+      "step": 1540
+    },
+    {
+      "grad_norm": 0.386528342962265,
+      "learning_rate": 9.999927064134949e-05,
+      "loss": 0.0546,
+      "step": 1550
+    },
+    {
+      "grad_norm": 0.49979427456855774,
+      "learning_rate": 9.999894256781095e-05,
+      "loss": 0.0558,
+      "step": 1560
+    },
+    {
+      "grad_norm": 0.49290731549263,
+      "learning_rate": 9.99985537408426e-05,
+      "loss": 0.0559,
+      "step": 1570
+    },
+    {
+      "grad_norm": 0.6054269671440125,
+      "learning_rate": 9.999810416091688e-05,
+      "loss": 0.0537,
+      "step": 1580
+    },
+    {
+      "grad_norm": 0.4739155173301697,
+      "learning_rate": 9.99975938285801e-05,
+      "loss": 0.0525,
+      "step": 1590
+    },
+    {
+      "grad_norm": 0.4727555215358734,
+      "learning_rate": 9.999702274445236e-05,
+      "loss": 0.0562,
+      "step": 1600
+    },
+    {
+      "grad_norm": 0.4145797789096832,
+      "learning_rate": 9.999639090922756e-05,
+      "loss": 0.0518,
+      "step": 1610
+    },
+    {
+      "grad_norm": 0.6100338697433472,
+      "learning_rate": 9.999569832367346e-05,
+      "loss": 0.0482,
+      "step": 1620
+    },
+    {
+      "grad_norm": 0.36332374811172485,
+      "learning_rate": 9.999494498863162e-05,
+      "loss": 0.0575,
+      "step": 1630
+    },
+    {
+      "grad_norm": 0.5447846055030823,
+      "learning_rate": 9.99941309050174e-05,
+      "loss": 0.0525,
+      "step": 1640
+    },
+    {
+      "grad_norm": 0.5260499119758606,
+      "learning_rate": 9.999325607381999e-05,
+      "loss": 0.0486,
+      "step": 1650
+    },
+    {
+      "grad_norm": 0.40509524941444397,
+      "learning_rate": 9.999232049610238e-05,
+      "loss": 0.047,
+      "step": 1660
+    },
+    {
+      "grad_norm": 0.5393689870834351,
+      "learning_rate": 9.999132417300142e-05,
+      "loss": 0.0483,
+      "step": 1670
+    },
+    {
+      "grad_norm": 0.45050477981567383,
+      "learning_rate": 9.99902671057277e-05,
+      "loss": 0.0516,
+      "step": 1680
+    },
+    {
+      "grad_norm": 0.513198971748352,
+      "learning_rate": 9.998914929556569e-05,
+      "loss": 0.0564,
+      "step": 1690
+    },
+    {
+      "grad_norm": 0.5798249840736389,
+      "learning_rate": 9.998797074387361e-05,
+      "loss": 0.055,
+      "step": 1700
+    },
+    {
+      "grad_norm": 0.4425696134567261,
+      "learning_rate": 9.99867314520835e-05,
+      "loss": 0.0515,
+      "step": 1710
+    },
+    {
+      "grad_norm": 0.4912015199661255,
+      "learning_rate": 9.998543142170126e-05,
+      "loss": 0.0494,
+      "step": 1720
+    },
+    {
+      "grad_norm": 0.4842802882194519,
+      "learning_rate": 9.99840706543065e-05,
+      "loss": 0.0473,
+      "step": 1730
+    },
+    {
+      "grad_norm": 0.37719348073005676,
+      "learning_rate": 9.998264915155274e-05,
+      "loss": 0.0478,
+      "step": 1740
+    },
+    {
+      "grad_norm": 0.44606560468673706,
+      "learning_rate": 9.998116691516718e-05,
+      "loss": 0.0496,
+      "step": 1750
+    },
+    {
+      "grad_norm": 0.4420314133167267,
+      "learning_rate": 9.997962394695091e-05,
+      "loss": 0.0489,
+      "step": 1760
+    },
+    {
+      "grad_norm": 0.5374616384506226,
+      "learning_rate": 9.997802024877875e-05,
+      "loss": 0.0508,
+      "step": 1770
+    },
+    {
+      "grad_norm": 0.3837665915489197,
+      "learning_rate": 9.99763558225994e-05,
+      "loss": 0.0491,
+      "step": 1780
+    },
+    {
+      "grad_norm": 0.5809544324874878,
+      "learning_rate": 9.997463067043526e-05,
+      "loss": 0.0531,
+      "step": 1790
+    },
+    {
+      "grad_norm": 0.5693655610084534,
+      "learning_rate": 9.997284479438253e-05,
+      "loss": 0.0495,
+      "step": 1800
+    },
+    {
+      "grad_norm": 0.5101096630096436,
+      "learning_rate": 9.997099819661127e-05,
+      "loss": 0.0498,
+      "step": 1810
+    },
+    {
+      "grad_norm": 0.4212728440761566,
+      "learning_rate": 9.996909087936524e-05,
+      "loss": 0.0521,
+      "step": 1820
+    },
+    {
+      "grad_norm": 0.3066025376319885,
+      "learning_rate": 9.996712284496201e-05,
+      "loss": 0.0497,
+      "step": 1830
+    },
+    {
+      "grad_norm": 0.5426850318908691,
+      "learning_rate": 9.996509409579293e-05,
+      "loss": 0.051,
+      "step": 1840
+    },
+    {
+      "grad_norm": 0.4050694704055786,
+      "learning_rate": 9.996300463432312e-05,
+      "loss": 0.053,
+      "step": 1850
+    },
+    {
+      "grad_norm": 0.3489755094051361,
+      "learning_rate": 9.996085446309148e-05,
+      "loss": 0.0505,
+      "step": 1860
+    },
+    {
+      "grad_norm": 0.31605497002601624,
+      "learning_rate": 9.995864358471066e-05,
+      "loss": 0.0474,
+      "step": 1870
+    },
+    {
+      "grad_norm": 0.352236807346344,
+      "learning_rate": 9.99563720018671e-05,
+      "loss": 0.047,
+      "step": 1880
+    },
+    {
+      "grad_norm": 0.4507584571838379,
+      "learning_rate": 9.995403971732098e-05,
+      "loss": 0.0504,
+      "step": 1890
+    },
+    {
+      "grad_norm": 0.42925962805747986,
+      "learning_rate": 9.995164673390625e-05,
+      "loss": 0.0473,
+      "step": 1900
+    },
+    {
+      "grad_norm": 0.353121280670166,
+      "learning_rate": 9.994919305453059e-05,
+      "loss": 0.0464,
+      "step": 1910
+    },
+    {
+      "grad_norm": 0.49948233366012573,
+      "learning_rate": 9.994667868217548e-05,
+      "loss": 0.0479,
+      "step": 1920
+    },
+    {
+      "grad_norm": 0.4217991828918457,
+      "learning_rate": 9.99441036198961e-05,
+      "loss": 0.0482,
+      "step": 1930
+    },
+    {
+      "grad_norm": 0.31471049785614014,
+      "learning_rate": 9.99414678708214e-05,
+      "loss": 0.0481,
+      "step": 1940
+    },
+    {
+      "grad_norm": 0.3589913845062256,
+      "learning_rate": 9.993877143815407e-05,
+      "loss": 0.0445,
+      "step": 1950
+    },
+    {
+      "grad_norm": 0.2932628095149994,
+      "learning_rate": 9.993601432517053e-05,
+      "loss": 0.046,
+      "step": 1960
+    },
+    {
+      "grad_norm": 0.35649144649505615,
+      "learning_rate": 9.993319653522091e-05,
+      "loss": 0.0456,
+      "step": 1970
+    },
+    {
+      "grad_norm": 0.4936084449291229,
+      "learning_rate": 9.993031807172911e-05,
+      "loss": 0.0475,
+      "step": 1980
+    },
+    {
+      "grad_norm": 0.33784282207489014,
+      "learning_rate": 9.992737893819273e-05,
+      "loss": 0.0485,
+      "step": 1990
+    },
+    {
+      "grad_norm": 0.5015108585357666,
+      "learning_rate": 9.992437913818312e-05,
+      "loss": 0.0449,
+      "step": 2000
+    },
+    {
+      "grad_norm": 0.42584389448165894,
+      "learning_rate": 9.992131867534526e-05,
+      "loss": 0.0444,
+      "step": 2010
+    },
+    {
+      "grad_norm": 0.41565531492233276,
+      "learning_rate": 9.991819755339796e-05,
+      "loss": 0.0483,
+      "step": 2020
+    },
+    {
+      "grad_norm": 0.380255788564682,
+      "learning_rate": 9.991501577613365e-05,
+      "loss": 0.0482,
+      "step": 2030
+    },
+    {
+      "grad_norm": 0.37106555700302124,
+      "learning_rate": 9.99117733474185e-05,
+      "loss": 0.0454,
+      "step": 2040
+    },
+    {
+      "grad_norm": 0.43988800048828125,
+      "learning_rate": 9.990847027119234e-05,
+      "loss": 0.0457,
+      "step": 2050
+    },
+    {
+      "grad_norm": 0.9517943263053894,
+      "learning_rate": 9.990510655146877e-05,
+      "loss": 0.0431,
+      "step": 2060
+    },
+    {
+      "grad_norm": 0.4155087172985077,
+      "learning_rate": 9.990168219233496e-05,
+      "loss": 0.0455,
+      "step": 2070
+    },
+    {
+      "grad_norm": 0.37841394543647766,
+      "learning_rate": 9.989819719795188e-05,
+      "loss": 0.0425,
+      "step": 2080
+    },
+    {
+      "grad_norm": 0.45210036635398865,
+      "learning_rate": 9.989465157255412e-05,
+      "loss": 0.047,
+      "step": 2090
+    },
+    {
+      "grad_norm": 0.3879967927932739,
+      "learning_rate": 9.989104532044994e-05,
+      "loss": 0.0448,
+      "step": 2100
+    },
+    {
+      "grad_norm": 0.3749665319919586,
+      "learning_rate": 9.988737844602128e-05,
+      "loss": 0.0408,
+      "step": 2110
+    },
+    {
+      "grad_norm": 0.38606101274490356,
+      "learning_rate": 9.988365095372372e-05,
+      "loss": 0.0422,
+      "step": 2120
+    },
+    {
+      "grad_norm": 0.45897138118743896,
+      "learning_rate": 9.987986284808654e-05,
+      "loss": 0.0458,
+      "step": 2130
+    },
+    {
+      "grad_norm": 0.3423418402671814,
+      "learning_rate": 9.987601413371264e-05,
+      "loss": 0.0425,
+      "step": 2140
+    },
+    {
+      "grad_norm": 0.43687859177589417,
+      "learning_rate": 9.987210481527855e-05,
+      "loss": 0.04,
+      "step": 2150
+    },
+    {
+      "grad_norm": 0.3930586278438568,
+      "learning_rate": 9.98681348975345e-05,
+      "loss": 0.042,
+      "step": 2160
+    },
+    {
+      "grad_norm": 0.4354424774646759,
+      "learning_rate": 9.986410438530427e-05,
+      "loss": 0.0409,
+      "step": 2170
+    },
+    {
+      "grad_norm": 0.4108027219772339,
+      "learning_rate": 9.986001328348534e-05,
+      "loss": 0.0394,
+      "step": 2180
+    },
+    {
+      "grad_norm": 0.32850104570388794,
+      "learning_rate": 9.985586159704878e-05,
+      "loss": 0.0407,
+      "step": 2190
+    },
+    {
+      "grad_norm": 0.45172032713890076,
+      "learning_rate": 9.985164933103929e-05,
+      "loss": 0.0475,
+      "step": 2200
+    },
+    {
+      "grad_norm": 0.30814841389656067,
+      "learning_rate": 9.984737649057513e-05,
+      "loss": 0.0438,
+      "step": 2210
+    },
+    {
+      "grad_norm": 0.43183210492134094,
+      "learning_rate": 9.984304308084827e-05,
+      "loss": 0.0485,
+      "step": 2220
+    },
+    {
+      "grad_norm": 0.4616735279560089,
+      "learning_rate": 9.983864910712416e-05,
+      "loss": 0.0414,
+      "step": 2230
+    },
+    {
+      "grad_norm": 0.36560776829719543,
+      "learning_rate": 9.98341945747419e-05,
+      "loss": 0.0398,
+      "step": 2240
+    },
+    {
+      "grad_norm": 0.43177369236946106,
+      "learning_rate": 9.98296794891142e-05,
+      "loss": 0.0445,
+      "step": 2250
+    },
+    {
+      "grad_norm": 0.4279302656650543,
+      "learning_rate": 9.982510385572725e-05,
+      "loss": 0.0467,
+      "step": 2260
+    },
+    {
+      "grad_norm": 0.40069276094436646,
+      "learning_rate": 9.982046768014094e-05,
+      "loss": 0.0445,
+      "step": 2270
+    },
+    {
+      "grad_norm": 0.4348323345184326,
+      "learning_rate": 9.981577096798863e-05,
+      "loss": 0.0369,
+      "step": 2280
+    },
+    {
+      "grad_norm": 0.41257351636886597,
+      "learning_rate": 9.981101372497727e-05,
+      "loss": 0.0445,
+      "step": 2290
+    },
+    {
+      "grad_norm": 0.4516085982322693,
+      "learning_rate": 9.980619595688737e-05,
+      "loss": 0.0438,
+      "step": 2300
+    },
+    {
+      "grad_norm": 0.30238786339759827,
+      "learning_rate": 9.980131766957295e-05,
+      "loss": 0.0421,
+      "step": 2310
+    },
+    {
+      "grad_norm": 0.47715190052986145,
+      "learning_rate": 9.979637886896163e-05,
+      "loss": 0.0423,
+      "step": 2320
+    },
+    {
+      "grad_norm": 0.29774099588394165,
+      "learning_rate": 9.979137956105447e-05,
+      "loss": 0.0464,
+      "step": 2330
+    },
+    {
+      "grad_norm": 0.29573264718055725,
+      "learning_rate": 9.978631975192613e-05,
+      "loss": 0.0376,
+      "step": 2340
+    },
+    {
+      "grad_norm": 0.32869309186935425,
+      "learning_rate": 9.978119944772475e-05,
+      "loss": 0.0467,
+      "step": 2350
+    },
+    {
+      "grad_norm": 0.3470063805580139,
+      "learning_rate": 9.977601865467197e-05,
+      "loss": 0.0405,
+      "step": 2360
+    },
+    {
+      "grad_norm": 0.3539649546146393,
+      "learning_rate": 9.977077737906297e-05,
+      "loss": 0.0403,
+      "step": 2370
+    },
+    {
+      "grad_norm": 0.3649725914001465,
+      "learning_rate": 9.976547562726636e-05,
+      "loss": 0.0414,
+      "step": 2380
+    },
+    {
+      "grad_norm": 0.32129499316215515,
+      "learning_rate": 9.976011340572429e-05,
+      "loss": 0.0423,
+      "step": 2390
+    },
+    {
+      "grad_norm": 0.2593826651573181,
+      "learning_rate": 9.975469072095237e-05,
+      "loss": 0.039,
+      "step": 2400
+    },
+    {
+      "grad_norm": 0.3159441351890564,
+      "learning_rate": 9.974920757953965e-05,
+      "loss": 0.0383,
+      "step": 2410
+    },
+    {
+      "grad_norm": 0.39688485860824585,
+      "learning_rate": 9.97436639881487e-05,
+      "loss": 0.046,
+      "step": 2420
+    },
+    {
+      "grad_norm": 0.3365003764629364,
+      "learning_rate": 9.973805995351545e-05,
+      "loss": 0.0404,
+      "step": 2430
+    },
+    {
+      "grad_norm": 0.3504849672317505,
+      "learning_rate": 9.973239548244939e-05,
+      "loss": 0.0431,
+      "step": 2440
+    },
+    {
+      "grad_norm": 0.4077778160572052,
+      "learning_rate": 9.972667058183333e-05,
+      "loss": 0.04,
+      "step": 2450
+    },
+    {
+      "grad_norm": 0.3008902966976166,
+      "learning_rate": 9.972088525862362e-05,
+      "loss": 0.0382,
+      "step": 2460
+    },
+    {
+      "grad_norm": 0.37545672059059143,
+      "learning_rate": 9.971503951984995e-05,
+      "loss": 0.0438,
+      "step": 2470
+    },
+    {
+      "grad_norm": 0.36859822273254395,
+      "learning_rate": 9.970913337261543e-05,
+      "loss": 0.0409,
+      "step": 2480
+    },
+    {
+      "grad_norm": 0.3510163128376007,
+      "learning_rate": 9.97031668240966e-05,
+      "loss": 0.0479,
+      "step": 2490
+    },
+    {
+      "grad_norm": 0.44913795590400696,
+      "learning_rate": 9.969713988154339e-05,
+      "loss": 0.0409,
+      "step": 2500
+    },
+    {
+      "grad_norm": 0.46766766905784607,
+      "learning_rate": 9.969105255227906e-05,
+      "loss": 0.0405,
+      "step": 2510
+    },
+    {
+      "grad_norm": 0.45442551374435425,
+      "learning_rate": 9.968490484370035e-05,
+      "loss": 0.0394,
+      "step": 2520
+    },
+    {
+      "grad_norm": 0.35867244005203247,
+      "learning_rate": 9.967869676327726e-05,
+      "loss": 0.0419,
+      "step": 2530
+    },
+    {
+      "grad_norm": 0.3887830674648285,
+      "learning_rate": 9.96724283185532e-05,
+      "loss": 0.0364,
+      "step": 2540
+    },
+    {
+      "grad_norm": 0.5115375518798828,
+      "learning_rate": 9.966609951714494e-05,
+      "loss": 0.0411,
+      "step": 2550
+    },
+    {
+      "grad_norm": 0.3326650559902191,
+      "learning_rate": 9.965971036674255e-05,
+      "loss": 0.0459,
+      "step": 2560
+    },
+    {
+      "grad_norm": 0.28854164481163025,
+      "learning_rate": 9.965326087510947e-05,
+      "loss": 0.0398,
+      "step": 2570
+    },
+    {
+      "grad_norm": 0.4342745244503021,
+      "learning_rate": 9.964675105008244e-05,
+      "loss": 0.0407,
+      "step": 2580
+    },
+    {
+      "grad_norm": 0.3539927899837494,
+      "learning_rate": 9.964018089957147e-05,
+      "loss": 0.0345,
+      "step": 2590
+    },
+    {
+      "grad_norm": 0.3186177611351013,
+      "learning_rate": 9.963355043155996e-05,
+      "loss": 0.0456,
+      "step": 2600
+    },
+    {
+      "grad_norm": 0.49563562870025635,
+      "learning_rate": 9.962685965410455e-05,
+      "loss": 0.0423,
+      "step": 2610
+    },
+    {
+      "grad_norm": 0.33249884843826294,
+      "learning_rate": 9.962010857533514e-05,
+      "loss": 0.0373,
+      "step": 2620
+    },
+    {
+      "grad_norm": 0.5187269449234009,
+      "learning_rate": 9.961329720345493e-05,
+      "loss": 0.0386,
+      "step": 2630
+    },
+    {
+      "grad_norm": 0.510249137878418,
+      "learning_rate": 9.96064255467404e-05,
+      "loss": 0.0428,
+      "step": 2640
+    },
+    {
+      "grad_norm": 0.35589292645454407,
+      "learning_rate": 9.959949361354126e-05,
+      "loss": 0.0392,
+      "step": 2650
+    },
+    {
+      "grad_norm": 0.48643842339515686,
+      "learning_rate": 9.959250141228045e-05,
+      "loss": 0.0406,
+      "step": 2660
+    },
+    {
+      "grad_norm": 0.4539337158203125,
+      "learning_rate": 9.958544895145414e-05,
+      "loss": 0.0404,
+      "step": 2670
+    },
+    {
+      "grad_norm": 0.29236212372779846,
+      "learning_rate": 9.957833623963177e-05,
+      "loss": 0.041,
+      "step": 2680
+    },
+    {
+      "grad_norm": 0.5335069894790649,
+      "learning_rate": 9.957116328545593e-05,
+      "loss": 0.0424,
+      "step": 2690
+    },
+    {
+      "grad_norm": 0.3695768117904663,
+      "learning_rate": 9.956393009764244e-05,
+      "loss": 0.0388,
+      "step": 2700
+    },
+    {
+      "grad_norm": 0.36372533440589905,
+      "learning_rate": 9.955663668498032e-05,
+      "loss": 0.0385,
+      "step": 2710
+    },
+    {
+      "grad_norm": 0.4580289125442505,
+      "learning_rate": 9.954928305633173e-05,
+      "loss": 0.0389,
+      "step": 2720
+    },
+    {
+      "grad_norm": 0.43339890241622925,
+      "learning_rate": 9.954186922063204e-05,
+      "loss": 0.0399,
+      "step": 2730
+    },
+    {
+      "grad_norm": 0.35029280185699463,
+      "learning_rate": 9.953439518688974e-05,
+      "loss": 0.0377,
+      "step": 2740
+    },
+    {
+      "grad_norm": 0.33518457412719727,
+      "learning_rate": 9.952686096418652e-05,
+      "loss": 0.0355,
+      "step": 2750
+    },
+    {
+      "grad_norm": 0.30202507972717285,
+      "learning_rate": 9.951926656167714e-05,
+      "loss": 0.0372,
+      "step": 2760
+    },
+    {
+      "grad_norm": 0.34633857011795044,
+      "learning_rate": 9.951161198858953e-05,
+      "loss": 0.0412,
+      "step": 2770
+    },
+    {
+      "grad_norm": 0.40867945551872253,
+      "learning_rate": 9.950389725422471e-05,
+      "loss": 0.0372,
+      "step": 2780
+    },
+    {
+      "grad_norm": 0.2682510316371918,
+      "learning_rate": 9.949612236795682e-05,
+      "loss": 0.0365,
+      "step": 2790
+    },
+    {
+      "grad_norm": 0.4570905566215515,
+      "learning_rate": 9.948828733923305e-05,
+      "loss": 0.0393,
+      "step": 2800
+    },
+    {
+      "grad_norm": 0.36744511127471924,
+      "learning_rate": 9.948039217757374e-05,
+      "loss": 0.0412,
+      "step": 2810
+    },
+    {
+      "grad_norm": 0.29317381978034973,
+      "learning_rate": 9.947243689257225e-05,
+      "loss": 0.038,
+      "step": 2820
+    },
+    {
+      "grad_norm": 0.3504505455493927,
+      "learning_rate": 9.946442149389497e-05,
+      "loss": 0.0348,
+      "step": 2830
+    },
+    {
+      "grad_norm": 0.31019940972328186,
+      "learning_rate": 9.945634599128139e-05,
+      "loss": 0.038,
+      "step": 2840
+    },
+    {
+      "grad_norm": 0.437590092420578,
+      "learning_rate": 9.944821039454402e-05,
+      "loss": 0.0387,
+      "step": 2850
+    },
+    {
+      "grad_norm": 0.36887866258621216,
+      "learning_rate": 9.944001471356835e-05,
+      "loss": 0.0405,
+      "step": 2860
+    },
+    {
+      "grad_norm": 0.2392815351486206,
+      "learning_rate": 9.94317589583129e-05,
+      "loss": 0.0435,
+      "step": 2870
+    },
+    {
+      "grad_norm": 0.21329613029956818,
+      "learning_rate": 9.942344313880922e-05,
+      "loss": 0.0365,
+      "step": 2880
+    },
+    {
+      "grad_norm": 0.3033916652202606,
+      "learning_rate": 9.941506726516179e-05,
+      "loss": 0.0385,
+      "step": 2890
+    },
+    {
+      "grad_norm": 0.2520839273929596,
+      "learning_rate": 9.94066313475481e-05,
+      "loss": 0.0363,
+      "step": 2900
+    },
+    {
+      "grad_norm": 0.43853965401649475,
+      "learning_rate": 9.939813539621857e-05,
+      "loss": 0.0363,
+      "step": 2910
+    },
+    {
+      "grad_norm": 0.36329326033592224,
+      "learning_rate": 9.93895794214966e-05,
+      "loss": 0.0402,
+      "step": 2920
+    },
+    {
+      "grad_norm": 0.26854604482650757,
+      "learning_rate": 9.938096343377852e-05,
+      "loss": 0.0459,
+      "step": 2930
+    },
+    {
+      "grad_norm": 0.44420111179351807,
+      "learning_rate": 9.937228744353353e-05,
+      "loss": 0.0432,
+      "step": 2940
+    },
+    {
+      "grad_norm": 0.363777756690979,
+      "learning_rate": 9.936355146130379e-05,
+      "loss": 0.0369,
+      "step": 2950
+    },
+    {
+      "grad_norm": 0.2704271972179413,
+      "learning_rate": 9.935475549770435e-05,
+      "loss": 0.0379,
+      "step": 2960
+    },
+    {
+      "grad_norm": 0.298907071352005,
+      "learning_rate": 9.934589956342315e-05,
+      "loss": 0.0356,
+      "step": 2970
+    },
+    {
+      "grad_norm": 0.31564566493034363,
+      "learning_rate": 9.933698366922093e-05,
+      "loss": 0.0339,
+      "step": 2980
+    },
+    {
+      "grad_norm": 0.36585021018981934,
+      "learning_rate": 9.93280078259314e-05,
+      "loss": 0.0373,
+      "step": 2990
+    },
+    {
+      "grad_norm": 0.33079367876052856,
+      "learning_rate": 9.931897204446104e-05,
+      "loss": 0.0372,
+      "step": 3000
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 30000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 3000,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-3000/wandb_config.json

@@ -0,0 +1 @@
+{"project": "finetune-gr00t-n1d6", "run_id": "reachy2"}

checkpoint-6000/config.json

@@ -0,0 +1,70 @@
+{
+  "action_horizon": 50,
+  "add_pos_embed": true,
+  "apply_sincos_state_encoding": true,
+  "architectures": [
+    "Gr00tN1d6"
+  ],
+  "attn_dropout": 0.2,
+  "attn_implementation": null,
+  "backbone_embedding_dim": 2048,
+  "backbone_model_type": "eagle",
+  "backbone_trainable_params_fp32": true,
+  "collator_overwrite_image_inputs": false,
+  "color_jitter_params": {
+    "brightness": 0.1,
+    "contrast": 0.1,
+    "hue": 0.1,
+    "saturation": 0.1
+  },
+  "crop_fraction": 0.95,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "eagle_collator": true,
+  "formalize_language": true,
+  "gemma_collator": false,
+  "hidden_size": 1024,
+  "image_crop_size": null,
+  "image_target_size": null,
+  "input_embedding_dim": 1536,
+  "load_bf16": true,
+  "max_action_dim": 128,
+  "max_num_embodiments": 32,
+  "max_seq_len": 1024,
+  "max_state_dim": 128,
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "model_type": "Gr00tN1d6",
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_inference_timesteps": 4,
+  "num_timestep_buckets": 1000,
+  "random_rotation_angle": null,
+  "reproject_vision": false,
+  "select_layer": 16,
+  "shortest_image_edge": 256,
+  "state_dropout_prob": 0.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "tune_diffusion_model": true,
+  "tune_llm": false,
+  "tune_projector": true,
+  "tune_top_llm_layers": 4,
+  "tune_visual": false,
+  "tune_vlln": true,
+  "use_albumentations_transforms": true,
+  "use_alternate_vl_dit": true,
+  "use_flash_attention": true,
+  "use_relative_action": true,
+  "use_vlln": true
+}

checkpoint-6000/embodiment_id.json

@@ -0,0 +1,12 @@
+{
+  "robocasa_panda_omron": 13,
+  "gr1": 20,
+  "behavior_r1_pro": 24,
+  "unitree_g1": 8,
+  "oxe_google": 0,
+  "oxe_widowx": 1,
+  "libero_panda": 2,
+  "oxe_droid": 16,
+  "new_embodiment": 10,
+  "reachy2": 11
+}

checkpoint-6000/experiment_cfg/conf.yaml

@@ -0,0 +1,203 @@
+load_config_path: null
+model:
+  model_type: Gr00tN1d6
+  model_dtype: bfloat16
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  backbone_model_type: eagle
+  model_revision: null
+  tune_top_llm_layers: 4
+  backbone_embedding_dim: 2048
+  tune_llm: false
+  tune_visual: false
+  select_layer: 16
+  reproject_vision: false
+  use_flash_attention: true
+  load_bf16: false
+  collator_overwrite_image_inputs: false
+  eagle_collator: true
+  backbone_trainable_params_fp32: true
+  image_crop_size: null
+  image_target_size: null
+  shortest_image_edge: 256
+  crop_fraction: 0.95
+  random_rotation_angle: null
+  color_jitter_params: null
+  use_albumentations_transforms: true
+  formalize_language: true
+  apply_sincos_state_encoding: false
+  use_relative_action: true
+  max_state_dim: 29
+  max_action_dim: 29
+  action_horizon: 16
+  hidden_size: 1024
+  input_embedding_dim: 1536
+  add_pos_embed: true
+  attn_dropout: 0.2
+  use_vlln: true
+  max_seq_len: 1024
+  use_alternate_vl_dit: true
+  attend_text_every_n_blocks: 2
+  diffusion_model_cfg:
+    positional_embeddings: null
+    num_layers: 32
+    num_attention_heads: 32
+    attention_head_dim: 48
+    norm_type: ada_norm
+    dropout: 0.2
+    final_dropout: true
+    output_dim: 1024
+    interleave_self_attention: true
+  num_inference_timesteps: 4
+  noise_beta_alpha: 1.5
+  noise_beta_beta: 1.0
+  noise_s: 0.999
+  num_timestep_buckets: 1000
+  tune_projector: true
+  tune_diffusion_model: true
+  tune_vlln: true
+  state_dropout_prob: 0.0
+  state_additive_noise_scale: 0.0
+  max_num_embodiments: 32
+data:
+  datasets:
+  - dataset_paths:
+    - ./datasets/reachy2_100
+    embodiment_tag: reachy2
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    reachy2:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front_cam
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - arm_joints
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      action:
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        modality_keys:
+        - arm_joints
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - rep: RELATIVE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: arm_joints
+        - rep: ABSOLUTE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: null
+      language:
+        delta_indices:
+        - 0
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+  download_cache: false
+  shard_size: 1024
+  episode_sampling_rate: 0.1
+  num_shards_per_epoch: 10000
+  override_pretraining_statistics: false
+  mode: single_turn
+  random_chop: 0.0
+  mock_dataset_mode: false
+  shuffle: true
+  seed: 42
+  multiprocessing_context: fork
+  allow_padding: false
+  subsample_ratio: 1.0
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  video_backend: decord
+training:
+  output_dir: /tmp/groot_output/reachy2
+  experiment_name: null
+  max_steps: 30000
+  global_batch_size: 64
+  batch_size: null
+  gradient_accumulation_steps: 1
+  learning_rate: 0.0001
+  lr_scheduler_type: cosine
+  weight_decay: 1.0e-05
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  max_grad_norm: 1.0
+  optim: adamw_torch
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  fp16: false
+  bf16: true
+  eval_bf16: true
+  logging_steps: 10
+  save_steps: 3000
+  save_total_limit: 5
+  save_vl_model: false
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  max_concurrent_uploads: 2
+  eval_strategy: 'no'
+  eval_steps: 500
+  eval_set_split_ratio: 0.1
+  eval_batch_size: 2
+  save_best_eval_metric_name: ''
+  save_best_eval_metric_greater_is_better: true
+  deepspeed_stage: 2
+  gradient_checkpointing: false
+  transformers_trust_remote_code: true
+  transformers_local_files_only: false
+  transformers_cache_dir: null
+  transformers_access_token: null
+  use_ddp: false
+  ddp_bucket_cap_mb: 100
+  num_gpus: 1
+  dataloader_num_workers: 8
+  remove_unused_columns: false
+  use_wandb: false
+  wandb_project: finetune-gr00t-n1d6
+  enable_profiling: false
+  max_retries: 3
+  assert_loss_less_than: null
+  add_rl_callback: false
+  enable_open_loop_eval: false
+  open_loop_eval_traj_ids:
+  - 0
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_plot_indices: null
+max_steps: 30000
+save_steps: 3000

checkpoint-6000/experiment_cfg/config.yaml

@@ -0,0 +1,237 @@
+!!python/object:gr00t.configs.base_config.Config
+data: !!python/object:gr00t.configs.data.data_config.DataConfig
+  allow_padding: false
+  datasets:
+  - !!python/object:gr00t.configs.data.data_config.SingleDatasetConfig
+    dataset_paths:
+    - ./datasets/reachy2_100
+    dataset_type: physical_embodiment
+    embodiment_tag: reachy2
+    mix_ratio: 1.0
+    val_dataset_path: null
+  download_cache: false
+  episode_sampling_rate: 0.1
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  mock_dataset_mode: false
+  modality_configs:
+    reachy2:
+      action: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs:
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: &id001 !!python/object/apply:gr00t.data.types.ActionFormat
+          - default
+          rep: !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - relative
+          state_key: arm_joints
+          type: &id002 !!python/object/apply:gr00t.data.types.ActionType
+          - non_eef
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: *id001
+          rep: !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - absolute
+          state_key: null
+          type: *id002
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        mean_std_embedding_keys: null
+        modality_keys:
+        - arm_joints
+        - gripper
+        sin_cos_embedding_keys: null
+      language: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+      state: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - arm_joints
+        sin_cos_embedding_keys: null
+      video: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - front_cam
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 10000
+  override_pretraining_statistics: false
+  random_chop: 0.0
+  seed: 42
+  shard_size: 1024
+  shuffle: true
+  subsample_ratio: 1.0
+  video_backend: decord
+load_config_path: null
+model: !!python/object:gr00t.configs.model.gr00t_n1d6.Gr00tN1d6Config
+  _attn_implementation_autoset: false
+  _attn_implementation_internal: null
+  _commit_hash: null
+  _name_or_path: ''
+  add_cross_attention: false
+  architectures: null
+  backbone_model_type: eagle
+  backbone_trainable_params_fp32: true
+  bad_words_ids: null
+  begin_suppress_tokens: null
+  bos_token_id: null
+  chunk_size_feed_forward: 0
+  color_jitter_params: null
+  cross_attention_hidden_size: null
+  decoder_start_token_id: null
+  diffusion_model_cfg:
+    attention_head_dim: 48
+    dropout: 0.2
+    final_dropout: true
+    interleave_self_attention: true
+    norm_type: ada_norm
+    num_attention_heads: 32
+    num_layers: 32
+    output_dim: 1024
+    positional_embeddings: null
+  diversity_penalty: 0.0
+  do_sample: false
+  eagle_collator: true
+  early_stopping: false
+  encoder_no_repeat_ngram_size: 0
+  eos_token_id: null
+  exponential_decay_length_penalty: null
+  finetuning_task: null
+  forced_bos_token_id: null
+  forced_eos_token_id: null
+  id2label:
+    0: LABEL_0
+    1: LABEL_1
+  is_decoder: false
+  is_encoder_decoder: false
+  label2id:
+    LABEL_0: 0
+    LABEL_1: 1
+  length_penalty: 1.0
+  load_bf16: false
+  max_length: 20
+  min_length: 0
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  no_repeat_ngram_size: 0
+  num_beam_groups: 1
+  num_beams: 1
+  num_return_sequences: 1
+  output_attentions: false
+  output_hidden_states: false
+  output_scores: false
+  pad_token_id: null
+  prefix: null
+  problem_type: null
+  pruned_heads: {}
+  random_rotation_angle: null
+  remove_invalid_values: false
+  repetition_penalty: 1.0
+  reproject_vision: false
+  return_dict: true
+  return_dict_in_generate: false
+  sep_token_id: null
+  state_dropout_prob: 0.0
+  suppress_tokens: null
+  task_specific_params: null
+  temperature: 1.0
+  tf_legacy_loss: false
+  tie_encoder_decoder: false
+  tie_word_embeddings: true
+  tokenizer_class: null
+  top_k: 50
+  top_p: 1.0
+  torch_dtype: null
+  torchscript: false
+  transformers_version: null
+  tune_diffusion_model: true
+  tune_llm: false
+  tune_projector: true
+  tune_visual: false
+  typical_p: 1.0
+  use_bfloat16: false
+  use_relative_action: true
+training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
+  add_rl_callback: false
+  assert_loss_less_than: null
+  batch_size: null
+  bf16: true
+  dataloader_num_workers: 8
+  ddp_bucket_cap_mb: 100
+  deepspeed_stage: 2
+  enable_open_loop_eval: false
+  enable_profiling: false
+  eval_batch_size: 2
+  eval_bf16: true
+  eval_set_split_ratio: 0.1
+  eval_steps: 500
+  eval_strategy: 'no'
+  experiment_name: null
+  fp16: false
+  global_batch_size: 64
+  gradient_accumulation_steps: 1
+  gradient_checkpointing: false
+  learning_rate: 0.0001
+  logging_steps: 10
+  lr_scheduler_type: cosine
+  max_concurrent_uploads: 2
+  max_grad_norm: 1.0
+  max_retries: 3
+  max_steps: 30000
+  num_gpus: 1
+  open_loop_eval_plot_indices: null
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_traj_ids:
+  - 0
+  optim: adamw_torch
+  output_dir: /tmp/groot_output/reachy2
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_steps: 3000
+  save_total_limit: 5
+  save_vl_model: false
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  transformers_access_token: null
+  transformers_cache_dir: null
+  transformers_local_files_only: false
+  transformers_trust_remote_code: true
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  use_ddp: false
+  use_wandb: false
+  wandb_project: finetune-gr00t-n1d6
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

checkpoint-6000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,1020 @@
+{
+  "reachy2": {
+    "state": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.316539764404297,
+          -0.3203781247138977,
+          2.417543411254883,
+          -109.398681640625,
+          -1.9847908020019531,
+          -3.484698534011841,
+          -0.3555510342121124
+        ],
+        "std": [
+          7.066186428070068,
+          3.507993459701538,
+          22.217893600463867,
+          14.851909637451172,
+          13.270956039428711,
+          11.239269256591797,
+          11.768560409545898
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723615646362305,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      }
+    },
+    "action": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.172462463378906,
+          -0.3404896557331085,
+          2.3248910903930664,
+          -109.23794555664062,
+          -1.812583327293396,
+          -3.5105223655700684,
+          -0.2696441113948822
+        ],
+        "std": [
+          7.037580966949463,
+          3.620150566101074,
+          22.205087661743164,
+          14.961851119995117,
+          13.164973258972168,
+          11.261707305908203,
+          11.734737396240234
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723625183105469,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.5416929721832275
+        ],
+        "max": [
+          0.9385063648223877
+        ],
+        "mean": [
+          0.8345981240272522
+        ],
+        "std": [
+          0.12303706258535363
+        ],
+        "q01": [
+          0.6467909216880798
+        ],
+        "q99": [
+          0.9350630044937134
+        ]
+      }
+    },
+    "relative_action": {
+      "arm_joints": {
+        "min": [
+          [
+            -20.440982818603516,
+            -11.665962219238281,
+            -14.02474594116211,
+            -20.817272186279297,
+            -10.09469223022461,
+            -6.401888370513916,
+            -8.767365455627441
+          ],
+          [
+            -20.449968338012695,
+            -11.677082061767578,
+            -25.868675231933594,
+            -30.582874298095703,
+            -17.461694717407227,
+            -11.940333366394043,
+            -16.13103485107422
+          ],
+          [
+            -19.697128295898438,
+            -10.979728698730469,
+            -35.53736114501953,
+            -32.2390022277832,
+            -22.624008178710938,
+            -17.247066497802734,
+            -22.225337982177734
+          ],
+          [
+            -20.74054527282715,
+            -10.334735870361328,
+            -44.72303009033203,
+            -33.42095947265625,
+            -25.056644439697266,
+            -22.322742462158203,
+            -27.965492248535156
+          ],
+          [
+            -24.075529098510742,
+            -10.338696479797363,
+            -50.480987548828125,
+            -34.134925842285156,
+            -26.33944320678711,
+            -25.707611083984375,
+            -31.17104721069336
+          ],
+          [
+            -26.567087173461914,
+            -10.34083366394043,
+            -54.44721603393555,
+            -34.8769645690918,
+            -26.375856399536133,
+            -29.134244918823242,
+            -33.48938751220703
+          ],
+          [
+            -27.659536361694336,
+            -10.340897560119629,
+            -57.30182647705078,
+            -36.380828857421875,
+            -26.375410079956055,
+            -32.329833984375,
+            -35.47602081298828
+          ],
+          [
+            -27.924762725830078,
+            -11.4292573928833,
+            -59.24040603637695,
+            -36.7144775390625,
+            -26.372446060180664,
+            -33.91874313354492,
+            -35.908573150634766
+          ],
+          [
+            -27.95497703552246,
+            -11.528078079223633,
+            -59.352848052978516,
+            -36.841575622558594,
+            -26.24493980407715,
+            -35.34848403930664,
+            -36.83350372314453
+          ],
+          [
+            -28.4979248046875,
+            -11.443194389343262,
+            -59.4117431640625,
+            -36.853206634521484,
+            -26.16107940673828,
+            -36.28956604003906,
+            -36.86201858520508
+          ],
+          [
+            -28.523822784423828,
+            -11.443158149719238,
+            -59.41852951049805,
+            -36.85321807861328,
+            -26.160951614379883,
+            -36.43997573852539,
+            -36.88083267211914
+          ],
+          [
+            -28.522306442260742,
+            -11.443062782287598,
+            -59.424468994140625,
+            -37.14503860473633,
+            -26.16082763671875,
+            -36.443668365478516,
+            -36.880828857421875
+          ],
+          [
+            -28.468467712402344,
+            -11.44290828704834,
+            -59.42174530029297,
+            -37.880104064941406,
+            -26.160717010498047,
+            -36.44268035888672,
+            -36.880828857421875
+          ],
+          [
+            -28.41590118408203,
+            -11.442895889282227,
+            -59.4072265625,
+            -38.47637939453125,
+            -26.160614013671875,
+            -36.4410400390625,
+            -36.880828857421875
+          ],
+          [
+            -28.362062454223633,
+            -11.442879676818848,
+            -59.38927459716797,
+            -41.26643371582031,
+            -26.160587310791016,
+            -36.43976974487305,
+            -36.88082504272461
+          ],
+          [
+            -28.244752883911133,
+            -11.442869186401367,
+            -59.369014739990234,
+            -44.968875885009766,
+            -26.160572052001953,
+            -36.433982849121094,
+            -36.880794525146484
+          ]
+        ],
+        "max": [
+          [
+            5.384329319000244,
+            11.368682861328125,
+            12.580580711364746,
+            21.664548873901367,
+            22.868919372558594,
+            13.739058494567871,
+            11.726442337036133
+          ],
+          [
+            9.6729097366333,
+            11.36864185333252,
+            23.997724533081055,
+            21.659494400024414,
+            22.868900299072266,
+            23.967060089111328,
+            14.994425773620605
+          ],
+          [
+            13.331961631774902,
+            11.190152168273926,
+            33.41366958618164,
+            21.635488510131836,
+            25.30813217163086,
+            26.21414566040039,
+            20.79837417602539
+          ],
+          [
+            16.202699661254883,
+            10.40443229675293,
+            40.09684753417969,
+            23.701637268066406,
+            28.603330612182617,
+            27.53230857849121,
+            25.446346282958984
+          ],
+          [
+            19.05048370361328,
+            10.40436840057373,
+            46.37767791748047,
+            28.047138214111328,
+            31.07068634033203,
+            28.03820037841797,
+            29.473896026611328
+          ],
+          [
+            20.245254516601562,
+            10.402230262756348,
+            51.10308837890625,
+            31.863624572753906,
+            31.10450553894043,
+            29.278892517089844,
+            32.1208381652832
+          ],
+          [
+            21.428510665893555,
+            10.398269653320312,
+            53.42000961303711,
+            34.682308197021484,
+            31.103679656982422,
+            30.250988006591797,
+            33.45231246948242
+          ],
+          [
+            21.98176383972168,
+            10.395852088928223,
+            55.17178726196289,
+            36.64023208618164,
+            31.101837158203125,
+            30.935718536376953,
+            34.229305267333984
+          ],
+          [
+            22.397876739501953,
+            11.4427490234375,
+            55.23532485961914,
+            37.58878707885742,
+            31.096826553344727,
+            31.260026931762695,
+            34.44834899902344
+          ],
+          [
+            22.6441650390625,
+            11.780808448791504,
+            55.34986114501953,
+            38.26850128173828,
+            31.096656799316406,
+            31.333267211914062,
+            34.469146728515625
+          ],
+          [
+            22.692224502563477,
+            11.812055587768555,
+            55.36343765258789,
+            38.49715805053711,
+            31.61040687561035,
+            31.37093162536621,
+            34.469146728515625
+          ],
+          [
+            22.72762107849121,
+            11.678536415100098,
+            55.36940383911133,
+            38.59236526489258,
+            35.083839416503906,
+            31.3709659576416,
+            34.46915054321289
+          ],
+          [
+            24.14029312133789,
+            11.678519248962402,
+            55.369537353515625,
+            38.63944625854492,
+            38.55842208862305,
+            31.37123680114746,
+            34.46915054321289
+          ],
+          [
+            24.973907470703125,
+            11.678510665893555,
+            55.36623001098633,
+            38.656288146972656,
+            39.96430587768555,
+            31.371545791625977,
+            34.46915054321289
+          ],
+          [
+            25.06302261352539,
+            11.678531646728516,
+            55.3525390625,
+            38.694881439208984,
+            40.0225944519043,
+            31.37181854248047,
+            34.468910217285156
+          ],
+          [
+            25.270801544189453,
+            11.678507804870605,
+            55.302616119384766,
+            38.694969177246094,
+            40.022586822509766,
+            31.37185287475586,
+            34.385284423828125
+          ]
+        ],
+        "mean": [
+          [
+            -0.09851349890232086,
+            -1.3889506362829707e-06,
+            -0.021199805662035942,
+            -0.0799250528216362,
+            0.1505303829908371,
+            0.11832442134618759,
+            0.13495467603206635
+          ],
+          [
+            -0.17057012021541595,
+            0.004252429585903883,
+            -0.008995133452117443,
+            -0.22457502782344818,
+            0.261105477809906,
+            0.27207958698272705,
+            0.24137930572032928
+          ],
+          [
+            -0.10028649866580963,
+            0.03141691908240318,
+            0.09048548340797424,
+            -0.5493719577789307,
+            0.19902515411376953,
+            0.47134673595428467,
+            0.2771521210670471
+          ],
+          [
+            -0.05444306880235672,
+            0.05378008633852005,
+            0.1796490103006363,
+            -0.879206120967865,
+            0.1520315706729889,
+            0.6866292357444763,
+            0.3248023986816406
+          ],
+          [
+            -0.013876532204449177,
+            0.07615811377763748,
+            0.2686113715171814,
+            -1.2192035913467407,
+            0.10656605660915375,
+            0.9128072261810303,
+            0.3744547367095947
+          ],
+          [
+            0.025331225246191025,
+            0.09861738234758377,
+            0.3578471839427948,
+            -1.5637919902801514,
+            0.06106000766158104,
+            1.1433591842651367,
+            0.4248007833957672
+          ],
+          [
+            0.0637386217713356,
+            0.12045339494943619,
+            0.44647839665412903,
+            -1.908605933189392,
+            0.016334885731339455,
+            1.3740676641464233,
+            0.4745999872684479
+          ],
+          [
+            0.10053692013025284,
+            0.13872262835502625,
+            0.5284987092018127,
+            -2.2462470531463623,
+            -0.02209978550672531,
+            1.5981824398040771,
+            0.518027663230896
+          ],
+          [
+            0.1247839704155922,
+            0.14948606491088867,
+            0.583483874797821,
+            -2.525705337524414,
+            -0.04458107799291611,
+            1.7731598615646362,
+            0.5445129871368408
+          ],
+          [
+            0.12165514379739761,
+            0.14827895164489746,
+            0.602064847946167,
+            -2.593496561050415,
+            -0.047096796333789825,
+            1.797343134880066,
+            0.5500790476799011
+          ],
+          [
+            0.1124163493514061,
+            0.14204099774360657,
+            0.6022222638130188,
+            -2.55938458442688,
+            -0.04651280865073204,
+            1.7794476747512817,
+            0.5512804985046387
+          ],
+          [
+            0.09131923317909241,
+            0.12857742607593536,
+            0.6059610843658447,
+            -2.4963152408599854,
+            -0.04961543530225754,
+            1.740905523300171,
+            0.5537145137786865
+          ],
+          [
+            0.054126936942338943,
+            0.10883964598178864,
+            0.6135755181312561,
+            -2.381096601486206,
+            -0.056390050798654556,
+            1.6682276725769043,
+            0.5572347044944763
+          ],
+          [
+            -0.001200970378704369,
+            0.08611582219600677,
+            0.6268700361251831,
+            -2.207465648651123,
+            -0.06850837916135788,
+            1.5567927360534668,
+            0.5623834133148193
+          ],
+          [
+            -0.07379738241434097,
+            0.06273127347230911,
+            0.6453807950019836,
+            -1.9786254167556763,
+            -0.08601187914609909,
+            1.4089199304580688,
+            0.5690953135490417
+          ],
+          [
+            -0.15925677120685577,
+            0.039619192481040955,
+            0.6695932149887085,
+            -1.7069525718688965,
+            -0.10912654548883438,
+            1.2322944402694702,
+            0.5775030851364136
+          ]
+        ],
+        "std": [
+          [
+            1.9917564392089844,
+            1.1211100816726685,
+            2.417202949523926,
+            2.94158935546875,
+            2.39365291595459,
+            1.874914526939392,
+            1.6537472009658813
+          ],
+          [
+            3.056582450866699,
+            1.6457527875900269,
+            4.458590507507324,
+            5.178750514984131,
+            3.765504837036133,
+            3.6351206302642822,
+            2.9263687133789062
+          ],
+          [
+            3.596614122390747,
+            1.8142179250717163,
+            6.294006824493408,
+            7.028163433074951,
+            4.514294147491455,
+            5.296512126922607,
+            4.0286865234375
+          ],
+          [
+            4.277676582336426,
+            2.040328025817871,
+            8.043697357177734,
+            8.868635177612305,
+            5.327646255493164,
+            6.84896993637085,
+            5.116058349609375
+          ],
+          [
+            4.939493179321289,
+            2.248694896697998,
+            9.646915435791016,
+            10.583718299865723,
+            6.0862040519714355,
+            8.271647453308105,
+            6.122095584869385
+          ],
+          [
+            5.552292823791504,
+            2.439715623855591,
+            11.097370147705078,
+            12.150360107421875,
+            6.7813005447387695,
+            9.56061840057373,
+            7.034397125244141
+          ],
+          [
+            6.109565258026123,
+            2.6163902282714844,
+            12.404134750366211,
+            13.574675559997559,
+            7.419105052947998,
+            10.72553539276123,
+            7.855382919311523
+          ],
+          [
+            6.608177661895752,
+            2.7779171466827393,
+            13.585582733154297,
+            14.875726699829102,
+            8.011663436889648,
+            11.787652015686035,
+            8.598040580749512
+          ],
+          [
+            7.046725273132324,
+            2.93743634223938,
+            14.688329696655273,
+            16.078651428222656,
+            8.60019302368164,
+            12.769647598266602,
+            9.277371406555176
+          ],
+          [
+            7.43002986907959,
+            3.1344540119171143,
+            15.875958442687988,
+            17.27137565612793,
+            9.292807579040527,
+            13.705930709838867,
+            9.92088794708252
+          ],
+          [
+            7.734869003295898,
+            3.2678585052490234,
+            17.0732421875,
+            18.39610481262207,
+            9.9412841796875,
+            14.559659957885742,
+            10.473697662353516
+          ],
+          [
+            7.963422775268555,
+            3.390157699584961,
+            18.116336822509766,
+            19.400726318359375,
+            10.5236177444458,
+            15.339380264282227,
+            10.94632339477539
+          ],
+          [
+            8.121000289916992,
+            3.5208446979522705,
+            19.046070098876953,
+            20.327335357666016,
+            11.053685188293457,
+            16.05907440185547,
+            11.358160972595215
+          ],
+          [
+            8.218527793884277,
+            3.651193857192993,
+            19.873519897460938,
+            21.182926177978516,
+            11.531027793884277,
+            16.716957092285156,
+            11.718016624450684
+          ],
+          [
+            8.26684284210205,
+            3.7774546146392822,
+            20.612743377685547,
+            21.96902847290039,
+            11.957940101623535,
+            17.311002731323242,
+            12.034592628479004
+          ],
+          [
+            8.276911735534668,
+            3.8994669914245605,
+            21.281835556030273,
+            22.691150665283203,
+            12.342350959777832,
+            17.84532928466797,
+            12.317776679992676
+          ]
+        ],
+        "q01": [
+          [
+            -9.697267532348633,
+            -3.6728341579437256,
+            -10.083383560180664,
+            -7.500851631164551,
+            -5.057114124298096,
+            -5.191580295562744,
+            -5.047628879547119
+          ],
+          [
+            -17.174776077270508,
+            -7.670596122741699,
+            -15.547335624694824,
+            -14.150874137878418,
+            -9.842418670654297,
+            -9.951410293579102,
+            -9.844037055969238
+          ],
+          [
+            -16.49785804748535,
+            -9.060977935791016,
+            -23.215665817260742,
+            -19.967872619628906,
+            -13.348301887512207,
+            -14.247873306274414,
+            -14.610755920410156
+          ],
+          [
+            -16.7535400390625,
+            -9.395726203918457,
+            -29.43296241760254,
+            -25.006587982177734,
+            -16.37109375,
+            -18.130178451538086,
+            -18.5999755859375
+          ],
+          [
+            -17.191621780395508,
+            -9.43246078491211,
+            -36.27531051635742,
+            -28.893749237060547,
+            -18.89956283569336,
+            -21.859458923339844,
+            -22.91756248474121
+          ],
+          [
+            -19.174758911132812,
+            -9.506808280944824,
+            -40.94317626953125,
+            -31.8222713470459,
+            -20.72899055480957,
+            -24.93976402282715,
+            -25.59789276123047
+          ],
+          [
+            -20.843280792236328,
+            -9.511780738830566,
+            -43.895530700683594,
+            -33.7496337890625,
+            -21.83380699157715,
+            -27.3231201171875,
+            -27.804710388183594
+          ],
+          [
+            -22.112207412719727,
+            -9.513916015625,
+            -47.63551330566406,
+            -34.7276611328125,
+            -22.691164016723633,
+            -29.399824142456055,
+            -30.01004981994629
+          ],
+          [
+            -23.385971069335938,
+            -9.525580406188965,
+            -49.76005172729492,
+            -35.31100082397461,
+            -23.10930824279785,
+            -30.991064071655273,
+            -31.236299514770508
+          ],
+          [
+            -23.984119415283203,
+            -9.732311248779297,
+            -51.624576568603516,
+            -35.76605224609375,
+            -23.42118263244629,
+            -31.96640396118164,
+            -32.20992660522461
+          ],
+          [
+            -24.42165184020996,
+            -10.062808990478516,
+            -52.54084777832031,
+            -36.21389389038086,
+            -23.58216094970703,
+            -32.878292083740234,
+            -33.002262115478516
+          ],
+          [
+            -24.806774139404297,
+            -10.06833267211914,
+            -53.43137741088867,
+            -36.29438781738281,
+            -23.675121307373047,
+            -33.4473876953125,
+            -33.2683219909668
+          ],
+          [
+            -24.799230575561523,
+            -10.070316314697266,
+            -54.24421691894531,
+            -36.37746810913086,
+            -23.685352325439453,
+            -33.9517822265625,
+            -33.42109298706055
+          ],
+          [
+            -24.732126235961914,
+            -10.07029914855957,
+            -54.64339828491211,
+            -36.4773063659668,
+            -23.682329177856445,
+            -34.212223052978516,
+            -33.55178451538086
+          ],
+          [
+            -24.3509578704834,
+            -10.070287704467773,
+            -54.56159210205078,
+            -36.60988998413086,
+            -23.682201385498047,
+            -34.409423828125,
+            -33.26326370239258
+          ],
+          [
+            -23.759092330932617,
+            -10.070272445678711,
+            -54.44050216674805,
+            -37.09710693359375,
+            -23.682085037231445,
+            -34.68149948120117,
+            -32.9384765625
+          ]
+        ],
+        "q99": [
+          [
+            3.203047752380371,
+            3.7505993843078613,
+            8.72464370727539,
+            9.948126792907715,
+            11.195253372192383,
+            5.477096080780029,
+            7.006125450134277
+          ],
+          [
+            5.7390570640563965,
+            6.635564804077148,
+            16.971813201904297,
+            18.642690658569336,
+            19.90561294555664,
+            10.401065826416016,
+            11.342881202697754
+          ],
+          [
+            8.532023429870605,
+            8.251216888427734,
+            24.75363540649414,
+            18.472192764282227,
+            19.774227142333984,
+            14.655866622924805,
+            15.70905590057373
+          ],
+          [
+            10.869905471801758,
+            9.230382919311523,
+            31.436893463134766,
+            19.715234756469727,
+            19.81361961364746,
+            18.759492874145508,
+            20.138126373291016
+          ],
+          [
+            12.716119766235352,
+            9.484061241149902,
+            37.04032516479492,
+            23.830785751342773,
+            19.81861114501953,
+            21.909305572509766,
+            23.74039649963379
+          ],
+          [
+            14.470975875854492,
+            9.510160446166992,
+            41.770057678222656,
+            27.173765182495117,
+            20.0618839263916,
+            24.20587921142578,
+            26.801916122436523
+          ],
+          [
+            16.024580001831055,
+            9.510558128356934,
+            45.34690475463867,
+            30.040668487548828,
+            21.872987747192383,
+            25.943317413330078,
+            28.912029266357422
+          ],
+          [
+            16.895124435424805,
+            9.513190269470215,
+            47.93039321899414,
+            32.24497985839844,
+            22.715166091918945,
+            27.396934509277344,
+            30.489883422851562
+          ],
+          [
+            17.761186599731445,
+            9.570911407470703,
+            50.055416107177734,
+            34.065670013427734,
+            23.788578033447266,
+            28.185653686523438,
+            31.546728134155273
+          ],
+          [
+            18.391948699951172,
+            9.712384223937988,
+            51.48924255371094,
+            35.40724182128906,
+            24.940738677978516,
+            28.75865364074707,
+            32.44330978393555
+          ],
+          [
+            19.106931686401367,
+            9.978495597839355,
+            52.564918518066406,
+            36.375064849853516,
+            25.661256790161133,
+            29.218664169311523,
+            33.03133010864258
+          ],
+          [
+            19.62432861328125,
+            10.027887344360352,
+            53.07897186279297,
+            36.976646423339844,
+            26.609840393066406,
+            29.49842071533203,
+            33.33045959472656
+          ],
+          [
+            19.768310546875,
+            10.027881622314453,
+            53.27336883544922,
+            37.44525909423828,
+            28.005691528320312,
+            29.92152214050293,
+            33.3731803894043
+          ],
+          [
+            19.806724548339844,
+            10.027848243713379,
+            53.31312942504883,
+            37.71003341674805,
+            29.430599212646484,
+            29.987842559814453,
+            33.3889274597168
+          ],
+          [
+            20.140544891357422,
+            10.027847290039062,
+            53.163578033447266,
+            37.966617584228516,
+            30.082942962646484,
+            30.24797248840332,
+            33.202571868896484
+          ],
+          [
+            20.164806365966797,
+            10.02784538269043,
+            52.79570770263672,
+            38.0666389465332,
+            30.232667922973633,
+            30.37464714050293,
+            32.94834899902344
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-6000/experiment_cfg/final_model_config.json

@@ -0,0 +1,53 @@
+{
+  "model_type": "Gr00tN1d6",
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "backbone_model_type": "eagle",
+  "model_revision": null,
+  "tune_top_llm_layers": 4,
+  "backbone_embedding_dim": 2048,
+  "tune_llm": false,
+  "tune_visual": false,
+  "select_layer": 16,
+  "reproject_vision": false,
+  "use_flash_attention": true,
+  "load_bf16": true,
+  "collator_overwrite_image_inputs": false,
+  "eagle_collator": true,
+  "backbone_trainable_params_fp32": true,
+  "apply_sincos_state_encoding": true,
+  "use_relative_action": true,
+  "max_state_dim": 128,
+  "max_action_dim": 128,
+  "action_horizon": 50,
+  "hidden_size": 1024,
+  "input_embedding_dim": 1536,
+  "add_pos_embed": true,
+  "attn_dropout": 0.2,
+  "use_vlln": true,
+  "max_seq_len": 1024,
+  "use_alternate_vl_dit": true,
+  "attend_text_every_n_blocks": 2,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "num_inference_timesteps": 4,
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_timestep_buckets": 1000,
+  "tune_projector": true,
+  "tune_diffusion_model": true,
+  "tune_vlln": true,
+  "state_dropout_prob": 0.0,
+  "state_additive_noise_scale": 0.0,
+  "max_num_embodiments": 32
+}

checkpoint-6000/processor_config.json

@@ -0,0 +1,453 @@
+{
+  "processor_class": "Gr00tN1d6Processor",
+  "processor_kwargs": {
+    "modality_configs": {
+      "behavior_r1_pro": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "observation.images.rgb.head_256_256",
+            "observation.images.rgb.left_wrist_256_256",
+            "observation.images.rgb.right_wrist_256_256"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "robot_pos",
+            "robot_ori_cos",
+            "robot_ori_sin",
+            "robot_2d_ori",
+            "robot_2d_ori_cos",
+            "robot_2d_ori_sin",
+            "robot_lin_vel",
+            "robot_ang_vel",
+            "arm_left_qpos",
+            "arm_left_qpos_sin",
+            "arm_left_qpos_cos",
+            "eef_left_pos",
+            "eef_left_quat",
+            "gripper_left_qpos",
+            "arm_right_qpos",
+            "arm_right_qpos_sin",
+            "arm_right_qpos_cos",
+            "eef_right_pos",
+            "eef_right_quat",
+            "gripper_right_qpos",
+            "trunk_qpos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31
+          ],
+          "modality_keys": [
+            "base",
+            "torso",
+            "left_arm",
+            "left_gripper",
+            "right_arm",
+            "right_gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "trunk_qpos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_left_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_right_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "gr1": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "ego_view_bg_crop_pad_res256_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "robocasa_panda_omron": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "res256_image_side_0",
+            "res256_image_side_1",
+            "res256_image_wrist_0"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "end_effector_position_relative",
+            "end_effector_rotation_relative",
+            "gripper_qpos",
+            "base_position",
+            "base_rotation"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "end_effector_position",
+            "end_effector_rotation",
+            "gripper_close",
+            "base_motion",
+            "control_mode"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.action.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "reachy2": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front_cam"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "arm_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "arm_joints",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      }
+    },
+    "image_crop_size": null,
+    "image_target_size": null,
+    "use_albumentations": true,
+    "random_rotation_angle": null,
+    "color_jitter_params": {
+      "brightness": 0.3,
+      "contrast": 0.4,
+      "saturation": 0.5,
+      "hue": 0.08
+    },
+    "shortest_image_edge": 256,
+    "crop_fraction": 0.95,
+    "model_name": "nvidia/Eagle-Block2A-2B-v2",
+    "model_type": "eagle",
+    "formalize_language": true,
+    "max_state_dim": 128,
+    "max_action_dim": 128,
+    "max_action_horizon": 50,
+    "use_percentiles": false,
+    "clip_outliers": true,
+    "apply_sincos_state_encoding": true,
+    "use_relative_action": true
+  }
+}

checkpoint-6000/trainer_state.json

@@ -0,0 +1,3634 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.2,
+  "eval_steps": 500,
+  "global_step": 6000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "grad_norm": 1.3009417057037354,
+      "learning_rate": 6.000000000000001e-07,
+      "loss": 1.1749,
+      "step": 10
+    },
+    {
+      "grad_norm": 1.2199763059616089,
+      "learning_rate": 1.2666666666666667e-06,
+      "loss": 1.1792,
+      "step": 20
+    },
+    {
+      "grad_norm": 1.2334861755371094,
+      "learning_rate": 1.9333333333333336e-06,
+      "loss": 1.1634,
+      "step": 30
+    },
+    {
+      "grad_norm": 0.49021822214126587,
+      "learning_rate": 2.6e-06,
+      "loss": 1.1531,
+      "step": 40
+    },
+    {
+      "grad_norm": 0.5607965588569641,
+      "learning_rate": 3.2666666666666666e-06,
+      "loss": 1.1421,
+      "step": 50
+    },
+    {
+      "grad_norm": 0.5553359389305115,
+      "learning_rate": 3.9333333333333335e-06,
+      "loss": 1.1477,
+      "step": 60
+    },
+    {
+      "grad_norm": 0.46306923031806946,
+      "learning_rate": 4.6e-06,
+      "loss": 1.1444,
+      "step": 70
+    },
+    {
+      "grad_norm": 0.37802454829216003,
+      "learning_rate": 5.266666666666667e-06,
+      "loss": 1.141,
+      "step": 80
+    },
+    {
+      "grad_norm": 0.30536743998527527,
+      "learning_rate": 5.933333333333334e-06,
+      "loss": 1.1446,
+      "step": 90
+    },
+    {
+      "grad_norm": 0.30887874960899353,
+      "learning_rate": 6.6e-06,
+      "loss": 1.1291,
+      "step": 100
+    },
+    {
+      "grad_norm": 0.3307710289955139,
+      "learning_rate": 7.266666666666668e-06,
+      "loss": 1.1313,
+      "step": 110
+    },
+    {
+      "grad_norm": 0.41364726424217224,
+      "learning_rate": 7.933333333333334e-06,
+      "loss": 1.1059,
+      "step": 120
+    },
+    {
+      "grad_norm": 0.6503350734710693,
+      "learning_rate": 8.599999999999999e-06,
+      "loss": 1.092,
+      "step": 130
+    },
+    {
+      "grad_norm": 0.4717787802219391,
+      "learning_rate": 9.266666666666667e-06,
+      "loss": 1.0759,
+      "step": 140
+    },
+    {
+      "grad_norm": 0.4680929481983185,
+      "learning_rate": 9.933333333333334e-06,
+      "loss": 1.0662,
+      "step": 150
+    },
+    {
+      "grad_norm": 0.5255718231201172,
+      "learning_rate": 1.06e-05,
+      "loss": 1.0681,
+      "step": 160
+    },
+    {
+      "grad_norm": 0.5515910983085632,
+      "learning_rate": 1.1266666666666667e-05,
+      "loss": 1.0722,
+      "step": 170
+    },
+    {
+      "grad_norm": 0.48872169852256775,
+      "learning_rate": 1.1933333333333333e-05,
+      "loss": 1.0508,
+      "step": 180
+    },
+    {
+      "grad_norm": 0.684299111366272,
+      "learning_rate": 1.2600000000000001e-05,
+      "loss": 1.0376,
+      "step": 190
+    },
+    {
+      "grad_norm": 1.1723419427871704,
+      "learning_rate": 1.3266666666666666e-05,
+      "loss": 1.0092,
+      "step": 200
+    },
+    {
+      "grad_norm": 1.824255347251892,
+      "learning_rate": 1.3933333333333334e-05,
+      "loss": 0.9294,
+      "step": 210
+    },
+    {
+      "grad_norm": 1.4841794967651367,
+      "learning_rate": 1.4599999999999999e-05,
+      "loss": 0.8456,
+      "step": 220
+    },
+    {
+      "grad_norm": 2.235757350921631,
+      "learning_rate": 1.5266666666666667e-05,
+      "loss": 0.7766,
+      "step": 230
+    },
+    {
+      "grad_norm": 2.0225579738616943,
+      "learning_rate": 1.5933333333333332e-05,
+      "loss": 0.7053,
+      "step": 240
+    },
+    {
+      "grad_norm": 2.08072829246521,
+      "learning_rate": 1.66e-05,
+      "loss": 0.6041,
+      "step": 250
+    },
+    {
+      "grad_norm": 1.6011135578155518,
+      "learning_rate": 1.726666666666667e-05,
+      "loss": 0.5229,
+      "step": 260
+    },
+    {
+      "grad_norm": 2.2321128845214844,
+      "learning_rate": 1.7933333333333337e-05,
+      "loss": 0.4281,
+      "step": 270
+    },
+    {
+      "grad_norm": 2.5934298038482666,
+      "learning_rate": 1.86e-05,
+      "loss": 0.3515,
+      "step": 280
+    },
+    {
+      "grad_norm": 1.9225094318389893,
+      "learning_rate": 1.926666666666667e-05,
+      "loss": 0.2963,
+      "step": 290
+    },
+    {
+      "grad_norm": 2.2619729042053223,
+      "learning_rate": 1.9933333333333334e-05,
+      "loss": 0.2743,
+      "step": 300
+    },
+    {
+      "grad_norm": 1.548316478729248,
+      "learning_rate": 2.06e-05,
+      "loss": 0.253,
+      "step": 310
+    },
+    {
+      "grad_norm": 2.221480369567871,
+      "learning_rate": 2.1266666666666667e-05,
+      "loss": 0.2073,
+      "step": 320
+    },
+    {
+      "grad_norm": 1.8750739097595215,
+      "learning_rate": 2.1933333333333332e-05,
+      "loss": 0.2117,
+      "step": 330
+    },
+    {
+      "grad_norm": 2.543257236480713,
+      "learning_rate": 2.26e-05,
+      "loss": 0.2149,
+      "step": 340
+    },
+    {
+      "grad_norm": 2.042208433151245,
+      "learning_rate": 2.326666666666667e-05,
+      "loss": 0.1861,
+      "step": 350
+    },
+    {
+      "grad_norm": 2.3000669479370117,
+      "learning_rate": 2.3933333333333337e-05,
+      "loss": 0.1862,
+      "step": 360
+    },
+    {
+      "grad_norm": 1.547419786453247,
+      "learning_rate": 2.46e-05,
+      "loss": 0.1628,
+      "step": 370
+    },
+    {
+      "grad_norm": 2.104421615600586,
+      "learning_rate": 2.5266666666666666e-05,
+      "loss": 0.161,
+      "step": 380
+    },
+    {
+      "grad_norm": 1.269059181213379,
+      "learning_rate": 2.5933333333333338e-05,
+      "loss": 0.1528,
+      "step": 390
+    },
+    {
+      "grad_norm": 1.5469577312469482,
+      "learning_rate": 2.6600000000000003e-05,
+      "loss": 0.1441,
+      "step": 400
+    },
+    {
+      "grad_norm": 1.4905346632003784,
+      "learning_rate": 2.7266666666666668e-05,
+      "loss": 0.1473,
+      "step": 410
+    },
+    {
+      "grad_norm": 1.8448561429977417,
+      "learning_rate": 2.7933333333333332e-05,
+      "loss": 0.1371,
+      "step": 420
+    },
+    {
+      "grad_norm": 1.6076915264129639,
+      "learning_rate": 2.86e-05,
+      "loss": 0.1409,
+      "step": 430
+    },
+    {
+      "grad_norm": 1.9364864826202393,
+      "learning_rate": 2.926666666666667e-05,
+      "loss": 0.1318,
+      "step": 440
+    },
+    {
+      "grad_norm": 1.8688479661941528,
+      "learning_rate": 2.9933333333333337e-05,
+      "loss": 0.1294,
+      "step": 450
+    },
+    {
+      "grad_norm": 1.5966501235961914,
+      "learning_rate": 3.06e-05,
+      "loss": 0.1204,
+      "step": 460
+    },
+    {
+      "grad_norm": 1.4227409362792969,
+      "learning_rate": 3.126666666666666e-05,
+      "loss": 0.1104,
+      "step": 470
+    },
+    {
+      "grad_norm": 1.5022822618484497,
+      "learning_rate": 3.1933333333333335e-05,
+      "loss": 0.1155,
+      "step": 480
+    },
+    {
+      "grad_norm": 1.629330039024353,
+      "learning_rate": 3.26e-05,
+      "loss": 0.1186,
+      "step": 490
+    },
+    {
+      "grad_norm": 1.404849648475647,
+      "learning_rate": 3.326666666666667e-05,
+      "loss": 0.1172,
+      "step": 500
+    },
+    {
+      "grad_norm": 1.5502809286117554,
+      "learning_rate": 3.3933333333333336e-05,
+      "loss": 0.1093,
+      "step": 510
+    },
+    {
+      "grad_norm": 1.4023650884628296,
+      "learning_rate": 3.46e-05,
+      "loss": 0.111,
+      "step": 520
+    },
+    {
+      "grad_norm": 1.4860507249832153,
+      "learning_rate": 3.526666666666667e-05,
+      "loss": 0.1104,
+      "step": 530
+    },
+    {
+      "grad_norm": 1.537412405014038,
+      "learning_rate": 3.593333333333334e-05,
+      "loss": 0.1084,
+      "step": 540
+    },
+    {
+      "grad_norm": 1.4830931425094604,
+      "learning_rate": 3.66e-05,
+      "loss": 0.0916,
+      "step": 550
+    },
+    {
+      "grad_norm": 1.4352984428405762,
+      "learning_rate": 3.726666666666667e-05,
+      "loss": 0.104,
+      "step": 560
+    },
+    {
+      "grad_norm": 1.3771204948425293,
+      "learning_rate": 3.793333333333334e-05,
+      "loss": 0.1079,
+      "step": 570
+    },
+    {
+      "grad_norm": 1.2277265787124634,
+      "learning_rate": 3.86e-05,
+      "loss": 0.1036,
+      "step": 580
+    },
+    {
+      "grad_norm": 1.1596946716308594,
+      "learning_rate": 3.926666666666667e-05,
+      "loss": 0.0972,
+      "step": 590
+    },
+    {
+      "grad_norm": 1.66548752784729,
+      "learning_rate": 3.993333333333333e-05,
+      "loss": 0.1048,
+      "step": 600
+    },
+    {
+      "grad_norm": 1.2037672996520996,
+      "learning_rate": 4.0600000000000004e-05,
+      "loss": 0.0939,
+      "step": 610
+    },
+    {
+      "grad_norm": 1.3335858583450317,
+      "learning_rate": 4.126666666666667e-05,
+      "loss": 0.0933,
+      "step": 620
+    },
+    {
+      "grad_norm": 1.2477325201034546,
+      "learning_rate": 4.1933333333333334e-05,
+      "loss": 0.0955,
+      "step": 630
+    },
+    {
+      "grad_norm": 1.316272497177124,
+      "learning_rate": 4.26e-05,
+      "loss": 0.0999,
+      "step": 640
+    },
+    {
+      "grad_norm": 0.7914609313011169,
+      "learning_rate": 4.3266666666666664e-05,
+      "loss": 0.1001,
+      "step": 650
+    },
+    {
+      "grad_norm": 0.918965756893158,
+      "learning_rate": 4.3933333333333335e-05,
+      "loss": 0.0919,
+      "step": 660
+    },
+    {
+      "grad_norm": 1.0262644290924072,
+      "learning_rate": 4.46e-05,
+      "loss": 0.084,
+      "step": 670
+    },
+    {
+      "grad_norm": 1.2336540222167969,
+      "learning_rate": 4.526666666666667e-05,
+      "loss": 0.0854,
+      "step": 680
+    },
+    {
+      "grad_norm": 0.9136370420455933,
+      "learning_rate": 4.5933333333333336e-05,
+      "loss": 0.0812,
+      "step": 690
+    },
+    {
+      "grad_norm": 1.1595962047576904,
+      "learning_rate": 4.660000000000001e-05,
+      "loss": 0.083,
+      "step": 700
+    },
+    {
+      "grad_norm": 1.0120071172714233,
+      "learning_rate": 4.726666666666667e-05,
+      "loss": 0.0828,
+      "step": 710
+    },
+    {
+      "grad_norm": 0.8206706643104553,
+      "learning_rate": 4.793333333333334e-05,
+      "loss": 0.0829,
+      "step": 720
+    },
+    {
+      "grad_norm": 0.9440912008285522,
+      "learning_rate": 4.86e-05,
+      "loss": 0.0848,
+      "step": 730
+    },
+    {
+      "grad_norm": 0.8843382596969604,
+      "learning_rate": 4.926666666666667e-05,
+      "loss": 0.0893,
+      "step": 740
+    },
+    {
+      "grad_norm": 1.190615177154541,
+      "learning_rate": 4.993333333333334e-05,
+      "loss": 0.0866,
+      "step": 750
+    },
+    {
+      "grad_norm": 0.9545953869819641,
+      "learning_rate": 5.0600000000000003e-05,
+      "loss": 0.0851,
+      "step": 760
+    },
+    {
+      "grad_norm": 1.1350778341293335,
+      "learning_rate": 5.1266666666666675e-05,
+      "loss": 0.0906,
+      "step": 770
+    },
+    {
+      "grad_norm": 0.9359301924705505,
+      "learning_rate": 5.193333333333333e-05,
+      "loss": 0.0807,
+      "step": 780
+    },
+    {
+      "grad_norm": 0.884219765663147,
+      "learning_rate": 5.2600000000000005e-05,
+      "loss": 0.0775,
+      "step": 790
+    },
+    {
+      "grad_norm": 1.1784193515777588,
+      "learning_rate": 5.326666666666666e-05,
+      "loss": 0.081,
+      "step": 800
+    },
+    {
+      "grad_norm": 1.332310438156128,
+      "learning_rate": 5.3933333333333334e-05,
+      "loss": 0.0772,
+      "step": 810
+    },
+    {
+      "grad_norm": 1.1053208112716675,
+      "learning_rate": 5.4600000000000006e-05,
+      "loss": 0.0773,
+      "step": 820
+    },
+    {
+      "grad_norm": 0.779946506023407,
+      "learning_rate": 5.5266666666666664e-05,
+      "loss": 0.0757,
+      "step": 830
+    },
+    {
+      "grad_norm": 0.8009672164916992,
+      "learning_rate": 5.5933333333333335e-05,
+      "loss": 0.0784,
+      "step": 840
+    },
+    {
+      "grad_norm": 1.1611248254776,
+      "learning_rate": 5.66e-05,
+      "loss": 0.0756,
+      "step": 850
+    },
+    {
+      "grad_norm": 0.9510448575019836,
+      "learning_rate": 5.726666666666667e-05,
+      "loss": 0.0793,
+      "step": 860
+    },
+    {
+      "grad_norm": 0.7078857421875,
+      "learning_rate": 5.7933333333333337e-05,
+      "loss": 0.0724,
+      "step": 870
+    },
+    {
+      "grad_norm": 1.2287640571594238,
+      "learning_rate": 5.86e-05,
+      "loss": 0.0745,
+      "step": 880
+    },
+    {
+      "grad_norm": 0.8060562014579773,
+      "learning_rate": 5.926666666666667e-05,
+      "loss": 0.0723,
+      "step": 890
+    },
+    {
+      "grad_norm": 0.6258071660995483,
+      "learning_rate": 5.9933333333333345e-05,
+      "loss": 0.0752,
+      "step": 900
+    },
+    {
+      "grad_norm": 0.5880326628684998,
+      "learning_rate": 6.06e-05,
+      "loss": 0.0688,
+      "step": 910
+    },
+    {
+      "grad_norm": 0.82889723777771,
+      "learning_rate": 6.126666666666667e-05,
+      "loss": 0.0704,
+      "step": 920
+    },
+    {
+      "grad_norm": 0.964141845703125,
+      "learning_rate": 6.193333333333333e-05,
+      "loss": 0.0734,
+      "step": 930
+    },
+    {
+      "grad_norm": 0.6693416237831116,
+      "learning_rate": 6.26e-05,
+      "loss": 0.0703,
+      "step": 940
+    },
+    {
+      "grad_norm": 0.7100448608398438,
+      "learning_rate": 6.326666666666667e-05,
+      "loss": 0.0731,
+      "step": 950
+    },
+    {
+      "grad_norm": 0.6688232421875,
+      "learning_rate": 6.393333333333333e-05,
+      "loss": 0.0689,
+      "step": 960
+    },
+    {
+      "grad_norm": 1.0207332372665405,
+      "learning_rate": 6.460000000000001e-05,
+      "loss": 0.0701,
+      "step": 970
+    },
+    {
+      "grad_norm": 0.649103045463562,
+      "learning_rate": 6.526666666666666e-05,
+      "loss": 0.0679,
+      "step": 980
+    },
+    {
+      "grad_norm": 0.8316885232925415,
+      "learning_rate": 6.593333333333334e-05,
+      "loss": 0.0755,
+      "step": 990
+    },
+    {
+      "grad_norm": 0.5805196762084961,
+      "learning_rate": 6.66e-05,
+      "loss": 0.0623,
+      "step": 1000
+    },
+    {
+      "grad_norm": 0.6480134129524231,
+      "learning_rate": 6.726666666666667e-05,
+      "loss": 0.0626,
+      "step": 1010
+    },
+    {
+      "grad_norm": 0.8561599850654602,
+      "learning_rate": 6.793333333333334e-05,
+      "loss": 0.0704,
+      "step": 1020
+    },
+    {
+      "grad_norm": 0.5639697313308716,
+      "learning_rate": 6.860000000000001e-05,
+      "loss": 0.0622,
+      "step": 1030
+    },
+    {
+      "grad_norm": 0.5264688730239868,
+      "learning_rate": 6.926666666666667e-05,
+      "loss": 0.0576,
+      "step": 1040
+    },
+    {
+      "grad_norm": 0.6329230666160583,
+      "learning_rate": 6.993333333333334e-05,
+      "loss": 0.0591,
+      "step": 1050
+    },
+    {
+      "grad_norm": 0.5251981019973755,
+      "learning_rate": 7.06e-05,
+      "loss": 0.0569,
+      "step": 1060
+    },
+    {
+      "grad_norm": 0.80583655834198,
+      "learning_rate": 7.126666666666667e-05,
+      "loss": 0.0619,
+      "step": 1070
+    },
+    {
+      "grad_norm": 0.7077835202217102,
+      "learning_rate": 7.193333333333334e-05,
+      "loss": 0.0638,
+      "step": 1080
+    },
+    {
+      "grad_norm": 0.6568991541862488,
+      "learning_rate": 7.26e-05,
+      "loss": 0.0635,
+      "step": 1090
+    },
+    {
+      "grad_norm": 0.6936503648757935,
+      "learning_rate": 7.326666666666667e-05,
+      "loss": 0.0586,
+      "step": 1100
+    },
+    {
+      "grad_norm": 0.8385834693908691,
+      "learning_rate": 7.393333333333333e-05,
+      "loss": 0.0563,
+      "step": 1110
+    },
+    {
+      "grad_norm": 0.6665258407592773,
+      "learning_rate": 7.46e-05,
+      "loss": 0.0598,
+      "step": 1120
+    },
+    {
+      "grad_norm": 0.7633928060531616,
+      "learning_rate": 7.526666666666668e-05,
+      "loss": 0.0682,
+      "step": 1130
+    },
+    {
+      "grad_norm": 0.6557093858718872,
+      "learning_rate": 7.593333333333334e-05,
+      "loss": 0.0622,
+      "step": 1140
+    },
+    {
+      "grad_norm": 0.7676181793212891,
+      "learning_rate": 7.66e-05,
+      "loss": 0.0554,
+      "step": 1150
+    },
+    {
+      "grad_norm": 0.49698543548583984,
+      "learning_rate": 7.726666666666667e-05,
+      "loss": 0.0619,
+      "step": 1160
+    },
+    {
+      "grad_norm": 0.658534824848175,
+      "learning_rate": 7.793333333333333e-05,
+      "loss": 0.0628,
+      "step": 1170
+    },
+    {
+      "grad_norm": 0.7802417874336243,
+      "learning_rate": 7.860000000000001e-05,
+      "loss": 0.0606,
+      "step": 1180
+    },
+    {
+      "grad_norm": 0.7362853288650513,
+      "learning_rate": 7.926666666666666e-05,
+      "loss": 0.057,
+      "step": 1190
+    },
+    {
+      "grad_norm": 0.7263256311416626,
+      "learning_rate": 7.993333333333334e-05,
+      "loss": 0.0658,
+      "step": 1200
+    },
+    {
+      "grad_norm": 0.605475902557373,
+      "learning_rate": 8.060000000000001e-05,
+      "loss": 0.0585,
+      "step": 1210
+    },
+    {
+      "grad_norm": 0.6713262796401978,
+      "learning_rate": 8.126666666666667e-05,
+      "loss": 0.0613,
+      "step": 1220
+    },
+    {
+      "grad_norm": 0.9524453282356262,
+      "learning_rate": 8.193333333333334e-05,
+      "loss": 0.0623,
+      "step": 1230
+    },
+    {
+      "grad_norm": 0.7799460887908936,
+      "learning_rate": 8.26e-05,
+      "loss": 0.0637,
+      "step": 1240
+    },
+    {
+      "grad_norm": 0.5529918670654297,
+      "learning_rate": 8.326666666666667e-05,
+      "loss": 0.0582,
+      "step": 1250
+    },
+    {
+      "grad_norm": 0.48169708251953125,
+      "learning_rate": 8.393333333333335e-05,
+      "loss": 0.0613,
+      "step": 1260
+    },
+    {
+      "grad_norm": 0.4488230049610138,
+      "learning_rate": 8.46e-05,
+      "loss": 0.0533,
+      "step": 1270
+    },
+    {
+      "grad_norm": 0.43133044242858887,
+      "learning_rate": 8.526666666666667e-05,
+      "loss": 0.0619,
+      "step": 1280
+    },
+    {
+      "grad_norm": 0.6242606043815613,
+      "learning_rate": 8.593333333333333e-05,
+      "loss": 0.056,
+      "step": 1290
+    },
+    {
+      "grad_norm": 0.48737120628356934,
+      "learning_rate": 8.66e-05,
+      "loss": 0.056,
+      "step": 1300
+    },
+    {
+      "grad_norm": 0.47289788722991943,
+      "learning_rate": 8.726666666666667e-05,
+      "loss": 0.0547,
+      "step": 1310
+    },
+    {
+      "grad_norm": 0.8228473663330078,
+      "learning_rate": 8.793333333333333e-05,
+      "loss": 0.0539,
+      "step": 1320
+    },
+    {
+      "grad_norm": 0.6020798087120056,
+      "learning_rate": 8.86e-05,
+      "loss": 0.0639,
+      "step": 1330
+    },
+    {
+      "grad_norm": 0.6205971837043762,
+      "learning_rate": 8.926666666666668e-05,
+      "loss": 0.0644,
+      "step": 1340
+    },
+    {
+      "grad_norm": 0.5030979514122009,
+      "learning_rate": 8.993333333333334e-05,
+      "loss": 0.0584,
+      "step": 1350
+    },
+    {
+      "grad_norm": 0.7016464471817017,
+      "learning_rate": 9.06e-05,
+      "loss": 0.0559,
+      "step": 1360
+    },
+    {
+      "grad_norm": 0.4915011525154114,
+      "learning_rate": 9.126666666666667e-05,
+      "loss": 0.0565,
+      "step": 1370
+    },
+    {
+      "grad_norm": 0.3934551775455475,
+      "learning_rate": 9.193333333333334e-05,
+      "loss": 0.052,
+      "step": 1380
+    },
+    {
+      "grad_norm": 0.5540144443511963,
+      "learning_rate": 9.260000000000001e-05,
+      "loss": 0.0523,
+      "step": 1390
+    },
+    {
+      "grad_norm": 0.5887146592140198,
+      "learning_rate": 9.326666666666667e-05,
+      "loss": 0.0528,
+      "step": 1400
+    },
+    {
+      "grad_norm": 0.6534768342971802,
+      "learning_rate": 9.393333333333334e-05,
+      "loss": 0.0518,
+      "step": 1410
+    },
+    {
+      "grad_norm": 0.5691227316856384,
+      "learning_rate": 9.46e-05,
+      "loss": 0.0536,
+      "step": 1420
+    },
+    {
+      "grad_norm": 0.3767082691192627,
+      "learning_rate": 9.526666666666667e-05,
+      "loss": 0.0568,
+      "step": 1430
+    },
+    {
+      "grad_norm": 0.47469252347946167,
+      "learning_rate": 9.593333333333334e-05,
+      "loss": 0.0569,
+      "step": 1440
+    },
+    {
+      "grad_norm": 0.5082659125328064,
+      "learning_rate": 9.66e-05,
+      "loss": 0.0549,
+      "step": 1450
+    },
+    {
+      "grad_norm": 0.6236725449562073,
+      "learning_rate": 9.726666666666667e-05,
+      "loss": 0.0636,
+      "step": 1460
+    },
+    {
+      "grad_norm": 0.5344048738479614,
+      "learning_rate": 9.793333333333333e-05,
+      "loss": 0.0559,
+      "step": 1470
+    },
+    {
+      "grad_norm": 0.5537348985671997,
+      "learning_rate": 9.86e-05,
+      "loss": 0.0573,
+      "step": 1480
+    },
+    {
+      "grad_norm": 0.6604304313659668,
+      "learning_rate": 9.926666666666668e-05,
+      "loss": 0.0559,
+      "step": 1490
+    },
+    {
+      "grad_norm": 0.6182398796081543,
+      "learning_rate": 9.993333333333334e-05,
+      "loss": 0.0526,
+      "step": 1500
+    },
+    {
+      "grad_norm": 0.4472421109676361,
+      "learning_rate": 9.999997539434007e-05,
+      "loss": 0.0554,
+      "step": 1510
+    },
+    {
+      "grad_norm": 0.5673554539680481,
+      "learning_rate": 9.999989033776898e-05,
+      "loss": 0.0498,
+      "step": 1520
+    },
+    {
+      "grad_norm": 0.631686806678772,
+      "learning_rate": 9.999974452661641e-05,
+      "loss": 0.0558,
+      "step": 1530
+    },
+    {
+      "grad_norm": 0.5880892276763916,
+      "learning_rate": 9.999953796105959e-05,
+      "loss": 0.0565,
+      "step": 1540
+    },
+    {
+      "grad_norm": 0.386528342962265,
+      "learning_rate": 9.999927064134949e-05,
+      "loss": 0.0546,
+      "step": 1550
+    },
+    {
+      "grad_norm": 0.49979427456855774,
+      "learning_rate": 9.999894256781095e-05,
+      "loss": 0.0558,
+      "step": 1560
+    },
+    {
+      "grad_norm": 0.49290731549263,
+      "learning_rate": 9.99985537408426e-05,
+      "loss": 0.0559,
+      "step": 1570
+    },
+    {
+      "grad_norm": 0.6054269671440125,
+      "learning_rate": 9.999810416091688e-05,
+      "loss": 0.0537,
+      "step": 1580
+    },
+    {
+      "grad_norm": 0.4739155173301697,
+      "learning_rate": 9.99975938285801e-05,
+      "loss": 0.0525,
+      "step": 1590
+    },
+    {
+      "grad_norm": 0.4727555215358734,
+      "learning_rate": 9.999702274445236e-05,
+      "loss": 0.0562,
+      "step": 1600
+    },
+    {
+      "grad_norm": 0.4145797789096832,
+      "learning_rate": 9.999639090922756e-05,
+      "loss": 0.0518,
+      "step": 1610
+    },
+    {
+      "grad_norm": 0.6100338697433472,
+      "learning_rate": 9.999569832367346e-05,
+      "loss": 0.0482,
+      "step": 1620
+    },
+    {
+      "grad_norm": 0.36332374811172485,
+      "learning_rate": 9.999494498863162e-05,
+      "loss": 0.0575,
+      "step": 1630
+    },
+    {
+      "grad_norm": 0.5447846055030823,
+      "learning_rate": 9.99941309050174e-05,
+      "loss": 0.0525,
+      "step": 1640
+    },
+    {
+      "grad_norm": 0.5260499119758606,
+      "learning_rate": 9.999325607381999e-05,
+      "loss": 0.0486,
+      "step": 1650
+    },
+    {
+      "grad_norm": 0.40509524941444397,
+      "learning_rate": 9.999232049610238e-05,
+      "loss": 0.047,
+      "step": 1660
+    },
+    {
+      "grad_norm": 0.5393689870834351,
+      "learning_rate": 9.999132417300142e-05,
+      "loss": 0.0483,
+      "step": 1670
+    },
+    {
+      "grad_norm": 0.45050477981567383,
+      "learning_rate": 9.99902671057277e-05,
+      "loss": 0.0516,
+      "step": 1680
+    },
+    {
+      "grad_norm": 0.513198971748352,
+      "learning_rate": 9.998914929556569e-05,
+      "loss": 0.0564,
+      "step": 1690
+    },
+    {
+      "grad_norm": 0.5798249840736389,
+      "learning_rate": 9.998797074387361e-05,
+      "loss": 0.055,
+      "step": 1700
+    },
+    {
+      "grad_norm": 0.4425696134567261,
+      "learning_rate": 9.99867314520835e-05,
+      "loss": 0.0515,
+      "step": 1710
+    },
+    {
+      "grad_norm": 0.4912015199661255,
+      "learning_rate": 9.998543142170126e-05,
+      "loss": 0.0494,
+      "step": 1720
+    },
+    {
+      "grad_norm": 0.4842802882194519,
+      "learning_rate": 9.99840706543065e-05,
+      "loss": 0.0473,
+      "step": 1730
+    },
+    {
+      "grad_norm": 0.37719348073005676,
+      "learning_rate": 9.998264915155274e-05,
+      "loss": 0.0478,
+      "step": 1740
+    },
+    {
+      "grad_norm": 0.44606560468673706,
+      "learning_rate": 9.998116691516718e-05,
+      "loss": 0.0496,
+      "step": 1750
+    },
+    {
+      "grad_norm": 0.4420314133167267,
+      "learning_rate": 9.997962394695091e-05,
+      "loss": 0.0489,
+      "step": 1760
+    },
+    {
+      "grad_norm": 0.5374616384506226,
+      "learning_rate": 9.997802024877875e-05,
+      "loss": 0.0508,
+      "step": 1770
+    },
+    {
+      "grad_norm": 0.3837665915489197,
+      "learning_rate": 9.99763558225994e-05,
+      "loss": 0.0491,
+      "step": 1780
+    },
+    {
+      "grad_norm": 0.5809544324874878,
+      "learning_rate": 9.997463067043526e-05,
+      "loss": 0.0531,
+      "step": 1790
+    },
+    {
+      "grad_norm": 0.5693655610084534,
+      "learning_rate": 9.997284479438253e-05,
+      "loss": 0.0495,
+      "step": 1800
+    },
+    {
+      "grad_norm": 0.5101096630096436,
+      "learning_rate": 9.997099819661127e-05,
+      "loss": 0.0498,
+      "step": 1810
+    },
+    {
+      "grad_norm": 0.4212728440761566,
+      "learning_rate": 9.996909087936524e-05,
+      "loss": 0.0521,
+      "step": 1820
+    },
+    {
+      "grad_norm": 0.3066025376319885,
+      "learning_rate": 9.996712284496201e-05,
+      "loss": 0.0497,
+      "step": 1830
+    },
+    {
+      "grad_norm": 0.5426850318908691,
+      "learning_rate": 9.996509409579293e-05,
+      "loss": 0.051,
+      "step": 1840
+    },
+    {
+      "grad_norm": 0.4050694704055786,
+      "learning_rate": 9.996300463432312e-05,
+      "loss": 0.053,
+      "step": 1850
+    },
+    {
+      "grad_norm": 0.3489755094051361,
+      "learning_rate": 9.996085446309148e-05,
+      "loss": 0.0505,
+      "step": 1860
+    },
+    {
+      "grad_norm": 0.31605497002601624,
+      "learning_rate": 9.995864358471066e-05,
+      "loss": 0.0474,
+      "step": 1870
+    },
+    {
+      "grad_norm": 0.352236807346344,
+      "learning_rate": 9.99563720018671e-05,
+      "loss": 0.047,
+      "step": 1880
+    },
+    {
+      "grad_norm": 0.4507584571838379,
+      "learning_rate": 9.995403971732098e-05,
+      "loss": 0.0504,
+      "step": 1890
+    },
+    {
+      "grad_norm": 0.42925962805747986,
+      "learning_rate": 9.995164673390625e-05,
+      "loss": 0.0473,
+      "step": 1900
+    },
+    {
+      "grad_norm": 0.353121280670166,
+      "learning_rate": 9.994919305453059e-05,
+      "loss": 0.0464,
+      "step": 1910
+    },
+    {
+      "grad_norm": 0.49948233366012573,
+      "learning_rate": 9.994667868217548e-05,
+      "loss": 0.0479,
+      "step": 1920
+    },
+    {
+      "grad_norm": 0.4217991828918457,
+      "learning_rate": 9.99441036198961e-05,
+      "loss": 0.0482,
+      "step": 1930
+    },
+    {
+      "grad_norm": 0.31471049785614014,
+      "learning_rate": 9.99414678708214e-05,
+      "loss": 0.0481,
+      "step": 1940
+    },
+    {
+      "grad_norm": 0.3589913845062256,
+      "learning_rate": 9.993877143815407e-05,
+      "loss": 0.0445,
+      "step": 1950
+    },
+    {
+      "grad_norm": 0.2932628095149994,
+      "learning_rate": 9.993601432517053e-05,
+      "loss": 0.046,
+      "step": 1960
+    },
+    {
+      "grad_norm": 0.35649144649505615,
+      "learning_rate": 9.993319653522091e-05,
+      "loss": 0.0456,
+      "step": 1970
+    },
+    {
+      "grad_norm": 0.4936084449291229,
+      "learning_rate": 9.993031807172911e-05,
+      "loss": 0.0475,
+      "step": 1980
+    },
+    {
+      "grad_norm": 0.33784282207489014,
+      "learning_rate": 9.992737893819273e-05,
+      "loss": 0.0485,
+      "step": 1990
+    },
+    {
+      "grad_norm": 0.5015108585357666,
+      "learning_rate": 9.992437913818312e-05,
+      "loss": 0.0449,
+      "step": 2000
+    },
+    {
+      "grad_norm": 0.42584389448165894,
+      "learning_rate": 9.992131867534526e-05,
+      "loss": 0.0444,
+      "step": 2010
+    },
+    {
+      "grad_norm": 0.41565531492233276,
+      "learning_rate": 9.991819755339796e-05,
+      "loss": 0.0483,
+      "step": 2020
+    },
+    {
+      "grad_norm": 0.380255788564682,
+      "learning_rate": 9.991501577613365e-05,
+      "loss": 0.0482,
+      "step": 2030
+    },
+    {
+      "grad_norm": 0.37106555700302124,
+      "learning_rate": 9.99117733474185e-05,
+      "loss": 0.0454,
+      "step": 2040
+    },
+    {
+      "grad_norm": 0.43988800048828125,
+      "learning_rate": 9.990847027119234e-05,
+      "loss": 0.0457,
+      "step": 2050
+    },
+    {
+      "grad_norm": 0.9517943263053894,
+      "learning_rate": 9.990510655146877e-05,
+      "loss": 0.0431,
+      "step": 2060
+    },
+    {
+      "grad_norm": 0.4155087172985077,
+      "learning_rate": 9.990168219233496e-05,
+      "loss": 0.0455,
+      "step": 2070
+    },
+    {
+      "grad_norm": 0.37841394543647766,
+      "learning_rate": 9.989819719795188e-05,
+      "loss": 0.0425,
+      "step": 2080
+    },
+    {
+      "grad_norm": 0.45210036635398865,
+      "learning_rate": 9.989465157255412e-05,
+      "loss": 0.047,
+      "step": 2090
+    },
+    {
+      "grad_norm": 0.3879967927932739,
+      "learning_rate": 9.989104532044994e-05,
+      "loss": 0.0448,
+      "step": 2100
+    },
+    {
+      "grad_norm": 0.3749665319919586,
+      "learning_rate": 9.988737844602128e-05,
+      "loss": 0.0408,
+      "step": 2110
+    },
+    {
+      "grad_norm": 0.38606101274490356,
+      "learning_rate": 9.988365095372372e-05,
+      "loss": 0.0422,
+      "step": 2120
+    },
+    {
+      "grad_norm": 0.45897138118743896,
+      "learning_rate": 9.987986284808654e-05,
+      "loss": 0.0458,
+      "step": 2130
+    },
+    {
+      "grad_norm": 0.3423418402671814,
+      "learning_rate": 9.987601413371264e-05,
+      "loss": 0.0425,
+      "step": 2140
+    },
+    {
+      "grad_norm": 0.43687859177589417,
+      "learning_rate": 9.987210481527855e-05,
+      "loss": 0.04,
+      "step": 2150
+    },
+    {
+      "grad_norm": 0.3930586278438568,
+      "learning_rate": 9.98681348975345e-05,
+      "loss": 0.042,
+      "step": 2160
+    },
+    {
+      "grad_norm": 0.4354424774646759,
+      "learning_rate": 9.986410438530427e-05,
+      "loss": 0.0409,
+      "step": 2170
+    },
+    {
+      "grad_norm": 0.4108027219772339,
+      "learning_rate": 9.986001328348534e-05,
+      "loss": 0.0394,
+      "step": 2180
+    },
+    {
+      "grad_norm": 0.32850104570388794,
+      "learning_rate": 9.985586159704878e-05,
+      "loss": 0.0407,
+      "step": 2190
+    },
+    {
+      "grad_norm": 0.45172032713890076,
+      "learning_rate": 9.985164933103929e-05,
+      "loss": 0.0475,
+      "step": 2200
+    },
+    {
+      "grad_norm": 0.30814841389656067,
+      "learning_rate": 9.984737649057513e-05,
+      "loss": 0.0438,
+      "step": 2210
+    },
+    {
+      "grad_norm": 0.43183210492134094,
+      "learning_rate": 9.984304308084827e-05,
+      "loss": 0.0485,
+      "step": 2220
+    },
+    {
+      "grad_norm": 0.4616735279560089,
+      "learning_rate": 9.983864910712416e-05,
+      "loss": 0.0414,
+      "step": 2230
+    },
+    {
+      "grad_norm": 0.36560776829719543,
+      "learning_rate": 9.98341945747419e-05,
+      "loss": 0.0398,
+      "step": 2240
+    },
+    {
+      "grad_norm": 0.43177369236946106,
+      "learning_rate": 9.98296794891142e-05,
+      "loss": 0.0445,
+      "step": 2250
+    },
+    {
+      "grad_norm": 0.4279302656650543,
+      "learning_rate": 9.982510385572725e-05,
+      "loss": 0.0467,
+      "step": 2260
+    },
+    {
+      "grad_norm": 0.40069276094436646,
+      "learning_rate": 9.982046768014094e-05,
+      "loss": 0.0445,
+      "step": 2270
+    },
+    {
+      "grad_norm": 0.4348323345184326,
+      "learning_rate": 9.981577096798863e-05,
+      "loss": 0.0369,
+      "step": 2280
+    },
+    {
+      "grad_norm": 0.41257351636886597,
+      "learning_rate": 9.981101372497727e-05,
+      "loss": 0.0445,
+      "step": 2290
+    },
+    {
+      "grad_norm": 0.4516085982322693,
+      "learning_rate": 9.980619595688737e-05,
+      "loss": 0.0438,
+      "step": 2300
+    },
+    {
+      "grad_norm": 0.30238786339759827,
+      "learning_rate": 9.980131766957295e-05,
+      "loss": 0.0421,
+      "step": 2310
+    },
+    {
+      "grad_norm": 0.47715190052986145,
+      "learning_rate": 9.979637886896163e-05,
+      "loss": 0.0423,
+      "step": 2320
+    },
+    {
+      "grad_norm": 0.29774099588394165,
+      "learning_rate": 9.979137956105447e-05,
+      "loss": 0.0464,
+      "step": 2330
+    },
+    {
+      "grad_norm": 0.29573264718055725,
+      "learning_rate": 9.978631975192613e-05,
+      "loss": 0.0376,
+      "step": 2340
+    },
+    {
+      "grad_norm": 0.32869309186935425,
+      "learning_rate": 9.978119944772475e-05,
+      "loss": 0.0467,
+      "step": 2350
+    },
+    {
+      "grad_norm": 0.3470063805580139,
+      "learning_rate": 9.977601865467197e-05,
+      "loss": 0.0405,
+      "step": 2360
+    },
+    {
+      "grad_norm": 0.3539649546146393,
+      "learning_rate": 9.977077737906297e-05,
+      "loss": 0.0403,
+      "step": 2370
+    },
+    {
+      "grad_norm": 0.3649725914001465,
+      "learning_rate": 9.976547562726636e-05,
+      "loss": 0.0414,
+      "step": 2380
+    },
+    {
+      "grad_norm": 0.32129499316215515,
+      "learning_rate": 9.976011340572429e-05,
+      "loss": 0.0423,
+      "step": 2390
+    },
+    {
+      "grad_norm": 0.2593826651573181,
+      "learning_rate": 9.975469072095237e-05,
+      "loss": 0.039,
+      "step": 2400
+    },
+    {
+      "grad_norm": 0.3159441351890564,
+      "learning_rate": 9.974920757953965e-05,
+      "loss": 0.0383,
+      "step": 2410
+    },
+    {
+      "grad_norm": 0.39688485860824585,
+      "learning_rate": 9.97436639881487e-05,
+      "loss": 0.046,
+      "step": 2420
+    },
+    {
+      "grad_norm": 0.3365003764629364,
+      "learning_rate": 9.973805995351545e-05,
+      "loss": 0.0404,
+      "step": 2430
+    },
+    {
+      "grad_norm": 0.3504849672317505,
+      "learning_rate": 9.973239548244939e-05,
+      "loss": 0.0431,
+      "step": 2440
+    },
+    {
+      "grad_norm": 0.4077778160572052,
+      "learning_rate": 9.972667058183333e-05,
+      "loss": 0.04,
+      "step": 2450
+    },
+    {
+      "grad_norm": 0.3008902966976166,
+      "learning_rate": 9.972088525862362e-05,
+      "loss": 0.0382,
+      "step": 2460
+    },
+    {
+      "grad_norm": 0.37545672059059143,
+      "learning_rate": 9.971503951984995e-05,
+      "loss": 0.0438,
+      "step": 2470
+    },
+    {
+      "grad_norm": 0.36859822273254395,
+      "learning_rate": 9.970913337261543e-05,
+      "loss": 0.0409,
+      "step": 2480
+    },
+    {
+      "grad_norm": 0.3510163128376007,
+      "learning_rate": 9.97031668240966e-05,
+      "loss": 0.0479,
+      "step": 2490
+    },
+    {
+      "grad_norm": 0.44913795590400696,
+      "learning_rate": 9.969713988154339e-05,
+      "loss": 0.0409,
+      "step": 2500
+    },
+    {
+      "grad_norm": 0.46766766905784607,
+      "learning_rate": 9.969105255227906e-05,
+      "loss": 0.0405,
+      "step": 2510
+    },
+    {
+      "grad_norm": 0.45442551374435425,
+      "learning_rate": 9.968490484370035e-05,
+      "loss": 0.0394,
+      "step": 2520
+    },
+    {
+      "grad_norm": 0.35867244005203247,
+      "learning_rate": 9.967869676327726e-05,
+      "loss": 0.0419,
+      "step": 2530
+    },
+    {
+      "grad_norm": 0.3887830674648285,
+      "learning_rate": 9.96724283185532e-05,
+      "loss": 0.0364,
+      "step": 2540
+    },
+    {
+      "grad_norm": 0.5115375518798828,
+      "learning_rate": 9.966609951714494e-05,
+      "loss": 0.0411,
+      "step": 2550
+    },
+    {
+      "grad_norm": 0.3326650559902191,
+      "learning_rate": 9.965971036674255e-05,
+      "loss": 0.0459,
+      "step": 2560
+    },
+    {
+      "grad_norm": 0.28854164481163025,
+      "learning_rate": 9.965326087510947e-05,
+      "loss": 0.0398,
+      "step": 2570
+    },
+    {
+      "grad_norm": 0.4342745244503021,
+      "learning_rate": 9.964675105008244e-05,
+      "loss": 0.0407,
+      "step": 2580
+    },
+    {
+      "grad_norm": 0.3539927899837494,
+      "learning_rate": 9.964018089957147e-05,
+      "loss": 0.0345,
+      "step": 2590
+    },
+    {
+      "grad_norm": 0.3186177611351013,
+      "learning_rate": 9.963355043155996e-05,
+      "loss": 0.0456,
+      "step": 2600
+    },
+    {
+      "grad_norm": 0.49563562870025635,
+      "learning_rate": 9.962685965410455e-05,
+      "loss": 0.0423,
+      "step": 2610
+    },
+    {
+      "grad_norm": 0.33249884843826294,
+      "learning_rate": 9.962010857533514e-05,
+      "loss": 0.0373,
+      "step": 2620
+    },
+    {
+      "grad_norm": 0.5187269449234009,
+      "learning_rate": 9.961329720345493e-05,
+      "loss": 0.0386,
+      "step": 2630
+    },
+    {
+      "grad_norm": 0.510249137878418,
+      "learning_rate": 9.96064255467404e-05,
+      "loss": 0.0428,
+      "step": 2640
+    },
+    {
+      "grad_norm": 0.35589292645454407,
+      "learning_rate": 9.959949361354126e-05,
+      "loss": 0.0392,
+      "step": 2650
+    },
+    {
+      "grad_norm": 0.48643842339515686,
+      "learning_rate": 9.959250141228045e-05,
+      "loss": 0.0406,
+      "step": 2660
+    },
+    {
+      "grad_norm": 0.4539337158203125,
+      "learning_rate": 9.958544895145414e-05,
+      "loss": 0.0404,
+      "step": 2670
+    },
+    {
+      "grad_norm": 0.29236212372779846,
+      "learning_rate": 9.957833623963177e-05,
+      "loss": 0.041,
+      "step": 2680
+    },
+    {
+      "grad_norm": 0.5335069894790649,
+      "learning_rate": 9.957116328545593e-05,
+      "loss": 0.0424,
+      "step": 2690
+    },
+    {
+      "grad_norm": 0.3695768117904663,
+      "learning_rate": 9.956393009764244e-05,
+      "loss": 0.0388,
+      "step": 2700
+    },
+    {
+      "grad_norm": 0.36372533440589905,
+      "learning_rate": 9.955663668498032e-05,
+      "loss": 0.0385,
+      "step": 2710
+    },
+    {
+      "grad_norm": 0.4580289125442505,
+      "learning_rate": 9.954928305633173e-05,
+      "loss": 0.0389,
+      "step": 2720
+    },
+    {
+      "grad_norm": 0.43339890241622925,
+      "learning_rate": 9.954186922063204e-05,
+      "loss": 0.0399,
+      "step": 2730
+    },
+    {
+      "grad_norm": 0.35029280185699463,
+      "learning_rate": 9.953439518688974e-05,
+      "loss": 0.0377,
+      "step": 2740
+    },
+    {
+      "grad_norm": 0.33518457412719727,
+      "learning_rate": 9.952686096418652e-05,
+      "loss": 0.0355,
+      "step": 2750
+    },
+    {
+      "grad_norm": 0.30202507972717285,
+      "learning_rate": 9.951926656167714e-05,
+      "loss": 0.0372,
+      "step": 2760
+    },
+    {
+      "grad_norm": 0.34633857011795044,
+      "learning_rate": 9.951161198858953e-05,
+      "loss": 0.0412,
+      "step": 2770
+    },
+    {
+      "grad_norm": 0.40867945551872253,
+      "learning_rate": 9.950389725422471e-05,
+      "loss": 0.0372,
+      "step": 2780
+    },
+    {
+      "grad_norm": 0.2682510316371918,
+      "learning_rate": 9.949612236795682e-05,
+      "loss": 0.0365,
+      "step": 2790
+    },
+    {
+      "grad_norm": 0.4570905566215515,
+      "learning_rate": 9.948828733923305e-05,
+      "loss": 0.0393,
+      "step": 2800
+    },
+    {
+      "grad_norm": 0.36744511127471924,
+      "learning_rate": 9.948039217757374e-05,
+      "loss": 0.0412,
+      "step": 2810
+    },
+    {
+      "grad_norm": 0.29317381978034973,
+      "learning_rate": 9.947243689257225e-05,
+      "loss": 0.038,
+      "step": 2820
+    },
+    {
+      "grad_norm": 0.3504505455493927,
+      "learning_rate": 9.946442149389497e-05,
+      "loss": 0.0348,
+      "step": 2830
+    },
+    {
+      "grad_norm": 0.31019940972328186,
+      "learning_rate": 9.945634599128139e-05,
+      "loss": 0.038,
+      "step": 2840
+    },
+    {
+      "grad_norm": 0.437590092420578,
+      "learning_rate": 9.944821039454402e-05,
+      "loss": 0.0387,
+      "step": 2850
+    },
+    {
+      "grad_norm": 0.36887866258621216,
+      "learning_rate": 9.944001471356835e-05,
+      "loss": 0.0405,
+      "step": 2860
+    },
+    {
+      "grad_norm": 0.2392815351486206,
+      "learning_rate": 9.94317589583129e-05,
+      "loss": 0.0435,
+      "step": 2870
+    },
+    {
+      "grad_norm": 0.21329613029956818,
+      "learning_rate": 9.942344313880922e-05,
+      "loss": 0.0365,
+      "step": 2880
+    },
+    {
+      "grad_norm": 0.3033916652202606,
+      "learning_rate": 9.941506726516179e-05,
+      "loss": 0.0385,
+      "step": 2890
+    },
+    {
+      "grad_norm": 0.2520839273929596,
+      "learning_rate": 9.94066313475481e-05,
+      "loss": 0.0363,
+      "step": 2900
+    },
+    {
+      "grad_norm": 0.43853965401649475,
+      "learning_rate": 9.939813539621857e-05,
+      "loss": 0.0363,
+      "step": 2910
+    },
+    {
+      "grad_norm": 0.36329326033592224,
+      "learning_rate": 9.93895794214966e-05,
+      "loss": 0.0402,
+      "step": 2920
+    },
+    {
+      "grad_norm": 0.26854604482650757,
+      "learning_rate": 9.938096343377852e-05,
+      "loss": 0.0459,
+      "step": 2930
+    },
+    {
+      "grad_norm": 0.44420111179351807,
+      "learning_rate": 9.937228744353353e-05,
+      "loss": 0.0432,
+      "step": 2940
+    },
+    {
+      "grad_norm": 0.363777756690979,
+      "learning_rate": 9.936355146130379e-05,
+      "loss": 0.0369,
+      "step": 2950
+    },
+    {
+      "grad_norm": 0.2704271972179413,
+      "learning_rate": 9.935475549770435e-05,
+      "loss": 0.0379,
+      "step": 2960
+    },
+    {
+      "grad_norm": 0.298907071352005,
+      "learning_rate": 9.934589956342315e-05,
+      "loss": 0.0356,
+      "step": 2970
+    },
+    {
+      "grad_norm": 0.31564566493034363,
+      "learning_rate": 9.933698366922093e-05,
+      "loss": 0.0339,
+      "step": 2980
+    },
+    {
+      "grad_norm": 0.36585021018981934,
+      "learning_rate": 9.93280078259314e-05,
+      "loss": 0.0373,
+      "step": 2990
+    },
+    {
+      "grad_norm": 0.33079367876052856,
+      "learning_rate": 9.931897204446104e-05,
+      "loss": 0.0372,
+      "step": 3000
+    },
+    {
+      "grad_norm": 0.40933725237846375,
+      "learning_rate": 9.930987633578915e-05,
+      "loss": 0.0366,
+      "step": 3010
+    },
+    {
+      "grad_norm": 0.4112361967563629,
+      "learning_rate": 9.93007207109679e-05,
+      "loss": 0.037,
+      "step": 3020
+    },
+    {
+      "grad_norm": 0.4142875075340271,
+      "learning_rate": 9.929150518112224e-05,
+      "loss": 0.0351,
+      "step": 3030
+    },
+    {
+      "grad_norm": 0.34691959619522095,
+      "learning_rate": 9.928222975744991e-05,
+      "loss": 0.0387,
+      "step": 3040
+    },
+    {
+      "grad_norm": 0.2551492750644684,
+      "learning_rate": 9.92728944512214e-05,
+      "loss": 0.0334,
+      "step": 3050
+    },
+    {
+      "grad_norm": 0.29876184463500977,
+      "learning_rate": 9.926349927378001e-05,
+      "loss": 0.0344,
+      "step": 3060
+    },
+    {
+      "grad_norm": 0.342134565114975,
+      "learning_rate": 9.925404423654174e-05,
+      "loss": 0.0347,
+      "step": 3070
+    },
+    {
+      "grad_norm": 0.37610647082328796,
+      "learning_rate": 9.924452935099537e-05,
+      "loss": 0.0355,
+      "step": 3080
+    },
+    {
+      "grad_norm": 0.28914567828178406,
+      "learning_rate": 9.92349546287024e-05,
+      "loss": 0.0377,
+      "step": 3090
+    },
+    {
+      "grad_norm": 0.2716347277164459,
+      "learning_rate": 9.9225320081297e-05,
+      "loss": 0.0381,
+      "step": 3100
+    },
+    {
+      "grad_norm": 0.2830624580383301,
+      "learning_rate": 9.921562572048606e-05,
+      "loss": 0.038,
+      "step": 3110
+    },
+    {
+      "grad_norm": 0.4740348756313324,
+      "learning_rate": 9.920587155804913e-05,
+      "loss": 0.0377,
+      "step": 3120
+    },
+    {
+      "grad_norm": 0.3273002505302429,
+      "learning_rate": 9.919605760583845e-05,
+      "loss": 0.0393,
+      "step": 3130
+    },
+    {
+      "grad_norm": 0.2517211139202118,
+      "learning_rate": 9.91861838757789e-05,
+      "loss": 0.0371,
+      "step": 3140
+    },
+    {
+      "grad_norm": 0.267072856426239,
+      "learning_rate": 9.917625037986798e-05,
+      "loss": 0.0372,
+      "step": 3150
+    },
+    {
+      "grad_norm": 0.33644410967826843,
+      "learning_rate": 9.916625713017583e-05,
+      "loss": 0.0359,
+      "step": 3160
+    },
+    {
+      "grad_norm": 0.3924078643321991,
+      "learning_rate": 9.915620413884519e-05,
+      "loss": 0.0368,
+      "step": 3170
+    },
+    {
+      "grad_norm": 0.26549962162971497,
+      "learning_rate": 9.914609141809139e-05,
+      "loss": 0.0347,
+      "step": 3180
+    },
+    {
+      "grad_norm": 0.3926424980163574,
+      "learning_rate": 9.913591898020235e-05,
+      "loss": 0.0355,
+      "step": 3190
+    },
+    {
+      "grad_norm": 0.23619306087493896,
+      "learning_rate": 9.912568683753853e-05,
+      "loss": 0.0375,
+      "step": 3200
+    },
+    {
+      "grad_norm": 0.342897891998291,
+      "learning_rate": 9.911539500253295e-05,
+      "loss": 0.036,
+      "step": 3210
+    },
+    {
+      "grad_norm": 0.2978026866912842,
+      "learning_rate": 9.910504348769118e-05,
+      "loss": 0.0346,
+      "step": 3220
+    },
+    {
+      "grad_norm": 0.3328244090080261,
+      "learning_rate": 9.909463230559127e-05,
+      "loss": 0.0316,
+      "step": 3230
+    },
+    {
+      "grad_norm": 0.2802560031414032,
+      "learning_rate": 9.908416146888376e-05,
+      "loss": 0.0342,
+      "step": 3240
+    },
+    {
+      "grad_norm": 0.34282174706459045,
+      "learning_rate": 9.907363099029175e-05,
+      "loss": 0.0332,
+      "step": 3250
+    },
+    {
+      "grad_norm": 0.40915417671203613,
+      "learning_rate": 9.906304088261073e-05,
+      "loss": 0.0398,
+      "step": 3260
+    },
+    {
+      "grad_norm": 0.25955238938331604,
+      "learning_rate": 9.905239115870872e-05,
+      "loss": 0.0362,
+      "step": 3270
+    },
+    {
+      "grad_norm": 0.333028107881546,
+      "learning_rate": 9.90416818315261e-05,
+      "loss": 0.0369,
+      "step": 3280
+    },
+    {
+      "grad_norm": 0.28226780891418457,
+      "learning_rate": 9.903091291407573e-05,
+      "loss": 0.0439,
+      "step": 3290
+    },
+    {
+      "grad_norm": 0.31488144397735596,
+      "learning_rate": 9.902008441944286e-05,
+      "loss": 0.0358,
+      "step": 3300
+    },
+    {
+      "grad_norm": 0.26195263862609863,
+      "learning_rate": 9.900919636078512e-05,
+      "loss": 0.0362,
+      "step": 3310
+    },
+    {
+      "grad_norm": 0.29751327633857727,
+      "learning_rate": 9.899824875133255e-05,
+      "loss": 0.037,
+      "step": 3320
+    },
+    {
+      "grad_norm": 0.4387319087982178,
+      "learning_rate": 9.898724160438749e-05,
+      "loss": 0.0349,
+      "step": 3330
+    },
+    {
+      "grad_norm": 0.2695530652999878,
+      "learning_rate": 9.89761749333247e-05,
+      "loss": 0.0376,
+      "step": 3340
+    },
+    {
+      "grad_norm": 0.36425113677978516,
+      "learning_rate": 9.896504875159122e-05,
+      "loss": 0.0361,
+      "step": 3350
+    },
+    {
+      "grad_norm": 0.29681363701820374,
+      "learning_rate": 9.89538630727064e-05,
+      "loss": 0.0359,
+      "step": 3360
+    },
+    {
+      "grad_norm": 0.2756250500679016,
+      "learning_rate": 9.894261791026189e-05,
+      "loss": 0.0402,
+      "step": 3370
+    },
+    {
+      "grad_norm": 0.3387982249259949,
+      "learning_rate": 9.893131327792165e-05,
+      "loss": 0.0335,
+      "step": 3380
+    },
+    {
+      "grad_norm": 0.3108982741832733,
+      "learning_rate": 9.891994918942182e-05,
+      "loss": 0.0331,
+      "step": 3390
+    },
+    {
+      "grad_norm": 0.4296475052833557,
+      "learning_rate": 9.890852565857092e-05,
+      "loss": 0.038,
+      "step": 3400
+    },
+    {
+      "grad_norm": 0.25025734305381775,
+      "learning_rate": 9.889704269924954e-05,
+      "loss": 0.0304,
+      "step": 3410
+    },
+    {
+      "grad_norm": 0.6248367428779602,
+      "learning_rate": 9.888550032541059e-05,
+      "loss": 0.0316,
+      "step": 3420
+    },
+    {
+      "grad_norm": 0.3411543667316437,
+      "learning_rate": 9.887389855107916e-05,
+      "loss": 0.0353,
+      "step": 3430
+    },
+    {
+      "grad_norm": 0.3803151845932007,
+      "learning_rate": 9.886223739035248e-05,
+      "loss": 0.0399,
+      "step": 3440
+    },
+    {
+      "grad_norm": 0.3013177812099457,
+      "learning_rate": 9.885051685739997e-05,
+      "loss": 0.0382,
+      "step": 3450
+    },
+    {
+      "grad_norm": 0.4035526216030121,
+      "learning_rate": 9.883873696646316e-05,
+      "loss": 0.0423,
+      "step": 3460
+    },
+    {
+      "grad_norm": 0.39729955792427063,
+      "learning_rate": 9.882689773185575e-05,
+      "loss": 0.0342,
+      "step": 3470
+    },
+    {
+      "grad_norm": 0.35385772585868835,
+      "learning_rate": 9.881499916796353e-05,
+      "loss": 0.0387,
+      "step": 3480
+    },
+    {
+      "grad_norm": 0.38120993971824646,
+      "learning_rate": 9.880304128924434e-05,
+      "loss": 0.0346,
+      "step": 3490
+    },
+    {
+      "grad_norm": 0.5219783782958984,
+      "learning_rate": 9.879102411022817e-05,
+      "loss": 0.0359,
+      "step": 3500
+    },
+    {
+      "grad_norm": 0.25719115138053894,
+      "learning_rate": 9.877894764551703e-05,
+      "loss": 0.0347,
+      "step": 3510
+    },
+    {
+      "grad_norm": 0.27055251598358154,
+      "learning_rate": 9.876681190978494e-05,
+      "loss": 0.0362,
+      "step": 3520
+    },
+    {
+      "grad_norm": 0.2928553521633148,
+      "learning_rate": 9.875461691777797e-05,
+      "loss": 0.0357,
+      "step": 3530
+    },
+    {
+      "grad_norm": 0.3087640702724457,
+      "learning_rate": 9.874236268431417e-05,
+      "loss": 0.0341,
+      "step": 3540
+    },
+    {
+      "grad_norm": 0.30907273292541504,
+      "learning_rate": 9.873004922428361e-05,
+      "loss": 0.0326,
+      "step": 3550
+    },
+    {
+      "grad_norm": 0.24913908541202545,
+      "learning_rate": 9.871767655264829e-05,
+      "loss": 0.0341,
+      "step": 3560
+    },
+    {
+      "grad_norm": 0.3109101355075836,
+      "learning_rate": 9.87052446844422e-05,
+      "loss": 0.036,
+      "step": 3570
+    },
+    {
+      "grad_norm": 0.2711428701877594,
+      "learning_rate": 9.869275363477122e-05,
+      "loss": 0.0389,
+      "step": 3580
+    },
+    {
+      "grad_norm": 0.41108351945877075,
+      "learning_rate": 9.868020341881312e-05,
+      "loss": 0.0378,
+      "step": 3590
+    },
+    {
+      "grad_norm": 0.23758524656295776,
+      "learning_rate": 9.866759405181765e-05,
+      "loss": 0.0356,
+      "step": 3600
+    },
+    {
+      "grad_norm": 0.27054616808891296,
+      "learning_rate": 9.865492554910633e-05,
+      "loss": 0.0363,
+      "step": 3610
+    },
+    {
+      "grad_norm": 0.32093489170074463,
+      "learning_rate": 9.864219792607262e-05,
+      "loss": 0.0319,
+      "step": 3620
+    },
+    {
+      "grad_norm": 0.28849124908447266,
+      "learning_rate": 9.862941119818177e-05,
+      "loss": 0.0351,
+      "step": 3630
+    },
+    {
+      "grad_norm": 0.3146919906139374,
+      "learning_rate": 9.861656538097086e-05,
+      "loss": 0.0335,
+      "step": 3640
+    },
+    {
+      "grad_norm": 0.2292512208223343,
+      "learning_rate": 9.860366049004877e-05,
+      "loss": 0.0322,
+      "step": 3650
+    },
+    {
+      "grad_norm": 0.34268423914909363,
+      "learning_rate": 9.859069654109615e-05,
+      "loss": 0.0365,
+      "step": 3660
+    },
+    {
+      "grad_norm": 0.29794055223464966,
+      "learning_rate": 9.857767354986545e-05,
+      "loss": 0.0386,
+      "step": 3670
+    },
+    {
+      "grad_norm": 0.3198016881942749,
+      "learning_rate": 9.856459153218078e-05,
+      "loss": 0.0345,
+      "step": 3680
+    },
+    {
+      "grad_norm": 0.27674734592437744,
+      "learning_rate": 9.855145050393808e-05,
+      "loss": 0.0329,
+      "step": 3690
+    },
+    {
+      "grad_norm": 0.30881065130233765,
+      "learning_rate": 9.85382504811049e-05,
+      "loss": 0.0344,
+      "step": 3700
+    },
+    {
+      "grad_norm": 0.26281753182411194,
+      "learning_rate": 9.852499147972054e-05,
+      "loss": 0.0349,
+      "step": 3710
+    },
+    {
+      "grad_norm": 0.25280869007110596,
+      "learning_rate": 9.851167351589592e-05,
+      "loss": 0.0363,
+      "step": 3720
+    },
+    {
+      "grad_norm": 0.4011027216911316,
+      "learning_rate": 9.849829660581363e-05,
+      "loss": 0.035,
+      "step": 3730
+    },
+    {
+      "grad_norm": 0.2938372492790222,
+      "learning_rate": 9.848486076572787e-05,
+      "loss": 0.0363,
+      "step": 3740
+    },
+    {
+      "grad_norm": 0.4511852562427521,
+      "learning_rate": 9.847136601196446e-05,
+      "loss": 0.0346,
+      "step": 3750
+    },
+    {
+      "grad_norm": 0.355217844247818,
+      "learning_rate": 9.845781236092078e-05,
+      "loss": 0.0348,
+      "step": 3760
+    },
+    {
+      "grad_norm": 0.3419903516769409,
+      "learning_rate": 9.844419982906583e-05,
+      "loss": 0.0333,
+      "step": 3770
+    },
+    {
+      "grad_norm": 0.5065255761146545,
+      "learning_rate": 9.843052843294008e-05,
+      "loss": 0.031,
+      "step": 3780
+    },
+    {
+      "grad_norm": 0.38208192586898804,
+      "learning_rate": 9.841679818915559e-05,
+      "loss": 0.0345,
+      "step": 3790
+    },
+    {
+      "grad_norm": 0.25567400455474854,
+      "learning_rate": 9.840300911439591e-05,
+      "loss": 0.034,
+      "step": 3800
+    },
+    {
+      "grad_norm": 0.3119286596775055,
+      "learning_rate": 9.838916122541603e-05,
+      "loss": 0.0331,
+      "step": 3810
+    },
+    {
+      "grad_norm": 0.27168720960617065,
+      "learning_rate": 9.837525453904246e-05,
+      "loss": 0.0339,
+      "step": 3820
+    },
+    {
+      "grad_norm": 0.26491183042526245,
+      "learning_rate": 9.836128907217314e-05,
+      "loss": 0.0352,
+      "step": 3830
+    },
+    {
+      "grad_norm": 0.2985789179801941,
+      "learning_rate": 9.834726484177743e-05,
+      "loss": 0.0317,
+      "step": 3840
+    },
+    {
+      "grad_norm": 0.2767351567745209,
+      "learning_rate": 9.833318186489609e-05,
+      "loss": 0.0335,
+      "step": 3850
+    },
+    {
+      "grad_norm": 0.3301824927330017,
+      "learning_rate": 9.831904015864126e-05,
+      "loss": 0.0322,
+      "step": 3860
+    },
+    {
+      "grad_norm": 0.45963019132614136,
+      "learning_rate": 9.830483974019645e-05,
+      "loss": 0.0336,
+      "step": 3870
+    },
+    {
+      "grad_norm": 0.30730369687080383,
+      "learning_rate": 9.82905806268165e-05,
+      "loss": 0.0327,
+      "step": 3880
+    },
+    {
+      "grad_norm": 0.2633878290653229,
+      "learning_rate": 9.82762628358276e-05,
+      "loss": 0.031,
+      "step": 3890
+    },
+    {
+      "grad_norm": 0.2650753855705261,
+      "learning_rate": 9.826188638462718e-05,
+      "loss": 0.0327,
+      "step": 3900
+    },
+    {
+      "grad_norm": 0.2208768129348755,
+      "learning_rate": 9.824745129068402e-05,
+      "loss": 0.0352,
+      "step": 3910
+    },
+    {
+      "grad_norm": 0.3218386471271515,
+      "learning_rate": 9.82329575715381e-05,
+      "loss": 0.0382,
+      "step": 3920
+    },
+    {
+      "grad_norm": 0.25913795828819275,
+      "learning_rate": 9.821840524480066e-05,
+      "loss": 0.0317,
+      "step": 3930
+    },
+    {
+      "grad_norm": 0.29261311888694763,
+      "learning_rate": 9.820379432815414e-05,
+      "loss": 0.0315,
+      "step": 3940
+    },
+    {
+      "grad_norm": 0.330141544342041,
+      "learning_rate": 9.81891248393522e-05,
+      "loss": 0.032,
+      "step": 3950
+    },
+    {
+      "grad_norm": 0.44203680753707886,
+      "learning_rate": 9.817439679621963e-05,
+      "loss": 0.0347,
+      "step": 3960
+    },
+    {
+      "grad_norm": 0.31121841073036194,
+      "learning_rate": 9.815961021665243e-05,
+      "loss": 0.0356,
+      "step": 3970
+    },
+    {
+      "grad_norm": 0.4357168674468994,
+      "learning_rate": 9.814476511861763e-05,
+      "loss": 0.0422,
+      "step": 3980
+    },
+    {
+      "grad_norm": 0.39395856857299805,
+      "learning_rate": 9.812986152015348e-05,
+      "loss": 0.036,
+      "step": 3990
+    },
+    {
+      "grad_norm": 0.24850888550281525,
+      "learning_rate": 9.811489943936922e-05,
+      "loss": 0.0334,
+      "step": 4000
+    },
+    {
+      "grad_norm": 0.294154554605484,
+      "learning_rate": 9.809987889444522e-05,
+      "loss": 0.0349,
+      "step": 4010
+    },
+    {
+      "grad_norm": 0.2648806571960449,
+      "learning_rate": 9.808479990363282e-05,
+      "loss": 0.0338,
+      "step": 4020
+    },
+    {
+      "grad_norm": 0.24719461798667908,
+      "learning_rate": 9.806966248525445e-05,
+      "loss": 0.033,
+      "step": 4030
+    },
+    {
+      "grad_norm": 0.2546530067920685,
+      "learning_rate": 9.805446665770348e-05,
+      "loss": 0.0306,
+      "step": 4040
+    },
+    {
+      "grad_norm": 0.25734084844589233,
+      "learning_rate": 9.803921243944429e-05,
+      "loss": 0.0279,
+      "step": 4050
+    },
+    {
+      "grad_norm": 0.27421101927757263,
+      "learning_rate": 9.802389984901218e-05,
+      "loss": 0.0336,
+      "step": 4060
+    },
+    {
+      "grad_norm": 0.2428167313337326,
+      "learning_rate": 9.80085289050134e-05,
+      "loss": 0.0325,
+      "step": 4070
+    },
+    {
+      "grad_norm": 0.3550518751144409,
+      "learning_rate": 9.799309962612508e-05,
+      "loss": 0.0336,
+      "step": 4080
+    },
+    {
+      "grad_norm": 0.29807665944099426,
+      "learning_rate": 9.797761203109527e-05,
+      "loss": 0.0346,
+      "step": 4090
+    },
+    {
+      "grad_norm": 0.4399183392524719,
+      "learning_rate": 9.796206613874283e-05,
+      "loss": 0.0342,
+      "step": 4100
+    },
+    {
+      "grad_norm": 0.25817739963531494,
+      "learning_rate": 9.794646196795754e-05,
+      "loss": 0.0337,
+      "step": 4110
+    },
+    {
+      "grad_norm": 0.5267719626426697,
+      "learning_rate": 9.793079953769987e-05,
+      "loss": 0.031,
+      "step": 4120
+    },
+    {
+      "grad_norm": 0.3641016185283661,
+      "learning_rate": 9.79150788670012e-05,
+      "loss": 0.0327,
+      "step": 4130
+    },
+    {
+      "grad_norm": 0.19937771558761597,
+      "learning_rate": 9.78992999749636e-05,
+      "loss": 0.0358,
+      "step": 4140
+    },
+    {
+      "grad_norm": 0.33499547839164734,
+      "learning_rate": 9.788346288075994e-05,
+      "loss": 0.031,
+      "step": 4150
+    },
+    {
+      "grad_norm": 0.36649659276008606,
+      "learning_rate": 9.786756760363373e-05,
+      "loss": 0.0367,
+      "step": 4160
+    },
+    {
+      "grad_norm": 0.2832200229167938,
+      "learning_rate": 9.78516141628993e-05,
+      "loss": 0.032,
+      "step": 4170
+    },
+    {
+      "grad_norm": 0.27441516518592834,
+      "learning_rate": 9.783560257794154e-05,
+      "loss": 0.0327,
+      "step": 4180
+    },
+    {
+      "grad_norm": 0.2544417977333069,
+      "learning_rate": 9.781953286821603e-05,
+      "loss": 0.0302,
+      "step": 4190
+    },
+    {
+      "grad_norm": 0.28273317217826843,
+      "learning_rate": 9.780340505324901e-05,
+      "loss": 0.0303,
+      "step": 4200
+    },
+    {
+      "grad_norm": 0.4371950328350067,
+      "learning_rate": 9.778721915263727e-05,
+      "loss": 0.0311,
+      "step": 4210
+    },
+    {
+      "grad_norm": 0.31068944931030273,
+      "learning_rate": 9.777097518604824e-05,
+      "loss": 0.0302,
+      "step": 4220
+    },
+    {
+      "grad_norm": 0.21536166965961456,
+      "learning_rate": 9.775467317321984e-05,
+      "loss": 0.0306,
+      "step": 4230
+    },
+    {
+      "grad_norm": 0.2604850232601166,
+      "learning_rate": 9.773831313396055e-05,
+      "loss": 0.0331,
+      "step": 4240
+    },
+    {
+      "grad_norm": 0.32531386613845825,
+      "learning_rate": 9.77218950881494e-05,
+      "loss": 0.0333,
+      "step": 4250
+    },
+    {
+      "grad_norm": 0.2731362283229828,
+      "learning_rate": 9.770541905573583e-05,
+      "loss": 0.0341,
+      "step": 4260
+    },
+    {
+      "grad_norm": 0.2093501091003418,
+      "learning_rate": 9.768888505673976e-05,
+      "loss": 0.029,
+      "step": 4270
+    },
+    {
+      "grad_norm": 0.42711615562438965,
+      "learning_rate": 9.767229311125162e-05,
+      "loss": 0.032,
+      "step": 4280
+    },
+    {
+      "grad_norm": 0.27278879284858704,
+      "learning_rate": 9.765564323943211e-05,
+      "loss": 0.033,
+      "step": 4290
+    },
+    {
+      "grad_norm": 0.2229641228914261,
+      "learning_rate": 9.763893546151244e-05,
+      "loss": 0.0364,
+      "step": 4300
+    },
+    {
+      "grad_norm": 0.3014525771141052,
+      "learning_rate": 9.762216979779412e-05,
+      "loss": 0.0361,
+      "step": 4310
+    },
+    {
+      "grad_norm": 0.24539710581302643,
+      "learning_rate": 9.760534626864902e-05,
+      "loss": 0.0325,
+      "step": 4320
+    },
+    {
+      "grad_norm": 0.2680601477622986,
+      "learning_rate": 9.758846489451931e-05,
+      "loss": 0.0325,
+      "step": 4330
+    },
+    {
+      "grad_norm": 0.3267567455768585,
+      "learning_rate": 9.757152569591748e-05,
+      "loss": 0.0337,
+      "step": 4340
+    },
+    {
+      "grad_norm": 0.25609248876571655,
+      "learning_rate": 9.75545286934262e-05,
+      "loss": 0.0316,
+      "step": 4350
+    },
+    {
+      "grad_norm": 0.3844185471534729,
+      "learning_rate": 9.753747390769847e-05,
+      "loss": 0.03,
+      "step": 4360
+    },
+    {
+      "grad_norm": 0.2432142198085785,
+      "learning_rate": 9.752036135945744e-05,
+      "loss": 0.0332,
+      "step": 4370
+    },
+    {
+      "grad_norm": 0.48657962679862976,
+      "learning_rate": 9.750319106949649e-05,
+      "loss": 0.0357,
+      "step": 4380
+    },
+    {
+      "grad_norm": 0.2904045581817627,
+      "learning_rate": 9.748596305867913e-05,
+      "loss": 0.0337,
+      "step": 4390
+    },
+    {
+      "grad_norm": 0.3355681598186493,
+      "learning_rate": 9.746867734793903e-05,
+      "loss": 0.0351,
+      "step": 4400
+    },
+    {
+      "grad_norm": 0.2996969223022461,
+      "learning_rate": 9.745133395827993e-05,
+      "loss": 0.0342,
+      "step": 4410
+    },
+    {
+      "grad_norm": 0.2818157374858856,
+      "learning_rate": 9.743393291077572e-05,
+      "loss": 0.0306,
+      "step": 4420
+    },
+    {
+      "grad_norm": 0.4198729395866394,
+      "learning_rate": 9.741647422657028e-05,
+      "loss": 0.0357,
+      "step": 4430
+    },
+    {
+      "grad_norm": 0.6452091336250305,
+      "learning_rate": 9.739895792687758e-05,
+      "loss": 0.0328,
+      "step": 4440
+    },
+    {
+      "grad_norm": 0.32312220335006714,
+      "learning_rate": 9.738138403298157e-05,
+      "loss": 0.0344,
+      "step": 4450
+    },
+    {
+      "grad_norm": 0.3436553478240967,
+      "learning_rate": 9.736375256623619e-05,
+      "loss": 0.0308,
+      "step": 4460
+    },
+    {
+      "grad_norm": 0.25958141684532166,
+      "learning_rate": 9.734606354806533e-05,
+      "loss": 0.0318,
+      "step": 4470
+    },
+    {
+      "grad_norm": 0.3319542706012726,
+      "learning_rate": 9.73283169999628e-05,
+      "loss": 0.0338,
+      "step": 4480
+    },
+    {
+      "grad_norm": 0.26436764001846313,
+      "learning_rate": 9.731051294349238e-05,
+      "loss": 0.0327,
+      "step": 4490
+    },
+    {
+      "grad_norm": 0.2518230080604553,
+      "learning_rate": 9.729265140028762e-05,
+      "loss": 0.0341,
+      "step": 4500
+    },
+    {
+      "grad_norm": 0.3020409047603607,
+      "learning_rate": 9.727473239205201e-05,
+      "loss": 0.0294,
+      "step": 4510
+    },
+    {
+      "grad_norm": 0.6439887285232544,
+      "learning_rate": 9.725675594055883e-05,
+      "loss": 0.0302,
+      "step": 4520
+    },
+    {
+      "grad_norm": 0.2171560525894165,
+      "learning_rate": 9.723872206765116e-05,
+      "loss": 0.0329,
+      "step": 4530
+    },
+    {
+      "grad_norm": 0.2762564420700073,
+      "learning_rate": 9.722063079524185e-05,
+      "loss": 0.0318,
+      "step": 4540
+    },
+    {
+      "grad_norm": 0.32399308681488037,
+      "learning_rate": 9.720248214531351e-05,
+      "loss": 0.0318,
+      "step": 4550
+    },
+    {
+      "grad_norm": 0.21756646037101746,
+      "learning_rate": 9.718427613991848e-05,
+      "loss": 0.032,
+      "step": 4560
+    },
+    {
+      "grad_norm": 0.3138534724712372,
+      "learning_rate": 9.716601280117873e-05,
+      "loss": 0.029,
+      "step": 4570
+    },
+    {
+      "grad_norm": 0.2731068730354309,
+      "learning_rate": 9.714769215128596e-05,
+      "loss": 0.0319,
+      "step": 4580
+    },
+    {
+      "grad_norm": 0.3360292613506317,
+      "learning_rate": 9.712931421250152e-05,
+      "loss": 0.0355,
+      "step": 4590
+    },
+    {
+      "grad_norm": 0.21750696003437042,
+      "learning_rate": 9.711087900715627e-05,
+      "loss": 0.0346,
+      "step": 4600
+    },
+    {
+      "grad_norm": 0.22192755341529846,
+      "learning_rate": 9.709238655765078e-05,
+      "loss": 0.0325,
+      "step": 4610
+    },
+    {
+      "grad_norm": 0.21116146445274353,
+      "learning_rate": 9.707383688645511e-05,
+      "loss": 0.0301,
+      "step": 4620
+    },
+    {
+      "grad_norm": 0.28577712178230286,
+      "learning_rate": 9.705523001610883e-05,
+      "loss": 0.0325,
+      "step": 4630
+    },
+    {
+      "grad_norm": 0.2641943097114563,
+      "learning_rate": 9.703656596922107e-05,
+      "loss": 0.0291,
+      "step": 4640
+    },
+    {
+      "grad_norm": 0.23968888819217682,
+      "learning_rate": 9.70178447684704e-05,
+      "loss": 0.0324,
+      "step": 4650
+    },
+    {
+      "grad_norm": 0.31251823902130127,
+      "learning_rate": 9.699906643660483e-05,
+      "loss": 0.0319,
+      "step": 4660
+    },
+    {
+      "grad_norm": 0.24932055175304413,
+      "learning_rate": 9.698023099644185e-05,
+      "loss": 0.0334,
+      "step": 4670
+    },
+    {
+      "grad_norm": 0.30251333117485046,
+      "learning_rate": 9.696133847086823e-05,
+      "loss": 0.0294,
+      "step": 4680
+    },
+    {
+      "grad_norm": 0.43989869952201843,
+      "learning_rate": 9.694238888284022e-05,
+      "loss": 0.0363,
+      "step": 4690
+    },
+    {
+      "grad_norm": 0.342578262090683,
+      "learning_rate": 9.692338225538333e-05,
+      "loss": 0.0329,
+      "step": 4700
+    },
+    {
+      "grad_norm": 0.2391144335269928,
+      "learning_rate": 9.690431861159241e-05,
+      "loss": 0.0324,
+      "step": 4710
+    },
+    {
+      "grad_norm": 0.3693285286426544,
+      "learning_rate": 9.688519797463161e-05,
+      "loss": 0.0311,
+      "step": 4720
+    },
+    {
+      "grad_norm": 0.39516782760620117,
+      "learning_rate": 9.686602036773426e-05,
+      "loss": 0.0331,
+      "step": 4730
+    },
+    {
+      "grad_norm": 0.3209303319454193,
+      "learning_rate": 9.684678581420302e-05,
+      "loss": 0.0325,
+      "step": 4740
+    },
+    {
+      "grad_norm": 0.24821114540100098,
+      "learning_rate": 9.682749433740962e-05,
+      "loss": 0.038,
+      "step": 4750
+    },
+    {
+      "grad_norm": 0.26262345910072327,
+      "learning_rate": 9.680814596079507e-05,
+      "loss": 0.0346,
+      "step": 4760
+    },
+    {
+      "grad_norm": 0.9238157272338867,
+      "learning_rate": 9.678874070786945e-05,
+      "loss": 0.0324,
+      "step": 4770
+    },
+    {
+      "grad_norm": 0.1953527331352234,
+      "learning_rate": 9.676927860221199e-05,
+      "loss": 0.0351,
+      "step": 4780
+    },
+    {
+      "grad_norm": 0.2232116311788559,
+      "learning_rate": 9.674975966747097e-05,
+      "loss": 0.0343,
+      "step": 4790
+    },
+    {
+      "grad_norm": 0.2524062395095825,
+      "learning_rate": 9.673018392736374e-05,
+      "loss": 0.0299,
+      "step": 4800
+    },
+    {
+      "grad_norm": 0.19976535439491272,
+      "learning_rate": 9.671055140567667e-05,
+      "loss": 0.0356,
+      "step": 4810
+    },
+    {
+      "grad_norm": 0.31831347942352295,
+      "learning_rate": 9.669086212626511e-05,
+      "loss": 0.0319,
+      "step": 4820
+    },
+    {
+      "grad_norm": 0.2573797404766083,
+      "learning_rate": 9.667111611305341e-05,
+      "loss": 0.0333,
+      "step": 4830
+    },
+    {
+      "grad_norm": 0.39954403042793274,
+      "learning_rate": 9.665131339003486e-05,
+      "loss": 0.0331,
+      "step": 4840
+    },
+    {
+      "grad_norm": 0.34243258833885193,
+      "learning_rate": 9.663145398127158e-05,
+      "loss": 0.0279,
+      "step": 4850
+    },
+    {
+      "grad_norm": 0.40026620030403137,
+      "learning_rate": 9.661153791089467e-05,
+      "loss": 0.0322,
+      "step": 4860
+    },
+    {
+      "grad_norm": 0.27648770809173584,
+      "learning_rate": 9.659156520310402e-05,
+      "loss": 0.0309,
+      "step": 4870
+    },
+    {
+      "grad_norm": 0.23704074323177338,
+      "learning_rate": 9.657153588216835e-05,
+      "loss": 0.0301,
+      "step": 4880
+    },
+    {
+      "grad_norm": 0.22470659017562866,
+      "learning_rate": 9.655144997242516e-05,
+      "loss": 0.0307,
+      "step": 4890
+    },
+    {
+      "grad_norm": 0.26196277141571045,
+      "learning_rate": 9.653130749828075e-05,
+      "loss": 0.031,
+      "step": 4900
+    },
+    {
+      "grad_norm": 0.43631988763809204,
+      "learning_rate": 9.65111084842101e-05,
+      "loss": 0.0314,
+      "step": 4910
+    },
+    {
+      "grad_norm": 0.2574961185455322,
+      "learning_rate": 9.649085295475695e-05,
+      "loss": 0.0345,
+      "step": 4920
+    },
+    {
+      "grad_norm": 0.26858606934547424,
+      "learning_rate": 9.647054093453365e-05,
+      "loss": 0.0322,
+      "step": 4930
+    },
+    {
+      "grad_norm": 0.23716683685779572,
+      "learning_rate": 9.645017244822123e-05,
+      "loss": 0.0282,
+      "step": 4940
+    },
+    {
+      "grad_norm": 0.2518445551395416,
+      "learning_rate": 9.642974752056931e-05,
+      "loss": 0.0325,
+      "step": 4950
+    },
+    {
+      "grad_norm": 0.3249128460884094,
+      "learning_rate": 9.640926617639613e-05,
+      "loss": 0.0331,
+      "step": 4960
+    },
+    {
+      "grad_norm": 0.22383485734462738,
+      "learning_rate": 9.638872844058843e-05,
+      "loss": 0.0303,
+      "step": 4970
+    },
+    {
+      "grad_norm": 0.26281267404556274,
+      "learning_rate": 9.63681343381015e-05,
+      "loss": 0.0313,
+      "step": 4980
+    },
+    {
+      "grad_norm": 0.238747239112854,
+      "learning_rate": 9.634748389395914e-05,
+      "loss": 0.029,
+      "step": 4990
+    },
+    {
+      "grad_norm": 0.36877527832984924,
+      "learning_rate": 9.632677713325353e-05,
+      "loss": 0.0358,
+      "step": 5000
+    },
+    {
+      "grad_norm": 0.3157718777656555,
+      "learning_rate": 9.63060140811454e-05,
+      "loss": 0.0331,
+      "step": 5010
+    },
+    {
+      "grad_norm": 0.3446662127971649,
+      "learning_rate": 9.628519476286379e-05,
+      "loss": 0.0344,
+      "step": 5020
+    },
+    {
+      "grad_norm": 0.2804613709449768,
+      "learning_rate": 9.626431920370612e-05,
+      "loss": 0.0307,
+      "step": 5030
+    },
+    {
+      "grad_norm": 0.2892847955226898,
+      "learning_rate": 9.624338742903819e-05,
+      "loss": 0.0295,
+      "step": 5040
+    },
+    {
+      "grad_norm": 0.5710871815681458,
+      "learning_rate": 9.622239946429406e-05,
+      "loss": 0.0276,
+      "step": 5050
+    },
+    {
+      "grad_norm": 0.26316478848457336,
+      "learning_rate": 9.620135533497609e-05,
+      "loss": 0.0302,
+      "step": 5060
+    },
+    {
+      "grad_norm": 0.3482988178730011,
+      "learning_rate": 9.61802550666549e-05,
+      "loss": 0.0286,
+      "step": 5070
+    },
+    {
+      "grad_norm": 0.26152899861335754,
+      "learning_rate": 9.615909868496928e-05,
+      "loss": 0.027,
+      "step": 5080
+    },
+    {
+      "grad_norm": 0.317668080329895,
+      "learning_rate": 9.613788621562622e-05,
+      "loss": 0.0296,
+      "step": 5090
+    },
+    {
+      "grad_norm": 0.25961193442344666,
+      "learning_rate": 9.611661768440093e-05,
+      "loss": 0.0312,
+      "step": 5100
+    },
+    {
+      "grad_norm": 0.25917866826057434,
+      "learning_rate": 9.609529311713661e-05,
+      "loss": 0.0304,
+      "step": 5110
+    },
+    {
+      "grad_norm": 0.25132909417152405,
+      "learning_rate": 9.607391253974466e-05,
+      "loss": 0.0296,
+      "step": 5120
+    },
+    {
+      "grad_norm": 0.2505704164505005,
+      "learning_rate": 9.605247597820448e-05,
+      "loss": 0.0316,
+      "step": 5130
+    },
+    {
+      "grad_norm": 0.23508676886558533,
+      "learning_rate": 9.603098345856354e-05,
+      "loss": 0.032,
+      "step": 5140
+    },
+    {
+      "grad_norm": 0.23048117756843567,
+      "learning_rate": 9.600943500693725e-05,
+      "loss": 0.0293,
+      "step": 5150
+    },
+    {
+      "grad_norm": 0.30994167923927307,
+      "learning_rate": 9.598783064950902e-05,
+      "loss": 0.0291,
+      "step": 5160
+    },
+    {
+      "grad_norm": 0.34612396359443665,
+      "learning_rate": 9.596617041253018e-05,
+      "loss": 0.0284,
+      "step": 5170
+    },
+    {
+      "grad_norm": 0.18393488228321075,
+      "learning_rate": 9.594445432231996e-05,
+      "loss": 0.0281,
+      "step": 5180
+    },
+    {
+      "grad_norm": 0.22680553793907166,
+      "learning_rate": 9.592268240526547e-05,
+      "loss": 0.0312,
+      "step": 5190
+    },
+    {
+      "grad_norm": 0.39471128582954407,
+      "learning_rate": 9.590085468782162e-05,
+      "loss": 0.0298,
+      "step": 5200
+    },
+    {
+      "grad_norm": 0.23328055441379547,
+      "learning_rate": 9.587897119651116e-05,
+      "loss": 0.0308,
+      "step": 5210
+    },
+    {
+      "grad_norm": 0.2956632375717163,
+      "learning_rate": 9.585703195792459e-05,
+      "loss": 0.0317,
+      "step": 5220
+    },
+    {
+      "grad_norm": 0.34203052520751953,
+      "learning_rate": 9.583503699872016e-05,
+      "loss": 0.0292,
+      "step": 5230
+    },
+    {
+      "grad_norm": 0.41844645142555237,
+      "learning_rate": 9.581298634562381e-05,
+      "loss": 0.0296,
+      "step": 5240
+    },
+    {
+      "grad_norm": 0.15062275528907776,
+      "learning_rate": 9.579088002542917e-05,
+      "loss": 0.0305,
+      "step": 5250
+    },
+    {
+      "grad_norm": 0.3013925552368164,
+      "learning_rate": 9.57687180649975e-05,
+      "loss": 0.0314,
+      "step": 5260
+    },
+    {
+      "grad_norm": 0.2418990582227707,
+      "learning_rate": 9.574650049125768e-05,
+      "loss": 0.0299,
+      "step": 5270
+    },
+    {
+      "grad_norm": 0.304177850484848,
+      "learning_rate": 9.572422733120614e-05,
+      "loss": 0.0283,
+      "step": 5280
+    },
+    {
+      "grad_norm": 0.23789408802986145,
+      "learning_rate": 9.570189861190689e-05,
+      "loss": 0.0266,
+      "step": 5290
+    },
+    {
+      "grad_norm": 0.28955498337745667,
+      "learning_rate": 9.56795143604914e-05,
+      "loss": 0.0308,
+      "step": 5300
+    },
+    {
+      "grad_norm": 0.16970686614513397,
+      "learning_rate": 9.565707460415869e-05,
+      "loss": 0.0286,
+      "step": 5310
+    },
+    {
+      "grad_norm": 0.33111685514450073,
+      "learning_rate": 9.563457937017515e-05,
+      "loss": 0.0292,
+      "step": 5320
+    },
+    {
+      "grad_norm": 0.2950333058834076,
+      "learning_rate": 9.56120286858746e-05,
+      "loss": 0.0282,
+      "step": 5330
+    },
+    {
+      "grad_norm": 0.2266271710395813,
+      "learning_rate": 9.558942257865829e-05,
+      "loss": 0.0303,
+      "step": 5340
+    },
+    {
+      "grad_norm": 0.27608922123908997,
+      "learning_rate": 9.556676107599472e-05,
+      "loss": 0.03,
+      "step": 5350
+    },
+    {
+      "grad_norm": 0.2176966667175293,
+      "learning_rate": 9.554404420541978e-05,
+      "loss": 0.0308,
+      "step": 5360
+    },
+    {
+      "grad_norm": 0.22792655229568481,
+      "learning_rate": 9.55212719945366e-05,
+      "loss": 0.0281,
+      "step": 5370
+    },
+    {
+      "grad_norm": 0.2198641151189804,
+      "learning_rate": 9.549844447101559e-05,
+      "loss": 0.0277,
+      "step": 5380
+    },
+    {
+      "grad_norm": 0.32253924012184143,
+      "learning_rate": 9.547556166259433e-05,
+      "loss": 0.0288,
+      "step": 5390
+    },
+    {
+      "grad_norm": 0.38950228691101074,
+      "learning_rate": 9.545262359707756e-05,
+      "loss": 0.0285,
+      "step": 5400
+    },
+    {
+      "grad_norm": 0.19937629997730255,
+      "learning_rate": 9.542963030233724e-05,
+      "loss": 0.0305,
+      "step": 5410
+    },
+    {
+      "grad_norm": 0.24631406366825104,
+      "learning_rate": 9.540658180631237e-05,
+      "loss": 0.0301,
+      "step": 5420
+    },
+    {
+      "grad_norm": 0.22192849218845367,
+      "learning_rate": 9.538347813700904e-05,
+      "loss": 0.0281,
+      "step": 5430
+    },
+    {
+      "grad_norm": 0.2845704257488251,
+      "learning_rate": 9.536031932250036e-05,
+      "loss": 0.0277,
+      "step": 5440
+    },
+    {
+      "grad_norm": 0.22213523089885712,
+      "learning_rate": 9.533710539092653e-05,
+      "loss": 0.0279,
+      "step": 5450
+    },
+    {
+      "grad_norm": 0.5807768106460571,
+      "learning_rate": 9.531383637049464e-05,
+      "loss": 0.0258,
+      "step": 5460
+    },
+    {
+      "grad_norm": 0.28621476888656616,
+      "learning_rate": 9.529051228947875e-05,
+      "loss": 0.0305,
+      "step": 5470
+    },
+    {
+      "grad_norm": 0.31888940930366516,
+      "learning_rate": 9.52671331762198e-05,
+      "loss": 0.0282,
+      "step": 5480
+    },
+    {
+      "grad_norm": 0.3353741466999054,
+      "learning_rate": 9.524369905912565e-05,
+      "loss": 0.0279,
+      "step": 5490
+    },
+    {
+      "grad_norm": 0.2915017604827881,
+      "learning_rate": 9.522020996667092e-05,
+      "loss": 0.0339,
+      "step": 5500
+    },
+    {
+      "grad_norm": 0.2807980477809906,
+      "learning_rate": 9.519666592739709e-05,
+      "loss": 0.0283,
+      "step": 5510
+    },
+    {
+      "grad_norm": 0.24497555196285248,
+      "learning_rate": 9.517306696991241e-05,
+      "loss": 0.0296,
+      "step": 5520
+    },
+    {
+      "grad_norm": 0.23863714933395386,
+      "learning_rate": 9.51494131228918e-05,
+      "loss": 0.0277,
+      "step": 5530
+    },
+    {
+      "grad_norm": 0.42585837841033936,
+      "learning_rate": 9.512570441507695e-05,
+      "loss": 0.0301,
+      "step": 5540
+    },
+    {
+      "grad_norm": 0.28657546639442444,
+      "learning_rate": 9.510194087527615e-05,
+      "loss": 0.0273,
+      "step": 5550
+    },
+    {
+      "grad_norm": 0.26880064606666565,
+      "learning_rate": 9.507812253236435e-05,
+      "loss": 0.0292,
+      "step": 5560
+    },
+    {
+      "grad_norm": 0.5710074305534363,
+      "learning_rate": 9.505424941528309e-05,
+      "loss": 0.0311,
+      "step": 5570
+    },
+    {
+      "grad_norm": 0.26207560300827026,
+      "learning_rate": 9.503032155304046e-05,
+      "loss": 0.0346,
+      "step": 5580
+    },
+    {
+      "grad_norm": 0.5927866101264954,
+      "learning_rate": 9.500633897471106e-05,
+      "loss": 0.0298,
+      "step": 5590
+    },
+    {
+      "grad_norm": 0.26272717118263245,
+      "learning_rate": 9.498230170943596e-05,
+      "loss": 0.0351,
+      "step": 5600
+    },
+    {
+      "grad_norm": 0.20593000948429108,
+      "learning_rate": 9.495820978642275e-05,
+      "loss": 0.0326,
+      "step": 5610
+    },
+    {
+      "grad_norm": 0.2362719476222992,
+      "learning_rate": 9.493406323494535e-05,
+      "loss": 0.0291,
+      "step": 5620
+    },
+    {
+      "grad_norm": 0.26528915762901306,
+      "learning_rate": 9.490986208434413e-05,
+      "loss": 0.0307,
+      "step": 5630
+    },
+    {
+      "grad_norm": 0.34711647033691406,
+      "learning_rate": 9.488560636402577e-05,
+      "loss": 0.0261,
+      "step": 5640
+    },
+    {
+      "grad_norm": 0.256011962890625,
+      "learning_rate": 9.486129610346321e-05,
+      "loss": 0.0286,
+      "step": 5650
+    },
+    {
+      "grad_norm": 0.2524960935115814,
+      "learning_rate": 9.483693133219576e-05,
+      "loss": 0.0285,
+      "step": 5660
+    },
+    {
+      "grad_norm": 0.3123167157173157,
+      "learning_rate": 9.481251207982888e-05,
+      "loss": 0.0318,
+      "step": 5670
+    },
+    {
+      "grad_norm": 0.23425480723381042,
+      "learning_rate": 9.47880383760343e-05,
+      "loss": 0.0261,
+      "step": 5680
+    },
+    {
+      "grad_norm": 0.25420740246772766,
+      "learning_rate": 9.476351025054983e-05,
+      "loss": 0.0306,
+      "step": 5690
+    },
+    {
+      "grad_norm": 0.3529684841632843,
+      "learning_rate": 9.473892773317952e-05,
+      "loss": 0.0253,
+      "step": 5700
+    },
+    {
+      "grad_norm": 0.31296902894973755,
+      "learning_rate": 9.471429085379338e-05,
+      "loss": 0.0292,
+      "step": 5710
+    },
+    {
+      "grad_norm": 0.2692904472351074,
+      "learning_rate": 9.468959964232757e-05,
+      "loss": 0.0272,
+      "step": 5720
+    },
+    {
+      "grad_norm": 0.3672389090061188,
+      "learning_rate": 9.466485412878425e-05,
+      "loss": 0.0334,
+      "step": 5730
+    },
+    {
+      "grad_norm": 0.3293924629688263,
+      "learning_rate": 9.464005434323154e-05,
+      "loss": 0.0298,
+      "step": 5740
+    },
+    {
+      "grad_norm": 0.2921695113182068,
+      "learning_rate": 9.461520031580352e-05,
+      "loss": 0.0266,
+      "step": 5750
+    },
+    {
+      "grad_norm": 0.29170626401901245,
+      "learning_rate": 9.459029207670019e-05,
+      "loss": 0.0293,
+      "step": 5760
+    },
+    {
+      "grad_norm": 0.33403706550598145,
+      "learning_rate": 9.456532965618737e-05,
+      "loss": 0.0285,
+      "step": 5770
+    },
+    {
+      "grad_norm": 0.24151206016540527,
+      "learning_rate": 9.454031308459681e-05,
+      "loss": 0.0287,
+      "step": 5780
+    },
+    {
+      "grad_norm": 0.23442505300045013,
+      "learning_rate": 9.451524239232595e-05,
+      "loss": 0.0289,
+      "step": 5790
+    },
+    {
+      "grad_norm": 0.22394013404846191,
+      "learning_rate": 9.449011760983809e-05,
+      "loss": 0.0316,
+      "step": 5800
+    },
+    {
+      "grad_norm": 0.162757009267807,
+      "learning_rate": 9.446493876766218e-05,
+      "loss": 0.03,
+      "step": 5810
+    },
+    {
+      "grad_norm": 0.24485236406326294,
+      "learning_rate": 9.44397058963929e-05,
+      "loss": 0.0308,
+      "step": 5820
+    },
+    {
+      "grad_norm": 0.30540770292282104,
+      "learning_rate": 9.441441902669056e-05,
+      "loss": 0.029,
+      "step": 5830
+    },
+    {
+      "grad_norm": 0.38131895661354065,
+      "learning_rate": 9.43890781892811e-05,
+      "loss": 0.0279,
+      "step": 5840
+    },
+    {
+      "grad_norm": 0.27463778853416443,
+      "learning_rate": 9.436368341495603e-05,
+      "loss": 0.0305,
+      "step": 5850
+    },
+    {
+      "grad_norm": 0.18571634590625763,
+      "learning_rate": 9.43382347345724e-05,
+      "loss": 0.0301,
+      "step": 5860
+    },
+    {
+      "grad_norm": 0.19644570350646973,
+      "learning_rate": 9.431273217905273e-05,
+      "loss": 0.0297,
+      "step": 5870
+    },
+    {
+      "grad_norm": 0.2669595181941986,
+      "learning_rate": 9.428717577938504e-05,
+      "loss": 0.0287,
+      "step": 5880
+    },
+    {
+      "grad_norm": 0.2687034010887146,
+      "learning_rate": 9.426156556662276e-05,
+      "loss": 0.0254,
+      "step": 5890
+    },
+    {
+      "grad_norm": 0.25875869393348694,
+      "learning_rate": 9.423590157188474e-05,
+      "loss": 0.0293,
+      "step": 5900
+    },
+    {
+      "grad_norm": 0.29127827286720276,
+      "learning_rate": 9.421018382635513e-05,
+      "loss": 0.0276,
+      "step": 5910
+    },
+    {
+      "grad_norm": 0.2213948369026184,
+      "learning_rate": 9.418441236128343e-05,
+      "loss": 0.0266,
+      "step": 5920
+    },
+    {
+      "grad_norm": 0.2627449631690979,
+      "learning_rate": 9.41585872079844e-05,
+      "loss": 0.0335,
+      "step": 5930
+    },
+    {
+      "grad_norm": 0.2908029854297638,
+      "learning_rate": 9.413270839783802e-05,
+      "loss": 0.0312,
+      "step": 5940
+    },
+    {
+      "grad_norm": 0.3014381229877472,
+      "learning_rate": 9.41067759622895e-05,
+      "loss": 0.0315,
+      "step": 5950
+    },
+    {
+      "grad_norm": 0.2169884592294693,
+      "learning_rate": 9.408078993284917e-05,
+      "loss": 0.0275,
+      "step": 5960
+    },
+    {
+      "grad_norm": 0.15191787481307983,
+      "learning_rate": 9.405475034109254e-05,
+      "loss": 0.0256,
+      "step": 5970
+    },
+    {
+      "grad_norm": 0.21545839309692383,
+      "learning_rate": 9.402865721866015e-05,
+      "loss": 0.0293,
+      "step": 5980
+    },
+    {
+      "grad_norm": 0.24423319101333618,
+      "learning_rate": 9.400251059725762e-05,
+      "loss": 0.0293,
+      "step": 5990
+    },
+    {
+      "grad_norm": 0.25787439942359924,
+      "learning_rate": 9.397631050865554e-05,
+      "loss": 0.0254,
+      "step": 6000
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 30000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 3000,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-6000/wandb_config.json

@@ -0,0 +1 @@
+{"project": "finetune-gr00t-n1d6", "run_id": "reachy2"}

checkpoint-9000/embodiment_id.json

@@ -0,0 +1,12 @@
+{
+  "robocasa_panda_omron": 13,
+  "gr1": 20,
+  "behavior_r1_pro": 24,
+  "unitree_g1": 8,
+  "oxe_google": 0,
+  "oxe_widowx": 1,
+  "libero_panda": 2,
+  "oxe_droid": 16,
+  "new_embodiment": 10,
+  "reachy2": 11
+}

checkpoint-9000/experiment_cfg/conf.yaml

@@ -0,0 +1,203 @@
+load_config_path: null
+model:
+  model_type: Gr00tN1d6
+  model_dtype: bfloat16
+  model_name: nvidia/Eagle-Block2A-2B-v2
+  backbone_model_type: eagle
+  model_revision: null
+  tune_top_llm_layers: 4
+  backbone_embedding_dim: 2048
+  tune_llm: false
+  tune_visual: false
+  select_layer: 16
+  reproject_vision: false
+  use_flash_attention: true
+  load_bf16: false
+  collator_overwrite_image_inputs: false
+  eagle_collator: true
+  backbone_trainable_params_fp32: true
+  image_crop_size: null
+  image_target_size: null
+  shortest_image_edge: 256
+  crop_fraction: 0.95
+  random_rotation_angle: null
+  color_jitter_params: null
+  use_albumentations_transforms: true
+  formalize_language: true
+  apply_sincos_state_encoding: false
+  use_relative_action: true
+  max_state_dim: 29
+  max_action_dim: 29
+  action_horizon: 16
+  hidden_size: 1024
+  input_embedding_dim: 1536
+  add_pos_embed: true
+  attn_dropout: 0.2
+  use_vlln: true
+  max_seq_len: 1024
+  use_alternate_vl_dit: true
+  attend_text_every_n_blocks: 2
+  diffusion_model_cfg:
+    positional_embeddings: null
+    num_layers: 32
+    num_attention_heads: 32
+    attention_head_dim: 48
+    norm_type: ada_norm
+    dropout: 0.2
+    final_dropout: true
+    output_dim: 1024
+    interleave_self_attention: true
+  num_inference_timesteps: 4
+  noise_beta_alpha: 1.5
+  noise_beta_beta: 1.0
+  noise_s: 0.999
+  num_timestep_buckets: 1000
+  tune_projector: true
+  tune_diffusion_model: true
+  tune_vlln: true
+  state_dropout_prob: 0.0
+  state_additive_noise_scale: 0.0
+  max_num_embodiments: 32
+data:
+  datasets:
+  - dataset_paths:
+    - ./datasets/reachy2_100
+    embodiment_tag: reachy2
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    reachy2:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front_cam
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - arm_joints
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      action:
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        modality_keys:
+        - arm_joints
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - rep: RELATIVE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: arm_joints
+        - rep: ABSOLUTE
+          type: NON_EEF
+          format: DEFAULT
+          state_key: null
+      language:
+        delta_indices:
+        - 0
+        modality_keys:
+        - annotation.human.task_description
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+  download_cache: false
+  shard_size: 1024
+  episode_sampling_rate: 0.1
+  num_shards_per_epoch: 10000
+  override_pretraining_statistics: false
+  mode: single_turn
+  random_chop: 0.0
+  mock_dataset_mode: false
+  shuffle: true
+  seed: 42
+  multiprocessing_context: fork
+  allow_padding: false
+  subsample_ratio: 1.0
+  image_crop_size:
+  - 244
+  - 244
+  image_target_size:
+  - 224
+  - 224
+  video_backend: decord
+training:
+  output_dir: /tmp/groot_output/reachy2
+  experiment_name: null
+  max_steps: 30000
+  global_batch_size: 64
+  batch_size: null
+  gradient_accumulation_steps: 1
+  learning_rate: 0.0001
+  lr_scheduler_type: cosine
+  weight_decay: 1.0e-05
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  max_grad_norm: 1.0
+  optim: adamw_torch
+  start_from_checkpoint: nvidia/GR00T-N1.6-3B
+  tf32: true
+  fp16: false
+  bf16: true
+  eval_bf16: true
+  logging_steps: 10
+  save_steps: 3000
+  save_total_limit: 5
+  save_vl_model: false
+  upload_checkpoints: false
+  upload_every: 1000
+  upload_last_n_checkpoints: 5
+  max_concurrent_uploads: 2
+  eval_strategy: 'no'
+  eval_steps: 500
+  eval_set_split_ratio: 0.1
+  eval_batch_size: 2
+  save_best_eval_metric_name: ''
+  save_best_eval_metric_greater_is_better: true
+  deepspeed_stage: 2
+  gradient_checkpointing: false
+  transformers_trust_remote_code: true
+  transformers_local_files_only: false
+  transformers_cache_dir: null
+  transformers_access_token: null
+  use_ddp: false
+  ddp_bucket_cap_mb: 100
+  num_gpus: 1
+  dataloader_num_workers: 8
+  remove_unused_columns: false
+  use_wandb: false
+  wandb_project: finetune-gr00t-n1d6
+  enable_profiling: false
+  max_retries: 3
+  assert_loss_less_than: null
+  add_rl_callback: false
+  enable_open_loop_eval: false
+  open_loop_eval_traj_ids:
+  - 0
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_plot_indices: null
+max_steps: 30000
+save_steps: 3000

checkpoint-9000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,1020 @@
+{
+  "reachy2": {
+    "state": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.316539764404297,
+          -0.3203781247138977,
+          2.417543411254883,
+          -109.398681640625,
+          -1.9847908020019531,
+          -3.484698534011841,
+          -0.3555510342121124
+        ],
+        "std": [
+          7.066186428070068,
+          3.507993459701538,
+          22.217893600463867,
+          14.851909637451172,
+          13.270956039428711,
+          11.239269256591797,
+          11.768560409545898
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723615646362305,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      }
+    },
+    "action": {
+      "arm_joints": {
+        "min": [
+          4.93005895614624,
+          -11.668736457824707,
+          -42.069515228271484,
+          -126.52813720703125,
+          -20.27557945251465,
+          -20.164306640625,
+          -21.239810943603516
+        ],
+        "max": [
+          41.02817153930664,
+          11.43976879119873,
+          42.18703079223633,
+          -74.72509765625,
+          20.33566665649414,
+          16.37257957458496,
+          21.077407836914062
+        ],
+        "mean": [
+          22.172462463378906,
+          -0.3404896557331085,
+          2.3248910903930664,
+          -109.23794555664062,
+          -1.812583327293396,
+          -3.5105223655700684,
+          -0.2696441113948822
+        ],
+        "std": [
+          7.037580966949463,
+          3.620150566101074,
+          22.205087661743164,
+          14.961851119995117,
+          13.164973258972168,
+          11.261707305908203,
+          11.734737396240234
+        ],
+        "q01": [
+          8.457822799682617,
+          -10.20227336883545,
+          -41.3212890625,
+          -126.43962860107422,
+          -20.06123924255371,
+          -20.072999954223633,
+          -20.989107131958008
+        ],
+        "q99": [
+          38.03876876831055,
+          9.723625183105469,
+          41.77986526489258,
+          -80.34276580810547,
+          20.061033248901367,
+          14.72492790222168,
+          20.985740661621094
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.5416929721832275
+        ],
+        "max": [
+          0.9385063648223877
+        ],
+        "mean": [
+          0.8345981240272522
+        ],
+        "std": [
+          0.12303706258535363
+        ],
+        "q01": [
+          0.6467909216880798
+        ],
+        "q99": [
+          0.9350630044937134
+        ]
+      }
+    },
+    "relative_action": {
+      "arm_joints": {
+        "min": [
+          [
+            -20.440982818603516,
+            -11.665962219238281,
+            -14.02474594116211,
+            -20.817272186279297,
+            -10.09469223022461,
+            -6.401888370513916,
+            -8.767365455627441
+          ],
+          [
+            -20.449968338012695,
+            -11.677082061767578,
+            -25.868675231933594,
+            -30.582874298095703,
+            -17.461694717407227,
+            -11.940333366394043,
+            -16.13103485107422
+          ],
+          [
+            -19.697128295898438,
+            -10.979728698730469,
+            -35.53736114501953,
+            -32.2390022277832,
+            -22.624008178710938,
+            -17.247066497802734,
+            -22.225337982177734
+          ],
+          [
+            -20.74054527282715,
+            -10.334735870361328,
+            -44.72303009033203,
+            -33.42095947265625,
+            -25.056644439697266,
+            -22.322742462158203,
+            -27.965492248535156
+          ],
+          [
+            -24.075529098510742,
+            -10.338696479797363,
+            -50.480987548828125,
+            -34.134925842285156,
+            -26.33944320678711,
+            -25.707611083984375,
+            -31.17104721069336
+          ],
+          [
+            -26.567087173461914,
+            -10.34083366394043,
+            -54.44721603393555,
+            -34.8769645690918,
+            -26.375856399536133,
+            -29.134244918823242,
+            -33.48938751220703
+          ],
+          [
+            -27.659536361694336,
+            -10.340897560119629,
+            -57.30182647705078,
+            -36.380828857421875,
+            -26.375410079956055,
+            -32.329833984375,
+            -35.47602081298828
+          ],
+          [
+            -27.924762725830078,
+            -11.4292573928833,
+            -59.24040603637695,
+            -36.7144775390625,
+            -26.372446060180664,
+            -33.91874313354492,
+            -35.908573150634766
+          ],
+          [
+            -27.95497703552246,
+            -11.528078079223633,
+            -59.352848052978516,
+            -36.841575622558594,
+            -26.24493980407715,
+            -35.34848403930664,
+            -36.83350372314453
+          ],
+          [
+            -28.4979248046875,
+            -11.443194389343262,
+            -59.4117431640625,
+            -36.853206634521484,
+            -26.16107940673828,
+            -36.28956604003906,
+            -36.86201858520508
+          ],
+          [
+            -28.523822784423828,
+            -11.443158149719238,
+            -59.41852951049805,
+            -36.85321807861328,
+            -26.160951614379883,
+            -36.43997573852539,
+            -36.88083267211914
+          ],
+          [
+            -28.522306442260742,
+            -11.443062782287598,
+            -59.424468994140625,
+            -37.14503860473633,
+            -26.16082763671875,
+            -36.443668365478516,
+            -36.880828857421875
+          ],
+          [
+            -28.468467712402344,
+            -11.44290828704834,
+            -59.42174530029297,
+            -37.880104064941406,
+            -26.160717010498047,
+            -36.44268035888672,
+            -36.880828857421875
+          ],
+          [
+            -28.41590118408203,
+            -11.442895889282227,
+            -59.4072265625,
+            -38.47637939453125,
+            -26.160614013671875,
+            -36.4410400390625,
+            -36.880828857421875
+          ],
+          [
+            -28.362062454223633,
+            -11.442879676818848,
+            -59.38927459716797,
+            -41.26643371582031,
+            -26.160587310791016,
+            -36.43976974487305,
+            -36.88082504272461
+          ],
+          [
+            -28.244752883911133,
+            -11.442869186401367,
+            -59.369014739990234,
+            -44.968875885009766,
+            -26.160572052001953,
+            -36.433982849121094,
+            -36.880794525146484
+          ]
+        ],
+        "max": [
+          [
+            5.384329319000244,
+            11.368682861328125,
+            12.580580711364746,
+            21.664548873901367,
+            22.868919372558594,
+            13.739058494567871,
+            11.726442337036133
+          ],
+          [
+            9.6729097366333,
+            11.36864185333252,
+            23.997724533081055,
+            21.659494400024414,
+            22.868900299072266,
+            23.967060089111328,
+            14.994425773620605
+          ],
+          [
+            13.331961631774902,
+            11.190152168273926,
+            33.41366958618164,
+            21.635488510131836,
+            25.30813217163086,
+            26.21414566040039,
+            20.79837417602539
+          ],
+          [
+            16.202699661254883,
+            10.40443229675293,
+            40.09684753417969,
+            23.701637268066406,
+            28.603330612182617,
+            27.53230857849121,
+            25.446346282958984
+          ],
+          [
+            19.05048370361328,
+            10.40436840057373,
+            46.37767791748047,
+            28.047138214111328,
+            31.07068634033203,
+            28.03820037841797,
+            29.473896026611328
+          ],
+          [
+            20.245254516601562,
+            10.402230262756348,
+            51.10308837890625,
+            31.863624572753906,
+            31.10450553894043,
+            29.278892517089844,
+            32.1208381652832
+          ],
+          [
+            21.428510665893555,
+            10.398269653320312,
+            53.42000961303711,
+            34.682308197021484,
+            31.103679656982422,
+            30.250988006591797,
+            33.45231246948242
+          ],
+          [
+            21.98176383972168,
+            10.395852088928223,
+            55.17178726196289,
+            36.64023208618164,
+            31.101837158203125,
+            30.935718536376953,
+            34.229305267333984
+          ],
+          [
+            22.397876739501953,
+            11.4427490234375,
+            55.23532485961914,
+            37.58878707885742,
+            31.096826553344727,
+            31.260026931762695,
+            34.44834899902344
+          ],
+          [
+            22.6441650390625,
+            11.780808448791504,
+            55.34986114501953,
+            38.26850128173828,
+            31.096656799316406,
+            31.333267211914062,
+            34.469146728515625
+          ],
+          [
+            22.692224502563477,
+            11.812055587768555,
+            55.36343765258789,
+            38.49715805053711,
+            31.61040687561035,
+            31.37093162536621,
+            34.469146728515625
+          ],
+          [
+            22.72762107849121,
+            11.678536415100098,
+            55.36940383911133,
+            38.59236526489258,
+            35.083839416503906,
+            31.3709659576416,
+            34.46915054321289
+          ],
+          [
+            24.14029312133789,
+            11.678519248962402,
+            55.369537353515625,
+            38.63944625854492,
+            38.55842208862305,
+            31.37123680114746,
+            34.46915054321289
+          ],
+          [
+            24.973907470703125,
+            11.678510665893555,
+            55.36623001098633,
+            38.656288146972656,
+            39.96430587768555,
+            31.371545791625977,
+            34.46915054321289
+          ],
+          [
+            25.06302261352539,
+            11.678531646728516,
+            55.3525390625,
+            38.694881439208984,
+            40.0225944519043,
+            31.37181854248047,
+            34.468910217285156
+          ],
+          [
+            25.270801544189453,
+            11.678507804870605,
+            55.302616119384766,
+            38.694969177246094,
+            40.022586822509766,
+            31.37185287475586,
+            34.385284423828125
+          ]
+        ],
+        "mean": [
+          [
+            -0.09851349890232086,
+            -1.3889506362829707e-06,
+            -0.021199805662035942,
+            -0.0799250528216362,
+            0.1505303829908371,
+            0.11832442134618759,
+            0.13495467603206635
+          ],
+          [
+            -0.17057012021541595,
+            0.004252429585903883,
+            -0.008995133452117443,
+            -0.22457502782344818,
+            0.261105477809906,
+            0.27207958698272705,
+            0.24137930572032928
+          ],
+          [
+            -0.10028649866580963,
+            0.03141691908240318,
+            0.09048548340797424,
+            -0.5493719577789307,
+            0.19902515411376953,
+            0.47134673595428467,
+            0.2771521210670471
+          ],
+          [
+            -0.05444306880235672,
+            0.05378008633852005,
+            0.1796490103006363,
+            -0.879206120967865,
+            0.1520315706729889,
+            0.6866292357444763,
+            0.3248023986816406
+          ],
+          [
+            -0.013876532204449177,
+            0.07615811377763748,
+            0.2686113715171814,
+            -1.2192035913467407,
+            0.10656605660915375,
+            0.9128072261810303,
+            0.3744547367095947
+          ],
+          [
+            0.025331225246191025,
+            0.09861738234758377,
+            0.3578471839427948,
+            -1.5637919902801514,
+            0.06106000766158104,
+            1.1433591842651367,
+            0.4248007833957672
+          ],
+          [
+            0.0637386217713356,
+            0.12045339494943619,
+            0.44647839665412903,
+            -1.908605933189392,
+            0.016334885731339455,
+            1.3740676641464233,
+            0.4745999872684479
+          ],
+          [
+            0.10053692013025284,
+            0.13872262835502625,
+            0.5284987092018127,
+            -2.2462470531463623,
+            -0.02209978550672531,
+            1.5981824398040771,
+            0.518027663230896
+          ],
+          [
+            0.1247839704155922,
+            0.14948606491088867,
+            0.583483874797821,
+            -2.525705337524414,
+            -0.04458107799291611,
+            1.7731598615646362,
+            0.5445129871368408
+          ],
+          [
+            0.12165514379739761,
+            0.14827895164489746,
+            0.602064847946167,
+            -2.593496561050415,
+            -0.047096796333789825,
+            1.797343134880066,
+            0.5500790476799011
+          ],
+          [
+            0.1124163493514061,
+            0.14204099774360657,
+            0.6022222638130188,
+            -2.55938458442688,
+            -0.04651280865073204,
+            1.7794476747512817,
+            0.5512804985046387
+          ],
+          [
+            0.09131923317909241,
+            0.12857742607593536,
+            0.6059610843658447,
+            -2.4963152408599854,
+            -0.04961543530225754,
+            1.740905523300171,
+            0.5537145137786865
+          ],
+          [
+            0.054126936942338943,
+            0.10883964598178864,
+            0.6135755181312561,
+            -2.381096601486206,
+            -0.056390050798654556,
+            1.6682276725769043,
+            0.5572347044944763
+          ],
+          [
+            -0.001200970378704369,
+            0.08611582219600677,
+            0.6268700361251831,
+            -2.207465648651123,
+            -0.06850837916135788,
+            1.5567927360534668,
+            0.5623834133148193
+          ],
+          [
+            -0.07379738241434097,
+            0.06273127347230911,
+            0.6453807950019836,
+            -1.9786254167556763,
+            -0.08601187914609909,
+            1.4089199304580688,
+            0.5690953135490417
+          ],
+          [
+            -0.15925677120685577,
+            0.039619192481040955,
+            0.6695932149887085,
+            -1.7069525718688965,
+            -0.10912654548883438,
+            1.2322944402694702,
+            0.5775030851364136
+          ]
+        ],
+        "std": [
+          [
+            1.9917564392089844,
+            1.1211100816726685,
+            2.417202949523926,
+            2.94158935546875,
+            2.39365291595459,
+            1.874914526939392,
+            1.6537472009658813
+          ],
+          [
+            3.056582450866699,
+            1.6457527875900269,
+            4.458590507507324,
+            5.178750514984131,
+            3.765504837036133,
+            3.6351206302642822,
+            2.9263687133789062
+          ],
+          [
+            3.596614122390747,
+            1.8142179250717163,
+            6.294006824493408,
+            7.028163433074951,
+            4.514294147491455,
+            5.296512126922607,
+            4.0286865234375
+          ],
+          [
+            4.277676582336426,
+            2.040328025817871,
+            8.043697357177734,
+            8.868635177612305,
+            5.327646255493164,
+            6.84896993637085,
+            5.116058349609375
+          ],
+          [
+            4.939493179321289,
+            2.248694896697998,
+            9.646915435791016,
+            10.583718299865723,
+            6.0862040519714355,
+            8.271647453308105,
+            6.122095584869385
+          ],
+          [
+            5.552292823791504,
+            2.439715623855591,
+            11.097370147705078,
+            12.150360107421875,
+            6.7813005447387695,
+            9.56061840057373,
+            7.034397125244141
+          ],
+          [
+            6.109565258026123,
+            2.6163902282714844,
+            12.404134750366211,
+            13.574675559997559,
+            7.419105052947998,
+            10.72553539276123,
+            7.855382919311523
+          ],
+          [
+            6.608177661895752,
+            2.7779171466827393,
+            13.585582733154297,
+            14.875726699829102,
+            8.011663436889648,
+            11.787652015686035,
+            8.598040580749512
+          ],
+          [
+            7.046725273132324,
+            2.93743634223938,
+            14.688329696655273,
+            16.078651428222656,
+            8.60019302368164,
+            12.769647598266602,
+            9.277371406555176
+          ],
+          [
+            7.43002986907959,
+            3.1344540119171143,
+            15.875958442687988,
+            17.27137565612793,
+            9.292807579040527,
+            13.705930709838867,
+            9.92088794708252
+          ],
+          [
+            7.734869003295898,
+            3.2678585052490234,
+            17.0732421875,
+            18.39610481262207,
+            9.9412841796875,
+            14.559659957885742,
+            10.473697662353516
+          ],
+          [
+            7.963422775268555,
+            3.390157699584961,
+            18.116336822509766,
+            19.400726318359375,
+            10.5236177444458,
+            15.339380264282227,
+            10.94632339477539
+          ],
+          [
+            8.121000289916992,
+            3.5208446979522705,
+            19.046070098876953,
+            20.327335357666016,
+            11.053685188293457,
+            16.05907440185547,
+            11.358160972595215
+          ],
+          [
+            8.218527793884277,
+            3.651193857192993,
+            19.873519897460938,
+            21.182926177978516,
+            11.531027793884277,
+            16.716957092285156,
+            11.718016624450684
+          ],
+          [
+            8.26684284210205,
+            3.7774546146392822,
+            20.612743377685547,
+            21.96902847290039,
+            11.957940101623535,
+            17.311002731323242,
+            12.034592628479004
+          ],
+          [
+            8.276911735534668,
+            3.8994669914245605,
+            21.281835556030273,
+            22.691150665283203,
+            12.342350959777832,
+            17.84532928466797,
+            12.317776679992676
+          ]
+        ],
+        "q01": [
+          [
+            -9.697267532348633,
+            -3.6728341579437256,
+            -10.083383560180664,
+            -7.500851631164551,
+            -5.057114124298096,
+            -5.191580295562744,
+            -5.047628879547119
+          ],
+          [
+            -17.174776077270508,
+            -7.670596122741699,
+            -15.547335624694824,
+            -14.150874137878418,
+            -9.842418670654297,
+            -9.951410293579102,
+            -9.844037055969238
+          ],
+          [
+            -16.49785804748535,
+            -9.060977935791016,
+            -23.215665817260742,
+            -19.967872619628906,
+            -13.348301887512207,
+            -14.247873306274414,
+            -14.610755920410156
+          ],
+          [
+            -16.7535400390625,
+            -9.395726203918457,
+            -29.43296241760254,
+            -25.006587982177734,
+            -16.37109375,
+            -18.130178451538086,
+            -18.5999755859375
+          ],
+          [
+            -17.191621780395508,
+            -9.43246078491211,
+            -36.27531051635742,
+            -28.893749237060547,
+            -18.89956283569336,
+            -21.859458923339844,
+            -22.91756248474121
+          ],
+          [
+            -19.174758911132812,
+            -9.506808280944824,
+            -40.94317626953125,
+            -31.8222713470459,
+            -20.72899055480957,
+            -24.93976402282715,
+            -25.59789276123047
+          ],
+          [
+            -20.843280792236328,
+            -9.511780738830566,
+            -43.895530700683594,
+            -33.7496337890625,
+            -21.83380699157715,
+            -27.3231201171875,
+            -27.804710388183594
+          ],
+          [
+            -22.112207412719727,
+            -9.513916015625,
+            -47.63551330566406,
+            -34.7276611328125,
+            -22.691164016723633,
+            -29.399824142456055,
+            -30.01004981994629
+          ],
+          [
+            -23.385971069335938,
+            -9.525580406188965,
+            -49.76005172729492,
+            -35.31100082397461,
+            -23.10930824279785,
+            -30.991064071655273,
+            -31.236299514770508
+          ],
+          [
+            -23.984119415283203,
+            -9.732311248779297,
+            -51.624576568603516,
+            -35.76605224609375,
+            -23.42118263244629,
+            -31.96640396118164,
+            -32.20992660522461
+          ],
+          [
+            -24.42165184020996,
+            -10.062808990478516,
+            -52.54084777832031,
+            -36.21389389038086,
+            -23.58216094970703,
+            -32.878292083740234,
+            -33.002262115478516
+          ],
+          [
+            -24.806774139404297,
+            -10.06833267211914,
+            -53.43137741088867,
+            -36.29438781738281,
+            -23.675121307373047,
+            -33.4473876953125,
+            -33.2683219909668
+          ],
+          [
+            -24.799230575561523,
+            -10.070316314697266,
+            -54.24421691894531,
+            -36.37746810913086,
+            -23.685352325439453,
+            -33.9517822265625,
+            -33.42109298706055
+          ],
+          [
+            -24.732126235961914,
+            -10.07029914855957,
+            -54.64339828491211,
+            -36.4773063659668,
+            -23.682329177856445,
+            -34.212223052978516,
+            -33.55178451538086
+          ],
+          [
+            -24.3509578704834,
+            -10.070287704467773,
+            -54.56159210205078,
+            -36.60988998413086,
+            -23.682201385498047,
+            -34.409423828125,
+            -33.26326370239258
+          ],
+          [
+            -23.759092330932617,
+            -10.070272445678711,
+            -54.44050216674805,
+            -37.09710693359375,
+            -23.682085037231445,
+            -34.68149948120117,
+            -32.9384765625
+          ]
+        ],
+        "q99": [
+          [
+            3.203047752380371,
+            3.7505993843078613,
+            8.72464370727539,
+            9.948126792907715,
+            11.195253372192383,
+            5.477096080780029,
+            7.006125450134277
+          ],
+          [
+            5.7390570640563965,
+            6.635564804077148,
+            16.971813201904297,
+            18.642690658569336,
+            19.90561294555664,
+            10.401065826416016,
+            11.342881202697754
+          ],
+          [
+            8.532023429870605,
+            8.251216888427734,
+            24.75363540649414,
+            18.472192764282227,
+            19.774227142333984,
+            14.655866622924805,
+            15.70905590057373
+          ],
+          [
+            10.869905471801758,
+            9.230382919311523,
+            31.436893463134766,
+            19.715234756469727,
+            19.81361961364746,
+            18.759492874145508,
+            20.138126373291016
+          ],
+          [
+            12.716119766235352,
+            9.484061241149902,
+            37.04032516479492,
+            23.830785751342773,
+            19.81861114501953,
+            21.909305572509766,
+            23.74039649963379
+          ],
+          [
+            14.470975875854492,
+            9.510160446166992,
+            41.770057678222656,
+            27.173765182495117,
+            20.0618839263916,
+            24.20587921142578,
+            26.801916122436523
+          ],
+          [
+            16.024580001831055,
+            9.510558128356934,
+            45.34690475463867,
+            30.040668487548828,
+            21.872987747192383,
+            25.943317413330078,
+            28.912029266357422
+          ],
+          [
+            16.895124435424805,
+            9.513190269470215,
+            47.93039321899414,
+            32.24497985839844,
+            22.715166091918945,
+            27.396934509277344,
+            30.489883422851562
+          ],
+          [
+            17.761186599731445,
+            9.570911407470703,
+            50.055416107177734,
+            34.065670013427734,
+            23.788578033447266,
+            28.185653686523438,
+            31.546728134155273
+          ],
+          [
+            18.391948699951172,
+            9.712384223937988,
+            51.48924255371094,
+            35.40724182128906,
+            24.940738677978516,
+            28.75865364074707,
+            32.44330978393555
+          ],
+          [
+            19.106931686401367,
+            9.978495597839355,
+            52.564918518066406,
+            36.375064849853516,
+            25.661256790161133,
+            29.218664169311523,
+            33.03133010864258
+          ],
+          [
+            19.62432861328125,
+            10.027887344360352,
+            53.07897186279297,
+            36.976646423339844,
+            26.609840393066406,
+            29.49842071533203,
+            33.33045959472656
+          ],
+          [
+            19.768310546875,
+            10.027881622314453,
+            53.27336883544922,
+            37.44525909423828,
+            28.005691528320312,
+            29.92152214050293,
+            33.3731803894043
+          ],
+          [
+            19.806724548339844,
+            10.027848243713379,
+            53.31312942504883,
+            37.71003341674805,
+            29.430599212646484,
+            29.987842559814453,
+            33.3889274597168
+          ],
+          [
+            20.140544891357422,
+            10.027847290039062,
+            53.163578033447266,
+            37.966617584228516,
+            30.082942962646484,
+            30.24797248840332,
+            33.202571868896484
+          ],
+          [
+            20.164806365966797,
+            10.02784538269043,
+            52.79570770263672,
+            38.0666389465332,
+            30.232667922973633,
+            30.37464714050293,
+            32.94834899902344
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-9000/experiment_cfg/final_model_config.json

@@ -0,0 +1,53 @@
+{
+  "model_type": "Gr00tN1d6",
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Eagle-Block2A-2B-v2",
+  "backbone_model_type": "eagle",
+  "model_revision": null,
+  "tune_top_llm_layers": 4,
+  "backbone_embedding_dim": 2048,
+  "tune_llm": false,
+  "tune_visual": false,
+  "select_layer": 16,
+  "reproject_vision": false,
+  "use_flash_attention": true,
+  "load_bf16": true,
+  "collator_overwrite_image_inputs": false,
+  "eagle_collator": true,
+  "backbone_trainable_params_fp32": true,
+  "apply_sincos_state_encoding": true,
+  "use_relative_action": true,
+  "max_state_dim": 128,
+  "max_action_dim": 128,
+  "action_horizon": 50,
+  "hidden_size": 1024,
+  "input_embedding_dim": 1536,
+  "add_pos_embed": true,
+  "attn_dropout": 0.2,
+  "use_vlln": true,
+  "max_seq_len": 1024,
+  "use_alternate_vl_dit": true,
+  "attend_text_every_n_blocks": 2,
+  "diffusion_model_cfg": {
+    "attention_head_dim": 48,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "interleave_self_attention": true,
+    "norm_type": "ada_norm",
+    "num_attention_heads": 32,
+    "num_layers": 32,
+    "output_dim": 1024,
+    "positional_embeddings": null
+  },
+  "num_inference_timesteps": 4,
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_timestep_buckets": 1000,
+  "tune_projector": true,
+  "tune_diffusion_model": true,
+  "tune_vlln": true,
+  "state_dropout_prob": 0.0,
+  "state_additive_noise_scale": 0.0,
+  "max_num_embodiments": 32
+}

checkpoint-9000/processor_config.json

@@ -0,0 +1,453 @@
+{
+  "processor_class": "Gr00tN1d6Processor",
+  "processor_kwargs": {
+    "modality_configs": {
+      "behavior_r1_pro": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "observation.images.rgb.head_256_256",
+            "observation.images.rgb.left_wrist_256_256",
+            "observation.images.rgb.right_wrist_256_256"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "robot_pos",
+            "robot_ori_cos",
+            "robot_ori_sin",
+            "robot_2d_ori",
+            "robot_2d_ori_cos",
+            "robot_2d_ori_sin",
+            "robot_lin_vel",
+            "robot_ang_vel",
+            "arm_left_qpos",
+            "arm_left_qpos_sin",
+            "arm_left_qpos_cos",
+            "eef_left_pos",
+            "eef_left_quat",
+            "gripper_left_qpos",
+            "arm_right_qpos",
+            "arm_right_qpos_sin",
+            "arm_right_qpos_cos",
+            "eef_right_pos",
+            "eef_right_quat",
+            "gripper_right_qpos",
+            "trunk_qpos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31
+          ],
+          "modality_keys": [
+            "base",
+            "torso",
+            "left_arm",
+            "left_gripper",
+            "right_arm",
+            "right_gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "trunk_qpos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_left_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_right_qpos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "gr1": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "ego_view_bg_crop_pad_res256_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "left_arm",
+            "right_arm",
+            "left_hand",
+            "right_hand",
+            "waist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "robocasa_panda_omron": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "res256_image_side_0",
+            "res256_image_side_1",
+            "res256_image_wrist_0"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "end_effector_position_relative",
+            "end_effector_rotation_relative",
+            "gripper_qpos",
+            "base_position",
+            "base_rotation"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "end_effector_position",
+            "end_effector_rotation",
+            "gripper_close",
+            "base_motion",
+            "control_mode"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.action.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "reachy2": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front_cam"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "arm_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "arm_joints",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "arm_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": null
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      }
+    },
+    "image_crop_size": null,
+    "image_target_size": null,
+    "use_albumentations": true,
+    "random_rotation_angle": null,
+    "color_jitter_params": {
+      "brightness": 0.3,
+      "contrast": 0.4,
+      "saturation": 0.5,
+      "hue": 0.08
+    },
+    "shortest_image_edge": 256,
+    "crop_fraction": 0.95,
+    "model_name": "nvidia/Eagle-Block2A-2B-v2",
+    "model_type": "eagle",
+    "formalize_language": true,
+    "max_state_dim": 128,
+    "max_action_dim": 128,
+    "max_action_horizon": 50,
+    "use_percentiles": false,
+    "clip_outliers": true,
+    "apply_sincos_state_encoding": true,
+    "use_relative_action": true
+  }
+}

checkpoint-9000/wandb_config.json

@@ -0,0 +1 @@
+{"project": "finetune-gr00t-n1d6", "run_id": "reachy2"}

experiment_cfg/conf.yaml

@@ -147,8 +147,8 @@ data:
 training:
   output_dir: /tmp/groot_output/reachy2
   experiment_name: null
-  max_steps: 500
-  global_batch_size: 32
+  max_steps: 30000
+  global_batch_size: 64
   batch_size: null
   gradient_accumulation_steps: 1
   learning_rate: 0.0001
@@ -164,8 +164,8 @@ training:
   bf16: true
   eval_bf16: true
   logging_steps: 10
-  save_steps: 500
-  save_total_limit: 2
+  save_steps: 3000
+  save_total_limit: 5
   save_vl_model: false
   upload_checkpoints: false
   upload_every: 1000
@@ -186,7 +186,7 @@ training:
   use_ddp: false
   ddp_bucket_cap_mb: 100
   num_gpus: 1
-  dataloader_num_workers: 4
+  dataloader_num_workers: 8
   remove_unused_columns: false
   use_wandb: false
   wandb_project: finetune-gr00t-n1d6
@@ -199,5 +199,5 @@ training:
   - 0
   open_loop_eval_steps_per_traj: 100
   open_loop_eval_plot_indices: null
-max_steps: 500
-save_steps: 500
+max_steps: 30000
+save_steps: 3000

experiment_cfg/config.yaml

@@ -185,7 +185,7 @@ training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
   assert_loss_less_than: null
   batch_size: null
   bf16: true
-  dataloader_num_workers: 4
+  dataloader_num_workers: 8
   ddp_bucket_cap_mb: 100
   deepspeed_stage: 2
   enable_open_loop_eval: false
@@ -197,7 +197,7 @@ training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
   eval_strategy: 'no'
   experiment_name: null
   fp16: false
-  global_batch_size: 32
+  global_batch_size: 64
   gradient_accumulation_steps: 1
   gradient_checkpointing: false
   learning_rate: 0.0001
@@ -206,7 +206,7 @@ training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
   max_concurrent_uploads: 2
   max_grad_norm: 1.0
   max_retries: 3
-  max_steps: 500
+  max_steps: 30000
   num_gpus: 1
   open_loop_eval_plot_indices: null
   open_loop_eval_steps_per_traj: 100
@@ -217,8 +217,8 @@ training: !!python/object:gr00t.configs.training.training_config.TrainingConfig
   remove_unused_columns: false
   save_best_eval_metric_greater_is_better: true
   save_best_eval_metric_name: ''
-  save_steps: 500
-  save_total_limit: 2
+  save_steps: 3000
+  save_total_limit: 5
   save_vl_model: false
   start_from_checkpoint: nvidia/GR00T-N1.6-3B
   tf32: true