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

## Abstract

Abstract Robots’ ability to follow language instructions and execute diverse 3D 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. These results demonstrate GravMAD’s strong multi-task learning and generalization in 3D manipulation. Video demonstrations are available at: https://gravmad.github.io .

## 1 Introduction

One of the ultimate goals of general-purpose robot manipulation learning is to enable robots to perform a wide range of tasks in real-world 3D environments based on natural language instructions (Hu et al., ). To achieve this, robots must understand task language instructions and align them with the spatial properties of relevant objects in the scene. Additionally, robots must effectively generalize across different tasks and environments; otherwise, their practical application will be limited Zhou et al. (). For example, if a robot has learned the policy for the task “Take the chicken off the grill”, it should also be able to perform the task “Put the chicken on the grill”. Without this generalization ability, its utility will be greatly reduced.
Recent research in robot learning for 3D manipulation tasks has focused on two mainstream approaches: imitation learning-based methods and pre-trained foundation model-based methods.
Imitation learning-based methods learn end-to-end policies from expert demonstrations in attempt to address 3D manipulation tasks (Walke et al., ; Padalkar et al., ; Argall et al., ). By designing various learning frameworks, such as incorporating different 3D representations (Shridhar et al., ; Chen et al., ; Goyal et al., ), policy representations (Ze et al., ; Ke et al., ; Yan et al., ), and multi-stage architectures (Gervet et al., ; Goyal et al., ), imitation learning-based policies can map perceptual information and language instructions to actions that complete complex 3D manipulation tasks.
However, these policies often overfit to specific tasks (Xie et al., ; Zhang et al., ), leading to significant performance degradation or even failure when applied to tasks that differ from those encountered during training (Brohan et al., ; Zitkovich et al., ).

Figure: 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.
Refer to caption: x1.png

Another line of cutting-edge research seeks to leverage foundation models trained on internet-scale data (OpenAI, ; Yang et al., ) to enhance policy generalization across a variety of tasks (Brohan et al., ; Hu et al., ; Huang et al., ). Unlike traditional imitation learning-based methods, approaches using pre-trained foundation models typically decouple perception, reasoning, and control during manipulation (Sharan et al., ). However, this decoupling often leads to a limited understanding of scenes and manipulation tasks (Huang et al., ), allowing robots to conceptually grasp tasks but failing to accurately complete tasks in 3D environments, resulting in failures.
This underscores a key challenge: both imitation learning-based and foundation model-based approaches struggle to balance precision and generalization when adapting to novel 3D manipulation tasks. Such a challenge raises a crucial question: Can the strengths of both approaches be combined to achieve precise yet generalized 3D manipulation?

To this end, we propose discovering key sub-goals for 3D manipulation tasks as a bridge between foundation models and learned policies, leading to the development of Grounded Spatial Value Maps-guided Action Diffusion (GravMAD), a novel sub-goals-driven, language-conditioned action diffusion framework.
Specifically, a new data distillation method called Sub-goal Keypose Discovery is introduced during the training phase. This method identifies the key sub-goals required for each sub-task stage from the demonstrations. In the inference phase, pre-trained foundation models are leveraged to interpret the robot’s 3D visual observations and task language instructions, directly identifying task sub-goals.
Once the task sub-goals are obtained, the voxel value maps introduced in Voxposer (Huang et al., ) are used to generate the corresponding Grounded Spatial Value Maps (GravMaps). These maps reflect both the cost associated with each sub-goal and the ideal gripper openness. The closer to the sub-goal, the lower (cooler) the cost; the farther away, the higher (warmer) the cost, while also indicating the gripper’s state within the sub-goal range.
Thus, they serve as intuitive tools for grounding language instructions into 3D robotic workspaces.
Finally, the generated GravMaps are integrated with the policy diffusion architecture proposed in 3D diffuser actor (Ke et al., ), forming the GravMAD framework. This enables the robot to utilize 3D visual observations, task language instructions, and GravMaps guidance to denoise random noise into precise end-effector poses.
As shown in Fig. 1, GravMAD effectively combines the precise manipulation capabilities of imitation learning-based methods with the reasoning and generalization abilities of foundation model-based approaches.
We extensively evaluate GravMAD on RLBench (James et al., ), a representative benchmark for instruction-following 3D manipulation tasks. The results show that GravMAD not only performs well on tasks encountered during training but also significantly outperforms state-of-the-art baseline methods in terms of generalization to novel tasks.

In summary, our contributions are: 1) We propose the concept of discovering key sub-goals in 3D manipulation tasks to bridge the gap between foundation models and learned policies. In the training phase, we introduce a data distillation method, Sub-goal Keypose Discovery, to identify task sub-goals. In the inference phase, foundation models are used for this purpose. 2) We generate GravMaps from these sub-goals, translating task language instructions into 3D spatial sub-goals and reflecting spatial relationships in the environment. 3) We propose a new action diffusion framework, GravMAD, guided by GravMaps. It is sub-goal-driven and language-conditioned, combining the precision of imitation learning with the generalization capabilities of foundation models. 4) We design two types of tasks on RLBench: 12 base tasks directly selected from RLBench and 8 novel tasks with modified scene configurations or task instructions. GravMAD achieves at least 13.36% higher success rates than state-of-the-art baselines on the 12 base tasks encountered during training, and surpasses them by 28.63% on the 8 novel tasks, highlighting its strong generalization capabilities.

## 2 Related Works

Learning 3D Manipulation Policies from Demonstrations.
Recent works have employed various perception methods to learn 3D manipulation policies from demonstrations to tackle the complexity of reasoning in 3D space. These methods include using 2D images (Brohan et al., ; Zitkovich et al., ; Jang et al., ), voxels (Shridhar et al., ; James et al., ), point clouds (Chen et al., ; Yuan et al., ), multi-view virtual images (Goyal et al., ; ), and feature fields (Gervet et al., ). To support policy learning, some studies (Ke et al., ; Xian et al., ; Yan et al., ; Ze et al., ) have integrated 3D scene representations with diffusion models (Ho et al., ). These approaches attempt to handle the multi-modality of actions, in contrast to behavior cloning methods that train deterministic policies.
By leveraging 3D representation learning, these policies can accurately complete tasks by accounting for the spatial properties of objects, such as orientation and position. This is especially effective for tasks that closely resemble those encountered during training Ze et al. (). However, these policies often lack the language understanding and generalization abilities of foundation models.
Our method extends the diffusion architecture (Ke et al., ) to fully leverage demonstration data and enhances policy generalization by integrating foundation models, combining the strengths of both approaches.

Foundation Models for 3D Manipulation.
Recent foundation models trained on internet-scale data have shown strong zero-shot and few-shot generalization, offering new opportunities for complex 3D manipulation tasks (Hu et al., ; Zhou et al., ). While some approaches fine-tune vision-language models with embodied data (Driess et al., ; Li et al., ), this increases computational costs due to the large data requirements. Alternatively, foundational vision models can generate visual representations for 3D tasks (Zhang et al., ; ), but they often lack the reasoning capabilities needed for complex tasks.
To address these challenges, some studies leverages large language models (LLMs) as high-level planners (Brohan et al., ; Hu et al., ; Huang et al., ), generating language-based plans executed by lower-level policies. Others utilize LLMs’ code-writing abilities to control robots via API calls or to create value maps for planning robot trajectories (Liang et al., ; Huang et al., ). However, these methods often sacrifice precision due to a rough understanding of complex 3D scenes.
Recent works have combined the reasoning capabilities of foundation models with fine-grained control in 3D manipulation to overcome this limitation (Huang et al., ; Sharan et al., ). For example, Huang et al. () uses pre-trained vision-language models (VLMs) to provide spatial constraints and a nonlinear solver to generate precise grasp poses.
Our method combines the learning power of diffusion architectures with the generalization of VLMs. VLMs generate spatial value maps that guide action diffusion, enabling precise control and multi-task generalization in 3D manipulation tasks.

## 3 Method

In this section, we introduce GravMAD (Grounded Spatial Value Maps Guided Action Diffusion), a multi-task, sub-goal-driven, language-conditioned diffusion framework for 3D manipulation, as shown in Fig. 2. We divide GravMAD’s design into three parts: Section 3.1 defines the 3D manipulation problem, Section 3.2 explains the definition and generation of GravMaps, and Section 3.3 details how GravMaps guide action diffusion in 3D manipulation.

Figure: 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.
Refer to caption: x2.png

### 3.1 Problem Formulation

