Add files using upload-large-folder tool

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.3/mgd_mask_ratio_ablation.yaml

@@ -0,0 +1,49 @@
+experiment_name: mgd_sweep_from_trained_base
+policy_type: groot_mgd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/MGD_Raw_VLM
+sweep:
+  mgd_token_mask_ratio:
+  - 0.5
+  - 0.3
+  - 0.7
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_mgd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    mgd_enabled: true
+    mgd_fm_loss_weight: 0.0
+    mgd_loss_weight: 1.0
+    mgd_sequence_hidden_dim: 512
+    mgd_token_mask_ratio: ${sweep.mgd_token_mask_ratio}
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+- name: phase3_mgd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    mgd_enabled: false
+    mgd_fm_loss_weight: 1.0
+    mgd_loss_weight: 0.0
+    mgd_token_mask_ratio: ${sweep.mgd_token_mask_ratio}
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.3/phase2/config.json

@@ -0,0 +1,75 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_MGD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "mgd_enabled": true,
+  "mgd_fm_loss_weight": 0.0,
+  "mgd_loss_weight": 1.0,
+  "mgd_pretrained_projector_path": null,
+  "mgd_sequence_hidden_dim": 512,
+  "mgd_target_dim": 512,
+  "mgd_target_pooling": "flatten",
+  "mgd_target_projection": "frozen_random",
+  "mgd_token_mask_ratio": 0.3,
+  "mgd_use_cosine_loss": true,
+  "mgd_use_mse_loss": false,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.3/phase2/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.3/phase3/config.json

@@ -0,0 +1,75 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_MGD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "mgd_enabled": false,
+  "mgd_fm_loss_weight": 1.0,
+  "mgd_loss_weight": 0.0,
+  "mgd_pretrained_projector_path": null,
+  "mgd_sequence_hidden_dim": 512,
+  "mgd_target_dim": 512,
+  "mgd_target_pooling": "flatten",
+  "mgd_target_projection": "frozen_random",
+  "mgd_token_mask_ratio": 0.3,
+  "mgd_use_cosine_loss": true,
+  "mgd_use_mse_loss": false,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.3/phase3/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.3/resolved_config.yaml

@@ -0,0 +1,51 @@
+experiment_name: mgd_sweep_from_trained_base
+policy_type: groot_mgd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/MGD_Raw_VLM
+sweep:
+  mgd_token_mask_ratio:
+  - 0.5
+  - 0.3
+  - 0.7
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_mgd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    mgd_enabled: true
+    mgd_fm_loss_weight: 0.0
+    mgd_loss_weight: 1.0
+    mgd_sequence_hidden_dim: 512
+    mgd_token_mask_ratio: 0.3
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+- name: phase3_mgd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    mgd_enabled: false
+    mgd_fm_loss_weight: 1.0
+    mgd_loss_weight: 0.0
+    mgd_token_mask_ratio: 0.3
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+resolved_sweep:
+  mgd_token_mask_ratio: 0.3

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.5/mgd_mask_ratio_ablation.yaml

@@ -0,0 +1,49 @@
+experiment_name: mgd_sweep_from_trained_base
+policy_type: groot_mgd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/MGD_Raw_VLM
+sweep:
+  mgd_token_mask_ratio:
+  - 0.5
+  - 0.3
+  - 0.7
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_mgd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    mgd_enabled: true
+    mgd_fm_loss_weight: 0.0
+    mgd_loss_weight: 1.0
+    mgd_sequence_hidden_dim: 512
+    mgd_token_mask_ratio: ${sweep.mgd_token_mask_ratio}
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+- name: phase3_mgd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    mgd_enabled: false
+    mgd_fm_loss_weight: 1.0
+    mgd_loss_weight: 0.0
+    mgd_token_mask_ratio: ${sweep.mgd_token_mask_ratio}
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.5/phase2/config.json

@@ -0,0 +1,75 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_MGD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "mgd_enabled": true,
+  "mgd_fm_loss_weight": 0.0,
+  "mgd_loss_weight": 1.0,
+  "mgd_pretrained_projector_path": null,
+  "mgd_sequence_hidden_dim": 512,
+  "mgd_target_dim": 512,
+  "mgd_target_pooling": "flatten",
+  "mgd_target_projection": "frozen_random",
+  "mgd_token_mask_ratio": 0.5,
+  "mgd_use_cosine_loss": true,
+  "mgd_use_mse_loss": false,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.5/phase2/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.5/phase3/config.json

