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

Shresth Grover; Andrey Kolobov; Dhruv Shah

cs.RONov 3, 2026

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

Authors: Shresth Grover, Andrey Kolobov, Dhruv Shah

Organizations: 1Princeton University · 2Microsoft Research

Abstract

Adapting a generalist vision-language-action (VLA) policy to a target task with limited demonstrations can improve nominal task success while reducing robustness to variations in appearance, objects, and instructions. We introduce reasoning-aligned fine-tuning (RAFT), a two-stage post-training procedure: RAFT first trains low-rank adapters on task-relevant reasoning annotations, then freezes them while fine-tuning the policy on action demonstrations. On a physical DROID robot, RAFT achieves $75.0\%$ out-of-distribution success, compared with $27.9\%$ for plain fine-tuning; on a multi-stage bimanual YAM task, it reaches $70.0\%$ mean success, compared with $36.0\%$ for the strongest baseline. RAFT also improves success under visual, object, and instruction shifts in simulated pick-and-place and open-drawer tasks. Compared with plain fine-tuning, RAFT policies place more attention on annotated task regions and yield higher linear-probe accuracy for object identity in later layers. A broader reasoning-data mixture further improves downstream success across demonstration budgets.

Figures & tables

Figure 1: Reasoning-aligned fine-tuning for robust VLA adaptation. (a) Plain FT, fine-tuned only on actions, succeeds in the training layout but fails under small scene changes. (b) RAFT first learns task-relevant reasoning through LoRA adapters, then freezes them while fine-tuning the policy on action demonstrations. (c) With the same demonstrations, RAFT succeeds under these shifts. On physical robots, RAFT raises success over Plain FT from $34.7\%$ to $70.0\%$ on YAM and from $27.9\%$ to $75.0\%$ on out-of-distribution DROID trials.

Figure 2: Overview of RAFT . Stage I updates only the reasoning adapters, using reasoning annotations. Stage II freezes the adapters and updates the language model, action expert, and projection layers on action demonstrations, so action gradients do not reach the reasoning adapters.

Figure 3: Example Stage-I reasoning annotation. In this YAM example, the left arm holds the lid open while the right arm reaches for the toy. The target combines language, object boxes, keypoints, motion traces, and per-arm state information. Blue and green mark the box and toy, yellow marks keypoints, and red marks the active arm’s trace.

Figure 4: Evaluation scenes. Representative in-distribution (ID) and out-of-distribution (OOD) observations, left to right: simulated pick and place, simulated open drawer, physical DROID pick and place, and the physical YAM two-stage task. The OOD panels show texture changes in simulation and background changes on the physical robots. The full set of evaluated shifts is shown in Fig. 20 .

Figure 5: Simulation results under distribution shift. Success rates for pick and place (left) and open drawer (right), with 106 trials per split. RAFT uses full reasoning for pick and place and bounding-box supervision for open drawer (Sec. 8.1 ). RAFT has the highest mean success on both tasks:

83\%

versus

80\%

for Plain FT on pick and place, and

39\%

versus

29\%

on open drawer.

Figure 6: Real-robot results under distribution shift. Left: DROID pick and place. Right: YAM box opening and place inside box. RAFT has the highest success on both platforms:

75\%

versus

28\%

for Plain FT on out-of-distribution DROID trials, and

70\%

versus

35\%

on YAM.

Figure 7: Successful physical-robot rollouts. Ordered frames show DROID picking up the brown toy and placing it in the bowl (top) and YAM opening the box and placing the toy inside (bottom). Numbers indicate temporal order; circles and arrows mark the manipulated object and motion.

Figure 8: Attention maps along one trajectory. Columns show successive observations; rows compare Base, Stage I, RAFT , and Plain FT. Overlays identify annotated task regions, and color scales are shared within each column. RAFT places more attention on the annotated regions than Plain FT; Fig. 10 quantifies this across layers.

