feat: 添加保存模型的功能和推理脚本

roboimi/demos/vla_scripts/train_vla.py

@@ -95,6 +95,17 @@ def main(cfg: DictConfig):
 
     log.info("✅ Training Loop with Real HDF5 Finished!")
 
+# --- 6. Save Checkpoint ---
+    save_dir = "checkpoints"
+    os.makedirs(save_dir, exist_ok=True)
+    save_path = os.path.join(save_dir, "vla_model_final.pt")
+    
+    # 保存整个 Agent 的 state_dict
+    torch.save(agent.state_dict(), save_path)
+    log.info(f"💾 Model saved to {save_path}")
+
+    log.info("✅ Training Loop Finished!")
+
 def recursive_to_device(data, device):
     if isinstance(data, torch.Tensor):
         return data.to(device)