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

Huy Ha$^{1,2}$, C. Karen Liu$^{1}$, Shuran Song$^{1,2}$

$^1$ Stanford University, $^2$ Columbia University

Project Page | Paper | Video

Not all robots are created equal — but what if you could design one for the task at hand? Transformer Transformer is 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 of its designs for cloth flinging on a bimanual ALOHA platform — it cut tracking error by 73% and peak joint speed by 30% compared to the original ALOHA (see it in action).

This repository contains the full stack behind those results: RoboTokens (a unified tokenization of robot embodiments, states, and actions from MuJoCo models), procedural robot generation over parameterized design spaces, data generation with Mink and RL controllers, Transformer Transformer training and inference for cross-embodiment control and motion-conditioned robot co-design, CMA-ES baselines, and the Blender pipelines behind all our robot visualizations.

If you have any questions, please contact Huy Ha at huyha [at] stanford [dot] edu.

Table of Contents

If you just want to start running some commands while skimming the paper, you should get started here, which downloads checkpoints and data, then evaluates a pretrained Transformer Transformer on control and co-design.

• 🏃‍♀️ Getting Started
  • ⚙️ Setup
  • 📍 Checkpoints & Data
  • 📊 Evaluation
• 🗄️ Data Generation
  • 🤖 RoboTokens
  • 🕹️ Controllers: Mink v.s. RL
  • 🧬 CMA-ES
  • ⚡ Parallelization
  • 📱 UMI Data Collection & Processing
  • 💾 Data Format
• 🚂 Model Training
  • 🎛️ Hydra Configs
  • 🎮 Control Only
  • 🦾 Hardware Generation Only
  • 🧠 Control + Hardware Generation
• 🔭 Extending
  • 🌱 Adding New Design Spaces
  • 🧩 Adding New RoboToken Fields
  • 🚧 Limitations
• 📽️ Visualizations
  • 🎬 Robot Rollouts
  • 🌪️ Diffusion Processes

Citation

If you find this work useful, please consider citing:

@article{ha2026transformer,
  title={Transformer Transformer: A Unified Model for Motion-Conditioned Robot Co-design},
  author={Ha, Huy and Liu, C. Karen and Song, Shuran},
  journal={arXiv preprint arXiv:2607.25798},
  year={2026}
}

Code Acknowledgements

Model:

• The Transformer Transformer backbone is modified from Meta's DiT, which also carries code from GLIDE and MAE. These derived files (t2/model/core.py and the positional embedding utilities in t2/model/utils.py) remain under their upstream Attribution-NonCommercial licenses; the rest of this repository is MIT licensed.
• Rotation conversion utilities in t2/train/augment.py are copied from PyTorch3D (BSD).

Simulation & Control:

• Physics runs on MuJoCo and MuJoCo MJX from Google DeepMind.
• Our MJX environments and RL training loop build on MuJoCo Playground and Brax's PPO implementation, with data augmentations adapted from JaxRL.
• The differential-IK oracle controllers are built on mink. Big shout out to Kevin Zakka for his consistently excellent open-source robotics work — go give him a few stars ⭐
• The CMA-ES baseline uses pycma by Nikolaus Hansen, and evosax powers evolution strategies on GPU.

Robot Models:

• Robot assets under assets/mjcf/ are modified from MuJoCo Menagerie — ALOHA 2 (Trossen Robotics), Agility Cassie, ANYmal B & C (ANYbotics), Unitree A1/Go1/Go2/H1, UR5e & UR10e (Universal Robots), and the Wonik Allegro hand. Each directory keeps its original license file.
• The quadruped manipulator assets combine the Unitree Go2 with the ARX5 arm, plus the Fin-Ray gripper and GoPro mount from UMI, assembled originally for UMI on Legs.

UMI Data:

• Human demonstrations were collected with the Universal Manipulation Interface and iPhUMI, UMI's iPhone-tracked extension.