Latest / Embodied AI 101 / RoboTTT: End-to-End Robotic Assembly with Test-Time Training

Transcript

RoboTTT: End-to-End Robotic Assembly with Test-Time Training · Embodied AI 101

• RoboTTT: End-to-End Robotic Assembly with Test-Time Training. Robotic assembly — tasks like joining gears, plugging together circuits, or snapping parts onto a model — is a long-standing challenge in robotics. Unlike short pick-and-place tasks, an assembly procedure can take minutes, involves many steps, and requires exquisite precision and force control. Traditional industrial robots execute assembly with carefully hand-engineered kinematic motions and force control loops. By contrast, an end-to-end learned policy must coordinate vision, language instructions, proprioception, and control over an extended interaction, often without pauses between sub-tasks. In practice, most learned “robot foundation models” today collapse this complexity by only looking at the current or a few past frames of sensor data. They essentially have no memory of what was done seconds ago. This limits them sharply: long-term dependencies in multi-stage tasks go unhandled. Enter RoboTTT, a new framework from NVIDIA’s GEAR Lab that scales the robot’s effective memory to thousands of timesteps while keeping inference fast. At 8,000 steps of context (on the order of five minutes of real operation), RoboTTT unlocks behaviors unseen before in a single learned policy: it can learn one-shot imitation from a human demonstration video, adapt its behavior on the fly if things go wrong, and robustly complete multi-stage assemblies that short-horizon policies cannot. In benchmark physical assembly tasks, RoboTTT dramatically outperforms existing baselines. For example, on a five-minute, ten-stage assembly sequence none of the prior policies ever complete, RoboTTT succeeds about 80% of the time — an 87% relative gain over the best prior model. Remarkably, simply increasing the pre-training context length (from 1,000 to 8,000 steps) boosts performance by 62–63%, suggesting that context length itself is a key new scaling axis for robot policies. Below we unpack how RoboTTT works, why it matters, and what new capabilities it delivers. The Long-Horizon Assembly Problem. Consider a task like assembling a toy car from many tiny parts (see Figure below). The robot must first pick up a chassis, then locate wheels, snap axles in place, install gears, bolt a motor on, attach decals, and so on. Each of these stages requires fine adjustments and contact-rich maneuvers. Critically, progress in later stages depends on what was done earlier – if a wheel didn’t seat correctly, the car won’t roll even if every later step is done right. Human workers naturally carry “memory” of what has happened: they see partially assembled products and plan future actions accordingly. Most robot policies, however, lack that long memory. Modern “foundation” models for robotics — like vision-language-action transformers (e.g., NVIDIA’s GR00T models) — usually condition on the current observation or maybe the last few frames. In other words, at each moment the policy decides on actions based mostly on the latest camera image and proprioceptive readings. That fixed-window context (often only a couple of time steps) severely limits reasoning about history. If the task sequence is long, the policy will “forget” what it achieved at the 10-second mark by the time it reaches the 300-second mark. This limitation shows up as poor closed-loop performance on multi-stage or long-horizon tasks. Even if the policy was perfect at each short step, small errors accumulate and the policy has no memory-based corrective mechanism. In practice, as one NVIDIA report notes, conventional models only unlock “one-shot in-context imitation from human video, on-the-fly policy improvement, robustness to perturbations, and stronger performance on multi-stage, long-horizon tasks” once the history length is extended dramatically. RoboTTT tackles this by fundamentally changing how context is handled. Instead of finite cached history, it uses Test-Time Training to compress arbitrarily long streams of observations into an internal memory (a set of fast weights) that never grows during inference. In effect, the robot “remembers” its entire interaction with the world, all while keeping per-step latency constant. Test-Time Training: Internalizing History. At the heart of RoboTTT is the idea of Test-Time Training (TTT) – an approach where the model continues to adapt its internal parameters as it acts. In traditional neural nets, a network has fixed “slow” weights learned during training, and those weights are never altered during inference. RoboTTT augments this with a smaller set of fast weights that are updated on the fly by gradient descent, using the robot’s own incoming data as a self-supervised training signal. The result is a model whose “memory” of history is encoded in these fast weights. Crucially, only one (or a fixed small number of) gradient update is done per timestep, so the inference cost stays constant no matter how long the robot has been running. Here’s roughly how it works: at each time step $t$, the robot takes in its observation (vision frames, proprioception, etc.). The policy has an inner network (the fast-weights module) with parameters $W_{t-1}$. The model computes a small “self-supervised” loss on the current data (for example, a reconstruction or contrastive objective that doesn’t need human labels). Then the fast weights are updated by one gradient step:. $W_t = W_{t-1} - \eta \nabla \ell_{\text{self}}(W_{t-1}; \text{current data})$. This single-step update encodes the new observation into $W_t$. Immediately after that, the updated fast-network $f(,\cdot;W_t)$ is used to help compute the policy’s action output for step $t$. In other words, the streaming data is written into the fast weights each step, and those weights act as a dynamic memory read-out. Because we always use freshly updated weights, the model’s output can be conditioned on all past data, no matter how far back, without explicitly reprocessing every past frame. Only the fast weight state vector is carried forward. Figure: Illustration of RoboTTT’s memory mechanism. At each time step, the incoming observation is used both to update the fast-weight memory (via one gradient descent step on a self-supervised loss) and to generate the action. As more data streams in, all of it is progressively compressed into $W_t$. This fast-weight memory idea is reminiscent of classic recurrent networks or Neural Turing Machines, but with a modern twist. Because the update is a true gradient step, the fast weights are learned (meta-learned) to be a good memory. In fact, during training RoboTTT samples whole trajectories and meta-learns the optimal initial fast weights $W_0$ so that fast-weight SGD will accumulate useful information. This “learn to learn” within each trajectory is inspired by Model-Agnostic Meta-Learning (MAML): the parameters $W_0$ and the projection matrices in the fast-weight layers are themselves optimized by the usual outer-loop learning. Importantly, RoboTTT does not destroy what the base model already knows. The fast weights are injected into the existing transformer policy via a gated mechanism. Each TTT layer’s output is combined with the original attention output using a learned tanh gate. At initialization this gate is nearly closed (tanh near zero), so the pretrained backbone functions normally. Over training it learns how much the new fast-weights branch should influence the computation. In equations, the combined output is $o = \tanh(\alpha)\odot o_{\text{TTT}} + o_{\text{attn}}$, blending the new memory-augmented signal with the original: this prevents the fast-weight read/write from overwhelming the pre-trained model. By integrating multiple such TTT layers into a vision-language-action (VLA) policy (in this case NVIDIA’s GR00T N1.7 model), RoboTTT effectively becomes a recurrent sequence model whose hidden state is the fast-weight memory. Because it is always being written to and read from, the policy can condition its decisions on thousands of steps of history, all at constant per-step cost. As the NVIDIA team puts it, “the stream is compressed into weights, not cached”. Model and Training Details. Concretely, the backbone VLA model is GR00T N1.7 – a multimodal transformer previously developed by NVIDIA for general-purpose robot control. Its input at each step includes a learned “register token” $R_t$, a vision-language token $\Phi_t$ (embedding robot’s camera + any instruction text), encoded proprioceptive token $q_t$, and a noised action token $\tilde{A}_t$ which the diffusion head denoises. RoboTTT inserts TTT layers into the diffusion transformer (DiT) action head, after the attention layers. This means the regular attention mechanism processes step-specific information, while the TTT layers process the temporal stream across steps. Training RoboTTT end-to-end requires special care, since naively backpropagating through thousands of steps would exhaust GPU memory. The authors use Truncated Backpropagation Through Time (TBPTT): training trajectories are split into manageable segments, and gradients are only propagated within each segment. Crucially, the fast-weight state is carried from segment to segment, but gradients into it are cut at the boundaries. In practice this lets them train on sequences length 8,000 by only storing a few hundred steps of activations at a time. For the action output, RoboTTT employs a diffusion (flow-matching) objective, as in the original GR00T. Each action is represented in “denoising” form and the network learns to predict the clean action from a noisy version. A subtlety emerged: if one applies the same diffusion noise level to an entire long trajectory, the model sees uniformly easy or hard noise across all steps, which destabilizes training. RoboTTT introduces Sequence Action Forcing: it samples a different noise level $\tau_t$ for each action chunk independently (at each timestep), so that the difficulty fluctuates and the model practices denoising at all levels in every sequence. The overall loss is the average flow-matching loss per step. During training, the fast-weights updates (the inner TTT loop) happen at every step, and the outer optimization adapts both the backbone parameters and the initial fast-weight $W_0$ to minimize the average trajectory loss. An additional clever trick involves leveraging imperfect demonstrations or failed robot actions. RoboTTT can mask the loss at selected timesteps, treating those as “pure context” that should still update the fast weights but carry no imitation target. This allows learning from heterogeneous sequences. In one scenario resembling DAgger, the model interleaves imperfect robot attempts with correct human corrections: the robot’s suboptimal moves are used only to update the fast-weight memory (no loss), while only the corrective moves contribute to the loss. During inference, the fast weights then encode a mapping from failure contexts to corrective actions, effectively teaching the policy to self-correct when it goes off track. Real-Robot Assembly Experiments. To test RoboTTT’s long-term control, the authors used a YAM robot setup (Yet-Another-Manipulator arms) and designed three challenging bimanual assembly tasks. Each task involves two arms interacting with many parts under camera guidance:. Pup-Go-Car: Assemble a plastic “sports racer” model from ~20 parts (wheels, chassis, motor, shell). This took about 5 minutes of real time to complete. Circuit: Connect components (like LEDs, wires, a battery holder) into an electrical circuit. Importantly, the exact layout is randomized – at test time the robot gets only the instruction “assemble circuit” plus (in some trials) a single human video demonstration of the new configuration. Completing this task requires inferring the goal from context. Gear-Bot: Build a remote-control car/robot with interlocking gears and wires (5 minutes). This is especially difficult: even small misalignments (like a gear slightly askew) will prevent the bot from functioning. Each task is multi-stage (peg insertions, alignments, screw assembly, etc.) and involves fine contact-rich maneuvers. The robot policy receives visual inputs (cameras above the workbench) and proprioception, plus a language instruction (“assemble X”). For baselines, the authors compare against several memory-limited policies: the original GR00T N1.7 with no history (single-step context), a “Hist” variant with one extra previous frame as context (short fixed history), and GDN (Gated DeltaNet), a state-of-the-art recurrent memory model. None of these baselines can accumulate thousands of steps of memory. Results: Mastering Assembly End-to-End. The quantitative outcomes are striking. Across all three assembly tasks, RoboTTT achieves the highest completion rates. In aggregate, RoboTTT’s policies achieve about 79% task completion on these long-horizon assemblies, compared to far lower scores for the baselines. This 79% is an 87% relative improvement over the single-step baseline policy and also 41% better than the best short-history baseline. In concrete terms, RoboTTT is often the only model that ever finishes the hardest assemblies. For example, on the five-minute Gear-Bot task, RoboTTT successfully finishes the task on 2 out of 10 trials, whereas none of the baselines ever record a full success (0/10). Similarly, on the Circuit task with random configurations, RoboTTT completes the task twice as often as the best baseline policy. Not only does RoboTTT outperform in mean score, but scaling up the context length itself directly boosts performance. Training the model with an 8,000-step context (instead of 1,000 steps) raises its completion scores by roughly 63%. In other words, a longer-trained memory yields dramatically better adulting of the tasks. This increase – steady gains in closed-loop performance with longer context – was demonstrated “for the first time” in this work, suggesting that context is a new scalable axis akin to model size or data quantity. As a simple illustration: on these assembly tasks the completion rate for RoboTTT rose from about 30% to over 70% when going from 1K to 8K context during training. Importantly, RoboTTT’s strengths go beyond raw task scores. It exhibits new capabilities enabled by the long context:. One-Shot Imitation from Human Video: The Circuit task included a test of generalization: at test time, the robot sees a previously unseen configuration of the circuit. It is only given one human video of that new configuration (no textual hints beyond “assemble circuit”). RoboTTT succeeds in about 65% of trials under this condition. In contrast, the baseline policy (GR00T N1.7 with short memory) scores only ~33% and never fully solves the unseen circuit. This gap – RoboTTT doubling the average completion rate – shows that the policy can indeed read the goal from context (the human demo) and replicate it. No other single model in this study could do one-shot imitation from an unseen demo. On-the-Fly Adaptation: Thanks to its test-time learning, RoboTTT can adapt mid-task if something goes wrong. In perturbation experiments (for example, nudging the assembly parts during execution), the model automatically “remembers” the pre-perturbation state via its fast weights and recovers from errors. On the Pup-Go-Car task, RoboTTT successfully recovered and completed the assembly in roughly 75–90% of trials even after artificial disturbances. The competitor models, lacking such a dynamic memory, often failed to recover at all. This real-time adaptation is essentially learning from its own mistakes as it operates – a form of online correction enabled by the fast-weight updates (somewhat analogous to DAgger-style recovery in classic IL). Robust Closed-Loop Control: Across all tests, RoboTTT consistently avoids “stage confusion” (mixing up sub-tasks) and achieves the fine alignments needed for hidden connections. In the data, it is the only method that ever registers a “full success” on every trial of the very difficult tasks, and it makes strategic corrections instead of giving up. Figure: Performance of RoboTTT vs baselines. (Left) Task completion percentages on each assembly. (Right) Self-improvement: success improves steadily as pretraining context grows. RoboTTT (green) dominates in all settings. These results were measured with no speed-up tricks or human interventions. The tasks were run in real time on physical hardware, and every pick, grasp, and alignment had to succeed under real friction, slight misplacements, and sensor noise. RoboTTT’s high 79% average completion means that in most trials it truly assembled the object from scratch, end to end – something no prior end-to-end policy had ever done in such complex assembly benchmarks. Discussion: Context as a Scaling Axis. The success of RoboTTT has broad implications. It demonstrates that temporal context is as important for robots as it is for language models. In NLP, we’ve seen how doubling context length lets GPT understand longer texts or documents. Similarly, these experiments show that doubling context (from 1K to 8K steps) nearly doubled the robot’s performance on long-horizon tasks. In fact, the authors report that simply training with longer context yields monotonic, unsaturated gains in real-world task success. This marks context length as a new “scaling axis” for robotic foundation models, on par with model size or data diversity. RoboTTT also introduces embedded learning in the action loop. By compressing history into fast weights, it effectively bypasses the quadratic cost of self-attention on long sequences. The fast-weight module serves as a powerful and compact memory: the researchers describe it as a “highly expressive compressor” that filters out redundant repetition and highlights relevant changes over thousands of steps. In practical terms, the robot policy can concentrate on the new, salient information in each frame (like a new part appearing) and rely on its internal state to recall the rest. This paradigm also blurs the line between offline training and online adaptation. RoboTTT is first pretrained on many hours of assembly demonstrations, but once deployed it continues to learn self-supervised from its observations. It is “aware” of its own history in a way conventional policies are not. This enables behaviors like the one-shot imitation and on-the-fly correction we saw. Conceptually, it’s as if the robot had a tiny “inner loop” of learning running in parallel with its control loop. For roboticists and ML researchers, RoboTTT opens new doors. One can imagine taking other vision-language robot models and equipping them with TTT layers to extend their reach. The fact that no extra sensors or datasets were needed – only clever model architecture and training – is appealing. In practice, using RoboTTT would mean collecting long (8K-step) demonstration trajectories (as the authors did in simulation or on the real YAM robot) and then fine-tuning the model with TBPTT. NVIDIA’s code release (and videos) should make that process accessible. That said, there are still challenges. The current policy was specialized on assembly tasks; generalizing RoboTTT to arbitrary tasks will require large and diverse datasets (analogous to how GPT was scaled from web text). The model’s capacity and fast-weight architecture have limits; extremely chaotic or unstructured histories might saturate the memory. Moreover, while RoboTTT ensures inference time is constant per step, the training is heavier because of TBPTT and the diffusion objective. Future work might explore making the fast-weight update more efficient or combining this with world-model planning. Nevertheless, RoboTTT stands out as a landmark demonstration of end-to-end assembly. It is, to our knowledge, the first single neural policy to autonomously put together multi-part objects in five real-time minutes with no scripted fallback. In doing so, it shows us that robotic learning models can be truly memory-ful. As foundation models evolve, incorporating mechanisms like Test-Time Training will likely become standard practice for any robot expected to ponder its long history and learn on the fly. References: The above is based on the Nvidia GEAR Lab’s RogueTTT paper and project page, which introduce and evaluate this approach on real robotic assembly benchmarks.