Paper ArXiv arXiv Video Code

Transformer Transformer

A Unified Model for Motion-Conditioned Robot Co-design

Huy Ha, Karen Liu, Shuran Song

Paper ArXiv arXiv Video Code

Not all robots are created equal

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 ALOHA2 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.

Technical Summary Video

A 17-minute walkthrough — or scroll for the figure-and-caption version.

What is the best robot
for a given manipulation task?

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 🙃

Demonstrate, Generate, Validate

Our framework reframes robot co-design as a three-step process: you demonstrate the desired end-effector motion (e.g., from a UMI3 human demonstration), our model generates an optimized embodiment, and the same model validates the design by directly controlling it. The same network plays three roles — generator, critic, and cross-embodiment controller — by training jointly on one unified tokenization of robots.

Three-step pipeline: demonstrate, generate, validate.
Three steps, one model. Given target end-effector motions (e.g., from a human demo) and a user-defined reward, Transformer Transformer generates a complete robot — geometry, kinematics, inertia, motors — then controls that robot to track the motion. No separate optimizer, critic, or controller pipeline.
The paper has three parts: a unified robot representation (RoboTokens), a unified architecture (Transformer Transformer), and a unified training objective (Dynamics Self-Guidance). You have until the architecture section to figure out the pun 🥸

PART 1

RoboTokens
A Unified Robot Representation

RoboTokens tokenization scheme.
Every robot becomes a typed token sequence. Blue tokens encode embodiment (links, joints, motors); orange tokens encode dynamics (states and actions). The same tokenizer ingests a quadruped, a humanoid, or a dexterous hand — no per-robot adapter — so a single model can learn across all of them.

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:

  • Complete. A RoboToken sequence captures everything about a rigid articulated robot: 5 embodiment token types (links, fixed joints, rotating/sliding joints, ball joints, motors) plus state and action tokens.
  • Flexible. Tokens reference each other by ID — a joint token knows the two link tokens it connects — so the same format handles a 6-DoF arm or a 35-DoF biped.
  • Consistent. The tokenizer canonicalizes redundant spatial offsets and stores inertial / transform data in a single frame, reducing the variance the model has to learn through.
  • Extensible. New token types slot in without changing the architecture — we add target end-effector pose tokens for trajectory tracking by pointing them at the relevant end-effector and timestep IDs, no retokenization of the rest of the robot required.
  • Optimizable. Unlike autoregressive text formats like MJCF, diffusing continuous-valued RoboTokens gives both global controllability (later tokens influence earlier ones across denoising steps) and differentiability — the two ingredients needed for gradient-based reward optimization at inference time.

We tokenized 11 robots from the MuJoCo Menagerie10, 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.

RoboTokens compactness vs MJCF text tokens.
RoboTokens are 27–110× more compact than MJCF text. A robot's full description fits in tens of typed tokens (blue counts) where a language tokenizer would need thousands (grey counts) — small enough that a transformer can diffuse them directly, rather than autoregressively writing XML.

PART 2

Transformer Transformer
A Unified Architecture

The model is a diffusion transformer (DiT1) trained on RoboTokens with a DDIM4 noise schedule. Noise is added to every input token except the ones we condition on — and by changing which tokens are masked, the same network wears different hats. Mask everything to get an unconditional robot generator. Mask the actions and you get a cross-embodiment controller5,6. Mask the embodiment and you get a motion-conditioned robot designer. Same weights, jointly trained on dynamics.

(And about the pun: the first “Transformer” refers to robots whose embodiments can change. The second refers to the self-attention architecture 🥸)

Transformer Transformer architecture.
One architecture, two use cases. (a) Motion-to-robot optimization jointly diffuses embodiment and dynamics tokens. (b) Cross-embodiment control diffuses the actions, conditioned on the embodiment and the target motion. At inference time, the model's own predictions are converted into a reward prediction whose gradients steer the generation toward higher-reward designs.

An embodiment-aware controller

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 experts7,8,9 trained one-per-design on the procedural quadruped space.

