checkpoint-5000 (16D, hybrid crop, batch128)

16d_30k/checkpoint-5000/config.json

@@ -0,0 +1,88 @@
+{
+  "action_horizon": 40,
+  "add_pos_embed": true,
+  "apply_sincos_state_encoding": false,
+  "architectures": [
+    "Gr00tN1d7"
+  ],
+  "attn_dropout": 0.2,
+  "attn_implementation": null,
+  "backbone_embedding_dim": 2048,
+  "backbone_trainable_params_fp32": true,
+  "color_jitter_params": {
+    "brightness": 0.3,
+    "contrast": 0.4,
+    "hue": 0.08,
+    "saturation": 0.5
+  },
+  "crop_fraction": 1.0,
+  "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
+  },
+  "dtype": "bfloat16",
+  "exclude_state": false,
+  "formalize_language": true,
+  "hidden_size": 1024,
+  "image_crop_size": [
+    230,
+    230
+  ],
+  "image_target_size": [
+    256,
+    256
+  ],
+  "letter_box_transform": false,
+  "load_bf16": false,
+  "max_action_dim": 132,
+  "max_num_embodiments": 32,
+  "max_seq_len": 1024,
+  "max_state_dim": 132,
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Cosmos-Reason2-2B",
+  "model_type": "Gr00tN1d7",
+  "noise_beta_alpha": 1.5,
+  "noise_beta_beta": 1.0,
+  "noise_s": 0.999,
+  "num_inference_timesteps": 4,
+  "num_timestep_buckets": 1000,
+  "random_history_crop": true,
+  "random_rotation_angle": 0,
+  "reproject_vision": false,
+  "rtc_ramp_rate": 6.0,
+  "select_layer": 16,
+  "shortest_image_edge": 256,
+  "state_dropout_prob": 0.2,
+  "state_gaussian_noise_std": 0.0,
+  "transformers_version": "4.57.3",
+  "tune_diffusion_model": true,
+  "tune_linear": true,
+  "tune_llm": false,
+  "tune_projector": true,
+  "tune_top_llm_layers": 0,
+  "tune_visual": false,
+  "tune_vlln": true,
+  "use_albumentations": true,
+  "use_alternate_vl_dit": true,
+  "use_flash_attention": true,
+  "use_future_tokens": false,
+  "use_mean_std": false,
+  "use_percentiles": false,
+  "use_vl_self_attention": true,
+  "use_vlln": true,
+  "vl_self_attention_cfg": {
+    "attention_head_dim": 64,
+    "dropout": 0.2,
+    "final_dropout": true,
+    "num_attention_heads": 32,
+    "num_layers": 4,
+    "positional_embeddings": null
+  }
+}

16d_30k/checkpoint-5000/embodiment_id.json

@@ -0,0 +1,57 @@
+{
+  "robocasa_panda_omron": 13,
+  "oxe_droid": 17,
+  "oxe_fractal": 18,
+  "oxe_language_table": 19,
+  "oxe_bridge": 20,
+  "unknown": 22,
+  "gr1_unified": 20,
+  "agibot": 26,
+  "sim_behavior_r1_pro": 23,
+  "xdof": 24,
+  "xdof_oss_data": 25,
+  "unitree_g1_full_body_with_waist_height_nav_cmd": 25,
+  "real_r1_pro_sharpa": 27,
+  "real_r1_pro_sharpa_add_view": 27,
+  "real_r1_pro_sharpa_relative_arm_joint": 26,
+  "real_r1_pro_sharpa_delta_eef": 26,
+  "real_r1_pro_sharpa_absolute_eef": 26,
+  "real_r1_pro_sharpa_meanstd": 26,
+  "real_r1_pro_sharpa_relative_eef": 26,
+  "real_r1_pro_sharpa_relative_eef_add_view": 26,
+  "real_r1_pro_sharpa_relative_eef_relative_hand": 26,
+  "real_r1_pro_sharpa_relative_eef_human": 26,
+  "real_r1_pro_sharpa_relative_eef_human_add_view": 26,
+  "real_r1_pro_sharpa_relative_eef_human_relative_hand": 26,
+  "real_r1_pro_sharpa_relative_eef_egodex": 26,
+  "real_r1_pro_sharpa_relative_eef_egodex_relative_hand": 26,
+  "real_r1_pro_sharpa_relative_eef_egodex_wrist_only": 26,
+  "real_r1_pro_sharpa_relative_eef_maxinsights": 26,
+  "real_r1_pro_sharpa_relative_eef_maxinsights_relative_hand": 26,
+  "real_r1_pro_sharpa_relative_eef_mecka": 26,
+  "real_r1_pro_sharpa_relative_eef_mecka_relative_hand": 26,
+  "real_g1_relative_eef_absolute_joints": 25,
+  "real_g1_relative_eef_absolute_joints_wrist_cam": 25,
+  "real_g1_relative_eef_relative_joints": 25,
+  "real_r1_pro_sharpa_relative_eef_relative_hand_relative_joint": 26,
+  "real_r1_pro_sharpa_relative_joint": 29,
+  "oxe_droid_relative_eef_relative_joint": 24,
+  "oxe_droid_relative_eef_relative_joint_swapped": 24,
+  "oxe_droid_relative_eef_relative_joint_upweight_z": 24,
+  "oxe_droid_relative_eef_relative_joint_upweight_z_swapped": 24,
+  "oxe_droid_relative_eef_relative_joint_3view": 24,
+  "oxe_droid_relative_eef_relative_joint_3view_swapped": 24,
+  "oxe_droid_relative_eef": 24,
+  "oxe_droid_joint_position_relative": 24,
+  "xdof_relative_eef_relative_joint": 27,
+  "xdof_relative_eef_relative_joint_subtask": 27,
+  "xdof_relative_eef": 27,
+  "xdof_relative_joint": 28,
+  "simpler_env_google": 0,
+  "simpler_env_widowx": 1,
+  "libero_sim": 2,
+  "droid_sim": 3,
+  "unitree_g1_sonic": 11,
+  "new_embodiment": 10,
+  "robocasa_gr1_tabletop": 10
+}

16d_30k/checkpoint-5000/experiment_cfg/conf.yaml

@@ -0,0 +1,214 @@
+load_config_path: null
+model:
+  model_type: Gr00tN1d7
+  model_dtype: bfloat16
+  model_name: nvidia/Cosmos-Reason2-2B
+  backbone_model_type: qwen
+  model_revision: null
+  tune_top_llm_layers: 0
+  backbone_embedding_dim: 2048
+  tune_llm: false
+  tune_visual: false
+  select_layer: 12
+  reproject_vision: false
+  use_flash_attention: true
+  load_bf16: false
+  backbone_trainable_params_fp32: true
+  image_crop_size:
+  - 230
+  - 230
+  image_target_size:
+  - 256
+  - 256
+  shortest_image_edge: null
+  crop_fraction: null
+  random_rotation_angle: null
+  color_jitter_params: null
+  use_albumentations_transforms: true
+  extra_augmentation_config: null
+  formalize_language: true
+  apply_sincos_state_encoding: false
+  use_percentiles: true
+  use_relative_action: true
+  max_state_dim: 132
+  max_action_dim: 132
+  action_horizon: 40
+  hidden_size: 1024
+  input_embedding_dim: 1536
+  state_history_length: 1
+  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: 16
+    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.2
+  exclude_state: false
+  use_mean_std: false
+  max_num_embodiments: 32
+data:
+  datasets:
+  - dataset_paths:
+    - /root/screwing_train/datasets/screwing35_follower_train
+    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:
+        - cam_left_head
+        - cam_left_wrist
+        - cam_right_wrist
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      state:
+        delta_indices:
+        - 0
+        modality_keys:
+        - arm_left
+        - arm_right
+        sin_cos_embedding_keys: null
+        mean_std_embedding_keys: null
+        action_configs: null
+      action:
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        modality_keys:
+        - arm_left
+        - arm_right
+        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
+      language:
+        delta_indices:
+        - 0
+        modality_keys:
+        - annotation.human.primitive_instruction
+        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: true
+  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: /root/screwing_train/runs/groot_16d_30k
+  experiment_name: null
+  max_steps: 30000
+  global_batch_size: 128
+  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: /root/screwing_train/base_model_patched
+  skip_weight_loading: false
+  tf32: true
+  fp16: false
+  bf16: true
+  eval_bf16: true
+  logging_steps: 10
+  save_steps: 2500
+  save_total_limit: 30
+  save_vl_model: false
+  save_only_model: true
+  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: 2
+  dataloader_num_workers: 24
+  remove_unused_columns: false
+  use_wandb: true
+  wandb_project: screwing_hil
+  enable_profiling: false
+  max_retries: 3
+  assert_loss_less_than: null
+  add_rl_callback: false
+  enable_open_loop_eval: false
+  open_loop_eval_traj_ids:
+  - 0
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_plot_indices: null
+max_steps: 30000
+save_steps: 2500

16d_30k/checkpoint-5000/experiment_cfg/config.yaml

@@ -0,0 +1,239 @@
+!!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:
+    - /root/screwing_train/datasets/screwing35_follower_train
+    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: &id002 !!python/object/apply:gr00t.data.types.ActionRepresentation
+          - absolute
+          state_key: null
+          type: &id003 !!python/object/apply:gr00t.data.types.ActionType
+          - non_eef
+        - !!python/object:gr00t.data.types.ActionConfig
+          format: *id001
+          rep: *id002
+          state_key: null
+          type: *id003
+        delta_indices:
+        - 0
+        - 1
+        - 2
+        - 3
+        - 4
+        - 5
+        - 6
+        - 7
+        - 8
+        - 9
+        - 10
+        - 11
+        - 12
+        - 13
+        - 14
+        - 15
+        mean_std_embedding_keys: null
+        modality_keys:
+        - arm_left
+        - arm_right
+        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.primitive_instruction
+        sin_cos_embedding_keys: null
+      state: !!python/object:gr00t.data.types.ModalityConfig
+        action_configs: null
+        delta_indices:
+        - 0
+        mean_std_embedding_keys: null
+        modality_keys:
+        - arm_left
+        - arm_right
+        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:
+        - cam_left_head
+        - cam_left_wrist
+        - cam_right_wrist
+        sin_cos_embedding_keys: null
+  mode: single_turn
+  multiprocessing_context: fork
+  num_shards_per_epoch: 100000
+  override_pretraining_statistics: true
+  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_n1d7.Gr00tN1d7Config
+  _attn_implementation_internal: null
+  _commit_hash: null
+  _name_or_path: ''
+  _output_attentions: false
+  add_cross_attention: false
+  architectures: null
+  backbone_trainable_params_fp32: true
+  bad_words_ids: null
+  begin_suppress_tokens: null
+  bos_token_id: null
+  chunk_size_feed_forward: 0
+  color_jitter_params: null
+  cross_attention_hidden_size: null
+  decoder_start_token_id: null
+  diffusion_model_cfg:
+    attention_head_dim: 48
+    dropout: 0.2
+    final_dropout: true
+    interleave_self_attention: true
+    norm_type: ada_norm
+    num_attention_heads: 32
+    num_layers: 16
+    output_dim: 1024
+    positional_embeddings: null
+  diversity_penalty: 0.0
+  do_sample: false
+  dtype: null
+  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/Cosmos-Reason2-2B
+  no_repeat_ngram_size: 0
+  num_beam_groups: 1
+  num_beams: 1
+  num_return_sequences: 1
+  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.2
+  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
+  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: 24
+  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: 128
+  gradient_accumulation_steps: 1
+  gradient_checkpointing: false
+  learning_rate: 0.0001
+  logging_steps: 10
+  lr_scheduler_type: cosine
+  max_concurrent_uploads: 2
+  max_grad_norm: 1.0
+  max_retries: 3
+  max_steps: 30000
+  num_gpus: 2
+  open_loop_eval_plot_indices: null
+  open_loop_eval_steps_per_traj: 100
+  open_loop_eval_traj_ids:
+  - 0
+  optim: adamw_torch
+  output_dir: /root/screwing_train/runs/groot_16d_30k
+  remove_unused_columns: false
+  save_best_eval_metric_greater_is_better: true
+  save_best_eval_metric_name: ''
+  save_only_model: true
+  save_steps: 2500
+  save_total_limit: 30
+  save_vl_model: false
+  skip_weight_loading: false
+  start_from_checkpoint: /root/screwing_train/base_model_patched
+  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: screwing_hil
+  warmup_ratio: 0.05
+  warmup_steps: 0
+  weight_decay: 1.0e-05

16d_30k/checkpoint-5000/experiment_cfg/dataset_statistics.json