We consider a problem setting where expert demonstrations consist of a robot trajectory $(o_{1},a_{1},o_{2},a_{2},\ldots)$ and a natural language instruction $\ell\in\mathcal{L}$ that describes the task goal. Each observation $o_{t}\in\mathcal{O}$ includes RGB-D images from one or more viewpoints. Each action $a_{t}\in\mathcal{A}$ contains the 3D position of the robot’s end-effector $a^{\text{pos}}\in\mathbb{R}^{3}$, a 6D rotation $a^{\text{rot}}\in\mathbb{R}^{6}$, and a binary gripper state $a^{\text{open}}\in\{0,1\}$. To avoid quaternion discontinuities, we use the 6D rotation representation (Ke et al., ).
In this setting, we assume that a robotic task is composed of multiple sub-tasks, with each sub-task completed when the robot reaches a sub-goal $g_{t}\in\mathcal{G}$, which specifies the 3D position $g^{\text{pos}}\in\mathbb{R}^{3}$, the gripper openness $g^{\text{open}}\in\{0,1\}$, and the 6D rotation $g^{\text{rot}}\in\mathbb{R}^{6}$. Based on this, we construct a new dataset $\mathcal{D}=\{\zeta_{1},\zeta_{2},\ldots\}$ from expert demonstrations. Each demonstration $\zeta$ consists of trajectories with sub-goals $\{\left(o_{1},g_{1},a_{1}\right),\left(o_{2},g_{2},a_{2}\right),\ldots\}$ and the corresponding language instruction $\ell$.
Our goal is to learn a policy $\pi:(\mathcal{O},\mathcal{L},\mathcal{G})\mapsto\mathcal{A}$, which maps observations $o_{t}$, sub-goals $g_{t}$, and instructions $\ell$ to actions $a_{t}$.
To facilitate sub-task segmentation and efficiently learn the policy, we frame the robot’s 3D manipulation learning problem as a keypose prediction problem following prior works (James & Davison, ; James et al., ; Goyal et al., ; Shridhar et al., ). Our model progressively predicts the next keypose based on current observations and uses a sampling-based motion planner (Klemm et al., ) to plan the trajectory between two keyposes.
In the existing keypose discovery method (James & Davison, ), a pose is identified as a keypose when the robot’s joint velocities are near zero, and the gripper state remains unchanged. Our work filters the task’s sub-goals based on these keyposes to facilitate sub-task segmentation and ensure efficient completion of the overall task.

Figure: 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.
Refer to caption: x3.png

### 3.2 GravMap: Grounded Spatial Value Maps

To tackle generalization challenges in 3D manipulation tasks, we introduce the spatial value maps (GravMap), an adaptation of the voxel value maps proposed by , denoted as $m$. GravMaps are adaptively synthesized based on task variations, translating language instructions into 3D spatial sub-goals and reflecting the spatial relationships within the environment. This provides precise guidance for robotic action diffusion. Each GravMap $m$ contains two voxel maps: (1) a spatial cost map $m_{c}$, with lower values near the sub-goal and higher costs further away, and (2) a gripper openness map $m_{o}$, indicating where the gripper should open or close.
As shown in Fig. 2(a), GravMaps are generated differently for training and inference. In training, they are identified from expert demonstrations using the sub-goal keypose discovery method. During inference, pre-trained models generate them from language instructions and observed images.

GravMap Synthesis with Sub-goal keypose Discovery during Training.
We define each sub-task stage in 3D manipulation as: (1) the process where the robotic end-effector transitions from not touching an object to making contact, or (2) the interaction between the end-effector or tool and a new object, where a series of operations are performed before disengaging. To efficiently segment these sub-task stages and find sub-goals, we build upon the existing keypose discovery method (James & Davison, ) and propose a novel data distillation method called sub-goal keypose discovery.

The sub-goal keypose discovery process iterates over each keypose $K_{p}^{i}\in\{K_{p}\}^{N_{k}}_{1}$, where $N_{k}$ is the number of keyposes in a task. For each keypose, the corresponding observation-action pair $(o_{K_{p}^{i}},a_{K_{p}^{i}})$ is passed to the function $S_{K}$, which outputs a Boolean value to determine whether the given keypose should be discovered as a sub-goal keypose. The decision is made based on whether the keypose satisfies the discovery constraint:
$S_{K}((o_{K_{p}^{i}},a_{K_{p}^{i}}))=\left\{\begin{array}[]{ll}1,&\text{if
discovery constraints are met}\\
0,&\text{otherwise}\end{array}\right.$.
The function $S_{K}$ can incorporate multiple constraints.
In our paper, we define two constraints for $S_{K}$, depending on the type of manipulation task, as shown in Fig. 3:
(1) For grasping tasks, such as “take the chicken off the grill", sub-goal keyposes are discovered based on the following constraints: a change in the gripper’s open/close state and a significant change in touch force.
(2) For contact-based tasks, such as “push the __ button", sub-goal keyposes are discovered solely based on significant changes in touch force.
For more details on sub-goal keypose discovery, please refer to Appendix A.1.

After discovering the sub-goal keyposes, the sub-task stages can be quickly segmented, and the corresponding sub-goals can be identified. The end-effector position $g^{\text{pos}}$ and gripper openness $g^{\text{open}}$ at these sub-goals are then input to the Gravmap generator to generate the GravMaps $m$ for training. The process of the GravMap generator is illustrated in Algorithm 1, adapted from Huang et al. ().

Figure: Algorithm 1 GravMap Generation Process

GravMap Synthesis with Foundation Model during Inference.
During the inference phase, we use pre-trained foundation models to synthesize GravMaps. First, to enable the robot to tie the task-related words with their manifestation in the 3D environment, we introduce a Set-of-Mark (SoM) (Yang et al., )-based Detector. This Detector uses Semantic-SAM (Li et al., ) to perform semantic segmentation on the observed RGB images and assigns numerical tags to the segmented regions. Next, the Detector uses GPT-4o to select task-relevant objects and their corresponding tags from the labeled images as contextual information $\mathcal{C}$.
Based on the task instructions $\ell$ and the context $\mathcal{C}$ provided by the Detector, we apply the LLM-based Planner proposed by to infer a series of text-based sub-goals.
Then, an LLM-based Composer (Huang et al., ) recursively generates code to parse each sub-goal. During execution, the code uses the context $\mathcal{C}$ to obtain the end-effector positions $g^{\text{pos}}$ and gripper openness states $g^{\text{open}}$ corresponding to each sub-goal. Finally, $g^{\text{pos}}$ and $g^{\text{open}}$ are fed into the GravMap generator shown in Algorithm 1, to generate the GravMaps.
Details of this process can be found in Appendix A.2.2.

We synthesize the GravMaps via sub-goal keypose discovery during training or foundation models during inference. GravMaps are then $m$ downsampled using farthest point sampling (FPS) and encoded into token $t_{m}$ with the DP3 (Ze et al., ) encoder, a lightweight MLP network.

### 3.3 GravMaps Guided Action Diffusion

After obtaining the GravMaps, they can be used to guide the action diffusion process, as shown in Fig. 2(b). Before the diffusion process begins, the robot should first perceive the 3D environment.

3D Scene Perception.
Building on previous works (Gervet et al., ; Ke et al., ), we use a 3D scene encoder to transform language instructions and multi-view RGB-D images into scene tokens, enhancing the robot’s 3D scene perception. RGB images are encoded using a pre-trained CLIP ResNet50 backbone (Radford et al., ) and a feature pyramid network. These features are lifted into 3D feature clouds using 3D positions derived from depth images and camera intrinsics. Simultaneously, the CLIP language encoder converts task instructions into language tokens. These tokens interact with the 3D feature cloud to generate scene tokens ($t_{s}$), enabling the robot to capture 3D environmental information.

GravMaps Guided Action Diffusion.
GravMAD builds upon the 3D trajectory diffusion architecture introduced by 3D Diffuser Actor (Ke et al., ) and further integrates GravMap tokens $t_{m}$ to guide the action diffusion process. Specifically, GravMAD models policy learning as the reconstruction of the robot’s end-effector pose using diffusion probabilistic models (DDPMs) (Ho et al., ). The end-effector pose is represented as $e=\left(a^{\text{pos}},a^{\text{rot}}\right)$. Starting with Gaussian noise $e_{K}=\left(a_{K}^{\text{pos}},a_{K}^{\text{rot}}\right)$, the denoising networks $\epsilon_{\theta}^{\text{pos}}$ and $\epsilon_{\theta}^{\text{rot}}$ perform $K$ iterative steps to progressively reconstruct the clean pose $e_{0}=\left(a^{\text{pos}}_{0},a^{\text{rot}}_{0}\right)$:

$$ $\displaystyle a_{k-1}^{\text{pos}}$ $\displaystyle=\alpha\left(a_{k}^{\text{pos}}-\gamma\epsilon_{\theta}^{\text{ pos}}\left(e_{k},k,p,t_{s},t_{m}\right)+\mathcal{N}\left(0,\sigma^{2}I\right) \right),$ (1) $\displaystyle a_{k-1}^{\text{rot}}$ $\displaystyle=\alpha\left(a_{k}^{\text{rot}}-\gamma\epsilon_{\theta}^{\text{ rot}}\left(e_{k},k,p,t_{s}\right)+\mathcal{N}\left(0,\sigma^{2}I\right)\right),$ $$

where $\alpha$, $\gamma$, and $\sigma$ are functions of the iteration step $k$, determined by the noise schedule. $\mathcal{N}\left(0,\sigma^{2}I\right)$ is Gaussian noise. Here, $p$ represents proprioceptive information (a short action history).
The denoising networks use 3D relative position attention layers (Gervet et al., ; Xian et al., ; Ke et al., ), with FiLM (Perez et al., ) conditioning applied to each layer based on proprioception $p$ and denoising step $k$. As shown in Fig. 2(b), after passing through linear layers, $a^{\text{pos}}_{K}$ and $a^{\text{rot}}_{K}$ are concatenated and attend to the 3D scene tokens $t_{s}$ via cross-attention. A self-attention layer then refines this representation to produce end-effector contextual features.
These features are processed by five prediction heads: the Position Head, Rotation Head, Openness Head, Auxiliary Openness Head, and Auxiliary Position Head.

The first two prediction heads predict the noise added to the original pose using the $L1$ norm, with the losses defined as:

$$ $\displaystyle\begin{aligned} &\mathcal{L}_{\text{pos}}=\|\epsilon_{k}^{\text{ pos}}-\epsilon_{\theta}^{\text{pos}}\left(e_{k},k,p,t_{s},t_{m}\right)\|,\\ &\mathcal{L}_{\text{rot}}=\|\epsilon_{k}^{\text{rot}}-\epsilon_{\theta}^{\text {rot}}\left(e_{k},k,p,t_{s}\right)\|,\end{aligned}$ (2) $$

where iteration $k$ is randomly selected, and $\epsilon_{k}^{\text{pos}}$ and $\epsilon_{k}^{\text{rot}}$ are randomly sampled as the ground truth noise.

The third prediction head is used to predict the gripper’s open/close state, and we use binary cross-entropy (BCE) loss for supervision:

$$ $\displaystyle\mathcal{L}_{\text{open}}=\text{BCE}\left(f_{\theta}^{\text{open} }\left(e_{k},k,p,t_{s},t_{m}\right),a^{\text{open}}\right)$ (3) $$

The last two prediction heads enable GravMAD to better focus on the ideal end-effector pose at sub-goals, with the loss functions defined as follows:

$$ $\displaystyle\begin{aligned} &\mathcal{L}_{\text{aux_pos}}=\|g^{\text{pos}}-f _{\theta}^{\text{aux_pos}}\left(e_{k},k,p,t_{s},t_{m}\right))\|,\\ &\mathcal{L}_{\text{aux_open}}=\text{BCE}\left(f_{\theta}^{\text{aux_open}}( e_{k},k,p,t_{s},t_{m}),g^{\text{open}}\right),\end{aligned}$ (4) $$

