RAFT: Reasoning-Aligned Fine-Tuning Improves Robustness of VLA Post-Training

Shresth Grover1Andrey Kolobov2Dhruv Shah1

1Princeton University2Microsoft Research

Paper Code Videos BibTeX

Problem Statement

Generalist vision-language-action (VLA) policies still need post-training to reach high success on a new target task. With only a few demonstrations, plain fine-tuning improves success in the demonstration setup but makes the policy brittle: on a physical DROID robot it succeeds in 60% of in-distribution trials but only 27.9% once the object color, distractors or background change.

TL;DR: Train a reasoning adapter first, freeze it, then fine-tune on actions. VLAs post-trained this way stay robust to new objects, clutter, poses and backgrounds that break plain fine-tuning.

Method

RAFT is a two-stage post-training procedure. Stage I trains a low-rank adapter (LoRA) on task-relevant reasoning annotations, such as object bounding boxes, keypoints and gripper traces, while the rest of the policy stays frozen. Stage II freezes that adapter and fine-tunes all remaining policy parameters on the action demonstrations.

RAFT two-stage diagram. Stage I trains a LoRA on the frozen VLM backbone with reasoning targets such as bounding boxes, keypoints and trace tokens. Stage II freezes that LoRA and trains the backbone and action expert on continuous actions.

Comparison with Baselines

Pick a robot and an out-of-distribution (OOD) axis to see all three baselines and RAFT on it. All policies were post-trained on the same demonstrations and start from a similar configuration.

Robot
Task
OOD axis

Results

Across two real robots and two simulated tasks, RAFT has the highest out-of-distribution success while matching or exceeding the baselines in distribution.

Setting

Ablations

Staging and a frozen adapter matter

On simulated pick and place, co-training reasoning and actions jointly, or leaving the adapter trainable in Stage II, both fall below RAFT.

Split

Broader reasoning data helps downstream

Adding reasoning data from other scenes to Stage I lowers held-out reasoning loss and improves in-distribution success by 0.9–3.8 points at every demonstration budget (106 evaluation episodes per budget).

Held-out reasoning loss: 0.62 → 0.29

RAFT with different reasoning modalities

Blue points are RAFT trained with different Stage-I reasoning modalities, plotted by in-distribution and out-of-distribution success; the red cross is plain fine-tuning without reasoning.

Task

RAFT is compatible with RETAIN weight merging

On Franka open drawer, using RAFT instead of plain fine-tuning as RETAIN's fine-tuned endpoint raises pooled success from 30.50% to 36.16% over ID, OOD-light and OOD-texture trials (merging at α = 0.9, 106 trials per condition).

Representation Analysis

RAFT attends more to the annotated task regions: at layer 11 they receive 3.2–3.4 times their share of the image, versus 1.9–2.1 for plain fine-tuning. Its late layers also keep more object information: linear probes on layers 14–17 decode object identity with 0.51–0.53 balanced accuracy, versus 0.45.

Frame
Compare RAFT with
↔
RAFT

Drag to compare attention maps on the same frame. White boxes mark the object and the target; white points mark annotated keypoints.

RAFT Rollouts

Autonomous RAFT policies on DROID and bimanual YAM, in the demonstration setup and under distribution shifts.

BibTeX

@article{grover2026raft,
  title   = {RAFT: Reasoning-Aligned Fine-Tuning Improves Robustness of VLA Post-Training},
  author  = {Grover, Shresth and Kolobov, Andrey and Shah, Dhruv},
  journal = {arXiv preprint arXiv:2611.04368},
  year    = {2026}
}