Figure 9: Layer-wise linear probes. Balanced accuracy for object identity (left; chance

1/7

) and episode-time quartile (right; chance

1/4

). Policies use the same 414 frames; episodes are held out from probe fitting, not policy training. Shading marks layers 14–17. In these layers, RAFT reaches 0.51–0.53 object-identity accuracy versus 0.45 for Plain FT, and 0.62–0.65 quartile accuracy versus 0.55.

Figure 10: Layer-wise attention enrichment. Enrichment for object boxes, keypoints, traces, and background; values above one indicate attention exceeding a region’s image-area share. Curves compare Base, Plain FT, Stage I, and RAFT ; shading marks layers 14–17. At layer 11, RAFT gives task regions 3.2–3.4 times their area share, versus 1.9–2.1 for Plain FT, while background attention stays near uniform.

Figure 11: Effect of reasoning data and demonstration budget. Left: held-out reasoning loss across corpus sizes and constructions. Middle: in-distribution pick-and-place success after 5,000 Stage-II steps from in-domain and broader-mixture Stage-I initializations. Right: RAFT versus Plain FT across demonstration budgets, with 106 evaluation episodes per budget. The broader mixture lowers held-out loss (0.29 versus 0.62) and improves this success by 0.9–3.8 points at every budget.

Figure 12: Training-scheme ablation on simulated pick and place. Mean success over trials. Staged training with a frozen adapter achieves the highest success (

83.2\%

), above Plain FT (

79.6\%

), a trainable adapter (

74.0\%

), and co-training (

63.2\%

	Stage I: reasoning alignment	Stage II: action fine-tuning
Trainable components	LoRA adapters only	Language model, action expert, and projection layers; SigLIP vision tower and Stage-I LoRA frozen
LoRA placement	All 18 language-model layers: q/k/v/o attention projections and gate/up/down MLP projections (126 modules); all 27 vision layers: q/k/v projections (81 modules)	The Stage-I LoRA remains attached and frozen; no adapter parameters are optimized in Stage II
LoRA configuration	Rank $r=64$ , scaling $\alpha=128$ ( $\alpha/r=2.0$ ), dropout $0.05$	—
Trainable parameters	$90.4$ M ( $78.4$ M language-model and $11.9$ M vision-adapter parameters)	$3.09$ B of $3.50$ B parameters
Optimizer	AdamW, $\beta=(0.9,0.999)$ , weight decay $10^{-6}$ , gradient clipping $1.0$	AdamW, $\beta=(0.9,0.95)$ , gradient clipping $1.0$
Learning rate	Peak $2\times 10^{-4}$ ; 100 warmup steps; cosine schedule (2,010-step horizon), 1,000 steps	Peak $5\times 10^{-5}$ ; 1,000 warmup steps; thereafter effectively constant at $5\times 10^{-5}$

Table 1: Training configuration for the two stages of RAFT .

Variant	Target response
VQA	The robot should grasp the glossy beige ceramic mug on the left side of the counter and place it into the white bowl on its right.
Spatial	The beige mug is left of the white bowl. The gripper begins above and to the right of the mug, moves to the mug, lifts it, then moves rightward to the bowl.
Bounding box	glossy beige ceramic mug: [310, 600, 380, 740]; white bowl: [410, 570, 520, 720]
Keypoints	grasp_point: [340, 670]; target_point: [460, 650]; lift_point: [340, 500]; pre_grasp_point: [340, 550]; pre_place_point: [460, 500]
Trajectory trace	[[600, 400], [470, 470], [340, 550], [340, 670], [400, 580], [460, 650]]
Full reasoning	The VQA plan together with the two object boxes, five semantic keypoints, and six-point trajectory trace shown above.

Table 2: Modality-specific Stage-I targets for a mug-into-bowl pick-and-place example. Coordinates use the normalized

[0,1000]

image frame.

Figure 13: DROID rollout comparison under an OOD background. Rows show RAFT , Plain FT, RETAIN, and DeLock at four ordered points in matched real-robot episodes. RAFT completes the manipulation, while the comparison policies fail before task completion.

Figure 14: DROID rollout comparison under OOD clutter. Rows show RAFT , Plain FT, RETAIN, and DeLock at four ordered points in matched real-robot episodes. RAFT completes the manipulation, while the comparison policies fail before task completion.

Figure 15: DROID rollout comparison with an OOD object. Rows show RAFT , Plain FT, RETAIN, and DeLock at four ordered points in matched real-robot episodes. RAFT completes the manipulation, while the comparison policies fail before task completion.

Figure 16: YAM rollout comparison under OOD clutter. Rows show RAFT , Plain FT, RETAIN, and DeLock on matched task instances; columns show ordered rollout frames. RAFT completes the two-stage task, while the comparison policies do not.

Figure 17: YAM rollout comparison under an OOD color shift. Rows show RAFT , Plain FT, RETAIN, and DeLock on matched task instances; columns show ordered rollout frames. RAFT completes the two-stage task, while the comparison policies do not.

Figure 18: YAM rollout comparison under an OOD initial pose. Rows show RAFT , Plain FT, RETAIN, and DeLock on matched task instances; columns show ordered rollout frames. RAFT completes the two-stage task, while the comparison policies do not.

Figure 19: Effect of reasoning modality on success and representation dimensionality. Left: ID and OOD-visual success on the open drawer task for six reasoning modalities. Center: ID and OOD-object success on pick and place for five modalities. Red crosses denote no-reasoning controls, and dashed lines mark Plain FT. Each policy is evaluated on 106 trials per axis. Right: Layer-wise mean Spearman correlation between participation ratio and success, pooled over three task-matched evaluation axes and six reasoning variants, for four participation-ratio measures. Filled markers denote significance after FDR correction; shading marks layers 14–17.

Adapter setting	ID	OOD-light	OOD-texture	OOD-object	OOD-prompt	Mean
LoRA on base (no Stage II)	72.64	60.38	9.43	33.02	66.04	48.30
LoRA trainable in Stage II	90.57	86.79	65.09	34.91	92.45	73.96
LoRA frozen in Stage II ( RAFT )	95.28	93.40	92.45	43.40	91.51	83.21

Table 3: Stage-II reasoning-adapter ablation on pick and place. Success rate (%) when the Stage-I reasoning LoRA is applied without Stage II, updated during Stage II, or frozen during Stage II ( RAFT ). Each split contains 106 episodes.

Method	ID	OOD-light	OOD-texture	OOD-object	OOD-prompt	Mean
Plain FT	90.57	87.74	92.45	35.85	91.51	79.62
Plain FT (frozen vision)	91.51	86.79	87.74	40.57	91.51	79.62
Co-train	79.25	70.75	72.64	17.92	75.47	63.21
RAFT	95.28	93.40	92.45	43.40	91.51	83.21

Table 4: Co-training and frozen-vision controls on simulated pick and place. Success rate (%) on five matched evaluation splits (

n=106

per cell).

Policy	ID	OOD-light	OOD-texture	Overall
RETAIN on Plain FT ( $\alpha=0.9$ )	38.68	34.91	17.92	30.50
RETAIN on RAFT (bbox, $\alpha=0.9$ )	41.51	37.74	29.25	36.16
RAFT (bbox) alone	44.34	41.51	28.30	38.05

Table 5: RETAIN applied to Plain FT and RAFT on the Franka open drawer task. Success rate (%) for each axis. RETAIN applied to RAFT outperforms RETAIN applied to Plain FT.

Figure 20: Evaluation setup across the four robot tasks. Each row begins with an in-distribution observation, followed by its OOD variants. Pick and place varies lighting, target object, instruction, and scene texture; the open drawer task varies lighting, instruction, and scene texture; DROID varies background, clutter, and object appearance; and YAM varies box pose, object color, clutter, and background. Instruction shifts reuse the ID image because only the task prompt changes.

Figure 21: Late-layer probes for all nine model arms. Leave-one-episode-out balanced accuracy for object identity and episode-time quartile across backbone layers. The shaded band marks layers 14–17. This expands the four-arm comparison in the main text to the full set of reasoning-modality variants.

VLA Model Adaptation

See all 48 papers

What to Ignore, What to React: Visually Robust RL Fine-Tuning of VLA Models
May 13, 2026Yuanfang Peng, Jingjing Fu, Chuheng Zhang +6VLM Robustness Domain Generalization
1Hong Kong University of Science and Technology · 2Microsoft Research Asia · 3Zhejiang University
Retrieve to Act: Motion Primitive Graph Retrieval for Robust VLA Execution
Oct 27, 2026Pham Tri Quang, Bao-Ngoc Dao, Minh Duc Nguyen +15Robustness of VLA Models Test-Time Adaptation of VLA Models
Center for AI Research, VinUniversity · VinRobotics · Hanoi University of Science and Technology +6
RL $^2$ -VLA: Adaptive RL Latent Compositional Steering with Test-Time Scaling for Vision-Language-Action Models
Jul 29, 2026Derek Ming Siang Tan, Shailesh Shailesh, Srikrishna Iyer +4Vision-Language-Action Models Test-Time Adaptation of VLA Models
1National University of Singapore · 3Singapore Technologies Engineering · University of Toronto

May 13, 2026J/K move · Enter open · S save

What to Ignore, What to React: Visually Robust RL Fine-Tuning of VLA Models

Yuanfang Peng, Jingjing Fu, Chuheng Zhang, Li Zhao+5

1Hong Kong University of Science and Technology · 2Microsoft Research Asia · 3Zhejiang University

VLM Robustness Domain Generalization Robotic RL Vision-Language-Action Models RL Fine-Tuning

Open paper AskPDF arXiv

Reinforcement learning (RL) fine-tuning has shown promise for Vision-Language-Action (VLA) models in robotic manipulation, but deployment-time visual shifts pose practical challenges. A key difficulty is that standard task rewards supervise task success, but offer limited guidance on whether a visual change is task-irrelevant or changes the behavior required for manipulation. We propose PAIR-VLA (Paired Action Invariance & Sensitivity for Visually Robust VLA), an RL fine-tuning framework to address this difficulty by adding two auxiliary objectives over paired visual variants during PPO optimization: an invariance term that reduces the discrepancy between action distributions for a task-preserving pair (e.g., different distractors), and a sensitivity objective that encourages separable action distributions for a task-altering pair (e.g., target object in a different pose). Together, these objectives turn visual variants from mere observation diversity into behavior-level guidance on policy responses during RL fine-tuning. We evaluate on ManiSkill3 across two representative VLA architectures, OpenVLA and

π_{0.5}

, under diverse out-of-distribution visual shifts including unseen distractors, texture changes, target object pose variation, viewpoint shifts, and lighting changes. Our method consistently improves over standard PPO, achieving average improvements of 16.62% on

π_{0.5}

and 9.10% on OpenVLA. Notably, ablations further show generalization across visual shifts: invariance guidance learned from distractor and texture variants transfers to target-pose and lighting shifts, while adding sensitivity guidance on target-pose variants further improves robustness to nuisance shifts, highlighting the broader transferability of behavior-level RL guidance.

Abstract

VLA Model Adaptation

Explore similar work

What to Ignore, What to React: Visually Robust RL Fine-Tuning of VLA Models

Retrieve to Act: Motion Primitive Graph Retrieval for Robust VLA Execution

RL2^22-VLA: Adaptive RL Latent Compositional Steering with Test-Time Scaling for Vision-Language-Action Models

RL $^2$ -VLA: Adaptive RL Latent Compositional Steering with Test-Time Scaling for Vision-Language-Action Models