where $f_{\theta}$ represents the pose prediction network in GravMAD, while $g^{\text{pos}}$ and $g^{\text{open}}$ denote the ground truth sub-goal positions and gripper openness, respectively.

In addition to the losses related to robot actions mentioned above, a contrastive learning loss is applied to enhance feature representations from GravMaps. Positive pairs are features from the same GravMap, while negative pairs come from different GravMaps. In each forward pass, one GravMap is extracted from the dataset, and $N-1$ different GravMaps are randomly generated. The loss maximizes similarity between positive pairs and minimizes it between negative pairs:

$$ $\displaystyle\mathcal{L}_{\text{con}}=-\frac{1}{N}\sum_{i=1}^{N}\log\frac{\exp ({f_{g}}_{i}\cdot{f_{g}}_{i}^{+}/T)}{\sum_{j=1}^{N}\exp({f_{g}}_{i}\cdot{f_{g} }_{j}/T)},$ (5) $$

where $T$ is the temperature parameter, ${f_{g}}_{i}$ represents the feature of the $i$-th sample, and ${f_{g}}_{i}^{+}$ represents the positive feature of the $i$-th sample.

At this stage, the training objective of GravMAD can be formulated by combining the losses from Eq.2, 3, 4, and 5 as follows:

$$ $\displaystyle\mathcal{L}_{\text{GravMAD}}=\mathcal{L}_{\text{open}}+\omega_{1} \cdot\mathcal{L}_{\text{pos}}+\omega_{2}\cdot\mathcal{L}_{\text{rot}}+\omega_{ 3}\cdot\mathcal{L}_{\text{aux_pos}}+\mathcal{L}_{\text{aux_open}}+\omega_{4} \cdot\mathcal{L}_{\text{con}},$ (6) $$

where $\omega_{1},\omega_{2},\omega_{3},\omega_{4}$ are adjustable hyperparameters. For more detailed implementation of GravMap and GravMAD, please refer to Appendix A.

## 4 Experiments

We aim to answer the following questions:
(i) Can GravMAD achieve superior generalization in novel 3D manipulation tasks compared to SOTA models? (See Sec. 4.2)
(ii) Is GravMAD’s performance competitive on the 3D manipulation tasks encountered during training? (See Sec. 4.3)
(iii) What key design elements contribute significantly to GravMAD’s overall performance? (See Sec. 4.4)

### 4.1 Environmental Setup

To thoroughly investigate these questions, we conduct our experiments on a representative instruction-following 3D manipulation benchmark, RLBench (James et al., ). We design two types of tasks to provide a comprehensive evaluation of GravMAD.
1) Base tasks. To evaluate GravMAD’s performance across 3D manipulation tasks encountered during training, we select 12 base tasks from RLBench’s 100 language-conditioned tasks, each featuring 2 to 60 variations in instructions, such as handling objects of different colors or quantities. For each base task, we collect 20 demonstrations for training and evaluate the final checkpoints using 3 random seeds over 25 episodes. Detailed descriptions of these tasks are provided in Appendix B.1.
2) Novel tasks. To further test GravMAD’s generalization capabilities, we modify the scene configurations or task instructions of several base tasks to create 8 novel tasks. These modifications introduce significant challenges for the robot regarding instruction comprehension, environmental perception, and policy generalization, as described in Appendix B.2. For each novel task, we evaluate the final checkpoints trained on the 12 base tasks. We use 3 random seeds over 25 episodes for each novel task.
For all tasks, we use a front-view $256\times 256$ RGB-D camera and a Franka Panda robot with parallel grippers.

Baselines.
We compare GravMAD against various baselines, covering both foundation model-based and imitation learning-based methods.
For the foundation model-based approach, we use VoxPoser (Huang et al., ) as the baseline. VoxPoser leverages GPT-4 to generate code for constructing value maps, which are then used by a heuristic-based motion planner to synthesize robotic arm trajectories. We reproduce this baseline in our tasks using prompt templates from and our SoM-based Detector, with five camera viewpoints in RLBench.
For the imitation learning-based baselines, we select: (1) 3D Diffuser Actor (Ke et al., ), which combines 3D scene representations with a diffusion policy for robotic manipulation tasks. To highlight instruction-following tasks, we use the enhanced language-conditioned version provided by ; and (2) Act3D (Gervet et al., ), which uses a 3D feature field within a policy transformer to represent the robot’s workspace.
Differences between GravMAD and these baselines are detailed in Appendix A.4.

Training and Evaluation Details.
GravMAD runs in a multi-task setting during both the training and testing phases. All models complete 600k training iterations on an NVIDIA RTX4090 GPU, with the final checkpoint selected using three random seeds for evaluation.
During testing, except for the novel task “Push Buttons Light”, which must be completed within 3 time steps, all other tasks must be finished within 25 time steps; otherwise, they are considered failures. Evaluation metrics include average success rate and rank. The success rate measures the proportion of tasks completed based on language instructions. Meanwhile, the average rank calculates the average of each model’s rankings across all tasks, reflecting the model’s overall performance in the tasks.
Two settings are used for generating context $\mathcal{C}$ during testing: Manual and VLM. In the manual setting, we manually provide the Detector with the precise 3D coordinates of task-related objects in the simulation to generate accurate context. In the VLM setting, we use a Detector implemented with SoM and GPT-4o to locate task-related objects and generate context.

**Table 1: Generalization to 8 novel RLBench tasks. Evaluations on 8 novel tasks are conducted using 3 seeds, with 25 test episodes per task, utilizing the final checkpoints from training on 12 base tasks. Performance gains are compared to the best-performing baselines, indicated by underlines.**
|  | Avg. | Avg. | Close | Close Jar | Close Jar | Condition | Meat On | Open Drawer | Stack cups | Push Buttons |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | Success $\uparrow$ | Rank $\downarrow$ | Drawer | Banana | Distractor | Block | Grill | Small | blocks | Light |
| Voxposer (Huang et al., ) | 34.29 | 2.6 | 96.00 ±4.00 | 17.33 ±19.73 | 22.67 ±10.07 | 25.00 ±23.26 | 38.67 ±12.22 | 6.67 ±2.31 | 0.00 ±0.00 | 68.00 ±18.33 |
| Act3D (Gervet et al., ) | 17.83 | 3.5 | 66.67±9.24 | 29.33±9.24 | 41.33±4.62 | 0.00±0.00 | 1.33±2.31 | 2.67±4.62 | 0.00±0.00 | 1.33±2.31 |
| 3D Diffuser Actor (Ke et al., ) | 29.38 | 2.9 | 81.33±6.11 | 48.00±4.00 | 42.67±4.62 | 27.00±10.15 | 0.00±0.00 | 2.67±4.62 | 2.67±2.31 | 30.67±12.86 |
| GravMAD (VLM) | 62.92 | 1.0 | 97.33±2.31 | 84.00±00.00 | 86.67±2.31 | 74.00±11.14 | 45.33±4.62 | 21.33±12.86 | 18.67±2.31 | 76.00±8.00 |
| Performance gain | +28.63 | - | +1.33 | +36.00 | +44.00 | +47.00 | +6.66 | +14.66 | +16.00 | +8.00 |

