Upload folder using huggingface_hub

checkpoint-2000/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-2000/embodiment_id.json

@@ -0,0 +1,11 @@
+{
+  "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
+}

checkpoint-2000/experiment_cfg/conf.yaml

@@ -0,0 +1,210 @@
+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:
+    brightness: 0.3
+    contrast: 0.4
+    saturation: 0.5
+    hue: 0.08
+  use_albumentations_transforms: true
+  extra_augmentation_config: null
+  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:
+    - ./demo_data/SOarmtest/viis0002/pick_lego
+    embodiment_tag: new_embodiment
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    new_embodiment:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front
+        - wrist
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - single_arm
+        - gripper
+        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:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - 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:
+        - 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: 100000
+  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: torchcodec
+training:
+  output_dir: ./demo_data/SOarmtest/saved_checkpoints
+  experiment_name: null
+  max_steps: 5000
+  global_batch_size: 32
+  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: 2000
+  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: 4
+  remove_unused_columns: false
+  use_wandb: true
+  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: 5000
+save_steps: 2000

checkpoint-2000/experiment_cfg/config.yaml

@@ -0,0 +1,244 @@
+!!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:
+    - ./demo_data/SOarmtest/viis0002/pick_lego
+    dataset_type: physical_embodiment
+    embodiment_tag: new_embodiment
+    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:
+    new_embodiment:
+      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: null
+          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:
+        - single_arm
+        - 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:
+        - single_arm
+        - gripper
+        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
+        - wrist
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 100000
+  override_pretraining_statistics: false
+  random_chop: 0.0
+  seed: 42
+  shard_size: 1024
+  shuffle: true
+  subsample_ratio: 1.0
+  video_backend: torchcodec
+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:
+    brightness: 0.3
+    contrast: 0.4
+    hue: 0.08
+    saturation: 0.5
+  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
+  extra_augmentation_config: 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: 4
+  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: 32
+  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: 5000
+  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: ./demo_data/SOarmtest/saved_checkpoints
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_steps: 2000
+  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: true
+  wandb_project: finetune-gr00t-n1d6
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

