💻 computer science

GeoAlign: Beyond Semantics with State-Guided Spatial Alignment in VLA Models

GeoAlign introduces a state-guided spatial alignment architecture for Vision-Language-Action models that enhances manipulation performance by post-training an RGB branch with robot-domain RGB-D supervision to generate geometry-aware features queried by proprioceptive states, achieving state-of-the-art results on both simulated and real-world benchmarks.

Original authors: Yizhi Chen, Zhanxiang Cao, Xinyi Peng, Yixiao Zheng, Xiaxi Si, Yiheng Li, Liyun Yan, Keqi Zhu, Xueyun Chen, Shengcheng Fu, Tianyue Zhan, Yufei Jia, Jinming Yao, Yan Xie, Kun Wang, Cewu Lu, Yue Gao

Published 2026-06-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yizhi Chen, Zhanxiang Cao, Xinyi Peng, Yixiao Zheng, Xiaxi Si, Yiheng Li, Liyun Yan, Keqi Zhu, Xueyun Chen, Shengcheng Fu, Tianyue Zhan, Yufei Jia, Jinming Yao, Yan Xie, Kun Wang, Cewu Lu, Yue Gao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to perform a delicate task, like picking up a clear glass bottle or sliding a piece of tape into a tight slot. You might think the robot just needs to "see" the object and "know" what to do. But as this paper explains, there's a hidden problem: Robots often get the "what" right, but fail at the "how."

Current robot brains (called VLA models) are great at understanding language and identifying objects ("That is a bottle"). However, they often struggle with the fine-grained geometry—the exact shape, the tiny gaps, and the precise alignment needed to actually move the object without crashing or dropping it. It's like knowing you need to thread a needle, but having eyes that only see the general shape of the needle, not the tiny hole in the eye.

Here is how GeoAlign fixes this, using a few simple analogies:

1. The Problem: The "Blurry Depth" Camera

Usually, to understand 3D space, robots use special depth cameras. But these cameras are terrible at seeing transparent things (like glass) or thin rings (like tape rolls). The depth data comes back broken, missing, or full of holes.

The Paper's Solution: Instead of relying on a broken depth camera, GeoAlign teaches the robot to imagine the 3D shape just by looking at a standard color photo (RGB).
The Analogy: Think of a master carpenter who can look at a 2D blueprint or a flat photo of a complex chair and instantly "feel" the 3D structure in their mind. GeoAlign trains the robot to do this: it learns to extract the "skeleton" of the object from a flat picture, creating a mental map of the geometry without needing a broken 3D sensor.

2. The Innovation: The "State-Guided" Flashlight

Once the robot has this mental 3D map, it still needs to know where to look. If the robot is reaching for a cup, it needs to focus on the handle. If it's about to pour, it needs to focus on the rim.

The Problem: Most robots look at the whole scene at once, which is too much information.
The Paper's Solution: GeoAlign uses the robot's own body position (its "proprioceptive state") as a flashlight.
The Analogy: Imagine you are in a dark room with a flashlight. You don't need to see the whole room to find a key; you just shine the light exactly where your hand is moving.
- If the robot's hand is reaching out, the "flashlight" shines on the space in front of the gripper.
- If the robot is aligning a ring, the light focuses on the ring's edge.
- The robot's body tells the brain, "I am in this specific pose, so look here for the geometry details I need right now."

3. The Result: Compact "Geometry Tokens"

Instead of feeding the robot a massive, heavy 3D map that slows it down, GeoAlign uses the "flashlight" to grab just the tiny, essential pieces of geometry needed for the next move.

The Analogy: Instead of carrying a whole library of books (the full 3D map) to the kitchen, the robot just grabs a single, sticky note with the exact instruction it needs ("Turn left 5 degrees"). These are called compact geometry tokens. They are small, fast, and contain exactly the spatial info needed to execute the move.

What Did They Prove?

The authors tested this system in three ways:

Simulated Games (LIBERO): The robot solved 99% of tasks, beating previous models. It was especially good at tasks requiring spatial reasoning (like stacking or sliding).
Simulated "Fractal" Tasks: It improved success rates by about 5-6% over standard models.
Real-World Robots (ALOHA): They put the robot on a real table. It successfully handled tricky items like clear tape and transparent bottles that usually confuse robots.
- Example: On a task involving clear tape, the standard robot succeeded only 20% of the time. GeoAlign succeeded 35% of the time.
- Example: For a transparent bottle, the standard robot got 35%, while GeoAlign got 75%.

The Bottom Line

GeoAlign is like giving a robot a super-powered imagination and a smart flashlight.

It learns to "see" 3D shapes from flat photos (even for transparent objects).
It uses its own body position to know exactly which part of that 3D shape matters right now.
It ignores the messy, broken data from depth cameras and focuses on the clean, geometric details needed to actually grab and move things.

The paper concludes that while this doesn't solve every problem (like predicting if a robot will bump into a wall it can't see), it significantly helps robots perform delicate, geometry-heavy tasks that they previously failed at.

Technical Summary: GeoAlign

Problem Statement
Current Vision–Language–Action (VLA) models primarily optimize for semantic grounding, mapping observations and instructions to actions. However, fine-grained robotic manipulation often fails not due to a lack of semantic understanding, but due to a lack of geometry-aware spatial alignment. Executable manipulation requires precise handling of tight clearances, contact-sensitive motions, and interactions with transparent or annular objects—scenarios where high-level semantic representations are insufficient. While some approaches attempt to enrich visual representations with depth or 3D context, reliance on measured depth inputs can fail around transparent objects or thin structures due to sensor limitations. Conversely, methods that incorporate spatial structure directly into action generation often struggle to dynamically select execution-relevant local geometry while maintaining a compact conditioning interface for continuous action decoders.

Methodology
GeoAlign introduces a State-Guided Spatial Alignment architecture designed to bridge the gap between semantic grounding and executable manipulation without relying on raw depth inputs during policy rollout. The framework operates in two distinct stages:

Geometry-Enhanced Post-Training (GEP):
- A Depth Anything V2-Small model is initialized and post-trained on a robot-domain RGB-D dataset using metric depth supervision (optimized via Scale-Invariant Logarithmic loss).
- Crucially, the depth prediction head is discarded after this stage. The retained encoder-side features, now termed GEP features, serve as RGB-derived geometry representations. This allows the model to learn geometric cues from RGB images that are robust to the sensor failures common in depth maps (e.g., transparency).
State-Guided Spatial Alignment:
- During policy training and rollout, the system processes multi-view RGB images and language instructions to generate semantic tokens.
- Simultaneously, a frozen GEP branch maps RGB images to an image-space geometry feature grid.
- Instead of using fixed learned queries or global pooling, GeoAlign generates proprioceptive-state queries. The robot's current state ( $s_t$ ) is encoded and projected into $K=8$ query slots.
- These state-guided queries perform cross-attention over the GEP feature grid to extract compact, phase-dependent geometry tokens. This mechanism allows the model to dynamically select local geometry cues relevant to the current manipulation phase (e.g., reaching vs. inserting).
- The resulting 8 compact geometry tokens are concatenated with semantic tokens to condition a flow-matching Diffusion Transformer (DiT) action decoder (based on the Isaac-GR00T N1.6-3B backbone).

Key Contributions

State-Guided Spatial Alignment: An architecture where the robot's proprioceptive state actively queries a geometry feature grid, extracting compact tokens that inform action generation based on the current phase of manipulation.
RGB-Derived GEP Features: A method to obtain geometry-enhanced features by post-training a depth model with robot-domain supervision, utilizing encoder features for policy conditioning rather than raw depth predictions.
Validation: Comprehensive evaluation across simulation (LIBERO, SimplerEnv-Fractal) and real-world (AgileX ALOHA) environments, with ablation studies isolating the contributions of geometry post-training and state-guided querying.

Results
GeoAlign demonstrates significant performance improvements over baselines across three distinct benchmarks:

LIBERO: Achieved 99.0% average success rate across four suites, outperforming the RGB-only baseline (97.0%). The largest gains were observed in the "Spatial" suite (100.0%) and "Long" suite (96.6%), which require complex spatial reasoning.
SimplerEnv-Fractal: Achieved 85.3% average success across three task families (Pick Coke Can, Move Near, Open/Close Drawer), a +5.7 percentage point improvement over the RGB-only baseline.
Real-World ALOHA: Achieved 78.8% average success on eight geometry-critical tasks involving transparent containers, annular tape rolls, and occluded objects. This outperformed the controlled RGB-only baseline (65.0%) and the $\pi_0.5$ model (67.5%) under the same setup.

Ablation studies confirmed that both the robot-domain geometry post-training and the proprioceptive-state-guided querying are essential for these gains. Specifically, removing post-training dropped LIBERO performance to 95.9%, while removing state-guided queries (using global pooling or learned queries instead) dropped performance to 91.6% and 96.2%, respectively.

Significance and Claims
The paper claims that GeoAlign successfully addresses the local geometry problem in VLA policies by decoupling geometric reasoning from raw depth inputs. By using RGB-derived geometry features shaped by depth supervision but queried dynamically by the robot's state, the model achieves robust spatial alignment for executable manipulation. The authors emphasize that this approach allows the action decoder to receive spatial features whose relevance changes across manipulation phases, improving success in tasks requiring precise alignment and handling of challenging visual geometries (like transparency) where standard depth sensors fail. The work positions itself as a complementary route to existing VLA research, focusing specifically on the spatial alignment required for fine-grained action decoding rather than general semantic understanding.