refactor：大重构

roboimi/demos/vla_scripts/train_vla.py

@@ -12,28 +12,28 @@ from torch.optim import AdamW
 from torch.optim.lr_scheduler import LambdaLR
 from pathlib import Path
 
-# Ensure correct import path
+# 确保正确的导入路径
 sys.path.append(os.getcwd())
 
 from hydra.utils import instantiate
 
 log = logging.getLogger(__name__)
 
-# Register resolver for list length in configs (e.g., ${len:${data.camera_names}})
+# 注册列表长度解析器（用于配置中如 ${len:${data.camera_names}}）
 if not OmegaConf.has_resolver("len"):
     OmegaConf.register_new_resolver("len", lambda x: len(x))
 
 
 def recursive_to_device(data, device):
     """
-    Recursively move nested dictionaries/lists of tensors to specified device.
+    递归地将嵌套字典/列表中的张量移动到指定设备。
 
     Args:
-        data: Dictionary, list, or tensor
-        device: Target device (e.g., 'cuda', 'cpu')
+        data: 字典、列表或张量
+        device: 目标设备 (例如 'cuda', 'cpu')
 
     Returns:
-        Data structure with all tensors moved to device
+        所有张量已移动到指定设备的数据结构
     """
     if isinstance(data, torch.Tensor):
         return data.to(device)
@@ -46,36 +46,36 @@ def recursive_to_device(data, device):
 
 def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_type='cosine', min_lr=0):
     """
-    Create a learning rate scheduler with warmup.
+    创建带预热的学习率调度器。
 
     Args:
-        optimizer: PyTorch optimizer
-        warmup_steps: Number of warmup steps
-        max_steps: Total training steps
-        scheduler_type: Type of scheduler after warmup ('cosine' or 'constant')
-        min_lr: Minimum learning rate (for cosine decay)
+        optimizer: PyTorch 优化器
+        warmup_steps: 预热步数
+        max_steps: 总训练步数
+        scheduler_type: 预热后的调度器类型 ('cosine' 或 'constant')
+        min_lr: 最小学习率（用于余弦衰减）
 
     Returns:
-        LambdaLR scheduler
+        LambdaLR 调度器
     """
     import math
-    # Capture initial lr before LambdaLR modifies it
+    # 在 LambdaLR 修改前捕获初始学习率
     base_lr = optimizer.param_groups[0]['lr']
     min_lr_ratio = min_lr / base_lr if base_lr > 0 else 0.0
 
     def lr_lambda(step):
-        # Warmup phase: linear increase from 0 to 1
+        # 预热阶段：从 0 线性增加到 1
         if step < warmup_steps:
             return float(step) / float(max(1, warmup_steps))
 
-        # Post-warmup phase
+        # 预热后阶段
         if scheduler_type == 'cosine':
-            # Cosine annealing from 1 to min_lr_ratio
+            # 从 1 到 min_lr_ratio 的余弦退火
             progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
             cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
             return max(min_lr_ratio, cosine_decay)
         else:
-            # Constant learning rate
+            # 恒定学习率
             return 1.0
 
     return LambdaLR(optimizer, lr_lambda)
@@ -84,40 +84,40 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
 @hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
 def main(cfg: DictConfig):
     """
-    VLA Training Script with ResNet Backbone and Diffusion Policy.
+    VLA 训练脚本（ResNet 骨干网络 + Diffusion 策略）
 
-    This script:
-    1. Loads dataset from HDF5 files
-    2. Instantiates VLAAgent with ResNet vision encoder
-    3. Trains diffusion-based action prediction
-    4. Saves checkpoints periodically
+    该脚本功能：
+    1. 从 HDF5 文件加载数据集
+    2. 实例化带 ResNet 视觉编码器的 VLAAgent
+    3. 训练基于扩散的动作预测模型
+    4. 定期保存检查点
     """
 
-    # Print configuration
+    # 打印配置
     print("=" * 80)
-    print("VLA Training Configuration:")
+    print("VLA 训练配置:")
     print("=" * 80)
     print(OmegaConf.to_yaml(cfg))
     print("=" * 80)
 
-    log.info(f"🚀 Starting VLA Training (Device: {cfg.train.device})")
+    log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})")
 
-    # Create checkpoint directory
+    # 创建检查点目录
     checkpoint_dir = Path("checkpoints")
     checkpoint_dir.mkdir(exist_ok=True)
 
     # =========================================================================
