Learning to Move Before Learning to Do: Task-Agnostic pretraining for VLAs

Decomposition Hypothesis§

The paper posits that VLA models conflate two distinct abilities: (1) physical competence (how to move) and (2) semantic alignment (what to do). Only semantic alignment requires language supervision. This motivates a two-stage framework, TAP (Task-Agnostic Pretraining).

TAP Framework§

Stage 1: Task-Agnostic Pretraining

• Input: Unlabeled interaction data (off-task trajectories, autonomous play)
• Objective: Self-supervised Inverse Dynamics (ID) to predict action $a_t$ given consecutive observations $o_t, o_{t+1}$:

$$\mathcal{L}_{ID} = -\log p_\theta(a_t | o_t, o_{t+1})$$ ID is trained to compress observations into a 64-dimensional latent $z_t$, then predict action via a small MLP.

Stage 2: Language Grounding

• Freeze the motor backbone from Stage 1.
• Train a lightweight language-conditioned policy head on minimal expert demonstrations to map instructions and observations to actions.

Implementation Details§

• Observations: 256x256 RGB images, optionally depth.
• Observation encoder: Clipped ResNet-18 (no gradual warmup needed due to frozen backbone).
• ID latent dimension: 64.
• Action space: 7-DOF joint positions (WidowX) or 4-DOF (SIMPLER).
• Training: Stage 1 uses AdamW, lr=1e-4, batch size 64 over 50k steps.

Code Snippet (PyTorch-style)§

class InverseDynamicsModel(nn.Module):
    def __init__(self, obs_encoder, action_dim=7):
        super().__init__()
        self.encoder = obs_encoder  # ResNet-18, frozen after pretraining
        self.projector = nn.Sequential(
            nn.Linear(512*2, 256),
            nn.ReLU(),
            nn.Linear(256, 64)  # latent z
        )
        self.action_head = nn.Sequential(
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim)
        )

    def forward(self, obs_t, obs_tp1):
        feat_t = self.encoder(obs_t).flatten(start_dim=1)
        feat_tp1 = self.encoder(obs_tp1).flatten(start_dim=1)
        z = self.projector(torch.cat([feat_t, feat_tp1], dim=-1))
        action = self.action_head(z)
        return action

Experimental Results§

• SIMPLER benchmark: TAP (trained on 5k expert demos + 50k unlabeled trajectories) matches performance of full VLA trained on >1M labeled trajectories (e.g., 72% vs 62% BC on the coffee task).
• WidowX real robot: Under camera perturbations (rotation, shift, blur), TAP retains 25% success, while baselines (Octo, RT-2-X) drop to 0%.

Key Equations§

• Inverse Dynamics loss: $\mathcal{L}_{ID} = \mathbb{E}_{(o_t, o_{t+1}, a_t)\sim\mathcal{D}_{unlab}}[\|\hat{a}_t - a_t\|^2]$ (mean squared error for continuous actions).
• Language Grounding loss: $\mathcal{L}_{LG} = \mathbb{E}_{(o_t, l, a_t)\sim\mathcal{D}_{expert}}[\|\pi_\phi(o_t, l) - a_t\|^2]$ where $\pi_\phi$ is a small MLP taking frozen features and instruction embeddings.

Takeaways§

• Decoupling motor priors from language alignment is both sample-efficient and robust to distribution shift.
• The ID objective captures temporal structure and physical dynamics without any task or language labels.
• The framework is scalable: unlabeled data from failed trials or autonomous play is cheap to obtain.