feat: add IMF AttnRes policy training path

roboimi/demos/vla_scripts/train_vla.py

@@ -14,8 +14,17 @@ from torch.optim import AdamW
 from torch.optim.lr_scheduler import LambdaLR
 from pathlib import Path
 
-# 确保正确的导入路径
-sys.path.append(os.getcwd())
+# 确保正确的导入路径（不能依赖 cwd，因为 Hydra 会在运行时切换 cwd）
+def _ensure_repo_root_on_syspath():
+    repo_root = Path(__file__).resolve().parents[3]
+    repo_root_str = str(repo_root)
+    if repo_root_str in sys.path:
+        sys.path.remove(repo_root_str)
+    sys.path.insert(0, repo_root_str)
+    return repo_root
+
+
+_REPO_ROOT = _ensure_repo_root_on_syspath()
 
 from hydra.utils import instantiate
 
@@ -26,6 +35,13 @@ if not OmegaConf.has_resolver("len"):
     OmegaConf.register_new_resolver("len", lambda x: len(x))
 
 
+def _configure_cuda_runtime(cfg):
+    """Apply process-level CUDA runtime switches required by this environment."""
+    if str(cfg.train.device).startswith('cuda') and bool(cfg.train.get('disable_cudnn', False)):
+        torch.backends.cudnn.enabled = False
+        log.warning('⚠️  已按配置禁用 cuDNN；GPU 卷积将回退到非-cuDNN 实现')
+
+
 def recursive_to_device(data, device):
     """
     递归地将嵌套字典/列表中的张量移动到指定设备。
@@ -113,14 +129,11 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
 
 
 def build_training_optimizer(agent, lr, weight_decay):
-    """为训练脚本构建优化器，优先复用 transformer head 自带的参数分组。"""
+    """为训练脚本构建优化器，优先复用任意 head 自带的参数分组。"""
     trainable_params = [param for param in agent.parameters() if param.requires_grad]
     noise_pred_net = getattr(agent, 'noise_pred_net', None)
     get_optim_groups = getattr(noise_pred_net, 'get_optim_groups', None)
-    use_head_groups = (
-        getattr(agent, 'head_type', None) == 'transformer'
-        and callable(get_optim_groups)
-    )
+    use_head_groups = callable(get_optim_groups)
 
     if not use_head_groups:
         return AdamW(trainable_params, lr=lr, weight_decay=weight_decay)
@@ -138,7 +151,7 @@ def build_training_optimizer(agent, lr, weight_decay):
         for param in params:
             param_id = id(param)
             if param_id in grouped_param_ids:
-                raise ValueError('Transformer optimizer groups contain duplicate parameters')
+                raise ValueError('Head optimizer groups contain duplicate parameters')
             grouped_param_ids.add(param_id)
 
     head_trainable_param_ids = {
@@ -146,7 +159,7 @@ def build_training_optimizer(agent, lr, weight_decay):
     }
     missing_head_param_ids = head_trainable_param_ids - grouped_param_ids
     if missing_head_param_ids:
-        raise ValueError('Transformer optimizer groups missed trainable head parameters')
+        raise ValueError('Head optimizer groups missed trainable head parameters')
 
     remaining_params = [
         param for param in trainable_params
@@ -258,6 +271,7 @@ def _run_training(cfg: DictConfig):
     print("=" * 80)
 
     log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})")
+    _configure_cuda_runtime(cfg)
     swanlab_module = _init_swanlab(cfg)
     try:
         # 创建检查点目录