-    # 1. Instantiate Dataset & DataLoader
+    # 1. 实例化数据集与 DataLoader
     # =========================================================================
-    log.info("📦 Loading dataset...")
+    log.info("📦 加载数据集...")
     try:
         dataset = instantiate(cfg.data)
-        log.info(f"✅ Dataset loaded successfully. Total samples: {len(dataset)}")
+        log.info(f"✅ 数据集加载成功。总样本数: {len(dataset)}")
     except Exception as e:
-        log.error(f"❌ Failed to load dataset: {e}")
+        log.error(f"❌ 数据集加载失败: {e}")
         raise
 
-    # Train/Val split
+    # 训练/验证集划分
     val_split = float(cfg.train.get('val_split', 0.1))
     seed = int(cfg.train.get('seed', 42))
     val_size = int(len(dataset) * val_split)
@@ -128,10 +128,10 @@ def main(cfg: DictConfig):
             [train_size, val_size],
             generator=torch.Generator().manual_seed(seed)
         )
-        log.info(f"✅ Dataset split: train={train_size}, val={val_size} (val_split={val_split})")
+        log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})")
     else:
         train_dataset, val_dataset = dataset, None
-        log.info("✅ Dataset split: train=all, val=0 (val_split=0)")
+        log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")
 
     train_loader = DataLoader(
         train_dataset,
@@ -139,7 +139,7 @@ def main(cfg: DictConfig):
         shuffle=True,
         num_workers=cfg.train.num_workers,
         pin_memory=(cfg.train.device != "cpu"),
-        drop_last=True  # Drop incomplete batches for stable training
+        drop_last=True  # 丢弃不完整批次以稳定训练
     )
 
     val_loader = None
@@ -153,34 +153,14 @@ def main(cfg: DictConfig):
             drop_last=False
         )
 
-    log.info(f"✅ Train loader batches per epoch: {len(train_loader)}")
+    log.info(f"✅ 训练加载器每轮批次数: {len(train_loader)}")
     if val_loader is not None:
-        log.info(f"✅ Val loader batches per epoch: {len(val_loader)}")
+        log.info(f"✅ 验证加载器每轮批次数: {len(val_loader)}")
 
     # =========================================================================
-    # 2. Instantiate VLA Agent
+    # 2. 加载数据集统计信息（将传递给 agent）
     # =========================================================================
-    log.info("🤖 Initializing VLA Agent...")
-    try:
-        agent = instantiate(cfg.agent)
-        agent.to(cfg.train.device)
-        agent.train()
-        log.info(f"✅ Agent initialized and moved to {cfg.train.device}")
-
-        # Count parameters
-        total_params = sum(p.numel() for p in agent.parameters())
-        trainable_params = sum(p.numel() for p in agent.parameters() if p.requires_grad)
-        log.info(f"📊 Total parameters: {total_params:,}")
-        log.info(f"📊 Trainable parameters: {trainable_params:,}")
-
-    except Exception as e:
-        log.error(f"❌ Failed to initialize agent: {e}")
-        raise
-
-    # =========================================================================
-    # 2.5. Load Dataset Statistics (will be saved into checkpoints)
-    # =========================================================================
-    log.info("💾 Loading dataset statistics...")
+    log.info("💾 加载数据集统计信息...")
     dataset_stats = None
     try:
         dataset_dir = cfg.data.get('dataset_dir', 'roboimi/demos/dataset/sim_transfer')
@@ -201,22 +181,43 @@ def main(cfg: DictConfig):
                 'qpos_min': stats['qpos']['min'].tolist(),
                 'qpos_max': stats['qpos']['max'].tolist(),
             }
-            log.info(f"✅ Dataset statistics loaded (normalization: {dataset_stats['normalization_type']})")
+            log.info(f"✅ 数据集统计信息加载完成 (归一化: {dataset_stats['normalization_type']})")
         else:
-            log.warning(f"⚠️  Statistics file not found: {stats_path}")
-            log.warning("⚠️  Actions will not be denormalized during inference!")
+            log.warning(f"⚠️  统计文件未找到: {stats_path}")
+            log.warning("⚠️  推理时动作将无法反归一化！")
 
     except Exception as e:
