EmbodiedMAE: A Unified 3D Multi-Modal Representation for Robot Manipulation

TL;DR

EmbodiedMAE introduces a unified 3D multi-modal representation for robot manipulation, improving performance and efficiency over existing models by effectively integrating RGB, depth, and point cloud data.

Contribution

The paper develops EmbodiedMAE, a novel multi-modal masked autoencoder that learns comprehensive 3D representations, and enhances the DROID dataset with high-quality 3D data for embodied vision tasks.

Findings

Outperforms state-of-the-art vision models in simulation and real-world tasks.

Exhibits strong scaling behavior with model size.

Enables effective policy learning from 3D inputs.

Abstract

We present EmbodiedMAE, a unified 3D multi-modal representation for robot manipulation. Current approaches suffer from significant domain gaps between training datasets and robot manipulation tasks, while also lacking model architectures that can effectively incorporate 3D information. To overcome these limitations, we enhance the DROID dataset with high-quality depth maps and point clouds, constructing DROID-3D as a valuable supplement for 3D embodied vision research. Then we develop EmbodiedMAE, a multi-modal masked autoencoder that simultaneously learns representations across RGB, depth, and point cloud modalities through stochastic masking and cross-modal fusion. Trained on DROID-3D, EmbodiedMAE consistently outperforms state-of-the-art vision foundation models (VFMs) in both training efficiency and final performance across 70 simulation tasks and 20 real-world robot manipulation tasks on two robot platforms. The model exhibits strong scaling behavior with size and promotes effective policy learning from 3D inputs. Experimental results establish EmbodiedMAE as a reliable unified 3D multi-modal VFM for embodied AI systems, particularly in precise tabletop manipulation settings where spatial perception is critical.