@@ -0,0 +1,75 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_MGD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "mgd_enabled": false,
+  "mgd_fm_loss_weight": 1.0,
+  "mgd_loss_weight": 0.0,
+  "mgd_pretrained_projector_path": null,
+  "mgd_sequence_hidden_dim": 512,
+  "mgd_target_dim": 512,
+  "mgd_target_pooling": "flatten",
+  "mgd_target_projection": "frozen_random",
+  "mgd_token_mask_ratio": 0.5,
+  "mgd_use_cosine_loss": true,
+  "mgd_use_mse_loss": false,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.5/phase3/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.5/resolved_config.yaml

@@ -0,0 +1,51 @@
+experiment_name: mgd_sweep_from_trained_base
+policy_type: groot_mgd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/MGD_Raw_VLM
+sweep:
+  mgd_token_mask_ratio:
+  - 0.5
+  - 0.3
+  - 0.7
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_mgd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    mgd_enabled: true
+    mgd_fm_loss_weight: 0.0
+    mgd_loss_weight: 1.0
+    mgd_sequence_hidden_dim: 512
+    mgd_token_mask_ratio: 0.5
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+- name: phase3_mgd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    mgd_enabled: false
+    mgd_fm_loss_weight: 1.0
+    mgd_loss_weight: 0.0
+    mgd_token_mask_ratio: 0.5
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+resolved_sweep:
+  mgd_token_mask_ratio: 0.5

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.7/mgd_mask_ratio_ablation.yaml

@@ -0,0 +1,49 @@
+experiment_name: mgd_sweep_from_trained_base
+policy_type: groot_mgd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/MGD_Raw_VLM
+sweep:
+  mgd_token_mask_ratio:
+  - 0.5
+  - 0.3
+  - 0.7
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_mgd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    mgd_enabled: true
+    mgd_fm_loss_weight: 0.0
+    mgd_loss_weight: 1.0
+    mgd_sequence_hidden_dim: 512
+    mgd_token_mask_ratio: ${sweep.mgd_token_mask_ratio}
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+- name: phase3_mgd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    mgd_enabled: false
+    mgd_fm_loss_weight: 1.0
+    mgd_loss_weight: 0.0
+    mgd_token_mask_ratio: ${sweep.mgd_token_mask_ratio}
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.7/phase2/config.json

@@ -0,0 +1,75 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_MGD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "mgd_enabled": true,
+  "mgd_fm_loss_weight": 0.0,
+  "mgd_loss_weight": 1.0,
+  "mgd_pretrained_projector_path": null,
+  "mgd_sequence_hidden_dim": 512,
+  "mgd_target_dim": 512,
+  "mgd_target_pooling": "flatten",
+  "mgd_target_projection": "frozen_random",
+  "mgd_token_mask_ratio": 0.7,
+  "mgd_use_cosine_loss": true,
+  "mgd_use_mse_loss": false,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.7/phase2/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.7/phase3/config.json

@@ -0,0 +1,75 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_MGD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "mgd_enabled": false,
+  "mgd_fm_loss_weight": 1.0,
+  "mgd_loss_weight": 0.0,
+  "mgd_pretrained_projector_path": null,
+  "mgd_sequence_hidden_dim": 512,
+  "mgd_target_dim": 512,
+  "mgd_target_pooling": "flatten",
+  "mgd_target_projection": "frozen_random",
+  "mgd_token_mask_ratio": 0.7,
+  "mgd_use_cosine_loss": true,
+  "mgd_use_mse_loss": false,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.7/phase3/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/MGD_Raw_VLM/mgd_token_mask_ratio_0.7/resolved_config.yaml

