Direct Action-Head Injection of A Grounded 3D Point Unlocks Spatial and Task Generalization

Technical Summary: Direct Action-Head Injection of A Grounded 3D Point Unlocks Spatial and Task Generalization

1. Problem Statement

Vision-Language-Action (VLA) models leverage large-scale vision-language pretraining to enable flexible robot manipulation. However, these models exhibit significant brittleness at test time along two critical axes:

• Spatial Generalization: Performance collapses when object positions differ from those encountered during training.
• Task Generalization: Models fail to adapt when a familiar scene is paired with a different language instruction, often replaying memorized trajectories rather than responding to the new command.

Existing approaches attempt to mitigate this by providing explicit spatial grounding (e.g., object locations) via language prompts (text tokens) or visual prompts (image overlays). The authors argue that these methods fail to systematically address how the spatial signal is represented and injected. Specifically, relying on 2D representations forces the policy to internally reconstruct 3D geometry from visual and linguistic features, a burden that limits robust generalization.

2. Methodology

The paper proposes a lightweight, model-agnostic module that fundamentally changes how grounding signals are processed: lifting a 2D grounding point into 3D and injecting it directly into the action head.

The pipeline consists of three stages:

1. 3D Lifting: Given a 2D target point (obtained from any off-the-shelf visual grounding module, such as a VLM or detector), the system lifts this point into 3D space using depth information and camera parameters. It calculates the relative displacement ($d$) between the target position ($p_t$) and the current gripper position ($p_g$) in the robot base frame:
   $$d = p_t - p_g$$
2. Spatial Embedding: This 3D displacement vector is encoded into a spatial embedding ($z_{spatial}$) using a simple two-layer Multi-Layer Perceptron (MLP). The output dimension matches the hidden size of the action head.
3. Direct Injection via AdaLN: Modern VLAs often utilize a Diffusion Transformer (DiT) as the action head, which employs Adaptive Layer Normalization (AdaLN) for conditioning. Instead of adding the signal to the input tokens or text prompts, the method injects the spatial embedding directly into the AdaLN mechanism. The spatial embedding is combined with the timestep embedding ($z_{time}$) via element-wise addition to modulate the scale ($\gamma$) and shift ($\beta$) parameters of the normalization layers:
   $$\gamma, \beta = \text{Linear}(z_{time} + z_{spatial})$$
   This allows the policy to incorporate task-relevant 3D geometric information directly during action prediction without modifying the VLA backbone or pretraining pipeline.

3. Key Contributions

• Identification of the Bottleneck: The authors demonstrate that the representation (2D vs. 3D) and injection mechanism (prompting vs. direct head injection) are the primary factors determining generalization, rather than the source of the grounding signal itself.
• Lightweight Architecture: The proposed solution requires no changes to the VLA backbone or pretraining data. It adds only a two-layer MLP and utilizes existing AdaLN conditioning mechanisms.
• Backbone Agnosticism: The method is validated on multiple distinct VLA backbones (GR00T-N1.6 and $\pi0.5$), proving its applicability across different model architectures.

4. Experimental Results

The method was evaluated on the LIBERO and LIBERO-PRO benchmarks, which include Task Perturbation (changing instructions) and Position Perturbation (changing object locations).

• Simulation Performance:

  • On GR00T-N1.6, the method improved the average success rate under Task Perturbation from 31.2% to 77.5% and under Position Perturbation from 28.1% to 60.2%.
  • On $\pi0.5$, similar gains were observed, increasing success rates from 37.3% to 75.9% (Task) and 47.0% to 72.2% (Position).
  • Controlled Comparisons: When provided with the identical oracle 2D target point, 2D variants (via text, visual, or AdaLN prompts) yielded only modest improvements. In contrast, the 3D representation injected via AdaLN produced dramatic gains, confirming that the 3D signal and direct injection are the critical factors.
  • Data Efficiency: A method trained on only 20 tasks using this approach outperformed baselines trained on 69 tasks, suggesting that architectural design outweighs data scaling for generalization.
  • Robustness: The method remained effective even when trained and tested with noisy depth inputs.

• Real-World Validation:

  • Experiments on a Franka Emika Panda robot with consumer-grade RGB-D sensors confirmed the simulation findings.
  • While baselines and 2D-prompt variants collapsed to near-zero success under perturbations in the real world, the proposed method retained high performance, demonstrating practical applicability despite noisy grounding inputs.

5. Significance and Claims

The paper claims that the central finding is a "simple but overlooked" insight: given adequate 2D grounding, lifting the signal into 3D and injecting it directly into the action head is the key to unlocking both spatial and task generalization in VLAs.

The authors emphasize that this approach achieves substantial performance gains without the need for complex 3D encoders, additional foundation models, or architectural overhauls. The work suggests that the field has been over-complicating the integration of spatial signals, and that a direct, geometry-aligned injection into the action head is the most effective path forward for robust robot manipulation. The authors note limitations regarding object-agnostic motions and the assumption of available 2D grounding, but position their contribution as a foundational step for improving VLA robustness.