# Held-out Episode Validation Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Add optional held-out episode validation to main training flow, including explicit val episode selection and periodic action MSE evaluation without merging LEWM-only model features.

**Architecture:** Extend the generic dataset with optional episode filtering metadata, extend train config with explicit held-out validation knobs, and wire train_vla to choose between random split and explicit episode split. Reuse agent.predict_action_chunk for action MSE so the metric stays model-agnostic.

**Tech Stack:** Python, Hydra/OmegaConf, PyTorch, unittest

---

### Task 1: Add failing tests for dataset episode filtering and config semantics

**Files:**
- Modify: `tests/test_simple_robot_dataset_image_loading.py`
- Modify: `tests/test_train_vla_rollout_validation.py`

- [ ] Add dataset tests for `episode_indices` and `available_episode_indices`.
- [ ] Add training tests for explicit held-out episode splitting and fail-fast config validation.
- [ ] Run focused tests and verify they fail for the expected missing behavior.

### Task 2: Implement minimal dataset and training support

**Files:**
- Modify: `roboimi/vla/data/simpe_robot_dataset.py`
- Modify: `roboimi/vla/conf/config.yaml`
- Modify: `roboimi/demos/vla_scripts/train_vla.py`

- [ ] Add config keys `train.val_episode_indices` and `train.action_mse_val_freq_epochs`.
- [ ] Add optional dataset filtering by episode index plus `available_episode_indices` metadata.
- [ ] Add explicit train/val dataset builder and held-out action MSE computation.
- [ ] Log `val/action_mse` only when explicit held-out episode validation is configured.

### Task 3: Verify focused coverage

**Files:**
- Test: `tests/test_simple_robot_dataset_image_loading.py`
- Test: `tests/test_train_vla_rollout_validation.py`
- Test: `tests/test_train_vla_swanlab_logging.py`

- [ ] Run focused unittest targets for dataset filtering, held-out MSE, and SwanLab logging.
- [ ] Fix any regressions with minimal code changes.