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
• feat(abc): ABC-130k mcap to zarr pipeline + yam_bimanual embodiment #606
• feat(val): nested validation dataset groups #605 👈 (View in Graphite)
• refactoring aciton transforms #604
• main

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

[github-actions]

Claude Code Review

Review: PR #605 — Nested validation dataset groups

Summary

Adds nested {group: {embodiment: dataset}} support for valid_datasets, with per-group dataloaders, metric namespacing, episode-boundary video cutting, and a norm-stats-wiring fix that closes a real correctness bug. Overall a careful, well-commented change.

Key concerns

1. Sequential CombinedLoader unwrap assumption is fragile. _unwrap_combined_loader_batch unwraps any 3-tuple in a while loop. If Lightning ever changes its outer wrapping (or an inner batch happens to itself be a length-3 tuple — unlikely for dicts, but possible for some collate paths), this over-unwraps silently. Consider gating on isinstance(batch[0], dict) or checking len(loaders) > 1 state, and asserting the final type is a dict before returning.

2. Episode-boundary cutting relies on ordering that is not enforced. The docstring says "MultiDataset lays its index map out episode by episode … samples arrive grouped by episode and in frame order." This is true for the current MultiDataset, but there is no assertion or test that shuffle is off and ordering holds. If someone flips shuffle: true in a group's valid_dataloader_params, _buffer_per_episode will produce garbled videos — episodes will alternate and each hash change will flush a partial file, then the second occurrence will be dropped as "already written." At minimum, warn or assert shuffle=False on val loaders when using the episode path, or detect out-of-order hashes and fall back to chunked.

3. _is_embodiment_name swallows too much. Only KeyError/AttributeError are caught; if get_embodiment_id ever raises ValueError (common for enum-style lookups), the shape-detection code will crash instead of routing to the grouped branch. Suggest broadening the except, or asking the embodiment module for a canonical is_known(name) predicate.

4. sync_dist=True with per-group namespaced metrics. Different DDP ranks iterate the same group index at the same step, so this is fine, but note: if a group has fewer batches on some ranks (uneven sharding), the log call still fires. Not new to this PR, worth flagging.

5. max_episode_frames truncation is silent. An episode longer than the cap is flushed and then all subsequent samples with that hash are dropped (they hit the episode in written branch). Worth a log.warning the first time this happens per epoch — a silently truncated episode looks identical to a correctly-written short one.

Suggestions

• Add a unit test for _as_valid_groups covering: flat, grouped, empty, mixed-keys (raises), non-mapping values (raises), and the None-dataset filter.
• Add a small integration test for MultiDataModuleWrapper.val_dataloader() with 1 vs 2 groups, verifying the return type (bare vs list) and valid_group_names ordering.
• Test the norm-stats sentinel: construct a datamodule where a val group dataset is not wired, confirm trainHydra raises.
• In _buffer_per_episode, assert monotonic hash grouping (open_episode never re-appears after being written within a single epoch/group) to catch shuffle-on regressions loudly.
• Consider logging (once) when the chunked fallback path is taken — episode-hash absence in a batch usually indicates an upstream data-pipeline regression, not intentional legacy usage.
• Minor: _namespace_metrics doesn't handle Valid_epoch/... or other Valid-prefixed-but-not-Valid/ keys — probably fine, but worth a comment that this is intentional.

Verdict

Comment — leaning approve. The norm-stats fix alone justifies landing this; the correctness of the multi-group path is sound. Please address (1) and (2) before merge — the unwrap loop should be tightened, and the episode-cutting path needs either an ordering assertion or explicit documentation that val groups must run with shuffle=False.

---

Reviewed by Claude · Review workflow