feat(vla): add SmolVLA conditioning and experiment artifacts

roboimi/demos/vla_scripts/train_vla.py

@@ -118,6 +118,109 @@ def recursive_to_device(data, device):
     return data
 
 
+def build_agent_input(batch_data):
+    images = {}
+    for key, value in batch_data.items():
+        if key.startswith('observation.') and key != 'observation.state':
+            images[key.replace('observation.', '')] = value
+
+    agent_input = {
+        'images': images,
+        'qpos': batch_data['observation.state'],
+        'action': batch_data['action'],
+    }
+    if 'action_is_pad' in batch_data:
+        agent_input['action_is_pad'] = batch_data['action_is_pad']
+    if 'task' in batch_data:
+        agent_input['task'] = batch_data['task']
+    return agent_input
+
+
+def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None):
+    kwargs = {'image_resize_shape': dataset_image_resize_shape}
+    if episode_indices is not None:
+        kwargs['episode_indices'] = episode_indices
+    return instantiate(cfg.data, **kwargs)
+
+
+def build_train_val_datasets(cfg, dataset_image_resize_shape):
+    val_episode_indices = cfg.train.get('val_episode_indices', None)
+    if val_episode_indices:
+        dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
+        available_episode_indices = list(getattr(dataset, 'available_episode_indices', []))
+        if not available_episode_indices:
+            raise ValueError('显式 val_episode_indices 需要数据集暴露 available_episode_indices')
+        requested_val_episode_indices = sorted(int(idx) for idx in val_episode_indices)
+        available_set = set(available_episode_indices)
+        missing = sorted(set(requested_val_episode_indices) - available_set)
+        if missing:
+            raise ValueError(
+                f'val_episode_indices {missing} 不存在于数据集可用 episodes {available_episode_indices}'
+            )
+        val_set = set(requested_val_episode_indices)
+        train_episode_indices = [
+            idx for idx in available_episode_indices
+            if idx not in val_set
+        ]
+        if not train_episode_indices:
+            raise ValueError('显式 val_episode_indices 不能覆盖全部 episodes，训练集将为空')
+
+        train_dataset = _instantiate_dataset(
+            cfg,
+            dataset_image_resize_shape,
+            episode_indices=train_episode_indices,
+        )
+        val_dataset = _instantiate_dataset(
+            cfg,
+            dataset_image_resize_shape,
+            episode_indices=requested_val_episode_indices,
+        )
+        return dataset, train_dataset, val_dataset, requested_val_episode_indices
+
+    dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
+    val_split = float(cfg.train.get('val_split', 0.1))
+    seed = int(cfg.train.get('seed', 42))
+    val_size = int(len(dataset) * val_split)
+    train_size = len(dataset) - val_size
+    if val_size > 0:
+        train_dataset, val_dataset = random_split(
+            dataset,
+            [train_size, val_size],
+            generator=torch.Generator().manual_seed(seed)
+        )
+    else:
+        train_dataset, val_dataset = dataset, None
+    return dataset, train_dataset, val_dataset, None
+
+
+def compute_action_mse_validation(agent, val_loader, device):
+    if val_loader is None:
+        return None
+
+    was_training = agent.training
+    agent.eval()
+    total_squared_error = 0.0
+    total_count = 0.0
+    with torch.no_grad():
+        for val_batch in val_loader:
+            val_batch = recursive_to_device(val_batch, device)
+            val_input = build_agent_input(val_batch)
+            pred_actions = agent.predict_action_chunk(val_input)
+            target_actions = val_input['action']
+            squared_error = (pred_actions - target_actions).pow(2)
+            action_is_pad = val_input.get('action_is_pad', None)
+            if action_is_pad is not None:
+                mask = (~action_is_pad).unsqueeze(-1).to(squared_error.dtype)
+                total_squared_error += (squared_error * mask).sum().item()
+                total_count += mask.sum().item() * squared_error.shape[-1]
+            else:
+                total_squared_error += squared_error.sum().item()
+                total_count += target_actions.numel()
+    if was_training:
+        agent.train()
+    return total_squared_error / max(total_count, 1.0)
+
+
 def resolve_resume_checkpoint(resume_ckpt, checkpoint_dir):
     """
     解析恢复训练用的 checkpoint 路径。
@@ -369,6 +472,11 @@ def _run_training(cfg: DictConfig):
     _configure_cuda_runtime(cfg)
     swanlab_module = _init_swanlab(cfg)
     try:
+        action_mse_val_freq_epochs = int(cfg.train.get('action_mse_val_freq_epochs', 0) or 0)
+        explicit_val_episode_indices = cfg.train.get('val_episode_indices', None)
+        if action_mse_val_freq_epochs > 0 and not explicit_val_episode_indices:
+            raise ValueError('action_mse_val_freq_epochs > 0 requires train.val_episode_indices')
+
         # 创建检查点目录
         run_output_dir = _resolve_run_output_dir()
         checkpoint_dir = run_output_dir / "checkpoints"
@@ -381,33 +489,35 @@ def _run_training(cfg: DictConfig):
         log.info("📦 加载数据集...")
         try:
             dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
-            vision_backbone_cfg = cfg.agent.get('vision_backbone', None)
-            if vision_backbone_cfg is not None and 'dataset_image_resize_shape' in vision_backbone_cfg:
-                dataset_image_resize_shape = vision_backbone_cfg.get('dataset_image_resize_shape')
-            dataset = instantiate(
-                cfg.data,
-                image_resize_shape=dataset_image_resize_shape,
+            for backbone_key in ('vision_backbone', 'condition_encoder'):
+                backbone_cfg = cfg.agent.get(backbone_key, None)
+                if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg:
+                    dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape')
+                    break
+            dataset, train_dataset, val_dataset, explicit_val_episode_indices = (
+                build_train_val_datasets(cfg, dataset_image_resize_shape)
             )
             log.info(f"✅ 数据集加载成功。总样本数: {len(dataset)}")
         except Exception as e:
             log.error(f"❌ 数据集加载失败: {e}")
             raise
 
-        # 训练/验证集划分
-        val_split = float(cfg.train.get('val_split', 0.1))
-        seed = int(cfg.train.get('seed', 42))
-        val_size = int(len(dataset) * val_split)
-        train_size = len(dataset) - val_size
-        if val_size > 0:
-            train_dataset, val_dataset = random_split(
-                dataset,
-                [train_size, val_size],
-                generator=torch.Generator().manual_seed(seed)
+        if explicit_val_episode_indices is not None:
+            log.info(
+                "✅ 数据集划分: 训练集=%s, 验证集=%s (显式 held-out episodes=%s)",
+                len(train_dataset),
+                len(val_dataset),
+                explicit_val_episode_indices,
             )
-            log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})")
         else:
-            train_dataset, val_dataset = dataset, None
-            log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")
+            val_split = float(cfg.train.get('val_split', 0.1))
+            val_size = len(val_dataset) if val_dataset is not None else 0
+            if val_size > 0:
+                log.info(
+                    f"✅ 数据集划分: 训练集={len(train_dataset)}, 验证集={val_size} (验证比例={val_split})"
+                )
+            else:
+                log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")
 
         train_batch_size = int(cfg.train.batch_size)
         train_drop_last = len(train_dataset) >= train_batch_size
@@ -643,22 +753,6 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         log.info("🏋️ 开始训练循环...")
 
-        def build_agent_input(batch_data):
-            """构建 agent 输入格式"""
-            images = {}
-            # SimpleRobotDataset 返回 observation.{cam_name} 格式
-            for cam_name in cfg.data.camera_names:
-                key = f"observation.{cam_name}"
-                if key in batch_data:
-                    images[cam_name] = batch_data[key]
-
-            return {
-                'images': images,
-                'qpos': batch_data['observation.state'],  # SimpleRobotDataset 使用 observation.state
-                'action': batch_data['action'],
-                'action_is_pad': batch_data.get('action_is_pad', None)  # 传递padding mask
-            }
-
         def save_checkpoint(checkpoint_path: Path, step: int, loss_value, val_loss=None, rollout_avg_reward=None):
             agent_stats = agent.get_normalization_stats()
             torch.save({
@@ -753,6 +847,9 @@ def _run_training(cfg: DictConfig):
         steps_per_epoch = len(train_loader)
         rollout_val_freq_epochs = int(cfg.train.get('rollout_val_freq_epochs', 0) or 0)
         rollout_validation_enabled = rollout_val_freq_epochs > 0
+        action_mse_validation_enabled = (
+            explicit_val_episode_indices is not None and action_mse_val_freq_epochs > 0
+        )
         best_loss = resume_best_loss
         best_rollout_reward = resume_best_rollout_reward
         last_loss = resume_loss
@@ -911,6 +1008,26 @@ def _run_training(cfg: DictConfig):
                 and completed_epoch > 0
                 and completed_epoch % rollout_val_freq_epochs == 0
             )
+            should_run_action_mse_validation = (
+                action_mse_validation_enabled
+                and steps_per_epoch > 0
+                and completed_steps % steps_per_epoch == 0
+                and completed_epoch > 0
+                and completed_epoch % action_mse_val_freq_epochs == 0
+            )
+            if should_run_action_mse_validation:
+                action_mse = compute_action_mse_validation(agent, val_loader, cfg.train.device)
+                if action_mse is not None:
+                    log.info(
+                        f"步骤 {step}/{cfg.train.max_steps} | Epoch {completed_epoch} held-out action MSE: "
+                        f"{action_mse:.6f}"
+                    )
+                    _log_to_swanlab(
+                        swanlab_module,
+                        {'val/action_mse': action_mse},
+                        step=step,
+                    )
+
             if should_run_epoch_rollout:
                 if checkpoint_path is None:
                     checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt"