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

TL;DR

This work tackles robustness gaps when adapting pretrained Vision-Language-Action policies for robotic manipulation, where spatial inductive bias tends to erode under RL fine-tuning. It introduces SA-VLA, a framework that fuses implicit spatial tokens with 2D visual tokens to produce geometry-aware embeddings, provides step-level dense rewards aligned with Reach-Place-Leave manipulation phases, and employs SCAN, a spatially-conditioned annealed exploration strategy, to preserve geometric priors during learning. Key contributions include spatial token fusion for zero-shot spatial generalization, dense phase-aware rewards for improved credit assignment, and SCAN for geometry-aligned exploration, all validated on LIBERO and LIBERO-PLUS with notable gains in robustness and transferability (e.g., overall zero-shot gain of +$2.25\%$, camera-view perturbations +$3.83\%$, initial-state +$0.52\%$). The results demonstrate that aligning representation learning, reward design, and exploration with task geometry yields more stable RL adaptation and better generalization in cluttered, multi-object manipulation tasks, with code available at the project page.

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.

Figures (9)

• Figure 1: Illustration of spatial inductive bias collapse during naive RL fine-tuning (left) and preserved spatial grounding with SA-VLA (right) under the same task and identical spatial perturbations. For each method, end-effector poses from three temporal phases of a single execution trajectory are rendered as semi-transparent red, green, and blue masks and overlaid to visualize how spatial behavior evolves over time.
• Figure 2: Overview of SA-VLA. Visual and spatial tokens are fused into geometry-aware embeddings, which are optimized via step-level dense rewards and spatially-conditioned exploration (SCAN) for robust RL adaptation.
• Figure 3: Spatial Token Fusion. Visual semantic tokens $\textbf{x}$ attend to spatial tokens $\textbf{z}$ augmented with positional and view embeddings via unidirectional cross-attention. A learnable channel-wise gate $\mathbf{g}$ modulates spatial contributions, followed by a residual MLP to produce fused embeddings for flow-matching policies.
• Figure 4: Phase-consistent geometric progress used for step-level dense rewards, decomposing manipulation into Reach, Place, and Leave phases.
• Figure 5: Training dynamics on the LIBERO-PLUS spatial-perturbation subset. Success rates are evaluated using SDE-based policy checkpoints saved every 10 training steps. Solid curves denote few-shot RL, and dashed curves denote zero-shot evaluation. Zero-shot evaluation uses 8 environments with a global batch size of 384, while few-shot RL uses 64 environments with a global batch size of 2048.