checkpoint-2000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,824 @@
+{
+  "new_embodiment": {
+    "state": {
+      "single_arm": {
+        "min": [
+          -6.8131866455078125,
+          -102.19779968261719,
+          -40.43955993652344,
+          39.9560432434082,
+          -66.32967376708984
+        ],
+        "max": [
+          34.593406677246094,
+          42.769229888916016,
+          96.26373291015625,
+          98.06593322753906,
+          96.74725341796875
+        ],
+        "mean": [
+          10.6347642478831,
+          -40.60343800114444,
+          40.109735021899446,
+          66.37540221377535,
+          1.723290164849921
+        ],
+        "std": [
+          13.663472797431913,
+          53.514238207466356,
+          47.089223703864675,
+          15.49883594624875,
+          29.442594278901634
+        ],
+        "q01": [
+          -2.8227532818225196,
+          -102.04170769592945,
+          -18.832328638595566,
+          47.303414007272124,
+          -23.538191580154248
+        ],
+        "q99": [
+          31.17592093949654,
+          27.78409835219628,
+          96.09969459877085,
+          91.52694173019303,
+          34.804566667323705
+        ]
+      },
+      "gripper": {
+        "min": [
+          1.1597938537597656
+        ],
+        "max": [
+          30.476804733276367
+        ],
+        "mean": [
+          4.725540477545816
+        ],
+        "std": [
+          5.719213168100463
+        ],
+        "q01": [
+          1.219114189586798
+        ],
+        "q99": [
+          18.975632466717574
+        ]
+      }
+    },
+    "action": {
+      "single_arm": {
+        "min": [
+          -7.340659141540527,
+          -103.73626708984375,
+          -43.736263275146484,
+          38.63736343383789,
+          -67.47252655029297
+        ],
+        "max": [
+          35.20878982543945,
+          42.5494499206543,
+          96.74725341796875,
+          100.61538696289062,
+          97.27472686767578
+        ],
+        "mean": [
+          10.654706232867026,
+          -41.43508459659817,
+          38.95115310430438,
+          66.17248404139285,
+          1.6835930980913996
+        ],
+        "std": [
+          13.792301015873905,
+          53.27088262273843,
+          48.06070434685447,
+          15.735264444521983,
+          29.50880036285619
+        ],
+        "q01": [
+          -3.007477344415986,
+          -102.62527508767714,
+          -21.606171980367048,
+          46.25193771491276,
+          -24.028302588862974
+        ],
+        "q99": [
+          31.658033587037096,
+          27.49569149494191,
+          96.40327760792744,
+          92.26730265898277,
+          35.153253393452815
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.0
+        ],
+        "max": [
+          30.952381134033203
+        ],
+        "mean": [
+          3.5887609717193083
+        ],
+        "std": [
+          5.951741507358969
+        ],
+        "q01": [
+          0.7785521832201806
+        ],
+        "q99": [
+          19.41596701897272
+        ]
+      }
+    },
+    "relative_action": {
+      "single_arm": {
+        "min": [
+          [
+            -6.153846740722656,
+            -12.615386962890625,
+            -13.406593322753906,
+            -9.406593322753906,
+            -13.450549125671387
+          ],
+          [
+            -7.824176788330078,
+            -15.692306518554688,
+            -16.39560317993164,
+            -11.86812973022461,
+            -13.450549125671387
+          ],
+          [
+            -9.582418441772461,
+            -18.681320190429688,
+            -19.296703338623047,
+            -14.241756439208984,
+            -15.472527503967285
+          ],
+          [
+            -11.252748489379883,
+            -21.670331954956055,
+            -22.19780158996582,
+            -16.615379333496094,
+            -16.96703338623047
+          ],
+          [
+            -12.923077583312988,
+            -24.65934181213379,
+            -24.923078536987305,
+            -18.901100158691406,
+            -18.461538314819336
+          ],
+          [
+            -14.593405723571777,
+            -27.648353576660156,
+            -27.736263275146484,
+            -21.186813354492188,
+            -19.516483306884766
+          ],
+          [
+            -16.087913513183594,
+            -30.637365341186523,
+            -30.46154022216797,
+            -23.384613037109375,
+            -21.010990142822266
+          ],
+          [
+            -17.58241844177246,
+            -33.45055389404297,
+            -33.09890365600586,
+            -25.406593322753906,
+            -23.032970428466797
+          ],
+          [
+            -18.813186645507812,
+            -36.351646423339844,
+            -35.648353576660156,
+            -27.428573608398438,
+            -24.967037200927734
+          ],
+          [
+            -20.043954849243164,
+            -39.164833068847656,
+            -38.109893798828125,
+            -29.09890365600586,
+            -27.076923370361328
+          ],
+          [
+            -21.098899841308594,
+            -42.06593322753906,
+            -40.57142639160156,
+            -30.769229888916016,
+            -29.274723052978516
+          ],
+          [
+            -22.153844833374023,
+            -44.879119873046875,
+            -42.857139587402344,
+            -32.351646423339844,
+            -31.38461685180664
+          ],
+          [
+            -23.208789825439453,
+            -47.69231033325195,
+            -45.05494689941406,
+            -33.75823974609375,
+            -33.494503021240234
+          ],
+          [
+            -24.08791160583496,
+            -50.68132019042969,
+            -47.164833068847656,
+            -34.989013671875,
+            -35.51648712158203
+          ],
+          [
+            -24.967031478881836,
+            -53.67033004760742,
+            -49.098899841308594,
+            -36.04396057128906,
+            -37.626373291015625
+          ],
+          [
+            -25.758241653442383,
+            -56.747249603271484,
+            -50.94505310058594,
+            -37.01099395751953,
+            -39.648353576660156
+          ]
+        ],
+        "max": [
+          [
+            5.8021979331970215,
+            14.021980285644531,
+            12.527473449707031,
+            8.703296661376953,
+            9.670330047607422
+          ],
+          [
+            6.505494117736816,
+            17.626373291015625,
+            16.131866455078125,
+            11.164833068847656,
+            12.219779968261719
+          ],
+          [
+            7.82417631149292,
+            21.142852783203125,
+            19.824176788330078,
+            13.626373291015625,
+            14.681318283081055
+          ],
+          [
+            9.05494499206543,
+            24.395599365234375,
+            23.428571701049805,
+            15.824176788330078,
+            17.142858505249023
+          ],
+          [
+            10.373626708984375,
+            27.384613037109375,
+            26.94505500793457,
+            18.021976470947266,
+            19.69230842590332
+          ],
+          [
+            11.69230842590332,
+            30.197803497314453,
+            30.461538314819336,
+            20.043956756591797,
+            22.153846740722656
+          ],
+          [
+            13.01098918914795,
+            32.83516311645508,
+            33.97801971435547,
+            22.241756439208984,
+            24.703298568725586
+          ],
+          [
+            14.153846740722656,
+            35.384613037109375,
+            37.406593322753906,
+            24.35165023803711,
+            27.07692527770996
+          ],
+          [
+            15.296703338623047,
+            37.75823974609375,
+            40.74725341796875,
+            26.197803497314453,
+            29.362638473510742
+          ],
+          [
+            16.351648330688477,
+            40.04395294189453,
+            44.08790969848633,
+            27.956043243408203,
+            31.560440063476562
+          ],
+          [
+            17.318681716918945,
+            42.241756439208984,
+            47.340660095214844,
+            29.71428680419922,
+            33.67033004760742
+          ],
+          [
+            18.197803497314453,
+            44.527469635009766,
+            50.68131637573242,
+            31.384613037109375,
+            35.95604705810547
+          ],
+          [
+            18.901100158691406,
+            46.4615364074707,
+            53.93406295776367,
+            32.96703338623047,
+            38.5054931640625
+          ],
+          [
+            19.516483306884766,
+            48.48351287841797,
+            57.010986328125,
+            34.28571319580078,
+            41.14285659790039
+          ],
+          [
+            20.043956756591797,
+            50.32966613769531,
+            60.0,
+            35.604393005371094,
+            43.516483306884766
+          ],
+          [
+            20.395605087280273,
+            52.26373291015625,
+            62.989013671875,
+            36.659339904785156,
+            45.80220031738281
+          ]
+        ],
+        "mean": [
+          [
+            0.024023909121751785,
+            -0.8380120396614075,
+            -1.1888031959533691,
+            -0.20351642370224,
+            -0.03874167427420616
+          ],
+          [
+            0.02381421998143196,
+            -0.8381950259208679,
+            -1.1877086162567139,
+            -0.20610879361629486,
+            -0.038331978023052216
+          ],
+          [
+            0.02358066476881504,
+            -0.8384636640548706,
+            -1.186354160308838,
+            -0.20878157019615173,
+            -0.03791269287467003
+          ],
+          [
+            0.02331378683447838,
+            -0.8388302326202393,
+            -1.1847561597824097,
+            -0.21150696277618408,
+            -0.03749352693557739
+          ],
+          [
+            0.022984733805060387,
+            -0.8392952680587769,
+            -1.182924509048462,
+            -0.21426069736480713,
+            -0.037074267864227295
+          ],
+          [
+            0.02258922904729843,
+            -0.8398675322532654,
+            -1.1808984279632568,
+            -0.21703894436359406,
+            -0.03665487840771675
+          ],
+          [
+            0.02210799790918827,
+            -0.840530514717102,
+            -1.1787015199661255,
+            -0.21982122957706451,
+            -0.03624051436781883
+          ],
+          [
+            0.021555287763476372,
+            -0.8412824869155884,
+            -1.176315188407898,
+            -0.2226090133190155,
+            -0.03585929051041603
+          ],
+          [
+            0.020935896784067154,
+            -0.8421463966369629,
+            -1.1736962795257568,
+            -0.22543004155158997,
+            -0.035549700260162354
+          ],
+          [
+            0.020240070298314095,
+            -0.8431134819984436,
+            -1.1708482503890991,
+            -0.2282552719116211,
+            -0.03533996269106865
+          ],
+          [
+            0.01948723942041397,
+            -0.8441924452781677,
+            -1.1677296161651611,
+            -0.23114195466041565,
+            -0.03522548824548721
+          ],
+          [
+            0.018677080050110817,
+            -0.8453867435455322,
+            -1.164296269416809,
+            -0.23409080505371094,
+            -0.03516348451375961
+          ],
+          [
+            0.017790833488106728,
+            -0.846683919429779,
+            -1.1604851484298706,
+            -0.23711234331130981,
+            -0.035125281661748886
+          ],
+          [
+            0.016804957762360573,
+            -0.8481130599975586,
+            -1.156294584274292,
+            -0.24017632007598877,
+            -0.03509661555290222
+          ],
+          [
+            0.015680411830544472,
+            -0.8497316241264343,
+            -1.151717185974121,
+            -0.24328306317329407,
+            -0.035068120807409286
+          ],
+          [
+            0.01436557900160551,
+            -0.8516368269920349,
+            -1.1466983556747437,
+            -0.24642817676067352,
+            -0.035063356161117554
+          ]
+        ],
+        "std": [
+          [
+            1.0472843647003174,
+            3.3955397605895996,
+            3.209542989730835,
+            1.5478397607803345,
+            1.5488325357437134
+          ],
+          [
+            1.3025773763656616,
+            4.228511333465576,
+            3.9392380714416504,
+            1.9040354490280151,
+            1.9508968591690063
+          ],
+          [
+            1.5567840337753296,
+            5.061123371124268,
+            4.678185939788818,
+            2.259927749633789,
+            2.3464455604553223
+          ],
+          [
+            1.8073352575302124,
+            5.890169143676758,
+            5.418741226196289,
+            2.6112139225006104,
+            2.734940528869629
+          ],
+          [
+            2.053133726119995,
+            6.713922023773193,
+            6.156966209411621,
+            2.9563238620758057,
+            3.1171467304229736
+          ],
+          [
+            2.2936792373657227,
+            7.531452178955078,
+            6.89060115814209,
+            3.2943806648254395,
+            3.4939496517181396
+          ],
+          [
+            2.528961658477783,
+            8.342284202575684,
+            7.61845064163208,
+            3.6247124671936035,
+            3.8658618927001953
+          ],
+          [
+            2.7590219974517822,
+            9.146015167236328,
+            8.339773178100586,
+            3.9470086097717285,
+            4.233205795288086
+          ],
+          [
+            2.9840755462646484,
+            9.942627906799316,
+            9.053919792175293,
+            4.261190414428711,
+            4.596113681793213
+          ],
+          [
+            3.2043535709381104,
+            10.731932640075684,
+            9.760570526123047,
+            4.567060947418213,
+            4.954709053039551
+          ],
+          [
+            3.420076608657837,
+            11.51407527923584,
+            10.459877967834473,
+            4.864653587341309,
+            5.3089189529418945
+          ],
+          [
+            3.6315865516662598,
+            12.289117813110352,
+            11.151065826416016,
+            5.154054641723633,
+            5.65879487991333
+          ],
+          [
+            3.8390471935272217,
+            13.056876182556152,
+            11.834860801696777,
+            5.435141086578369,
+            6.004306316375732
+          ],
+          [
+            4.042593002319336,
+            13.817543029785156,
+            12.511225700378418,
+            5.708395004272461,
+            6.345395565032959
+          ],
+          [
+            4.24238395690918,
+            14.571209907531738,
+            13.179614067077637,
+            5.973555564880371,
+            6.682039737701416
+          ],
+          [
+            4.438711166381836,
+            15.317753791809082,
+            13.840078353881836,
+            6.230597496032715,
+            7.01429557800293
+          ]
+        ],
+        "q01": [
+          [
+            -3.472527027130127,
+            -9.978021621704102,
+            -9.978019714355469,
+            -5.098899841308594,
+            -5.36263370513916
+          ],
+          [
+            -4.483516693115234,
+            -12.175825119018555,
+            -12.263736724853516,
+            -6.417581558227539,
+            -6.857143402099609
+          ],
+          [
+            -5.450550079345703,
+            -14.373626708984375,
+            -14.373626708984375,
+            -7.736263275146484,
+            -8.263736724853516
+          ],
+          [
+            -6.4175825119018555,
+            -16.571428298950195,
+            -16.483516693115234,
+            -9.054946899414062,
+            -9.846153259277344
+          ],
+          [
+            -7.296703338623047,
+            -18.85714340209961,
+            -18.681316375732422,
+            -10.285713195800781,
+            -11.296703338623047
+          ],
+          [
+            -8.175825119018555,
+            -21.054946899414062,
+            -20.79120635986328,
+            -11.516483306884766,
+            -12.747251510620117
+          ],
+          [
+            -9.05494499206543,
+            -23.340660095214844,
+            -22.813186645507812,
+            -12.703296661376953,
+            -14.19780158996582
+          ],
+          [
+            -10.021978378295898,
+            -25.582416534423828,
+            -24.791210174560547,
+            -13.80219841003418,
+            -15.560441017150879
+          ],
+          [
+            -10.813186645507812,
+            -27.648353576660156,
+            -26.769229888916016,
+            -14.813186645507812,
+            -16.96703338623047
+          ],
+          [
+            -11.780220031738281,
+            -29.934066772460938,
+            -28.659339904785156,
+            -15.824176788330078,
+            -18.373626708984375
+          ],
+          [
+            -12.615385055541992,
+            -32.21977996826172,
+            -30.593406677246094,
+            -16.703292846679688,
+            -19.516483306884766
+          ],
+          [
+            -13.450550079345703,
+            -34.32966995239258,
+            -32.395599365234375,
+            -17.670326232910156,
+            -20.791208267211914
+          ],
+          [
+            -14.329669952392578,
+            -36.52747344970703,
+            -34.19780349731445,
+            -18.549448013305664,
+            -21.802196502685547
+          ],
+          [
+            -15.164834976196289,
+            -38.68132019042969,
+            -35.82417297363281,
+            -19.20879364013672,
+            -22.901100158691406
+          ],
+          [
+            -15.95604419708252,
+            -40.92308044433594,
+            -37.626373291015625,
+            -20.043954849243164,
+            -24.131868362426758
+          ],
+          [
+            -16.703296661376953,
+            -42.989009857177734,
+            -39.25274658203125,
+            -20.74725341796875,
+            -25.450550079345703
+          ]
+        ],
+        "q99": [
+          [
+            3.4285709857940674,
+            7.5604400634765625,
+            7.076923847198486,
+            4.131866455078125,
+            5.494504928588867
+          ],
+          [
+            4.307692527770996,
+            9.802200317382812,
+            9.538460731506348,
+            5.274726867675781,
+            7.032966613769531
+          ],
+          [
+            5.1868133544921875,
+            12.087913513183594,
+            11.868131637573242,
+            6.5054931640625,
+            8.527473449707031
+          ],
+          [
+            6.065934181213379,
+            14.197805404663086,
+            14.285714149475098,
+            7.648347854614258,
+            9.934066772460938
+          ],
+          [
+            6.857143402099609,
+            16.483516693115234,
+            16.571428298950195,
+            8.703292846679688,
+            11.340660095214844
+          ],
+          [
+            7.736264705657959,
+            18.593406677246094,
+            18.85714340209961,
+            9.846153259277344,
+            12.835165023803711
+          ],
+          [
+            8.615385055541992,
+            20.52747344970703,
+            21.14285659790039,
+            10.813188552856445,
+            14.241758346557617
+          ],
+          [
+            9.494504928588867,
+            22.725276947021484,
+            23.428571701049805,
+            11.824176788330078,
+            15.648351669311523
+          ],
+          [
+            10.285714149475098,
+            24.879121780395508,
+            25.714282989501953,
+            12.835166931152344,
+            17.05494499206543
+          ],
+          [
+            11.076923370361328,
+            26.85714340209961,
+            27.868133544921875,
+            13.758243560791016,
+            18.461538314819336
+          ],
+          [
+            11.7362642288208,
+            28.70329475402832,
+            30.109888076782227,
+            14.725275039672852,
+            19.868133544921875
+          ],
+          [
+            12.395605087280273,
+            30.813186645507812,
+            32.17582321166992,
+            15.560440063476562,
+            21.186813354492188
+          ],
+          [
+            13.098901748657227,
+            32.52747344970703,
+            34.19779968261719,
+            16.351646423339844,
+            22.549449920654297
+          ],
+          [
+            13.714285850524902,
+            34.54945373535156,
+            36.21977996826172,
+            17.142860412597656,
+            24.0
+          ],
+          [
+            14.329669952392578,
+            36.57143020629883,
+            38.153846740722656,
+            17.846155166625977,
+            25.230770111083984
+          ],
+          [
+            14.769231796264648,
+            38.28571319580078,
+            40.17582702636719,
+            18.637367248535156,
+            26.41758155822754
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-2000/experiment_cfg/final_model_config.json

@@ -0,0 +1,54 @@
+{
+  "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,
+  "extra_augmentation_config": null,
+  "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-2000/model-00001-of-00002.safetensors

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

checkpoint-2000/model-00002-of-00002.safetensors

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

checkpoint-2000/optimizer.pt

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

checkpoint-2000/processor_config.json

@@ -0,0 +1,455 @@
+{
+  "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
+        }
+      },
+      "new_embodiment": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front",
+            "wrist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "single_arm",
+            "gripper"
+          ],
+          "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": [
+            "single_arm",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "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": [
+            "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-2000/rng_state.pth

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

checkpoint-2000/scheduler.pt

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

checkpoint-2000/trainer_state.json

@@ -0,0 +1,1234 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.4,
+  "eval_steps": 500,
+  "global_step": 2000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "grad_norm": 2.497777223587036,
+      "learning_rate": 3.6e-06,
+      "loss": 1.1476,
+      "step": 10
+    },
+    {
+      "grad_norm": 0.5899417996406555,
+      "learning_rate": 7.6e-06,
+      "loss": 1.1021,
+      "step": 20
+    },
+    {
+      "grad_norm": 0.6255532503128052,
+      "learning_rate": 1.16e-05,
+      "loss": 1.0779,
+      "step": 30
+    },
+    {
+      "grad_norm": 0.5589399933815002,
+      "learning_rate": 1.56e-05,
+      "loss": 1.0599,
+      "step": 40
+    },
+    {
+      "grad_norm": 0.6063516736030579,
+      "learning_rate": 1.9600000000000002e-05,
+      "loss": 1.0856,
+      "step": 50
+    },
+    {
+      "grad_norm": 0.5875065326690674,
+      "learning_rate": 2.36e-05,
+      "loss": 1.0636,
+      "step": 60
+    },
+    {
+      "grad_norm": 0.6206650733947754,
+      "learning_rate": 2.7600000000000003e-05,
+      "loss": 1.0586,
+      "step": 70
+    },
+    {
+      "grad_norm": 0.6976884007453918,
+      "learning_rate": 3.16e-05,
+      "loss": 1.0581,
+      "step": 80
+    },
+    {
+      "grad_norm": 0.4975197911262512,
+      "learning_rate": 3.56e-05,
+      "loss": 1.0338,
+      "step": 90
+    },
+    {
+      "grad_norm": 0.5175010561943054,
+      "learning_rate": 3.960000000000001e-05,
+      "loss": 1.0363,
+      "step": 100
+    },
+    {
+      "grad_norm": 3.203188180923462,
+      "learning_rate": 4.36e-05,
+      "loss": 0.9943,
+      "step": 110
+    },
+    {
+      "grad_norm": 1.0046463012695312,
+      "learning_rate": 4.76e-05,
+      "loss": 0.9498,
+      "step": 120
+    },
+    {
+      "grad_norm": 0.824455976486206,
+      "learning_rate": 5.16e-05,
+      "loss": 0.9141,
+      "step": 130
+    },
+    {
+      "grad_norm": 0.933220386505127,
+      "learning_rate": 5.560000000000001e-05,
+      "loss": 0.8966,
+      "step": 140
+    },
+    {
+      "grad_norm": 1.6803524494171143,
+      "learning_rate": 5.96e-05,
+      "loss": 0.8456,
+      "step": 150
+    },
+    {
+      "grad_norm": 1.3335356712341309,
+      "learning_rate": 6.36e-05,
+      "loss": 0.7857,
+      "step": 160
+    },
+    {
+      "grad_norm": 1.7373239994049072,
+      "learning_rate": 6.76e-05,
+      "loss": 0.6714,
+      "step": 170
+    },
+    {
+      "grad_norm": 1.2455403804779053,
+      "learning_rate": 7.16e-05,
+      "loss": 0.5578,
+      "step": 180
+    },
+    {
+      "grad_norm": 2.673060178756714,
+      "learning_rate": 7.560000000000001e-05,
+      "loss": 0.4206,
+      "step": 190
+    },
+    {
+      "grad_norm": 1.7698984146118164,
+      "learning_rate": 7.960000000000001e-05,
+      "loss": 0.3168,
+      "step": 200
+    },
+    {
+      "grad_norm": 1.393641471862793,
+      "learning_rate": 8.36e-05,
+      "loss": 0.2192,
+      "step": 210
+    },
+    {
+      "grad_norm": 1.3259373903274536,
+      "learning_rate": 8.76e-05,
+      "loss": 0.1876,
+      "step": 220
+    },
+    {
+      "grad_norm": 0.9707556366920471,
+      "learning_rate": 9.16e-05,
+      "loss": 0.1548,
+      "step": 230
+    },
+    {
+      "grad_norm": 0.8586066961288452,
+      "learning_rate": 9.56e-05,
+      "loss": 0.1401,
+      "step": 240
+    },
+    {
+      "grad_norm": 0.8358578085899353,
+      "learning_rate": 9.960000000000001e-05,
+      "loss": 0.1216,
+      "step": 250
+    },
+    {
+      "grad_norm": 0.6485460996627808,
+      "learning_rate": 9.999911419878559e-05,
+      "loss": 0.1228,
+      "step": 260
+    },
+    {
+      "grad_norm": 1.0369858741760254,
+      "learning_rate": 9.999605221019081e-05,
+      "loss": 0.1396,
+      "step": 270
+    },
+    {
+      "grad_norm": 0.8534221053123474,
+      "learning_rate": 9.999080323230761e-05,
+      "loss": 0.1208,
+      "step": 280
+    },
+    {
+      "grad_norm": 0.8548014760017395,
+      "learning_rate": 9.998336749474329e-05,
+      "loss": 0.1341,
+      "step": 290
+    },
+    {
+      "grad_norm": 0.9402936100959778,
+      "learning_rate": 9.997374532276107e-05,
+      "loss": 0.102,
+      "step": 300
+    },
+    {
+      "grad_norm": 0.7149423956871033,
+      "learning_rate": 9.996193713726596e-05,
+      "loss": 0.1092,
+      "step": 310
+    },
+    {
+      "grad_norm": 0.7244396805763245,
+      "learning_rate": 9.994794345478624e-05,
+      "loss": 0.0924,
+      "step": 320
+    },
+    {
+      "grad_norm": 0.5502927899360657,
+      "learning_rate": 9.99317648874509e-05,
+      "loss": 0.0838,
+      "step": 330
+    },
+    {
+      "grad_norm": 0.6544800400733948,
+      "learning_rate": 9.991340214296292e-05,
+      "loss": 0.0873,
+      "step": 340
+    },
+    {
+      "grad_norm": 0.7180399894714355,
+      "learning_rate": 9.989285602456819e-05,
+      "loss": 0.0734,
+      "step": 350
+    },
+    {
+      "grad_norm": 0.6445586681365967,
+      "learning_rate": 9.98701274310205e-05,
+      "loss": 0.0794,
+      "step": 360
+    },
+    {
+      "grad_norm": 0.5352494716644287,
+      "learning_rate": 9.984521735654218e-05,
+      "loss": 0.0907,
+      "step": 370
+    },
+    {
+      "grad_norm": 0.6672635674476624,
+      "learning_rate": 9.981812689078057e-05,
+      "loss": 0.0848,
+      "step": 380
+    },
+    {
+      "grad_norm": 0.5291720628738403,
+      "learning_rate": 9.978885721876041e-05,
+      "loss": 0.086,
+      "step": 390
+    },
+    {
+      "grad_norm": 0.39044567942619324,
+      "learning_rate": 9.975740962083198e-05,
+      "loss": 0.0787,
+      "step": 400
+    },
+    {
+      "grad_norm": 0.7575682997703552,
+      "learning_rate": 9.972378547261504e-05,
+      "loss": 0.0836,
+      "step": 410
+    },
+    {
+      "grad_norm": 0.6050339341163635,
+      "learning_rate": 9.968798624493885e-05,
+      "loss": 0.0738,
+      "step": 420
+    },
+    {
+      "grad_norm": 0.594001054763794,
+      "learning_rate": 9.965001350377753e-05,
+      "loss": 0.0817,
+      "step": 430
+    },
+    {
+      "grad_norm": 0.5727436542510986,
+      "learning_rate": 9.960986891018183e-05,
+      "loss": 0.0728,
+      "step": 440
+    },
+    {
+      "grad_norm": 0.4483703076839447,
+      "learning_rate": 9.95675542202063e-05,
+      "loss": 0.0747,
+      "step": 450
+    },
+    {
+      "grad_norm": 0.3992336094379425,
+      "learning_rate": 9.952307128483256e-05,
+      "loss": 0.0796,
+      "step": 460
+    },
+    {
+      "grad_norm": 0.5178577303886414,
+      "learning_rate": 9.947642204988835e-05,
+      "loss": 0.0833,
+      "step": 470
+    },
+    {
+      "grad_norm": 0.4497125446796417,
+      "learning_rate": 9.942760855596226e-05,
+      "loss": 0.0787,
+      "step": 480
+    },
+    {
+      "grad_norm": 0.5343725085258484,
+      "learning_rate": 9.937663293831471e-05,
+      "loss": 0.0808,
+      "step": 490
+    },
+    {
+      "grad_norm": 0.47977861762046814,
+      "learning_rate": 9.932349742678433e-05,
+      "loss": 0.0701,
+      "step": 500
+    },
+    {
+      "grad_norm": 0.5831079483032227,
+      "learning_rate": 9.926820434569051e-05,
+      "loss": 0.0742,
+      "step": 510
+    },
+    {
+      "grad_norm": 0.47160646319389343,
+      "learning_rate": 9.921075611373179e-05,
+      "loss": 0.067,
+      "step": 520
+    },
+    {
+      "grad_norm": 0.35774683952331543,
+      "learning_rate": 9.915115524387988e-05,
+      "loss": 0.0709,
+      "step": 530
+    },
+    {
+      "grad_norm": 0.5747243762016296,
+      "learning_rate": 9.908940434326997e-05,
+      "loss": 0.0668,
+      "step": 540
+    },
+    {
+      "grad_norm": 0.42538684606552124,
+      "learning_rate": 9.902550611308645e-05,
+      "loss": 0.0676,
+      "step": 550
+    },
+    {
+      "grad_norm": 0.5689398646354675,
+      "learning_rate": 9.895946334844494e-05,
+      "loss": 0.0631,
+      "step": 560
+    },
+    {
+      "grad_norm": 0.648316502571106,
+      "learning_rate": 9.889127893826989e-05,
+      "loss": 0.0643,
+      "step": 570
+    },
+    {
+      "grad_norm": 0.5649459362030029,
+      "learning_rate": 9.882095586516831e-05,
+      "loss": 0.0679,
+      "step": 580
+    },
+    {
+      "grad_norm": 0.42304474115371704,
+      "learning_rate": 9.874849720529921e-05,
+      "loss": 0.0744,
+      "step": 590
+    },
+    {
+      "grad_norm": 0.5694190263748169,
+      "learning_rate": 9.867390612823914e-05,
+      "loss": 0.0683,
+      "step": 600
+    },
+    {
+      "grad_norm": 0.5124210715293884,
+      "learning_rate": 9.859718589684344e-05,
+      "loss": 0.0797,
+      "step": 610
+    },
+    {
+      "grad_norm": 0.6514238119125366,
+      "learning_rate": 9.851833986710353e-05,
+      "loss": 0.0644,
+      "step": 620
+    },
+    {
+      "grad_norm": 0.49233269691467285,
+      "learning_rate": 9.843737148800023e-05,
+      "loss": 0.0568,
+      "step": 630
+    },
+    {
+      "grad_norm": 0.40201443433761597,
+      "learning_rate": 9.835428430135271e-05,
+      "loss": 0.063,
+      "step": 640
+    },
+    {
+      "grad_norm": 0.4645858407020569,
+      "learning_rate": 9.82690819416637e-05,
+      "loss": 0.0641,
+      "step": 650
+    },
+    {
+      "grad_norm": 0.5921373963356018,
+      "learning_rate": 9.818176813596041e-05,
+      "loss": 0.0622,
+      "step": 660
+    },
+    {
+      "grad_norm": 0.5843824148178101,
+      "learning_rate": 9.809234670363159e-05,
+      "loss": 0.0724,
+      "step": 670
+    },
+    {
+      "grad_norm": 0.6411054730415344,
+      "learning_rate": 9.800082155626034e-05,
+      "loss": 0.0697,
+      "step": 680
+    },
+    {
+      "grad_norm": 0.4261399507522583,
+      "learning_rate": 9.790719669745312e-05,
+      "loss": 0.0641,
+      "step": 690
+    },
+    {
+      "grad_norm": 0.3344235122203827,
+      "learning_rate": 9.781147622266455e-05,
+      "loss": 0.0642,
+      "step": 700
+    },
+    {
+      "grad_norm": 0.4839865565299988,
+      "learning_rate": 9.771366431901831e-05,
+      "loss": 0.0734,
+      "step": 710
+    },
+    {
+      "grad_norm": 0.3950216472148895,
+      "learning_rate": 9.761376526512394e-05,
+      "loss": 0.063,
+      "step": 720
+    },
+    {
+      "grad_norm": 0.45863887667655945,
+      "learning_rate": 9.751178343088963e-05,
+      "loss": 0.0598,
+      "step": 730
+    },
+    {
+      "grad_norm": 0.414181113243103,
+      "learning_rate": 9.740772327733123e-05,
+      "loss": 0.0592,
+      "step": 740
+    },
+    {
+      "grad_norm": 0.48416581749916077,
+      "learning_rate": 9.730158935637697e-05,
+      "loss": 0.0576,
+      "step": 750
+    },
+    {
+      "grad_norm": 0.40759462118148804,
+      "learning_rate": 9.719338631066834e-05,
+      "loss": 0.0522,
+      "step": 760
+    },
+    {
+      "grad_norm": 0.4843514859676361,
+      "learning_rate": 9.708311887335713e-05,
+      "loss": 0.0551,
+      "step": 770
+    },
+    {
+      "grad_norm": 0.37705329060554504,
+      "learning_rate": 9.697079186789823e-05,
+      "loss": 0.0601,
+      "step": 780
+    },
+    {
+      "grad_norm": 0.41809579730033875,
+      "learning_rate": 9.685641020783876e-05,
+      "loss": 0.0567,
+      "step": 790
+    },
+    {
+      "grad_norm": 0.4055332839488983,
+      "learning_rate": 9.67399788966031e-05,
+      "loss": 0.0569,
+      "step": 800
+    },
+    {
+      "grad_norm": 0.384504109621048,
+      "learning_rate": 9.662150302727395e-05,
+      "loss": 0.0586,
+      "step": 810
+    },
+    {
+      "grad_norm": 0.38468998670578003,
+      "learning_rate": 9.650098778236968e-05,
+      "loss": 0.0531,
+      "step": 820
+    },
+    {
+      "grad_norm": 0.5237128734588623,
+      "learning_rate": 9.637843843361749e-05,
+      "loss": 0.0552,
+      "step": 830
+    },
+    {
+      "grad_norm": 0.4317338466644287,
+      "learning_rate": 9.62538603417229e-05,
+      "loss": 0.0509,
+      "step": 840
+    },
+    {
+      "grad_norm": 0.3871803879737854,
+      "learning_rate": 9.612725895613526e-05,
+      "loss": 0.0625,
+      "step": 850
+    },
+    {
+      "grad_norm": 0.4323013126850128,
+      "learning_rate": 9.599863981480926e-05,
+      "loss": 0.0604,
+      "step": 860
+    },
+    {
+      "grad_norm": 0.36297470331192017,
+      "learning_rate": 9.586800854396283e-05,
+      "loss": 0.0545,
+      "step": 870
+    },
+    {
+      "grad_norm": 0.37466663122177124,
+      "learning_rate": 9.573537085783095e-05,
+      "loss": 0.0529,
+      "step": 880
+    },
+    {
+      "grad_norm": 0.3929672837257385,
+      "learning_rate": 9.560073255841571e-05,
+      "loss": 0.05,
+      "step": 890
+    },
+    {
+      "grad_norm": 0.4040852189064026,
+      "learning_rate": 9.546409953523247e-05,
+      "loss": 0.0603,
+      "step": 900
+    },
+    {
+      "grad_norm": 0.6248286366462708,
+      "learning_rate": 9.532547776505229e-05,
+      "loss": 0.0492,
+      "step": 910
+    },
+    {
+      "grad_norm": 0.5975190997123718,
+      "learning_rate": 9.518487331164048e-05,
+      "loss": 0.0538,
+      "step": 920
+    },
+    {
+      "grad_norm": 0.42542824149131775,
+      "learning_rate": 9.504229232549134e-05,
+      "loss": 0.0523,
+      "step": 930
+    },
+    {
+      "grad_norm": 0.328101247549057,
+      "learning_rate": 9.489774104355909e-05,
+      "loss": 0.0519,
+      "step": 940
+    },
+    {
+      "grad_norm": 0.5055563449859619,
+      "learning_rate": 9.475122578898507e-05,
+      "loss": 0.0603,
+      "step": 950
+    },
+    {
+      "grad_norm": 0.5114820003509521,
+      "learning_rate": 9.460275297082119e-05,
+      "loss": 0.0486,
+      "step": 960
+    },
+    {
+      "grad_norm": 0.30745837092399597,
+      "learning_rate": 9.445232908374948e-05,
+      "loss": 0.0515,
+      "step": 970
+    },
+    {
+      "grad_norm": 0.5132862329483032,
+      "learning_rate": 9.429996070779808e-05,
+      "loss": 0.0534,
+      "step": 980
+    },
+    {
+      "grad_norm": 0.3776204586029053,
+      "learning_rate": 9.414565450805333e-05,
+      "loss": 0.0543,
+      "step": 990
+    },
+    {
+      "grad_norm": 0.39591264724731445,
+      "learning_rate": 9.398941723436831e-05,
+      "loss": 0.051,
+      "step": 1000
+    },
+    {
+      "grad_norm": 0.3092094957828522,
+      "learning_rate": 9.383125572106752e-05,
+      "loss": 0.0485,
+      "step": 1010
+    },
+    {
+      "grad_norm": 0.4166562855243683,
+      "learning_rate": 9.367117688664791e-05,
+      "loss": 0.054,
+      "step": 1020
+    },
+    {
+      "grad_norm": 0.38996878266334534,
+      "learning_rate": 9.35091877334763e-05,
+      "loss": 0.0506,
+      "step": 1030
+    },
+    {
+      "grad_norm": 0.4977191686630249,
+      "learning_rate": 9.334529534748297e-05,
+      "loss": 0.0543,
+      "step": 1040
+    },
+    {
+      "grad_norm": 0.47602540254592896,
+      "learning_rate": 9.317950689785188e-05,
+      "loss": 0.048,
+      "step": 1050
+    },
+    {
+      "grad_norm": 0.3852064907550812,
+      "learning_rate": 9.301182963670688e-05,
+      "loss": 0.0563,
+      "step": 1060
+    },
+    {
+      "grad_norm": 0.2994341254234314,
+      "learning_rate": 9.284227089879456e-05,
+      "loss": 0.0528,
+      "step": 1070
+    },
+    {
+      "grad_norm": 0.4053724706172943,
+      "learning_rate": 9.26708381011634e-05,
+      "loss": 0.0546,
+      "step": 1080
+    },
+    {
+      "grad_norm": 0.28891879320144653,
+      "learning_rate": 9.249753874283937e-05,
+      "loss": 0.0435,
+      "step": 1090
+    },
+    {
+      "grad_norm": 0.2600522041320801,
+      "learning_rate": 9.232238040449779e-05,
+      "loss": 0.0459,
+      "step": 1100
+    },
+    {
+      "grad_norm": 0.5400200486183167,
+      "learning_rate": 9.214537074813181e-05,
+      "loss": 0.0553,
+      "step": 1110
+    },
+    {
+      "grad_norm": 0.4278520941734314,
+      "learning_rate": 9.196651751671724e-05,
+      "loss": 0.0543,
+      "step": 1120
+    },
+    {
+      "grad_norm": 0.43360674381256104,
+      "learning_rate": 9.178582853387384e-05,
+      "loss": 0.058,
+      "step": 1130
+    },
+    {
+      "grad_norm": 0.5356864929199219,
+      "learning_rate": 9.160331170352304e-05,
+      "loss": 0.0607,
+      "step": 1140
+    },
+    {
+      "grad_norm": 0.37267163395881653,
+      "learning_rate": 9.141897500954229e-05,
+      "loss": 0.0527,
+      "step": 1150
+    },
+    {
+      "grad_norm": 0.4231492877006531,
+      "learning_rate": 9.123282651541576e-05,
+      "loss": 0.0515,
+      "step": 1160
+    },
+    {
+      "grad_norm": 0.41324272751808167,
+      "learning_rate": 9.104487436388161e-05,
+      "loss": 0.0477,
+      "step": 1170
+    },
+    {
+      "grad_norm": 0.2980903089046478,
+      "learning_rate": 9.085512677657582e-05,
+      "loss": 0.0545,
+      "step": 1180
+    },
+    {
+      "grad_norm": 0.3594469428062439,
+      "learning_rate": 9.066359205367258e-05,
+      "loss": 0.0519,
+      "step": 1190
+    },
+    {
+      "grad_norm": 0.3295166492462158,
+      "learning_rate": 9.047027857352112e-05,
+      "loss": 0.0489,
+      "step": 1200
+    },
+    {
+      "grad_norm": 0.35804498195648193,
+      "learning_rate": 9.027519479227935e-05,
+      "loss": 0.0542,
+      "step": 1210
+    },
+    {
+      "grad_norm": 0.5596129298210144,
+      "learning_rate": 9.007834924354383e-05,
+      "loss": 0.0492,
+      "step": 1220
+    },
+    {
+      "grad_norm": 0.32689106464385986,
+      "learning_rate": 8.987975053797655e-05,
+      "loss": 0.0478,
+      "step": 1230
+    },
+    {
+      "grad_norm": 0.3586772680282593,
+      "learning_rate": 8.967940736292825e-05,
+      "loss": 0.0518,
+      "step": 1240
+    },
+    {
+      "grad_norm": 0.4046182632446289,
+      "learning_rate": 8.947732848205846e-05,
+      "loss": 0.0561,
+      "step": 1250
+    },
+    {
+      "grad_norm": 0.2832528352737427,
+      "learning_rate": 8.927352273495204e-05,
+      "loss": 0.0481,
+      "step": 1260
+    },
+    {
+      "grad_norm": 0.38299012184143066,
+      "learning_rate": 8.906799903673265e-05,
+      "loss": 0.0543,
+      "step": 1270
+    },
+    {
+      "grad_norm": 0.39338234066963196,
+      "learning_rate": 8.88607663776726e-05,
+      "loss": 0.053,
+      "step": 1280
+    },
+    {
+      "grad_norm": 0.326089471578598,
+      "learning_rate": 8.865183382279978e-05,
+      "loss": 0.0584,
+      "step": 1290
+    },
+    {
+      "grad_norm": 0.5139309167861938,
+      "learning_rate": 8.844121051150096e-05,
+      "loss": 0.0508,
+      "step": 1300
+    },
+    {
+      "grad_norm": 0.2685113251209259,
+      "learning_rate": 8.822890565712211e-05,
+      "loss": 0.0473,
+      "step": 1310
+    },
+    {
+      "grad_norm": 0.2396756261587143,
+      "learning_rate": 8.801492854656536e-05,
+      "loss": 0.047,
+      "step": 1320
+    },
+    {
+      "grad_norm": 0.2990108132362366,
+      "learning_rate": 8.779928853988268e-05,
+      "loss": 0.0422,
+      "step": 1330
+    },
+    {
+      "grad_norm": 0.25585711002349854,
+      "learning_rate": 8.758199506986655e-05,
+      "loss": 0.0443,
+      "step": 1340
+    },
+    {
+      "grad_norm": 0.2606576681137085,
+      "learning_rate": 8.73630576416373e-05,
+      "loss": 0.0399,
+      "step": 1350
+    },
+    {
+      "grad_norm": 0.43027952313423157,
+      "learning_rate": 8.714248583222726e-05,
+      "loss": 0.0449,
+      "step": 1360
+    },
+    {
+      "grad_norm": 0.3107929229736328,
+      "learning_rate": 8.692028929016196e-05,
+      "loss": 0.039,
+      "step": 1370
+    },
+    {
+      "grad_norm": 0.26511290669441223,
+      "learning_rate": 8.669647773503797e-05,
+      "loss": 0.0466,
+      "step": 1380
+    },
+    {
+      "grad_norm": 0.40483543276786804,
+      "learning_rate": 8.647106095709773e-05,
+      "loss": 0.04,
+      "step": 1390
+    },
+    {
+      "grad_norm": 0.3609035909175873,
+      "learning_rate": 8.624404881680139e-05,
+      "loss": 0.0468,
+      "step": 1400
+    },
+    {
+      "grad_norm": 0.3585107624530792,
+      "learning_rate": 8.601545124439535e-05,
+      "loss": 0.0396,
+      "step": 1410
+    },
+    {
+      "grad_norm": 0.39873549342155457,
+      "learning_rate": 8.5785278239478e-05,
+      "loss": 0.0506,
+      "step": 1420
+    },
+    {
+      "grad_norm": 0.26610010862350464,
+      "learning_rate": 8.555353987056224e-05,
+      "loss": 0.0446,
+      "step": 1430
+    },
+    {
+      "grad_norm": 0.29464486241340637,
+      "learning_rate": 8.532024627463505e-05,
+      "loss": 0.0449,
+      "step": 1440
+    },
+    {
+      "grad_norm": 0.3142518699169159,
+      "learning_rate": 8.508540765671407e-05,
+      "loss": 0.0449,
+      "step": 1450
+    },
+    {
+      "grad_norm": 0.3424524962902069,
+      "learning_rate": 8.484903428940121e-05,
+      "loss": 0.0462,
+      "step": 1460
+    },
+    {
+      "grad_norm": 0.4664880037307739,
+      "learning_rate": 8.461113651243334e-05,
+      "loss": 0.0499,
+      "step": 1470
+    },
+    {
+      "grad_norm": 0.28201568126678467,
+      "learning_rate": 8.437172473222987e-05,
+      "loss": 0.0467,
+      "step": 1480
+    },
+    {
+      "grad_norm": 0.4245680868625641,
+      "learning_rate": 8.413080942143767e-05,
+      "loss": 0.0439,
+      "step": 1490
+    },
+    {
+      "grad_norm": 0.42034053802490234,
+      "learning_rate": 8.388840111847288e-05,
+      "loss": 0.0532,
+      "step": 1500
+    },
+    {
+      "grad_norm": 0.3836017847061157,
+      "learning_rate": 8.364451042705998e-05,
+      "loss": 0.049,
+      "step": 1510
+    },
+    {
+      "grad_norm": 0.26328060030937195,
+      "learning_rate": 8.33991480157679e-05,
+      "loss": 0.0427,
+      "step": 1520
+    },
+    {
+      "grad_norm": 0.38004162907600403,
+      "learning_rate": 8.315232461754338e-05,
+      "loss": 0.0473,
+      "step": 1530
+    },
+    {
+      "grad_norm": 0.45609551668167114,
+      "learning_rate": 8.290405102924144e-05,
+      "loss": 0.0533,
+      "step": 1540
+    },
+    {
+      "grad_norm": 0.20536701381206512,
+      "learning_rate": 8.265433811115316e-05,
+      "loss": 0.0436,
+      "step": 1550
+    },
+    {
+      "grad_norm": 0.3798106908798218,
+      "learning_rate": 8.240319678653049e-05,
+      "loss": 0.0439,
+      "step": 1560
+    },
+    {
+      "grad_norm": 0.28516557812690735,
+      "learning_rate": 8.215063804110857e-05,
+      "loss": 0.0425,
+      "step": 1570
+    },
+    {
+      "grad_norm": 0.27644988894462585,
+      "learning_rate": 8.189667292262512e-05,
+      "loss": 0.0403,
+      "step": 1580
+    },
+    {
+      "grad_norm": 0.40204596519470215,
+      "learning_rate": 8.164131254033716e-05,
+      "loss": 0.0474,
+      "step": 1590
+    },
+    {
+      "grad_norm": 0.43881192803382874,
+      "learning_rate": 8.138456806453503e-05,
+      "loss": 0.0403,
+      "step": 1600
+    },
+    {
+      "grad_norm": 0.2845113277435303,
+      "learning_rate": 8.112645072605386e-05,
+      "loss": 0.0392,
+      "step": 1610
+    },
+    {
+      "grad_norm": 0.29198166728019714,
+      "learning_rate": 8.086697181578222e-05,
+      "loss": 0.0348,
+      "step": 1620
+    },
+    {
+      "grad_norm": 0.3617711663246155,
+      "learning_rate": 8.060614268416823e-05,
+      "loss": 0.0376,
+      "step": 1630
+    },
+    {
+      "grad_norm": 0.404875785112381,
+      "learning_rate": 8.034397474072309e-05,
+      "loss": 0.0454,
+      "step": 1640
+    },
+    {
+      "grad_norm": 0.22489239275455475,
+      "learning_rate": 8.008047945352193e-05,
+      "loss": 0.0425,
+      "step": 1650
+    },
+    {
+      "grad_norm": 0.27072638273239136,
+      "learning_rate": 7.981566834870225e-05,
+      "loss": 0.0383,
+      "step": 1660
+    },
+    {
+      "grad_norm": 0.24278584122657776,
+      "learning_rate": 7.954955300995961e-05,
+      "loss": 0.0421,
+      "step": 1670
+    },
+    {
+      "grad_norm": 0.3021833002567291,
+      "learning_rate": 7.928214507804104e-05,
+      "loss": 0.0429,
+      "step": 1680
+    },
+    {
+      "grad_norm": 0.38383153080940247,
+      "learning_rate": 7.901345625023576e-05,
+      "loss": 0.0436,
+      "step": 1690
+    },
+    {
+      "grad_norm": 0.510800838470459,
+      "learning_rate": 7.874349827986354e-05,
+      "loss": 0.0414,
+      "step": 1700
+    },
+    {
+      "grad_norm": 0.3297543525695801,
+      "learning_rate": 7.847228297576053e-05,
+      "loss": 0.0436,
+      "step": 1710
+    },
+    {
+      "grad_norm": 0.244502454996109,
+      "learning_rate": 7.819982220176276e-05,
+      "loss": 0.0453,
+      "step": 1720
+    },
+    {
+      "grad_norm": 0.4956775903701782,
+      "learning_rate": 7.792612787618714e-05,
+      "loss": 0.0466,
+      "step": 1730
+    },
+    {
+      "grad_norm": 0.25343015789985657,
+      "learning_rate": 7.765121197131009e-05,
+      "loss": 0.0424,
+      "step": 1740
+    },
+    {
+      "grad_norm": 0.26113173365592957,
+      "learning_rate": 7.737508651284391e-05,
+      "loss": 0.0416,
+      "step": 1750
+    },
+    {
+      "grad_norm": 0.2182377576828003,
+      "learning_rate": 7.709776357941069e-05,
+      "loss": 0.0441,
+      "step": 1760
+    },
+    {
+      "grad_norm": 0.3971315026283264,
+      "learning_rate": 7.681925530201392e-05,
+      "loss": 0.0438,
+      "step": 1770
+    },
+    {
+      "grad_norm": 0.28176161646842957,
+      "learning_rate": 7.65395738635079e-05,
+      "loss": 0.0408,
+      "step": 1780
+    },
+    {
+      "grad_norm": 0.34179630875587463,
+      "learning_rate": 7.62587314980648e-05,
+      "loss": 0.0479,
+      "step": 1790
+    },
+    {
+      "grad_norm": 0.32931986451148987,
+      "learning_rate": 7.597674049063947e-05,
+      "loss": 0.0429,
+      "step": 1800
+    },
+    {
+      "grad_norm": 0.39432066679000854,
+      "learning_rate": 7.569361317643211e-05,
+      "loss": 0.0431,
+      "step": 1810
+    },
+    {
+      "grad_norm": 0.3576849699020386,
+      "learning_rate": 7.540936194034865e-05,
+      "loss": 0.0384,
+      "step": 1820
+    },
+    {
+      "grad_norm": 0.23606491088867188,
+      "learning_rate": 7.512399921645901e-05,
+      "loss": 0.0357,
+      "step": 1830
+    },
+    {
+      "grad_norm": 0.2227422446012497,
+      "learning_rate": 7.483753748745317e-05,
+      "loss": 0.0367,
+      "step": 1840
+    },
+    {
+      "grad_norm": 0.2424551248550415,
+      "learning_rate": 7.454998928409516e-05,
+      "loss": 0.0385,
+      "step": 1850
+    },
+    {
+      "grad_norm": 0.33607348799705505,
+      "learning_rate": 7.426136718467493e-05,
+      "loss": 0.0412,
+      "step": 1860
+    },
+    {
+      "grad_norm": 0.2656790018081665,
+      "learning_rate": 7.397168381445812e-05,
+      "loss": 0.0402,
+      "step": 1870
+    },
+    {
+      "grad_norm": 0.27483707666397095,
+      "learning_rate": 7.368095184513377e-05,
+      "loss": 0.0332,
+      "step": 1880
+    },
+    {
+      "grad_norm": 0.32864516973495483,
+      "learning_rate": 7.338918399426005e-05,
+      "loss": 0.0352,
+      "step": 1890
+    },
+    {
+      "grad_norm": 0.2956995964050293,
+      "learning_rate": 7.309639302470801e-05,
+      "loss": 0.0464,
+      "step": 1900
+    },
+    {
+      "grad_norm": 0.3073318302631378,
+      "learning_rate": 7.280259174410312e-05,
+      "loss": 0.0378,
+      "step": 1910
+    },
+    {
+      "grad_norm": 0.3206455409526825,
+      "learning_rate": 7.250779300426517e-05,
+      "loss": 0.0428,
+      "step": 1920
+    },
+    {
+      "grad_norm": 0.40320953726768494,
+      "learning_rate": 7.22120097006461e-05,
+      "loss": 0.0434,
+      "step": 1930
+    },
+    {
+      "grad_norm": 0.2006569355726242,
+      "learning_rate": 7.191525477176577e-05,
+      "loss": 0.0367,
+      "step": 1940
+    },
+    {
+      "grad_norm": 0.26202642917633057,
+      "learning_rate": 7.161754119864616e-05,
+      "loss": 0.0387,
+      "step": 1950
+    },
+    {
+      "grad_norm": 0.36526966094970703,
+      "learning_rate": 7.131888200424339e-05,
+      "loss": 0.0344,
+      "step": 1960
+    },
+    {
+      "grad_norm": 0.3327760696411133,
+      "learning_rate": 7.101929025287816e-05,
+      "loss": 0.0401,
+      "step": 1970
+    },
+    {
+      "grad_norm": 0.36563345789909363,
+      "learning_rate": 7.071877904966423e-05,
+      "loss": 0.0474,
+      "step": 1980
+    },
+    {
+      "grad_norm": 0.27546799182891846,
+      "learning_rate": 7.04173615399351e-05,
+      "loss": 0.0376,
+      "step": 1990
+    },
+    {
+      "grad_norm": 0.3782268464565277,
+      "learning_rate": 7.011505090866913e-05,
+      "loss": 0.0324,
+      "step": 2000
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 5000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 2000,
+  "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": 32,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-2000/training_args.bin

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

checkpoint-2000/wandb_config.json

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

checkpoint-4000/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-4000/embodiment_id.json

@@ -0,0 +1,11 @@
+{
+  "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
+}

checkpoint-4000/experiment_cfg/conf.yaml

@@ -0,0 +1,210 @@
+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:
+    brightness: 0.3
+    contrast: 0.4
+    saturation: 0.5
+    hue: 0.08
+  use_albumentations_transforms: true
+  extra_augmentation_config: null
+  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:
+    - ./demo_data/SOarmtest/viis0002/pick_lego
+    embodiment_tag: new_embodiment
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    new_embodiment:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front
+        - wrist
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - single_arm
+        - gripper
+        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:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - 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:
+        - 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: 100000
+  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: torchcodec
+training:
+  output_dir: ./demo_data/SOarmtest/saved_checkpoints
+  experiment_name: null
+  max_steps: 5000
+  global_batch_size: 32
+  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: 2000
+  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: 4
+  remove_unused_columns: false
+  use_wandb: true
+  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: 5000
+save_steps: 2000

checkpoint-4000/experiment_cfg/config.yaml

@@ -0,0 +1,244 @@
+!!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:
+    - ./demo_data/SOarmtest/viis0002/pick_lego
+    dataset_type: physical_embodiment
+    embodiment_tag: new_embodiment
+    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:
+    new_embodiment:
+      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: null
+          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:
+        - single_arm
+        - 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:
+        - single_arm
+        - gripper
+        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
+        - wrist
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 100000
+  override_pretraining_statistics: false
+  random_chop: 0.0
+  seed: 42
+  shard_size: 1024
+  shuffle: true
+  subsample_ratio: 1.0
+  video_backend: torchcodec
+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:
+    brightness: 0.3
+    contrast: 0.4
+    hue: 0.08
+    saturation: 0.5
+  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
+  extra_augmentation_config: 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: 4
+  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: 32
+  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: 5000
+  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: ./demo_data/SOarmtest/saved_checkpoints
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_steps: 2000
+  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: true
+  wandb_project: finetune-gr00t-n1d6
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

checkpoint-4000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,824 @@
+{
+  "new_embodiment": {
+    "state": {
+      "single_arm": {
+        "min": [
+          -6.8131866455078125,
+          -102.19779968261719,
+          -40.43955993652344,
+          39.9560432434082,
+          -66.32967376708984
+        ],
+        "max": [
+          34.593406677246094,
+          42.769229888916016,
+          96.26373291015625,
+          98.06593322753906,
+          96.74725341796875
+        ],
+        "mean": [
+          10.6347642478831,
+          -40.60343800114444,
+          40.109735021899446,
+          66.37540221377535,
+          1.723290164849921
+        ],
+        "std": [
+          13.663472797431913,
+          53.514238207466356,
+          47.089223703864675,
+          15.49883594624875,
+          29.442594278901634
+        ],
+        "q01": [
+          -2.8227532818225196,
+          -102.04170769592945,
+          -18.832328638595566,
+          47.303414007272124,
+          -23.538191580154248
+        ],
+        "q99": [
+          31.17592093949654,
+          27.78409835219628,
+          96.09969459877085,
+          91.52694173019303,
+          34.804566667323705
+        ]
+      },
+      "gripper": {
+        "min": [
+          1.1597938537597656
+        ],
+        "max": [
+          30.476804733276367
+        ],
+        "mean": [
+          4.725540477545816
+        ],
+        "std": [
+          5.719213168100463
+        ],
+        "q01": [
+          1.219114189586798
+        ],
+        "q99": [
+          18.975632466717574
+        ]
+      }
+    },
+    "action": {
+      "single_arm": {
+        "min": [
+          -7.340659141540527,
+          -103.73626708984375,
+          -43.736263275146484,
+          38.63736343383789,
+          -67.47252655029297
+        ],
+        "max": [
+          35.20878982543945,
+          42.5494499206543,
+          96.74725341796875,
+          100.61538696289062,
+          97.27472686767578
+        ],
+        "mean": [
+          10.654706232867026,
+          -41.43508459659817,
+          38.95115310430438,
+          66.17248404139285,
+          1.6835930980913996
+        ],
+        "std": [
+          13.792301015873905,
+          53.27088262273843,
+          48.06070434685447,
+          15.735264444521983,
+          29.50880036285619
+        ],
+        "q01": [
+          -3.007477344415986,
+          -102.62527508767714,
+          -21.606171980367048,
+          46.25193771491276,
+          -24.028302588862974
+        ],
+        "q99": [
+          31.658033587037096,
+          27.49569149494191,
+          96.40327760792744,
+          92.26730265898277,
+          35.153253393452815
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.0
+        ],
+        "max": [
+          30.952381134033203
+        ],
+        "mean": [
+          3.5887609717193083
+        ],
+        "std": [
+          5.951741507358969
+        ],
+        "q01": [
+          0.7785521832201806
+        ],
+        "q99": [
+          19.41596701897272
+        ]
+      }
+    },
+    "relative_action": {
+      "single_arm": {
+        "min": [
+          [
+            -6.153846740722656,
+            -12.615386962890625,
+            -13.406593322753906,
+            -9.406593322753906,
+            -13.450549125671387
+          ],
+          [
+            -7.824176788330078,
+            -15.692306518554688,
+            -16.39560317993164,
+            -11.86812973022461,
+            -13.450549125671387
+          ],
+          [
+            -9.582418441772461,
+            -18.681320190429688,
+            -19.296703338623047,
+            -14.241756439208984,
+            -15.472527503967285
+          ],
+          [
+            -11.252748489379883,
+            -21.670331954956055,
+            -22.19780158996582,
+            -16.615379333496094,
+            -16.96703338623047
+          ],
+          [
+            -12.923077583312988,
+            -24.65934181213379,
+            -24.923078536987305,
+            -18.901100158691406,
+            -18.461538314819336
+          ],
+          [
+            -14.593405723571777,
+            -27.648353576660156,
+            -27.736263275146484,
+            -21.186813354492188,
+            -19.516483306884766
+          ],
+          [
+            -16.087913513183594,
+            -30.637365341186523,
+            -30.46154022216797,
+            -23.384613037109375,
+            -21.010990142822266
+          ],
+          [
+            -17.58241844177246,
+            -33.45055389404297,
+            -33.09890365600586,
+            -25.406593322753906,
+            -23.032970428466797
+          ],
+          [
+            -18.813186645507812,
+            -36.351646423339844,
+            -35.648353576660156,
+            -27.428573608398438,
+            -24.967037200927734
+          ],
+          [
+            -20.043954849243164,
+            -39.164833068847656,
+            -38.109893798828125,
+            -29.09890365600586,
+            -27.076923370361328
+          ],
+          [
+            -21.098899841308594,
+            -42.06593322753906,
+            -40.57142639160156,
+            -30.769229888916016,
+            -29.274723052978516
+          ],
+          [
+            -22.153844833374023,
+            -44.879119873046875,
+            -42.857139587402344,
+            -32.351646423339844,
+            -31.38461685180664
+          ],
+          [
+            -23.208789825439453,
+            -47.69231033325195,
+            -45.05494689941406,
+            -33.75823974609375,
+            -33.494503021240234
+          ],
+          [
+            -24.08791160583496,
+            -50.68132019042969,
+            -47.164833068847656,
+            -34.989013671875,
+            -35.51648712158203
+          ],
+          [
+            -24.967031478881836,
+            -53.67033004760742,
+            -49.098899841308594,
+            -36.04396057128906,
+            -37.626373291015625
+          ],
+          [
+            -25.758241653442383,
+            -56.747249603271484,
+            -50.94505310058594,
+            -37.01099395751953,
+            -39.648353576660156
+          ]
+        ],
+        "max": [
+          [
+            5.8021979331970215,
+            14.021980285644531,
+            12.527473449707031,
+            8.703296661376953,
+            9.670330047607422
+          ],
+          [
+            6.505494117736816,
+            17.626373291015625,
+            16.131866455078125,
+            11.164833068847656,
+            12.219779968261719
+          ],
+          [
+            7.82417631149292,
+            21.142852783203125,
+            19.824176788330078,
+            13.626373291015625,
+            14.681318283081055
+          ],
+          [
+            9.05494499206543,
+            24.395599365234375,
+            23.428571701049805,
+            15.824176788330078,
+            17.142858505249023
+          ],
+          [
+            10.373626708984375,
+            27.384613037109375,
+            26.94505500793457,
+            18.021976470947266,
+            19.69230842590332
+          ],
+          [
+            11.69230842590332,
+            30.197803497314453,
+            30.461538314819336,
+            20.043956756591797,
+            22.153846740722656
+          ],
+          [
+            13.01098918914795,
+            32.83516311645508,
+            33.97801971435547,
+            22.241756439208984,
+            24.703298568725586
+          ],
+          [
+            14.153846740722656,
+            35.384613037109375,
+            37.406593322753906,
+            24.35165023803711,
+            27.07692527770996
+          ],
+          [
+            15.296703338623047,
+            37.75823974609375,
+            40.74725341796875,
+            26.197803497314453,
+            29.362638473510742
+          ],
+          [
+            16.351648330688477,
+            40.04395294189453,
+            44.08790969848633,
+            27.956043243408203,
+            31.560440063476562
+          ],
+          [
+            17.318681716918945,
+            42.241756439208984,
+            47.340660095214844,
+            29.71428680419922,
+            33.67033004760742
+          ],
+          [
+            18.197803497314453,
+            44.527469635009766,
+            50.68131637573242,
+            31.384613037109375,
+            35.95604705810547
+          ],
+          [
+            18.901100158691406,
+            46.4615364074707,
+            53.93406295776367,
+            32.96703338623047,
+            38.5054931640625
+          ],
+          [
+            19.516483306884766,
+            48.48351287841797,
+            57.010986328125,
+            34.28571319580078,
+            41.14285659790039
+          ],
+          [
+            20.043956756591797,
+            50.32966613769531,
+            60.0,
+            35.604393005371094,
+            43.516483306884766
+          ],
+          [
+            20.395605087280273,
+            52.26373291015625,
+            62.989013671875,
+            36.659339904785156,
+            45.80220031738281
+          ]
+        ],
+        "mean": [
+          [
+            0.024023909121751785,
+            -0.8380120396614075,
+            -1.1888031959533691,
+            -0.20351642370224,
+            -0.03874167427420616
+          ],
+          [
+            0.02381421998143196,
+            -0.8381950259208679,
+            -1.1877086162567139,
+            -0.20610879361629486,
+            -0.038331978023052216
+          ],
+          [
+            0.02358066476881504,
+            -0.8384636640548706,
+            -1.186354160308838,
+            -0.20878157019615173,
+            -0.03791269287467003
+          ],
+          [
+            0.02331378683447838,
+            -0.8388302326202393,
+            -1.1847561597824097,
+            -0.21150696277618408,
+            -0.03749352693557739
+          ],
+          [
+            0.022984733805060387,
+            -0.8392952680587769,
+            -1.182924509048462,
+            -0.21426069736480713,
+            -0.037074267864227295
+          ],
+          [
+            0.02258922904729843,
+            -0.8398675322532654,
+            -1.1808984279632568,
+            -0.21703894436359406,
+            -0.03665487840771675
+          ],
+          [
+            0.02210799790918827,
+            -0.840530514717102,
+            -1.1787015199661255,
+            -0.21982122957706451,
+            -0.03624051436781883
+          ],
+          [
+            0.021555287763476372,
+            -0.8412824869155884,
+            -1.176315188407898,
+            -0.2226090133190155,
+            -0.03585929051041603
+          ],
+          [
+            0.020935896784067154,
+            -0.8421463966369629,
+            -1.1736962795257568,
+            -0.22543004155158997,
+            -0.035549700260162354
+          ],
+          [
+            0.020240070298314095,
+            -0.8431134819984436,
+            -1.1708482503890991,
+            -0.2282552719116211,
+            -0.03533996269106865
+          ],
+          [
+            0.01948723942041397,
+            -0.8441924452781677,
+            -1.1677296161651611,
+            -0.23114195466041565,
+            -0.03522548824548721
+          ],
+          [
+            0.018677080050110817,
+            -0.8453867435455322,
+            -1.164296269416809,
+            -0.23409080505371094,
+            -0.03516348451375961
+          ],
+          [
+            0.017790833488106728,
+            -0.846683919429779,
+            -1.1604851484298706,
+            -0.23711234331130981,
+            -0.035125281661748886
+          ],
+          [
+            0.016804957762360573,
+            -0.8481130599975586,
+            -1.156294584274292,
+            -0.24017632007598877,
+            -0.03509661555290222
+          ],
+          [
+            0.015680411830544472,
+            -0.8497316241264343,
+            -1.151717185974121,
+            -0.24328306317329407,
+            -0.035068120807409286
+          ],
+          [
+            0.01436557900160551,
+            -0.8516368269920349,
+            -1.1466983556747437,
+            -0.24642817676067352,
+            -0.035063356161117554
+          ]
+        ],
+        "std": [
+          [
+            1.0472843647003174,
+            3.3955397605895996,
+            3.209542989730835,
+            1.5478397607803345,
+            1.5488325357437134
+          ],
+          [
+            1.3025773763656616,
+            4.228511333465576,
+            3.9392380714416504,
+            1.9040354490280151,
+            1.9508968591690063
+          ],
+          [
+            1.5567840337753296,
+            5.061123371124268,
+            4.678185939788818,
+            2.259927749633789,
+            2.3464455604553223
+          ],
+          [
+            1.8073352575302124,
+            5.890169143676758,
+            5.418741226196289,
+            2.6112139225006104,
+            2.734940528869629
+          ],
+          [
+            2.053133726119995,
+            6.713922023773193,
+            6.156966209411621,
+            2.9563238620758057,
+            3.1171467304229736
+          ],
+          [
+            2.2936792373657227,
+            7.531452178955078,
+            6.89060115814209,
+            3.2943806648254395,
+            3.4939496517181396
+          ],
+          [
+            2.528961658477783,
+            8.342284202575684,
+            7.61845064163208,
+            3.6247124671936035,
+            3.8658618927001953
+          ],
+          [
+            2.7590219974517822,
+            9.146015167236328,
+            8.339773178100586,
+            3.9470086097717285,
+            4.233205795288086
+          ],
+          [
+            2.9840755462646484,
+            9.942627906799316,
+            9.053919792175293,
+            4.261190414428711,
+            4.596113681793213
+          ],
+          [
+            3.2043535709381104,
+            10.731932640075684,
+            9.760570526123047,
+            4.567060947418213,
+            4.954709053039551
+          ],
+          [
+            3.420076608657837,
+            11.51407527923584,
+            10.459877967834473,
+            4.864653587341309,
+            5.3089189529418945
+          ],
+          [
+            3.6315865516662598,
+            12.289117813110352,
+            11.151065826416016,
+            5.154054641723633,
+            5.65879487991333
+          ],
+          [
+            3.8390471935272217,
+            13.056876182556152,
+            11.834860801696777,
+            5.435141086578369,
+            6.004306316375732
+          ],
+          [
+            4.042593002319336,
+            13.817543029785156,
+            12.511225700378418,
+            5.708395004272461,
+            6.345395565032959
+          ],
+          [
+            4.24238395690918,
+            14.571209907531738,
+            13.179614067077637,
+            5.973555564880371,
+            6.682039737701416
+          ],
+          [
+            4.438711166381836,
+            15.317753791809082,
+            13.840078353881836,
+            6.230597496032715,
+            7.01429557800293
+          ]
+        ],
+        "q01": [
+          [
+            -3.472527027130127,
+            -9.978021621704102,
+            -9.978019714355469,
+            -5.098899841308594,
+            -5.36263370513916
+          ],
+          [
+            -4.483516693115234,
+            -12.175825119018555,
+            -12.263736724853516,
+            -6.417581558227539,
+            -6.857143402099609
+          ],
+          [
+            -5.450550079345703,
+            -14.373626708984375,
+            -14.373626708984375,
+            -7.736263275146484,
+            -8.263736724853516
+          ],
+          [
+            -6.4175825119018555,
+            -16.571428298950195,
+            -16.483516693115234,
+            -9.054946899414062,
+            -9.846153259277344
+          ],
+          [
+            -7.296703338623047,
+            -18.85714340209961,
+            -18.681316375732422,
+            -10.285713195800781,
+            -11.296703338623047
+          ],
+          [
+            -8.175825119018555,
+            -21.054946899414062,
+            -20.79120635986328,
+            -11.516483306884766,
+            -12.747251510620117
+          ],
+          [
+            -9.05494499206543,
+            -23.340660095214844,
+            -22.813186645507812,
+            -12.703296661376953,
+            -14.19780158996582
+          ],
+          [
+            -10.021978378295898,
+            -25.582416534423828,
+            -24.791210174560547,
+            -13.80219841003418,
+            -15.560441017150879
+          ],
+          [
+            -10.813186645507812,
+            -27.648353576660156,
+            -26.769229888916016,
+            -14.813186645507812,
+            -16.96703338623047
+          ],
+          [
+            -11.780220031738281,
+            -29.934066772460938,
+            -28.659339904785156,
+            -15.824176788330078,
+            -18.373626708984375
+          ],
+          [
+            -12.615385055541992,
+            -32.21977996826172,
+            -30.593406677246094,
+            -16.703292846679688,
+            -19.516483306884766
+          ],
+          [
+            -13.450550079345703,
+            -34.32966995239258,
+            -32.395599365234375,
+            -17.670326232910156,
+            -20.791208267211914
+          ],
+          [
+            -14.329669952392578,
+            -36.52747344970703,
+            -34.19780349731445,
+            -18.549448013305664,
+            -21.802196502685547
+          ],
+          [
+            -15.164834976196289,
+            -38.68132019042969,
+            -35.82417297363281,
+            -19.20879364013672,
+            -22.901100158691406
+          ],
+          [
+            -15.95604419708252,
+            -40.92308044433594,
+            -37.626373291015625,
+            -20.043954849243164,
+            -24.131868362426758
+          ],
+          [
+            -16.703296661376953,
+            -42.989009857177734,
+            -39.25274658203125,
+            -20.74725341796875,
+            -25.450550079345703
+          ]
+        ],
+        "q99": [
+          [
+            3.4285709857940674,
+            7.5604400634765625,
+            7.076923847198486,
+            4.131866455078125,
+            5.494504928588867
+          ],
+          [
+            4.307692527770996,
+            9.802200317382812,
+            9.538460731506348,
+            5.274726867675781,
+            7.032966613769531
+          ],
+          [
+            5.1868133544921875,
+            12.087913513183594,
+            11.868131637573242,
+            6.5054931640625,
+            8.527473449707031
+          ],
+          [
+            6.065934181213379,
+            14.197805404663086,
+            14.285714149475098,
+            7.648347854614258,
+            9.934066772460938
+          ],
+          [
+            6.857143402099609,
+            16.483516693115234,
+            16.571428298950195,
+            8.703292846679688,
+            11.340660095214844
+          ],
+          [
+            7.736264705657959,
+            18.593406677246094,
+            18.85714340209961,
+            9.846153259277344,
+            12.835165023803711
+          ],
+          [
+            8.615385055541992,
+            20.52747344970703,
+            21.14285659790039,
+            10.813188552856445,
+            14.241758346557617
+          ],
+          [
+            9.494504928588867,
+            22.725276947021484,
+            23.428571701049805,
+            11.824176788330078,
+            15.648351669311523
+          ],
+          [
+            10.285714149475098,
+            24.879121780395508,
+            25.714282989501953,
+            12.835166931152344,
+            17.05494499206543
+          ],
+          [
+            11.076923370361328,
+            26.85714340209961,
+            27.868133544921875,
+            13.758243560791016,
+            18.461538314819336
+          ],
+          [
+            11.7362642288208,
+            28.70329475402832,
+            30.109888076782227,
+            14.725275039672852,
+            19.868133544921875
+          ],
+          [
+            12.395605087280273,
+            30.813186645507812,
+            32.17582321166992,
+            15.560440063476562,
+            21.186813354492188
+          ],
+          [
+            13.098901748657227,
+            32.52747344970703,
+            34.19779968261719,
+            16.351646423339844,
+            22.549449920654297
+          ],
+          [
+            13.714285850524902,
+            34.54945373535156,
+            36.21977996826172,
+            17.142860412597656,
+            24.0
+          ],
+          [
+            14.329669952392578,
+            36.57143020629883,
+            38.153846740722656,
+            17.846155166625977,
+            25.230770111083984
+          ],
+          [
+            14.769231796264648,
+            38.28571319580078,
+            40.17582702636719,
+            18.637367248535156,
+            26.41758155822754
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-4000/experiment_cfg/final_model_config.json

@@ -0,0 +1,54 @@
+{
+  "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,
+  "extra_augmentation_config": null,
+  "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-4000/model-00001-of-00002.safetensors

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

checkpoint-4000/model-00002-of-00002.safetensors

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

checkpoint-4000/optimizer.pt

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

checkpoint-4000/processor_config.json

@@ -0,0 +1,455 @@
+{
+  "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
+        }
+      },
+      "new_embodiment": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front",
+            "wrist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "single_arm",
+            "gripper"
+          ],
+          "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": [
+            "single_arm",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "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": [
+            "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-4000/rng_state.pth

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

checkpoint-4000/scheduler.pt

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

checkpoint-4000/trainer_state.json

@@ -0,0 +1,2434 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.8,
+  "eval_steps": 500,
+  "global_step": 4000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "grad_norm": 2.497777223587036,
+      "learning_rate": 3.6e-06,
+      "loss": 1.1476,
+      "step": 10
+    },
+    {
+      "grad_norm": 0.5899417996406555,
+      "learning_rate": 7.6e-06,
+      "loss": 1.1021,
+      "step": 20
+    },
+    {
+      "grad_norm": 0.6255532503128052,
+      "learning_rate": 1.16e-05,
+      "loss": 1.0779,
+      "step": 30
+    },
+    {
+      "grad_norm": 0.5589399933815002,
+      "learning_rate": 1.56e-05,
+      "loss": 1.0599,
+      "step": 40
+    },
+    {
+      "grad_norm": 0.6063516736030579,
+      "learning_rate": 1.9600000000000002e-05,
+      "loss": 1.0856,
+      "step": 50
+    },
+    {
+      "grad_norm": 0.5875065326690674,
+      "learning_rate": 2.36e-05,
+      "loss": 1.0636,
+      "step": 60
+    },
+    {
+      "grad_norm": 0.6206650733947754,
+      "learning_rate": 2.7600000000000003e-05,
+      "loss": 1.0586,
+      "step": 70
+    },
+    {
+      "grad_norm": 0.6976884007453918,
+      "learning_rate": 3.16e-05,
+      "loss": 1.0581,
+      "step": 80
+    },
+    {
+      "grad_norm": 0.4975197911262512,
+      "learning_rate": 3.56e-05,
+      "loss": 1.0338,
+      "step": 90
+    },
+    {
+      "grad_norm": 0.5175010561943054,
+      "learning_rate": 3.960000000000001e-05,
+      "loss": 1.0363,
+      "step": 100
+    },
+    {
+      "grad_norm": 3.203188180923462,
+      "learning_rate": 4.36e-05,
+      "loss": 0.9943,
+      "step": 110
+    },
+    {
+      "grad_norm": 1.0046463012695312,
+      "learning_rate": 4.76e-05,
+      "loss": 0.9498,
+      "step": 120
+    },
+    {
+      "grad_norm": 0.824455976486206,
+      "learning_rate": 5.16e-05,
+      "loss": 0.9141,
+      "step": 130
+    },
+    {
+      "grad_norm": 0.933220386505127,
+      "learning_rate": 5.560000000000001e-05,
+      "loss": 0.8966,
+      "step": 140
+    },
+    {
+      "grad_norm": 1.6803524494171143,
+      "learning_rate": 5.96e-05,
+      "loss": 0.8456,
+      "step": 150
+    },
+    {
+      "grad_norm": 1.3335356712341309,
+      "learning_rate": 6.36e-05,
+      "loss": 0.7857,
+      "step": 160
+    },
+    {
+      "grad_norm": 1.7373239994049072,
+      "learning_rate": 6.76e-05,
+      "loss": 0.6714,
+      "step": 170
+    },
+    {
+      "grad_norm": 1.2455403804779053,
+      "learning_rate": 7.16e-05,
+      "loss": 0.5578,
+      "step": 180
+    },
+    {
+      "grad_norm": 2.673060178756714,
+      "learning_rate": 7.560000000000001e-05,
+      "loss": 0.4206,
+      "step": 190
+    },
+    {
+      "grad_norm": 1.7698984146118164,
+      "learning_rate": 7.960000000000001e-05,
+      "loss": 0.3168,
+      "step": 200
+    },
+    {
+      "grad_norm": 1.393641471862793,
+      "learning_rate": 8.36e-05,
+      "loss": 0.2192,
+      "step": 210
+    },
+    {
+      "grad_norm": 1.3259373903274536,
+      "learning_rate": 8.76e-05,
+      "loss": 0.1876,
+      "step": 220
+    },
+    {
+      "grad_norm": 0.9707556366920471,
+      "learning_rate": 9.16e-05,
+      "loss": 0.1548,
+      "step": 230
+    },
+    {
+      "grad_norm": 0.8586066961288452,
+      "learning_rate": 9.56e-05,
+      "loss": 0.1401,
+      "step": 240
+    },
+    {
+      "grad_norm": 0.8358578085899353,
+      "learning_rate": 9.960000000000001e-05,
+      "loss": 0.1216,
+      "step": 250
+    },
+    {
+      "grad_norm": 0.6485460996627808,
+      "learning_rate": 9.999911419878559e-05,
+      "loss": 0.1228,
+      "step": 260
+    },
+    {
+      "grad_norm": 1.0369858741760254,
+      "learning_rate": 9.999605221019081e-05,
+      "loss": 0.1396,
+      "step": 270
+    },
+    {
+      "grad_norm": 0.8534221053123474,
+      "learning_rate": 9.999080323230761e-05,
+      "loss": 0.1208,
+      "step": 280
+    },
+    {
+      "grad_norm": 0.8548014760017395,
+      "learning_rate": 9.998336749474329e-05,
+      "loss": 0.1341,
+      "step": 290
+    },
+    {
+      "grad_norm": 0.9402936100959778,
+      "learning_rate": 9.997374532276107e-05,
+      "loss": 0.102,
+      "step": 300
+    },
+    {
+      "grad_norm": 0.7149423956871033,
+      "learning_rate": 9.996193713726596e-05,
+      "loss": 0.1092,
+      "step": 310
+    },
+    {
+      "grad_norm": 0.7244396805763245,
+      "learning_rate": 9.994794345478624e-05,
+      "loss": 0.0924,
+      "step": 320
+    },
+    {
+      "grad_norm": 0.5502927899360657,
+      "learning_rate": 9.99317648874509e-05,
+      "loss": 0.0838,
+      "step": 330
+    },
+    {
+      "grad_norm": 0.6544800400733948,
+      "learning_rate": 9.991340214296292e-05,
+      "loss": 0.0873,
+      "step": 340
+    },
+    {
+      "grad_norm": 0.7180399894714355,
+      "learning_rate": 9.989285602456819e-05,
+      "loss": 0.0734,
+      "step": 350
+    },
+    {
+      "grad_norm": 0.6445586681365967,
+      "learning_rate": 9.98701274310205e-05,
+      "loss": 0.0794,
+      "step": 360
+    },
+    {
+      "grad_norm": 0.5352494716644287,
+      "learning_rate": 9.984521735654218e-05,
+      "loss": 0.0907,
+      "step": 370
+    },
+    {
+      "grad_norm": 0.6672635674476624,
+      "learning_rate": 9.981812689078057e-05,
+      "loss": 0.0848,
+      "step": 380
+    },
+    {
+      "grad_norm": 0.5291720628738403,
+      "learning_rate": 9.978885721876041e-05,
+      "loss": 0.086,
+      "step": 390
+    },
+    {
+      "grad_norm": 0.39044567942619324,
+      "learning_rate": 9.975740962083198e-05,
+      "loss": 0.0787,
+      "step": 400
+    },
+    {
+      "grad_norm": 0.7575682997703552,
+      "learning_rate": 9.972378547261504e-05,
+      "loss": 0.0836,
+      "step": 410
+    },
+    {
+      "grad_norm": 0.6050339341163635,
+      "learning_rate": 9.968798624493885e-05,
+      "loss": 0.0738,
+      "step": 420
+    },
+    {
+      "grad_norm": 0.594001054763794,
+      "learning_rate": 9.965001350377753e-05,
+      "loss": 0.0817,
+      "step": 430
+    },
+    {
+      "grad_norm": 0.5727436542510986,
+      "learning_rate": 9.960986891018183e-05,
+      "loss": 0.0728,
+      "step": 440
+    },
+    {
+      "grad_norm": 0.4483703076839447,
+      "learning_rate": 9.95675542202063e-05,
+      "loss": 0.0747,
+      "step": 450
+    },
+    {
+      "grad_norm": 0.3992336094379425,
+      "learning_rate": 9.952307128483256e-05,
+      "loss": 0.0796,
+      "step": 460
+    },
+    {
+      "grad_norm": 0.5178577303886414,
+      "learning_rate": 9.947642204988835e-05,
+      "loss": 0.0833,
+      "step": 470
+    },
+    {
+      "grad_norm": 0.4497125446796417,
+      "learning_rate": 9.942760855596226e-05,
+      "loss": 0.0787,
+      "step": 480
+    },
+    {
+      "grad_norm": 0.5343725085258484,
+      "learning_rate": 9.937663293831471e-05,
+      "loss": 0.0808,
+      "step": 490
+    },
+    {
+      "grad_norm": 0.47977861762046814,
+      "learning_rate": 9.932349742678433e-05,
+      "loss": 0.0701,
+      "step": 500
+    },
+    {
+      "grad_norm": 0.5831079483032227,
+      "learning_rate": 9.926820434569051e-05,
+      "loss": 0.0742,
+      "step": 510
+    },
+    {
+      "grad_norm": 0.47160646319389343,
+      "learning_rate": 9.921075611373179e-05,
+      "loss": 0.067,
+      "step": 520
+    },
+    {
+      "grad_norm": 0.35774683952331543,
+      "learning_rate": 9.915115524387988e-05,
+      "loss": 0.0709,
+      "step": 530
+    },
+    {
+      "grad_norm": 0.5747243762016296,
+      "learning_rate": 9.908940434326997e-05,
+      "loss": 0.0668,
+      "step": 540
+    },
+    {
+      "grad_norm": 0.42538684606552124,
+      "learning_rate": 9.902550611308645e-05,
+      "loss": 0.0676,
+      "step": 550
+    },
+    {
+      "grad_norm": 0.5689398646354675,
+      "learning_rate": 9.895946334844494e-05,
+      "loss": 0.0631,
+      "step": 560
+    },
+    {
+      "grad_norm": 0.648316502571106,
+      "learning_rate": 9.889127893826989e-05,
+      "loss": 0.0643,
+      "step": 570
+    },
+    {
+      "grad_norm": 0.5649459362030029,
+      "learning_rate": 9.882095586516831e-05,
+      "loss": 0.0679,
+      "step": 580
+    },
+    {
+      "grad_norm": 0.42304474115371704,
+      "learning_rate": 9.874849720529921e-05,
+      "loss": 0.0744,
+      "step": 590
+    },
+    {
+      "grad_norm": 0.5694190263748169,
+      "learning_rate": 9.867390612823914e-05,
+      "loss": 0.0683,
+      "step": 600
+    },
+    {
+      "grad_norm": 0.5124210715293884,
+      "learning_rate": 9.859718589684344e-05,
+      "loss": 0.0797,
+      "step": 610
+    },
+    {
+      "grad_norm": 0.6514238119125366,
+      "learning_rate": 9.851833986710353e-05,
+      "loss": 0.0644,
+      "step": 620
+    },
+    {
+      "grad_norm": 0.49233269691467285,
+      "learning_rate": 9.843737148800023e-05,
+      "loss": 0.0568,
+      "step": 630
+    },
+    {
+      "grad_norm": 0.40201443433761597,
+      "learning_rate": 9.835428430135271e-05,
+      "loss": 0.063,
+      "step": 640
+    },
+    {
+      "grad_norm": 0.4645858407020569,
+      "learning_rate": 9.82690819416637e-05,
+      "loss": 0.0641,
+      "step": 650
+    },
+    {
+      "grad_norm": 0.5921373963356018,
+      "learning_rate": 9.818176813596041e-05,
+      "loss": 0.0622,
+      "step": 660
+    },
+    {
+      "grad_norm": 0.5843824148178101,
+      "learning_rate": 9.809234670363159e-05,
+      "loss": 0.0724,
+      "step": 670
+    },
+    {
+      "grad_norm": 0.6411054730415344,
+      "learning_rate": 9.800082155626034e-05,
+      "loss": 0.0697,
+      "step": 680
+    },
+    {
+      "grad_norm": 0.4261399507522583,
+      "learning_rate": 9.790719669745312e-05,
+      "loss": 0.0641,
+      "step": 690
+    },
+    {
+      "grad_norm": 0.3344235122203827,
+      "learning_rate": 9.781147622266455e-05,
+      "loss": 0.0642,
+      "step": 700
+    },
+    {
+      "grad_norm": 0.4839865565299988,
+      "learning_rate": 9.771366431901831e-05,
+      "loss": 0.0734,
+      "step": 710
+    },
+    {
+      "grad_norm": 0.3950216472148895,
+      "learning_rate": 9.761376526512394e-05,
+      "loss": 0.063,
+      "step": 720
+    },
+    {
+      "grad_norm": 0.45863887667655945,
+      "learning_rate": 9.751178343088963e-05,
+      "loss": 0.0598,
+      "step": 730
+    },
+    {
+      "grad_norm": 0.414181113243103,
+      "learning_rate": 9.740772327733123e-05,
+      "loss": 0.0592,
+      "step": 740
+    },
+    {
+      "grad_norm": 0.48416581749916077,
+      "learning_rate": 9.730158935637697e-05,
+      "loss": 0.0576,
+      "step": 750
+    },
+    {
+      "grad_norm": 0.40759462118148804,
+      "learning_rate": 9.719338631066834e-05,
+      "loss": 0.0522,
+      "step": 760
+    },
+    {
+      "grad_norm": 0.4843514859676361,
+      "learning_rate": 9.708311887335713e-05,
+      "loss": 0.0551,
+      "step": 770
+    },
+    {
+      "grad_norm": 0.37705329060554504,
+      "learning_rate": 9.697079186789823e-05,
+      "loss": 0.0601,
+      "step": 780
+    },
+    {
+      "grad_norm": 0.41809579730033875,
+      "learning_rate": 9.685641020783876e-05,
+      "loss": 0.0567,
+      "step": 790
+    },
+    {
+      "grad_norm": 0.4055332839488983,
+      "learning_rate": 9.67399788966031e-05,
+      "loss": 0.0569,
+      "step": 800
+    },
+    {
+      "grad_norm": 0.384504109621048,
+      "learning_rate": 9.662150302727395e-05,
+      "loss": 0.0586,
+      "step": 810
+    },
+    {
+      "grad_norm": 0.38468998670578003,
+      "learning_rate": 9.650098778236968e-05,
+      "loss": 0.0531,
+      "step": 820
+    },
+    {
+      "grad_norm": 0.5237128734588623,
+      "learning_rate": 9.637843843361749e-05,
+      "loss": 0.0552,
+      "step": 830
+    },
+    {
+      "grad_norm": 0.4317338466644287,
+      "learning_rate": 9.62538603417229e-05,
+      "loss": 0.0509,
+      "step": 840
+    },
+    {
+      "grad_norm": 0.3871803879737854,
+      "learning_rate": 9.612725895613526e-05,
+      "loss": 0.0625,
+      "step": 850
+    },
+    {
+      "grad_norm": 0.4323013126850128,
+      "learning_rate": 9.599863981480926e-05,
+      "loss": 0.0604,
+      "step": 860
+    },
+    {
+      "grad_norm": 0.36297470331192017,
+      "learning_rate": 9.586800854396283e-05,
+      "loss": 0.0545,
+      "step": 870
+    },
+    {
+      "grad_norm": 0.37466663122177124,
+      "learning_rate": 9.573537085783095e-05,
+      "loss": 0.0529,
+      "step": 880
+    },
+    {
+      "grad_norm": 0.3929672837257385,
+      "learning_rate": 9.560073255841571e-05,
+      "loss": 0.05,
+      "step": 890
+    },
+    {
+      "grad_norm": 0.4040852189064026,
+      "learning_rate": 9.546409953523247e-05,
+      "loss": 0.0603,
+      "step": 900
+    },
+    {
+      "grad_norm": 0.6248286366462708,
+      "learning_rate": 9.532547776505229e-05,
+      "loss": 0.0492,
+      "step": 910
+    },
+    {
+      "grad_norm": 0.5975190997123718,
+      "learning_rate": 9.518487331164048e-05,
+      "loss": 0.0538,
+      "step": 920
+    },
+    {
+      "grad_norm": 0.42542824149131775,
+      "learning_rate": 9.504229232549134e-05,
+      "loss": 0.0523,
+      "step": 930
+    },
+    {
+      "grad_norm": 0.328101247549057,
+      "learning_rate": 9.489774104355909e-05,
+      "loss": 0.0519,
+      "step": 940
+    },
+    {
+      "grad_norm": 0.5055563449859619,
+      "learning_rate": 9.475122578898507e-05,
+      "loss": 0.0603,
+      "step": 950
+    },
+    {
+      "grad_norm": 0.5114820003509521,
+      "learning_rate": 9.460275297082119e-05,
+      "loss": 0.0486,
+      "step": 960
+    },
+    {
+      "grad_norm": 0.30745837092399597,
+      "learning_rate": 9.445232908374948e-05,
+      "loss": 0.0515,
+      "step": 970
+    },
+    {
+      "grad_norm": 0.5132862329483032,
+      "learning_rate": 9.429996070779808e-05,
+      "loss": 0.0534,
+      "step": 980
+    },
+    {
+      "grad_norm": 0.3776204586029053,
+      "learning_rate": 9.414565450805333e-05,
+      "loss": 0.0543,
+      "step": 990
+    },
+    {
+      "grad_norm": 0.39591264724731445,
+      "learning_rate": 9.398941723436831e-05,
+      "loss": 0.051,
+      "step": 1000
+    },
+    {
+      "grad_norm": 0.3092094957828522,
+      "learning_rate": 9.383125572106752e-05,
+      "loss": 0.0485,
+      "step": 1010
+    },
+    {
+      "grad_norm": 0.4166562855243683,
+      "learning_rate": 9.367117688664791e-05,
+      "loss": 0.054,
+      "step": 1020
+    },
+    {
+      "grad_norm": 0.38996878266334534,
+      "learning_rate": 9.35091877334763e-05,
+      "loss": 0.0506,
+      "step": 1030
+    },
+    {
+      "grad_norm": 0.4977191686630249,
+      "learning_rate": 9.334529534748297e-05,
+      "loss": 0.0543,
+      "step": 1040
+    },
+    {
+      "grad_norm": 0.47602540254592896,
+      "learning_rate": 9.317950689785188e-05,
+      "loss": 0.048,
+      "step": 1050
+    },
+    {
+      "grad_norm": 0.3852064907550812,
+      "learning_rate": 9.301182963670688e-05,
+      "loss": 0.0563,
+      "step": 1060
+    },
+    {
+      "grad_norm": 0.2994341254234314,
+      "learning_rate": 9.284227089879456e-05,
+      "loss": 0.0528,
+      "step": 1070
+    },
+    {
+      "grad_norm": 0.4053724706172943,
+      "learning_rate": 9.26708381011634e-05,
+      "loss": 0.0546,
+      "step": 1080
+    },
+    {
+      "grad_norm": 0.28891879320144653,
+      "learning_rate": 9.249753874283937e-05,
+      "loss": 0.0435,
+      "step": 1090
+    },
+    {
+      "grad_norm": 0.2600522041320801,
+      "learning_rate": 9.232238040449779e-05,
+      "loss": 0.0459,
+      "step": 1100
+    },
+    {
+      "grad_norm": 0.5400200486183167,
+      "learning_rate": 9.214537074813181e-05,
+      "loss": 0.0553,
+      "step": 1110
+    },
+    {
+      "grad_norm": 0.4278520941734314,
+      "learning_rate": 9.196651751671724e-05,
+      "loss": 0.0543,
+      "step": 1120
+    },
+    {
+      "grad_norm": 0.43360674381256104,
+      "learning_rate": 9.178582853387384e-05,
+      "loss": 0.058,
+      "step": 1130
+    },
+    {
+      "grad_norm": 0.5356864929199219,
+      "learning_rate": 9.160331170352304e-05,
+      "loss": 0.0607,
+      "step": 1140
+    },
+    {
+      "grad_norm": 0.37267163395881653,
+      "learning_rate": 9.141897500954229e-05,
+      "loss": 0.0527,
+      "step": 1150
+    },
+    {
+      "grad_norm": 0.4231492877006531,
+      "learning_rate": 9.123282651541576e-05,
+      "loss": 0.0515,
+      "step": 1160
+    },
+    {
+      "grad_norm": 0.41324272751808167,
+      "learning_rate": 9.104487436388161e-05,
+      "loss": 0.0477,
+      "step": 1170
+    },
+    {
+      "grad_norm": 0.2980903089046478,
+      "learning_rate": 9.085512677657582e-05,
+      "loss": 0.0545,
+      "step": 1180
+    },
+    {
+      "grad_norm": 0.3594469428062439,
+      "learning_rate": 9.066359205367258e-05,
+      "loss": 0.0519,
+      "step": 1190
+    },
+    {
+      "grad_norm": 0.3295166492462158,
+      "learning_rate": 9.047027857352112e-05,
+      "loss": 0.0489,
+      "step": 1200
+    },
+    {
+      "grad_norm": 0.35804498195648193,
+      "learning_rate": 9.027519479227935e-05,
+      "loss": 0.0542,
+      "step": 1210
+    },
+    {
+      "grad_norm": 0.5596129298210144,
+      "learning_rate": 9.007834924354383e-05,
+      "loss": 0.0492,
+      "step": 1220
+    },
+    {
+      "grad_norm": 0.32689106464385986,
+      "learning_rate": 8.987975053797655e-05,
+      "loss": 0.0478,
+      "step": 1230
+    },
+    {
+      "grad_norm": 0.3586772680282593,
+      "learning_rate": 8.967940736292825e-05,
+      "loss": 0.0518,
+      "step": 1240
+    },
+    {
+      "grad_norm": 0.4046182632446289,
+      "learning_rate": 8.947732848205846e-05,
+      "loss": 0.0561,
+      "step": 1250
+    },
+    {
+      "grad_norm": 0.2832528352737427,
+      "learning_rate": 8.927352273495204e-05,
+      "loss": 0.0481,
+      "step": 1260
+    },
+    {
+      "grad_norm": 0.38299012184143066,
+      "learning_rate": 8.906799903673265e-05,
+      "loss": 0.0543,
+      "step": 1270
+    },
+    {
+      "grad_norm": 0.39338234066963196,
+      "learning_rate": 8.88607663776726e-05,
+      "loss": 0.053,
+      "step": 1280
+    },
+    {
+      "grad_norm": 0.326089471578598,
+      "learning_rate": 8.865183382279978e-05,
+      "loss": 0.0584,
+      "step": 1290
+    },
+    {
+      "grad_norm": 0.5139309167861938,
+      "learning_rate": 8.844121051150096e-05,
+      "loss": 0.0508,
+      "step": 1300
+    },
+    {
+      "grad_norm": 0.2685113251209259,
+      "learning_rate": 8.822890565712211e-05,
+      "loss": 0.0473,
+      "step": 1310
+    },
+    {
+      "grad_norm": 0.2396756261587143,
+      "learning_rate": 8.801492854656536e-05,
+      "loss": 0.047,
+      "step": 1320
+    },
+    {
+      "grad_norm": 0.2990108132362366,
+      "learning_rate": 8.779928853988268e-05,
+      "loss": 0.0422,
+      "step": 1330
+    },
+    {
+      "grad_norm": 0.25585711002349854,
+      "learning_rate": 8.758199506986655e-05,
+      "loss": 0.0443,
+      "step": 1340
+    },
+    {
+      "grad_norm": 0.2606576681137085,
+      "learning_rate": 8.73630576416373e-05,
+      "loss": 0.0399,
+      "step": 1350
+    },
+    {
+      "grad_norm": 0.43027952313423157,
+      "learning_rate": 8.714248583222726e-05,
+      "loss": 0.0449,
+      "step": 1360
+    },
+    {
+      "grad_norm": 0.3107929229736328,
+      "learning_rate": 8.692028929016196e-05,
+      "loss": 0.039,
+      "step": 1370
+    },
+    {
+      "grad_norm": 0.26511290669441223,
+      "learning_rate": 8.669647773503797e-05,
+      "loss": 0.0466,
+      "step": 1380
+    },
+    {
+      "grad_norm": 0.40483543276786804,
+      "learning_rate": 8.647106095709773e-05,
+      "loss": 0.04,
+      "step": 1390
+    },
+    {
+      "grad_norm": 0.3609035909175873,
+      "learning_rate": 8.624404881680139e-05,
+      "loss": 0.0468,
+      "step": 1400
+    },
+    {
+      "grad_norm": 0.3585107624530792,
+      "learning_rate": 8.601545124439535e-05,
+      "loss": 0.0396,
+      "step": 1410
+    },
+    {
+      "grad_norm": 0.39873549342155457,
+      "learning_rate": 8.5785278239478e-05,
+      "loss": 0.0506,
+      "step": 1420
+    },
+    {
+      "grad_norm": 0.26610010862350464,
+      "learning_rate": 8.555353987056224e-05,
+      "loss": 0.0446,
+      "step": 1430
+    },
+    {
+      "grad_norm": 0.29464486241340637,
+      "learning_rate": 8.532024627463505e-05,
+      "loss": 0.0449,
+      "step": 1440
+    },
+    {
+      "grad_norm": 0.3142518699169159,
+      "learning_rate": 8.508540765671407e-05,
+      "loss": 0.0449,
+      "step": 1450
+    },
+    {
+      "grad_norm": 0.3424524962902069,
+      "learning_rate": 8.484903428940121e-05,
+      "loss": 0.0462,
+      "step": 1460
+    },
+    {
+      "grad_norm": 0.4664880037307739,
+      "learning_rate": 8.461113651243334e-05,
+      "loss": 0.0499,
+      "step": 1470
+    },
+    {
+      "grad_norm": 0.28201568126678467,
+      "learning_rate": 8.437172473222987e-05,
+      "loss": 0.0467,
+      "step": 1480
+    },
+    {
+      "grad_norm": 0.4245680868625641,
+      "learning_rate": 8.413080942143767e-05,
+      "loss": 0.0439,
+      "step": 1490
+    },
+    {
+      "grad_norm": 0.42034053802490234,
+      "learning_rate": 8.388840111847288e-05,
+      "loss": 0.0532,
+      "step": 1500
+    },
+    {
+      "grad_norm": 0.3836017847061157,
+      "learning_rate": 8.364451042705998e-05,
+      "loss": 0.049,
+      "step": 1510
+    },
+    {
+      "grad_norm": 0.26328060030937195,
+      "learning_rate": 8.33991480157679e-05,
+      "loss": 0.0427,
+      "step": 1520
+    },
+    {
+      "grad_norm": 0.38004162907600403,
+      "learning_rate": 8.315232461754338e-05,
+      "loss": 0.0473,
+      "step": 1530
+    },
+    {
+      "grad_norm": 0.45609551668167114,
+      "learning_rate": 8.290405102924144e-05,
+      "loss": 0.0533,
+      "step": 1540
+    },
+    {
+      "grad_norm": 0.20536701381206512,
+      "learning_rate": 8.265433811115316e-05,
+      "loss": 0.0436,
+      "step": 1550
+    },
+    {
+      "grad_norm": 0.3798106908798218,
+      "learning_rate": 8.240319678653049e-05,
+      "loss": 0.0439,
+      "step": 1560
+    },
+    {
+      "grad_norm": 0.28516557812690735,
+      "learning_rate": 8.215063804110857e-05,
+      "loss": 0.0425,
+      "step": 1570
+    },
+    {
+      "grad_norm": 0.27644988894462585,
+      "learning_rate": 8.189667292262512e-05,
+      "loss": 0.0403,
+      "step": 1580
+    },
+    {
+      "grad_norm": 0.40204596519470215,
+      "learning_rate": 8.164131254033716e-05,
+      "loss": 0.0474,
+      "step": 1590
+    },
+    {
+      "grad_norm": 0.43881192803382874,
+      "learning_rate": 8.138456806453503e-05,
+      "loss": 0.0403,
+      "step": 1600
+    },
+    {
+      "grad_norm": 0.2845113277435303,
+      "learning_rate": 8.112645072605386e-05,
+      "loss": 0.0392,
+      "step": 1610
+    },
+    {
+      "grad_norm": 0.29198166728019714,
+      "learning_rate": 8.086697181578222e-05,
+      "loss": 0.0348,
+      "step": 1620
+    },
+    {
+      "grad_norm": 0.3617711663246155,
+      "learning_rate": 8.060614268416823e-05,
+      "loss": 0.0376,
+      "step": 1630
+    },
+    {
+      "grad_norm": 0.404875785112381,
+      "learning_rate": 8.034397474072309e-05,
+      "loss": 0.0454,
+      "step": 1640
+    },
+    {
+      "grad_norm": 0.22489239275455475,
+      "learning_rate": 8.008047945352193e-05,
+      "loss": 0.0425,
+      "step": 1650
+    },
+    {
+      "grad_norm": 0.27072638273239136,
+      "learning_rate": 7.981566834870225e-05,
+      "loss": 0.0383,
+      "step": 1660
+    },
+    {
+      "grad_norm": 0.24278584122657776,
+      "learning_rate": 7.954955300995961e-05,
+      "loss": 0.0421,
+      "step": 1670
+    },
+    {
+      "grad_norm": 0.3021833002567291,
+      "learning_rate": 7.928214507804104e-05,
+      "loss": 0.0429,
+      "step": 1680
+    },
+    {
+      "grad_norm": 0.38383153080940247,
+      "learning_rate": 7.901345625023576e-05,
+      "loss": 0.0436,
+      "step": 1690
+    },
+    {
+      "grad_norm": 0.510800838470459,
+      "learning_rate": 7.874349827986354e-05,
+      "loss": 0.0414,
+      "step": 1700
+    },
+    {
+      "grad_norm": 0.3297543525695801,
+      "learning_rate": 7.847228297576053e-05,
+      "loss": 0.0436,
+      "step": 1710
+    },
+    {
+      "grad_norm": 0.244502454996109,
+      "learning_rate": 7.819982220176276e-05,
+      "loss": 0.0453,
+      "step": 1720
+    },
+    {
+      "grad_norm": 0.4956775903701782,
+      "learning_rate": 7.792612787618714e-05,
+      "loss": 0.0466,
+      "step": 1730
+    },
+    {
+      "grad_norm": 0.25343015789985657,
+      "learning_rate": 7.765121197131009e-05,
+      "loss": 0.0424,
+      "step": 1740
+    },
+    {
+      "grad_norm": 0.26113173365592957,
+      "learning_rate": 7.737508651284391e-05,
+      "loss": 0.0416,
+      "step": 1750
+    },
+    {
+      "grad_norm": 0.2182377576828003,
+      "learning_rate": 7.709776357941069e-05,
+      "loss": 0.0441,
+      "step": 1760
+    },
+    {
+      "grad_norm": 0.3971315026283264,
+      "learning_rate": 7.681925530201392e-05,
+      "loss": 0.0438,
+      "step": 1770
+    },
+    {
+      "grad_norm": 0.28176161646842957,
+      "learning_rate": 7.65395738635079e-05,
+      "loss": 0.0408,
+      "step": 1780
+    },
+    {
+      "grad_norm": 0.34179630875587463,
+      "learning_rate": 7.62587314980648e-05,
+      "loss": 0.0479,
+      "step": 1790
+    },
+    {
+      "grad_norm": 0.32931986451148987,
+      "learning_rate": 7.597674049063947e-05,
+      "loss": 0.0429,
+      "step": 1800
+    },
+    {
+      "grad_norm": 0.39432066679000854,
+      "learning_rate": 7.569361317643211e-05,
+      "loss": 0.0431,
+      "step": 1810
+    },
+    {
+      "grad_norm": 0.3576849699020386,
+      "learning_rate": 7.540936194034865e-05,
+      "loss": 0.0384,
+      "step": 1820
+    },
+    {
+      "grad_norm": 0.23606491088867188,
+      "learning_rate": 7.512399921645901e-05,
+      "loss": 0.0357,
+      "step": 1830
+    },
+    {
+      "grad_norm": 0.2227422446012497,
+      "learning_rate": 7.483753748745317e-05,
+      "loss": 0.0367,
+      "step": 1840
+    },
+    {
+      "grad_norm": 0.2424551248550415,
+      "learning_rate": 7.454998928409516e-05,
+      "loss": 0.0385,
+      "step": 1850
+    },
+    {
+      "grad_norm": 0.33607348799705505,
+      "learning_rate": 7.426136718467493e-05,
+      "loss": 0.0412,
+      "step": 1860
+    },
+    {
+      "grad_norm": 0.2656790018081665,
+      "learning_rate": 7.397168381445812e-05,
+      "loss": 0.0402,
+      "step": 1870
+    },
+    {
+      "grad_norm": 0.27483707666397095,
+      "learning_rate": 7.368095184513377e-05,
+      "loss": 0.0332,
+      "step": 1880
+    },
+    {
+      "grad_norm": 0.32864516973495483,
+      "learning_rate": 7.338918399426005e-05,
+      "loss": 0.0352,
+      "step": 1890
+    },
+    {
+      "grad_norm": 0.2956995964050293,
+      "learning_rate": 7.309639302470801e-05,
+      "loss": 0.0464,
+      "step": 1900
+    },
+    {
+      "grad_norm": 0.3073318302631378,
+      "learning_rate": 7.280259174410312e-05,
+      "loss": 0.0378,
+      "step": 1910
+    },
+    {
+      "grad_norm": 0.3206455409526825,
+      "learning_rate": 7.250779300426517e-05,
+      "loss": 0.0428,
+      "step": 1920
+    },
+    {
+      "grad_norm": 0.40320953726768494,
+      "learning_rate": 7.22120097006461e-05,
+      "loss": 0.0434,
+      "step": 1930
+    },
+    {
+      "grad_norm": 0.2006569355726242,
+      "learning_rate": 7.191525477176577e-05,
+      "loss": 0.0367,
+      "step": 1940
+    },
+    {
+      "grad_norm": 0.26202642917633057,
+      "learning_rate": 7.161754119864616e-05,
+      "loss": 0.0387,
+      "step": 1950
+    },
+    {
+      "grad_norm": 0.36526966094970703,
+      "learning_rate": 7.131888200424339e-05,
+      "loss": 0.0344,
+      "step": 1960
+    },
+    {
+      "grad_norm": 0.3327760696411133,
+      "learning_rate": 7.101929025287816e-05,
+      "loss": 0.0401,
+      "step": 1970
+    },
+    {
+      "grad_norm": 0.36563345789909363,
+      "learning_rate": 7.071877904966423e-05,
+      "loss": 0.0474,
+      "step": 1980
+    },
+    {
+      "grad_norm": 0.27546799182891846,
+      "learning_rate": 7.04173615399351e-05,
+      "loss": 0.0376,
+      "step": 1990
+    },
+    {
+      "grad_norm": 0.3782268464565277,
+      "learning_rate": 7.011505090866913e-05,
+      "loss": 0.0324,
+      "step": 2000
+    },
+    {
+      "grad_norm": 0.43527549505233765,
+      "learning_rate": 6.981186037991271e-05,
+      "loss": 0.0358,
+      "step": 2010
+    },
+    {
+      "grad_norm": 0.327829509973526,
+      "learning_rate": 6.950780321620174e-05,
+      "loss": 0.0396,
+      "step": 2020
+    },
+    {
+      "grad_norm": 0.3114089369773865,
+      "learning_rate": 6.920289271798157e-05,
+      "loss": 0.04,
+      "step": 2030
+    },
+    {
+      "grad_norm": 0.34222158789634705,
+      "learning_rate": 6.889714222302517e-05,
+      "loss": 0.0417,
+      "step": 2040
+    },
+    {
+      "grad_norm": 0.19661477208137512,
+      "learning_rate": 6.85905651058497e-05,
+      "loss": 0.0395,
+      "step": 2050
+    },
+    {
+      "grad_norm": 0.28478899598121643,
+      "learning_rate": 6.82831747771314e-05,
+      "loss": 0.0343,
+      "step": 2060
+    },
+    {
+      "grad_norm": 0.32954296469688416,
+      "learning_rate": 6.797498468311907e-05,
+      "loss": 0.0331,
+      "step": 2070
+    },
+    {
+      "grad_norm": 0.21166834235191345,
+      "learning_rate": 6.766600830504585e-05,
+      "loss": 0.0328,
+      "step": 2080
+    },
+    {
+      "grad_norm": 0.2743442952632904,
+      "learning_rate": 6.735625915853942e-05,
+      "loss": 0.0374,
+      "step": 2090
+    },
+    {
+      "grad_norm": 0.23425710201263428,
+      "learning_rate": 6.70457507930309e-05,
+      "loss": 0.0299,
+      "step": 2100
+    },
+    {
+      "grad_norm": 0.26980066299438477,
+      "learning_rate": 6.673449679116215e-05,
+      "loss": 0.0343,
+      "step": 2110
+    },
+    {
+      "grad_norm": 0.3023163080215454,
+      "learning_rate": 6.642251076819148e-05,
+      "loss": 0.0357,
+      "step": 2120
+    },
+    {
+      "grad_norm": 0.22990819811820984,
+      "learning_rate": 6.610980637139827e-05,
+      "loss": 0.0357,
+      "step": 2130
+    },
+    {
+      "grad_norm": 0.20696483552455902,
+      "learning_rate": 6.579639727948583e-05,
+      "loss": 0.0368,
+      "step": 2140
+    },
+    {
+      "grad_norm": 0.2658112645149231,
+      "learning_rate": 6.548229720198315e-05,
+      "loss": 0.0391,
+      "step": 2150
+    },
+    {
+      "grad_norm": 0.3075551986694336,
+      "learning_rate": 6.516751987864517e-05,
+      "loss": 0.0337,
+      "step": 2160
+    },
+    {
+      "grad_norm": 0.3634116053581238,
+      "learning_rate": 6.485207907885175e-05,
+      "loss": 0.0353,
+      "step": 2170
+    },
+    {
+      "grad_norm": 0.2967609763145447,
+      "learning_rate": 6.453598860100536e-05,
+      "loss": 0.0386,
+      "step": 2180
+    },
+    {
+      "grad_norm": 0.19808591902256012,
+      "learning_rate": 6.421926227192749e-05,
+      "loss": 0.0316,
+      "step": 2190
+    },
+    {
+      "grad_norm": 0.29046621918678284,
+      "learning_rate": 6.390191394625381e-05,
+      "loss": 0.0348,
+      "step": 2200
+    },
+    {
+      "grad_norm": 0.3192194700241089,
+      "learning_rate": 6.358395750582817e-05,
+      "loss": 0.0332,
+      "step": 2210
+    },
+    {
+      "grad_norm": 0.4521180987358093,
+      "learning_rate": 6.326540685909532e-05,
+      "loss": 0.0378,
+      "step": 2220
+    },
+    {
+      "grad_norm": 0.2575896680355072,
+      "learning_rate": 6.294627594049249e-05,
+      "loss": 0.0402,
+      "step": 2230
+    },
+    {
+      "grad_norm": 0.19322504103183746,
+      "learning_rate": 6.262657870983989e-05,
+      "loss": 0.0302,
+      "step": 2240
+    },
+    {
+      "grad_norm": 0.29969140887260437,
+      "learning_rate": 6.230632915173009e-05,
+      "loss": 0.039,
+      "step": 2250
+    },
+    {
+      "grad_norm": 0.3442073464393616,
+      "learning_rate": 6.198554127491622e-05,
+      "loss": 0.0356,
+      "step": 2260
+    },
+    {
+      "grad_norm": 0.4116600453853607,
+      "learning_rate": 6.166422911169923e-05,
+      "loss": 0.0318,
+      "step": 2270
+    },
+    {
+      "grad_norm": 0.2713433504104614,
+      "learning_rate": 6.1342406717314e-05,
+      "loss": 0.0391,
+      "step": 2280
+    },
+    {
+      "grad_norm": 0.30095675587654114,
+      "learning_rate": 6.102008816931466e-05,
+      "loss": 0.0376,
+      "step": 2290
+    },
+    {
+      "grad_norm": 0.332207590341568,
+      "learning_rate": 6.069728756695866e-05,
+      "loss": 0.0313,
+      "step": 2300
+    },
+    {
+      "grad_norm": 0.2808273136615753,
+      "learning_rate": 6.037401903059008e-05,
+      "loss": 0.0356,
+      "step": 2310
+    },
+    {
+      "grad_norm": 0.2504296898841858,
+      "learning_rate": 6.005029670102195e-05,
+      "loss": 0.0277,
+      "step": 2320
+    },
+    {
+      "grad_norm": 0.3681153953075409,
+      "learning_rate": 5.972613473891766e-05,
+      "loss": 0.0427,
+      "step": 2330
+    },
+    {
+      "grad_norm": 0.42227885127067566,
+      "learning_rate": 5.940154732417158e-05,
+      "loss": 0.0338,
+      "step": 2340
+    },
+    {
+      "grad_norm": 0.4492588937282562,
+      "learning_rate": 5.907654865528876e-05,
+      "loss": 0.0382,
+      "step": 2350
+    },
+    {
+      "grad_norm": 0.3388882577419281,
+      "learning_rate": 5.875115294876381e-05,
+      "loss": 0.0309,
+      "step": 2360
+    },
+    {
+      "grad_norm": 0.34087541699409485,
+      "learning_rate": 5.842537443845908e-05,
+      "loss": 0.0346,
+      "step": 2370
+    },
+    {
+      "grad_norm": 0.3469429910182953,
+      "learning_rate": 5.809922737498198e-05,
+      "loss": 0.0284,
+      "step": 2380
+    },
+    {
+      "grad_norm": 0.2134752869606018,
+      "learning_rate": 5.777272602506165e-05,
+      "loss": 0.0365,
+      "step": 2390
+    },
+    {
+      "grad_norm": 0.3818526864051819,
+      "learning_rate": 5.744588467092483e-05,
+      "loss": 0.0383,
+      "step": 2400
+    },
+    {
+      "grad_norm": 0.2888863980770111,
+      "learning_rate": 5.7118717609671194e-05,
+      "loss": 0.0301,
+      "step": 2410
+    },
+    {
+      "grad_norm": 0.2873191237449646,
+      "learning_rate": 5.679123915264786e-05,
+      "loss": 0.033,
+      "step": 2420
+    },
+    {
+      "grad_norm": 0.23120439052581787,
+      "learning_rate": 5.646346362482342e-05,
+      "loss": 0.0311,
+      "step": 2430
+    },
+    {
+      "grad_norm": 0.34765511751174927,
+      "learning_rate": 5.613540536416132e-05,
+      "loss": 0.0289,
+      "step": 2440
+    },
+    {
+      "grad_norm": 0.39259421825408936,
+      "learning_rate": 5.5807078720992645e-05,
+      "loss": 0.0354,
+      "step": 2450
+    },
+    {
+      "grad_norm": 0.257591187953949,
+      "learning_rate": 5.547849805738836e-05,
+      "loss": 0.0368,
+      "step": 2460
+    },
+    {
+      "grad_norm": 0.25859981775283813,
+      "learning_rate": 5.514967774653118e-05,
+      "loss": 0.0302,
+      "step": 2470
+    },
+    {
+      "grad_norm": 0.20701059699058533,
+      "learning_rate": 5.482063217208674e-05,
+      "loss": 0.0294,
+      "step": 2480
+    },
+    {
+      "grad_norm": 0.351088285446167,
+      "learning_rate": 5.449137572757439e-05,
+      "loss": 0.027,
+      "step": 2490
+    },
+    {
+      "grad_norm": 0.2970297932624817,
+      "learning_rate": 5.4161922815737696e-05,
+      "loss": 0.0352,
+      "step": 2500
+    },
+    {
+      "grad_norm": 0.22500845789909363,
+      "learning_rate": 5.3832287847914276e-05,
+      "loss": 0.0344,
+      "step": 2510
+    },
+    {
+      "grad_norm": 0.1831464320421219,
+      "learning_rate": 5.35024852434055e-05,
+      "loss": 0.0277,
+      "step": 2520
+    },
+    {
+      "grad_norm": 0.2711646258831024,
+      "learning_rate": 5.317252942884567e-05,
+      "loss": 0.0286,
+      "step": 2530
+    },
+    {
+      "grad_norm": 0.30330950021743774,
+      "learning_rate": 5.284243483757109e-05,
+      "loss": 0.0346,
+      "step": 2540
+    },
+    {
+      "grad_norm": 0.293664813041687,
+      "learning_rate": 5.2512215908988484e-05,
+      "loss": 0.0343,
+      "step": 2550
+    },
+    {
+      "grad_norm": 0.24688895046710968,
+      "learning_rate": 5.218188708794357e-05,
+      "loss": 0.0287,
+      "step": 2560
+    },
+    {
+      "grad_norm": 0.3329557776451111,
+      "learning_rate": 5.18514628240891e-05,
+      "loss": 0.0291,
+      "step": 2570
+    },
+    {
+      "grad_norm": 0.2306559830904007,
+      "learning_rate": 5.1520957571252795e-05,
+      "loss": 0.0305,
+      "step": 2580
+    },
+    {
+      "grad_norm": 0.2929078936576843,
+      "learning_rate": 5.1190385786805106e-05,
+      "loss": 0.0349,
+      "step": 2590
+    },
+    {
+      "grad_norm": 0.18607696890830994,
+      "learning_rate": 5.085976193102677e-05,
+      "loss": 0.0286,
+      "step": 2600
+    },
+    {
+      "grad_norm": 0.2675177752971649,
+      "learning_rate": 5.052910046647634e-05,
+      "loss": 0.0302,
+      "step": 2610
+    },
+    {
+      "grad_norm": 0.20251218974590302,
+      "learning_rate": 5.0198415857357464e-05,
+      "loss": 0.0285,
+      "step": 2620
+    },
+    {
+      "grad_norm": 0.19350293278694153,
+      "learning_rate": 4.9867722568886223e-05,
+      "loss": 0.0292,
+      "step": 2630
+    },
+    {
+      "grad_norm": 0.28456851840019226,
+      "learning_rate": 4.9537035066658314e-05,
+      "loss": 0.0263,
+      "step": 2640
+    },
+    {
+      "grad_norm": 0.26157236099243164,
+      "learning_rate": 4.920636781601638e-05,
+      "loss": 0.0321,
+      "step": 2650
+    },
+    {
+      "grad_norm": 0.25256747007369995,
+      "learning_rate": 4.88757352814172e-05,
+      "loss": 0.0359,
+      "step": 2660
+    },
+    {
+      "grad_norm": 0.1635654717683792,
+      "learning_rate": 4.8545151925798924e-05,
+      "loss": 0.0292,
+      "step": 2670
+    },
+    {
+      "grad_norm": 0.2764965891838074,
+      "learning_rate": 4.821463220994848e-05,
+      "loss": 0.0349,
+      "step": 2680
+    },
+    {
+      "grad_norm": 0.22119125723838806,
+      "learning_rate": 4.788419059186895e-05,
+      "loss": 0.0282,
+      "step": 2690
+    },
+    {
+      "grad_norm": 0.274278461933136,
+      "learning_rate": 4.7553841526147205e-05,
+      "loss": 0.0318,
+      "step": 2700
+    },
+    {
+      "grad_norm": 0.26557064056396484,
+      "learning_rate": 4.722359946332156e-05,
+      "loss": 0.0308,
+      "step": 2710
+    },
+    {
+      "grad_norm": 0.2745562493801117,
+      "learning_rate": 4.6893478849249654e-05,
+      "loss": 0.027,
+      "step": 2720
+    },
+    {
+      "grad_norm": 0.23878148198127747,
+      "learning_rate": 4.656349412447664e-05,
+      "loss": 0.0292,
+      "step": 2730
+    },
+    {
+      "grad_norm": 0.211318701505661,
+      "learning_rate": 4.623365972360337e-05,
+      "loss": 0.0297,
+      "step": 2740
+    },
+    {
+      "grad_norm": 0.1693108081817627,
+      "learning_rate": 4.590399007465503e-05,
+      "loss": 0.0269,
+      "step": 2750
+    },
+    {
+      "grad_norm": 0.24284867942333221,
+      "learning_rate": 4.557449959845005e-05,
+      "loss": 0.0293,
+      "step": 2760
+    },
+    {
+      "grad_norm": 0.3161594271659851,
+      "learning_rate": 4.524520270796927e-05,
+      "loss": 0.0254,
+      "step": 2770
+    },
+    {
+      "grad_norm": 0.3319292664527893,
+      "learning_rate": 4.491611380772545e-05,
+      "loss": 0.0254,
+      "step": 2780
+    },
+    {
+      "grad_norm": 0.18327000737190247,
+      "learning_rate": 4.458724729313318e-05,
+      "loss": 0.0255,
+      "step": 2790
+    },
+    {
+      "grad_norm": 0.23681236803531647,
+      "learning_rate": 4.42586175498792e-05,
+      "loss": 0.0299,
+      "step": 2800
+    },
+    {
+      "grad_norm": 0.19555073976516724,
+      "learning_rate": 4.3930238953293094e-05,
+      "loss": 0.0296,
+      "step": 2810
+    },
+    {
+      "grad_norm": 0.3946215510368347,
+      "learning_rate": 4.360212586771847e-05,
+      "loss": 0.0288,
+      "step": 2820
+    },
+    {
+      "grad_norm": 0.24904312193393707,
+      "learning_rate": 4.327429264588463e-05,
+      "loss": 0.0245,
+      "step": 2830
+    },
+    {
+      "grad_norm": 0.1730395257472992,
+      "learning_rate": 4.2946753628278725e-05,
+      "loss": 0.0268,
+      "step": 2840
+    },
+    {
+      "grad_norm": 0.24282491207122803,
+      "learning_rate": 4.2619523142518474e-05,
+      "loss": 0.0258,
+      "step": 2850
+    },
+    {
+      "grad_norm": 0.17457225918769836,
+      "learning_rate": 4.229261550272539e-05,
+      "loss": 0.0273,
+      "step": 2860
+    },
+    {
+      "grad_norm": 0.23483504354953766,
+      "learning_rate": 4.196604500889868e-05,
+      "loss": 0.0322,
+      "step": 2870
+    },
+    {
+      "grad_norm": 0.2152663916349411,
+      "learning_rate": 4.163982594628969e-05,
+      "loss": 0.0266,
+      "step": 2880
+    },
+    {
+      "grad_norm": 0.17791348695755005,
+      "learning_rate": 4.131397258477702e-05,
+      "loss": 0.0264,
+      "step": 2890
+    },
+    {
+      "grad_norm": 0.22545894980430603,
+      "learning_rate": 4.0988499178242315e-05,
+      "loss": 0.028,
+      "step": 2900
+    },
+    {
+      "grad_norm": 0.3103232681751251,
+      "learning_rate": 4.066341996394678e-05,
+      "loss": 0.0389,
+      "step": 2910
+    },
+    {
+      "grad_norm": 0.337921679019928,
+      "learning_rate": 4.033874916190833e-05,
+      "loss": 0.0273,
+      "step": 2920
+    },
+    {
+      "grad_norm": 0.21746782958507538,
+      "learning_rate": 4.001450097427966e-05,
+      "loss": 0.0265,
+      "step": 2930
+    },
+    {
+      "grad_norm": 0.16620582342147827,
+      "learning_rate": 3.9690689584726894e-05,
+      "loss": 0.0251,
+      "step": 2940
+    },
+    {
+      "grad_norm": 0.3499850332736969,
+      "learning_rate": 3.936732915780923e-05,
+      "loss": 0.0254,
+      "step": 2950
+    },
+    {
+      "grad_norm": 0.2770783007144928,
+      "learning_rate": 3.904443383835929e-05,
+      "loss": 0.0274,
+      "step": 2960
+    },
+    {
+      "grad_norm": 0.2933313250541687,
+      "learning_rate": 3.872201775086437e-05,
+      "loss": 0.0277,
+      "step": 2970
+    },
+    {
+      "grad_norm": 0.2841816246509552,
+      "learning_rate": 3.8400094998848616e-05,
+      "loss": 0.0244,
+      "step": 2980
+    },
+    {
+      "grad_norm": 0.22859936952590942,
+      "learning_rate": 3.807867966425611e-05,
+      "loss": 0.0238,
+      "step": 2990
+    },
+    {
+      "grad_norm": 0.2738678753376007,
+      "learning_rate": 3.775778580683481e-05,
+      "loss": 0.027,
+      "step": 3000
+    },
+    {
+      "grad_norm": 0.31196361780166626,
+      "learning_rate": 3.743742746352156e-05,
+      "loss": 0.0255,
+      "step": 3010
+    },
+    {
+      "grad_norm": 0.2890107035636902,
+      "learning_rate": 3.711761864782817e-05,
+      "loss": 0.0303,
+      "step": 3020
+    },
+    {
+      "grad_norm": 0.29215967655181885,
+      "learning_rate": 3.679837334922825e-05,
+      "loss": 0.0282,
+      "step": 3030
+    },
+    {
+      "grad_norm": 0.17950771749019623,
+      "learning_rate": 3.647970553254538e-05,
+      "loss": 0.0244,
+      "step": 3040
+    },
+    {
+      "grad_norm": 0.20166261494159698,
+      "learning_rate": 3.61616291373422e-05,
+      "loss": 0.0277,
+      "step": 3050
+    },
+    {
+      "grad_norm": 0.18279321491718292,
+      "learning_rate": 3.584415807731065e-05,
+      "loss": 0.0252,
+      "step": 3060
+    },
+    {
+      "grad_norm": 0.24165138602256775,
+      "learning_rate": 3.552730623966337e-05,
+      "loss": 0.0269,
+      "step": 3070
+    },
+    {
+      "grad_norm": 0.32261815667152405,
+      "learning_rate": 3.521108748452617e-05,
+      "loss": 0.0254,
+      "step": 3080
+    },
+    {
+      "grad_norm": 0.3094899654388428,
+      "learning_rate": 3.489551564433186e-05,
+      "loss": 0.0279,
+      "step": 3090
+    },
+    {
+      "grad_norm": 0.2739303708076477,
+      "learning_rate": 3.4580604523215006e-05,
+      "loss": 0.0234,
+      "step": 3100
+    },
+    {
+      "grad_norm": 0.20867018401622772,
+      "learning_rate": 3.4266367896408216e-05,
+      "loss": 0.0314,
+      "step": 3110
+    },
+    {
+      "grad_norm": 0.2572309374809265,
+      "learning_rate": 3.3952819509639534e-05,
+      "loss": 0.0259,
+      "step": 3120
+    },
+    {
+      "grad_norm": 0.24716201424598694,
+      "learning_rate": 3.3639973078531165e-05,
+      "loss": 0.0284,
+      "step": 3130
+    },
+    {
+      "grad_norm": 0.28539377450942993,
+      "learning_rate": 3.332784228799947e-05,
+      "loss": 0.0251,
+      "step": 3140
+    },
+    {
+      "grad_norm": 0.27438855171203613,
+      "learning_rate": 3.301644079165638e-05,
+      "loss": 0.0271,
+      "step": 3150
+    },
+    {
+      "grad_norm": 0.35537853837013245,
+      "learning_rate": 3.27057822112122e-05,
+      "loss": 0.0245,
+      "step": 3160
+    },
+    {
+      "grad_norm": 0.3076242208480835,
+      "learning_rate": 3.239588013587958e-05,
+      "loss": 0.0241,
+      "step": 3170
+    },
+    {
+      "grad_norm": 0.2089366912841797,
+      "learning_rate": 3.208674812177926e-05,
+      "loss": 0.0235,
+      "step": 3180
+    },
+    {
+      "grad_norm": 0.28236523270606995,
+      "learning_rate": 3.177839969134698e-05,
+      "loss": 0.0248,
+      "step": 3190
+    },
+    {
+      "grad_norm": 0.26007798314094543,
+      "learning_rate": 3.1470848332742e-05,
+      "loss": 0.0269,
+      "step": 3200
+    },
+    {
+      "grad_norm": 0.2209538072347641,
+      "learning_rate": 3.116410749925708e-05,
+      "loss": 0.0247,
+      "step": 3210
+    },
+    {
+      "grad_norm": 0.16867880523204803,
+      "learning_rate": 3.085819060872995e-05,
+      "loss": 0.02,
+      "step": 3220
+    },
+    {
+      "grad_norm": 0.20084629952907562,
+      "learning_rate": 3.055311104295648e-05,
+      "loss": 0.0217,
+      "step": 3230
+    },
+    {
+      "grad_norm": 0.1993807554244995,
+      "learning_rate": 3.024888214710517e-05,
+      "loss": 0.0218,
+      "step": 3240
+    },
+    {
+      "grad_norm": 0.1759999543428421,
+      "learning_rate": 2.994551722913349e-05,
+      "loss": 0.0209,
+      "step": 3250
+    },
+    {
+      "grad_norm": 0.2935018837451935,
+      "learning_rate": 2.9643029559205727e-05,
+      "loss": 0.0303,
+      "step": 3260
+    },
+    {
+      "grad_norm": 0.1872728019952774,
+      "learning_rate": 2.934143236911248e-05,
+      "loss": 0.0259,
+      "step": 3270
+    },
+    {
+      "grad_norm": 0.34932175278663635,
+      "learning_rate": 2.90407388516919e-05,
+      "loss": 0.0239,
+      "step": 3280
+    },
+    {
+      "grad_norm": 0.16118347644805908,
+      "learning_rate": 2.8740962160252495e-05,
+      "loss": 0.0232,
+      "step": 3290
+    },
+    {
+      "grad_norm": 0.23605194687843323,
+      "learning_rate": 2.844211540799797e-05,
+      "loss": 0.0225,
+      "step": 3300
+    },
+    {
+      "grad_norm": 0.13263435661792755,
+      "learning_rate": 2.8144211667453368e-05,
+      "loss": 0.0219,
+      "step": 3310
+    },
+    {
+      "grad_norm": 0.2748834788799286,
+      "learning_rate": 2.7847263969893344e-05,
+      "loss": 0.0248,
+      "step": 3320
+    },
+    {
+      "grad_norm": 0.2700659930706024,
+      "learning_rate": 2.7551285304772206e-05,
+      "loss": 0.024,
+      "step": 3330
+    },
+    {
+      "grad_norm": 0.23235876858234406,
+      "learning_rate": 2.7256288619155567e-05,
+      "loss": 0.0222,
+      "step": 3340
+    },
+    {
+      "grad_norm": 0.3042518198490143,
+      "learning_rate": 2.6962286817154158e-05,
+      "loss": 0.0294,
+      "step": 3350
+    },
+    {
+      "grad_norm": 0.2259960025548935,
+      "learning_rate": 2.6669292759359166e-05,
+      "loss": 0.0202,
+      "step": 3360
+    },
+    {
+      "grad_norm": 0.274394690990448,
+      "learning_rate": 2.637731926227993e-05,
+      "loss": 0.0325,
+      "step": 3370
+    },
+    {
+      "grad_norm": 0.2576564848423004,
+      "learning_rate": 2.6086379097783033e-05,
+      "loss": 0.0248,
+      "step": 3380
+    },
+    {
+      "grad_norm": 0.29949021339416504,
+      "learning_rate": 2.579648499253377e-05,
+      "loss": 0.0258,
+      "step": 3390
+    },
+    {
+      "grad_norm": 0.2665102481842041,
+      "learning_rate": 2.5507649627439466e-05,
+      "loss": 0.0279,
+      "step": 3400
+    },
+    {
+      "grad_norm": 0.21922795474529266,
+      "learning_rate": 2.5219885637094653e-05,
+      "loss": 0.0239,
+      "step": 3410
+    },
+    {
+      "grad_norm": 0.13632823526859283,
+      "learning_rate": 2.4933205609228533e-05,
+      "loss": 0.024,
+      "step": 3420
+    },
+    {
+      "grad_norm": 0.2548131048679352,
+      "learning_rate": 2.464762208415419e-05,
+      "loss": 0.0221,
+      "step": 3430
+    },
+    {
+      "grad_norm": 0.3076316714286804,
+      "learning_rate": 2.4363147554220213e-05,
+      "loss": 0.023,
+      "step": 3440
+    },
+    {
+      "grad_norm": 0.22425901889801025,
+      "learning_rate": 2.407979446326411e-05,
+      "loss": 0.0216,
+      "step": 3450
+    },
+    {
+      "grad_norm": 0.2095557600259781,
+      "learning_rate": 2.379757520606799e-05,
+      "loss": 0.023,
+      "step": 3460
+    },
+    {
+      "grad_norm": 0.17088909447193146,
+      "learning_rate": 2.3516502127816455e-05,
+      "loss": 0.0216,
+      "step": 3470
+    },
+    {
+      "grad_norm": 0.21437765657901764,
+      "learning_rate": 2.323658752355647e-05,
+      "loss": 0.0229,
+      "step": 3480
+    },
+    {
+      "grad_norm": 0.17913414537906647,
+      "learning_rate": 2.2957843637659654e-05,
+      "loss": 0.0223,
+      "step": 3490
+    },
+    {
+      "grad_norm": 0.2331356257200241,
+      "learning_rate": 2.2680282663286552e-05,
+      "loss": 0.0286,
+      "step": 3500
+    },
+    {
+      "grad_norm": 0.20977725088596344,
+      "learning_rate": 2.2403916741853364e-05,
+      "loss": 0.0249,
+      "step": 3510
+    },
+    {
+      "grad_norm": 0.17635555565357208,
+      "learning_rate": 2.2128757962500817e-05,
+      "loss": 0.0213,
+      "step": 3520
+    },
+    {
+      "grad_norm": 0.1792505383491516,
+      "learning_rate": 2.1854818361565275e-05,
+      "loss": 0.0207,
+      "step": 3530
+    },
+    {
+      "grad_norm": 0.19617965817451477,
+      "learning_rate": 2.1582109922052364e-05,
+      "loss": 0.0268,
+      "step": 3540
+    },
+    {
+      "grad_norm": 0.22589251399040222,
+      "learning_rate": 2.1310644573112635e-05,
+      "loss": 0.0219,
+      "step": 3550
+    },
+    {
+      "grad_norm": 0.16943573951721191,
+      "learning_rate": 2.1040434189519924e-05,
+      "loss": 0.0219,
+      "step": 3560
+    },
+    {
+      "grad_norm": 0.2124646008014679,
+      "learning_rate": 2.0771490591151733e-05,
+      "loss": 0.0267,
+      "step": 3570
+    },
+    {
+      "grad_norm": 0.2132735699415207,
+      "learning_rate": 2.0503825542472317e-05,
+      "loss": 0.0248,
+      "step": 3580
+    },
+    {
+      "grad_norm": 0.2320823222398758,
+      "learning_rate": 2.023745075201805e-05,
+      "loss": 0.0211,
+      "step": 3590
+    },
+    {
+      "grad_norm": 0.19216440618038177,
+      "learning_rate": 1.9972377871885157e-05,
+      "loss": 0.0254,
+      "step": 3600
+    },
+    {
+      "grad_norm": 0.15354987978935242,
+      "learning_rate": 1.970861849722017e-05,
+      "loss": 0.0186,
+      "step": 3610
+    },
+    {
+      "grad_norm": 0.15961474180221558,
+      "learning_rate": 1.9446184165712587e-05,
+      "loss": 0.0185,
+      "step": 3620
+    },
+    {
+      "grad_norm": 0.18211418390274048,
+      "learning_rate": 1.9185086357090214e-05,
+      "loss": 0.0231,
+      "step": 3630
+    },
+    {
+      "grad_norm": 0.21647849678993225,
+      "learning_rate": 1.8925336492617057e-05,
+      "loss": 0.0228,
+      "step": 3640
+    },
+    {
+      "grad_norm": 0.18972018361091614,
+      "learning_rate": 1.8666945934593666e-05,
+      "loss": 0.021,
+      "step": 3650
+    },
+    {
+      "grad_norm": 0.13877266645431519,
+      "learning_rate": 1.8409925985860126e-05,
+      "loss": 0.0177,
+      "step": 3660
+    },
+    {
+      "grad_norm": 0.15036696195602417,
+      "learning_rate": 1.8154287889301603e-05,
+      "loss": 0.021,
+      "step": 3670
+    },
+    {
+      "grad_norm": 0.1516561061143875,
+      "learning_rate": 1.7900042827356612e-05,
+      "loss": 0.0233,
+      "step": 3680
+    },
+    {
+      "grad_norm": 0.21299676597118378,
+      "learning_rate": 1.76472019215278e-05,
+      "loss": 0.0245,
+      "step": 3690
+    },
+    {
+      "grad_norm": 0.24672634899616241,
+      "learning_rate": 1.739577623189545e-05,
+      "loss": 0.0186,
+      "step": 3700
+    },
+    {
+      "grad_norm": 0.15656235814094543,
+      "learning_rate": 1.7145776756633768e-05,
+      "loss": 0.0254,
+      "step": 3710
+    },
+    {
+      "grad_norm": 0.23428286612033844,
+      "learning_rate": 1.6897214431529646e-05,
+      "loss": 0.0209,
+      "step": 3720
+    },
+    {
+      "grad_norm": 0.1686718761920929,
+      "learning_rate": 1.6650100129504475e-05,
+      "loss": 0.0197,
+      "step": 3730
+    },
+    {
+      "grad_norm": 0.2660563886165619,
+      "learning_rate": 1.6404444660138335e-05,
+      "loss": 0.0256,
+      "step": 3740
+    },
+    {
+      "grad_norm": 0.24974121153354645,
+      "learning_rate": 1.616025876919725e-05,
+      "loss": 0.0168,
+      "step": 3750
+    },
+    {
+      "grad_norm": 0.19375593960285187,
+      "learning_rate": 1.5917553138163172e-05,
+      "loss": 0.0244,
+      "step": 3760
+    },
+    {
+      "grad_norm": 0.23583021759986877,
+      "learning_rate": 1.5676338383766632e-05,
+      "loss": 0.0284,
+      "step": 3770
+    },
+    {
+      "grad_norm": 0.20768067240715027,
+      "learning_rate": 1.5436625057522447e-05,
+      "loss": 0.0214,
+      "step": 3780
+    },
+    {
+      "grad_norm": 0.2571140229701996,
+      "learning_rate": 1.519842364526804e-05,
+      "loss": 0.02,
+      "step": 3790
+    },
+    {
+      "grad_norm": 0.2075309306383133,
+      "learning_rate": 1.4961744566704855e-05,
+      "loss": 0.0198,
+      "step": 3800
+    },
+    {
+      "grad_norm": 0.16042163968086243,
+      "learning_rate": 1.4726598174942551e-05,
+      "loss": 0.0158,
+      "step": 3810
+    },
+    {
+      "grad_norm": 0.19813798367977142,
+      "learning_rate": 1.4492994756046035e-05,
+      "loss": 0.0205,
+      "step": 3820
+    },
+    {
+      "grad_norm": 0.18999746441841125,
+      "learning_rate": 1.4260944528585645e-05,
+      "loss": 0.0226,
+      "step": 3830
+    },
+    {
+      "grad_norm": 0.20053531229496002,
+      "learning_rate": 1.4030457643190048e-05,
+      "loss": 0.0182,
+      "step": 3840
+    },
+    {
+      "grad_norm": 0.18872573971748352,
+      "learning_rate": 1.3801544182102311e-05,
+      "loss": 0.0178,
+      "step": 3850
+    },
+    {
+      "grad_norm": 0.22023218870162964,
+      "learning_rate": 1.3574214158738763e-05,
+      "loss": 0.022,
+      "step": 3860
+    },
+    {
+      "grad_norm": 0.3135218024253845,
+      "learning_rate": 1.3348477517251101e-05,
+      "loss": 0.0206,
+      "step": 3870
+    },
+    {
+      "grad_norm": 0.14147350192070007,
+      "learning_rate": 1.312434413209131e-05,
+      "loss": 0.0183,
+      "step": 3880
+    },
+    {
+      "grad_norm": 0.1488725244998932,
+      "learning_rate": 1.2901823807579727e-05,
+      "loss": 0.0185,
+      "step": 3890
+    },
+    {
+      "grad_norm": 0.2105879932641983,
+      "learning_rate": 1.2680926277476245e-05,
+      "loss": 0.0215,
+      "step": 3900
+    },
+    {
+      "grad_norm": 0.24176077544689178,
+      "learning_rate": 1.2461661204554397e-05,
+      "loss": 0.0167,
+      "step": 3910
+    },
+    {
+      "grad_norm": 0.1613641083240509,
+      "learning_rate": 1.2244038180178835e-05,
+      "loss": 0.0189,
+      "step": 3920
+    },
+    {
+      "grad_norm": 0.1966574788093567,
+      "learning_rate": 1.2028066723885612e-05,
+      "loss": 0.0203,
+      "step": 3930
+    },
+    {
+      "grad_norm": 0.2452186942100525,
+      "learning_rate": 1.1813756282965888e-05,
+      "loss": 0.0226,
+      "step": 3940
+    },
+    {
+      "grad_norm": 0.27143678069114685,
+      "learning_rate": 1.1601116232052638e-05,
+      "loss": 0.0241,
+      "step": 3950
+    },
+    {
+      "grad_norm": 0.26432090997695923,
+      "learning_rate": 1.1390155872710517e-05,
+      "loss": 0.0185,
+      "step": 3960
+    },
+    {
+      "grad_norm": 0.20945382118225098,
+      "learning_rate": 1.1180884433029087e-05,
+      "loss": 0.0229,
+      "step": 3970
+    },
+    {
+      "grad_norm": 0.16551558673381805,
+      "learning_rate": 1.097331106721904e-05,
+      "loss": 0.0197,
+      "step": 3980
+    },
+    {
+      "grad_norm": 0.16200856864452362,
+      "learning_rate": 1.0767444855211862e-05,
+      "loss": 0.0197,
+      "step": 3990
+    },
+    {
+      "grad_norm": 0.24197158217430115,
+      "learning_rate": 1.0563294802262558e-05,
+      "loss": 0.0222,
+      "step": 4000
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 5000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 2000,
+  "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": 32,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-4000/training_args.bin

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

checkpoint-4000/wandb_config.json

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

checkpoint-5000/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-5000/embodiment_id.json

@@ -0,0 +1,11 @@
+{
+  "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
+}

checkpoint-5000/experiment_cfg/conf.yaml

@@ -0,0 +1,210 @@
+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:
+    brightness: 0.3
+    contrast: 0.4
+    saturation: 0.5
+    hue: 0.08
+  use_albumentations_transforms: true
+  extra_augmentation_config: null
+  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:
+    - ./demo_data/SOarmtest/viis0002/pick_lego
+    embodiment_tag: new_embodiment
+    mix_ratio: 1.0
+    dataset_type: physical_embodiment
+    val_dataset_path: null
+  modality_configs:
+    new_embodiment:
+      video:
+        delta_indices:
+        - 0
+        modality_keys:
+        - front
+        - wrist
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - single_arm
+        - gripper
+        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:
+        - single_arm
+        - gripper
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs:
+        - 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:
+        - 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: 100000
+  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: torchcodec
+training:
+  output_dir: ./demo_data/SOarmtest/saved_checkpoints
+  experiment_name: null
+  max_steps: 5000
+  global_batch_size: 32
+  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: 2000
+  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: 4
+  remove_unused_columns: false
+  use_wandb: true
+  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: 5000
+save_steps: 2000

checkpoint-5000/experiment_cfg/config.yaml

@@ -0,0 +1,244 @@
+!!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:
+    - ./demo_data/SOarmtest/viis0002/pick_lego
+    dataset_type: physical_embodiment
+    embodiment_tag: new_embodiment
+    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:
+    new_embodiment:
+      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: null
+          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:
+        - single_arm
+        - 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:
+        - single_arm
+        - gripper
+        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
+        - wrist
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 100000
+  override_pretraining_statistics: false
+  random_chop: 0.0
+  seed: 42
+  shard_size: 1024
+  shuffle: true
+  subsample_ratio: 1.0
+  video_backend: torchcodec
+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:
+    brightness: 0.3
+    contrast: 0.4
+    hue: 0.08
+    saturation: 0.5
+  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
+  extra_augmentation_config: 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: 4
+  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: 32
+  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: 5000
+  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: ./demo_data/SOarmtest/saved_checkpoints
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_steps: 2000
+  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: true
+  wandb_project: finetune-gr00t-n1d6
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

checkpoint-5000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,824 @@
+{
+  "new_embodiment": {
+    "state": {
+      "single_arm": {
+        "min": [
+          -6.8131866455078125,
+          -102.19779968261719,
+          -40.43955993652344,
+          39.9560432434082,
+          -66.32967376708984
+        ],
+        "max": [
+          34.593406677246094,
+          42.769229888916016,
+          96.26373291015625,
+          98.06593322753906,
+          96.74725341796875
+        ],
+        "mean": [
+          10.6347642478831,
+          -40.60343800114444,
+          40.109735021899446,
+          66.37540221377535,
+          1.723290164849921
+        ],
+        "std": [
+          13.663472797431913,
+          53.514238207466356,
+          47.089223703864675,
+          15.49883594624875,
+          29.442594278901634
+        ],
+        "q01": [
+          -2.8227532818225196,
+          -102.04170769592945,
+          -18.832328638595566,
+          47.303414007272124,
+          -23.538191580154248
+        ],
+        "q99": [
+          31.17592093949654,
+          27.78409835219628,
+          96.09969459877085,
+          91.52694173019303,
+          34.804566667323705
+        ]
+      },
+      "gripper": {
+        "min": [
+          1.1597938537597656
+        ],
+        "max": [
+          30.476804733276367
+        ],
+        "mean": [
+          4.725540477545816
+        ],
+        "std": [
+          5.719213168100463
+        ],
+        "q01": [
+          1.219114189586798
+        ],
+        "q99": [
+          18.975632466717574
+        ]
+      }
+    },
+    "action": {
+      "single_arm": {
+        "min": [
+          -7.340659141540527,
+          -103.73626708984375,
+          -43.736263275146484,
+          38.63736343383789,
+          -67.47252655029297
+        ],
+        "max": [
+          35.20878982543945,
+          42.5494499206543,
+          96.74725341796875,
+          100.61538696289062,
+          97.27472686767578
+        ],
+        "mean": [
+          10.654706232867026,
+          -41.43508459659817,
+          38.95115310430438,
+          66.17248404139285,
+          1.6835930980913996
+        ],
+        "std": [
+          13.792301015873905,
+          53.27088262273843,
+          48.06070434685447,
+          15.735264444521983,
+          29.50880036285619
+        ],
+        "q01": [
+          -3.007477344415986,
+          -102.62527508767714,
+          -21.606171980367048,
+          46.25193771491276,
+          -24.028302588862974
+        ],
+        "q99": [
+          31.658033587037096,
+          27.49569149494191,
+          96.40327760792744,
+          92.26730265898277,
+          35.153253393452815
+        ]
+      },
+      "gripper": {
+        "min": [
+          0.0
+        ],
+        "max": [
+          30.952381134033203
+        ],
+        "mean": [
+          3.5887609717193083
+        ],
+        "std": [
+          5.951741507358969
+        ],
+        "q01": [
+          0.7785521832201806
+        ],
+        "q99": [
+          19.41596701897272
+        ]
+      }
+    },
+    "relative_action": {
+      "single_arm": {
+        "min": [
+          [
+            -6.153846740722656,
+            -12.615386962890625,
+            -13.406593322753906,
+            -9.406593322753906,
+            -13.450549125671387
+          ],
+          [
+            -7.824176788330078,
+            -15.692306518554688,
+            -16.39560317993164,
+            -11.86812973022461,
+            -13.450549125671387
+          ],
+          [
+            -9.582418441772461,
+            -18.681320190429688,
+            -19.296703338623047,
+            -14.241756439208984,
+            -15.472527503967285
+          ],
+          [
+            -11.252748489379883,
+            -21.670331954956055,
+            -22.19780158996582,
+            -16.615379333496094,
+            -16.96703338623047
+          ],
+          [
+            -12.923077583312988,
+            -24.65934181213379,
+            -24.923078536987305,
+            -18.901100158691406,
+            -18.461538314819336
+          ],
+          [
+            -14.593405723571777,
+            -27.648353576660156,
+            -27.736263275146484,
+            -21.186813354492188,
+            -19.516483306884766
+          ],
+          [
+            -16.087913513183594,
+            -30.637365341186523,
+            -30.46154022216797,
+            -23.384613037109375,
+            -21.010990142822266
+          ],
+          [
+            -17.58241844177246,
+            -33.45055389404297,
+            -33.09890365600586,
+            -25.406593322753906,
+            -23.032970428466797
+          ],
+          [
+            -18.813186645507812,
+            -36.351646423339844,
+            -35.648353576660156,
+            -27.428573608398438,
+            -24.967037200927734
+          ],
+          [
+            -20.043954849243164,
+            -39.164833068847656,
+            -38.109893798828125,
+            -29.09890365600586,
+            -27.076923370361328
+          ],
+          [
+            -21.098899841308594,
+            -42.06593322753906,
+            -40.57142639160156,
+            -30.769229888916016,
+            -29.274723052978516
+          ],
+          [
+            -22.153844833374023,
+            -44.879119873046875,
+            -42.857139587402344,
+            -32.351646423339844,
+            -31.38461685180664
+          ],
+          [
+            -23.208789825439453,
+            -47.69231033325195,
+            -45.05494689941406,
+            -33.75823974609375,
+            -33.494503021240234
+          ],
+          [
+            -24.08791160583496,
+            -50.68132019042969,
+            -47.164833068847656,
+            -34.989013671875,
+            -35.51648712158203
+          ],
+          [
+            -24.967031478881836,
+            -53.67033004760742,
+            -49.098899841308594,
+            -36.04396057128906,
+            -37.626373291015625
+          ],
+          [
+            -25.758241653442383,
+            -56.747249603271484,
+            -50.94505310058594,
+            -37.01099395751953,
+            -39.648353576660156
+          ]
+        ],
+        "max": [
+          [
+            5.8021979331970215,
+            14.021980285644531,
+            12.527473449707031,
+            8.703296661376953,
+            9.670330047607422
+          ],
+          [
+            6.505494117736816,
+            17.626373291015625,
+            16.131866455078125,
+            11.164833068847656,
+            12.219779968261719
+          ],
+          [
+            7.82417631149292,
+            21.142852783203125,
+            19.824176788330078,
+            13.626373291015625,
+            14.681318283081055
+          ],
+          [
+            9.05494499206543,
+            24.395599365234375,
+            23.428571701049805,
+            15.824176788330078,
+            17.142858505249023
+          ],
+          [
+            10.373626708984375,
+            27.384613037109375,
+            26.94505500793457,
+            18.021976470947266,
+            19.69230842590332
+          ],
+          [
+            11.69230842590332,
+            30.197803497314453,
+            30.461538314819336,
+            20.043956756591797,
+            22.153846740722656
+          ],
+          [
+            13.01098918914795,
+            32.83516311645508,
+            33.97801971435547,
+            22.241756439208984,
+            24.703298568725586
+          ],
+          [
+            14.153846740722656,
+            35.384613037109375,
+            37.406593322753906,
+            24.35165023803711,
+            27.07692527770996
+          ],
+          [
+            15.296703338623047,
+            37.75823974609375,
+            40.74725341796875,
+            26.197803497314453,
+            29.362638473510742
+          ],
+          [
+            16.351648330688477,
+            40.04395294189453,
+            44.08790969848633,
+            27.956043243408203,
+            31.560440063476562
+          ],
+          [
+            17.318681716918945,
+            42.241756439208984,
+            47.340660095214844,
+            29.71428680419922,
+            33.67033004760742
+          ],
+          [
+            18.197803497314453,
+            44.527469635009766,
+            50.68131637573242,
+            31.384613037109375,
+            35.95604705810547
+          ],
+          [
+            18.901100158691406,
+            46.4615364074707,
+            53.93406295776367,
+            32.96703338623047,
+            38.5054931640625
+          ],
+          [
+            19.516483306884766,
+            48.48351287841797,
+            57.010986328125,
+            34.28571319580078,
+            41.14285659790039
+          ],
+          [
+            20.043956756591797,
+            50.32966613769531,
+            60.0,
+            35.604393005371094,
+            43.516483306884766
+          ],
+          [
+            20.395605087280273,
+            52.26373291015625,
+            62.989013671875,
+            36.659339904785156,
+            45.80220031738281
+          ]
+        ],
+        "mean": [
+          [
+            0.024023909121751785,
+            -0.8380120396614075,
+            -1.1888031959533691,
+            -0.20351642370224,
+            -0.03874167427420616
+          ],
+          [
+            0.02381421998143196,
+            -0.8381950259208679,
+            -1.1877086162567139,
+            -0.20610879361629486,
+            -0.038331978023052216
+          ],
+          [
+            0.02358066476881504,
+            -0.8384636640548706,
+            -1.186354160308838,
+            -0.20878157019615173,
+            -0.03791269287467003
+          ],
+          [
+            0.02331378683447838,
+            -0.8388302326202393,
+            -1.1847561597824097,
+            -0.21150696277618408,
+            -0.03749352693557739
+          ],
+          [
+            0.022984733805060387,
+            -0.8392952680587769,
+            -1.182924509048462,
+            -0.21426069736480713,
+            -0.037074267864227295
+          ],
+          [
+            0.02258922904729843,
+            -0.8398675322532654,
+            -1.1808984279632568,
+            -0.21703894436359406,
+            -0.03665487840771675
+          ],
+          [
+            0.02210799790918827,
+            -0.840530514717102,
+            -1.1787015199661255,
+            -0.21982122957706451,
+            -0.03624051436781883
+          ],
+          [
+            0.021555287763476372,
+            -0.8412824869155884,
+            -1.176315188407898,
+            -0.2226090133190155,
+            -0.03585929051041603
+          ],
+          [
+            0.020935896784067154,
+            -0.8421463966369629,
+            -1.1736962795257568,
+            -0.22543004155158997,
+            -0.035549700260162354
+          ],
+          [
+            0.020240070298314095,
+            -0.8431134819984436,
+            -1.1708482503890991,
+            -0.2282552719116211,
+            -0.03533996269106865
+          ],
+          [
+            0.01948723942041397,
+            -0.8441924452781677,
+            -1.1677296161651611,
+            -0.23114195466041565,
+            -0.03522548824548721
+          ],
+          [
+            0.018677080050110817,
+            -0.8453867435455322,
+            -1.164296269416809,
+            -0.23409080505371094,
+            -0.03516348451375961
+          ],
+          [
+            0.017790833488106728,
+            -0.846683919429779,
+            -1.1604851484298706,
+            -0.23711234331130981,
+            -0.035125281661748886
+          ],
+          [
+            0.016804957762360573,
+            -0.8481130599975586,
+            -1.156294584274292,
+            -0.24017632007598877,
+            -0.03509661555290222
+          ],
+          [
+            0.015680411830544472,
+            -0.8497316241264343,
+            -1.151717185974121,
+            -0.24328306317329407,
+            -0.035068120807409286
+          ],
+          [
+            0.01436557900160551,
+            -0.8516368269920349,
+            -1.1466983556747437,
+            -0.24642817676067352,
+            -0.035063356161117554
+          ]
+        ],
+        "std": [
+          [
+            1.0472843647003174,
+            3.3955397605895996,
+            3.209542989730835,
+            1.5478397607803345,
+            1.5488325357437134
+          ],
+          [
+            1.3025773763656616,
+            4.228511333465576,
+            3.9392380714416504,
+            1.9040354490280151,
+            1.9508968591690063
+          ],
+          [
+            1.5567840337753296,
+            5.061123371124268,
+            4.678185939788818,
+            2.259927749633789,
+            2.3464455604553223
+          ],
+          [
+            1.8073352575302124,
+            5.890169143676758,
+            5.418741226196289,
+            2.6112139225006104,
+            2.734940528869629
+          ],
+          [
+            2.053133726119995,
+            6.713922023773193,
+            6.156966209411621,
+            2.9563238620758057,
+            3.1171467304229736
+          ],
+          [
+            2.2936792373657227,
+            7.531452178955078,
+            6.89060115814209,
+            3.2943806648254395,
+            3.4939496517181396
+          ],
+          [
+            2.528961658477783,
+            8.342284202575684,
+            7.61845064163208,
+            3.6247124671936035,
+            3.8658618927001953
+          ],
+          [
+            2.7590219974517822,
+            9.146015167236328,
+            8.339773178100586,
+            3.9470086097717285,
+            4.233205795288086
+          ],
+          [
+            2.9840755462646484,
+            9.942627906799316,
+            9.053919792175293,
+            4.261190414428711,
+            4.596113681793213
+          ],
+          [
+            3.2043535709381104,
+            10.731932640075684,
+            9.760570526123047,
+            4.567060947418213,
+            4.954709053039551
+          ],
+          [
+            3.420076608657837,
+            11.51407527923584,
+            10.459877967834473,
+            4.864653587341309,
+            5.3089189529418945
+          ],
+          [
+            3.6315865516662598,
+            12.289117813110352,
+            11.151065826416016,
+            5.154054641723633,
+            5.65879487991333
+          ],
+          [
+            3.8390471935272217,
+            13.056876182556152,
+            11.834860801696777,
+            5.435141086578369,
+            6.004306316375732
+          ],
+          [
+            4.042593002319336,
+            13.817543029785156,
+            12.511225700378418,
+            5.708395004272461,
+            6.345395565032959
+          ],
+          [
+            4.24238395690918,
+            14.571209907531738,
+            13.179614067077637,
+            5.973555564880371,
+            6.682039737701416
+          ],
+          [
+            4.438711166381836,
+            15.317753791809082,
+            13.840078353881836,
+            6.230597496032715,
+            7.01429557800293
+          ]
+        ],
+        "q01": [
+          [
+            -3.472527027130127,
+            -9.978021621704102,
+            -9.978019714355469,
+            -5.098899841308594,
+            -5.36263370513916
+          ],
+          [
+            -4.483516693115234,
+            -12.175825119018555,
+            -12.263736724853516,
+            -6.417581558227539,
+            -6.857143402099609
+          ],
+          [
+            -5.450550079345703,
+            -14.373626708984375,
+            -14.373626708984375,
+            -7.736263275146484,
+            -8.263736724853516
+          ],
+          [
+            -6.4175825119018555,
+            -16.571428298950195,
+            -16.483516693115234,
+            -9.054946899414062,
+            -9.846153259277344
+          ],
+          [
+            -7.296703338623047,
+            -18.85714340209961,
+            -18.681316375732422,
+            -10.285713195800781,
+            -11.296703338623047
+          ],
+          [
+            -8.175825119018555,
+            -21.054946899414062,
+            -20.79120635986328,
+            -11.516483306884766,
+            -12.747251510620117
+          ],
+          [
+            -9.05494499206543,
+            -23.340660095214844,
+            -22.813186645507812,
+            -12.703296661376953,
+            -14.19780158996582
+          ],
+          [
+            -10.021978378295898,
+            -25.582416534423828,
+            -24.791210174560547,
+            -13.80219841003418,
+            -15.560441017150879
+          ],
+          [
+            -10.813186645507812,
+            -27.648353576660156,
+            -26.769229888916016,
+            -14.813186645507812,
+            -16.96703338623047
+          ],
+          [
+            -11.780220031738281,
+            -29.934066772460938,
+            -28.659339904785156,
+            -15.824176788330078,
+            -18.373626708984375
+          ],
+          [
+            -12.615385055541992,
+            -32.21977996826172,
+            -30.593406677246094,
+            -16.703292846679688,
+            -19.516483306884766
+          ],
+          [
+            -13.450550079345703,
+            -34.32966995239258,
+            -32.395599365234375,
+            -17.670326232910156,
+            -20.791208267211914
+          ],
+          [
+            -14.329669952392578,
+            -36.52747344970703,
+            -34.19780349731445,
+            -18.549448013305664,
+            -21.802196502685547
+          ],
+          [
+            -15.164834976196289,
+            -38.68132019042969,
+            -35.82417297363281,
+            -19.20879364013672,
+            -22.901100158691406
+          ],
+          [
+            -15.95604419708252,
+            -40.92308044433594,
+            -37.626373291015625,
+            -20.043954849243164,
+            -24.131868362426758
+          ],
+          [
+            -16.703296661376953,
+            -42.989009857177734,
+            -39.25274658203125,
+            -20.74725341796875,
+            -25.450550079345703
+          ]
+        ],
+        "q99": [
+          [
+            3.4285709857940674,
+            7.5604400634765625,
+            7.076923847198486,
+            4.131866455078125,
+            5.494504928588867
+          ],
+          [
+            4.307692527770996,
+            9.802200317382812,
+            9.538460731506348,
+            5.274726867675781,
+            7.032966613769531
+          ],
+          [
+            5.1868133544921875,
+            12.087913513183594,
+            11.868131637573242,
+            6.5054931640625,
+            8.527473449707031
+          ],
+          [
+            6.065934181213379,
+            14.197805404663086,
+            14.285714149475098,
+            7.648347854614258,
+            9.934066772460938
+          ],
+          [
+            6.857143402099609,
+            16.483516693115234,
+            16.571428298950195,
+            8.703292846679688,
+            11.340660095214844
+          ],
+          [
+            7.736264705657959,
+            18.593406677246094,
+            18.85714340209961,
+            9.846153259277344,
+            12.835165023803711
+          ],
+          [
+            8.615385055541992,
+            20.52747344970703,
+            21.14285659790039,
+            10.813188552856445,
+            14.241758346557617
+          ],
+          [
+            9.494504928588867,
+            22.725276947021484,
+            23.428571701049805,
+            11.824176788330078,
+            15.648351669311523
+          ],
+          [
+            10.285714149475098,
+            24.879121780395508,
+            25.714282989501953,
+            12.835166931152344,
+            17.05494499206543
+          ],
+          [
+            11.076923370361328,
+            26.85714340209961,
+            27.868133544921875,
+            13.758243560791016,
+            18.461538314819336
+          ],
+          [
+            11.7362642288208,
+            28.70329475402832,
+            30.109888076782227,
+            14.725275039672852,
+            19.868133544921875
+          ],
+          [
+            12.395605087280273,
+            30.813186645507812,
+            32.17582321166992,
+            15.560440063476562,
+            21.186813354492188
+          ],
+          [
+            13.098901748657227,
+            32.52747344970703,
+            34.19779968261719,
+            16.351646423339844,
+            22.549449920654297
+          ],
+          [
+            13.714285850524902,
+            34.54945373535156,
+            36.21977996826172,
+            17.142860412597656,
+            24.0
+          ],
+          [
+            14.329669952392578,
+            36.57143020629883,
+            38.153846740722656,
+            17.846155166625977,
+            25.230770111083984
+          ],
+          [
+            14.769231796264648,
+            38.28571319580078,
+            40.17582702636719,
+            18.637367248535156,
+            26.41758155822754
+          ]
+        ]
+      }
+    }
+  }
+}

checkpoint-5000/experiment_cfg/final_model_config.json

@@ -0,0 +1,54 @@
+{
+  "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,
+  "extra_augmentation_config": null,
+  "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-5000/model-00001-of-00002.safetensors

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

checkpoint-5000/model-00002-of-00002.safetensors

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

checkpoint-5000/optimizer.pt

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

checkpoint-5000/processor_config.json

@@ -0,0 +1,455 @@
+{
+  "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
+        }
+      },
+      "new_embodiment": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "front",
+            "wrist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "single_arm",
+            "gripper"
+          ],
+          "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": [
+            "single_arm",
+            "gripper"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "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": [
+            "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-5000/rng_state.pth

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

checkpoint-5000/scheduler.pt

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