-        log.warning(f"⚠️  Failed to load statistics: {e}")
-        log.warning("⚠️  Training will continue, but inference may not work correctly")
+        log.warning(f"⚠️  统计信息加载失败: {e}")
+        log.warning("⚠️  训练将继续，但推理可能无法正常工作")
 
     # =========================================================================
-    # 3. Setup Optimizer & LR Scheduler
+    # 3. 实例化 VLA Agent
+    # =========================================================================
+    log.info("🤖 初始化 VLA Agent...")
+    try:
+        # 将 dataset_stats 和 normalization_type 传递给 agent
+        agent = instantiate(cfg.agent, dataset_stats=dataset_stats)
+        agent.to(cfg.train.device)
+        agent.train()
+        log.info(f"✅ Agent 初始化完成并已移至 {cfg.train.device}")
+
+        # 统计参数量
+        total_params = sum(p.numel() for p in agent.parameters())
+        trainable_params = sum(p.numel() for p in agent.parameters() if p.requires_grad)
+        log.info(f"📊 总参数量: {total_params:,}")
+        log.info(f"📊 可训练参数量: {trainable_params:,}")
+
+    except Exception as e:
+        log.error(f"❌ Agent 初始化失败: {e}")
+        raise
+
+    # =========================================================================
+    # 4. 设置优化器与学习率调度器
     # =========================================================================
     optimizer = AdamW(agent.parameters(), lr=cfg.train.lr, weight_decay=1e-5)
-    log.info(f"🔧 Optimizer: AdamW (lr={cfg.train.lr})")
+    log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr})")
 
-    # Setup learning rate scheduler with warmup
+    # 设置带预热的学習率调度器
     warmup_steps = int(cfg.train.get('warmup_steps', 500))
     scheduler_type = cfg.train.get('scheduler_type', 'cosine')
     min_lr = float(cfg.train.get('min_lr', 1e-6))
@@ -228,33 +229,36 @@ def main(cfg: DictConfig):
         scheduler_type=scheduler_type,
         min_lr=min_lr
     )
-    log.info(f"📈 LR Scheduler: {scheduler_type} with {warmup_steps} warmup steps (min_lr={min_lr})")
+    log.info(f"📈 学习率调度器: {scheduler_type}，{warmup_steps} 步预热 (最小学习率={min_lr})")
 
     # =========================================================================
-    # 4. Training Loop
+    # 5. 训练循环
     # =========================================================================
-    log.info("🏋️ Starting training loop...")
+    log.info("🏋️ 开始训练循环...")
 
     def build_agent_input(batch_data):
+        """构建 agent 输入格式"""
         images = {}
+        # SimpleRobotDataset 返回 observation.{cam_name} 格式
         for cam_name in cfg.data.camera_names:
-            key = f"image_{cam_name}"
+            key = f"observation.{cam_name}"
             if key in batch_data:
                 images[cam_name] = batch_data[key]
 
         return {
             'images': images,
-            'qpos': batch_data['qpos'],
+            'qpos': batch_data['observation.state'],  # SimpleRobotDataset 使用 observation.state
             'action': batch_data['action']
         }
 
     def run_validation():
+        """运行验证"""
         if val_loader is None:
             return None
         agent.eval()
 
-        # 🔧 FIX: Set deterministic seed for validation to get reproducible loss
-        # This ensures validation loss is comparable across different steps
+        # 设置确定性种子以获得可重现的损失
+        # 这确保验证损失在不同步骤之间可比较
         torch.manual_seed(42)
         if torch.cuda.is_available():
             torch.cuda.manual_seed(42)
@@ -272,7 +276,7 @@ def main(cfg: DictConfig):
         return total_loss / max(num_batches, 1)
 
     data_iter = iter(train_loader)
-    pbar = tqdm(range(cfg.train.max_steps), desc="Training", ncols=100)
+    pbar = tqdm(range(cfg.train.max_steps), desc="训练中", ncols=100)
 
     best_loss = float('inf')
 
@@ -280,47 +284,47 @@ def main(cfg: DictConfig):
         try:
             batch = next(data_iter)
         except StopIteration:
-            # Restart iterator when epoch ends
+            # 轮次结束时重启迭代器
             data_iter = iter(train_loader)
             batch = next(data_iter)
 
         # =====================================================================
-        # Move batch to device
+        # 将批次移至设备
         # =====================================================================
         batch = recursive_to_device(batch, cfg.train.device)
 
         # =====================================================================
-        # Prepare agent input
+        # 准备 agent 输入
         # =====================================================================
