One-sentence Summary

Researchers at Stanford University and Columbia University propose Transformer Transformer, a diffusion transformer trained on a unified RoboTokens tokenization of robot embodiments, states, and actions, which functions as a reward-agnostic dynamics model and employs Dynamics Self-Guidance to steer embodiment diffusion toward high-value designs, achieving zero-shot optimization of unseen rewards and trajectories and enabling a fabricated ALOHA design that reduced tracking error by over 70%.

Key Contributions

• RoboTokens provides a unified tokenization of robot embodiments, states, and actions, enabling a single diffusion transformer to span diverse morphologies such as wheeled bimanual, quadruped, and humanoid robots.
• Dynamics Self-Guidance converts the reward-agnostic dynamics predictions of Transformer Transformer into reward-specific value scores at inference time, steering embodiment diffusion toward high-value designs without retraining.
• Across three design spaces, the model zero-shot optimizes unseen rewards and trajectories, and a fabricated optimized ALOHA design reduces real-world tracking error by over 70% compared to the original.

Introduction

The authors address robot co-design, the problem of automatically generating a robot’s physical embodiment to maximize performance on a given task. Prior work often relies on limited design representations, overfits to a single reward function, or separates the generator, critic, and controller into disjoint components that are difficult to optimize jointly. The authors propose Transformer Transformer, a diffusion transformer that unifies these three roles into a single model. It introduces RoboTokens, a flexible tokenization scheme that represents any articulated robot’s embodiment, states, and actions, and it is trained to model general dynamics across diverse embodiments. At inference time, the model zero-shot optimizes for user-specified reward functions by converting its reward-agnostic dynamics predictions into reward-specific guidance for the diffusion process, enabling fast, GPU-parallelized co-design.

Dataset

The authors construct a dataset that pairs diverse robot designs with target motion trajectories, enabling a model to evaluate how well a given design can track a desired skill. The dataset is built from three main components: a set of recorded end-effector motions, three combinatorial robot design spaces, and controller-generated state-action logs.

Key details of the dataset composition and sources:

• Target Motion Dataset: 76 motion trajectories are collected using an UMI gripper extension built on ARKit. These trajectories cover single end-effector skills such as tossing, screwing, opening drawers, and scrubbing. The split uses 56 trajectories for training and 20 for validation. Bimanual dishwashing demonstrations from UMI are also included for the mobile bimanual design space.
• Kinematic Design Space: Based on the Trossen ViperX 300S, this space varies mounting pose, degrees of freedom count, joint orientations, and link lengths. It tests reasoning about kinematic reachability for fixed-base arms.
• Dynamic Control Design Space: A quadruped manipulator space that varies leg design (serial vs. spring-loaded linkage), arm mounts (fixed vs. sliding rail), and lengths for the torso, legs, and arms. This space stresses whole-body control performance, stability, and inertia trade-offs.
• Task Complexity Design Space: A mobile bimanual space with spine variations (fixed, sliding, bending) and length changes. It captures long-horizon, coordinated bimanual tasks like dishwashing.

How the data is generated and processed:

• For non-legged robots (fixed-base and mobile bimanual spaces), the authors use Mink, a differential inverse kinematics library for MuJoCo. A fixed 80 ms look-ahead time compensates for joint position control lag, with terms for tracking, posture regularization, and damping.
• For legged robots, reinforcement learning policies are trained to serve as whole-body controllers. Instead of training one policy per robot, one policy is trained per discrete design variation (e.g., leg type) and receives continuous variations (e.g., leg lengths) as observations. This yields 128 RL expert policies covering all combinations of 7 binary design choices. These policies generate the state and action data used for both data collection and design validation.

The resulting dataset provides state and action RoboTokens for each robot design tracking a target motion, forming the training foundation for the co-design evaluation model.

Method

The authors propose Transformer Transformer, a unified framework designed to model the joint distribution of robot embodiment and dynamics. The overarching pipeline involves demonstrating tasks via end-effector trajectories, generating optimized robot designs using reward-guided diffusion, and validating these designs with a cross-embodiment controller.

As shown in the figure below:

To enable this, the authors first introduce RoboToken, a tokenization scheme that converts any robot description into a learning-ready representation. RoboToken captures both time-invariant embodiment attributes and time-varying dynamics. The tokenizer parses robot simulations to extract attributes for links, joints, and motors, concatenating them into continuous-valued vectors. Embodiment tokens represent physical parameters such as primitive type, size, and inertia, while dynamics tokens capture states and actions at specific timesteps. To support diverse kinematics, tokens include IDs that point to connected components, allowing the representation of arbitrary connectivity and passive joints.

The tokenization process is illustrated in the figure below:

This representation is highly efficient. As shown in the figure below, RoboToken sequences are significantly more compact than traditional text-based formats like MJCF, often requiring 27 to 110 times fewer tokens.

To train a model over these tokens, the authors employ a Diffusion Transformer (DiT) architecture capable of handling variable sequence lengths of high-precision continuous data. For each token, a token-type-specific linear projection maps the vector into the DiT latent space, augmented with learned embeddings for token IDs. Since the domain exhibits SE(2) equivariance due to gravity, planar transform augmentations are applied. Each token type sequence is padded to a maximum length and concatenated into a single multi-modal sequence. The model is trained using the DDIM scheduler, with episode timesteps subsampled to form fixed-length sequences.

The unified architecture and its applications are detailed in the figure below:

The model is trained with different masked modeling schemes to support two primary applications. First, for Motion-to-Robot Optimization, the model simultaneously diffuses embodiment and dynamics tokens conditioned on a target motion trajectory. This non-autoregressive formulation allows the model to assess newly generated designs by evaluating reward functions that depend on both embodiment and dynamics, avoiding the error accumulation typical of autoregressive models. An embodiment generator runs parallel diffusion processes and ranks candidates using a specified reward function. Second, for Cross-Embodiment Control, the model acts as an embodiment-aware controller. Conditioned on the embodiment, current state, and target motion, it predicts expert actions. RoboToken's flexibility allows the model to handle heterogeneous state and action spaces by varying the number of tokens in the context, enabling generalization to unseen embodiments.

To improve sampling efficiency in challenging optimization landscapes, the authors introduce Dynamics Self-Guidance. This first-order diffusion guidance method utilizes gradients from the model's noisy reward predictions to optimize noisy embodiment predictions. At each diffusion step, the DiT predicts noise to denoise tokens. Even if the full diffusion process is not differentiable, per-step reward gradients can guide samples toward higher-reward candidates. This is incorporated following classifier-guided DDIM, allowing the model to optimize for unseen reward functions without training additional networks.

The model's capability to generate diverse designs across different optimization landscapes is shown in the figure below:

Experiment

The evaluation covers multiple robot design spaces, including fixed-base arms, quadrupeds, bimanual mobile manipulators, and humanoids, to assess how well a diffusion-based tokenization and generation framework unifies embodiment representation, produces physically plausible designs, and optimizes for unseen reward functions. The model generates designs that remain on the training manifold, respects cross-attribute coherence, and supports zero-shot optimization with parallelized evaluation that outperforms CMA-ES in both speed and reward. It also demonstrates diffusion composition for multi-trajectory optimization, test-time compute scaling, and the ability to self-validate designs through a unified controller, with a real-world ALOHA setup confirming improved tracking performance and reduced joint speeds.