Learning Action Priors for Cross-embodiment Robot Manipulation

Methodology Summary§

The paper proposes a two-stage framework to inject motion priors into Vision-Language-Action (VLA) models without requiring vision or language data in the first stage.

Stage 1: Motion Prior Pretraining§

A lightweight encoder-decoder action module is trained using flow matching on unconditioned action trajectories. Given a trajectory $\tau = (a_1, a_2, \dots, a_T)$, the encoder compresses the history into a latent $z$, and the decoder predicts future actions via a flow from noise to target action distribution. The encoder $E_\phi$ maps a historical sequence $\tau_{1:t}$ to a latent $z_t$. The decoder $D_\psi$ learns a conditional flow $\psi: [0,1] \times \mathcal{A} \times \mathcal{Z} \to \mathcal{V}$ (vector field) that generates actions $a_{t+1}$ from noise $\epsilon$ by solving the ODE $d\mathbf{x} = \psi(\mathbf{x}, z_t) dt$, optimized via the flow-matching objective:

$$ \mathcal{L}_{\text{FM}} = \mathbb{E}_{t, p_1(\mathbf{x}_1|z), p_t(\mathbf{x}_t|z)} \left[ \left\| v_t(\mathbf{x}_t|z) - \psi(\mathbf{x}_t, z) \right\|^2 \right] $$

where $v_t$ is the target vector field and $p_t$ is the probability path between noise and data.

Stage 2: VLA Alignment§

The pretrained decoder is reused as the action head in a VLA model (e.g., based on a VLM backbone). The encoder is used as a history compressor, outputting a single temporal context token per step. Training involves two losses: 1. Behavioral Cloning (BC) loss on $(o_t, l, a_t)$ pairs (image, language, action). 2. Latent distillation loss to align the VLM's latent space with the action prior's latent space, using KL divergence: $$ \mathcal{L}_{\text{KD}} = \text{KL}(p_{\text{VLM}}(z|o,l) \,||\, p_{\text{prior}}(z|\tau_{1:t-1})) $$

This avoids catastrophic forgetting of motion priors while allowing end-to-end fine-tuning.

Implementation Details§

• Action module: 4-layer transformer encoder + 4-layer flow-matching decoder (MLP-based).
• Stage 1 data: 1M+ trajectory snippets from diverse robots (Franka, UR5, etc.) without any visual or language annotations.
• Stage 2 VLA: Built on OpenVLA (7B) with LoRA fine-tuning for efficiency.
• Flow matching: 10-step ODE solver during inference (Euler method).

Code Snippet (PyTorch-style)§

# Stage 1: Motion Prior Training
class ActionPrior(nn.Module):
    def __init__(self, latent_dim=64, hidden_dim=256):
        super().__init__()
        self.encoder = TransformerEncoder(d_model=latent_dim, nhead=4)
        self.decoder = FlowMatchingDecoder(input_dim=latent_dim, hidden_dim=hidden_dim)
    
    def forward(self, traj_hist, noise, t):
        # traj_hist: (B, L, D) historical actions
        z = self.encoder(traj_hist).mean(dim=1)  # (B, D)
        v_pred = self.decoder(noise, z, t)  # (B, D)
        return v_pred, z

# Loss: MSE between predicted and target vector field
loss = F.mse_loss(v_pred, target_v)

# Stage 2: VLA with Prior
class VLAwithPrior(nn.Module):
    def __init__(self, vlm_backbone, prior_encoder, prior_decoder):
        super().__init__()
        self.vlm = vlm_backbone
        self.history_encoder = prior_encoder  # frozen later? (distillation)
        self.action_decoder = prior_decoder   # trainable
    
    def forward(self, image, text, hist_actions):
        vlm_features = self.vlm(image, text)
        hist_latent = self.history_encoder(hist_actions)  # from prior
        # alignment via KL loss
        kl_loss = kl_div(vlm_features, hist_latent)
        action = self.action_decoder(vlm_features)
        return action, kl_loss

Key Results§

• On 13 cross-embodiment tasks (simulation + real), method achieves 10-20% higher success rate over vanilla VLA fine-tuning.
• Faster convergence: reaches 80% success in 50K steps vs 100K+ steps without prior.
• Data efficiency: on a real-world task with only 50 demonstrations, success rate improves from 35% to 65%.
• Scaling Stage 1 data (from 200K to 1M trajectories) yields monotonic improvements.

Significance§

This work decouples motion learning from cross-modal learning, enabling large-scale pretraining of physical priors from cheap, unlabeled action data alone. The framework is model-agnostic and can potentially benefit any VLA architecture.