Task adaptation of Vision-Language-Action model: 1st Place Solution for the 2025 BEHAVIOR Challenge

AI summary

Overview

Research area: Robot learning / embodied AI — specifically Vision-Language-Action (VLA) models for long-horizon household manipulation.

Technical level: Advanced. The paper assumes familiarity with flow matching, diffusion, transformer attention patterns, KV caches, and action chunking, though the prose is largely readable.

Scope: A competition report describing the architecture, training, and inference modifications that produced the 1st place entry in the 2025 BEHAVIOR Challenge, achieving a 26% q-score across 50 long-horizon household tasks.

What This Paper Is About

The BEHAVIOR Challenge asks a single robot policy to perform 50 different long-horizon household tasks — from turning on a radio to cooking a hotdog — in photo-realistic simulation, with two arms, mobile navigation, and three cameras. The core problem is that these tasks run for thousands of timesteps (6.6 minutes on average, up to 14 minutes), so small errors compound, many visually identical states mean different things, and the training data contains only successful demonstrations with no recovery examples. The authors build on the Pi0.5 vision-language-action model and modify it to handle these specific difficulties.

Key Contributions

1. Correlated noise for flow matching. Instead of sampling independent Gaussian noise during training, the model samples noise from a distribution matching the empirical action covariance, regularized with a shrinkage parameter of 0.5. This makes denoising steps more uniform in difficulty and enables correlation-aware inpainting.

2. Correlation-aware soft inpainting at inference. A rolling window predicts 30 actions, executes the first 26, and saves 4 as initial conditions for the next prediction. Corrections applied to the constrained dimensions are propagated to free dimensions using the learned action correlation matrix, rather than applied as hard constraints.

3. Learnable mixed-layer attention. A KV cache transformation lets each of 18 action expert layers attend to learned linear combinations of all VLM layers, instead of a fixed layer-to-layer mapping. The model is initialized to Pi0.5's behavior (identity weights, zero bias) and learns deviations.

4. System 2 stage tracking and correction rules. The model predicts the current task stage through a linear classifier on the VLM output, filters noisy predictions with majority voting over the three most recent predictions, and feeds the stable stage estimate back as context — resolving visually ambiguous non-Markovian states. Simple heuristics recover from failure modes such as closing an empty gripper.

Main Findings

• Leaderboard result: 26% q-score across all 50 tasks on both the public and private leaderboards.
• Task scale: The benchmark features 50 tasks averaging 6.6 minutes per episode, with the longest tasks averaging 14 minutes. Each task is evaluated over 10 episodes with randomized initial conditions.
• Stage prediction accuracy: System 2 stage prediction achieves approximately 99% accuracy on training data.
• Gripper correction impact: The gripper recovery rule alone approximately doubled the success rate in selected tasks where grabbing objects was a common failure mode.
• Emergent recovery: After training on 50 tasks, the robot sometimes learned to open a gripper even without explicit recovery demonstrations in the data, though these cases were rare and did not substantially improve total success rate.
• Action compression: Compressing 26 predicted actions into 20 execution steps via cubic spline interpolation yields a 1.3× speedup, with base velocity dimensions scaled by 1.3×.
• Training budget: Total competition budget was approximately $13k — about $3k of personal money and $10k sponsored by Nebius for main training runs on 8× H200 GPUs.
• Self-reported limitation: The authors state this is a competition entry rather than a research paper and that many design choices were guided by intuition or quick experiments rather than systematic evaluation. They explicitly do not claim every component is necessary or optimal, and per-component ablation results are not reported.

Methodology in Plain English

The team started from Pi0.5, a model that uses flow matching to turn random noise into robot action sequences, and modified it in three areas.

Architecture. Rather than accepting a natural-language task description, the model uses 50 trainable task embeddings, one per task, each 2048-dimensional — justified because the challenge uses a fixed set of 50 tasks present in both training and evaluation, with no requirement to generalize to new language-described tasks. A small auxiliary "System 2" classifier reads the VLM output and predicts which stage of the task the robot is in, with invalid stages masked before softmax. Because individual stage predictions are noisy, a sliding window of three predictions governs transitions: two of three votes advance a stage, three unanimous votes jumping two stages trigger a one-stage catch-up, and three unanimous votes for the previous stage trigger a rollback. For attention, the model computes KV caches for all VLM layers and transforms them with learned scalar combinations plus biases, so the action expert can choose which layers to consult. A hierarchical attention mask keeps image and task tokens from attending to noisier inputs like robot state or predicted stage, preventing target leakage into the stage classifier.

