[RyanPCo]

Adds an optional long-range visual memory to RDT, and the loader support it needs. Without trunk.memory and key_map.image_memory, models and keymaps are unchanged.

Loader

• Human.get_keymap(image_memory=N, image_memory_stride_s=s) adds <front>_mem: N front frames s seconds apart, ending at the current frame, plus a (N,) <front>_mem_mask of the steps inside the episode. Steps before the episode start repeat the oldest real frame and are masked.
• history_stride_s spaces proprio history in seconds, read in each episode's fps, so a window spans the same time at any fps. It replaces history_stride when set.
• History windows are now point reads instead of a contiguous span thinned afterwards. On JPEG keys (one frame per chunk) a span would fetch every frame between steps; plan_reads still merges them on numeric keys. Each distinct frame is read and decoded once.

Model

• RDTMemory encodes each memory frame to one token with a second DINOv3 tower (DINOv3Stem(pool="cls")), fine-tuned at the backbone LR scale, and tags it with a sin-cos code of the frame's age in seconds. fuse_proprio adds a projection of that step's proprio token, which needs the proprio history on the same grid.
• The N tokens join the DiT's self-attention prefix [t; freq; state; memory; actions] under a key mask, so padded steps are never attended.
• short_dropout replaces a sample's whole short image stream by a learned null token during training, so the policy cannot lean on the latest frames alone.
• DINOv3Stem gains pool (patch / cls) and grad_checkpointing. Backprop through the fine-tuned memory tower (320 frames per GPU) OOMed an H100 without it.
• The backbone LR group now walks every module with backbone_parameters, not just the encoders, so the memory tower lands in it.
• Model and data config must agree: memory without image_memory in the batch (or the reverse), a frame-count mismatch, or an encoder that gives more than one token per frame raises.

Fixes

• HPT._stem_history_len called .get on the built nn.ModuleDict, so every proprio-history (K > 1) run crashed at its first batch. The unit tests used a plain dict.
• The last-step eval MSE handed torchmetrics a strided [:, -1] slice, which it .view(-1)s. On GPU .cpu() made it contiguous; on a CPU run it raised.

The memory ablation train configs (a 2x2 over proprio history and visual memory on the Mecka 20% split) and a cluster launcher are kept out of this PR.

Tests

• tests/unit/test_image_memory.py: frame spacing in seconds, front padding and masking, the first frame, one-frame windows, one read per distinct frame, no memory by default, norm mode dropping the memory camera.
• tests/unit/test_rdt.py: masked memory tokens ignored, age tagging, short dropout train-only, loss reaching every trainable parameter (with and without fuse_proprio), the memory tower in the backbone LR group, and the config/batch agreement checks.
• tests/unit/test_proprio_history.py: stride in seconds following the episode fps, and history_len read from a built ModuleDict.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

[RyanPCo]

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