@@ -0,0 +1,257 @@
+{
+  "new_embodiment": {
+    "state": {
+      "arm_left": {
+        "min": [
+          -0.3512935936450958,
+          -0.0003355582884978503,
+          -0.4693501889705658,
+          -2.0985336303710938,
+          -0.30987611413002014,
+          -0.390062540769577,
+          -0.756526529788971,
+          -0.0015339808305725455
+        ],
+        "max": [
+          0.4587441384792328,
+          0.4674566686153412,
+          0.2403915673494339,
+          -0.8097261190414429,
+          0.30678418278694153,
+          0.7219297289848328,
+          0.09214304387569427,
+          1.0001554489135742
+        ],
+        "mean": [
+          0.045786481350660324,
+          0.1028616651892662,
+          -0.08757976442575455,
+          -1.7214374542236328,
+          0.016063803806900978,
+          0.22246477007865906,
+          -0.3082488775253296,
+          0.23499731719493866
+        ],
+        "std": [
+          0.20334461331367493,
+          0.0962500348687172,
+          0.16793721914291382,
+          0.24059325456619263,
+          0.12082602083683014,
+          0.2635464370250702,
+          0.22770243883132935,
+          0.38598382472991943
+        ],
+        "q01": [
+          -0.28486502170562744,
+          0.0,
+          -0.4279566705226898,
+          -2.013697385787964,
+          -0.2523278594017029,
+          -0.2879210114479065,
+          -0.7209579944610596,
+          0.0
+        ],
+        "q99": [
+          0.4115143120288849,
+          0.4564431607723236,
+          0.18591128289699554,
+          -0.9433622360229492,
+          0.2618792951107025,
+          0.7118869423866272,
+          0.06595534831285477,
+          0.9970875382423401
+        ]
+      },
+      "arm_right": {
+        "min": [
+          -0.9125747680664062,
+          -0.46813976764678955,
+          -0.08915065228939056,
+          -2.0727076530456543,
+          -0.9404261112213135,
+          -0.15235544741153717,
+          -0.053947482258081436,
+          0.0
+        ],
+        "max": [
+          0.3974807858467102,
+          -0.031806133687496185,
+          1.0656253099441528,
+          -1.032512903213501,
+          0.13881327211856842,
+          1.1778095960617065,
+          0.8715579509735107,
+          1.277806043624878
+        ],
+        "mean": [
+          -0.25570186972618103,
+          -0.2676304280757904,
+          0.3035861551761627,
+          -1.7064963579177856,
+          -0.25106674432754517,
+          0.5091767311096191,
+          0.43237197399139404,
+          0.3860172629356384
+        ],
+        "std": [
+          0.25905025005340576,
+          0.0594896525144577,
+          0.17308898270130157,
+          0.18367764353752136,
+          0.19315017759799957,
+          0.2880626618862152,
+          0.1807917058467865,
+          0.33443471789360046
+        ],
+        "q01": [
+          -0.8353484272956848,
+          -0.43082091212272644,
+          -0.01853959634900093,
+          -2.0081846714019775,
+          -0.7463775277137756,
+          -0.06663229316473007,
+          0.07988695055246353,
+          0.0015339808305725455
+        ],
+        "q99": [
+          0.28748956322669983,
+          -0.14874820411205292,
+          0.8437014222145081,
+          -1.2486603260040283,
+          0.12267052382230759,
+          1.0415489673614502,
+          0.8113635182380676,
+          0.8682331442832947
+        ]
+      }
+    },
+    "action": {
+      "arm_left": {
+        "min": [
+          -0.3512935936450958,
+          -0.0003355582884978503,
+          -0.4693501889705658,
+          -2.0985336303710938,
+          -0.30987611413002014,
+          -0.390062540769577,
+          -0.756526529788971,
+          -0.0015339808305725455
+        ],
+        "max": [
+          0.4587441384792328,
+          0.4674566686153412,
+          0.2403915673494339,
+          -0.8097261190414429,
+          0.30678418278694153,
+          0.7219297289848328,
+          0.09214304387569427,
+          1.0001554489135742
+        ],
+        "mean": [
+          0.04590541869401932,
+          0.10304273664951324,
+          -0.08747139573097229,
+          -1.721885085105896,
+          0.016054483130574226,
+          0.22288048267364502,
+          -0.3089201748371124,
+          0.23478840291500092
+        ],
+        "std": [
+          0.20350636541843414,
+          0.09618024528026581,
+          0.16804173588752747,
+          0.24068261682987166,
+          0.1208517923951149,
+          0.26317518949508667,
+          0.22766965627670288,
+          0.3860923647880554
+        ],
+        "q01": [
+          -0.28486502170562744,
+          0.0,
+          -0.4279566705226898,
+          -2.013697385787964,
+          -0.2523278594017029,
+          -0.2877771854400635,
+          -0.7209579944610596,
+          0.0
+        ],
+        "q99": [
+          0.41167011857032776,
+          0.4564431607723236,
+          0.1860191375017166,
+          -0.9433622360229492,
+          0.2618792951107025,
+          0.7118869423866272,
+          0.06595534831285477,
+          0.9970875382423401
+        ]
+      },
+      "arm_right": {
+        "min": [
+          -0.9125747680664062,
+          -0.46813976764678955,
+          -0.08915065228939056,
+          -2.0727076530456543,
+          -0.9404261112213135,
+          -0.15235544741153717,
+          -0.053947482258081436,
+          0.0
+        ],
+        "max": [
+          0.3974807858467102,
+          -0.031806133687496185,
+          1.0656253099441528,
+          -1.032512903213501,
+          0.13881327211856842,
+          1.1778095960617065,
+          0.8715579509735107,
+          1.277806043624878
+        ],
+        "mean": [
+          -0.255352258682251,
+          -0.2675997316837311,
+          0.3033587634563446,
+          -1.7069568634033203,
+          -0.2512829005718231,
+          0.5092591047286987,
+          0.43265601992607117,
+          0.3857439160346985
+        ],
+        "std": [
+          0.25960078835487366,
+          0.059564266353845596,
+          0.17328619956970215,
+          0.1838947683572763,
+          0.19315995275974274,
+          0.2880189120769501,
+          0.18076331913471222,
+          0.3346700072288513
+        ],
+        "q01": [
+          -0.8353484272956848,
+          -0.43082091212272644,
+          -0.01853959634900093,
+          -2.0086400508880615,
+          -0.7463775277137756,
+          -0.06663229316473007,
+          0.07988695055246353,
+          0.0015339808305725455
+        ],
+        "q99": [
+          0.288448303937912,
+          -0.14868827164173126,
+          0.8437014222145081,
+          -1.2486603260040283,
+          0.12267052382230759,
+          1.0415489673614502,
+          0.8113635182380676,
+          0.8682331442832947
+        ]
+      }
+    },
+    "relative_action": {}
+  }
+}

16d_30k/checkpoint-5000/experiment_cfg/final_model_config.json

@@ -0,0 +1,55 @@
+{
+  "model_type": "Gr00tN1d7",
+  "model_dtype": "bfloat16",
+  "model_name": "nvidia/Cosmos-Reason2-2B",
+  "backbone_model_type": "qwen",
+  "model_revision": null,
+  "tune_top_llm_layers": 0,
+  "backbone_embedding_dim": 2048,
+  "tune_llm": false,
+  "tune_visual": false,
+  "select_layer": 16,
+  "reproject_vision": false,
+  "use_flash_attention": true,
+  "load_bf16": false,
+  "backbone_trainable_params_fp32": true,
+  "extra_augmentation_config": null,
+  "apply_sincos_state_encoding": false,
+  "use_percentiles": false,
+  "use_relative_action": false,
+  "max_state_dim": 132,
+  "max_action_dim": 132,
+  "action_horizon": 40,
+  "hidden_size": 1024,
+  "input_embedding_dim": 1536,
+  "state_history_length": 1,
+  "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.2,
+  "exclude_state": false,
+  "use_mean_std": false,
+  "max_num_embodiments": 32
+}

16d_30k/checkpoint-5000/model-00001-of-00002.safetensors

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

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

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

16d_30k/checkpoint-5000/processor_config.json