### 4.2 Generalization performance of GravMAD to novel tasks

In Table 1, we present the generalization performance of models trained on 12 base tasks when tested on 8 novel tasks, along with visualized trajectories from two of these tasks. The results show that changes in task scenarios and instructions negatively impact the test performance of all pre-trained models to some extent. However, GravMAD exhibits superior generalization across all 8 novel tasks compared to the baseline models.
In terms of average success rate, GravMAD outperforms VoxPoser, Act3D, and 3D Diffuser Actor by 28.63%, 45.09%, and 33.54%, respectively. VoxPoser leverages large models to achieve a certain level of performance on novel tasks, but its heuristic motion planner fails to grasp object properties and task interaction conditions, leading to poor results on tasks requiring fine manipulation, as shown in the trajectory visualizations.
Similarly, 3D Diffuser Actor and Act3D struggle to transfer skills from training to novel tasks, primarily due to overfitting to training-specific tasks, which hampers generalization. In contrast, GravMAD uses VLM-generated GravMaps to guide action diffusion, enabling effective object interaction and strong performance on novel tasks. These results clearly demonstrate GravMAD’s superior generalization, surpassing all baseline models.

### 4.3 Test Performance of GravMAD on Base Tasks

**Table 2: Multi-task test results on 12 base tasks. All models are trained on 12 base tasks with 20 demonstrations each. Final checkpoints are evaluated across 3 seeds with 25 test episodes per task. Performance gains are compared to the best-performing baselines.**
| Models | Avg. Success $\uparrow$ | Avg. Rank $\downarrow$ | Close Jar | Open Drawer | Meat off Grill | Slide Block | Put in Drawer |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Voxposer (Huang et al., ) | 15.11 | 4.5 | 12.00±10.58 | 10.67±8.33 | 45.33±24.44 | 0.00±0.00 | 0.00±0.00 |
| Act3D (Gervet et al., ) | 34.11 | 4.3 | 61.33±4.62 | 41.33±4.62 | 60.0±6.92 | 78.67±2.31 | 49.33±10.07 |
| 3D Diffuser Actor (Ke et al., ) | 55.81 | 2.3 | 66.67±2.31 | 88.00±6.93 | 88.00±4.00 | 84.00±0.00 | 94.67±2.31 |
| GravMAD (Manual) | 69.17 | 1.3 | 100.00±0.00 | 76.67±4.62 | 89.33±2.31 | 93.33±2.31 | 78.67±6.11 |
| Performance gain | +13.36 | - | +33.33 | -13.33 | +1.33 | +9.33 | -16.00 |
| GravMAD (VLM) | 56.72 | 2.1 | 100.00±0.00 | 58.67±2.31 | 70.67±2.31 | 80.00±0.00 | 61.33±9.24 |
| Performance gain | +0.91 | - | +33.33 | -29.33 | -17.33 | -4.00 | -33.34 |
| Models | Push Buttons | Stack Blocks | Place Cups | Place Wine | Screw Bulb | Insert Peg | Stack Cups |
| Voxposer (Huang et al., ) | 80.00±13.86 | 16.00±12.00 | 6.67±8.33 | 5.33±2.31 | 4.00±4.00 | 0.00±0.00 | 1.33±2.31 |
| Act3D (Gervet et al., ) | 66.67±2.31 | 0.00±0.00 | 0.00±0.00 | 45.33±2.31 | 6.67±2.31 | 0.00±0.00 | 0.00±0.00 |
| 3D Diffuser Actor (Ke et al., ) | 94.67±2.31 | 13.67±2.89 | 5.33±6.11 | 82.67±2.31 | 29.33±2.31 | 2.67±4.62 | 20.00±0.00 |
| GravMAD (Manual) | 98.67±2.31 | 56.67±4.62 | 5.33±2.31 | 77.33±4.62 | 66.67±6.11 | 32.00±6.93 | 57.33±2.31 |
| Performance gain | +4.00 | +40.67 | -1.34 | -5.34 | +37.34 | +29.33 | +37.33 |
| GravMAD (VLM) | 97.33±2.31 | 51.33±6.11 | 5.33±4.62 | 33.33±4.62 | 54.67±6.11 | 18.67±4.62 | 49.33±2.31 |
| Performance gain | +2.66 | +35.33 | -1.34 | -49.34 | +25.34 | +16.00 | +29.33 |

Table 2 compares the performance of all models on 12 base tasks. GravMAD (Manual) outperforms Act3D and Voxposer across all tasks and exceeds the best baseline, 3D Diffuser Actor, in 9 out of 12 tasks, with an average success rate improvement of 13.36%. Despite the Detector’s coarse SoM positioning affecting GravMAD (VLM)’s performance, it still outperforms Act3D and Voxposer on all tasks, with a 0.91% higher average success rate than 3D Diffuser Actor. These results clearly show that GravMAD remains highly competitive even on previously seen tasks. As long as task-related object positions are accurate, the generated GravMap effectively reflects sub-goals and guides action diffusion, enabling precise execution by GravMAD.
GravMAD (Manual) underperforms 3D Diffuser Actor in the “Open Drawer”, “Put in Drawer”, and “Place Wine” tasks due to slight deviations between the manually provided object positions and the sub-goals. In high-precision tasks, even small deviations can impact performance. For example, in the “Open Drawer” task, the robot needs to grasp the center of the small handle for optimal performance. After manually adjusting the sub-goal to better align with the handle, performance improved. GravMAD (VLM) also struggles in tasks like “Place Wine” due to inaccuracies in the object positions provided by the Detector, especially when Semantic SAM fails to provide precise locations or the camera doesn’t capture the full scene.
For further analysis of failure cases, please refer to Appendix B.3.

### 4.4 Ablations

Figure: 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".
Refer to caption: x5.png

Extensive ablation studies are conducted to analyze the role of each key design element in GravMAD, with the results shown in Fig 4. The following findings are revealed:
1) Impact of replacing GravMaps with specific sub-goal position and openness: Training policies only with sub-goals $g^{\text{pos}}$ and $g^{\text{open}}$ (w/o. GravMap) leads to a significant drop in task success rates. These policies tend to focus on specific points rather than broader regions, making them more prone to inaccuracies from slight positional changes.
2) Importance of both cost map and gripper map in GravMaps: The combination of the cost map and gripper map within GravMaps is essential for guiding the model’s attention to sub-goal locations and ensuring effective gripper usage. The absence of the gripper map causes a moderate decline in performance (w/o. Grip. map). In contrast, omitting the cost map leads to complete task failure, as the model predicts actions with non-unit quaternions, highlighting the cost map’s critical importance (w/o. Cost map).
3) Significance of contrastive learning loss and auxiliary losses: Removing the contrastive learning loss $\mathcal{L}_{\text{con}}$ results in highly similar features from the point cloud encoder, diminishing their effectiveness in action denoising and leading to a decline in model performance (w/o. Contra. loss). Similarly, the absence of auxiliary losses $\mathcal{L}_{\text{aux_pos}}$ and $\mathcal{L}_{\text{aux_open}}$ weakens the model’s focus on sub-goals, leading to a noticeable drop in performance (w/o Aux. loss).
4) Effect of GravMap tokens on guiding rotation actions: Conditioning rotation actions with GravMap tokens in the action diffusion process results in a performance drop, likely due to the inherent nature of rotation actions, which makes them difficult to be guided explicitly through value maps (w. Guided Rot.).
5) Impact of different point cloud encoders on GravMap performance. Replacing the DP3 encoder in GravMAD with PointNet (Qi et al., ) (DP3 Encoder $\rightarrow$ PointNet) or PointNet++ (Qi et al., ) (DP3 Encoder $\rightarrow$ PointNet++) leads to a performance decline. We suspect that lightweight encoders help prevent overfitting to training data details, enhancing GravMAD’s generalization ability across different tasks or unseen data.

## 5 Conclusion and Discussion