@@ -0,0 +1,51 @@
+experiment_name: mgd_sweep_from_trained_base
+policy_type: groot_mgd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/MGD_Raw_VLM
+sweep:
+  mgd_token_mask_ratio:
+  - 0.5
+  - 0.3
+  - 0.7
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_mgd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    mgd_enabled: true
+    mgd_fm_loss_weight: 0.0
+    mgd_loss_weight: 1.0
+    mgd_sequence_hidden_dim: 512
+    mgd_token_mask_ratio: 0.7
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+- name: phase3_mgd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    mgd_enabled: false
+    mgd_fm_loss_weight: 1.0
+    mgd_loss_weight: 0.0
+    mgd_token_mask_ratio: 0.7
+    mgd_use_cosine_loss: true
+    mgd_use_mse_loss: false
+resolved_sweep:
+  mgd_token_mask_ratio: 0.7

VLM_Only/RKD_Raw_VLM/rkd_temp_0.05/phase3/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.05/rkd_temp_ablation.yaml

@@ -0,0 +1,48 @@
+experiment_name: rkd_temp_ablation_from_trained_base
+policy_type: groot_rkd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/RKD_Raw_VLM
+sweep:
+  rkd_temp:
+  - 0.05
+  - 0.1
+  - 0.2
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_rkd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    rkd_enabled: true
+    rkd_fm_loss_weight: 0.0
+    rkd_loss_weight: 1.0
+    rkd_action_temp: ${sweep.rkd_temp}
+    rkd_vlm_temp: ${sweep.rkd_temp}
+    rkd_relation_mode: token_pair_mean
+- name: phase3_rkd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    rkd_enabled: false
+    rkd_fm_loss_weight: 1.0
+    rkd_loss_weight: 0.0
+    rkd_action_temp: ${sweep.rkd_temp}
+    rkd_vlm_temp: ${sweep.rkd_temp}
+    rkd_relation_mode: token_pair_mean

VLM_Only/RKD_Raw_VLM/rkd_temp_0.1/phase2/config.json

@@ -0,0 +1,70 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_RKD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "rkd_action_temp": 0.1,
+  "rkd_enabled": true,
+  "rkd_fm_loss_weight": 0.0,
+  "rkd_loss_weight": 1.0,
+  "rkd_relation_mode": "token_pair_mean",
+  "rkd_vlm_temp": 0.1,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.1/phase2/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.1/phase3/config.json

@@ -0,0 +1,70 @@
+{
+  "action_dim": 32,
+  "action_head_cfg": {
+    "action_dim": 32,
+    "action_horizon": 16,
+    "add_pos_embed": true,
+    "backbone_embedding_dim": 2048,
+    "diffusion_model_cfg": {
+      "attention_head_dim": 48,
+      "cross_attention_dim": 2048,
+      "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
+    },
+    "hidden_size": 1024,
+    "input_embedding_dim": 1536,
+    "max_action_dim": 32,
+    "max_state_dim": 64,
+    "model_dtype": "float32",
+    "noise_beta_alpha": 1.5,
+    "noise_beta_beta": 1.0,
+    "noise_s": 0.999,
+    "num_inference_timesteps": 4,
+    "num_target_vision_tokens": 32,
+    "num_timestep_buckets": 1000,
+    "tune_diffusion_model": true,
+    "tune_projector": 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
+    }
+  },
+  "action_horizon": 16,
+  "architectures": [
+    "GR00T_N1_5_RKD"
+  ],
+  "attn_implementation": null,
+  "backbone_cfg": {
+    "eagle_path": "NVEagle/eagle_er-qwen3_1_7B-Siglip2_400M_stage1_5_128gpu_er_v7_1mlp_nops",
+    "load_bf16": false,
+    "project_to_dim": null,
+    "reproject_vision": false,
+    "select_layer": 12,
+    "tune_llm": false,
+    "tune_visual": true,
+    "use_flash_attention": true
+  },
+  "compute_dtype": "bfloat16",
+  "hidden_size": 2048,
+  "model_dtype": "float32",
+  "model_type": "gr00t_n1_5",
+  "rkd_action_temp": 0.1,
+  "rkd_enabled": false,
+  "rkd_fm_loss_weight": 1.0,
+  "rkd_loss_weight": 0.0,
+  "rkd_relation_mode": "token_pair_mean",
+  "rkd_vlm_temp": 0.1,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3"
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.1/phase3/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.1/resolved_config.yaml

