TTT-Linear across-time layer (RoboTTT-style): closed-form inner SGD on fast weights, per-episode reset, tanh-gated ~identity at init; DualTrunkLevel ttt knob; arms fix2ada+ttt (H-Net) and txar_sdpada(+ttt) flat pair

[ElmoPA]

TTT-Linear across-time layer (RoboTTT-style): closed-form inner SGD on fast weights, per-episode reset, tanh-gated ~identity at init; DualTrunkLevel ttt knob; arms fix2ada+ttt (H-Net) and txar_sdpada(+ttt) flat pair

TTT-Linear stability fix (q/k L2-norm + out-LN + eta 0.02, after both TTT arms NaN'd); SDPHead detach_offregime + FlowHead denoiser_arch knobs; new arms dpclone_flow/ddpm_regoff499/sdp_srfix/sr2c15/sr2opt; 250-cadence launcher

[ElmoPA]

• feat(rldb): h264 video-backed images + role-in-key annotations + span selection #541 : 2 dependent PRs (#542 , #550 )
• TTT-Linear across-time layer (RoboTTT-style): closed-form inner SGD on fast weights, per-episode reset, tanh-gated ~identity at init; DualTrunkLevel ttt knob; arms fix2ada+ttt (H-Net) and txar_sdpada(+ttt) flat pair #540 👈 (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 #540

Summary

Large research PR adding a TTT-Linear across-time layer, Streaming Diffusion Policy (SDP) head, Flow head, and a batch of new flat/hierarchical model configs for pushshapes sim experiments. Almost entirely new files — low collision risk with existing training, but the diff is enormous and largely untestable from what's shown.

Key concerns

1. TTTLinearLayer performance/correctness

   • The inner scan is a Python double loop over episodes and chunks with in-place out[c0:c1] = ... assignments on a preallocated out. This works but (a) is slow for long packs, (b) creates a long autograd graph across every chunk of every episode in the batch. Have you profiled memory at typical cu_seqlens sizes? The comment says "no inner autograd" but the outer graph still backprops through every W update.
   • W = W - eta * grad with W starting as self.W0 means the first chunk's forward has Qc @ W0.T where W0 is zero-initialized (torch.zeros(d,d)) — so out[c0:c1] for the first chunk of every episode is exactly zero at init. That's fine given the tanh gate (~1e-3), but worth a comment; readers will assume W0 is trainable and nonzero.
   • int(cu_seqlens[b]) inside the loop forces a GPU→CPU sync per episode per forward. Minor but adds up in DDP.

2. TTT stability guarantees are empirical, not proven

   • The docstring notes "both TTT arms NaN'd" and lists three fixes (L2-norm q/k, eta 0.02, out-LN). Good instinct, but there's no test asserting finite outputs across e.g. 10k random forwards or long sequences. Given the history of NaNs, a unit test in tests/ exercising TTTLinearLayer(d=128, chunk_size=8) on (T=2048, cu_seqlens=[0,512,1024,2048]) and asserting torch.isfinite(out).all() after a backward pass would be cheap insurance.
   • eta = self.log_eta.exp() is unbounded above. If AdamW pushes log_eta up, stability regresses silently. Consider eta = eta_max * torch.sigmoid(self.log_eta) or a hard clamp.

3. SDPHead _buffer_targets pusher_hold assumption

   • assert self.D == 2, "pusher_hold pad assumes 2D xy actions" — fine for pushshapes, but the head is being introduced as a general dualstream head. If anyone tries end_mode=pusher_hold on bimanual arms this asserts at construction time, which is OK, but the coupling of "action space == state[:2]" to normalization is subtle: the comment says "ws512 action bounds", i.e. this only works if norm stats were computed with matching state/action bounds. Recommend either (a) reading the pad from a config field rather than state[:, :2], or (b) adding a runtime check.

4. SDPHead rollout state is on self._stream

   • Storing rollout buffer state on the module (self._stream) means the module is not stateless across evaluation runs, and DDP eval sharing modules will race. Also breaks if two eval instances share the same module (unlikely in your setup but worth a comment). At minimum _stream should be reset on train()/eval() transitions, or passed through batch.
   • The t_env == 0 or self._stream is None reset condition is right for single-rollout eval but fragile — please add a note in the eval loop / algo.step that documents this contract.

5. num_inference_steps divisibility

   • assert self.S % self.K == 0 in SDPHead.__init__ — good, but the configs shipping (buffer_chunks: 4, num_inference_steps: 16) satisfy it. Just make sure the launcher docs mention this so someone doesn't set K=3.

6. Configs

   • Six near-identical 150+ line YAML files with the ObsEncoder block duplicated verbatim per config. This is going to be a maintenance nightmare. Strongly recommend refactoring the common obs-encoder block into a Hydra defaults fragment (e.g. model/parts/pushshapes_obs.yaml). Not blocking for this PR, but flag it.
   • bf_nopre_sdp_fix2ada_ttt.yaml: both DualTrunkLevels carry ttt: {chunk_size: 8} but no other TTT knob (eta, gate_init) is exposed. If you later want to sweep those you'll have to touch the config schema — consider forwarding a full ttt dict now.
   • The sbatch scripts hardcode /coc/flash7/paphiwetsa3/... paths. These shouldn't be in the repo root; move to scripts/slurm/ or user-scoped configs. bf_eval_par.sbatch and bf_prdec_abl.sbatch at repo root are noise for other collaborators.

7. No tests

   • Diff shows zero test files. This is a large addition of numerical machinery (TTT scan, DDIM ladder math, per-episode buffer resets, hetero-D latent wrapper hinted at with LatentRHDenoiser) with prior NaN incidents. At least smoke tests for:
     • TTTLinearLayer finite outputs + backward
     • SDPHead train loss on a fake batch of (T=8, C=4, D=2, K=4)
     • _buffer_targets with end_mode="pusher_hold" and hand-constructed cu_seqlens

8. Training regressions

   • New model configs are all bf_* (batchflow) and don't touch existing configs — good, no regression surface.
   • No changes to norm stats, data loaders, or trainHydra.py — good.

Suggestions

• Add a `tests/test_ttt_layer

---

Reviewed by Claude · Review workflow