In this paper, we introduce GravMAD, a novel action diffusion framework that facilitates generalized 3D manipulation using sub-goals. GravMAD grounds language instructions into spatial subgoals within the 3D workspace through grounded spatial value maps (GravMaps). During training, these GravMaps are generated from demonstrations by Sub-goal Keyposes Discovery . In the inference phase, GravMaps are constructed by leveraging foundational models to directly predict sub-goals. Consequently, GravMAD seamlessly integrates the precision of imitation learning with the strong generalization capabilities of foundational models, leading to superior performance across a variety of manipulation tasks.
Extensive experiments on the RLBench benchmark show that GravMAD achieves competitive performance on training tasks. It also generalizes well to novel tasks, demonstrating its potential for practical use across diverse 3D environments.
Despite its promising results, GravMAD has some limitations. First, its effectiveness relies heavily on prompt engineering, which can be challenging for inexperienced users. Additionally, the visual-language models (VLMs) have limited detection capabilities and are sensitive to changes in camera perspective, affecting performance and preventing optimal efficiency and accuracy. Future work will address these issues to enhance the model’s performance, expand its applicability, and validate its use on real-world robots.

## Appendix A Additional Implementation Details

### A.1 Heuristics for Sub-goal Keypose Discovery

Building on keypose discovery (James & Davison, ), we propose the Sub-goal Keypose Discovery method to identify sub-goal keyposes from demonstrations, focusing on changes in the gripper’s state and touch forces. This is particularly relevant for object manipulation tasks, where the robot’s interactions with objects can be segmented into discrete sub-goals.

The implementation of the Sub-goal Keypose Discovery algorithm starts with a set of pre-computed keyposes, which are frames selected from the demonstration sequence through an initial keypose discovery process. We introduce two functions: touch_change, shown in Algorithm 2, and gripper_change, shown in Algorithm 3, to evaluate whether a keypose qualifies as a sub-goal. The first function checks for significant changes in the gripper’s touch forces, while the second evaluates changes in the gripper’s open/close state. The pseudocode in Algorithm 4 outlines the heuristic steps for identifying sub-goal keyposes.

One current limitation of the Sub-goal Keypose Discovery method is its inability to effectively handle tasks involving tool use, which we plan to address in future research.

Figure: Algorithm 2 touch_change Function

Figure: Algorithm 3 gripper_change Function

Figure: Algorithm 4 Heuristics for Sub-goal Keypose Discovery

### A.2 Details of GravMap Synthesis

#### A.2.1 Training Phase

To facilitate GravMap synthesis, we assign a goal action to each keypose by linking it to the action performed at the nearest future sub-goal. This association enables us to determine the relevant cost and gripper state for different regions of the GravMap. In the first map, $m_{c}\in\mathbb{R}^{w\times h\times d}$, the cost is lower near the positions of the robotic end-effector at these sub-goal keyposes and higher as the distance increases. In the second map, $m_{o}\in\mathbb{R}^{w\times h\times d}$, areas near the end-effector’s position at the sub-goal keyposes reflect the gripper state at the sub-goal, while other areas reflect the gripper state at the current frame.

#### A.2.2 Inference Phase

Figure: Algorithm 5 GravMap Generation

In this section, we introduce the complete pipeline for GravMap generation, as outlined in Algorithm 5. This includes Algorithm 1, which details the process of generating a GravMap from a language instruction $\ell$ and an observed RGB image $\mathcal{O}$. The GravMap generation pipeline integrates VLMs to interpret instructions, ground them in the visual context, and translate them into coarse 3D voxel representations, i.e., the GravMap.

Our pipeline consists of the following three components:

- •
Detector. Starting with an instruction $\ell$ and an observed RGB image $\mathcal{O}$, the RGB image is passed through the Semantic-SAM segmentation model, which labels each object with a numerical tag, producing a labeled image $\mathcal{O}^{\prime}$. The GPT4o-based Detector uses the prompts $\mathcal{P}_{\text{det}}$ and $\mathcal{P}^{\prime}_{\text{det}}$ (adapted from Huang et al. ()) to select relevant objects and obtain their 3D positions. The output is a set of selected objects, or context $\mathcal{C}$, which includes the objects’ identities and their spatial coordinates in the 3D environment. In the VLM setting, the Detector accesses initial RGB images from four views: wrist, left shoulder, right shoulder, and front camera. In the manual setting, precise 3D object attributes are provided from the simulation.
- •
Planner. The GPT4o-based Planner takes the instruction $\ell$, context $\mathcal{C}$, and planner-specific prompts $\mathcal{P}_{\text{plan}}$ and $\mathcal{P}^{\prime}_{\text{plan}}$ (adapted from Huang et al. ()) to infer a sequence of sub-tasks $(st_{1},st_{2},\ldots,st_{i})$. Each sub-task describes an action or interaction needed to fulfill the instruction $\ell$. Progress is tracked based on the robot’s gripper state (open/closed) and whether it is holding an object. The current sub-task is then passed to the Composer for further processing.
- •
Composer. Following  Huang et al. (), the GPT4o-based Composer parses each inferred sub-task $st_{i}$ using corresponding prompts $\mathcal{P}_{\text{com}}$ and $\mathcal{P}^{\prime}_{\text{com}}$. The Composer generates the sub-goal position $g^{\text{pos}}$ and sub-goal openness $g^{\text{open}}$ by recursively generating code. This includes calls to get_cost_map and get_gripper_map, which are triggered by cost map prompt $\mathcal{P}_{\text{cost}}$ and gripper map prompt $\mathcal{P}_{\text{gripper}}$. For example, for a sub-task like “push close the topmost drawer,“ the Composer might generate:
get_cost_map(’a point 30cm into the topmost drawer handle’) and get_gripper_map(’close everywhere’). Natural language parameters are parsed by GPT to generate code that assigns values to $g^{\text{pos}}$ and $g^{\text{open}}$. The final GravMap generator in Algorithm 1 then processes $g^{\text{pos}}$ and $g^{\text{open}}$ to generate the GravMap $m$.