@@ -0,0 +1,50 @@
+experiment_name: rkd_temp_ablation_from_trained_base
+policy_type: groot_rkd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/RKD_Raw_VLM
+sweep:
+  rkd_temp:
+  - 0.05
+  - 0.1
+  - 0.2
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_rkd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    rkd_enabled: true
+    rkd_fm_loss_weight: 0.0
+    rkd_loss_weight: 1.0
+    rkd_action_temp: 0.1
+    rkd_vlm_temp: 0.1
+    rkd_relation_mode: token_pair_mean
+- name: phase3_rkd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    rkd_enabled: false
+    rkd_fm_loss_weight: 1.0
+    rkd_loss_weight: 0.0
+    rkd_action_temp: 0.1
+    rkd_vlm_temp: 0.1
+    rkd_relation_mode: token_pair_mean
+resolved_sweep:
+  rkd_temp: 0.1

VLM_Only/RKD_Raw_VLM/rkd_temp_0.1/rkd_temp_ablation.yaml

@@ -0,0 +1,48 @@
+experiment_name: rkd_temp_ablation_from_trained_base
+policy_type: groot_rkd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/RKD_Raw_VLM
+sweep:
+  rkd_temp:
+  - 0.05
+  - 0.1
+  - 0.2
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_rkd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    rkd_enabled: true
+    rkd_fm_loss_weight: 0.0
+    rkd_loss_weight: 1.0
+    rkd_action_temp: ${sweep.rkd_temp}
+    rkd_vlm_temp: ${sweep.rkd_temp}
+    rkd_relation_mode: token_pair_mean
+- name: phase3_rkd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    rkd_enabled: false
+    rkd_fm_loss_weight: 1.0
+    rkd_loss_weight: 0.0
+    rkd_action_temp: ${sweep.rkd_temp}
+    rkd_vlm_temp: ${sweep.rkd_temp}
+    rkd_relation_mode: token_pair_mean

VLM_Only/RKD_Raw_VLM/rkd_temp_0.2/phase2/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.2/phase3/experiment_cfg/metadata.json

