roboimi/docs/superpowers/plans/2026-05-06-held-out-episode-validation.md

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.