@@ -0,0 +1,1124 @@
+{
+  "processor_class": "Gr00tN1d7Processor",
+  "processor_kwargs": {
+    "modality_configs": {
+      "real_g1_relative_eef_relative_joints": {
+        "video": {
+          "delta_indices": [
+            -20,
+            0
+          ],
+          "modality_keys": [
+            "ego_view"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef_9d",
+            "right_wrist_eef_9d",
+            "left_hand",
+            "right_hand",
+            "left_arm",
+            "right_arm",
+            "waist"
+          ],
+          "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,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef_9d",
+            "right_wrist_eef_9d",
+            "left_hand",
+            "right_hand",
+            "left_arm",
+            "right_arm",
+            "waist",
+            "base_height_command",
+            "navigate_command"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef_9d"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef_9d"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_hand"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_hand"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_arm"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_arm"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "waist"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "base_height_command"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "navigate_command"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.task_description"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "real_r1_pro_sharpa_relative_eef_mecka": {
+        "video": {
+          "delta_indices": [
+            -30,
+            0
+          ],
+          "modality_keys": [
+            "ego_view_cropratio_res320x240_freq30"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_hand_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_hand_joints"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "real_r1_pro_sharpa_relative_eef_human": {
+        "video": {
+          "delta_indices": [
+            -20,
+            0
+          ],
+          "modality_keys": [
+            "ego_view_res320x240_freq20",
+            "left_wrist_view_res320x240_freq20",
+            "right_wrist_view_res320x240_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_hand_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_hand_joints"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "real_r1_pro_sharpa_relative_eef": {
+        "video": {
+          "delta_indices": [
+            -20,
+            0
+          ],
+          "modality_keys": [
+            "ego_view_res320x240_freq20",
+            "left_wrist_view_res320x240_freq20",
+            "right_wrist_view_res320x240_freq20"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_hand_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_hand_joints"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "xdof_relative_eef_relative_joint": {
+        "video": {
+          "delta_indices": [
+            -30,
+            0
+          ],
+          "modality_keys": [
+            "top_camera-images-rgb_320_240",
+            "left_camera-images-rgb_320_240",
+            "right_camera-images-rgb_320_240"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_gripper_pos",
+            "right_gripper_pos",
+            "left_joint_pos",
+            "right_joint_pos"
+          ],
+          "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,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_gripper_pos",
+            "right_gripper_pos",
+            "left_joint_pos",
+            "right_joint_pos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_gripper_pos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_gripper_pos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_joint_pos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_joint_pos"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "real_r1_pro_sharpa_relative_eef_maxinsights": {
+        "video": {
+          "delta_indices": [
+            -30,
+            0
+          ],
+          "modality_keys": [
+            "ego_view_cropratio_res320x240_freq30"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18,
+            19,
+            20,
+            21,
+            22,
+            23,
+            24,
+            25,
+            26,
+            27,
+            28,
+            29,
+            30,
+            31,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_hand_joints",
+            "right_hand_joints"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_hand_joints"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_hand_joints"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.coarse_action"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "xdof_relative_eef_relative_joint_subtask": {
+        "video": {
+          "delta_indices": [
+            -30,
+            0
+          ],
+          "modality_keys": [
+            "top_camera-images-rgb_320_240",
+            "left_camera-images-rgb_320_240",
+            "right_camera-images-rgb_320_240"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_gripper_pos",
+            "right_gripper_pos",
+            "left_joint_pos",
+            "right_joint_pos"
+          ],
+          "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,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "left_wrist_eef",
+            "right_wrist_eef",
+            "left_gripper_pos",
+            "right_gripper_pos",
+            "left_joint_pos",
+            "right_joint_pos"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "left_wrist_eef"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "right_wrist_eef"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_gripper_pos"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_gripper_pos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "left_joint_pos"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "right_joint_pos"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.sub_task"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "oxe_droid_relative_eef_relative_joint": {
+        "video": {
+          "delta_indices": [
+            -15,
+            0
+          ],
+          "modality_keys": [
+            "exterior_image_1_left",
+            "wrist_image_left"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "eef_9d",
+            "gripper_position",
+            "joint_position"
+          ],
+          "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,
+            32,
+            33,
+            34,
+            35,
+            36,
+            37,
+            38,
+            39
+          ],
+          "modality_keys": [
+            "eef_9d",
+            "gripper_position",
+            "joint_position"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": [
+            {
+              "rep": "RELATIVE",
+              "type": "EEF",
+              "format": "XYZ_ROT6D",
+              "state_key": "eef_9d"
+            },
+            {
+              "rep": "ABSOLUTE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "gripper_position"
+            },
+            {
+              "rep": "RELATIVE",
+              "type": "NON_EEF",
+              "format": "DEFAULT",
+              "state_key": "joint_position"
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.language.language_instruction"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      },
+      "new_embodiment": {
+        "video": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "cam_left_head",
+            "cam_left_wrist",
+            "cam_right_wrist"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "state": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "arm_left",
+            "arm_right"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        },
+        "action": {
+          "delta_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15
+          ],
+          "modality_keys": [
+            "arm_left",
+            "arm_right"
+          ],
+          "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
+            }
+          ]
+        },
+        "language": {
+          "delta_indices": [
+            0
+          ],
+          "modality_keys": [
+            "annotation.human.primitive_instruction"
+          ],
+          "sin_cos_embedding_keys": null,
+          "mean_std_embedding_keys": null,
+          "action_configs": null
+        }
+      }
+    },
+    "image_crop_size": [
+      230,
+      230
+    ],
+    "image_target_size": [
+      256,
+      256
+    ],
+    "use_albumentations": true,
+    "random_rotation_angle": 0,
+    "color_jitter_params": {
+      "brightness": 0.3,
+      "contrast": 0.4,
+      "saturation": 0.5,
+      "hue": 0.08
+    },
+    "shortest_image_edge": 256,
+    "crop_fraction": 1.0,
+    "letter_box_transform": false,
+    "model_name": "nvidia/Cosmos-Reason2-2B",
+    "model_type": "qwen",
+    "formalize_language": true,
+    "max_state_dim": 132,
+    "max_action_dim": 132,
+    "max_action_horizon": 40,
+    "use_percentiles": false,
+    "use_mean_std": false,
+    "clip_outliers": true,
+    "apply_sincos_state_encoding": false,
+    "use_relative_action": true,
+    "exclude_state": false,
+    "state_dropout_prob": 0.2
+  }
+}

16d_30k/checkpoint-5000/trainer_state.json

@@ -0,0 +1,3034 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.16666666666666666,
+  "eval_steps": 500,
+  "global_step": 5000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "grad_norm": 0.6208589673042297,
+      "learning_rate": 6.000000000000001e-07,
+      "loss": 1.2746,
+      "step": 10
+    },
+    {
+      "grad_norm": 0.6478228569030762,
+      "learning_rate": 1.2666666666666667e-06,
+      "loss": 1.282,
+      "step": 20
+    },
+    {
+      "grad_norm": 0.5199011564254761,
+      "learning_rate": 1.9333333333333336e-06,
+      "loss": 1.2766,
+      "step": 30
+    },
+    {
+      "grad_norm": 0.6404500603675842,
+      "learning_rate": 2.6e-06,
+      "loss": 1.2254,
+      "step": 40
+    },
+    {
+      "grad_norm": 0.5328770279884338,
+      "learning_rate": 3.2666666666666666e-06,
+      "loss": 1.1938,
+      "step": 50
+    },
+    {
+      "grad_norm": 0.5413777232170105,
+      "learning_rate": 3.9333333333333335e-06,
+      "loss": 1.1359,
+      "step": 60
+    },
+    {
+      "grad_norm": 0.49827340245246887,
+      "learning_rate": 4.6e-06,
+      "loss": 1.0902,
+      "step": 70
+    },
+    {
+      "grad_norm": 0.46059373021125793,
+      "learning_rate": 5.266666666666667e-06,
+      "loss": 1.0633,
+      "step": 80
+    },
+    {
+      "grad_norm": 0.4852186441421509,
+      "learning_rate": 5.933333333333334e-06,
+      "loss": 1.048,
+      "step": 90
+    },
+    {
+      "grad_norm": 0.644527018070221,
+      "learning_rate": 6.6e-06,
+      "loss": 1.0359,
+      "step": 100
+    },
+    {
+      "grad_norm": 0.6214720606803894,
+      "learning_rate": 7.266666666666668e-06,
+      "loss": 1.0328,
+      "step": 110
+    },
+    {
+      "grad_norm": 0.38992390036582947,
+      "learning_rate": 7.933333333333334e-06,
+      "loss": 1.0285,
+      "step": 120
+    },
+    {
+      "grad_norm": 0.5468058586120605,
+      "learning_rate": 8.599999999999999e-06,
+      "loss": 1.0145,
+      "step": 130
+    },
+    {
+      "grad_norm": 0.743439257144928,
+      "learning_rate": 9.266666666666667e-06,
+      "loss": 1.0,
+      "step": 140
+    },
+    {
+      "grad_norm": 0.9581028819084167,
+      "learning_rate": 9.933333333333334e-06,
+      "loss": 0.9418,
+      "step": 150
+    },
+    {
+      "grad_norm": 1.0746005773544312,
+      "learning_rate": 1.06e-05,
+      "loss": 0.8699,
+      "step": 160
+    },
+    {
+      "grad_norm": 1.3099654912948608,
+      "learning_rate": 1.1266666666666667e-05,
+      "loss": 0.7877,
+      "step": 170
+    },
+    {
+      "grad_norm": 1.502665400505066,
+      "learning_rate": 1.1933333333333333e-05,
+      "loss": 0.7059,
+      "step": 180
+    },
+    {
+      "grad_norm": 1.873738408088684,
+      "learning_rate": 1.2600000000000001e-05,
+      "loss": 0.6436,
+      "step": 190
+    },
+    {
+      "grad_norm": 2.2341554164886475,
+      "learning_rate": 1.3266666666666666e-05,
+      "loss": 0.5734,
+      "step": 200
+    },
+    {
+      "grad_norm": 1.5233055353164673,
+      "learning_rate": 1.3933333333333334e-05,
+      "loss": 0.5299,
+      "step": 210
+    },
+    {
+      "grad_norm": 1.5989943742752075,
+      "learning_rate": 1.4599999999999999e-05,
+      "loss": 0.482,
+      "step": 220
+    },
+    {
+      "grad_norm": 1.6445683240890503,
+      "learning_rate": 1.5266666666666667e-05,
+      "loss": 0.4228,
+      "step": 230
+    },
+    {
+      "grad_norm": 1.5990145206451416,
+      "learning_rate": 1.5933333333333332e-05,
+      "loss": 0.3652,
+      "step": 240
+    },
+    {
+      "grad_norm": 1.929441213607788,
+      "learning_rate": 1.66e-05,
+      "loss": 0.3081,
+      "step": 250
+    },
+    {
+      "grad_norm": 2.241440773010254,
+      "learning_rate": 1.726666666666667e-05,
+      "loss": 0.2768,
+      "step": 260
+    },
+    {
+      "grad_norm": 1.6238600015640259,
+      "learning_rate": 1.7933333333333337e-05,
+      "loss": 0.2393,
+      "step": 270
+    },
+    {
+      "grad_norm": 1.9927496910095215,
+      "learning_rate": 1.86e-05,
+      "loss": 0.1933,
+      "step": 280
+    },
+    {
+      "grad_norm": 2.0286266803741455,
+      "learning_rate": 1.926666666666667e-05,
+      "loss": 0.1622,
+      "step": 290
+    },
+    {
+      "grad_norm": 1.611731767654419,
+      "learning_rate": 1.9933333333333334e-05,
+      "loss": 0.1558,
+      "step": 300
+    },
+    {
+      "grad_norm": 1.767244815826416,
+      "learning_rate": 2.06e-05,
+      "loss": 0.1541,
+      "step": 310
+    },
+    {
+      "grad_norm": 1.6201092004776,
+      "learning_rate": 2.1266666666666667e-05,
+      "loss": 0.1285,
+      "step": 320
+    },
+    {
+      "grad_norm": 1.7983145713806152,
+      "learning_rate": 2.1933333333333332e-05,
+      "loss": 0.1288,
+      "step": 330
+    },
+    {
+      "grad_norm": 2.1310415267944336,
+      "learning_rate": 2.26e-05,
+      "loss": 0.1355,
+      "step": 340
+    },
+    {
+      "grad_norm": 1.7930408716201782,
+      "learning_rate": 2.326666666666667e-05,
+      "loss": 0.124,
+      "step": 350
+    },
+    {
+      "grad_norm": 1.5554205179214478,
+      "learning_rate": 2.3933333333333337e-05,
+      "loss": 0.1252,
+      "step": 360
+    },
+    {
+      "grad_norm": 1.9930179119110107,
+      "learning_rate": 2.46e-05,
+      "loss": 0.1045,
+      "step": 370
+    },
+    {
+      "grad_norm": 1.8832621574401855,
+      "learning_rate": 2.5266666666666666e-05,
+      "loss": 0.1069,
+      "step": 380
+    },
+    {
+      "grad_norm": 1.9651679992675781,
+      "learning_rate": 2.5933333333333338e-05,
+      "loss": 0.1077,
+      "step": 390
+    },
+    {
+      "grad_norm": 1.5526353120803833,
+      "learning_rate": 2.6600000000000003e-05,
+      "loss": 0.1046,
+      "step": 400
+    },
+    {
+      "grad_norm": 1.75920569896698,
+      "learning_rate": 2.7266666666666668e-05,
+      "loss": 0.1031,
+      "step": 410
+    },
+    {
+      "grad_norm": 1.6312711238861084,
+      "learning_rate": 2.7933333333333332e-05,
+      "loss": 0.0985,
+      "step": 420
+    },
+    {
+      "grad_norm": 1.9239977598190308,
+      "learning_rate": 2.86e-05,
+      "loss": 0.1029,
+      "step": 430
+    },
+    {
+      "grad_norm": 1.6071902513504028,
+      "learning_rate": 2.926666666666667e-05,
+      "loss": 0.092,
+      "step": 440
+    },
+    {
+      "grad_norm": 1.3962411880493164,
+      "learning_rate": 2.9933333333333337e-05,
+      "loss": 0.0927,
+      "step": 450
+    },
+    {
+      "grad_norm": 1.8652002811431885,
+      "learning_rate": 3.06e-05,
+      "loss": 0.0897,
+      "step": 460
+    },
+    {
+      "grad_norm": 1.4422093629837036,
+      "learning_rate": 3.126666666666666e-05,
+      "loss": 0.0838,
+      "step": 470
+    },
+    {
+      "grad_norm": 1.1935328245162964,
+      "learning_rate": 3.1933333333333335e-05,
+      "loss": 0.0873,
+      "step": 480
+    },
+    {
+      "grad_norm": 1.549795389175415,
+      "learning_rate": 3.26e-05,
+      "loss": 0.0846,
+      "step": 490
+    },
+    {
+      "grad_norm": 1.3405718803405762,
+      "learning_rate": 3.326666666666667e-05,
+      "loss": 0.0866,
+      "step": 500
+    },
+    {
+      "grad_norm": 1.535412311553955,
+      "learning_rate": 3.3933333333333336e-05,
+      "loss": 0.0806,
+      "step": 510
+    },
+    {
+      "grad_norm": 1.7111966609954834,
+      "learning_rate": 3.46e-05,
+      "loss": 0.084,
+      "step": 520
+    },
+    {
+      "grad_norm": 1.424939513206482,
+      "learning_rate": 3.526666666666667e-05,
+      "loss": 0.0822,
+      "step": 530
+    },
+    {
+      "grad_norm": 2.0077829360961914,
+      "learning_rate": 3.593333333333334e-05,
+      "loss": 0.0831,
+      "step": 540
+    },
+    {
+      "grad_norm": 1.394214391708374,
+      "learning_rate": 3.66e-05,
+      "loss": 0.0741,
+      "step": 550
+    },
+    {
+      "grad_norm": 1.670121431350708,
+      "learning_rate": 3.726666666666667e-05,
+      "loss": 0.0842,
+      "step": 560
+    },
+    {
+      "grad_norm": 1.5677956342697144,
+      "learning_rate": 3.793333333333334e-05,
+      "loss": 0.0877,
+      "step": 570
+    },
+    {
+      "grad_norm": 1.7964739799499512,
+      "learning_rate": 3.86e-05,
+      "loss": 0.0815,
+      "step": 580
+    },
+    {
+      "grad_norm": 1.4517905712127686,
+      "learning_rate": 3.926666666666667e-05,
+      "loss": 0.0792,
+      "step": 590
+    },
+    {
+      "grad_norm": 1.4527804851531982,
+      "learning_rate": 3.993333333333333e-05,
+      "loss": 0.0793,
+      "step": 600
+    },
+    {
+      "grad_norm": 1.163316011428833,
+      "learning_rate": 4.0600000000000004e-05,
+      "loss": 0.0745,
+      "step": 610
+    },
+    {
+      "grad_norm": 1.5932477712631226,
+      "learning_rate": 4.126666666666667e-05,
+      "loss": 0.0754,
+      "step": 620
+    },
+    {
+      "grad_norm": 1.437664270401001,
+      "learning_rate": 4.1933333333333334e-05,
+      "loss": 0.0756,
+      "step": 630
+    },
+    {
+      "grad_norm": 1.3479504585266113,
+      "learning_rate": 4.26e-05,
+      "loss": 0.0778,
+      "step": 640
+    },
+    {
+      "grad_norm": 1.1809881925582886,
+      "learning_rate": 4.3266666666666664e-05,
+      "loss": 0.0791,
+      "step": 650
+    },
+    {
+      "grad_norm": 1.592273235321045,
+      "learning_rate": 4.3933333333333335e-05,
+      "loss": 0.0726,
+      "step": 660
+    },
+    {
+      "grad_norm": 1.2641804218292236,
+      "learning_rate": 4.46e-05,
+      "loss": 0.0692,
+      "step": 670
+    },
+    {
+      "grad_norm": 1.5784950256347656,
+      "learning_rate": 4.526666666666667e-05,
+      "loss": 0.0671,
+      "step": 680
+    },
+    {
+      "grad_norm": 1.8079859018325806,
+      "learning_rate": 4.5933333333333336e-05,
+      "loss": 0.0648,
+      "step": 690
+    },
+    {
+      "grad_norm": 1.452563762664795,
+      "learning_rate": 4.660000000000001e-05,
+      "loss": 0.0725,
+      "step": 700
+    },
+    {
+      "grad_norm": 1.6882253885269165,
+      "learning_rate": 4.726666666666667e-05,
+      "loss": 0.0678,
+      "step": 710
+    },
+    {
+      "grad_norm": 1.1656196117401123,
+      "learning_rate": 4.793333333333334e-05,
+      "loss": 0.0683,
+      "step": 720
+    },
+    {
+      "grad_norm": 1.4506734609603882,
+      "learning_rate": 4.86e-05,
+      "loss": 0.071,
+      "step": 730
+    },
+    {
+      "grad_norm": 1.171644687652588,
+      "learning_rate": 4.926666666666667e-05,
+      "loss": 0.0742,
+      "step": 740
+    },
+    {
+      "grad_norm": 1.3164993524551392,
+      "learning_rate": 4.993333333333334e-05,
+      "loss": 0.0717,
+      "step": 750
+    },
+    {
+      "grad_norm": 1.38760244846344,
+      "learning_rate": 5.0600000000000003e-05,
+      "loss": 0.0654,
+      "step": 760
+    },
+    {
+      "grad_norm": 1.2424696683883667,
+      "learning_rate": 5.1266666666666675e-05,
+      "loss": 0.0722,
+      "step": 770
+    },
+    {
+      "grad_norm": 1.150290608406067,
+      "learning_rate": 5.193333333333333e-05,
+      "loss": 0.0657,
+      "step": 780
+    },
+    {
+      "grad_norm": 0.8061581254005432,
+      "learning_rate": 5.2600000000000005e-05,
+      "loss": 0.0622,
+      "step": 790
+    },
+    {
+      "grad_norm": 1.1061843633651733,
+      "learning_rate": 5.326666666666666e-05,
+      "loss": 0.0637,
+      "step": 800
+    },
+    {
+      "grad_norm": 1.2182507514953613,
+      "learning_rate": 5.3933333333333334e-05,
+      "loss": 0.0589,
+      "step": 810
+    },
+    {
+      "grad_norm": 1.4238909482955933,
+      "learning_rate": 5.4600000000000006e-05,
+      "loss": 0.0614,
+      "step": 820
+    },
+    {
+      "grad_norm": 0.9367620944976807,
+      "learning_rate": 5.5266666666666664e-05,
+      "loss": 0.0622,
+      "step": 830
+    },
+    {
+      "grad_norm": 1.3980387449264526,
+      "learning_rate": 5.5933333333333335e-05,
+      "loss": 0.065,
+      "step": 840
+    },
+    {
+      "grad_norm": 1.3358958959579468,
+      "learning_rate": 5.66e-05,
+      "loss": 0.068,
+      "step": 850
+    },
+    {
+      "grad_norm": 0.9675244092941284,
+      "learning_rate": 5.726666666666667e-05,
+      "loss": 0.0634,
+      "step": 860
+    },
+    {
+      "grad_norm": 0.9360224008560181,
+      "learning_rate": 5.7933333333333337e-05,
+      "loss": 0.0586,
+      "step": 870
+    },
+    {
+      "grad_norm": 1.0572683811187744,
+      "learning_rate": 5.86e-05,
+      "loss": 0.0641,
+      "step": 880
+    },
+    {
+      "grad_norm": 1.307563066482544,
+      "learning_rate": 5.926666666666667e-05,
+      "loss": 0.0593,
+      "step": 890
+    },
+    {
+      "grad_norm": 0.8857334852218628,
+      "learning_rate": 5.9933333333333345e-05,
+      "loss": 0.0634,
+      "step": 900
+    },
+    {
+      "grad_norm": 1.2183889150619507,
+      "learning_rate": 6.06e-05,
+      "loss": 0.0571,
+      "step": 910
+    },
+    {
+      "grad_norm": 1.2149014472961426,
+      "learning_rate": 6.126666666666667e-05,
+      "loss": 0.0586,
+      "step": 920
+    },
+    {
+      "grad_norm": 0.971682071685791,
+      "learning_rate": 6.193333333333333e-05,
+      "loss": 0.0609,
+      "step": 930
+    },
+    {
+      "grad_norm": 1.2368853092193604,
+      "learning_rate": 6.26e-05,
+      "loss": 0.0595,
+      "step": 940
+    },
+    {
+      "grad_norm": 0.9683165550231934,
+      "learning_rate": 6.326666666666667e-05,
+      "loss": 0.0616,
+      "step": 950
+    },
+    {
+      "grad_norm": 0.8244102001190186,
+      "learning_rate": 6.393333333333333e-05,
+      "loss": 0.0575,
+      "step": 960
+    },
+    {
+      "grad_norm": 0.970955491065979,
+      "learning_rate": 6.460000000000001e-05,
+      "loss": 0.0581,
+      "step": 970
+    },
+    {
+      "grad_norm": 1.0662046670913696,
+      "learning_rate": 6.526666666666666e-05,
+      "loss": 0.0557,
+      "step": 980
+    },
+    {
+      "grad_norm": 0.9161382913589478,
+      "learning_rate": 6.593333333333334e-05,
+      "loss": 0.0622,
+      "step": 990
+    },
+    {
+      "grad_norm": 0.8752316832542419,
+      "learning_rate": 6.66e-05,
+      "loss": 0.0539,
+      "step": 1000
+    },
+    {
+      "grad_norm": 0.7521458268165588,
+      "learning_rate": 6.726666666666667e-05,
+      "loss": 0.0506,
+      "step": 1010
+    },
+    {
+      "grad_norm": 0.8169746994972229,
+      "learning_rate": 6.793333333333334e-05,
+      "loss": 0.0603,
+      "step": 1020
+    },
+    {
+      "grad_norm": 0.8718522191047668,
+      "learning_rate": 6.860000000000001e-05,
+      "loss": 0.0516,
+      "step": 1030
+    },
+    {
+      "grad_norm": 0.7131385207176208,
+      "learning_rate": 6.926666666666667e-05,
+      "loss": 0.0492,
+      "step": 1040
+    },
+    {
+      "grad_norm": 0.8163769841194153,
+      "learning_rate": 6.993333333333334e-05,
+      "loss": 0.0475,
+      "step": 1050
+    },
+    {
+      "grad_norm": 0.6880462169647217,
+      "learning_rate": 7.06e-05,
+      "loss": 0.0457,
+      "step": 1060
+    },
+    {
+      "grad_norm": 1.1300956010818481,
+      "learning_rate": 7.126666666666667e-05,
+      "loss": 0.0508,
+      "step": 1070
+    },
+    {
+      "grad_norm": 1.0228790044784546,
+      "learning_rate": 7.193333333333334e-05,
+      "loss": 0.0524,
+      "step": 1080
+    },
+    {
+      "grad_norm": 1.0921752452850342,
+      "learning_rate": 7.26e-05,
+      "loss": 0.0499,
+      "step": 1090
+    },
+    {
+      "grad_norm": 0.5290265679359436,
+      "learning_rate": 7.326666666666667e-05,
+      "loss": 0.0495,
+      "step": 1100
+    },
+    {
+      "grad_norm": 1.0428764820098877,
+      "learning_rate": 7.393333333333333e-05,
+      "loss": 0.0481,
+      "step": 1110
+    },
+    {
+      "grad_norm": 0.6707554459571838,
+      "learning_rate": 7.46e-05,
+      "loss": 0.0487,
+      "step": 1120
+    },
+    {
+      "grad_norm": 0.7440798282623291,
+      "learning_rate": 7.526666666666668e-05,
+      "loss": 0.0509,
+      "step": 1130
+    },
+    {
+      "grad_norm": 0.9219576120376587,
+      "learning_rate": 7.593333333333334e-05,
+      "loss": 0.0501,
+      "step": 1140
+    },
+    {
+      "grad_norm": 1.2103952169418335,
+      "learning_rate": 7.66e-05,
+      "loss": 0.0483,
+      "step": 1150
+    },
+    {
+      "grad_norm": 0.8884477615356445,
+      "learning_rate": 7.726666666666667e-05,
+      "loss": 0.0512,
+      "step": 1160
+    },
+    {
+      "grad_norm": 0.9153897762298584,
+      "learning_rate": 7.793333333333333e-05,
+      "loss": 0.0508,
+      "step": 1170
+    },
+    {
+      "grad_norm": 0.6855103969573975,
+      "learning_rate": 7.860000000000001e-05,
+      "loss": 0.0479,
+      "step": 1180
+    },
+    {
+      "grad_norm": 0.6617424488067627,
+      "learning_rate": 7.926666666666666e-05,
+      "loss": 0.0438,
+      "step": 1190
+    },
+    {
+      "grad_norm": 0.8509991765022278,
+      "learning_rate": 7.993333333333334e-05,
+      "loss": 0.0502,
+      "step": 1200
+    },
+    {
+      "grad_norm": 0.8399519920349121,
+      "learning_rate": 8.060000000000001e-05,
+      "loss": 0.0466,
+      "step": 1210
+    },
+    {
+      "grad_norm": 0.6967006325721741,
+      "learning_rate": 8.126666666666667e-05,
+      "loss": 0.0458,
+      "step": 1220
+    },
+    {
+      "grad_norm": 0.7278540134429932,
+      "learning_rate": 8.193333333333334e-05,
+      "loss": 0.047,
+      "step": 1230
+    },
+    {
+      "grad_norm": 0.8349273204803467,
+      "learning_rate": 8.26e-05,
+      "loss": 0.0448,
+      "step": 1240
+    },
+    {
+      "grad_norm": 0.8303181529045105,
+      "learning_rate": 8.326666666666667e-05,
+      "loss": 0.0431,
+      "step": 1250
+    },
+    {
+      "grad_norm": 0.9701292514801025,
+      "learning_rate": 8.393333333333335e-05,
+      "loss": 0.0469,
+      "step": 1260
+    },
+    {
+      "grad_norm": 0.8359465003013611,
+      "learning_rate": 8.46e-05,
+      "loss": 0.0444,
+      "step": 1270
+    },
+    {
+      "grad_norm": 0.6769126057624817,
+      "learning_rate": 8.526666666666667e-05,
+      "loss": 0.0483,
+      "step": 1280
+    },
+    {
+      "grad_norm": 0.7515086531639099,
+      "learning_rate": 8.593333333333333e-05,
+      "loss": 0.0463,
+      "step": 1290
+    },
+    {
+      "grad_norm": 0.8276355266571045,
+      "learning_rate": 8.66e-05,
+      "loss": 0.0438,
+      "step": 1300
+    },
+    {
+      "grad_norm": 0.5742467045783997,
+      "learning_rate": 8.726666666666667e-05,
+      "loss": 0.0426,
+      "step": 1310
+    },
+    {
+      "grad_norm": 0.8865119814872742,
+      "learning_rate": 8.793333333333333e-05,
+      "loss": 0.0389,
+      "step": 1320
+    },
+    {
+      "grad_norm": 0.6985632181167603,
+      "learning_rate": 8.86e-05,
+      "loss": 0.0482,
+      "step": 1330
+    },
+    {
+      "grad_norm": 0.7202398180961609,
+      "learning_rate": 8.926666666666668e-05,
+      "loss": 0.0491,
+      "step": 1340
+    },
+    {
+      "grad_norm": 0.6817957758903503,
+      "learning_rate": 8.993333333333334e-05,
+      "loss": 0.0467,
+      "step": 1350
+    },
+    {
+      "grad_norm": 0.6327383518218994,
+      "learning_rate": 9.06e-05,
+      "loss": 0.0445,
+      "step": 1360
+    },
+    {
+      "grad_norm": 0.4578998386859894,
+      "learning_rate": 9.126666666666667e-05,
+      "loss": 0.0437,
+      "step": 1370
+    },
+    {
+      "grad_norm": 0.7874047160148621,
+      "learning_rate": 9.193333333333334e-05,
+      "loss": 0.0381,
+      "step": 1380
+    },
+    {
+      "grad_norm": 0.6632944345474243,
+      "learning_rate": 9.260000000000001e-05,
+      "loss": 0.0412,
+      "step": 1390
+    },
+    {
+      "grad_norm": 0.7300984263420105,
+      "learning_rate": 9.326666666666667e-05,
+      "loss": 0.0386,
+      "step": 1400
+    },
+    {
+      "grad_norm": 0.7691074013710022,
+      "learning_rate": 9.393333333333334e-05,
+      "loss": 0.0412,
+      "step": 1410
+    },
+    {
+      "grad_norm": 0.9861852526664734,
+      "learning_rate": 9.46e-05,
+      "loss": 0.0404,
+      "step": 1420
+    },
+    {
+      "grad_norm": 0.8763015866279602,
+      "learning_rate": 9.526666666666667e-05,
+      "loss": 0.0415,
+      "step": 1430
+    },
+    {
+      "grad_norm": 0.7889045476913452,
+      "learning_rate": 9.593333333333334e-05,
+      "loss": 0.0401,
+      "step": 1440
+    },
+    {
+      "grad_norm": 0.7392305731773376,
+      "learning_rate": 9.66e-05,
+      "loss": 0.0391,
+      "step": 1450
+    },
+    {
+      "grad_norm": 0.7114829421043396,
+      "learning_rate": 9.726666666666667e-05,
+      "loss": 0.0453,
+      "step": 1460
+    },
+    {
+      "grad_norm": 0.5600028038024902,
+      "learning_rate": 9.793333333333333e-05,
+      "loss": 0.0385,
+      "step": 1470
+    },
+    {
+      "grad_norm": 0.7662580013275146,
+      "learning_rate": 9.86e-05,
+      "loss": 0.0399,
+      "step": 1480
+    },
+    {
+      "grad_norm": 0.7539925575256348,
+      "learning_rate": 9.926666666666668e-05,
+      "loss": 0.0407,
+      "step": 1490
+    },
+    {
+      "grad_norm": 0.7231246829032898,
+      "learning_rate": 9.993333333333334e-05,
+      "loss": 0.04,
+      "step": 1500
+    },
+    {
+      "grad_norm": 0.6617307662963867,
+      "learning_rate": 9.999997539434007e-05,
+      "loss": 0.0429,
+      "step": 1510
+    },
+    {
+      "grad_norm": 0.7075979709625244,
+      "learning_rate": 9.999989033776898e-05,
+      "loss": 0.0377,
+      "step": 1520
+    },
+    {
+      "grad_norm": 0.6477856636047363,
+      "learning_rate": 9.999974452661641e-05,
+      "loss": 0.0417,
+      "step": 1530
+    },
+    {
+      "grad_norm": 0.5997700691223145,
+      "learning_rate": 9.999953796105959e-05,
+      "loss": 0.0394,
+      "step": 1540
+    },
+    {
+      "grad_norm": 0.5923762917518616,
+      "learning_rate": 9.999927064134949e-05,
+      "loss": 0.0372,
+      "step": 1550
+    },
+    {
+      "grad_norm": 0.46595168113708496,
+      "learning_rate": 9.999894256781095e-05,
+      "loss": 0.0385,
+      "step": 1560
+    },
+    {
+      "grad_norm": 0.583235502243042,
+      "learning_rate": 9.99985537408426e-05,
+      "loss": 0.0388,
+      "step": 1570
+    },
+    {
+      "grad_norm": 0.7547820806503296,
+      "learning_rate": 9.999810416091688e-05,
+      "loss": 0.0378,
+      "step": 1580
+    },
+    {
+      "grad_norm": 0.49323832988739014,
+      "learning_rate": 9.99975938285801e-05,
+      "loss": 0.0386,
+      "step": 1590
+    },
+    {
+      "grad_norm": 0.6207570433616638,
+      "learning_rate": 9.999702274445236e-05,
+      "loss": 0.0412,
+      "step": 1600
+    },
+    {
+      "grad_norm": 0.7229306697845459,
+      "learning_rate": 9.999639090922756e-05,
+      "loss": 0.038,
+      "step": 1610
+    },
+    {
+      "grad_norm": 0.6029353141784668,
+      "learning_rate": 9.999569832367346e-05,
+      "loss": 0.0358,
+      "step": 1620
+    },
+    {
+      "grad_norm": 0.5695170164108276,
+      "learning_rate": 9.999494498863162e-05,
+      "loss": 0.038,
+      "step": 1630
+    },
+    {
+      "grad_norm": 0.4581107795238495,
+      "learning_rate": 9.99941309050174e-05,
+      "loss": 0.0336,
+      "step": 1640
+    },
+    {
+      "grad_norm": 0.5888658761978149,
+      "learning_rate": 9.999325607381999e-05,
+      "loss": 0.0363,
+      "step": 1650
+    },
+    {
+      "grad_norm": 0.42355862259864807,
+      "learning_rate": 9.999232049610238e-05,
+      "loss": 0.0329,
+      "step": 1660
+    },
+    {
+      "grad_norm": 0.44192200899124146,
+      "learning_rate": 9.999132417300142e-05,
+      "loss": 0.0318,
+      "step": 1670
+    },
+    {
+      "grad_norm": 0.582292377948761,
+      "learning_rate": 9.99902671057277e-05,
+      "loss": 0.037,
+      "step": 1680
+    },
+    {
+      "grad_norm": 0.640484631061554,
+      "learning_rate": 9.998914929556569e-05,
+      "loss": 0.0415,
+      "step": 1690
+    },
+    {
+      "grad_norm": 0.5324658155441284,
+      "learning_rate": 9.998797074387361e-05,
+      "loss": 0.0406,
+      "step": 1700
+    },
+    {
+      "grad_norm": 0.4396641254425049,
+      "learning_rate": 9.99867314520835e-05,
+      "loss": 0.0358,
+      "step": 1710
+    },
+    {
+      "grad_norm": 0.5642094016075134,
+      "learning_rate": 9.998543142170126e-05,
+      "loss": 0.0355,
+      "step": 1720
+    },
+    {
+      "grad_norm": 0.4254641830921173,
+      "learning_rate": 9.99840706543065e-05,
+      "loss": 0.0347,
+      "step": 1730
+    },
+    {
+      "grad_norm": 0.6551932096481323,
+      "learning_rate": 9.998264915155274e-05,
+      "loss": 0.0331,
+      "step": 1740
+    },
+    {
+      "grad_norm": 0.6991862654685974,
+      "learning_rate": 9.998116691516718e-05,
+      "loss": 0.0369,
+      "step": 1750
+    },
+    {
+      "grad_norm": 0.5255552530288696,
+      "learning_rate": 9.997962394695091e-05,
+      "loss": 0.0357,
+      "step": 1760
+    },
+    {
+      "grad_norm": 0.5381750464439392,
+      "learning_rate": 9.997802024877875e-05,
+      "loss": 0.0369,
+      "step": 1770
+    },
+    {
+      "grad_norm": 0.46699705719947815,
+      "learning_rate": 9.99763558225994e-05,
+      "loss": 0.0332,
+      "step": 1780
+    },
+    {
+      "grad_norm": 0.5546779632568359,
+      "learning_rate": 9.997463067043526e-05,
+      "loss": 0.035,
+      "step": 1790
+    },
+    {
+      "grad_norm": 0.5206839442253113,
+      "learning_rate": 9.997284479438253e-05,
+      "loss": 0.0332,
+      "step": 1800
+    },
+    {
+      "grad_norm": 0.5379693508148193,
+      "learning_rate": 9.997099819661127e-05,
+      "loss": 0.0309,
+      "step": 1810
+    },
+    {
+      "grad_norm": 0.4955759346485138,
+      "learning_rate": 9.996909087936524e-05,
+      "loss": 0.0353,
+      "step": 1820
+    },
+    {
+      "grad_norm": 0.42540910840034485,
+      "learning_rate": 9.996712284496201e-05,
+      "loss": 0.0327,
+      "step": 1830
+    },
+    {
+      "grad_norm": 0.3907061219215393,
+      "learning_rate": 9.996509409579293e-05,
+      "loss": 0.0323,
+      "step": 1840
+    },
+    {
+      "grad_norm": 0.6000369787216187,
+      "learning_rate": 9.996300463432312e-05,
+      "loss": 0.0336,
+      "step": 1850
+    },
+    {
+      "grad_norm": 0.5673982501029968,
+      "learning_rate": 9.996085446309148e-05,
+      "loss": 0.0369,
+      "step": 1860
+    },
+    {
+      "grad_norm": 0.49209660291671753,
+      "learning_rate": 9.995864358471066e-05,
+      "loss": 0.0341,
+      "step": 1870
+    },
+    {
+      "grad_norm": 0.7095867991447449,
+      "learning_rate": 9.99563720018671e-05,
+      "loss": 0.0316,
+      "step": 1880
+    },
+    {
+      "grad_norm": 0.6636500358581543,
+      "learning_rate": 9.995403971732098e-05,
+      "loss": 0.035,
+      "step": 1890
+    },
+    {
+      "grad_norm": 0.4743417799472809,
+      "learning_rate": 9.995164673390625e-05,
+      "loss": 0.0311,
+      "step": 1900
+    },
+    {
+      "grad_norm": 0.5287648439407349,
+      "learning_rate": 9.994919305453059e-05,
+      "loss": 0.031,
+      "step": 1910
+    },
+    {
+      "grad_norm": 0.45866692066192627,
+      "learning_rate": 9.994667868217548e-05,
+      "loss": 0.032,
+      "step": 1920
+    },
+    {
+      "grad_norm": 0.5198220610618591,
+      "learning_rate": 9.99441036198961e-05,
+      "loss": 0.0326,
+      "step": 1930
+    },
+    {
+      "grad_norm": 0.5405005812644958,
+      "learning_rate": 9.99414678708214e-05,
+      "loss": 0.0322,
+      "step": 1940
+    },
+    {
+      "grad_norm": 0.42117366194725037,
+      "learning_rate": 9.993877143815407e-05,
+      "loss": 0.0314,
+      "step": 1950
+    },
+    {
+      "grad_norm": 0.4138369858264923,
+      "learning_rate": 9.993601432517053e-05,
+      "loss": 0.0311,
+      "step": 1960
+    },
+    {
+      "grad_norm": 0.5241675972938538,
+      "learning_rate": 9.993319653522091e-05,
+      "loss": 0.0322,
+      "step": 1970
+    },
+    {
+      "grad_norm": 0.48230016231536865,
+      "learning_rate": 9.993031807172911e-05,
+      "loss": 0.0321,
+      "step": 1980
+    },
+    {
+      "grad_norm": 0.4062332808971405,
+      "learning_rate": 9.992737893819273e-05,
+      "loss": 0.0317,
+      "step": 1990
+    },
+    {
+      "grad_norm": 0.5321586728096008,
+      "learning_rate": 9.992437913818312e-05,
+      "loss": 0.0307,
+      "step": 2000
+    },
+    {
+      "grad_norm": 0.4237384498119354,
+      "learning_rate": 9.992131867534526e-05,
+      "loss": 0.0288,
+      "step": 2010
+    },
+    {
+      "grad_norm": 0.5707708597183228,
+      "learning_rate": 9.991819755339796e-05,
+      "loss": 0.0318,
+      "step": 2020
+    },
+    {
+      "grad_norm": 0.3146398961544037,
+      "learning_rate": 9.991501577613365e-05,
+      "loss": 0.0312,
+      "step": 2030
+    },
+    {
+      "grad_norm": 0.5046992897987366,
+      "learning_rate": 9.99117733474185e-05,
+      "loss": 0.0305,
+      "step": 2040
+    },
+    {
+      "grad_norm": 0.4503960609436035,
+      "learning_rate": 9.990847027119234e-05,
+      "loss": 0.032,
+      "step": 2050
+    },
+    {
+      "grad_norm": 0.6307269930839539,
+      "learning_rate": 9.990510655146877e-05,
+      "loss": 0.0314,
+      "step": 2060
+    },
+    {
+      "grad_norm": 0.5946677327156067,
+      "learning_rate": 9.990168219233496e-05,
+      "loss": 0.031,
+      "step": 2070
+    },
+    {
+      "grad_norm": 0.478351891040802,
+      "learning_rate": 9.989819719795188e-05,
+      "loss": 0.0289,
+      "step": 2080
+    },
+    {
+      "grad_norm": 0.5508173108100891,
+      "learning_rate": 9.989465157255412e-05,
+      "loss": 0.0334,
+      "step": 2090
+    },
+    {
+      "grad_norm": 0.4113427996635437,
+      "learning_rate": 9.989104532044994e-05,
+      "loss": 0.0293,
+      "step": 2100
+    },
+    {
+      "grad_norm": 0.47436055541038513,
+      "learning_rate": 9.988737844602128e-05,
+      "loss": 0.0264,
+      "step": 2110
+    },
+    {
+      "grad_norm": 0.6576429009437561,
+      "learning_rate": 9.988365095372372e-05,
+      "loss": 0.029,
+      "step": 2120
+    },
+    {
+      "grad_norm": 0.4389994740486145,
+      "learning_rate": 9.987986284808654e-05,
+      "loss": 0.0313,
+      "step": 2130
+    },
+    {
+      "grad_norm": 0.5478070378303528,
+      "learning_rate": 9.987601413371264e-05,
+      "loss": 0.0303,
+      "step": 2140
+    },
+    {
+      "grad_norm": 0.4985005557537079,
+      "learning_rate": 9.987210481527855e-05,
+      "loss": 0.0284,
+      "step": 2150
+    },
+    {
+      "grad_norm": 0.4611583650112152,
+      "learning_rate": 9.98681348975345e-05,
+      "loss": 0.0283,
+      "step": 2160
+    },
+    {
+      "grad_norm": 0.43860509991645813,
+      "learning_rate": 9.986410438530427e-05,
+      "loss": 0.0269,
+      "step": 2170
+    },
+    {
+      "grad_norm": 0.541072428226471,
+      "learning_rate": 9.986001328348534e-05,
+      "loss": 0.0273,
+      "step": 2180
+    },
+    {
+      "grad_norm": 0.4160540997982025,
+      "learning_rate": 9.985586159704878e-05,
+      "loss": 0.0267,
+      "step": 2190
+    },
+    {
+      "grad_norm": 0.5684108138084412,
+      "learning_rate": 9.985164933103929e-05,
+      "loss": 0.0331,
+      "step": 2200
+    },
+    {
+      "grad_norm": 0.60076904296875,
+      "learning_rate": 9.984737649057513e-05,
+      "loss": 0.0315,
+      "step": 2210
+    },
+    {
+      "grad_norm": 0.4548479914665222,
+      "learning_rate": 9.984304308084827e-05,
+      "loss": 0.0332,
+      "step": 2220
+    },
+    {
+      "grad_norm": 0.32628098130226135,
+      "learning_rate": 9.983864910712416e-05,
+      "loss": 0.0277,
+      "step": 2230
+    },
+    {
+      "grad_norm": 0.4243181645870209,
+      "learning_rate": 9.98341945747419e-05,
+      "loss": 0.0271,
+      "step": 2240
+    },
+    {
+      "grad_norm": 0.4711628258228302,
+      "learning_rate": 9.98296794891142e-05,
+      "loss": 0.0283,
+      "step": 2250
+    },
+    {
+      "grad_norm": 0.3072936534881592,
+      "learning_rate": 9.982510385572725e-05,
+      "loss": 0.0294,
+      "step": 2260
+    },
+    {
+      "grad_norm": 0.37349826097488403,
+      "learning_rate": 9.982046768014094e-05,
+      "loss": 0.0271,
+      "step": 2270
+    },
+    {
+      "grad_norm": 0.41922029852867126,
+      "learning_rate": 9.981577096798863e-05,
+      "loss": 0.0255,
+      "step": 2280
+    },
+    {
+      "grad_norm": 0.41333553194999695,
+      "learning_rate": 9.981101372497727e-05,
+      "loss": 0.0276,
+      "step": 2290
+    },
+    {
+      "grad_norm": 0.49039316177368164,
+      "learning_rate": 9.980619595688737e-05,
+      "loss": 0.0264,
+      "step": 2300
+    },
+    {
+      "grad_norm": 0.45840826630592346,
+      "learning_rate": 9.980131766957295e-05,
+      "loss": 0.0289,
+      "step": 2310
+    },
+    {
+      "grad_norm": 0.43983516097068787,
+      "learning_rate": 9.979637886896163e-05,
+      "loss": 0.0295,
+      "step": 2320
+    },
+    {
+      "grad_norm": 0.519990861415863,
+      "learning_rate": 9.979137956105447e-05,
+      "loss": 0.0351,
+      "step": 2330
+    },
+    {
+      "grad_norm": 0.37182801961898804,
+      "learning_rate": 9.978631975192613e-05,
+      "loss": 0.0272,
+      "step": 2340
+    },
+    {
+      "grad_norm": 0.3941374123096466,
+      "learning_rate": 9.978119944772475e-05,
+      "loss": 0.0338,
+      "step": 2350
+    },
+    {
+      "grad_norm": 0.49252980947494507,
+      "learning_rate": 9.977601865467197e-05,
+      "loss": 0.029,
+      "step": 2360
+    },
+    {
+      "grad_norm": 0.36538147926330566,
+      "learning_rate": 9.977077737906297e-05,
+      "loss": 0.0271,
+      "step": 2370
+    },
+    {
+      "grad_norm": 0.48212969303131104,
+      "learning_rate": 9.976547562726636e-05,
+      "loss": 0.0286,
+      "step": 2380
+    },
+    {
+      "grad_norm": 0.3857809603214264,
+      "learning_rate": 9.976011340572429e-05,
+      "loss": 0.0299,
+      "step": 2390
+    },
+    {
+      "grad_norm": 0.3532734215259552,
+      "learning_rate": 9.975469072095237e-05,
+      "loss": 0.029,
+      "step": 2400
+    },
+    {
+      "grad_norm": 0.3491647243499756,
+      "learning_rate": 9.974920757953965e-05,
+      "loss": 0.0246,
+      "step": 2410
+    },
+    {
+      "grad_norm": 0.40819358825683594,
+      "learning_rate": 9.97436639881487e-05,
+      "loss": 0.0322,
+      "step": 2420
+    },
+    {
+      "grad_norm": 0.4187079071998596,
+      "learning_rate": 9.973805995351545e-05,
+      "loss": 0.0284,
+      "step": 2430
+    },
+    {
+      "grad_norm": 0.47072187066078186,
+      "learning_rate": 9.973239548244939e-05,
+      "loss": 0.0287,
+      "step": 2440
+    },
+    {
+      "grad_norm": 0.4486375153064728,
+      "learning_rate": 9.972667058183333e-05,
+      "loss": 0.0274,
+      "step": 2450
+    },
+    {
+      "grad_norm": 0.387708455324173,
+      "learning_rate": 9.972088525862362e-05,
+      "loss": 0.0246,
+      "step": 2460
+    },
+    {
+      "grad_norm": 0.45445823669433594,
+      "learning_rate": 9.971503951984995e-05,
+      "loss": 0.0262,
+      "step": 2470
+    },
+    {
+      "grad_norm": 0.5244600176811218,
+      "learning_rate": 9.970913337261543e-05,
+      "loss": 0.026,
+      "step": 2480
+    },
+    {
+      "grad_norm": 0.33774614334106445,
+      "learning_rate": 9.97031668240966e-05,
+      "loss": 0.0327,
+      "step": 2490
+    },
+    {
+      "grad_norm": 0.45767754316329956,
+      "learning_rate": 9.969713988154339e-05,
+      "loss": 0.0278,
+      "step": 2500
+    },
+    {
+      "grad_norm": 0.452639639377594,
+      "learning_rate": 9.969105255227906e-05,
+      "loss": 0.0281,
+      "step": 2510
+    },
+    {
+      "grad_norm": 0.4824736416339874,
+      "learning_rate": 9.968490484370035e-05,
+      "loss": 0.0268,
+      "step": 2520
+    },
+    {
+      "grad_norm": 0.42104530334472656,
+      "learning_rate": 9.967869676327726e-05,
+      "loss": 0.0277,
+      "step": 2530
+    },
+    {
+      "grad_norm": 0.3387914299964905,
+      "learning_rate": 9.96724283185532e-05,
+      "loss": 0.0256,
+      "step": 2540
+    },
+    {
+      "grad_norm": 0.46041038632392883,
+      "learning_rate": 9.966609951714494e-05,
+      "loss": 0.0283,
+      "step": 2550
+    },
+    {
+      "grad_norm": 0.482858806848526,
+      "learning_rate": 9.965971036674255e-05,
+      "loss": 0.0302,
+      "step": 2560
+    },
+    {
+      "grad_norm": 0.4021144509315491,
+      "learning_rate": 9.965326087510947e-05,
+      "loss": 0.025,
+      "step": 2570
+    },
+    {
+      "grad_norm": 0.4823286533355713,
+      "learning_rate": 9.964675105008244e-05,
+      "loss": 0.0258,
+      "step": 2580
+    },
+    {
+      "grad_norm": 0.4102964401245117,
+      "learning_rate": 9.964018089957147e-05,
+      "loss": 0.0236,
+      "step": 2590
+    },
+    {
+      "grad_norm": 0.4453105330467224,
+      "learning_rate": 9.963355043155996e-05,
+      "loss": 0.0293,
+      "step": 2600
+    },
+    {
+      "grad_norm": 0.42528659105300903,
+      "learning_rate": 9.962685965410455e-05,
+      "loss": 0.0284,
+      "step": 2610
+    },
+    {
+      "grad_norm": 0.4291960597038269,
+      "learning_rate": 9.962010857533514e-05,
+      "loss": 0.0271,
+      "step": 2620
+    },
+    {
+      "grad_norm": 0.38245514035224915,
+      "learning_rate": 9.961329720345493e-05,
+      "loss": 0.0258,
+      "step": 2630
+    },
+    {
+      "grad_norm": 0.5179672241210938,
+      "learning_rate": 9.96064255467404e-05,
+      "loss": 0.0281,
+      "step": 2640
+    },
+    {
+      "grad_norm": 0.38057994842529297,
+      "learning_rate": 9.959949361354126e-05,
+      "loss": 0.0281,
+      "step": 2650
+    },
+    {
+      "grad_norm": 0.42559677362442017,
+      "learning_rate": 9.959250141228045e-05,
+      "loss": 0.0262,
+      "step": 2660
+    },
+    {
+      "grad_norm": 0.3978665769100189,
+      "learning_rate": 9.958544895145414e-05,
+      "loss": 0.0265,
+      "step": 2670
+    },
+    {
+      "grad_norm": 0.42062318325042725,
+      "learning_rate": 9.957833623963177e-05,
+      "loss": 0.0272,
+      "step": 2680
+    },
+    {
+      "grad_norm": 0.36069512367248535,
+      "learning_rate": 9.957116328545593e-05,
+      "loss": 0.0279,
+      "step": 2690
+    },
+    {
+      "grad_norm": 0.3319280743598938,
+      "learning_rate": 9.956393009764244e-05,
+      "loss": 0.0244,
+      "step": 2700
+    },
+    {
+      "grad_norm": 0.38565364480018616,
+      "learning_rate": 9.955663668498032e-05,
+      "loss": 0.0242,
+      "step": 2710
+    },
+    {
+      "grad_norm": 0.44830113649368286,
+      "learning_rate": 9.954928305633173e-05,
+      "loss": 0.025,
+      "step": 2720
+    },
+    {
+      "grad_norm": 0.3731986880302429,
+      "learning_rate": 9.954186922063204e-05,
+      "loss": 0.0254,
+      "step": 2730
+    },
+    {
+      "grad_norm": 0.36890801787376404,
+      "learning_rate": 9.953439518688974e-05,
+      "loss": 0.0242,
+      "step": 2740
+    },
+    {
+      "grad_norm": 0.47049516439437866,
+      "learning_rate": 9.952686096418652e-05,
+      "loss": 0.0223,
+      "step": 2750
+    },
+    {
+      "grad_norm": 0.2939774990081787,
+      "learning_rate": 9.951926656167714e-05,
+      "loss": 0.023,
+      "step": 2760
+    },
+    {
+      "grad_norm": 0.4406392574310303,
+      "learning_rate": 9.951161198858953e-05,
+      "loss": 0.0264,
+      "step": 2770
+    },
+    {
+      "grad_norm": 0.41060757637023926,
+      "learning_rate": 9.950389725422471e-05,
+      "loss": 0.0226,
+      "step": 2780
+    },
+    {
+      "grad_norm": 0.3619372844696045,
+      "learning_rate": 9.949612236795682e-05,
+      "loss": 0.0233,
+      "step": 2790
+    },
+    {
+      "grad_norm": 0.36144009232521057,
+      "learning_rate": 9.948828733923305e-05,
+      "loss": 0.0252,
+      "step": 2800
+    },
+    {
+      "grad_norm": 0.4079873561859131,
+      "learning_rate": 9.948039217757374e-05,
+      "loss": 0.0262,
+      "step": 2810
+    },
+    {
+      "grad_norm": 0.47260719537734985,
+      "learning_rate": 9.947243689257225e-05,
+      "loss": 0.0272,
+      "step": 2820
+    },
+    {
+      "grad_norm": 0.3860573172569275,
+      "learning_rate": 9.946442149389497e-05,
+      "loss": 0.0241,
+      "step": 2830
+    },
+    {
+      "grad_norm": 0.3537745773792267,
+      "learning_rate": 9.945634599128139e-05,
+      "loss": 0.0239,
+      "step": 2840
+    },
+    {
+      "grad_norm": 0.37954181432724,
+      "learning_rate": 9.944821039454402e-05,
+      "loss": 0.0247,
+      "step": 2850
+    },
+    {
+      "grad_norm": 0.46963128447532654,
+      "learning_rate": 9.944001471356835e-05,
+      "loss": 0.0287,
+      "step": 2860
+    },
+    {
+      "grad_norm": 0.37109649181365967,
+      "learning_rate": 9.94317589583129e-05,
+      "loss": 0.0281,
+      "step": 2870
+    },
+    {
+      "grad_norm": 0.2578738033771515,
+      "learning_rate": 9.942344313880922e-05,
+      "loss": 0.0236,
+      "step": 2880
+    },
+    {
+      "grad_norm": 0.33481988310813904,
+      "learning_rate": 9.941506726516179e-05,
+      "loss": 0.0246,
+      "step": 2890
+    },
+    {
+      "grad_norm": 0.2952929437160492,
+      "learning_rate": 9.94066313475481e-05,
+      "loss": 0.024,
+      "step": 2900
+    },
+    {
+      "grad_norm": 0.33066704869270325,
+      "learning_rate": 9.939813539621857e-05,
+      "loss": 0.0227,
+      "step": 2910
+    },
+    {
+      "grad_norm": 0.4288365840911865,
+      "learning_rate": 9.93895794214966e-05,
+      "loss": 0.0268,
+      "step": 2920
+    },
+    {
+      "grad_norm": 0.37286025285720825,
+      "learning_rate": 9.938096343377852e-05,
+      "loss": 0.0284,
+      "step": 2930
+    },
+    {
+      "grad_norm": 0.35089603066444397,
+      "learning_rate": 9.937228744353353e-05,
+      "loss": 0.0285,
+      "step": 2940
+    },
+    {
+      "grad_norm": 0.49157941341400146,
+      "learning_rate": 9.936355146130379e-05,
+      "loss": 0.0254,
+      "step": 2950
+    },
+    {
+      "grad_norm": 0.5023393630981445,
+      "learning_rate": 9.935475549770435e-05,
+      "loss": 0.0258,
+      "step": 2960
+    },
+    {
+      "grad_norm": 0.4146648645401001,
+      "learning_rate": 9.934589956342315e-05,
+      "loss": 0.0244,
+      "step": 2970
+    },
+    {
+      "grad_norm": 0.4460250735282898,
+      "learning_rate": 9.933698366922093e-05,
+      "loss": 0.0237,
+      "step": 2980
+    },
+    {
+      "grad_norm": 0.44778910279273987,
+      "learning_rate": 9.93280078259314e-05,
+      "loss": 0.025,
+      "step": 2990
+    },
+    {
+      "grad_norm": 0.3062107563018799,
+      "learning_rate": 9.931897204446104e-05,
+      "loss": 0.0241,
+      "step": 3000
+    },
+    {
+      "grad_norm": 0.3562763035297394,
+      "learning_rate": 9.930987633578915e-05,
+      "loss": 0.0235,
+      "step": 3010
+    },
+    {
+      "grad_norm": 0.37201783061027527,
+      "learning_rate": 9.93007207109679e-05,
+      "loss": 0.0231,
+      "step": 3020
+    },
+    {
+      "grad_norm": 0.3234574496746063,
+      "learning_rate": 9.929150518112224e-05,
+      "loss": 0.0215,
+      "step": 3030
+    },
+    {
+      "grad_norm": 0.4681793451309204,
+      "learning_rate": 9.928222975744991e-05,
+      "loss": 0.0244,
+      "step": 3040
+    },
+    {
+      "grad_norm": 0.4070403277873993,
+      "learning_rate": 9.92728944512214e-05,
+      "loss": 0.0225,
+      "step": 3050
+    },
+    {
+      "grad_norm": 0.44380709528923035,
+      "learning_rate": 9.926349927378001e-05,
+      "loss": 0.0222,
+      "step": 3060
+    },
+    {
+      "grad_norm": 0.43774229288101196,
+      "learning_rate": 9.925404423654174e-05,
+      "loss": 0.023,
+      "step": 3070
+    },
+    {
+      "grad_norm": 0.35820233821868896,
+      "learning_rate": 9.924452935099537e-05,
+      "loss": 0.0238,
+      "step": 3080
+    },
+    {
+      "grad_norm": 0.29745739698410034,
+      "learning_rate": 9.92349546287024e-05,
+      "loss": 0.024,
+      "step": 3090
+    },
+    {
+      "grad_norm": 0.46426740288734436,
+      "learning_rate": 9.9225320081297e-05,
+      "loss": 0.0245,
+      "step": 3100
+    },
+    {
+      "grad_norm": 0.3944103419780731,
+      "learning_rate": 9.921562572048606e-05,
+      "loss": 0.0238,
+      "step": 3110
+    },
+    {
+      "grad_norm": 0.35884836316108704,
+      "learning_rate": 9.920587155804913e-05,
+      "loss": 0.0247,
+      "step": 3120
+    },
+    {
+      "grad_norm": 0.3041065037250519,
+      "learning_rate": 9.919605760583845e-05,
+      "loss": 0.0246,
+      "step": 3130
+    },
+    {
+      "grad_norm": 0.263491690158844,
+      "learning_rate": 9.91861838757789e-05,
+      "loss": 0.0227,
+      "step": 3140
+    },
+    {
+      "grad_norm": 0.24135278165340424,
+      "learning_rate": 9.917625037986798e-05,
+      "loss": 0.0231,
+      "step": 3150
+    },
+    {
+      "grad_norm": 0.42098352313041687,
+      "learning_rate": 9.916625713017583e-05,
+      "loss": 0.021,
+      "step": 3160
+    },
+    {
+      "grad_norm": 0.3001592457294464,
+      "learning_rate": 9.915620413884519e-05,
+      "loss": 0.0208,
+      "step": 3170
+    },
+    {
+      "grad_norm": 0.3589301109313965,
+      "learning_rate": 9.914609141809139e-05,
+      "loss": 0.021,
+      "step": 3180
+    },
+    {
+      "grad_norm": 0.4449957311153412,
+      "learning_rate": 9.913591898020235e-05,
+      "loss": 0.0231,
+      "step": 3190
+    },
+    {
+      "grad_norm": 0.49094122648239136,
+      "learning_rate": 9.912568683753853e-05,
+      "loss": 0.0242,
+      "step": 3200
+    },
+    {
+      "grad_norm": 0.31530943512916565,
+      "learning_rate": 9.911539500253295e-05,
+      "loss": 0.0253,
+      "step": 3210
+    },
+    {
+      "grad_norm": 0.3312493860721588,
+      "learning_rate": 9.910504348769118e-05,
+      "loss": 0.0219,
+      "step": 3220
+    },
+    {
+      "grad_norm": 0.25662603974342346,
+      "learning_rate": 9.909463230559127e-05,
+      "loss": 0.0202,
+      "step": 3230
+    },
+    {
+      "grad_norm": 0.35309165716171265,
+      "learning_rate": 9.908416146888376e-05,
+      "loss": 0.0226,
+      "step": 3240
+    },
+    {
+      "grad_norm": 0.3551497161388397,
+      "learning_rate": 9.907363099029175e-05,
+      "loss": 0.0218,
+      "step": 3250
+    },
+    {
+      "grad_norm": 0.3065674602985382,
+      "learning_rate": 9.906304088261073e-05,
+      "loss": 0.0259,
+      "step": 3260
+    },
+    {
+      "grad_norm": 0.34306737780570984,
+      "learning_rate": 9.905239115870872e-05,
+      "loss": 0.0246,
+      "step": 3270
+    },
+    {
+      "grad_norm": 0.25299543142318726,
+      "learning_rate": 9.90416818315261e-05,
+      "loss": 0.022,
+      "step": 3280
+    },
+    {
+      "grad_norm": 0.3820056915283203,
+      "learning_rate": 9.903091291407573e-05,
+      "loss": 0.027,
+      "step": 3290
+    },
+    {
+      "grad_norm": 0.3862016201019287,
+      "learning_rate": 9.902008441944286e-05,
+      "loss": 0.0231,
+      "step": 3300
+    },
+    {
+      "grad_norm": 0.3469122648239136,
+      "learning_rate": 9.900919636078512e-05,
+      "loss": 0.0226,
+      "step": 3310
+    },
+    {
+      "grad_norm": 0.3083968758583069,
+      "learning_rate": 9.899824875133255e-05,
+      "loss": 0.0242,
+      "step": 3320
+    },
+    {
+      "grad_norm": 0.46409106254577637,
+      "learning_rate": 9.898724160438749e-05,
+      "loss": 0.0219,
+      "step": 3330
+    },
+    {
+      "grad_norm": 0.2895025908946991,
+      "learning_rate": 9.89761749333247e-05,
+      "loss": 0.0232,
+      "step": 3340
+    },
+    {
+      "grad_norm": 0.33076900243759155,
+      "learning_rate": 9.896504875159122e-05,
+      "loss": 0.0216,
+      "step": 3350
+    },
+    {
+      "grad_norm": 0.30149585008621216,
+      "learning_rate": 9.89538630727064e-05,
+      "loss": 0.0212,
+      "step": 3360
+    },
+    {
+      "grad_norm": 0.26640868186950684,
+      "learning_rate": 9.894261791026189e-05,
+      "loss": 0.0231,
+      "step": 3370
+    },
+    {
+      "grad_norm": 0.43421223759651184,
+      "learning_rate": 9.893131327792165e-05,
+      "loss": 0.0211,
+      "step": 3380
+    },
+    {
+      "grad_norm": 0.34027573466300964,
+      "learning_rate": 9.891994918942182e-05,
+      "loss": 0.0202,
+      "step": 3390
+    },
+    {
+      "grad_norm": 0.3347245156764984,
+      "learning_rate": 9.890852565857092e-05,
+      "loss": 0.0229,
+      "step": 3400
+    },
+    {
+      "grad_norm": 0.35026970505714417,
+      "learning_rate": 9.889704269924954e-05,
+      "loss": 0.0206,
+      "step": 3410
+    },
+    {
+      "grad_norm": 0.3942737579345703,
+      "learning_rate": 9.888550032541059e-05,
+      "loss": 0.022,
+      "step": 3420
+    },
+    {
+      "grad_norm": 0.39353707432746887,
+      "learning_rate": 9.887389855107916e-05,
+      "loss": 0.0226,
+      "step": 3430
+    },
+    {
+      "grad_norm": 0.3677908182144165,
+      "learning_rate": 9.886223739035248e-05,
+      "loss": 0.0272,
+      "step": 3440
+    },
+    {
+      "grad_norm": 0.37933892011642456,
+      "learning_rate": 9.885051685739997e-05,
+      "loss": 0.0243,
+      "step": 3450
+    },
+    {
+      "grad_norm": 0.4893406629562378,
+      "learning_rate": 9.883873696646316e-05,
+      "loss": 0.027,
+      "step": 3460
+    },
+    {
+      "grad_norm": 0.300083190202713,
+      "learning_rate": 9.882689773185575e-05,
+      "loss": 0.0207,
+      "step": 3470
+    },
+    {
+      "grad_norm": 0.40003201365470886,
+      "learning_rate": 9.881499916796353e-05,
+      "loss": 0.0243,
+      "step": 3480
+    },
+    {
+      "grad_norm": 0.3844235837459564,
+      "learning_rate": 9.880304128924434e-05,
+      "loss": 0.0225,
+      "step": 3490
+    },
+    {
+      "grad_norm": 0.31904521584510803,
+      "learning_rate": 9.879102411022817e-05,
+      "loss": 0.0223,
+      "step": 3500
+    },
+    {
+      "grad_norm": 0.34493687748908997,
+      "learning_rate": 9.877894764551703e-05,
+      "loss": 0.0201,
+      "step": 3510
+    },
+    {
+      "grad_norm": 0.41117796301841736,
+      "learning_rate": 9.876681190978494e-05,
+      "loss": 0.024,
+      "step": 3520
+    },
+    {
+      "grad_norm": 0.3274706304073334,
+      "learning_rate": 9.875461691777797e-05,
+      "loss": 0.0258,
+      "step": 3530
+    },
+    {
+      "grad_norm": 0.28383010625839233,
+      "learning_rate": 9.874236268431417e-05,
+      "loss": 0.0219,
+      "step": 3540
+    },
+    {
+      "grad_norm": 0.3705231845378876,
+      "learning_rate": 9.873004922428361e-05,
+      "loss": 0.0213,
+      "step": 3550
+    },
+    {
+      "grad_norm": 0.3192824423313141,
+      "learning_rate": 9.871767655264829e-05,
+      "loss": 0.022,
+      "step": 3560
+    },
+    {
+      "grad_norm": 0.25980308651924133,
+      "learning_rate": 9.87052446844422e-05,
+      "loss": 0.0211,
+      "step": 3570
+    },
+    {
+      "grad_norm": 0.2518915832042694,
+      "learning_rate": 9.869275363477122e-05,
+      "loss": 0.024,
+      "step": 3580
+    },
+    {
+      "grad_norm": 0.3379889726638794,
+      "learning_rate": 9.868020341881312e-05,
+      "loss": 0.0209,
+      "step": 3590
+    },
+    {
+      "grad_norm": 0.39600417017936707,
+      "learning_rate": 9.866759405181765e-05,
+      "loss": 0.0237,
+      "step": 3600
+    },
+    {
+      "grad_norm": 0.33133968710899353,
+      "learning_rate": 9.865492554910633e-05,
+      "loss": 0.0209,
+      "step": 3610
+    },
+    {
+      "grad_norm": 0.32698678970336914,
+      "learning_rate": 9.864219792607262e-05,
+      "loss": 0.0196,
+      "step": 3620
+    },
+    {
+      "grad_norm": 0.3060948848724365,
+      "learning_rate": 9.862941119818177e-05,
+      "loss": 0.0192,
+      "step": 3630
+    },
+    {
+      "grad_norm": 0.3851982355117798,
+      "learning_rate": 9.861656538097086e-05,
+      "loss": 0.0208,
+      "step": 3640
+    },
+    {
+      "grad_norm": 0.49290308356285095,
+      "learning_rate": 9.860366049004877e-05,
+      "loss": 0.022,
+      "step": 3650
+    },
+    {
+      "grad_norm": 0.3673649728298187,
+      "learning_rate": 9.859069654109615e-05,
+      "loss": 0.0233,
+      "step": 3660
+    },
+    {
+      "grad_norm": 0.2948131561279297,
+      "learning_rate": 9.857767354986545e-05,
+      "loss": 0.023,
+      "step": 3670
+    },
+    {
+      "grad_norm": 0.3540187478065491,
+      "learning_rate": 9.856459153218078e-05,
+      "loss": 0.0196,
+      "step": 3680
+    },
+    {
+      "grad_norm": 0.27228233218193054,
+      "learning_rate": 9.855145050393808e-05,
+      "loss": 0.0211,
+      "step": 3690
+    },
+    {
+      "grad_norm": 0.24771757423877716,
+      "learning_rate": 9.85382504811049e-05,
+      "loss": 0.021,
+      "step": 3700
+    },
+    {
+      "grad_norm": 0.26397937536239624,
+      "learning_rate": 9.852499147972054e-05,
+      "loss": 0.0201,
+      "step": 3710
+    },
+    {
+      "grad_norm": 0.28861305117607117,
+      "learning_rate": 9.851167351589592e-05,
+      "loss": 0.0215,
+      "step": 3720
+    },
+    {
+      "grad_norm": 0.3233491778373718,
+      "learning_rate": 9.849829660581363e-05,
+      "loss": 0.0201,
+      "step": 3730
+    },
+    {
+      "grad_norm": 0.3898829519748688,
+      "learning_rate": 9.848486076572787e-05,
+      "loss": 0.0221,
+      "step": 3740
+    },
+    {
+      "grad_norm": 0.38374295830726624,
+      "learning_rate": 9.847136601196446e-05,
+      "loss": 0.0223,
+      "step": 3750
+    },
+    {
+      "grad_norm": 0.2840521037578583,
+      "learning_rate": 9.845781236092078e-05,
+      "loss": 0.0221,
+      "step": 3760
+    },
+    {
+      "grad_norm": 0.3588821589946747,
+      "learning_rate": 9.844419982906583e-05,
+      "loss": 0.0212,
+      "step": 3770
+    },
+    {
+      "grad_norm": 0.34918320178985596,
+      "learning_rate": 9.843052843294008e-05,
+      "loss": 0.019,
+      "step": 3780
+    },
+    {
+      "grad_norm": 0.3414880633354187,
+      "learning_rate": 9.841679818915559e-05,
+      "loss": 0.0216,
+      "step": 3790
+    },
+    {
+      "grad_norm": 0.35619089007377625,
+      "learning_rate": 9.840300911439591e-05,
+      "loss": 0.0213,
+      "step": 3800
+    },
+    {
+      "grad_norm": 0.3069482147693634,
+      "learning_rate": 9.838916122541603e-05,
+      "loss": 0.019,
+      "step": 3810
+    },
+    {
+      "grad_norm": 0.31180906295776367,
+      "learning_rate": 9.837525453904246e-05,
+      "loss": 0.0196,
+      "step": 3820
+    },
+    {
+      "grad_norm": 0.29071903228759766,
+      "learning_rate": 9.836128907217314e-05,
+      "loss": 0.0201,
+      "step": 3830
+    },
+    {
+      "grad_norm": 0.336658239364624,
+      "learning_rate": 9.834726484177743e-05,
+      "loss": 0.0198,
+      "step": 3840
+    },
+    {
+      "grad_norm": 0.3188108503818512,
+      "learning_rate": 9.833318186489609e-05,
+      "loss": 0.0205,
+      "step": 3850
+    },
+    {
+      "grad_norm": 0.2784384787082672,
+      "learning_rate": 9.831904015864126e-05,
+      "loss": 0.0196,
+      "step": 3860
+    },
+    {
+      "grad_norm": 0.3234076201915741,
+      "learning_rate": 9.830483974019645e-05,
+      "loss": 0.0211,
+      "step": 3870
+    },
+    {
+      "grad_norm": 0.3331151008605957,
+      "learning_rate": 9.82905806268165e-05,
+      "loss": 0.0191,
+      "step": 3880
+    },
+    {
+      "grad_norm": 0.31715860962867737,
+      "learning_rate": 9.82762628358276e-05,
+      "loss": 0.0173,
+      "step": 3890
+    },
+    {
+      "grad_norm": 0.3170963227748871,
+      "learning_rate": 9.826188638462718e-05,
+      "loss": 0.018,
+      "step": 3900
+    },
+    {
+      "grad_norm": 0.2912008762359619,
+      "learning_rate": 9.824745129068402e-05,
+      "loss": 0.022,
+      "step": 3910
+    },
+    {
+      "grad_norm": 0.48704785108566284,
+      "learning_rate": 9.82329575715381e-05,
+      "loss": 0.0242,
+      "step": 3920
+    },
+    {
+      "grad_norm": 0.28424951434135437,
+      "learning_rate": 9.821840524480066e-05,
+      "loss": 0.0203,
+      "step": 3930
+    },
+    {
+      "grad_norm": 0.2560309171676636,
+      "learning_rate": 9.820379432815414e-05,
+      "loss": 0.0203,
+      "step": 3940
+    },
+    {
+      "grad_norm": 0.3647140562534332,
+      "learning_rate": 9.81891248393522e-05,
+      "loss": 0.0187,
+      "step": 3950
+    },
+    {
+      "grad_norm": 0.31988921761512756,
+      "learning_rate": 9.817439679621963e-05,
+      "loss": 0.0207,
+      "step": 3960
+    },
+    {
+      "grad_norm": 0.3253921866416931,
+      "learning_rate": 9.815961021665243e-05,
+      "loss": 0.0212,
+      "step": 3970
+    },
+    {
+      "grad_norm": 0.29170772433280945,
+      "learning_rate": 9.814476511861763e-05,
+      "loss": 0.0263,
+      "step": 3980
+    },
+    {
+      "grad_norm": 0.4017032980918884,
+      "learning_rate": 9.812986152015348e-05,
+      "loss": 0.0265,
+      "step": 3990
+    },
+    {
+      "grad_norm": 0.28655076026916504,
+      "learning_rate": 9.811489943936922e-05,
+      "loss": 0.0205,
+      "step": 4000
+    },
+    {
+      "grad_norm": 0.23081740736961365,
+      "learning_rate": 9.809987889444522e-05,
+      "loss": 0.021,
+      "step": 4010
+    },
+    {
+      "grad_norm": 0.34190860390663147,
+      "learning_rate": 9.808479990363282e-05,
+      "loss": 0.0203,
+      "step": 4020
+    },
+    {
+      "grad_norm": 0.2937399744987488,
+      "learning_rate": 9.806966248525445e-05,
+      "loss": 0.0188,
+      "step": 4030
+    },
+    {
+      "grad_norm": 0.2845223546028137,
+      "learning_rate": 9.805446665770348e-05,
+      "loss": 0.0191,
+      "step": 4040
+    },
+    {
+      "grad_norm": 0.2973966896533966,
+      "learning_rate": 9.803921243944429e-05,
+      "loss": 0.0175,
+      "step": 4050
+    },
+    {
+      "grad_norm": 0.3171347379684448,
+      "learning_rate": 9.802389984901218e-05,
+      "loss": 0.0195,
+      "step": 4060
+    },
+    {
+      "grad_norm": 0.2669548988342285,
+      "learning_rate": 9.80085289050134e-05,
+      "loss": 0.0183,
+      "step": 4070
+    },
+    {
+      "grad_norm": 0.42913228273391724,
+      "learning_rate": 9.799309962612508e-05,
+      "loss": 0.0196,
+      "step": 4080
+    },
+    {
+      "grad_norm": 0.4162735342979431,
+      "learning_rate": 9.797761203109527e-05,
+      "loss": 0.0215,
+      "step": 4090
+    },
+    {
+      "grad_norm": 0.31903430819511414,
+      "learning_rate": 9.796206613874283e-05,
+      "loss": 0.0195,
+      "step": 4100
+    },
+    {
+      "grad_norm": 0.3554883301258087,
+      "learning_rate": 9.794646196795754e-05,
+      "loss": 0.0201,
+      "step": 4110
+    },
+    {
+      "grad_norm": 0.31217727065086365,
+      "learning_rate": 9.793079953769987e-05,
+      "loss": 0.019,
+      "step": 4120
+    },
+    {
+      "grad_norm": 0.28474947810173035,
+      "learning_rate": 9.79150788670012e-05,
+      "loss": 0.0184,
+      "step": 4130
+    },
+    {
+      "grad_norm": 0.32388654351234436,
+      "learning_rate": 9.78992999749636e-05,
+      "loss": 0.0216,
+      "step": 4140
+    },
+    {
+      "grad_norm": 0.20659787952899933,
+      "learning_rate": 9.788346288075994e-05,
+      "loss": 0.0181,
+      "step": 4150
+    },
+    {
+      "grad_norm": 0.2864404022693634,
+      "learning_rate": 9.786756760363373e-05,
+      "loss": 0.0207,
+      "step": 4160
+    },
+    {
+      "grad_norm": 0.2995869815349579,
+      "learning_rate": 9.78516141628993e-05,
+      "loss": 0.0182,
+      "step": 4170
+    },
+    {
+      "grad_norm": 0.31789812445640564,
+      "learning_rate": 9.783560257794154e-05,
+      "loss": 0.0188,
+      "step": 4180
+    },
+    {
+      "grad_norm": 0.28194430470466614,
+      "learning_rate": 9.781953286821603e-05,
+      "loss": 0.0168,
+      "step": 4190
+    },
+    {
+      "grad_norm": 0.31114333868026733,
+      "learning_rate": 9.780340505324901e-05,
+      "loss": 0.0172,
+      "step": 4200
+    },
+    {
+      "grad_norm": 0.2011519819498062,
+      "learning_rate": 9.778721915263727e-05,
+      "loss": 0.0177,
+      "step": 4210
+    },
+    {
+      "grad_norm": 0.22189873456954956,
+      "learning_rate": 9.777097518604824e-05,
+      "loss": 0.0167,
+      "step": 4220
+    },
+    {
+      "grad_norm": 0.30673474073410034,
+      "learning_rate": 9.775467317321984e-05,
+      "loss": 0.0161,
+      "step": 4230
+    },
+    {
+      "grad_norm": 0.3615068197250366,
+      "learning_rate": 9.773831313396055e-05,
+      "loss": 0.0201,
+      "step": 4240
+    },
+    {
+      "grad_norm": 0.25624963641166687,
+      "learning_rate": 9.77218950881494e-05,
+      "loss": 0.0176,
+      "step": 4250
+    },
+    {
+      "grad_norm": 0.32759591937065125,
+      "learning_rate": 9.770541905573583e-05,
+      "loss": 0.0182,
+      "step": 4260
+    },
+    {
+      "grad_norm": 0.23994171619415283,
+      "learning_rate": 9.768888505673976e-05,
+      "loss": 0.0157,
+      "step": 4270
+    },
+    {
+      "grad_norm": 0.29931554198265076,
+      "learning_rate": 9.767229311125162e-05,
+      "loss": 0.0173,
+      "step": 4280
+    },
+    {
+      "grad_norm": 0.3035125732421875,
+      "learning_rate": 9.765564323943211e-05,
+      "loss": 0.0181,
+      "step": 4290
+    },
+    {
+      "grad_norm": 0.3667411804199219,
+      "learning_rate": 9.763893546151244e-05,
+      "loss": 0.0192,
+      "step": 4300
+    },
+    {
+      "grad_norm": 0.2947080731391907,
+      "learning_rate": 9.762216979779412e-05,
+      "loss": 0.0203,
+      "step": 4310
+    },
+    {
+      "grad_norm": 0.2096947431564331,
+      "learning_rate": 9.760534626864902e-05,
+      "loss": 0.0185,
+      "step": 4320
+    },
+    {
+      "grad_norm": 0.301518976688385,
+      "learning_rate": 9.758846489451931e-05,
+      "loss": 0.0203,
+      "step": 4330
+    },
+    {
+      "grad_norm": 0.3711138367652893,
+      "learning_rate": 9.757152569591748e-05,
+      "loss": 0.0199,
+      "step": 4340
+    },
+    {
+      "grad_norm": 0.23832212388515472,
+      "learning_rate": 9.75545286934262e-05,
+      "loss": 0.0184,
+      "step": 4350
+    },
+    {
+      "grad_norm": 0.3737817108631134,
+      "learning_rate": 9.753747390769847e-05,
+      "loss": 0.018,
+      "step": 4360
+    },
+    {
+      "grad_norm": 0.2761482000350952,
+      "learning_rate": 9.752036135945744e-05,
+      "loss": 0.0194,
+      "step": 4370
+    },
+    {
+      "grad_norm": 0.40149354934692383,
+      "learning_rate": 9.750319106949649e-05,
+      "loss": 0.0199,
+      "step": 4380
+    },
+    {
+      "grad_norm": 0.2389083057641983,
+      "learning_rate": 9.748596305867913e-05,
+      "loss": 0.0193,
+      "step": 4390
+    },
+    {
+      "grad_norm": 0.2904573380947113,
+      "learning_rate": 9.746867734793903e-05,
+      "loss": 0.0194,
+      "step": 4400
+    },
+    {
+      "grad_norm": 0.28680405020713806,
+      "learning_rate": 9.745133395827993e-05,
+      "loss": 0.0213,
+      "step": 4410
+    },
+    {
+      "grad_norm": 0.3064373731613159,
+      "learning_rate": 9.743393291077572e-05,
+      "loss": 0.0178,
+      "step": 4420
+    },
+    {
+      "grad_norm": 0.23364102840423584,
+      "learning_rate": 9.741647422657028e-05,
+      "loss": 0.02,
+      "step": 4430
+    },
+    {
+      "grad_norm": 0.2293742299079895,
+      "learning_rate": 9.739895792687758e-05,
+      "loss": 0.0173,
+      "step": 4440
+    },
+    {
+      "grad_norm": 0.4201240539550781,
+      "learning_rate": 9.738138403298157e-05,
+      "loss": 0.0203,
+      "step": 4450
+    },
+    {
+      "grad_norm": 0.30462443828582764,
+      "learning_rate": 9.736375256623619e-05,
+      "loss": 0.0173,
+      "step": 4460
+    },
+    {
+      "grad_norm": 0.26824820041656494,
+      "learning_rate": 9.734606354806533e-05,
+      "loss": 0.0163,
+      "step": 4470
+    },
+    {
+      "grad_norm": 0.3436148762702942,
+      "learning_rate": 9.73283169999628e-05,
+      "loss": 0.0196,
+      "step": 4480
+    },
+    {
+      "grad_norm": 0.2987440526485443,
+      "learning_rate": 9.731051294349238e-05,
+      "loss": 0.0185,
+      "step": 4490
+    },
+    {
+      "grad_norm": 0.3387744724750519,
+      "learning_rate": 9.729265140028762e-05,
+      "loss": 0.0208,
+      "step": 4500
+    },
+    {
+      "grad_norm": 0.22033269703388214,
+      "learning_rate": 9.727473239205201e-05,
+      "loss": 0.0154,
+      "step": 4510
+    },
+    {
+      "grad_norm": 0.3045775890350342,
+      "learning_rate": 9.725675594055883e-05,
+      "loss": 0.0156,
+      "step": 4520
+    },
+    {
+      "grad_norm": 0.3046495020389557,
+      "learning_rate": 9.723872206765116e-05,
+      "loss": 0.0172,
+      "step": 4530
+    },
+    {
+      "grad_norm": 0.28411006927490234,
+      "learning_rate": 9.722063079524185e-05,
+      "loss": 0.0176,
+      "step": 4540
+    },
+    {
+      "grad_norm": 0.2794414758682251,
+      "learning_rate": 9.720248214531351e-05,
+      "loss": 0.017,
+      "step": 4550
+    },
+    {
+      "grad_norm": 0.28699204325675964,
+      "learning_rate": 9.718427613991848e-05,
+      "loss": 0.0183,
+      "step": 4560
+    },
+    {
+      "grad_norm": 0.23602908849716187,
+      "learning_rate": 9.716601280117873e-05,
+      "loss": 0.015,
+      "step": 4570
+    },
+    {
+      "grad_norm": 0.3002392053604126,
+      "learning_rate": 9.714769215128596e-05,
+      "loss": 0.0187,
+      "step": 4580
+    },
+    {
+      "grad_norm": 0.2982144355773926,
+      "learning_rate": 9.712931421250152e-05,
+      "loss": 0.0204,
+      "step": 4590
+    },
+    {
+      "grad_norm": 0.33681154251098633,
+      "learning_rate": 9.711087900715627e-05,
+      "loss": 0.0186,
+      "step": 4600
+    },
+    {
+      "grad_norm": 0.2547285854816437,
+      "learning_rate": 9.709238655765078e-05,
+      "loss": 0.0193,
+      "step": 4610
+    },
+    {
+      "grad_norm": 0.2880988121032715,
+      "learning_rate": 9.707383688645511e-05,
+      "loss": 0.0169,
+      "step": 4620
+    },
+    {
+      "grad_norm": 0.2253848910331726,
+      "learning_rate": 9.705523001610883e-05,
+      "loss": 0.0173,
+      "step": 4630
+    },
+    {
+      "grad_norm": 0.24737682938575745,
+      "learning_rate": 9.703656596922107e-05,
+      "loss": 0.0152,
+      "step": 4640
+    },
+    {
+      "grad_norm": 0.26355868577957153,
+      "learning_rate": 9.70178447684704e-05,
+      "loss": 0.0171,
+      "step": 4650
+    },
+    {
+      "grad_norm": 0.33829301595687866,
+      "learning_rate": 9.699906643660483e-05,
+      "loss": 0.0177,
+      "step": 4660
+    },
+    {
+      "grad_norm": 0.27919113636016846,
+      "learning_rate": 9.698023099644185e-05,
+      "loss": 0.018,
+      "step": 4670
+    },
+    {
+      "grad_norm": 0.3479491174221039,
+      "learning_rate": 9.696133847086823e-05,
+      "loss": 0.0165,
+      "step": 4680
+    },
+    {
+      "grad_norm": 0.23865260183811188,
+      "learning_rate": 9.694238888284022e-05,
+      "loss": 0.0206,
+      "step": 4690
+    },
+    {
+      "grad_norm": 0.2442196160554886,
+      "learning_rate": 9.692338225538333e-05,
+      "loss": 0.0168,
+      "step": 4700
+    },
+    {
+      "grad_norm": 0.2947649657726288,
+      "learning_rate": 9.690431861159241e-05,
+      "loss": 0.0169,
+      "step": 4710
+    },
+    {
+      "grad_norm": 0.336140900850296,
+      "learning_rate": 9.688519797463161e-05,
+      "loss": 0.0184,
+      "step": 4720
+    },
+    {
+      "grad_norm": 0.39535579085350037,
+      "learning_rate": 9.686602036773426e-05,
+      "loss": 0.018,
+      "step": 4730
+    },
+    {
+      "grad_norm": 0.28196772933006287,
+      "learning_rate": 9.684678581420302e-05,
+      "loss": 0.0175,
+      "step": 4740
+    },
+    {
+      "grad_norm": 0.33065393567085266,
+      "learning_rate": 9.682749433740962e-05,
+      "loss": 0.0201,
+      "step": 4750
+    },
+    {
+      "grad_norm": 0.298825740814209,
+      "learning_rate": 9.680814596079507e-05,
+      "loss": 0.0194,
+      "step": 4760
+    },
+    {
+      "grad_norm": 0.28574132919311523,
+      "learning_rate": 9.678874070786945e-05,
+      "loss": 0.018,
+      "step": 4770
+    },
+    {
+      "grad_norm": 0.2778089940547943,
+      "learning_rate": 9.676927860221199e-05,
+      "loss": 0.0191,
+      "step": 4780
+    },
+    {
+      "grad_norm": 0.2691532373428345,
+      "learning_rate": 9.674975966747097e-05,
+      "loss": 0.018,
+      "step": 4790
+    },
+    {
+      "grad_norm": 0.28231796622276306,
+      "learning_rate": 9.673018392736374e-05,
+      "loss": 0.0161,
+      "step": 4800
+    },
+    {
+      "grad_norm": 0.3621669113636017,
+      "learning_rate": 9.671055140567667e-05,
+      "loss": 0.0189,
+      "step": 4810
+    },
+    {
+      "grad_norm": 0.3124409019947052,
+      "learning_rate": 9.669086212626511e-05,
+      "loss": 0.0182,
+      "step": 4820
+    },
+    {
+      "grad_norm": 0.30056077241897583,
+      "learning_rate": 9.667111611305341e-05,
+      "loss": 0.0183,
+      "step": 4830
+    },
+    {
+      "grad_norm": 0.3483680486679077,
+      "learning_rate": 9.665131339003486e-05,
+      "loss": 0.018,
+      "step": 4840
+    },
+    {
+      "grad_norm": 0.2676101624965668,
+      "learning_rate": 9.663145398127158e-05,
+      "loss": 0.0156,
+      "step": 4850
+    },
+    {
+      "grad_norm": 0.23675771057605743,
+      "learning_rate": 9.661153791089467e-05,
+      "loss": 0.0185,
+      "step": 4860
+    },
+    {
+      "grad_norm": 0.2168475240468979,
+      "learning_rate": 9.659156520310402e-05,
+      "loss": 0.0173,
+      "step": 4870
+    },
+    {
+      "grad_norm": 0.2815901041030884,
+      "learning_rate": 9.657153588216835e-05,
+      "loss": 0.0161,
+      "step": 4880
+    },
+    {
+      "grad_norm": 0.2556694746017456,
+      "learning_rate": 9.655144997242516e-05,
+      "loss": 0.0156,
+      "step": 4890
+    },
+    {
+      "grad_norm": 0.21819503605365753,
+      "learning_rate": 9.653130749828075e-05,
+      "loss": 0.0157,
+      "step": 4900
+    },
+    {
+      "grad_norm": 0.2848703861236572,
+      "learning_rate": 9.65111084842101e-05,
+      "loss": 0.0164,
+      "step": 4910
+    },
+    {
+      "grad_norm": 0.3069203794002533,
+      "learning_rate": 9.649085295475695e-05,
+      "loss": 0.0201,
+      "step": 4920
+    },
+    {
+      "grad_norm": 0.29833701252937317,
+      "learning_rate": 9.647054093453365e-05,
+      "loss": 0.0199,
+      "step": 4930
+    },
+    {
+      "grad_norm": 0.20307645201683044,
+      "learning_rate": 9.645017244822123e-05,
+      "loss": 0.0167,
+      "step": 4940
+    },
+    {
+      "grad_norm": 0.21525900065898895,
+      "learning_rate": 9.642974752056931e-05,
+      "loss": 0.018,
+      "step": 4950
+    },
+    {
+      "grad_norm": 0.33719441294670105,
+      "learning_rate": 9.640926617639613e-05,
+      "loss": 0.019,
+      "step": 4960
+    },
+    {
+      "grad_norm": 0.3345821797847748,
+      "learning_rate": 9.638872844058843e-05,
+      "loss": 0.0174,
+      "step": 4970
+    },
+    {
+      "grad_norm": 0.29934078454971313,
+      "learning_rate": 9.63681343381015e-05,
+      "loss": 0.019,
+      "step": 4980
+    },
+    {
+      "grad_norm": 0.31826502084732056,
+      "learning_rate": 9.634748389395914e-05,
+      "loss": 0.017,
+      "step": 4990
+    },
+    {
+      "grad_norm": 0.3035266101360321,
+      "learning_rate": 9.632677713325353e-05,
+      "loss": 0.018,
+      "step": 5000
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 30000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 9223372036854775807,
+  "save_steps": 2500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

16d_30k/checkpoint-5000/training_args.bin

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

16d_30k/checkpoint-5000/wandb_config.json

@@ -0,0 +1 @@
+{"project": "screwing_hil", "run_id": "groot_16d_30k"}