@@ -0,0 +1,431 @@
+{
+    "new_embodiment": {
+        "statistics": {
+            "state": {
+                "base_position": {
+                    "max": [
+                        7.3139495849609375,
+                        0.4876587688922882,
+                        0.7196521759033203
+                    ],
+                    "min": [
+                        -4.94293737411499,
+                        -6.890198230743408,
+                        0.6996827125549316
+                    ],
+                    "mean": [
+                        2.749288365724937,
+                        -2.0455216239909983,
+                        0.700532551020533
+                    ],
+                    "std": [
+                        1.6786295897046262,
+                        1.312897969434244,
+                        0.00133751758906258
+                    ],
+                    "q01": [
+                        -1.4241907881148037,
+                        -5.1716547935950885,
+                        0.6999670898768614
+                    ],
+                    "q99": [
+                        6.198111545831555,
+                        -0.5873531334104489,
+                        0.7038833014242587
+                    ]
+                },
+                "base_rotation": {
+                    "max": [
+                        0.0,
+                        0.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        0.0,
+                        0.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.0,
+                        0.0,
+                        0.23084999587450247,
+                        0.6238431088581847
+                    ],
+                    "std": [
+                        0.0,
+                        0.0,
+                        0.6556404371644047,
+                        0.3572252105732042
+                    ],
+                    "q01": [
+                        0.0,
+                        0.0,
+                        -0.9999999999999999,
+                        1.7482354593273349e-06
+                    ],
+                    "q99": [
+                        0.0,
+                        0.0,
+                        0.9999999999999999,
+                        0.9999999999999999
+                    ]
+                },
+                "end_effector_position_relative": {
+                    "max": [
+                        0.9014528393745422,
+                        0.8003877401351929,
+                        0.9829942584037781
+                    ],
+                    "min": [
+                        -0.37471598386764526,
+                        -0.8472502827644348,
+                        -0.25070279836654663
+                    ],
+                    "mean": [
+                        0.28667332601863704,
+                        -0.038315473170876704,
+                        0.45974658906777444
+                    ],
+                    "std": [
+                        0.17067058828305154,
+                        0.23569952037784195,
+                        0.21950702354181487
+                    ],
+                    "q01": [
+                        0.015172948963784228,
+                        -0.44450502063220615,
+                        0.23314444103329338
+                    ],
+                    "q99": [
+                        0.5609069689572412,
+                        0.38454179128009547,
+                        0.7028102736473852
+                    ]
+                },
+                "end_effector_rotation_relative": {
+                    "max": [
+                        0.9999998807907104,
+                        0.9984455108642578,
+                        0.9434727430343628,
+                        0.9062229990959167
+                    ],
+                    "min": [
+                        -0.9999930262565613,
+                        -0.9988337755203247,
+                        -0.9618149995803833,
+                        2.1872274658107926e-07
+                    ],
+                    "mean": [
+                        -0.25672297401336,
+                        0.024425335806687216,
+                        -0.08740977164483847,
+                        0.16608739644018397
+                    ],
+                    "std": [
+                        0.7932585208392383,
+                        0.3102260195580416,
+                        0.3772472576315148,
+                        0.17451959300158773
+                    ],
+                    "q01": [
+                        -0.99379006987882,
+                        -0.5478102050038828,
+                        -0.6101777952971636,
+                        0.002274085014134748
+                    ],
+                    "q99": [
+                        0.8804856038896288,
+                        0.5910611092873037,
+                        0.5216058042993562,
+                        0.5231965974012255
+                    ]
+                },
+                "gripper_qpos": {
+                    "max": [
+                        0.055869169533252716,
+                        0.010916369967162609
+                    ],
+                    "min": [
+                        -0.011436971835792065,
+                        -0.05664053186774254
+                    ],
+                    "mean": [
+                        0.031651551516052555,
+                        -0.03161481387920421
+                    ],
+                    "std": [
+                        0.013119769222217526,
+                        0.01306078225192402
+                    ],
+                    "q01": [
+                        0.0060999246446777336,
+                        -0.04062601653964198
+                    ],
+                    "q99": [
+                        0.04054891621517283,
+                        -0.0059163567113456345
+                    ]
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0,
+                        0.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0,
+                        0.0
+                    ],
+                    "mean": [
+                        0.008144896753718373,
+                        -0.00018893135146078263,
+                        -0.0008739062727221845,
+                        0.0
+                    ],
+                    "std": [
+                        0.11030045424364411,
+                        0.10148082570313594,
+                        0.08975698373467506,
+                        0.0
+                    ],
+                    "q01": [
+                        -0.07287373067581518,
+                        -0.09446994199569948,
+                        -0.07702818259249572,
+                        0.0
+                    ],
+                    "q99": [
+                        0.04485485414407898,
+                        0.09626941605965177,
+                        0.07610665906122742,
+                        0.0
+                    ]
+                },
+                "control_mode": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.9221974313425939
+                    ],
+                    "std": [
+                        0.38670194342612674
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        -0.384693946883099
+                    ]
+                },
+                "end_effector_position": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.01976129923017491,
+                        -0.020870631344490034,
+                        -0.05993476517349181
+                    ],
+                    "std": [
+                        0.4347024990804053,
+                        0.4221662976569997,
+                        0.3845506843044066
+                    ],
+                    "q01": [
+                        -0.8835792939767807,
+                        -0.9280919251713778,
+                        -0.783361071470956
+                    ],
+                    "q99": [
+                        0.7622084438449307,
+                        0.8625125252609077,
+                        0.7470974753250706
+                    ]
+                },
+                "end_effector_rotation": {
+                    "max": [
+                        1.0,
+                        1.0,
+                        1.0
+                    ],
+                    "min": [
+                        -1.0,
+                        -1.0,
+                        -1.0
+                    ],
+                    "mean": [
+                        0.008089673068544335,
+                        -0.026498254627927348,
+                        0.0029083551743059005
+                    ],
+                    "std": [
+                        0.11216734503733773,
+                        0.13206002613798629,
+                        0.12615870412692864
+                    ],
+                    "q01": [
+                        -0.2814627972846672,
+                        -0.4342386516115439,
+                        -0.34489755628340574
+                    ],
+                    "q99": [
+                        0.31879353266939386,
+                        0.28411797683220563,
+                        0.3597718671669894
+                    ]
+                },
+                "gripper_close": {
+                    "max": [
+                        1.0
+                    ],
+                    "min": [
+                        -1.0
+                    ],
+                    "mean": [
+                        -0.3824228874769045
+                    ],
+                    "std": [
+                        0.9240015096469786
+                    ],
+                    "q01": [
+                        -1.0
+                    ],
+                    "q99": [
+                        0.5597272305813592
+                    ]
+                }
+            }
+        },
+        "modalities": {
+            "video": {
+                "robot0_eye_in_hand": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_left": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                },
+                "robot0_agentview_right": {
+                    "resolution": [
+                        256,
+                        256
+                    ],
+                    "channels": 3,
+                    "fps": 20.0
+                }
+            },
+            "state": {
+                "base_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "base_rotation": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position_relative": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation_relative": {
+                    "absolute": true,
+                    "rotation_type": "quaternion",
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "gripper_qpos": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        2
+                    ],
+                    "continuous": true
+                }
+            },
+            "action": {
+                "base_motion": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        4
+                    ],
+                    "continuous": true
+                },
+                "control_mode": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                },
+                "end_effector_position": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "end_effector_rotation": {
+                    "absolute": true,
+                    "rotation_type": "axis_angle",
+                    "shape": [
+                        3
+                    ],
+                    "continuous": true
+                },
+                "gripper_close": {
+                    "absolute": true,
+                    "rotation_type": null,
+                    "shape": [
+                        1
+                    ],
+                    "continuous": true
+                }
+            }
+        },
+        "embodiment_tag": "new_embodiment"
+    }
+}

