Transformer Transformer: A Unified Model for Motion-Conditioned Robot Co-design

TL;DR: This paper treats a robot's body the same way we treat states and actions — as tokens in a sequence — and trains one diffusion transformer over all of them jointly. Because the model learns the joint distribution of (embodiment, states, actions) rather than any particular reward, you can hand it a reward function it has never seen at test time, differentiate through its own denoised predictions, and steer the diffusion toward robot designs that score well — no separate critic, no simulator in the loop, no per-reward retraining. It beats CMA-ES on most design spaces while being orders of magnitude faster (0.5 s vs. 47 s on one benchmark; 21 s vs. 3.2 hours on another), and the authors fabricated one of its designs: a re-mounted, re-proportioned ALOHA that cut tracking error on dynamic cloth-flinging from 13.0 cm to 3.5 cm.

Why co-design is stuck in an inner-loop trap

Robot co-design — jointly optimizing a robot's hardware and its controller — has an ugly computational structure. To evaluate whether a candidate body is any good, you have to control it: get a controller (often train one), roll it out in simulation for thousands of sequential timesteps, and score the trajectory. Only then does the outer loop (typically an evolutionary algorithm like CMA-ES) get one fitness value and propose the next candidate. Every design evaluation is a full episode; every episode is inherently sequential; and legged robots need a fresh RL policy per morphology. The design spaces are also nasty: mixtures of discrete choices (how many joints? spring-loaded or serial legs? which mounting orientation?) and continuous ones (link lengths, mounting positions), so gradients don't exist in the usual sense.

Prior learned approaches attack pieces of this — learned design generators, learned critics, learned controllers — but each is trained on its own representation with its own objective, and stitching them into a pipeline keeps the iterative structure. Methods that use dynamics gradients for design either need a separately trained dynamics model or a differentiable simulator, and most sidestep the control problem by assuming fixed open-loop or purely kinematic actions.

The paper's framing question is concrete and slightly subversive: given end-effector trajectories from human demonstrations (the now-standard UMI-style task representation for manipulation), which body tracks them best? The observation motivating all of it: end-effector trajectories are embodiment-agnostic, but their execution is not — so imperfect cross-embodiment transfer is not just a nuisance, it's a design signal.

The idea: the body is just more tokens

The foundational move is a tokenization scheme, RoboToken, that puts a robot's morphology in the same sequence as its behavior. A robot description (anything expressible in MJCF: fixed-base arms, quadrupeds, humanoids, hands) becomes a set of continuous-valued embodiment tokens — one per link, joint, and motor — carrying physical parameters: primitive geometry and size, mass and inertia (log-encoded), joint ranges, damping, motor gains and force limits. An episode of behavior becomes state tokens (link poses, joint positions/velocities) and action tokens, one per entity per timestep.

