fix: support headless rollout on remote training hosts

roboimi/demos/vla_scripts/train_vla.py

@@ -76,6 +76,18 @@ if not OmegaConf.has_resolver("len"):
     OmegaConf.register_new_resolver("len", lambda x: len(x))
 
 
+def _resolve_run_output_dir() -> Path:
+    try:
+        from hydra.core.hydra_config import HydraConfig
+        if HydraConfig.initialized():
+            output_dir = HydraConfig.get().runtime.output_dir
+            if output_dir:
+                return Path(output_dir).resolve()
+    except Exception:
+        pass
+    return Path.cwd().resolve()
+
+
 _maybe_reexec_without_problematic_ld_preload()
 
 
@@ -319,8 +331,9 @@ def _run_training(cfg: DictConfig):
     swanlab_module = _init_swanlab(cfg)
     try:
         # 创建检查点目录
-        checkpoint_dir = Path("checkpoints")
-        checkpoint_dir.mkdir(exist_ok=True)
+        run_output_dir = _resolve_run_output_dir()
+        checkpoint_dir = run_output_dir / "checkpoints"
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
         default_best_model_path = checkpoint_dir / "vla_model_best.pt"
 
         # =========================================================================