VLM_Only/RKD_Raw_VLM/rkd_temp_0.2/resolved_config.yaml

@@ -0,0 +1,50 @@
+experiment_name: rkd_temp_ablation_from_trained_base
+policy_type: groot_rkd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/RKD_Raw_VLM
+sweep:
+  rkd_temp:
+  - 0.05
+  - 0.1
+  - 0.2
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_rkd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    rkd_enabled: true
+    rkd_fm_loss_weight: 0.0
+    rkd_loss_weight: 1.0
+    rkd_action_temp: 0.2
+    rkd_vlm_temp: 0.2
+    rkd_relation_mode: token_pair_mean
+- name: phase3_rkd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    rkd_enabled: false
+    rkd_fm_loss_weight: 1.0
+    rkd_loss_weight: 0.0
+    rkd_action_temp: 0.2
+    rkd_vlm_temp: 0.2
+    rkd_relation_mode: token_pair_mean
+resolved_sweep:
+  rkd_temp: 0.2

VLM_Only/RKD_Raw_VLM/rkd_temp_0.2/rkd_temp_ablation.yaml

@@ -0,0 +1,48 @@
+experiment_name: rkd_temp_ablation_from_trained_base
+policy_type: groot_rkd_raw
+base_model_path: ${HOME}/groot_robocasa/Atomic26_baseline/checkpoint-60000
+output_root: ${HOME}/groot_robocasa/robocasa_v2/VLM_Only/RKD_Raw_VLM
+sweep:
+  rkd_temp:
+  - 0.05
+  - 0.1
+  - 0.2
+dataset:
+  dataset_soup: my_atomic26_human
+training:
+  num_gpus: 2
+  batch_size: 32
+model: null
+phases:
+- name: phase2_rkd_only
+  max_steps: 60000
+  save_steps: 0
+  trainable:
+    preset: freeze_processing_line
+    tune_llm: true
+    tune_visual: true
+    tune_projector: false
+    tune_diffusion_model: false
+  losses:
+    rkd_enabled: true
+    rkd_fm_loss_weight: 0.0
+    rkd_loss_weight: 1.0
+    rkd_action_temp: ${sweep.rkd_temp}
+    rkd_vlm_temp: ${sweep.rkd_temp}
+    rkd_relation_mode: token_pair_mean
+- name: phase3_rkd_fm
+  max_steps: 30000
+  batch_size: 64
+  save_steps: 0
+  trainable:
+    tune_llm: false
+    tune_visual: false
+    tune_projector: true
+    tune_diffusion_model: true
+  losses:
+    rkd_enabled: false
+    rkd_fm_loss_weight: 1.0
+    rkd_loss_weight: 0.0
+    rkd_action_temp: ${sweep.rkd_temp}
+    rkd_vlm_temp: ${sweep.rkd_temp}
+    rkd_relation_mode: token_pair_mean