Connectivity is handled by pointers: a joint token stores the IDs of the two links it connects; a motor token stores the ID of the joint it drives; each state token points to its embodiment token and its timestep. These IDs become learned positional embeddings. Arbitrary kinematic trees, variable DoF counts, even passive joints (a joint with no motor pointing to it — think Cassie's leg linkage) all fit naturally.

Why not just have an LLM emit MJCF/URDF text? The paper's answer is sharp and worth internalizing. First, compactness: RoboToken sequences are 27–110× shorter than the equivalent MJCF text tokens; eleven Menagerie robots spanning two orders of magnitude in mass (0.65 kg Allegro hand to 67.5 kg Anymal C) and 6–35 joints all fit in 28–101 tokens. Second, consistency: MJCF allows infinitely many equivalent transform decompositions, so a text model must waste capacity learning that redundancy; the tokenizer canonicalizes transforms and splits lumped inertias so there's one representation per robot. Third — the crucial one — optimizability: autoregressive categorical generation gives you no global controllability (early tokens can't respond to late ones) and no gradients (sampling is non-differentiable). Diffusion over continuous tokens gives you both.

The dynamics model is the critic

A diffusion transformer (DiT) is trained to denoise these sequences. Two masking schemes share the one architecture: condition on target end-effector poses and diffuse everything else, and you have a motion-to-robot generator that dreams up a body together with the states and actions it would exhibit while tracking the motion. Condition on embodiment + current state + targets and diffuse only actions, and you have a cross-embodiment controller.

Here is the aha. Because the model diffuses embodiment and dynamics jointly, every sample comes with everything a reward function could want: the body (for mass/size penalties) and the predicted trajectory (for tracking error, torque, joint velocity). The training objective is completely reward-agnostic — it's just dynamics modeling — but at inference you can plug the model's own denoised predictions into any differentiable reward you invent on the spot and get a value estimate. That value estimate can be used two ways:

Zeroth-order: run n parallel diffusion chains, score each candidate with the reward applied to its own predicted dynamics, return the best. Pure test-time-compute scaling — and indeed reward improves monotonically with more parallel samples across all design spaces.

Dynamics Self-Guidance (DSG): at each diffusion step k, take the model's noise prediction \epsilon_k, form the implied denoised tokens, evaluate the reward R(\text{embodiment}, \text{states}, \text{actions}), and backpropagate \nabla R into the noisy embodiment tokens, folding the gradient in via classifier-guided DDIM. The full sampling chain need not be differentiable — per-step gradients suffice to steer. Prior dynamics-guided design needed a separately trained dynamics model or a differentiable simulator; here the generator guides itself.

There's a second, less flashy insight buried here that explains the speed numbers. The model is non-autoregressive in time: it evaluates a design's long-horizon behavior by attending over just 8 timesteps sampled anywhere in the episode, in parallel. CMA-ES must run a controller and simulator through thousands of sequential steps per candidate. So this method parallelizes along two axes at once — batch (128 candidates per GPU pass) and time — and the axes multiply. Diffusing and scoring 128 candidates finishes before CMA-ES has evaluated 5.

Training data is procedurally generated: three design spaces (a fixed-base ViperX arm space stressing kinematic reachability, a quadruped-manipulator space stressing dynamic whole-body control, a mobile bimanual space stressing task complexity), controlled by differential IK (Mink) for the non-legged spaces and by 128 RL expert policies — one per discrete quadruped design combination, continuous parameters observed by the policy — for the legged one. Target motions come from 76 UMI human demonstrations plus a bimanual dishwashing set.

The evidence

The headline comparison is against CMA-ES with the discrete choices sampled and continuous ones evolved — the standard co-design baseline. Rewards and trajectories at evaluation are held out.

The pattern: on the arm and bimanual spaces, the diffusion model matches or beats CMA-ES at a small fraction of the wall-clock. The multi-trajectory setting is the killer app — one design must serve all 26 dishwashing trajectories simultaneously, handled zero-shot via diffusion composition despite single-trajectory training, and CMA-ES's cost scales linearly in trajectories (3.2 hours vs. 21–31 seconds). On the quadruped, though, CMA-ES wins on tracking (1.9 vs. 2.6 cm). The authors are upfront about why: predicted-vs-actual reward correlation is visibly worse for the legged space (Fig. 7), because the dynamics model must implicitly simulate a learned whole-body controller navigating contact discontinuities and fall-terminations. The method is exactly as good as its dynamics model — a clean, honest failure mode.

An interesting negative-space result: DSG and zeroth-order search are substitutes, not complements. Guidance improves each sample; search picks the best of many; with a large sample budget they converge to the same performance. Guidance wins by 242–341 reward only in the single-sample regime — i.e., it matters when samples are expensive (bigger models, longer episodes, interactive use), not when you can brute-force.

The same model, conditioned on embodiment, controls unseen robots. Validation-by-self-control correlates with expert-controller performance at Pearson's r = 0.53 on diffused quadrupeds, and in the bimanual space a 63.6M-parameter unified model reaches 5.9 cm / 9.6° tracking against the Mink oracle's 4.8 cm / 7.6° — promising but not yet a replacement for the 128 RL experts.

And then the part that separates this from pure-simulation co-design papers: they built the thing. Optimizing ALOHA for a tracking-velocity reward on dynamic cloth-flinging, the model chose to mount the arms upside down behind the workspace — converting an overhead fling into a more efficient underarm swing — with link lengths long enough to cover the motion but light enough for the Dynamixels.

That's a 73% error reduction on a task deliberately chosen to stress unmodeled physics — aerodynamic drag, cloth friction and weight — which is decent evidence the designs aren't just exploiting simulator quirks.

What to be skeptical about

"Unified" is aspirational in the main results. The headline optimization experiments train one model per design space (explicitly, to isolate embodiment-optimization performance); the truly unified multi-task model appears only in the bimanual appendix experiments. The architecture and representation demonstrably span embodiment classes, but a single model trained across all three spaces and both tasks isn't the thing that produced the main numbers. The selector's question — how well one learned dynamics model generalizes across genuinely different embodiment classes — remains mostly open.

No extrapolation, by construction. Generated designs interpolate within the training manifold: no hexapods from quadruped-only data, no link lengths beyond the training range. The authors frame this as feature-plus-limitation, fairly. But it means the "generative co-design" is really amortized search over a procedurally defined design space — the interesting design priors (spring-loaded linkages, bending spines, mounting options) were authored by humans into the procedural grammar. The upside-down ALOHA mount was one of three discrete mounting options in the grammar, not an invention.

The data bill is real. 128 RL experts × 16 A100-hours ≈ 2,000 GPU-hours just for the quadruped controllers, plus 3.8M/1.3M/50K episodes across the spaces. And for ViperX/ALOHA, training data was biased toward good designs by running CMA-ES first — so the evolutionary baseline is partly inside the model's training distribution. Amortization pays off if you query many rewards/trajectories per design space; for a one-off design problem, CMA-ES may still be the pragmatic choice.

Guidance needs hand-tuning. Guidance scales are 50 / 100 / 0.2 / 500 per setting, and embodiment-reward terms need a separately downscaled coefficient because self-attention makes embodiment-token gradients much larger than dynamics-token gradients — a subtle, un-principled patch (§8.4). Stochastic DDIM (\eta=1) was also needed specifically for DSG to escape local minima. On the fabrication question the triage flagged: the tracking-velocity reward used for ALOHA is one of the paper's standard, simple reward forms (Eq. 2 with a velocity penalty), not obviously cherry-shaped — but the data-quality biasing above does apply to the ALOHA space.

Why it matters anyway

The durable contribution may be less "diffusion beats CMA-ES" and more the representational claim: embodiment as first-class tokens alongside states and actions, in a format a single network can read, write, and differentiate through. Cross-embodiment control work already conditions on morphology, but usually with lossy embeddings; a complete, canonical, compact, continuous robot representation is the kind of infrastructure that outlives the paper it debuted in — plausibly the substrate for a future "robot foundation model" where the body is an output, not a constraint. The test-time-scaling result (more parallel diffusion samples → better designs, across all spaces and rewards, emerging purely from dynamics modeling) is a nice echo of language-model inference scaling in an entirely different domain, though the authors note it doesn't yet scale as robustly.

If you read one section, make it §2.3 plus appendix §8.4/§9.4 — the self-guidance mechanism and its tuning warts are where you'll form your own judgment about how general this really is. §2.1 (RoboToken) is the part most likely to be reused by others.