The prompts mentioned above can be found on the website: [https://gravmad.github.io](https://gravmad.github.io).

### A.3 Detail of Model Architecture and Hyper-parameters for GravMAD

**Table 3: Hyper-parameters for GravMAD, including Sub-goal Keypose Discovery, GravMap, model configuration, training, and evaluation.**
|  | Values |
| --- | --- |
| Sub-goal Keypose Discovery |  |
| touch_threshold | 2 |
| Tolerance: delta | 0.005 |
| gripper_threshold | 4 |
| GravMap |  |
| map_size: $S_{m}$ | 100 |
| offset_range: $\beta_{o}$ | 3 |
| radius: $\beta_{r}$ | 3 |
| downsample ratio: $\beta_{d}$ | 4 |
| number of sampled points: $N_{p}$ | 1024 |
| Model |  |
| image_size | 256 |
| token_dim | 120 |
| diffusion_timestep | 100 |
| noise_scheduler: position | scaled_linear |
| noise_scheduler: rotation | squaredcos |
| action_space | absolute pose |
| Train |  |
| batch_size | 8 |
| optimizer | Adam |
| train_iters | 600K |
| learning_rate | 1e^-4 |
| weight_decay | 5e^-4 |
| loss weight: $\omega_{1}$ | 30 |
| loss weight: $\omega_{2}$ | 10 |
| loss weight: $\omega_{3}$ | 30 |
| loss weight: $\omega_{4}$ | 10 |
| Evaluation |  |
| maximal step except push_button_light | 25 |
| maximal step of push_button_light | 3 |

The detailed hyperparameters of GravMAD are listed in Table 3. Additionally, Fig. 5 provides a detailed overview of GravMAD’s modules, including the 3D Scene Encoder and the prediction heads.

(a) The 3D Scene Encoder processes visual and language information separately, merging them via a cross-attention mechanism, with proprioception integrated through FiLM. This allows the model to understand tasks like “Take the chicken off the grill" in a 3D environment.
First, the visual input is processed by a 2D Visual Encoder, transforming image data into feature representations. These 2D features are then passed through a 3D lifting module, converting them into 3D representations using depth information.
Simultaneously, the language input, such as the instruction “Take the chicken off the grill", is encoded into language tokens by the Language Encoder.
Finally, the 3D visual features and language tokens are combined through cross-attention, producing 3D Scene tokens.

(b) Each prediction head consists of Attention layers and an MLP.
The Auxiliary Position Head receives tokens from the previous layer, which first go through cross-attention with GravMap tokens, followed by self-attention to refine the features. The tokens are then passed through an MLP to output the sub-goal end-effector position.
Similarly, the Auxiliary Openness Head takes tokens from the self-attention layer of the Auxiliary Position Head and uses an MLP to predict the sub-goal gripper openness.
The Position Head follows the same process as the Auxiliary Position Head, while the Openness Head mirrors the Auxiliary Openness Head.
The Rotation Head processes tokens with self-attention and an MLP to predict rotation error.

Figure: Figure 5: Detailed description of the modules in GravMAD, including the 3D Scene Encoder and the prediction heads
Refer to caption: x6.png

### A.4 Comparison between GravMAD and other baseline models

Figure: Figure 6: Comparison of GravMAD with Voxposer and 3D Diffuser Actor. Unlike Voxposer, which uses planning, GravMAD leverages GravMaps for learning. Compared to 3D Diffuser Actor, GravMAD employs GravMap tokens to guide the action diffusion process and introduces auxiliary position and openness heads to improve representation learning.
Refer to caption: x7.png

We compare GravMAD with Voxposer (Huang et al., ) and 3D Diffuser Actor (Ke et al., ) in Fig. 6.

(a) Voxposer. We describe our reproduction of Voxposer on RLBench. Voxposer uses our SOM-driven Detector to process the input observation and instruction, generating context information. The Planner then receives this context and outputs a sub-goal, representing an intermediate step necessary for the overall motion plan. The Composer processes this sub-goal, producing three maps: Cost Map, Rotation Map, and Gripper Map. These maps guide the robot’s movement toward the target in the environment. Note that Voxposer’s testing process involves a different number of steps compared to 3D Diffuser Actor and GravMAD, completing only after all LLM inferences are executed.

(b) 3D Diffuser Actor. We use the language-enhanced version of 3D Diffuser Actor as a baseline. 3D Diffuser employs a 3D Scene Encoder to transform visual and language inputs into 3D Scene tokens, providing an understanding of the 3D environment. An MLP encodes noisy estimates of position and rotation into corresponding tokens, which are then fed, along with the 3D Scene tokens, into a denoising network for action diffusion. This network, conditioned on proprioception and the denoising step, includes attention layers, Openness Head, Position Head, and Rotation Head. During diffusion, noisy position/rotation tokens attend to 3D Scene tokens, and cross-attention with instruction tokens enhances language understanding. These instruction tokens are also used in the prediction processes of the Openness, Position, and Rotation heads.

(c) GravMAD (ours). GravMAD shares components with Voxposer, such as the Detector, Planner, and Composer, but incorporates task-specific prompt engineering. Unlike Voxposer, which uses maps for planning, GravMAD encodes these maps into tokens using a point cloud encoder, which are then employed in the action diffusion process. Compared to 3D Diffuser Actor, the key difference is that GravMAD uses GravMap tokens instead of language tokens, improving generalization. Additionally, GravMAD introduces two auxiliary tasks to predict sub-goals, enhancing representation learning.

## Appendix B Addtional Experimental Details

### B.1 Base Task

Figure: Figure 7: Visualization of 12 base tasks.
Refer to caption: x8.png

As shown in Fig. 7, we select 12 tasks from RLBench (James et al., ), each containing 2 to 60 variants in the instructions, covering differences in color, placement, category, and count. In addition to instruction variations, the objects, distractors, and their positions and scenes are randomly initialized in the environment. The templates representing task goals in the instructions are also modified while maintaining their semantic meaning. A summary of the 12 tasks is provided in Table 4.

**Table 4: The 12 base tasks selected from RLBench (James et al., )**
| Task | Variation Type | # of Variations | Avg. Keyposes | Language Template |
| --- | --- | --- | --- | --- |
| close jar | color | 20 | 6.0 | “close the — jar” |
| open drawer | placement | 3 | 3.0 | “open the — drawer” |
| screw bulb | color | 20 | 7.0 | “screw in the — light bulb” |
| meat off grill | category | 2 | 5.0 | “take the — off the grill” |
| slide block | color | 4 | 4.7 | “slide the block to — target” |
| put in drawer | placement | 3 | 12.0 | “put the item in the — drawer” |
| push buttons | color | 50 | 3.8 | “push the — button, [then the — button]” |
| stack blocks | color, count | 60 | 14.6 | “stack — — blocks” |
| insert peg | color | 20 | 5.0 | “put the ring on the — spoke” |
| stack cups | color | 20 | 10.0 | “stack the other cups on top of the — cup” |
| place cups | count | 3 | 11.5 | “place — cups on the cup holder” |
| place wine | count | 3 | 5.0 | “stack the wine bottle to the — of the rack” |

We provide a detailed description of each task below and explain modifications from RLBench origin codebase.

#### B.1.1 Close Jar

Task: Close the jar by placing the lid on the jar.
filename: close_jar.py
Modified: The modified success condition registers a single DetectedCondition to check if the jar lid is correctly placed on the jar using a proximity sensor, discarding the previous condition of checking if nothing is grasped by the gripper.
Success Metric: The jar lid is successfully placed on the jar as detected by the proximity sensor.

#### B.1.2 Open Drawer

Task: Open the drawer by gripping the handle and pulling it open.
filename: open_drawer.py
Modified: The cam_over_shoulder_left camera’s position and orientation were modified to better observe the drawer. The camera was repositioned to [0.2, 0.90, 1.10] and reoriented to [0.5*math.pi, 0, 0].
Success Metric: The drawer is successfully opened to the desired position as detected by the joint condition on the drawer’s joint.

#### B.1.3 Screw Bulb

Task: Screw in the light bulb by picking it up from the holder and placing it into the lamp.
filename: light_bulb_in.py
Modified: No.
Success Metric: The light bulb is successfully screwed into the lamp and detected by the proximity sensor.

#### B.1.4 Meat Off Grill

Task: Take the specified meat off the grill and place it next to the grill.
filename: meat_off_grill.py
Modified: The cam_over_shoulder_right camera’s position and orientation were modified to better observe the drawer. The camera was repositioned to [0.20,-0.36,1.85] and reoriented to [-0.85*math.pi, 0, math.pi].
Success Metric: The specified meat is successfully removed from the grill and detected by the proximity sensor.

#### B.1.5 Slide Block

Task: Slide the block to the target of a specified color.
filename: slide_block_to_color_target.py
Modified: No.
Success Metric: The block is successfully detected on top of the target color as indicated by the proximity sensor.

#### B.1.6 Put In Drawer

Task: Put the item in the specified drawer.
filename: put_item_in_drawer.py
Modified: The cam_over_shoulder_left camera’s position and orientation were modified to better observe the drawer. The camera was repositioned to [0.2, 0.90, 1.15] and reoriented to [0.5*math.pi, 0, 0].
Success Metric: The item is successfully placed in the drawer as detected by the proximity sensor.

#### B.1.7 Push Buttons

Task: Press the buttons of the specified color in order
filename: push_buttons.py
Modified: No.
Success Metric: The buttons are successfully pushed in order.

#### B.1.8 Stack Blocks

Task: Stack a specified number of blocks of the same color in a vertical stack.
filename: stack_blocks.py
Modified: No.
Success Metric: The blocks are successfully stacked according to the specified color and number.

#### B.1.9 Insert Peg

Task: Insert a square ring onto the spoke with the specified color.
filename: insert_onto_square_peg.py
Modified: No.
Success Metric: The square ring is successfully placed onto the correctly colored spoke.

#### B.1.10 Stack Cups

Task: Stack two cups on top of the cup with the specified color.
filename: stack_cups.py
Modified: No.
Success Metric: The cups are successfully stacked with the correct cup as the base.

#### B.1.11 Place Cups

Task: Place a specified number of cups onto a cup holder.
filename: place_cups.py
Modified: No.
Success Metric: The cups are successfully placed onto the holder according to the task instructions.

#### B.1.12 Place Wine at Rack Location

Task: Place the wine bottle onto the specified location on the wine rack.
filename: place_wine_at_rack_location.py
Modified: No.
Success Metric: The wine bottle is successfully placed at the correct rack location and released from the gripper.

### B.2 Novel Task

Figure: Figure 8: Visualization of 8 novel tasks.
Refer to caption: x9.png

As shown in Fig. 8, we create 8 novel tasks that differ from the original training tasks to test policy generalization. These tasks feature scenes and objects similar to those in the training tasks. A summary of the seven tasks is provided in Table 5.

**Table 5: The 8 novel tasks changed based on base tasks.**
| Task | Variation Type | # of Variations | Avg. Keyposes | Language Template |
| --- | --- | --- | --- | --- |
| close drawer | placement | 3 | 2.0 | “close the — drawer” |
| close jar banana | placement | 2 | 6.0 | “close the jar closer to the banana” |
| close jar distractors | color | 20 | 6.0 | “close the — jar” |
| condition block | color, count | 72 | 11.0 | “build a tall tower out of — — cubes, and add a blue block if it exists” |
| meat on grill | category | 2 | 5.0 | “put the<br>— on the grill” |
| open small drawer | placement | 3 | 3.0 | “open the — drawer” |
| stack cups blocks | color | 20 | 10.0 | “Identify the most common color in the block pile, and stack the other cups |
|  |  |  |  | on the cup that matches that color” |
| push button light | color | 20 | 2.0 | “push the button with the same color as the light” |

We provide a detailed description of each novel task below and explain the modifications from the base tasks.

#### B.2.1 Meat on Grill

Task: Place either a chicken or a steak on the grill depending on the variation.
filename: meat_on_grill.py
Base task: meat off grill.
Modified: The task requires placing meat onto the grill, whereas the base task involves removing it. The cam_over_shoulder_right camera’s position and orientation were modified to better observe the drawer. The camera was repositioned to [0.20,-0.36,1.85] and reoriented to [-0.85*math.pi, 0, math.pi].
Success Metric: The selected meat (chicken or steak) is successfully placed on the grill and released from the gripper.

#### B.2.2 Stack Cups Blocks

Task: Identify the most common color in the block pile, and stack the other cups on the cup that matches that color.
filename: stack_cups_blocks.py
Base task: Stack cups.
Modified: The task involves identifying the cup that matches the most common color among the distractor blocks, then stacking the other two cups on top. The base task is simply stacking the cups without considering block colors.
Success Metric: Success is measured when the correct cup is stacked with the other cups based on the color identification and all cups are within the target area defined by the proximity sensor.

#### B.2.3 Close Jar Banana

Task: Close the jar that is closer to the banana by screwing on its lid.
filename: close_jar_banana.py
Base task: close jar.
Modified: The task involves identifying the jar closer to the banana and screwing its lid on, while the base task only requires closing a jar without proximity consideration.
Success Metric: The lid is successfully placed on the jar closest to the banana, confirmed by the proximity sensor.

#### B.2.4 Close Jar Distractor

Task: Close the jar by screwing on the lid, while distractor objects are present.
filename: close_jar_distractor.py
Base task: close jar.
Modified: The task includes distractor objects, such as a button and block, which are colored and placed near the jars. These objects have been encountered during training, adding complexity compared to the base task.
Success Metric: The jar lid is successfully placed on the target jar, confirmed by the proximity sensor.

#### B.2.5 Close Drawer

Task: Close one of the drawers (bottom, middle, or top) by sliding it shut.
filename: close_drawer.py
Base task: open drawer.
Modified: The task involves closing the drawer instead of opening it.
Success Metric: The selected drawer is closed successfully, confirmed by the joint position of the drawer.

#### B.2.6 Open Drawer Small

Task: Open one of the smaller drawers (bottom, middle, or top) by sliding it open.
filename: open_drawer_small.py
Base task: open drawer.
Modified: The task involves opening a smaller drawer compared to the base task, with adjusted camera settings for better visibility.
Success Metric: The selected drawer is opened successfully, verified by the joint position of the drawer.

#### B.2.7 Condition Block

Task: Stack a specified number of blocks and, if the blue block is present, add it to the stack.
filename: condition_block.py
Base task: stack blocks.
Modified: The task involves stacking a specified number of blocks, with an additional requirement to include the blue block if it is present.
Success Metric: The correct number of target blocks are stacked, and if the blue block is present, it is also correctly added to the stack.

#### B.2.8 Push Button Light

Task: Push the button that matches the color of a light bulb on the first attempt.
filename: push_buttons_light.py
Base task: push button.
Modified: The task involves pressing a single button that matches the color of a light bulb. The button must be pressed correctly on the first attempt; repeated attempts are not allowed.
Success Metric: The correct button matching the light bulb’s color is pressed on the first attempt.

### B.3 Failure cases of GravMAD

Figure: Figure 9: Failure cause analysis, including (a) visualization of failure examples; (b) comparison of imprecise labels and expected labels.
Refer to caption: x10.png

In this section, we analyze why GravMAD underperforms compared to the baseline model 3D Diffuser Actor on certain base tasks, particularly in the “Place Wine” task and drawer-related tasks.

As discussed in the main paper, GravMaps represent spatial relationships in 3D space, but this introduces a challenge: areas close to the sub-goal often share the same cost value, as seen in the value map on the right side of Fig. 9 (a). This uniform cost value can mislead the robot into assuming it should complete the sub-goal within that area. For tasks requiring precise actions, such as the “Open Drawer” task, GravMaps’ coarse guidance may lead to suboptimal performance compared to 3D Diffuser Actor.
In the left schematic of Fig. 9(a), the robot must grasp the center of a small handle to achieve optimal performance in the “Open Drawer” task. This high precision demand on the end-effector results in a lower success rate for GravMAD. This limitation extends to the “Put in Drawer” task, which depends on the successful completion of “Open Drawer”. Similarly, in the “Place Wine” task, insufficient predictive accuracy causes the robot to misalign the bottle with the correct slot by one unit, leading to failure.

In the VLM setting, sub-goal accuracy often suffers, as shown in Fig. 9(b), further reducing model performance. These inaccuracies typically arise from two factors: (1) SAM may fail to accurately identify ideal areas, leading to imprecise contextual information from the Detector module for tasks like “Place Wine”, “Open Drawer”, and “Put in Drawer”; (2) the camera’s positioning may not capture the full scene, leaving some task-relevant objects out of view, as seen in tasks like “Meat off Grill”.
To overcome these VLM limitations, potential solutions include: (1) integrating multi-view information into the Detector for a more comprehensive scene observation; and (2) using a more granular segmentation model to provide GPT-4 with a wider range of labels, improving the quality of the context generated by the Detector.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Can you provide more details on the selection strategy of the 12 tasks over 100 language-conditioned tasks from RLBench? Why do you choose 12 and why these 12? The currently designed novel tasks feature scenes and objects similar to those in the training tasks. Can you design more complex novel tasks?
2. Can you evaluate the performance of the proposed method in a real-world experiment? Most setups should be similar to 3D Diffuser Actor and not that hard.
soundness: 2
strengths: 1. The paper is well-written, easy to follow, and presents its ideas clearly. 
2. The approach of combining VoxPoser's planning and perception ability with precise subtask imitation learning from 3D Diffuser Actor is well-motivated and clever to enhance the model’s long horizon task capability to achieve both VoxPoser's generalization and 3D Diffuser Actor's preciseness.
3. The reported scores are evaluated across 3 seeds, which are statically robust.
summary: The paper proposes a pipeline of Grounded Spatial Value Maps-guided Action Diffusion (*GravMAD*), which is a subgoal-driven, language-conditioned action diffusion framework combining VoxPoser (Huang et al., 2023) with subtask imitation learning based on 3D Diffuser Actor (Ke et al., 2024). The idea is essentially to break tasks into sub-goals based on language instructions, allowing auxiliary guidance to help robots conceptually grasp tasks more accurately than directly using pre-trained foundation models.

The division of subtasks is conducted via Sub-goal Keypose Discovery, which is defined to happen when end effectors make contact or disengage with the objects. The proposed method uses pre-computed keyposes and touch&gripper change detection to segment subtasks for training epochs. During inference, the proposed method uses the LLM-based Planner and Composer from VoxPoser (Huang et al., 2023).

With the segmented subtasks, the proposed method first uses VoxPoser (Huang et al., 2023) to generate 3D value maps *GravMaps* and then uses 3D Diffuser Actor (Ke et al., 2024) to conduct subtask imitation learning and generate action sequences with the condition of tokens from *GravMaps*.

The authors test *GravMAD* on 20 RLBench (James et al., 2020) simulated tasks, including 12 original tasks selected from RLBench and  8 modified new tasks. Empirical results show that the proposed framework outperforms both VoxPoser and 3D Diffuser Actor in most seen and novel tasks, which validates GravMAD's superiority in both generalization and preciseness.
weaknesses: 1. The proposed Sub-goal Keypose Discover method is kind of naive. The proposed method simply uses the change in the gripper’s open/close state and the change in touch force to segment a whole task, which could suit simple tasks like moving the chicken but might not work for tasks with longer horizons. This segmenting requires complex reasoning over the concrete task description and comprehensive understanding. During inference, foundation models like VoxPoser can be better at dividing subtasks due to the reasoning ability of LLM. However, for this proposed pipeline where different segmenting methods are used in training and testing, there could even be a domain shift for the subtasks' distribution itself due to unaligned segmenting.
2. The comparison and evaluation are limited. The selection strategy of the 12 tasks over 100 language-conditioned tasks from RLBench is not mentioned. The 8 modified tasks are not justified for their novelty, which might cause a doubt on opportunistic choices. 
3. Since both baselines of VoxPoser and 3D Diffuser Actor have conducted real-world evaluations, a real-world comparison for *GravMAD* will be more convincing.

Though conducted in an A+B paradigm on VoxPoser and 3D Diffuser Actor, this work's engineering contribution is solid and the idea of subtask imitation learning is novel. I will raise my score if more validations are provided.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: * The use of a foundation model, such as GPT-4o, during inference might be time-consuming and could potentially hinder real-time responsiveness in robotic manipulation. Could the authors provide the average inference time for GravMAD at each sub-task stage?
* During training, why not generate GravMap in the same way as during inference? (One possible reason may be that the foundation model’s inference time is too lengthy?) The GravMap generated from demonstration ground-truths may differ in distribution from the GravMap produced by the foundation model, which could impact inference performance.
soundness: 2
strengths: * The authors conduct extensive experiments, including comparisons with different types of baselines and a detailed ablation study.
* The authors provide diverse visualization in the paper and on the website, enhancing readers' understanding.
* The paper is well-organized and easy to comprehend.
summary: This paper introduces GravMAD, a new approach that decomposes 3D manipulation tasks into key sub-goals, generates 3D spatial GravMaps to capture spatial relationships from these sub-goals, and employs an action diffusion policy to learn actions through imitation learning. This design aims to combine the precision of imitation learning with the generalization capabilities of foundation models. Extensive experiments conducted on RLBench show that GravMAD achieves improvements in both training tasks and novel tasks.
weaknesses: * The novelty of the proposed method is limited, as it resembles a combination of previous works, Voxposer and 3D Diffuser Actor. Specifically, the authors use a similar approach to Voxposer by decomposing tasks and generating 3D spatial heatmaps for each sub-task, which are then integrated into the 3D Diffuser Actor architecture. The primary innovation appears limited to the proposed GravMAD framework.
* The comparison between GravMAD and 3D Diffuser Actor on training tasks indicates that, although the introduction of GravMap improves performance on most tasks, GravMAD's performance heavily relies on both the Detector and GravMap. If the Detector’s predictions lack precision, it can even degrade the method's performance, potentially compromising its robustness in scenarios where camera capture is unclear or ambiguous.
* In the ablation study w/o. Cost Map, it is surprising that the method achieves a zero score. Could the authors provide further explanation on why the predicted actions result in non-unit quaternions in this scenario? (This seems to conflict with the statement in Section 3.1, which mentions, "to avoid quaternion discontinuities, we use the 6D rotation representation.") Additionally, if the predicted quaternion is not-unit, we can normalize and convert it into a unit quternion. 
* It is advisable to conduct real-world experiments, as many related works, such as Voxposer and 3D Diffuser Actor, also include real-world demonstrations.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What are the specific criteria for “significant changes” in Sub-goal Keypose Discovery?
2. From Table 2, it is evident that the algorithm exhibits significant differences in performance on similar tasks, such as placing a cup and placing wine. Can you explain the reasons for this?
soundness: 2
strengths: 1. GravMAD shows strong generalization capabilities, achieving at least 13.36% higher success rates than state-of-the-art baselines on 12 base tasks encountered during training and surpassing them by 28.63% on 8 novel tasks.
  2. The manuscript is well-written with a clear structure, logical flow, detailed technical descriptions, and intuitive chart presentations.
summary: This manuscript investigates 3D manipulation tasks through robot learning. Traditional imitation learning excels with familiar tasks but struggles with new, unseen ones. Recent approaches use large foundation models to mitigate this but lack task-specific learning, often leading to failures in 3D environments. To overcome this, the author introduces a novel framework termed Grounded Space Value Map-guided Action Diffusion (GravMAD). Due to its promising results, I am tentatively inclined to accept the paper. However, I find certain aspects unclear, especially the evaluation of design choices and overall performance. I am open to revising my evaluation if the authors can address these concerns comprehensively and convincingly detail the contributions of their work.
weaknesses: - Insufficient argumentation
  1. The manuscript mentions in section 3.2 that different sub-goal keypose discovery rules were designed for different types of manipulation tasks, which seems to lack a reasonable explanation and generalizability.
  2. The manuscript does not adequately justify the exclusion of GravMap’s guidance in rotation prediction, stating that “rotation actions are difficult to be guided explicitly through value maps.”
- Inadequate Evaluation: 
  1.  Although it is claimed that this method is generalized, all experiments were conducted in the RLBench environment, and there is a lack of validation on other simulation environments or real robots.
  2. The main evaluation metric used is success rate, overlooking other important dimensions such as completion time, smoothness of action, etc.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Some are raised in the Weaknesses section.
- What is the difference between GravMaps and the cost maps described in VoxPoser? Is it the downsampling part?
- Why are rotation maps from VoxPoser not used?
- Do you have any intuition on how well your proposed method would scale with more data? Would Imitation Learning catch up at some point (with a lot of training data)?
- Could you expand a bit more about the limitations of your proposed method?
- GravMaps are just an overparameterization of sub-goals at a high level. The fact that the ablation variant (w/o. GravMap) performs so poorly is not very intuitive. Could you provide more details about this ablation variant and maybe you have some intuition about its performance?
soundness: 3
strengths: - The motivation of the paper is strong and well-explained.
- Figures (both in the main paper and the appendix) are clear and very helpful in understanding the proposed method.
- The technical challenge of integrating many components together in a unified framework is well-executed and appreciated.
- Ablation studies are well-executed and provide valuable insights into the method.
- Overall, utilising foundation models together with some demonstrations to achieve both, task generalisation and high precision seems like a great direction.
summary: This paper proposes a novel method for learning from demonstrations while incorporating knowledge from foundation models to tackle robotic manipulation. It learns to predict robot end-effector actions using a diffusion process and enhances the representation using 3D cost maps introduced in previous work. These cost maps are computed using heuristically extracted sub-goals during training, while foundation models are used instead at inference. In this way, the proposed method is able to combine strong generalisation capabilities from foundation models and learn precise robot actions from available demonstrations. The experiments show strong performance when compared to the considered baselines, while ablation studies provide valuable insights into different components of the method.
weaknesses: - Novelty and contribution claims are not very clear:
    - 1) Discovering sub-goals and predicting them to guide robot actions has been done before [1,2,3,4]. Is the novelty here mainly the 3D part combined with foundation models? Maybe the difference could be emphasised more.
    - 2) From the paper, it seems that GravMaps are Cost and Gripper openness maps used in VoxPoser, just renamed. Without properly explaining the difference, it shouldn’t be called a novel contribution.
    - 4) Claiming to design two types of tasks (line 124) and then using base tasks from RLBench for one type can come off as being disingenuous about the contribution claims. Rephrasing the 4th contribution to focus more on the evaluation rather than task design would be beneficial.