One policy, many quadrupeds. The same controller tracks the same trajectory across quadrupeds with continuous variations (link dimensions) and discrete variations (DoF, knee direction, spring-loaded vs serial leg).

A generator over diverse robot classes

Run the model with no conditioning at all and it diffuses out complete robots from Gaussian noise. Trained on diverse robots from Menagerie10, a single Transformer Transformer produces fixed-base arms, quadrupeds, humanoids, and dexterous hands. Because diffusion is stochastic, every run gives different robots.

Robots, diffused. A single model generates fixed-base arms, quadrupeds, humanoids, and dexterous hands from the same noise schedule.

PART 3

Dynamics Self-Guidance
Zero-shotting Rewards

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 DDIM11. The model is, in effect, asking itself: “how should this robot change to increase the reward?” We call this Dynamics Self-Guidance.

Why “self-guidance”: related work that incorporates dynamics into diffusion guidance trains a separate dynamics model12 or relies on a differentiable simulator13. Here, the model uses its own dynamics predictions to steer itself.
Climbing the landscape.
One gradient, the whole episode. Transformer Transformer predicts state and action tokens across the episode's horizon in parallel; the reward gradient then flows back through all of it to the embodiment tokens. Hardware gets optimized over long-horizon dynamics, not one timestep at a time.

Reward in, robot out

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.

Diverse generated designs for different conditioning.
Same model, different conditioning, different robots. Swapping the target trajectory (Floor Scrub → Dynamic Toss) or adding a size penalty reshapes the optimization landscape, and the generated embodiments shift accordingly — without retraining. Discrete choices (leg design, DoF) and continuous choices (link length, mounting point) both vary in response.

Composing diffusion models for whole datasets

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 composition14,15 lets us do exactly this: run the model in parallel with different trajectories as conditioning, average the predicted noise at each diffusion step, and you get a generator optimized for all of them simultaneously.

We tested this on the publicly-available UMI bimanual dish-washing dataset3, 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.

One robot, a whole task. Composed across the 26 held-out UMI dish-washing trajectories, the model returns a single embodiment optimized for the entire distribution — and the same network then drives it through each motion.

Better Designs, Faster

We compare against two baselines: Random (sample a robot from the procedural design space) and CMA-ES16, a black-box optimizer widely used in co-design. CMA-ES rolls out candidate robots in MuJoCo17 with the design space's controller (Mink18 for fixed-base / bimanual; an RL expert7 for legged robots), then biases the next generation toward higher-reward samples.

We give all methods the same task: generate a robot for a target trajectory and reward, control it, report the achieved reward.

Test-time compute scales the results

Like large language models that get better with longer thinking19, 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.

Test-time scaling results across three design spaces.
Higher reward, orders of magnitude faster. Across three design spaces and several reward functions, our Zeroth Order and Dynamics Self-Guidance variants both produce higher-reward designs as more parallel samples are allowed — and both reach CMA-ES-level quality in seconds rather than hours. For multi-trajectory bimanual optimization (rightmost), CMA-ES took 3+ hours; our model finished in under a minute.
Why so much faster? Two reasons multiply. (1) GPU parallelization over candidate designs. (2) Non-autoregressive diffusion evaluates an entire rollout in parallel rather than one timestep at a time. Typical learned dynamics models predict one step at a time; we predict the whole episode at once.

Real-World: Cloth Flinging on ALOHA

Original ALOHA vs Transformer Transformer-optimized ALOHA. Same task, same target trajectory, two robots. The optimized design unfolds the cloth where the original cannot.

To stress-test the whole system on hardware, we picked dynamic cloth unfolding on ALOHA2. Flinging20 punishes embodiment mistakes through both kinematics (the arm has to reach over the cloth at high speed) and unmodeled dynamics (aerodynamic drag, cloth weight, friction). The original ALOHA design struggled to even track the trajectory.

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:

  • Longer link lengths. Just enough to kinematically reach across the flinging trajectory, but light enough to stay within the Dynamixel motor torque envelope.
  • Inverted mounting. The arms hang upside-down behind the workspace, enabling an energy-efficient underarm swing instead of an overhead fling.
