config(flow-transfer): prepare new-data H16 cotrain

[ElmoPA]

No description provided.

[ElmoPA]

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

• config(flow-transfer): prepare new-data H16 cotrain #592 : 4 dependent PRs (#593 , #594 , #595 and 1 other) 👈 (View in Graphite)
• feat(pipeline): add U-socket arc-length policy #589 : 1 other dependent PR (#590 )
• feat(pipeline): port arc-length action tokens #585 : 1 other dependent PR (#586 )
• feat(pipeline): consolidate Fold, PushT, overlays, and MoE #579 : 2 other dependent PRs (#580 , #583 )
• main

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

[github-actions]

Claude Code Review

Review

Summary

Adds new-data H16 cotrain configs (DP + latent) for Flow Transfer combining U-Socket + Chain (base + gen), with step-level metrics support in ModelWrapper, and pinned smoke/norm-precompute Slurm launchers. Config-only PR plus one small Lightning behavior change.

Key concerns

1. _log_train_metric swallows metrics when both flags are off. The new early return if not self.train_metrics_on_step and not self.train_metrics_on_epoch: return means turning both off silently drops all Timing/Train/Optimizer/grad-norm logs. Existing behavior (both on_step=False, on_epoch=True) was equivalent to train_metrics_on_epoch=True, which is the default — so defaults are preserved. But this is a footgun: someone passing train_metrics_on_epoch=False alone (without intending to disable step logging) will lose all metrics. Consider asserting at least one is true, or documenting explicitly.

2. Metric key rename risk under train_metrics_on_step=True. When step logging is on, epoch-aggregated metrics are renamed with _epoch suffix (Train/Loss_epoch, Optimizer/param_group_0_lr_epoch, etc.). Any existing W&B dashboards, alerting, or downstream analysis keyed on the old names (Train/Loss, Optimizer/param_group_0_lr) will silently break for runs that opt in. The verifier explicitly filters _epoch suffix, so it's aware — but confirm this doesn't break existing plotting/eval tooling.

3. on_save_checkpoint overwrites hyper_parameters. Directly mutating checkpoint["hyper_parameters"] bypasses Lightning's save_hyperparameters() mechanism. If save_hyperparameters() is used elsewhere in ModelWrapper (I don't see it here, but check), this could clobber values on resume. Lower risk since only three keys are set, but prefer overriding via self.save_hyperparameters(...) in __init__ if these need to round-trip.

4. Norm stats reused across DP and latent arms via byte-equality check. The norm-precompute script asserts dp.data == latent.data after resolution and reuses one artifact for both. This is correct only if both experiments resolve to the identical data tree — the assertion enforces that, which is good. But note that if either experiment config diverges later (e.g., different transform, different action_horizon), the shared artifact silently becomes invalid. Consider adding a norm-stats hash pin to each experiment config, not just the launcher.

5. chain_gen_effective filter contract. The filter lambda row.get('episode_hash') != 'episode_T_chain_gripper_obs7_000050' excludes by episode_hash. Verify that LocalEpisodeResolverManyWithEmbodimentOverride populates episode_hash from the directory name without .zarr suffix — the comment says it does, but this should be covered by a unit test since a rename or resolver change would silently un-exclude the audited idle-heavy episode. The smoke launcher's frame-count check would catch it, but a direct filter test would fail earlier.

6. limit_val_batches: 0 in experiment configs vs. 1 in smoke override. Production runs have limit_val_batches=0 (no validation), which means val_check_interval: 10000 is dead config. Confirm this is intentional (offline eval only via HumanRobotOverlayEval?). If validation is meant to run, this needs to be nonzero.

Suggestions

• Add an assertion or warning if both train_metrics_on_step and train_metrics_on_epoch are False.
• Add a unit test that the chain filter lambda correctly excludes episode_T_chain_gripper_obs7_000050 given a fake resolver output row.
• Consider pinning norm_stats.precomputed_norm_path SHA in the experiment configs (as a comment or in the launcher contract already, which you do).
• Confirm existing W&B dashboards for pushshapes-flow-transfer project handle the _epoch suffix rename.
• The verifier's new expected_world_size=1 default preserves backward compat for existing smoke callers ✓.

Verdict: Comment

Configs and launchers look carefully engineered with strong provenance/pinning. The pl_model.py metric-logging changes are the main non-config surface — they preserve default behavior but introduce metric-key rename semantics under opt-in that deserve a quick sanity check against downstream tooling before merging.

---

Reviewed by Claude · Review workflow