- Some parts of the method and design choices need to be better explained:
    - How GravMap tokens are used in the network architecture is not described in the main text of the paper.
    - Explaining in more detail why sub-goals are extracted differently during training and inference would greatly benefit the clarity of the paper. It becomes obvious while reading the paper, but initially, it can cause confusion.
    - It is not clear why rotation maps from VoxPoser are not used. Justifying it would make the overall approach clearer.
    - GravMap Synthesis during Inference is not explained clearly. Lines 305-309 are very dense with information, but it’s very difficult to understand how the sub-goals are actually extracted. If it is done just using existing methods, then contribution claims should be adjusted accordingly. 
- It is not clear if the proposed sub-goal keypose discovery is actually needed. There are no experiments using keyposes, without filtering them. If there are other reasons to filter keyposes, they should be explained.
- Overall, sub-goal keypose discovery seems very tailored towards specific types of tasks. Limitations of it (and how they could be lifted) are not discussed.
- What are 3 random seeds used for? Are the models re-trained using different seeds or just to change object poses and the stochastic diffusion process? If it’s the latter, then it is the same as having 3x episodes and the reported statistics don’t hold a lot of value.
- GravMAD (Manual) seems like an unfair comparison to the baselines. Maybe providing the same ground truth object pose information to other methods would be more valuable for such comparison.
- Given that the proposed method uses sub-goals to guide the diffusion process, selected baselines are not the most appropriate. Better baselines would be ChainedDiffuser [3] and/or Hierarchical Diffusion Policy [4].
- There are some contradictory and not complete claims in the paper:
    - Line 144 - Described methods already leverage demonstration data.
    - Line 205 - Quaternions, although they have the double cover property, are continuous unless additional restrictions are imposed (e.g. by requiring that the scalar part w≥0).
    - Line 206 stated that 6D rotation representation is used, but line 510 states that the model can predict non-unit quaternions. How is this possible?
    - What are “3D tasks”? Tasks that involve 3D observations or 6D manipulation? It is not a very clear term.
