feat(pipeline): consolidate Fold, PushT, overlays, and MoE

[ElmoPA]

No description provided.

[ElmoPA]

• feat(pipeline): consolidate Fold, PushT, overlays, and MoE #579 : 2 dependent PRs (#580 , #583 ) 👈 (View in Graphite)
• main

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

[github-actions]

Claude Code Review

Review of PR #579

Summary

Large PR consolidating Fold (RH cotrain, keypoints), PushT (circle + small_circle cotrain), teacher-forced overlay eval, and a shared-residual MoE denoiser variant on top of the Pipeline algo. Adds new configs, a new eval class, an MoE denoiser, and a stems/ module tree (obs encoders, VisualCore port).

Key concerns

1. Missing PR description. For a change this broad (new eval, new model variant, new data configs, new stems package, denoiser changes) the description shouldn't be empty. Please add: motivation, which experiments this unblocks, and what was validated (norm stats stability, existing configs still load, checkpoint compatibility for CrossTransformer).

2. CrossTransformer behavior change is checkpoint-incompatible when time_conditioning="additive". The proj_u and pos_emb shapes now depend on the new arg. Default remains "concat" so existing ckpts still load, but this is worth an explicit note in the PR body and ideally a test that loads a legacy state_dict. Also — the action_embedding_dim < act_dim guard is good, but consider whether the "concat" default with hidden_dim//2 < act_dim used to silently proceed for any existing config (i.e. does this raise on a config that used to work?).

3. HumanRobotOverlayEval._unnormalized_target shape check is too strict for packed batches. You raise if target.ndim != 3 only in the non-packed branch, but downstream _unnormalize_prediction unconditionally assumes (B, H, D). If a packed eval loader ever routes here, chunk_targets produces (B, chunk_len, D) — fine — but confirm predictions also come out (B, chunk_len, D) for packed inputs. Add a shape assert after chunking.

4. DDPSafeMoEFFN anchor trick. Summing parameter.reshape(-1)[0] * 0.0 across every expert param each forward to keep DDP happy works, but it's fragile and O(num_params). Consider find_unused_parameters=True on the DDP wrapper (documented cost), or static_graph=False with an explicit torch.zeros(1, requires_grad=True) anchor multiplied by a single sum. Also please add a comment explaining why this exists (top-k routing => some experts unused on a step => DDP hang) — a future reader will delete it.

5. skip_bounds_check: true on fold configs. The header comment justifies this well (quantile filter drops extreme-reach frames), but this contradicts the codebase convention of validated data loading. Please:

   • File a follow-up to fix MultiDataset.__getitem__ to skip rather than resample out-of-bounds frames.
   • Or at least log a count of what would have been filtered, so we know the tail we're keeping isn't dominated by malformed frames.

6. mode: total on fold configs. Header notes this trains DP on val episodes and makes val a fit metric, not generalization. That's a real research hazard — a reader six weeks from now will read "Valid/…_mse" as generalization. Please:

   • Rename the logged metric to something like Fit/…_mse when mode: total, or
   • Add a WandB tag / config field that surfaces this in the run summary.

7. eval_pipeline_action_mse.yaml sets max_frames: 0. In HumanRobotOverlayEval, limit > 0 gates the [:max(0, limit-rendered)] slice; with limit == 0 and limit is not None, max_frames_by_embodiment.get(...) returns 0, then the if limit is not None and limit > 0 branch is skipped, meaning all frames get rendered, not zero. That's the opposite of what the config name suggests (action-MSE only, no viz). If the intent is "no viz," set viz_func: null (already done) and fix the guard to treat limit == 0 as "render nothing." Right now viz_func: null alone makes it work by accident because viz = None short-circuits.

8. crop_scope: "episode" relies on self._episode_cu being stamped by an outer stage. VisualCore mentions "the packed pipeline's ObsEncoders stage stamps self._episode_cu" — but this diff doesn't show that stamping code, and none of the new configs use episode scope. If it's not exercised, drop the code path or add a smoke test; dead conditional branches in visual encoders bite later.

9. VisualCore provenance. Copying SpatialSoftmax + VisualCore verbatim from EgoVerse2 is fine and the docstring is honest about it, but please add the source commit SHA (not just "branch hpt-hnet-pusher-nc3") so we can reproduce the port. Also verify licensing/attribution matches robomimic upstream if this ultimately traces back there.

10. No tests. For a PR this size, at minimum I'd expect:

    • A smoke test that instantiates each new Hydra experiment config (pipeline_sampler_kp, pipeline_sampler_kp_moe, pipeline_sampler_pusht_h16) and does one forward+backward on a tiny synthetic batch.
    • A unit test for MoEFFN verifying (a) top_k routing shape, (b) aux loss is finite, (c) all experts receive grads across a batch when top_k < num_experts (the DDP-safety property).
    • A unit test for CrossTransformer(time_conditioning="additive") matching the concat variant on a fixed seed when hidden_dim matches.

Suggestions

• Rename HumanRobotOverlayEval — the alias FoldOverlayEval = HumanRobotOverlayEval at the bottom suggests you weren't sure. If the primary use is Fold + PushT overlay, `

---

Reviewed by Claude · Review workflow