Joint speed comparison.
Smoother joint trajectories. Joint speed over the flinging motion. The optimized design has fewer peaks and a lower maximum than the original — the kinematic and dynamic gains transfer to hardware.

Limitations & Future Work

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 LLMs19, 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.

References

  1. Peebles, W., & Xie, S. Scalable diffusion models with transformers. ICCV, 2023.
  2. Zhao, T. Z., Kumar, V., Levine, S., & Finn, C. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv:2304.13705, 2023.
  3. Chi, C., Xu, Z., Pan, C., Cousineau, E., Burchfiel, B., Feng, S., Tedrake, R., & Song, S. Universal Manipulation Interface: In-The-Wild Robot Teaching Without In-The-Wild Robots. Proceedings of Robotics: Science and Systems (RSS), 2024.
  4. Song, J., Meng, C., & Ermon, S. Denoising Diffusion Implicit Models. ICLR, 2021.
  5. Furuta, H., Iwasawa, Y., Matsuo, Y., & Gu, S. S. A system for morphology-task generalization via unified representation and behavior distillation. arXiv:2211.14296, 2022.
  6. Gupta, A., Fan, L., Ganguli, S., & Fei-Fei, L. MetaMorph: Learning universal controllers with transformers. arXiv:2203.11931, 2022.
  7. Ha, H., Gao, Y., Fu, Z., Tan, J., & Song, S. UMI on Legs: Making Manipulation Policies Mobile with Manipulation-Centric Whole-body Controllers. Conference on Robot Learning (CoRL), 2024.
  8. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. Proximal policy optimization algorithms. arXiv:1707.06347, 2017.
  9. Rudin, N., Hoeller, D., Reist, P., & Hutter, M. Learning to walk in minutes using massively parallel deep reinforcement learning. CoRL, 2022.
  10. Zakka, K., Tassa, Y., & MuJoCo Menagerie Contributors. MuJoCo Menagerie: A collection of high-quality simulation models for MuJoCo. 2022. github.com/google-deepmind/mujoco_menagerie
  11. Dhariwal, P., & Nichol, A. Diffusion models beat GANs on image synthesis. NeurIPS, 2021.
  12. Xu, X., Ha, H., & Song, S. Dynamics-Guided Diffusion Model for Sensor-less Robot Manipulator Design. arXiv:2402.15038, 2024.
  13. Wang, T.-H. J., Zheng, J., Ma, P., Du, Y., Kim, B., Spielberg, A., Tenenbaum, J., Gan, C., & Rus, D. DiffuseBot: Breeding soft robots with physics-augmented generative diffusion models. NeurIPS, 2023.
  14. Liu, N., Li, S., Du, Y., Torralba, A., & Tenenbaum, J. B. Compositional visual generation with composable diffusion models. ECCV, 2022.
  15. Du, Y., Durkan, C., Strudel, R., Tenenbaum, J. B., Dieleman, S., Fergus, R., Sohl-Dickstein, J., Doucet, A., & Grathwohl, W. S. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. ICML, 2023.
  16. Hansen, N., & Ostermeier, A. Completely derandomized self-adaptation in evolution strategies. Evolutionary Computation 9(2): 159–195, 2001.
  17. Todorov, E., Erez, T., & Tassa, Y. MuJoCo: A physics engine for model-based control. IROS, 2012.
  18. Zakka, K. Mink: Python inverse kinematics based on MuJoCo. 2025. github.com/kevinzakka/mink
  19. Brown, B., Juravsky, J., Ehrlich, R., Clark, R., Le, Q. V., Ré, C., & Mirhoseini, A. Large language monkeys: Scaling inference compute with repeated sampling. arXiv:2407.21787, 2024.
  20. Ha, H., & Song, S. FlingBot: The unreasonable effectiveness of dynamic manipulation for cloth unfolding. CoRL, 2021.

Full reference list available in the paper. The selection above covers the citations used on this page.