feat(val): nested validation dataset groups

[AnikethCheluva]

valid_datasets now accepts {group: {embodiment: dataset}} as well as the old
flat {embodiment: dataset}. Each group gets its own dataloader, metric
namespace and video directory, so a held-out split can be evaluated alongside
the in-distribution one in a single run.

Videos are cut on episode boundaries and named by episode_hash instead of
fixed-size chunks, with a per-episode frame cap and a fallback to the old
chunking when a batch carries no episode_hash.

Norm stats are wired through MultiDataModuleWrapper.iter_valid_datasets() and
guarded. valid_datasets is a back-compat alias for a single group, so iterating
it left every other group unnormalised while the evaluator unnormalised it
anyway, silently corrupting that group's overlays and metrics.

Also unwraps the extra tuple Lightning leaks when CombinedLoaders nest, which
otherwise fails with 'tuple object has no attribute items'.

[AnikethCheluva]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• Aniketh/arc #573 : 2 dependent PRs (#610 , #611 )
• feat(abc): ABC-130k mcap to zarr pipeline + yam_bimanual embodiment #606
• feat: add HPT Qwen text stem #612
• feat(val): nested validation dataset groups #605 👈 (View in Graphite)
• refactoring aciton transforms #604 : 1 other dependent PR (#176 )
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Claude Code Review

Review

Summary

Adds nested {group: {embodiment: dataset}} support for valid_datasets, giving each val group its own dataloader, metric namespace, and video directory. Fixes a real norm-stats bug where iterating valid_datasets (a single-group alias) left other groups unnormalized, and handles Lightning's tuple-wrapping when returning multiple CombinedLoaders.

Key concerns

1. on_validation_start reset breaks multi-group runs. Lightning fires on_validation_start once per val loop (i.e., once per dataloader in the list), not once per epoch. Resetting val_written and val_image_buffer at the start of every group's val loop is fine, but it also means the previous group's final episode (still open, flushed only in on_validation_end) will never get written — because on_validation_end only fires after the last group. Please verify this: if on_validation_start/end fire per-loop, the last-episode flush for groups 0..N-2 is lost; if they fire once per epoch (as the code comments assume), this is fine. This is worth an explicit test with 2 groups and confirming both groups' final episodes land on disk.

2. Episode ordering assumption is load-bearing but undocumented at the data layer. _buffer_per_episode assumes samples arrive grouped by episode and in frame order (validation shuffle=False + MultiDataset index layout). If anyone ever adds a val-time sampler, sets shuffle=True, or changes MultiDataset's index construction, videos will silently interleave episodes and the episode in written guard will start dropping legitimate frames. Consider an assertion (e.g., detect episode-hash flip-flop within a group) or at least a # LOAD-BEARING comment on the relevant val dataloader params.

3. episode_hash provenance. The docstring claims episode_hash "survives process_batch_for_training" because it's not a registered zarr key. This is true for the current ACT/HPT path, but Pi 0.5 and any algo that re-collates or filters batch keys could drop it. Since the fallback silently reverts to chunked mode, a regression here is invisible except in the video filenames. Suggest logging a one-time warning per (group, key) when the fallback triggers.

4. _is_embodiment_name shape detection. Catching KeyError, AttributeError is narrow — if get_embodiment_id ever raises ValueError (common for enum lookups), all keys will be misclassified as group names and a flat config will be interpreted as grouped with garbage embodiment children, producing a confusing downstream error. Recommend except Exception here or explicitly documenting/testing which exceptions get_embodiment_id raises.

5. Test coverage. No tests included. Given this touches the norm-stats wiring (which the PR itself calls out as a source of a silent correctness bug), at minimum please add:

   • a unit test that _as_valid_groups handles both shapes and rejects mixed keys,
   • an integration-ish test that iter_valid_datasets visits every dataset for a grouped config,
   • the wiring assertion in trainHydra.py exercised on a 2-group config.

Suggestions

• Add an assertion in _buffer_per_episode that detects out-of-order episodes (open episode reappears mid-stream) and raises rather than silently dropping.
• The RuntimeError for unwired norm stats is great — consider adding it as a proper test since it's the guard against the exact class of bug this PR fixes.
• _unwrap_combined_loader_batch: the while loop can over-unwrap if a legitimate batch happens to be a 3-tuple. Tighten to a single if unwrap, or check that batch[1] and batch[2] are ints before stripping.
• Comment in val_dataloader that returning a bare loader for single-group is deliberate for back-compat with dataloader_idx=0 — currently the comment is there but worth cross-referencing from _unwrap_combined_loader_batch since these two behaviors are coupled.
• The f-string, not %-args comment is useful but this is the third+ place I've seen it — consider fixing the pylogger wrapper instead.

Verdict: Request Changes

Primarily to (a) confirm/fix the on_validation_start/end lifecycle for the multi-group case (concern #1 is potentially data-losing) and (b) add at least minimal tests for the norm-stats wiring and shape normalization, given the PR's own framing that silent norm-stats bugs already shipped once here.

---

Reviewed by Claude · Review workflow