SA-VLA: Spatially-Aware Flow-Matching for Vision-Language-Action Reinforcement Learning

TL;DR

SA-VLA introduces a spatially-aware reinforcement learning framework that maintains spatial grounding during policy adaptation, leading to improved robustness and transferability in robotic manipulation tasks.

Contribution

It proposes a novel spatially-aware RL adaptation method that preserves spatial inductive bias by aligning representation, reward, and exploration with task geometry.

Findings

Enhanced robustness in manipulation benchmarks

Improved zero-shot spatial generalization

Stable RL fine-tuning across tasks

Abstract

Vision-Language-Action (VLA) models exhibit strong generalization in robotic manipulation, yet reinforcement learning (RL) fine-tuning often degrades robustness under spatial distribution shifts. For flow-matching VLA policies, this degradation is closely associated with the erosion of spatial inductive bias during RL adaptation, as sparse rewards and spatially agnostic exploration increasingly favor short-horizon visual cues. To address this issue, we propose \textbf{SA-VLA}, a spatially-aware RL adaptation framework that preserves spatial grounding during policy optimization by aligning representation learning, reward design, and exploration with task geometry. SA-VLA fuses implicit spatial representations with visual tokens, provides dense rewards that reflect geometric progress, and employs \textbf{SCAN}, a spatially-conditioned annealed exploration strategy tailored to flow-matching dynamics. Across challenging multi-object and cluttered manipulation benchmarks, SA-VLA enables stable RL fine-tuning and improves zero-shot spatial generalization, yielding more robust and transferable behaviors. Code and project page are available at https://xupan.top/Projects/savla.