RayViT: Ray-Conditioned Visual Representations for Viewpoint-Robust Imitation Learning
2026-07-31 • Robotics
RoboticsComputer Vision and Pattern Recognition
AI summaryⓘ
The authors designed a new method called RayViT to help robots learn tasks better from camera images. Regular images can confuse robots if the camera moves a bit, because they lack clear 3D clues. RayViT adds information about the camera's geometry as a special map to a vision transformer, helping the robot understand where things are in space. Experiments show this makes robot learning more accurate and stable, especially when the camera angle changes.
visual imitation learningvision transformer (ViT)camera geometryPlücker ray mapcross-attentionpositional embeddingscosine similarity lossrobot learningcamera perturbationsRoboCasa benchmark
Authors
Qian Wang, Longrui Chen, Peiran Sun, Aleksandar Taranovic, Niklas Freymuth, Ge Li, Weiran Liao, C. F. Maximilian Nagy, Yucheng Tan, Tao Chen, Gerhard Neumann
Abstract
Visual imitation learning enables robots to acquire visuomotor skills directly from images, yet RGB observations lack explicit geometric cues, making learned policies brittle to camera perturbations. To address this, we propose \textbf{Ray-conditioned Vision Transformer Encoder (RayViT)}, a lightweight architecture that injects camera geometry into pretrained ViT backbones. RayViT represents camera geometry as a Plücker ray map, patchifies it into ray features, and uses gated cross-attention to produce a ray-conditioned class token. These ray features are added as dense positional embeddings, while the ray class token replaces the original ViT class token to provide a geometry-aware summary representation. We combine this approach with an auxiliary cosine similarity loss to consistently improve the performance and robustness for geometry-aware tokens. Experiments on sim- and real-robot tasks demonstrate that RayViT improves robustness by approximately 13 percentage points under camera perturbations in multi-task RoboCasa benchmark and by 1.78 average completed stages in real-world multi-task success rate compared to baselines.