Training. Actions are predicted as deltas from the current joint positions and normalized per timestep, because early actions in a chunk have small deltas while later ones vary more. The key training trick is sampling noise from a regularized covariance matrix rather than from independent Gaussian noise — the covariance is estimated from training data, giving a 690×690 matrix (30 timesteps × 23 action dimensions), then shrunk by blending 0.5 of the empirical covariance with 0.5 of the identity. The team also computes 15 flow-matching predictions per single VLM forward pass to reduce gradient variance while amortizing the expensive vision-language computation. Total loss combines action loss with a stage loss weighted 0.1 and a FAST auxiliary loss weighted 0.05. Initial multi-task training on all 50 tasks took 15 days on 8× H200 GPUs; tasks were then split into best, good, and bad groups for roughly one week of fine-tuning per group, and the final submission uses 4 task-specific checkpoints selected by task ID.

Inference. Beyond correlated inpainting, the model applies compression and correction heuristics: it disables compression when gripper states change significantly, giving the policy more time during grasps. It opens the gripper whenever a closed state was never seen at that stage in training data, treating it as a failed grasp. Correction is applied only during early denoising steps (flow time greater than 0.3), leaving late steps free to adapt to observations.

Why This Matters

This work shows that a small independent team can win a large-scale embodied AI benchmark by combining a strong pretrained VLA backbone with targeted engineering — model correlations explicitly, add a lightweight memory mechanism for ambiguous states, and patch known failure modes with heuristics. It also argues that when training data contains only successes, pure learning struggles to recover from mistakes, and that explicit correlation-aware generation is a route to smoother open-loop action execution.

Real-world applications:

• Household and service robotics: the 50 tasks span cooking, cleaning, and rearrangement, mapping directly onto domestic assistant use cases.
• Warehouse and logistics manipulation: bimanual coordination with parallel-jaw grippers on long multi-step procedures transfers to picking and packing workflows.
• Assistive robotics: long-horizon task execution with recovery from failed grasps is central to helping users with multi-step daily activities.
• Industrial inspection or lab automation: structured, fixed repertoires of long procedures are exactly where task embeddings and stage tracking apply, since those settings also have a known, closed set of tasks.

Industry relevance: The paper's honest budget accounting ($13k total, 8× H200 training) and released code and weights make the approach reproducible at a scale many labs and startups can afford. The finding that task-specific fine-tuning with checkpoint switching beat a single multi-task policy is directly relevant to deployment decisions. The gripper-recovery heuristic illustrates a broader industrial pattern: when demonstration data is clean and success-only, hand-coded safety and recovery rules can be more reliable than hoping the policy learns to recover.

Future Directions

• Isolate component contributions. The authors explicitly state they could not run rigorous ablations; systematically measuring how much each of correlated noise, learnable mixed-layer attention, stage tracking, and correction rules contributes is the most obvious next step.
• Test whether learnable mixed-layer attention matters more from scratch. The paper suggests the effect may be larger for models trained from scratch or initialized from non-robotics VLMs, since identity initialization from an already fine-tuned Pi0.5 was already a good prior.
• Scale past hand-written rules. The authors added only one task-specific rule (for the radio task) and say the approach is neither scalable nor generalizable, leaving open how to replace heuristics with learned recovery behavior.
• Generalize beyond fixed task sets. Task embeddings were justified by the closed 50-task set; whether the approach extends to natural-language task specification, unseen objects, or new tasks is not tested.

Target Audience

Robotics and embodied AI researchers working on VLA models, flow matching, and long-horizon manipulation; competition participants in BEHAVIOR and similar benchmarks; and practitioners who need practical, budget-conscious techniques for task adaptation of pretrained robot policies. Readers looking for controlled ablation studies or per-task quantitative breakdowns will need to go elsewhere — the paper reports a 26% q-score, an approximately 99% stage prediction training accuracy, and an approximately doubled success rate from gripper correction in selected tasks, but does not provide systematic component-level results.

Authors’ abstract

We present a vision-action policy that won 1st place in the 2025 BEHAVIOR Challenge - a large-scale benchmark featuring 50 diverse long-horizon household tasks in photo-realistic simulation, requiring bimanual manipulation, navigation, and context-aware decision making. Building on the Pi0.5 architecture, we introduce several innovations. Our primary contribution is correlated noise for flow matching, which improves training efficiency and enables correlation-aware inpainting for smooth action sequences. We also apply learnable mixed-layer attention and System 2 stage tracking for ambiguity resolution. Training employs multi-sample flow matching to reduce variance, while inference uses action compression and challenge-specific correction rules. Our approach achieves 26% q-score across all 50 tasks on both public and private leaderboards.