-        # Dataset returns: {action, qpos, image_<cam_name>, ...}
-        # Agent expects: {images: dict, qpos: tensor, action: tensor}
+        # 数据集返回: {action, qpos, image_<cam_name>, ...}
+        # Agent 期望: {images: dict, qpos: tensor, action: tensor}
 
-        # Prepare agent input
+        # 准备 agent 输入
         agent_input = build_agent_input(batch)
 
         # =====================================================================
-        # Forward pass & compute loss
+        # 前向传播与损失计算
         # =====================================================================
         try:
             loss = agent.compute_loss(agent_input)
         except Exception as e:
-            log.error(f"❌ Forward pass failed at step {step}: {e}")
+            log.error(f"❌ 步骤 {step} 前向传播失败: {e}")
             raise
 
         # =====================================================================
-        # Backward pass & optimization
+        # 反向传播与优化
         # =====================================================================
         optimizer.zero_grad()
         loss.backward()
 
-        # Gradient clipping for stable training
+        # 梯度裁剪以稳定训练
         torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm=1.0)
 
         optimizer.step()
         scheduler.step()
 
         # =====================================================================
-        # Logging
+        # 日志记录
         # =====================================================================
         if step % cfg.train.log_freq == 0:
             current_lr = optimizer.param_groups[0]['lr']
@@ -329,16 +333,16 @@ def main(cfg: DictConfig):
                 "lr": f"{current_lr:.2e}",
                 "best_loss": f"{best_loss:.4f}"
             })
-            log.info(f"Step {step}/{cfg.train.max_steps} | Loss: {loss.item():.4f} | LR: {current_lr:.2e}")
+            log.info(f"步骤 {step}/{cfg.train.max_steps} | 损失: {loss.item():.4f} | 学习率: {current_lr:.2e}")
 
         # =====================================================================
-        # Checkpoint saving & Validation
+        # 检查点保存与验证
         # =====================================================================
         if step > 0 and step % cfg.train.save_freq == 0:
-            # Run validation
+            # 运行验证
             val_loss = run_validation()
             if val_loss is not None:
-                log.info(f"Step {step}/{cfg.train.max_steps} | Val Loss: {val_loss:.4f}")
+                log.info(f"步骤 {step}/{cfg.train.max_steps} | 验证损失: {val_loss:.4f}")
 
             checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt"
             torch.save({
@@ -351,9 +355,9 @@ def main(cfg: DictConfig):
                 'dataset_stats': dataset_stats,
                 'current_lr': optimizer.param_groups[0]['lr'],
             }, checkpoint_path)
-            log.info(f"💾 Checkpoint saved: {checkpoint_path}")
+            log.info(f"💾 检查点已保存: {checkpoint_path}")
 
-            # Save best model based on validation loss
+            # 根据验证损失保存最佳模型
             eval_loss = val_loss if val_loss is not None else loss.item()
             if eval_loss < best_loss:
                 best_loss = eval_loss
@@ -368,10 +372,10 @@ def main(cfg: DictConfig):
                     'dataset_stats': dataset_stats,
                     'current_lr': optimizer.param_groups[0]['lr'],
                 }, best_model_path)
-                log.info(f"🌟 Best model updated: {best_model_path} (val_loss: {best_loss:.4f})")
+                log.info(f"🌟 最佳模型已更新: {best_model_path} (验证损失: {best_loss:.4f})")
 
     # =========================================================================
-    # 5. Save Final Model
+    # 6. 保存最终模型
     # =========================================================================
     final_model_path = checkpoint_dir / "vla_model_final.pt"
     torch.save({
@@ -383,11 +387,11 @@ def main(cfg: DictConfig):
         'dataset_stats': dataset_stats,
         'current_lr': optimizer.param_groups[0]['lr'],
     }, final_model_path)
-    log.info(f"💾 Final model saved: {final_model_path}")
+    log.info(f"💾 最终模型已保存: {final_model_path}")
 
-    log.info("✅ Training completed successfully!")
-    log.info(f"📊 Final Loss: {loss.item():.4f}")
-    log.info(f"📊 Best Loss: {best_loss:.4f}")
+    log.info("✅ 训练成功完成!")
+    log.info(f"📊 最终损失: {loss.item():.4f}")
+    log.info(f"📊 最佳损失: {best_loss:.4f}")
 
 
 if __name__ == "__main__":