Research project · 2026

SFM: Spatially-Aware Flow Matching for Embodied Reinforcement Learning

Xu Pan^1,2Zhenglin Wan³Xingrui Yu^2,*Xianwei Zheng⁴Youkai Ke⁴Ming Sun⁵Rui Wang⁵Ziwei Wang¹Ivor Tsang²

School of Electrical and Electronic Engineering, Nanyang Technological University
Centre for Frontier AI Research (CFAR), Institute of High Performance Computing (IHPC), Agency for Science, Technology and Research (A*STAR)
Department of Computer Science, National University of Singapore
State Key Laboratory of Information Engineering in Surveying, Mapping and Remote Sensing (LIESMARS), Wuhan University
Institute of Information Engineering, Chinese Academy of Sciences

^* Corresponding author

Under Review · 2026

Abstract

Reinforcement learning can improve a flow-matching Vision-Language-Action policy on its training distribution while making its geometry less reliable under new viewpoints and scene layouts. SFM treats this as a coupled failure of representation, credit assignment, and exploration. It aligns all three in a shared geometric latent space through implicit spatial token fusion, Spatially Grounded Reward (SGR), and Spatially Conditioned Annealed Exploration (SCAE). On LIBERO^[3] and LIBERO-Plus^[4], this unified design improves robustness to spatial distribution shifts without sacrificing in-distribution control.

Comparison of naive reinforcement learning and SFM under spatial distribution shifts — **Fig. 1.** Conventional RL fine-tuning can drift toward viewpoint-specific visual correlations. SFM preserves spatially consistent behavior by grounding optimization in geometry.

Method

SFM training framework — **Fig. 2.** SFM shares a geometry-aware latent across three optimization paths: spatial token fusion for observation encoding, phase-aware SGR for credit assignment, and SCAE for structured action exploration.

Spatial token fusion

Visual Geometry Grounded Transformer features^[1] inject implicit camera and scene geometry into the VLA observation stream through cross-attention.

Grounded reward

SGR decomposes manipulation into approach, relation–actuation, and stabilization phases, providing dense geometric progress signals beyond sparse task success.

Annealed exploration

SCAE conditions stochastic action perturbations on spatial state and training progress, then anneals them as the policy becomes more certain.

Principal component visualization of spatial features — **Fig. 3.** Spatial features remain responsive to geometric changes after fusion; SFM produces a more structured response than the RL baseline.

Baseline spatial response — **Fig. 3.** Spatial features remain responsive to geometric changes after fusion; SFM produces a more structured response than the RL baseline.

Cross-attention spatial token fusion — **Fig. 4.** Geometry tokens are projected and fused with visual tokens through cross-attention before the policy predicts an action flow.

Spatially conditioned exploration residuals — **Fig. 5.** SCAE learns spatially structured residuals instead of applying uniform action noise.

Three phases of spatially grounded reward — **Fig. 6.** SGR changes its geometric objective across approach, relation–actuation, and stabilization, matching the causal structure of manipulation.

Evaluation

SFM initializes from the frozen 7B SigLIP–Gemma backbone of π_0.5^[2] and is optimized with PPO and GAE in RLinf-VLA^[7]. Training uses 64 parallel environments, batches of 1,024 transitions, and a 240-step episode horizon. The spatial robustness suite covers few-shot camera-view (CV) and robot-initial-state (RIS) shifts, plus zero-shot background texture, lighting, layout, object-location, and sensor-noise shifts.

64parallel environments

1,024transitions per batch

≈38.5 htraining on four H800 GPUs

Method	CV	RIS	Total
π_0.5	78.47	83.77	81.00
SFM w/ spatial fusion	82.30	84.29	83.25

Table 1. Success rate (%) under few-shot spatial perturbations. Implicit spatial representation improves both camera-view and robot-state robustness.

SFM training curves — **Fig. 7.** SFM improves optimization stability and converges to stronger task performance than reward-only and conventional RL baselines.

Evaluation of spatially grounded and sparse rewards — **Fig. 8.** Spatially grounded reward provides more reliable few-shot and zero-shot transfer than sparse task-level feedback.

Results

The complete system reaches 90.14% aggregate robustness, outperforming both ReinFlow^[5] and Flow-GRPO^[6]. Its clearest gain appears under camera-view perturbation, where success rises to 87.21%. Ablations show that the three components are complementary: removing SGR is especially damaging, while spatial fusion and SCAE each add robustness beyond dense reward alone.

Method	Total	BT	CV	LI	LC	OL	RIS	SN
π_0.5	88.96	100.00	82.16	88.63	98.54	100.00	83.74	87.93
ReinFlow	87.92±0.81	97.53±0.53	84.32±0.94	84.20±1.45	97.32±1.69	100.00±0.00	79.15±1.58	92.72±2.32
Flow-GRPO	89.29±0.46	99.07±0.93	85.23±1.13	87.20±0.47	97.81±0.73	100.00±0.00	81.12±0.28	92.53±1.72
SFM	90.14±0.35	100.00±0.00	87.21±0.31	86.10±1.52	98.54±1.26	100.00±0.00	83.09±0.57	93.30±1.85

Table 2. LIBERO-Plus success rate (%). BT: background texture; CV: camera view; LI: lighting; LC: layout; OL: object location; RIS: robot initial state; SN: sensor noise.

Variant	SCAE	SGR	Fusion	Total
π_0.5	—	—	—	81.00
w/o fusion	—	✓	—	81.75
w/o SGR	—	—	✓	77.50
w/o SCAE	—	✓	✓	83.00
SFM	✓	✓	✓	83.75

Table 3. Component ablation under few-shot CV and RIS perturbations. The full geometry-aligned system performs best.

Comparison of learned exploration noise and stochastic differential equation noise — **Fig. 9.** Learned, spatially conditioned exploration complements SGR more effectively than generic SDE noise.

Qualitative SFM robot manipulation rollout — **Fig. 10.** A qualitative rollout shows stable approach, interaction, and completion across the manipulation sequence.

Citation

@misc{pan2026sfm,
  title={SFM: Spatially-Aware Flow Matching for Embodied Reinforcement Learning},
  author={Xu Pan and Zhenglin Wan and Xingrui Yu and Xianwei Zheng and Youkai Ke and Ming Sun and Rui Wang and Ziwei Wang and Ivor Tsang},
  year={2026},
  note={Under review}
}

References

J. Wang, M. Chen, N. Karaev, et al. “VGGT: Visual Geometry Grounded Transformer.” CVPR, 2025.arXiv
Physical Intelligence, K. Black, N. Brown, et al. “π_0.5: a Vision-Language-Action Model with Open-World Generalization.” 2025.arXiv
B. Liu, Y. Zhu, C. Gao, et al. “LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning.” NeurIPS, 2023.arXiv
S. Fei, S. Wang, J. Shi, et al. “LIBERO-Plus: In-depth Robustness Analysis of Vision-Language-Action Models.” 2025.arXiv
T. Zhang, C. Yu, S. Su, and Y. Wang. “ReinFlow: Fine-tuning Flow Matching Policy with Online Reinforcement Learning.” 2025.arXiv
J. Liu, G. Liu, J. Liang, et al. “Flow-GRPO: Training Flow Matching Models via Online RL.” 2025.arXiv
H. Zang, M. Wei, S. Xu, et al. “RLinf-VLA: A Unified and Efficient Framework for VLA + RL Training.” 2025.arXiv