# Transformer Transformer: A Unified Model for Motion-Conditioned Robot Co-Design

> Source: <https://transformer-transformer.github.io/>
> Published: 2026-07-29 03:52:29+00:00

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.
