Spatial understanding is essential for robots to perceive their surrounding environments and execute actions. Existing approaches typically encode 3D awareness into vision-language-action (VLA) models by directly aligning their representations with pre-trained 3D foundation models. However, we observe that such direct geometric feature alignment essentially exploits 3D foundation priors as a form of low-rank regularization in the feature space, resulting in severe rank collapse of the pre-trained VLA representations. In contrast, we shift the focus from enforcing point-wise equality to aligning pair-wise relationships. We propose a representation-alignment method that enables the injection of 3D spatial awareness without compromising the model's original expressivity. Extensive experiments on public benchmarks, our new spatial reasoning benchmark, and real-world manipulation scenarios demonstrate that our approach effectively incorporates 3D spatial representations into VLA reasoning, achieving robust performance and enhanced spatial awareness across diverse manipulation tasks.
Figures & tables
Figure 1: Spectral and rank analysis of different alignment methods. (a) Traditional 3D knowledge distillation aligns the point-wise magnitudes (i.e, vector norm) or angles (i.e, cos-similarity) across feature spaces, which often leads to rank collapse (b). In contrast, we align pair-wise relationships to achieve a rotation-invariant alignment that effectively prevents feature collapse. See Sec. 3.1 for details. VGGT [ 50 ] is used as the 3D teacher model.
Figure 2: Overview of our framework. We extract self- and cross-attention (attn.) maps from both the VLA and the 3D foundation models to serve as relational affinity matrices described in Sec. 3.1 . The 3D foundation model is kept frozen, and the VLA model is trained via LoRA. Other parts are fully trainable. More views can be naturally incorporated.
Table 3
Figure 3: Robustness to camera view changes. We evaluate results on camera view change by level on LIBERO-Spatial. Our method outperforms the Spatial Forcing across most levels, showing particular robustness at higher difficulty tiers, mitigating the performance degradation.
Figure 4: Proposed LIBERO-Probing benchmark. Our benchmark generates the semantic segmentation (Seg.), depth estimation (Depth), and correspondence matching (Corr.) results in the LIBERO training environments to analyze the VLA capabilities.
Figure 5: Cross-view attention across VGGT layers. The red markers indicate the query point in the wrist view and its correspondence in the agent view. Attention maps from layers (6,12,18,24) are overlaid on the agent view.
Figure 6: Component analysis. Left side: success rates (%) on LIBERO-Plus Spatial; higher is better. Bold indicates the best result. Right side: success rates across VGGT layers. Layer 0 means the output from the VGGT encoder.
Figure 7: Our robot environment and task examples. We use a 7-DoF robotic arm with a two-finger gripper, a side-view camera, a wrist camera, and an agent-view camera (used only for visualization). We evaluate on six tasks in a complex kitchen environment.
Figure 8: Performance comparison on real-world tasks. We report the success rates and the average rates on the six real robot tasks. All the models took 10 trials on each task. (a) – (f) are consistent with the labels described in Fig. 7 and Sec. 5.1 .
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameters
Simulation (LIBERO) Config
Real-Robot Config
Base VLA
OpenVLA (7B)
π0 -DROID (3B)
Optimizer
AdamW ( β1=0.9,β2=0.95 )
AdamW ( β1=0.9,β2=0.95 )
Learning rate
5e-4 / 5e-5 (after 20k steps)
2.5e-5 (peak)
Lr Scheduler
MultiStepLR
Cosine Decay
Warmup Steps
1,000
1,000
Batch Size
16
16
Appendix
Table S1: Combined hyperparameters for simulation and real-robot training configurations.
Figure S1: Dimensionality analysis of visual features.
Method
Alignment
Eff. Rank
Perf. (%)
Baseline
Point-wise
248.0
72.9
+ preventing rank collapse
Point-wise
258.8 ↑
78.6 ↑
Ours
Relational
306.3
85.2
+ inducing rank collapse
Relational
278.5 ↓
80.8 ↓
Appendix
Table S2: Rank intervention on LIBERO-Plus Spatial. We report effective rank and policy success rate after 30k training steps.
Model
Cam
L1
L2
L3
L4
L5
Ro
Lang
Light
Back
Noise
Lay
Avg
Loss
MSE
60.1
73.1
66.9
53.5
55.0
22.7
38.0
99.0
100.0
99.6
75.8
84.4
78.5
KL
73.9
83.6
76.7
69.3
70.0
59.1
52.9
88.7
97.9
98.1
88.6
95.6
84.4
L1
68.6
95.5
79.7
55.3
42.5
36.4
52.0
89.5
100.0
99.2
81.5
94.5
82.7
MMD
60.9
85.1
74.4
51.8
30.0
9.1
47.7
97.2
99.7
98.4
78.1
91.4
81.0
Wasserstein
83.2
97.0
88.0
74.6
67.5
86.4
48.1
83.6
99.7
94.2
90.3
96.1
85.2
Appendix
Table S3: Component analysis. We systematically analyze different component choices. KEY– EMA: Exponential Moving Average, coe.: coefficient.
Figure S2: Success rate across VLA integration layers. Layer 0 denotes the output of the VLA vision encoders.
Method
Train time (s/step)
FLOPs/step
Peak mem.
+ affinity distillation, KL
0.6312
32.91 T
27.05 GB
+ Sinkhorn-OT, T=10
0.6501 (+3.0%)
32.92 T (+0.03%)
27.05 GB
+ Sinkhorn-OT, T=20
0.6580 (+4.2%)
32.92 T (+0.03%)
27.05 GB
+ Sinkhorn-OT, T=50 (ours)
0.6814 (+8.0%)
32.94 T (+0.09%)
27.05 GB
Appendix
Table S4: Computational cost of attention-map distillation. Measurements use a batch size of 1 on a single H100 GPU. T denotes the number of Sinkhorn iterations. Percentages are relative to the KL baseline.
EMA decay rate
0.0
0.9
0.999
0.9999
Avg. successful rate (%)
85.2
85.1
86.6
82.1
Appendix
Table S5: Performance according to the decay parameter β on LIBERO-Plus Spatial.
Representation alignment studies how learned feature spaces can be compared or mapped to transfer meaning, capabilities, and information across models and tasks.
Nanjing University, Nanjing, China · Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China