- Limitations of the proposed method are only superficially discussed.

[1] Black, Kevin, et al. "Zero-shot robotic manipulation with pretrained image-editing diffusion models." arXiv preprint arXiv:2310.10639 (2023).

[2] Kang, Xuhui, Wenqian Ye, and Yen-Ling Kuo. "Imagined subgoals for hierarchical goal-conditioned policies." CoRL 2023 Workshop on Learning Effective Abstractions for Planning (LEAP). 2023.

[3] Xian, Zhou, et al. "Chaineddiffuser: Unifying trajectory diffusion and keypose prediction for robotic manipulation." 7th Annual Conference on Robot Learning. 2023.

[4] Ma, Xiao, et al. "Hierarchical Diffusion Policy for Kinematics-Aware Multi-Task Robotic Manipulation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

Meta-review: The paper presents a framework for robotic manipulation that combines prior work on action diffusion and robotic control with foundation models. The approach demonstrates solid generalization capabilities and outperforms state-of-the-art baselines. 

Strengths:

- Clear motivation and robust experimental design.

- Notable performance improvements on training and unseen tasks.

- Extensive ablation studies and real-world validations.

Weaknesses:

- Limited novelty beyond combining existing frameworks.

- Challenges with robustness and precision of certain components.

- Experiments are limited to RLBench, raising concerns about broader applicability, which was partially addressed by the authors with additional real-world experiments in the rebuttal.

While the paper has areas that could be improved, such as enhancing methodological contributions and broadening experimental validations, it provides useful insights into robotic manipulation. The authors have addressed most of the reviewers' concerns through detailed rebuttals and revisions, including real-world experiments and clarifications. Therefore, I recommend acceptance.

Decision: Accept (Poster)