algo: RDT (Robotics Diffusion Transformer) on the HPT chassis by RyanPCo · Pull Request #680 · GaTech-RL2/EgoVerse

RyanPCo · 2026-09-20T16:15:24Z

Adds RDT (Robotics Diffusion Transformer) as an algo on the HPT chassis. RDT(HPT) keeps HPT's batch handling, stems, encoders, denoising heads, eval and config layout, and swaps the policy network for RDT's DiT.

Behaviour changes

None for existing recipes. HPT gains a _build_policy hook and a frames argument on its image augs, and at the default frames=1 both do what they did before.

Model

No Perceiver pooling. The stems hand their full token sequences to the head.
The DiT blocks sit in the trunk slot and are shared across embodiments. They self-attend over [timestep; ctrl freq; state; noisy action chunk] and cross-attend to language tokens (even blocks) and image tokens (odd blocks).
Each FMPolicy head holds one embodiment's action adaptor, position table and output layer, so co-training shares the blocks at native action widths.
Conditions:
- images from DINOv3 ViT-B/16 (vision: DINOv3Stem, a timm DINOv3 ViT image stem (patch tokens, any aspect ratio) #679), fine-tuned at 0.1x LR, at 224x400;
- language from frozen Qwen3-Embedding-0.6B per-token states;
- a 2-frame image history, the past frame 0.1 s back (loader: image history (a past front frame, gap in seconds) and the episode fps in every sample #678);
- a ctrl-freq token from the per-sample fps;
- per-stream condition dropout (state and language 0.1).
Differences from the paper, with the reasons, are listed at the top of rdt_bc_mecka_6d_dinov3_vitb.yaml. The main ones: flow matching instead of DDPM, DINOv3 instead of SigLIP, Qwen3 instead of T5-XXL, and per-embodiment adaptors instead of the 128-D unified vector.
ot and freeze_repr are HPT-only and raise.

Configs

model/rdt_bc_mecka_6d_dinov3_vitb.yaml is the default: width 1024 x 28 blocks x 16 heads, with 392.7M trainable (DiT 296.6M, DINOv3 85.6M, plus adaptors). It uses backbone_lr_scale: 0.1 and AdamW 1e-4 with a 1k-step warmup, then a constant LR.
model/rdt_bc_mecka_6d_dinov3_vitl.yaml swaps in the ViT-L tower.
model/rdt_1b_bc_mecka_6d_dinov3_vitb.yaml is the RDT-1B shape (2048 wide, 32 heads).
train_zarr_mecka_flagship_6d_rdt.yaml is train_zarr_mecka_flagship_6d_hpt with the RDT model. It sets image_history_gap_s: 0.1 and gradient_clip_val: 1.0. It keeps proprio history at K=1 / stride 1, because the RDT state adaptor is declared at history_len: 1 and the two must agree.

python egomimic/trainHydra.py --config-name=train_zarr_mecka_flagship_6d_rdt \
  norm_stats.precomputed_norm_path=<out>/norm_stats

Prefetch the timm checkpoint on a login node before an offline or multi-rank run (command in the model config).

Step time, as recorded in the model config (H100, batch 32, bf16):

config	ms / step	memory
ViT-B, 1024-wide DiT	212	30 GiB
ViT-L	365	49 GiB
ViT-B, 1B DiT	359	57 GiB

Sampling takes 0.3 s per batch.

Image history through the augs

The (past, current) pair goes through the train augs as (B * 2, 3, H, W) with frames=2. One jitter draw per sample covers both frames, and every op runs per frame. Laid side by side instead, ColorJitter's contrast took one mean over both frames, which shifted the past frame by up to 0.20 depending on the current frame's content. A Resize or CenterCrop would also have mixed pixels across the seam.
BatchedColorJitter and PerSampleAugs take frames. The per-image fallback applies the aug once per clip.
The pair also gets HPT's per-sample absent-camera guard, since both frames go through _apply_image_augs.

Compile

compile.enabled=true leaves the trunk off HPT's compile list (no_trunk), but the DiT is still compiled. Every head's denoiser calls it, so dynamo traces its blocks into the head's graph.

Tests

tests/unit/test_rdt.py covers:
- the denoiser: shapes, zero init, masked language tokens, and one backbone shared across widths;
- the model: streams, parameters reached by the loss (plain DDP needs this), and the frame-count check;
- condition and image-history dropout;
- the fps requirement;
- pairing and ordering of the history frames under one aug draw on both aug paths, a per-frame contrast mean, a per-frame resize, and an absent past frame staying zero;
- the DiT being traced through the heads, asserted with a capturing backend.
tests/unit/test_train_step.py::test_rdt_train_step runs a train step on synthetic mecka episodes with a tiny DiT and a random tiny DINOv3 tower (rdt_small_overrides in tests/fixtures/recipes.py). It checks that <front>_hist and fps reach the batch.

Folded in (to keep the stack under Graphite's 50-branch cap):

loader: image history (a past front frame, gap in seconds) and the episode fps in every sample #678 — loader: image history (a past front frame, gap in seconds) and the episode fps in every sample
vision: DINOv3Stem, a timm DINOv3 ViT image stem (patch tokens, any aspect ratio) #679 — vision: DINOv3Stem, a timm DINOv3 ViT image stem (patch tokens, any aspect ratio)

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

RyanPCo · 2026-09-20T18:53:05Z

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

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): 3 low, 2 nit.

RyanPCo · 2026-10-30T20:55:41Z

+        cond_drop: dict | None = None,
+        image_history: int = 1,
+        image_history_dropout: float = 0.0,
+        **kwargs,


low (correctness): **kwargs is accepted but never used or forwarded, and RDT._build_policy passes the trunk: config straight in. A typo or an HPT-style key (e.g. +...trunk.num_blocks=12, HPT's name for depth) is silently dropped, and the model trains at depth 28. No RDT config passes extra keys, so remove **kwargs and let unknown keys fail at construction.

RyanPCo · 2026-10-30T20:55:41Z

+                "model.safetensors",
+                local_files_only=True,
+            )
+        except Exception:  # not cached, or not an HF-hosted timm model


low (error-handling): except Exception turns any failure, including a bad pretrained_cfg or a coding error, into a silent skip of the pretrained-weight hash check. Catch only the expected cases: huggingface_hub.errors.LocalEntryNotFoundError for a cache miss, and KeyError/TypeError when pretrained_cfg has no hf_hub_id. Log at debug level when skipping.

RyanPCo · 2026-10-30T20:55:41Z

+    def _stream(self, modality, stem) -> str:
+        if modality in self.encoders:
+            return "img"
+        return "lang" if hasattr(stem, "forward_with_mask") else "state"


low (correctness): _stream labels a stem as language only if it has forward_with_mask. A text stem without it, such as QwenPooledEncoder (HPT's usual text stem), is therefore routed to state. The prompt then silently becomes a state token, gets cond_drop['state'] instead of the lang dropout, and any token-count error blames proprio history. Detect text input explicitly (e.g. isinstance(data[modality], list)), and raise when such a stem lacks forward_with_mask, pointing the user to QwenPerTokenEncoder.

RyanPCo · 2026-10-30T20:55:41Z

+    assert torch.equal(large[0, D // 2 :], large[25, D // 2 :])
+
+
+def _dino(**kwargs):


nit (test-quality): _dino and TINY_TOWER are copied verbatim from tests/unit/test_dinov3_stem.py, also new in this PR, and rdt_small_overrides repeats the same tower kwargs. Move TINY_TOWER and a tiny_dinov3_stem(output_dim, **kw) helper into tests/fixtures and import them in both test files.

RyanPCo · 2026-10-30T20:55:41Z

+    assert (algo.train_image_augs.vectorized is not None) == vectorized
+    cam = "observations.images.front_img_1"
+    frame = torch.rand(1, 3, 8, 8).expand(B, -1, -1, -1).contiguous()
+    batch = {


nit (test-quality): This batch dict, and the one in test_algo_keeps_past_and_current_in_order, is exactly what _pair_batch returns a few lines below. Move _pair_batch above both tests and call it from them. Building the key as cam + IMAGE_HISTORY_SUFFIX instead of f"{cam}_hist" also keeps the tests tied to the real constant.

…isode fps in every sample Human.get_keymap(image_history_gap_s=...) adds a <front key>_hist camera entry; ZarrDataset reads it round(gap x fps) frames back, clamped at the episode start. Every sample also carries the episode's fps (zarr attr). The past frame names the front camera's zarr key, so the loader's one-call-per- key read covers both windows. JPEG keys are one frame per chunk, and a plain union span would read every frame in between; the read planner keeps them two rows. Per sample on a synthetic 30 fps episode of 360x640 JPEGs (median of 7x40 samples, login node, noisy): no history 11-12 ms; gap 0.1 s 21-24 -> 19-23 ms (the second decode, which history needs); 0.5 s 30-34 -> 20-21 ms; 2 s 53-61 -> 19 ms, flat in the gap. A test counts the JPEG rows read at 0.1, 0.5 and 2 s. Nothing consumes <front>_hist here; the keymap docstring says so. Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw

…spect ratio) Same call convention as ResNet. Loaded through timm: the locked transformers (4.53.2) has no DINOv3 class, and the timm hub repos are ungated. RoPE-only, so image_size may be [H, W]; exposes grid_size. pretrained_reference_state_dict reads the checkpoint back from the local HF cache (never the network), so the pretrained-weight hash check covers this tower as it does the Qwen stems; an uncached checkpoint or a random tower is skipped. Against the real timm/vit_base_patch16_dinov3.lvd1689m file the key sets match (162 tensors) and every value is equal. Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw

RDT(HPT) swaps the policy network: no Perceiver pooling; stems hand full token sequences to the head; the DiT blocks live in the trunk slot, shared across embodiments, and cross-attend to language / image tokens in alternating blocks; each FMPolicy head holds one embodiment's action adaptor, position table and output layer. DINOv3 ViT-B/L (fine-tuned at 0.1x) + frozen Qwen3-Embedding, 2-frame image history, ctrl-freq token, per-stream condition dropout. Configs: model/rdt_*; train_zarr_mecka_flagship_6d_rdt. The (past, current) pair goes through the train augs as (B * 2, 3, H, W) with frames=2: one jitter draw per sample covers both frames, and every op runs per frame. Laid side by side instead, ColorJitter's contrast took one mean over both frames (the past frame shifted by up to 0.20 with the current frame's content) and a Resize or CenterCrop would have mixed pixels across the seam. The pair also gets HPT's per-sample absent-camera guard, since both go through _apply_image_augs. compile.enabled=true leaves the trunk off HPT's compile list (no_trunk), but the DiT is still compiled: every head's denoiser calls it, so dynamo traces its blocks into the head's graph (asserted in test_rdt with a capturing backend). Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw

RyanPCo · 2026-10-30T23:09:49Z

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

Fixed

RDTModel swallows unknown trunk keys via **kwargs: removed **kwargs; test that num_blocks raises TypeError. Fixup targets 6dbadd0 (PR689), which last edited the adjacent signature lines, so autosquash applies cleanly
Blanket except Exception in pretrained_reference_state_dict: returns None only for a missing hf_hub_id or LocalEntryNotFoundError; other errors propagate. No module logger exists, so no debug log was added
_dino / TINY_TOWER fixture duplicated: new tests/fixtures/models.py (TINY_TOWER, tiny_dinov3_stem with a lazy import); used by test_dinov3_stem, test_rdt and rdt_small_overrides
Text stem without forward_with_mask treated as state: _stream raises for QwenPooledEncoder/T5Encoder (lazy import); forward raises if a non-tensor reaches the state branch; test added

Partial

Hand-built pair batches duplicate _pair_batch; _hist hard-coded: _pair_batch moved up and reused, CAM/PAST built from IMAGE_HISTORY_SUFFIX in test_rdt, test_image_history and test_train_step; one later memory test (PR689 code) still spells _hist

RyanPCo force-pushed the ryanco/rdt branch from fbf30ff to 7313a1a Compare September 20, 2026 19:09

RyanPCo force-pushed the ryanco/dinov3-stem branch 2 times, most recently from 360a16d to 4f33675 Compare September 21, 2026 18:45

RyanPCo force-pushed the ryanco/rdt branch 2 times, most recently from 1a86f11 to f87f4a1 Compare September 22, 2026 13:18

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 1cb4602 to 5980661 Compare October 12, 2026 23:19

RyanPCo force-pushed the ryanco/rdt branch from 6a8a7b9 to 23b0df3 Compare October 12, 2026 23:19

RyanPCo mentioned this pull request Oct 19, 2026

lance migration #716

Open

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 5980661 to 5e662b3 Compare October 19, 2026 19:50

RyanPCo force-pushed the ryanco/rdt branch from 23b0df3 to eef4d01 Compare October 19, 2026 19:51

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 5e662b3 to a4b4fd0 Compare October 19, 2026 20:22

RyanPCo force-pushed the ryanco/rdt branch from eef4d01 to a7ad86d Compare October 19, 2026 20:22

RyanPCo force-pushed the ryanco/pi-scratch-init branch from a4b4fd0 to 7e3a333 Compare October 30, 2026 18:27

RyanPCo force-pushed the ryanco/rdt branch from a7ad86d to 96c6bbd Compare October 30, 2026 18:27

RyanPCo commented Oct 30, 2026

View reviewed changes

RyanPCo and others added 3 commits October 30, 2026 18:51

RyanPCo force-pushed the ryanco/rdt branch from 96c6bbd to 6b43eda Compare October 30, 2026 23:05

RyanPCo force-pushed the ryanco/pi-scratch-init branch from 7e3a333 to 3c557cc Compare October 30, 2026 23:05

RyanPCo mentioned this pull request Oct 30, 2026

augs: per-sample image augmentation for HPT train-mode augs #651

Draft

Navigation Menu

Repository navigation

algo: RDT (Robotics Diffusion Transformer) on the HPT chassis - #680

algo: RDT (Robotics Diffusion Transformer) on the HPT chassis#680
RyanPCo wants to merge 3 commits into
ryanco/pi-scratch-initfrom
ryanco/rdt

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo left a comment

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		assert torch.equal(large[0, D // 2 :], large[25, D // 2 :])


		def _dino(**kwargs):

Conversation

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Behaviour changes

Model

Configs

Image history through the augs

Compile

Tests

Uh oh!

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

RyanPCo left a comment

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

RyanPCo commented Sep 20, 2026 •

edited

Loading

RyanPCo commented Sep 20, 2026 •

edited

Loading