GravMAD: Grounded Spatial Value Maps Guided Action Diffusion for Generalized 3D Manipulation

TL;DR

GravMAD addresses the challenge of generalizing 3D manipulation from language by bridging imitation-learning policies with foundation-model reasoning through GravMaps, which ground task sub-goals in 3D space. It introduces Sub-goal Keypose Discovery for training and leverages pre-trained foundation models to infer GravMaps during inference, enabling a diffusion-based policy to denoise toward precise end-effector poses conditioned on 3D spatial grounding. Empirical results on RLBench show strong generalization to novel tasks (up to a 28.63% improvement in average success) and competitive performance on base tasks, with additional validation on real robots. This work demonstrates a practical pathway for language-conditioned, generalizable 3D manipulation by combining spatial grounding with diffusion-based control.

Abstract

Robots' ability to follow language instructions and execute diverse 3D manipulation tasks is vital in robot learning. Traditional imitation learning-based methods perform well on seen tasks but struggle with novel, unseen ones due to variability. Recent approaches leverage large foundation models to assist in understanding novel tasks, thereby mitigating this issue. However, these methods lack a task-specific learning process, which is essential for an accurate understanding of 3D environments, often leading to execution failures. In this paper, we introduce GravMAD, a sub-goal-driven, language-conditioned action diffusion framework that combines the strengths of imitation learning and foundation models. Our approach breaks tasks into sub-goals based on language instructions, allowing auxiliary guidance during both training and inference. During training, we introduce Sub-goal Keypose Discovery to identify key sub-goals from demonstrations. Inference differs from training, as there are no demonstrations available, so we use pre-trained foundation models to bridge the gap and identify sub-goals for the current task. In both phases, GravMaps are generated from sub-goals, providing GravMAD with more flexible 3D spatial guidance compared to fixed 3D positions. Empirical evaluations on RLBench show that GravMAD significantly outperforms state-of-the-art methods, with a 28.63% improvement on novel tasks and a 13.36% gain on tasks encountered during training. Evaluations on real-world robotic tasks further show that GravMAD can reason about real-world tasks, associate them with relevant visual information, and generalize to novel tasks. These results demonstrate GravMAD's strong multi-task learning and generalization in 3D manipulation. Video demonstrations are available at: https://gravmad.github.io.

Figures (17)

• Figure 1: Comparison of Pipelines. (a) Imitation learning-based methods learn end-to-end policies that map language and 3D observations to actions for precise manipulation. (b) Foundation models-based methods use LLMs/VLMs to process inputs, generate plans, and execute actions with predefined primitives for task generalization. (c)(d) GravMAD combines both, using sub-goal guidance to leverage the language understanding of foundation models and the policy learning of imitation learning for precise and generalized manipulation.
• Figure 2: GravMAD Overview.(a) GravMap Synthesis: During training, we use Sub-goal Keypose Discovery to obtain sub-goals $g^\text{pos}$ and $g^\text{open}$. During inference, the Detector, Planner, and Composer pipeline interprets visual observations and language instructions to derive $g^\text{pos}$ and $g^\text{open}$, which are processed into a GravMap and encoded as a GravMap token. (b) GravMaps Guided Action Diffusion: The policy network perceives the scene and denoises noisy actions guided by the GravMap token. After $K$ denoising steps, the clean actions are executed by the robot.
• Figure 3: Visualization of sub-goal keyposes and sub-task stages. The left sub-figure shows image-based sub-goal keyposes and sub-task stages for "take the chicken off the grill" and "push the __ button" tasks. The right shows the sub-goal key poses and sub-task stages in the trajectory for the "take the chicken off the grill” task.
• Figure 4: Ablation Studies. We evaluate the impact of key design elements by reporting the average success rates across 12 base tasks and 8 novel tasks. In the results, "$\rightarrow$” denotes replacement, "w/o” indicates "without", and "w.” signifies "with".
• Figure 5: Detailed description of the modules in GravMAD, including the 3D Scene Encoder and the prediction heads