RoboTTT: Context Scaling for Robot Policies

Core Method: Test-Time Training with Fast Weights§

The key innovation is treating the parameters of a neural network as a recurrent state that can be updated via gradient descent during both training and inference. This compresses long observation histories into weight space, enabling extremely long context (up to 8K timesteps) without increasing inference latency.

Let $\theta$ be the base policy parameters (e.g., a Vision-Language-Action model). At each timestep $t$, the policy receives observation $o_t$ and outputs action $a_t$ using a set of fast weights $\phi_t$ that are updated online:

$$\phi_t = \phi_{t-1} - \eta \nabla_{\phi} \mathcal{L}(a_{t-1}, \hat{a}_{t-1})$$

where $\mathcal{L}$ is a loss function (e.g., action prediction MSE) and $\eta$ is a learning rate. The base parameters $\theta$ are shared across timesteps, while $\phi_t$ act as a compressed memory of the recent history.

Training Recipe: Sequence Action Forcing + Truncated BPTT§

To scale training context length, the authors propose:

1. Sequence Action Forcing: Train on long sequences (up to 8K steps) by unrolling the model and applying the fast weight update at each step. The loss is computed only on the final action prediction, not every step.

2. Truncated Backpropagation Through Time (TBPTT): To avoid memory blowup, backpropagation is truncated: gradients flow only through a fixed window (e.g., 32 steps) while the fast weights carry information across windows.

Formally, for a sequence of length $T$, the update is:

$$\phi_{t+1} = \phi_t - \eta \nabla_{\phi} \mathcal{L}(\pi(o_t; \theta, \phi_t), a_t)$$

$$\mathcal{L}_{\text{total}} = \sum_{t=T-W+1}^{T} \mathcal{L}_t$$

where $W$ is the truncation window size.

Architecture Detail§

The base policy $\pi$ is a transformer-based VLA model (e.g., based on RT-2). Fast weights are added as a small adapter network (e.g., a single linear layer or a tiny MLP) that modulates the transformer's activations. During training, the fast weights are initialized randomly and updated via gradient descent. At inference, the fast weights are initialized to zero or from a learned prior.

PyTorch Code Snippet (Simplified Training Loop)§

import torch
import torch.nn as nn

class FastWeightPolicy(nn.Module):
    def __init__(self, base_policy, fast_weight_dim):
        super().__init__()
        self.base = base_policy  # VLA transformer
        self.fast_net = nn.Linear(fast_weight_dim, fast_weight_dim)  # learnable
        self.fast_weight = nn.Parameter(torch.zeros(fast_weight_dim))  # recurrent state
    
    def forward(self, obs, lang, phi):
        # phi is current fast weight
        features = self.base.encode(obs, lang)
        phi_update = self.fast_net(phi)
        phi_new = phi + phi_update  # simple additive update
        action = self.base.decode(features + phi_new)
        return action, phi_new

# Training loop (truncated BPTT)
for seq in dataloader:  # seq of length T
    phi = torch.zeros(fast_weight_dim, device=device)
    for t in range(len(seq)):
        obs, lang, action_target = seq[t]
        action_pred, phi = model(obs, lang, phi)
        loss = mse_loss(action_pred, action_target)
        if t >= T - W:  # only accumulate gradients in window
            loss.backward(retain_graph=True)
    optimizer.step()
    optimizer.zero_grad()

Key Results§

• Context scaling: Steady improvement in closed-loop success rate as pretraining context length increases from 256 to 8K steps (62% improvement).
• One-shot imitation: After seeing a single human video demonstration (converted to actions via inverse kinematics), RoboTTT can perform the task without finetuning.
• On-the-fly improvement: During a long-horizon task, the policy can correct its behavior based on recent failures.
• Real-robot tasks: 87% improvement over single-step baseline, and first to complete a 10-stage assembly task (5 minutes).