• pi: keep per-sample intrinsics so overlays use the episode camera #690 : 3 dependent PRs (#691 , #707 , #716 )
• rdt: long-range visual memory; image memory window and history stride in seconds #689 👈 (View in Graphite)
• norm stats: cache sampled rows per episode so splits share them #687 : 1 other dependent PR (#688 )
• data: exclude episodes by hash; scan Mecka egomotion for failed head tracking #686
• eval: log validation videos to wandb #684 : 1 other dependent PR (#685 )
• slurm: stop cancelled jobs from leaking whole GPUs #682 : 1 other dependent PR (#683 )
• algo: RDT (Robotics Diffusion Transformer) on the HPT chassis #680 : 1 other dependent PR (#681 )
• pi: initialize from PaliGemma with a scratch action expert #677 : 1 other dependent PR (#678 )
• perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi #676
• loader: one zarr read per key, not one per keymap entry #674
• eval: build the pinned-episode video batches only where they are used #673
• checkpoints: last.ckpt as a symlink, not a second full copy #672
• trainHydra: metric_frames_per_episode: auto, and warn on empty prompts #671
• transforms: per-step wrist framing for proprio history #670
• viz: one keypoint renderer, the horizon trace #669
• configs: history and backbone-LR knobs on the keypoint wrist HPT model #668
• trainer: ddp_bf16_compress, DDP with a bf16 gradient all-reduce #667
• eval: overlay videos play at the source rate under DDP #666
• hpt: build the policy on CPU and let Lightning place it #664 : 1 other dependent PR (#665 )
• hpt: never re-initialise pretrained stems; hash-check them; mask text padding #645 : 1 other dependent PR (#646 )
• obs: proprio history (K proprio steps, one token each, learned time embedding) #650
• loss: mask the repeat-last action padding out of the HPT loss #649
• resnet: keep a frozen backbone in eval mode; give backbones their own LR group #652
• augs: per-sample image augmentation for HPT train-mode augs #651
• eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD #644
• data: mecka_fold_flagship_top3_hpt_6d, the 3-operator flagship train split #643
• val: subsampled metric loaders + pinned contiguous video loaders #642
• infra: mecka 6D / keypoint recipes, 300M models, defaults, and the rollout revert #640
• eval: layout-aware frame-invariant metrics, and a train-split viz head #659
• data: stop letting one bad episode or a rotation channel take a run down #658
• norm: guard degenerate channels, horizon pooling, and no rot6d normalization #657
• keypoints: the 144-D wrist-first hand-keypoint action becomes the human default #656
• mecka: fix the double-mirrored LEFT wrist frame, and correct existing zarrs #655
• 6d: PI goes 6D-only; Eva and Human gain the continuous-rotation modes #654
• 6d: continuous rotation representation, bimanual layouts, shared norm helpers #653
• train: give each Slurm slurm job gets its own cache dir to avoid conflict #638 : 1 other dependent PR (#639 )
• write git version into the run dir for reproducibility #637
• train: don't load PI base weights on resume and verify that checkpoint path exists #636
• configs: declare proprio and actions dims once per HPT model config #634 : 1 other dependent PR (#635 )
• cache norm stats for data configs to minimize recomputation #633
• import cloud packages lazily to support local training #631 : 1 other dependent PR (#632 )
• cache dataset resolver for faster training initialization #630
• validate episode hashes in data config before training #628
• modify pi camera map to match data #629
• add tests to ensure training works between every model and data vendor #626
• run CI tests on every PR #625
• add env variables for user-specific paths #624
• build: set pyproject and uv as the dependency managers #623
• fix(configs): fix paths for data hydra yaml configs #622
• refactor(algo): delete ACT #627 : 1 other dependent PR (#621 )
• tests: create dedicated test folder with unit and integration tests #620
• add pyright type checker #619
• main

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

[RyanPCo]

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 2 low, 1 nit.

[RyanPCo]

low (dead-code): Nothing reads RDTMemory.n_tokens, and each head's n_memory_tokens has to be set by hand to match trunk.memory.frames. Either derive the heads' n_memory_tokens from it (a hydra interpolation, or set it while building the heads) or delete the property.

[RyanPCo]

low (efficiency): This comprehension still decodes every entry. Front-padding duplicates are decoded and then overwritten, so the comment's "decode each distinct frame once" is not true. Dedupe before decoding: by_idx = dict(zip(map(int, window), data[k])), then unique = {i: simplejpeg.decode_jpeg(b, colorspace="RGB") for i, b in by_idx.items()}.

[RyanPCo]

nit (efficiency): The comment says each distinct frame is decoded once, but the dict comprehension calls simplejpeg.decode_jpeg for every window entry and keeps only the last result per index. Front-padded repeats are decoded K times. Check int(i) not in unique before decoding.

[RyanPCo]

Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed.

Fixed

• RDTMemory.n_tokens unused; n_memory_tokens set by hand: RDTModel.finalize_modules sets each head's n_memory_tokens from long_memory.n_tokens via the new RDTDenoiser.set_n_memory_tokens; test helper no longer passes it; test added
• Decode-once comprehension still decodes every frame: bytes deduped by index before decoding; test counts decode calls on a front-padded window
• JPEG window decode runs on every entry: duplicate of the decode-once item above; same fix