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

TL;DR

EmbodiedMAE tackles the domain gap in 3D embodied perception by pretraining a unified RGB-depth-point cloud representation on the large-scale DROID-3D dataset and distilling it into scalable backbones. It employs a multi-modal masked autoencoder with stochastic, per-modality masking and cross-modal decoding to learn robust cross-modal fusion, enabling strong policy learning for tabletop manipulation. Empirical results across 70 simulation tasks and 20 real-world tasks on two robot platforms show state-of-the-art performance, improved training efficiency, and favorable scaling, with RGBD inputs offering the best real-world robustness. The work provides a practical 3D VFM foundation for embodied AI and introduces a high-quality 3D resource (DROID-3D) to accelerate future research, while outlining avenues to extend into language-grounded, instruction-following embodied agents.

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.

Figures (11)

• Figure 1: Overview of EmbodiedMAE Pre-training. We pre-train a ViT-Giant scale multi-modal MAE on the large-scale DROID-3D robot manipulation dataset. We fix the total number of unmasked patches across RGB, depth, and point cloud modalities. The mask ratio allocated to each modality is stochastically sampled. After the Giant model pre-training, we distill it to obtain our Small/Base/Large scale models.
• Figure 2: Depth Quality Comparison. We evaluate depth data quality across several mainstream large-scale embodied AI datasets. Both BridgeDataV2 and RH20T exhibit unreliable and noisy depth information. While prior work has explored the use of AI models for depth estimation, we observe that such methods lack temporal consistency. In contrast, our solution, ZED SDK processing, achieves superior and consistent depth quality.
• Figure 3: EmbodiedMAE Visual Predictions. We evaluate its visual predictions under three settings: (a) Two modalities are almost masked, leaving one modality as the major infer source (column 1-9). (b) Model predicts one modality from another one (column 10-11). (c) Model is allowed to see a modified RGB patch during depth-to-RGB prediction, where the color of the visible patch is altered (column 12).
• Figure 4: Usage Example. We follow the Huggingface Transformers convention to make EmbodiedMAE highly user-friendly and easy to integrate.
• Figure 5: Policy Network for All VFMs. We adopt a compact RDT as the policy network, in which only VFMs are modular.