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

[ElmoPA]

No description provided.

[ElmoPA]

• feat(pipeline): consolidate Fold, PushT, overlays, and MoE #579 : 3 dependent PRs (#580 , #583 , #585 ) 👈 (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, multi-directional PR that consolidates Fold (RH keypoint cotrain), PushT (in-domain cotrain), overlay eval, and a new shared-residual MoE denoiser variant. Adds new Hydra configs, a Pipeline sampler model family, an obs/visual-core "stems" package, and MoE plumbing.

Key concerns

1. PR hygiene / scope

• No description on a diff this large. This is at minimum four independent features (Fold KP cotrain, PushT cotrain, MoE denoiser, overlay eval, plus a new stems/ package that looks like a port from another repo). These should be separate PRs — a single reviewer cannot reasonably certify all of it at once, and bisecting a training regression will be painful.
• Diff is truncated at 80k chars, so I cannot see stages_sampler.py, stages_sampler_moe.py, pipeline/algo.py, packed.py, tests, or the tail of visual_core.py. Anything below is contingent on those files.

2. Data integrity — Fold config

fold_rh_normal_cotrain_kp.yaml has an important, explicitly-documented behavior change:

• mode: total for both train and valid, meaning the same episodes appear in both splits. The comment acknowledges this ("validation episodes are NOT held out from training … val chunk-MSE is a fit metric, not a generalization metric").
• This is defensible for parity with the H-Net cells, but it must be surfaced in W&B run notes and any comparison table. A future reader looking at Valid/*_mse in wandb will assume it's a held-out metric. Consider renaming the logged metric prefix to Fit/… or logging a big warning at DataModule setup.
• skip_bounds_check: true — reasonable per comment (avoids silently dropping the exact tail-reach frames the policy needs), but this is a footgun for anyone copying the config. Consider a shared YAML anchor or a comment at the resolver level.

3. HumanRobotOverlayEval correctness

• _unnormalize_prediction reshapes (B,H,D) -> (B*H,D) before calling norm_stats.unnormalize. If the normalizer has any per-timestep behavior or does slicing on the last dim only, that's fine; but the target path calls unnormalize(batch, emb_id) on the whole batch (preserves 3D). The two paths use different code paths on the same tensor shape — is that guaranteed equivalent? Worth a unit test asserting unnormalize(x.reshape(...)).reshape(...) == unnormalize(x).
• The bare except Exception around viz(...) silently swallows failures and only prints. In a training run over many val steps this can hide a broken viz function forever. At minimum log via self.log/wandb, or fail loud on the first N steps and only suppress later.
• torch.quantile on very large flattened tensors allocates a sorted copy; if Valid/limit_val_batches grows this could OOM. Not a blocker, but consider .float().cpu() for the p95/median metric computation.
• The nested def take(value) closes over selected mutably — fine here, but consider selected = selected.to(device) once.

4. CrossTransformer back-compat

Adding time_conditioning="concat" as default preserves shape for existing checkpoints ✓. But the new if action_embedding_dim < act_dim: raise ValueError will now trip existing configs where hidden_dim // 2 < act_dim — previously this was silently rank-deficient and "worked." Please:

• Grep existing configs for hidden_dim/act_dim pairs and confirm none regress. bf_pipeline_sampler_pusht.yaml uses hidden_dim=256, act_dim=64 (128 ≥ 64 ✓). bf_pipeline_sampler_kp*.yaml uses hidden_dim=512, act_dim=128 (256 ≥ 128 ✓). But please also check any experiment configs not touched in this PR.

5. MoE — DDP correctness

DDPSafeMoEFFN.forward adds sum(param.reshape(-1)[0] * 0.0 for expert in experts for parameter in expert.parameters()) as a graph anchor. This works, but:

• It's O(num_experts × params_per_expert) Python-side per forward. With 8 experts × several linear layers × 16 blocks, that's ~hundreds of tiny ops per forward. Prefer find_unused_parameters=True on DDP, or the standard trick of a single sum(p.sum() * 0 for p in ...) accumulated once — or better, register experts under a module that DDP can be told to ignore.
• Confirm this works with gradient checkpointing (gradient_checkpointing: true in the configs) — the anchor must be inside the checkpointed region or DDP still complains on the outer graph.
• No test for the auxiliary loss actually being consumed. MoEFFN.last_aux_loss is stashed on the module but I can't see (diff truncated) whether stages_sampler_moe.py sums it into the training loss. If not, moe_aux_weight does nothing and load balancing will collapse to a few experts.

6. stems/ — new copy of visual encoders

visual_core.py header says it was ported verbatim from EgoVerse2 because EgoVerse-pact-2 only has SimpleConv. This is a real risk:

• Now there are two VisualCore-like classes in the repo. Which one should new configs point at? Please add a note in the CODEOWNERS/README and mark the old one as legacy, or delete it.
• The new stems/__init__.py re-exports SpatialSoftmax, VisualCore, ObsToken, CondEncoderModule, MultiEmbodimentCondEncoder — this becomes a new public API surface. Any tests?
• `crop

---

Reviewed by Claude · Review workflow