Not all robots are created equal — but what if you could design one for a specific task? We propose Transformer Transformer, a unified model that does exactly this: hand it a manipulation demonstration, and it generates a complete robot — every link, joint, motor, and inertial property — optimized for that motion. We fabricated one such design for cloth flinging on an ALOHA 2 bimanual platform; it reduced tracking error by 73% and max joint speed by 30% versus the original.
Behind that result is a diffusion transformer trained on RoboTokens, a unified tokenization of robot embodiments, states, and actions. The same architecture spans embodiment spaces (wheeled bimanual, quadrupeds, humanoids) and use cases (embodiment generation, cross-embodiment control). Rather than overfitting to one reward function, it is a dynamics model whose reward-agnostic predictions are converted into reward-specific value predictions at inference time, then used to steer embodiment diffusion through a procedure we call Dynamics Self-Guidance. Experiments across three design spaces show zero-shot optimization of unseen rewards and trajectories, improving performance and runtime over an evolutionary baseline.
A 17-minute walkthrough — or scroll for the figure-and-caption version.
A robot's embodiment decides what tasks it can do well. We pour effort into making robots smarter — bigger datasets, better algorithms, more capable policy architectures — but a bad embodiment can hamper even a great policy. You can collect all the tossing data in the world, but if your robot's shape is far from optimal, it might toss itself instead of the ball.
So let's frame the problem concretely. Given as input a target end-effector motion and a reward function, we want to generate a complete robot embodiment for the task — by complete we mean every link, joint, motor, and inertial property, plus a controller that drives the whole thing. We call this
motion-conditioned robot co-design. A randomly procedurally-generated robot attempting the target motion: it falls before it can finish 🙃
Our framework reframes robot co-design as a three-step process: you demonstrate the desired end-effector motion (e.g., from a UMI 3 human demonstration), our model
PART 1
If we want one model to co-design any robot, the network needs a shared vocabulary that spans every embodiment. RoboTokens are that vocabulary: typed tokens for every link, joint, motor, state, and action — organized so a single sequence describes both a robot's embodiment (time-invariant) and
its dynamics (time-varying). RoboTokens has the following key design goals:
We tokenized 11 robots from the MuJoCo Menagerie 10, spanning two orders of magnitude in mass (0.65 kg dexterous hand to 67.5 kg quadruped) and from 6 to 35 active joints. Each became a sequence of 28–101 RoboTokens — small enough to diffuse end-to-end.
PART 2
The model is a diffusion transformer (DiT 1) trained on RoboTokens with a DDIM
(And about the pun: the first “Transformer” refers to robots whose embodiments can change. The second refers to the self-attention architecture 🥸)
Because the controller sees the full embodiment as input, it can adjust its control strategy to the robot it's driving. Change link inertias, joint ranges, or motor gains, and the same policy keeps tracking — without any retraining per robot. We behavior-cloned whole-body controllers from RL experts 7,
Run the model with no conditioning at all and it diffuses out complete robots from Gaussian noise. Trained on diverse robots from Menagerie 10, a single Transformer Transformer produces fixed-base arms, quadrupeds, humanoids, and dexterous hands. Because diffusion is stochastic, every run gives different robots.
PART 3
Now the actual co-design problem: how do we generate a high-reward robot for a reward function the model has never seen during training?
Here is the trick. The same model predicts both the embodiment and its dynamics, so we can feed those predictions into the user's reward function and get a predicted reward. With predicted rewards in hand, we can run the model many times in parallel, predict the reward for each candidate, and
return the best one. That is best-of-N sampling, but at GPU speed. We can do better. Because the predicted reward is a differentiable function of the embodiment tokens, we can take its gradient and inject it back into the diffusion process — pushing samples toward higher reward at every denoising step, following classifier-guided DDIM 11. The model is, in effect, asking itself: “how should this robot change to increase the reward?” We call this
Take a randomly-generated quadruped that falls during a dynamic toss. Ask our model to generate a robot for the same trajectory with only a tracking reward and it lowers the center of mass and widens the support polygon — stable design choices it discovered as a side-effect of trying to maximize tracking reward it has never seen.
Now add a torque-penalty term and the optimization landscape shifts. The model returns a smaller, lighter robot — still tracking the motion, but using 30% less average torque. Add a size penalty and the landscape shifts again. This is what people might call zero-shotting rewards: train on dynamics modeling, optimize at inference for any reward you can write down.
Manipulation isn't one motion; it's a distribution of motions. We want to optimize a robot for an entire task — say, dish washing — not a single trajectory. Diffusion composition 14,
We tested this on the publicly-available UMI bimanual dish-washing dataset 3, holding out 26 validation trajectories. Transformer Transformer generated robots optimized over the entire held-out distribution — navigating both the discrete space (spine designs: fixed, sliding, bending) and the continuous space (arm dimensions, mounting offsets) — and the same model then served as the cross-embodiment controller to validate each generated design.
We compare against two baselines: Random (sample a robot from the procedural design space) and CMA-ES 16, a black-box optimizer widely used in co-design. CMA-ES rolls out candidate robots in MuJoCo
We give all methods the same task: generate a robot for a target trajectory and reward, control it, report the achieved reward.
Like large language models that get better with longer thinking 19, Transformer Transformer gets better designs when given more inference seeds to sample over. Performance does eventually plateau after about a minute — unlike the seemingly endless gains LLMs get from longer reasoning. We discuss this in the limitations.
To stress-test the whole system on hardware, we picked dynamic cloth unfolding on ALOHA 2. Flinging
We asked our model to optimize the ALOHA design for the “Tracking Velocity” reward, then fabricated the resulting design. Two things changed compared to the original:
Geometry and scene. RoboTokens currently cover primitive-based geometry, not arbitrary meshes or deformables, and don't yet encode the surrounding scene or contact targets. Extending to complex geometry and scene context is a clear next step toward fully task-conditioned co-design.
Controller transfer. Our cross-embodiment controller correlates with the RL oracle at Pearson r = 0.53 on diffused embodiments — promising but with headroom. Some outlier embodiments expose where the controller fails to generalize.
Test-time scaling plateau. Like other guided-diffusion methods, our test-time compute gains plateau after about a minute. Unlike LLMs 19, more inference time stops helping past a point. Improving reward-prediction accuracy and alternative inference schemes are open directions.
Data generation cost. RL-based whole-body controller data generation is expensive (one RL expert per discrete design choice, 16 hours of A100 per policy). Scaling Transformer Transformer to
new design spaces hinges on reducing this cost. Full reference list available in the paper. The selection above covers the citations used on this page.