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

Start with a target end-effector trajectory and a reward function. The model generates a robot’s links, joints and motors together with predicted states and actions for that motion. A reward can then score both hardware and behavior, such as tracking accuracy versus torque use. Unlike π₀.₅, the design itself is an output, not a fixed platform on which to predict actions.

The small configuration uses 256 hidden dimensions, 8 layers and 4 heads; the large version uses 512 dimensions, 12 layers and 8 heads. The unified-model experiments report 11.6M and 63.6M parameters respectively. Both tasks can share weights, but the main design-optimization comparisons train a separate motion-to-robot model per design space. The paper does not demonstrate one universal checkpoint solving every robot family and task.

Embodiment tokens describe links, fixed joints, sliding/rotating joints, ball joints and motors. Their attributes include primitive geometry, mass, inertia, joint ranges and actuation properties. Dynamics tokens describe component states and controls at particular times. Target-pose tokens reference the end-effector and timestep they constrain.

These are vectors of physical attributes, not language tokens. IDs connect a joint to its links and a motor to its joint, allowing different connectivity and degrees of freedom. Canonical transforms reduce redundant coordinate descriptions. Eleven Menagerie robots fit into 28–101 embodiment tokens, 27–110× more compact than their MJCF text tokenizations. That count excludes a full motion rollout.

For design generation, the network explicitly predicts special padding values, allowing variable component counts. In control mode, padding is masked. Continuous diffusion allows later components to influence earlier ones during denoising and provides gradients for design optimization.

Training learns dynamics rather than a critic tied to one reward. At inference, score the predicted embodiment, states and actions with a differentiable reward, then backpropagate that score to the noisy embodiment tokens. Classifier-guided DDIM uses the gradient to shift the next denoising update toward better designs. The model supplies its own dynamics predictions; guidance needs neither a separate learned dynamics network nor a differentiable simulator.

The zeroth-order baseline samples candidates in parallel, ranks their predicted rewards and returns the best. Self-guidance also improves candidates during sampling before ranking them. For example, adding a torque or size penalty can favor a smaller, lighter design without retraining the model. “Unseen reward” means a new scoring function over modeled quantities, not arbitrary new physical capabilities.

To design for several demonstrations, diffusion composition averages noise predictions conditioned on different trajectories at each step. Candidate designs and sampled episode times are evaluated in parallel. This avoids sequentially simulating every timestep during search, although simulation or hardware evaluation is still needed to check the resulting design.

Motion targets: 76 single-end-effector demonstrations captured with an ARKit-based UMI extension, split into 56 training and 20 validation trajectories. Bimanual experiments additionally use UMI dishwashing demonstrations, with 26 held-out validation trajectories.

Supervision: procedural robot descriptions plus simulator-recorded states and expert actions while tracking those targets. Mink differential inverse kinematics controls fixed-base and mobile bimanual robots. Quadrupeds use 128 PPO experts, one per combination of seven binary design choices; each expert observes continuous design variations. Perturbed actions diversify rollouts while unperturbed expert actions remain the supervision targets.

Episodes shorter than 100 steps are removed. Brief CMA-ES searches enrich fixed-base training data with more reachable designs. Each design is used for ten episodes. These are simulated behavior records, not manually annotated captions or segmentation masks.

Training subsamples eight timesteps: spread across an episode for design, consecutive for control. Planar position/heading augmentation respects gravity. The default recipe uses batch size 64, learning rate 5 × 10⁻³ with 500 warmup steps, and 50 epochs of 16,384 steps; the large model lowers the learning rate to 10⁻⁴. At inference, design uses 100 DDIM steps for accurate dynamics, while control uses five.

Across the three design spaces, parallel sampling and self-guidance achieve competitive or better rewards than CMA-ES with much shorter search times. Main comparisons validate designs with the design space’s expert controller, which separates hardware quality from learned-controller errors. Unified control is evaluated separately; on generated quadrupeds, its reward correlation with the RL oracle is Pearson r = 0.53.

For real cloth flinging, an optimized ALOHA used longer links and inverted mounting. Tracking position error fell from 13.0 to 3.5 cm and maximum joint speed from 2.57 to 1.82 rad/s. This validates one fabricated design, rather than proving all generated designs work on hardware.

Generation stays within the learned design distribution: quadruped-only training does not invent hexapods or extend link lengths beyond the training range. Geometry is primitive-based; scene context, deformables and structural-strength rewards are outside the demonstrated representation. Prediction errors limit guidance and controller transfer, gains plateau with more search, and each quadruped expert costs about 16 A100 hours to train.