GST-VLA: Structured Gaussian Spatial Tokens for 3D Depth-Aware Vision-Language-Action Models

TL;DR

Blations isolate the contribution of each GST component, each DA-CoT thought, and each training stage, confirming independent and synergistic gains concentrated on precision demanding tasks.

Abstract

VLA models encode visual observations as 2D patch tokens with no intrinsic geometric structure. We introduce GST-VLA with two contributions. First, the Gaussian Spatial Tokenizer (GST) converts frozen dense depth and frozen semantic patch features into $N_g{=}128$ anisotropic 3D Gaussian primitives, each parameterized by a metric residual mean $μ\in \mathbb{R}^3$, log-scale covariance $\log σ\in \mathbb{R}^3$, and learned opacity $α\in (0,1)$. The covariance eigenstructure encodes local surface orientation, and opacity provides per-primitive geometric confidence, both inaccessible from scalar depth. Spatial attention pooling with learned queries concentrates the fixed token budget on geometrically salient regions rather than distributing uniformly. Second, 3D Depth-Aware Chain-of-Thought (DA-CoT) reasoning supervises four structured intermediate spatial thoughts, covering 3D object grounding, grasp affordance contact geometry, pairwise metric distances, and coarse SE(3) waypoints, as explicit generation targets in the training loss. A cross-attention sublayer at every VLM transformer block provides direct access to the raw 256-primitive Gaussian field during DA-CoT generation. A 300M-parameter flow-matching action expert with mixture-of-experts feedforward sublayers decodes 7-DoF delta action chunks via conditional ODE integration, conditioned on both VLM hidden states and DA-CoT outputs through dual cross-attention. Trained with composite $\mathcal{L}_\mathrm{flow} + \mathcal{L}_\mathrm{CoT} + \mathcal{L}_\mathrm{depth}$ across three progressive stages, GST-VLA achieves 96.4% on LIBERO (+2.0%), and 80.2% on SimplerEnv (+5.4%). Ablations isolate the contribution of each GST component, each DA-CoT thought, and each training stage, confirming independent and synergistic gains concentrated on precision demanding tasks.

Figures (2)

• Figure 1: The proposed GST-VLA pipeline integrates five sequential stages to ground robot actions in structured 3D spatial reasoning. A frozen semantic encoder and a frozen depth expert process the RGB observation in parallel, extracting dense patch features and affine-invariant metric depth respectively. The novel trainable Gaussian Spatial Tokenizer (GST) fuses these streams by back-projecting depth into 3D, estimating per-patch Gaussian parameters $(\mu, \sigma, \alpha)$ from visual features, applying 3D Fourier positional encoding, and aggregating to $N_g$ structured spatial tokens via spatial attention pooling. These tokens are projected into the VLM reasoning core through a cross-attention projector, which generates supervised Depth-Aware Chain-of-Thought (DA-CoT) intermediate reasoning over 3D object grounding, grasp affordance, metric spatial relations, and SE(3) motion plan waypoints before producing action conditioning tokens.
• Figure 2: GST-VLA framework. Two frozen encoders produce semantic features and metric depth. The LoRA adapted GST lifts these into $N_g{=}128$ anisotropic 3D Gaussian tokens via four operations. A cross-attention projector injects spatial tokens into the VLM, where DA-CoT sublayers attend to the raw 256-primitive Gaussian field. The action expert receives dual conditioning: VLM hidden states (semantic and visual context) and DA-CoT action tokens (3D geometric reasoning).