IMAGINATION POLICY: Using Generative Point Cloud Models for Learning Manipulation Policies
Authors: Haojie Huang
Karl Schmeckpeper
Dian Wang
Ondrej Biza
Yaoyao Qian
Haotian Liu
Mingxi Jia
Robert Platt
Robin Walters
Venue: CoRL 2024 (2024)
Keywords: Manipulation policy learning
Generative model
Geometric learning
TL;DR: Using generative point cloud models for learning manipulation policies

Abstract: Humans can imagine goal states during planning and perform actions to match those goals. In this work, we propose IMAGINATION POLICY, a novel multi-task key-frame policy network for solving high-precision pick and place tasks. Instead of learning actions directly, IMAGINATION POLICY generates point clouds to imagine desired states which are then translated to actions using rigid action estimation. This transforms action inference into a local generative task. We leverage pick and place symmetries underlying the tasks in the generation process and achieve extremely high sample efficiency and generalizability to unseen configurations. Finally, we demonstrate state-of-the-art performance across various tasks on the RLbench benchmark compared with several strong baselines and validate our approach on a real robot.

IMAGINATION POLICY: Using Generative Point Cloud Models for Learning Manipulation Policies
Haojie Huang', Karl Schmeckpeper Dian Wang"', Ondrej Biza Yaoyao Qian"!,
Haotian Liu 23, Mingxi Jia Robert Platt 12, Robin Walters'
1.4 Equal Contribution, 'Northeastern University, Boston, MA 02115, USA
Boston Dynamics Al Institute, *Worcester Polytechnic Institute, BBrown University
fhuang baoj; plpaat:; v.watrers Jnootheasterr.ede https///haojuuang http::/haaojuunggggiubiioiimmme eithubio/mmmiineeppge
Abstract: Humans can imagine goal states during planning and perform actions to match those goals. In this work, we propose IMAGINATION POLICY, a novel multi-task key-frame policy network for solving high-precision pick and place tasks. Instead of learning actions directly, IMAGINATION POLICY generates point clouds to imagine desired states which are then translated actions using rigid action estimation. This transforms action inference into local generative task. We leverage pick and place symmetries underlying the tasks in the generation process and achieve extremely high sample efficiency and generalizability to un- seen configurations. Finally, we demonstrate state-of-the-an performance across various tasks on the RLbench benchmark compared with several strong baselines and validate our approach on a real robot.
Keywords: Manipulation policy learning, Generative model, Geometric learning
Introduction
Humans look scene and imagine how it would look with the objects in it rearranged. For example, given a flower and bottle on the table, can imagine the flower placed in the bottle. Using this mental picture, we can then manipulate the objects to match the imagined scene. How- ever, most robotic policy learning algorithms shortcut this process and map observations directly to actions (3) poses or displacements) [1, 2, 4, 5, 6J. These approaches lose important informa- tion about the desired geometry of the scene and are therefore less sample efficient and less precise than they might be.
Inspired by how humans solve tasks, we propose IMAGINATION POLICY which takes two point clouds as input and generates a new point cloud combining the inputs into a desirable configuration using a conditional point flow model (see Figure 1). Given the generated point cloud, we use point cloud registration methods to match the observed input point clouds with the "imagined" scene. This gives rigid body transformations which can be used to command robot arm to manipulate the objects. IMAGINATION POLICY consists of generative processes, each of which method. As shown in Figure 1a, the pick generator generates the points of the object positioned rel- ative to the gripper point cloud. The place generator generates pair of objects rearranged together as shown in Figure 1b. Compared to directly generating actions, this adds many degrees of freedom to the generative process which aids optimization and sensitivity to geometric interactions.
the above
IMAGINATION POLICY addresses two key challenges in current multitask manipulation policy learning: high precision manipulation and sample efficient learning. Methods like PerAct [1], RVT [2] and Diffuser Actor [3] struggle to Iearn high precision manipulation policies such those required to solve the RLBench [7] tasks Plug-Charger and Insert-Knife. IMAGINATION POLICY outperforms on these tasks by enabling the model to reason about detailed geometric interaction,
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

GGosp the flower*
Put the flower in the mug"
(4) Pick/Single Generation
(b) Place/Puir Generation
Figure f. Illustration of pick generation and place generation. pick generator generates the points of the object to picked conditioned the gripper point cloud. The place generator generates two objects repositioned together. The geuerated points are colored in orange.
such as how different parts an object's surface should displaced, which turn facilitates pre- cise reasoning about the movement of a tool tip to align with a mating surface. Our method also excels in sample efficiency, the ability to learn good policies with relatively few expert demon- strations. Because IMAGINATION POLICY reasons about the desired relative configuration of two objects, it can more easily incorporate symmetries of two object systems, called bi-equivariance, into the model. This significantly improves the sample efficiency of the model. While previous work [8, 9, 10, 11] has also used this kind of bi-equivariant structure to improve sample efficiency, our method the first to apply the idea outside of the pick-place setting in the more general key- frame multitask setting. While IMAGINATION POLICY can still solve pick-place tasks, it can also solve more general manipulation tasks like Plug- Shug-Charger, Insert-Knife, and Open-Microwave.
Our contributions in this paper are follows. 1) We are the first to propose generative point cloud model to estimate desired rigid body motion in keyframe manipulation setting. 2) We show how to implement SE(3) bi-equivariant constraints in this setting. 3) We demonstrate that the resulting method, IMAGINATION POLICY, achieves state-of-the-art performance on several RLBench tasks against several strong baselines.
2 Related Work
Point Cloud Generation. Previous works have explored point cloud generation using VAEs [12, 13] and GANs [14, 15]. Recently, score- based denoising models and normalizing flows [16, 17, 18, 19, 20, 21] have demonstrated the power and flexibility to generate high-quality point clouds. For example, Zhou et al. [18] proposed a probabilistic diffusion approach for unconditional point cloud generation. Luo and Hu [19] and Vahdat et al. [20] formulated conditional point cloud diffusion. PSF [22] achieved fast point cloud generation with rectified flow [21]. Ours, however, generates pick and place point clouds conditioned on the observation that can be used to estimate rigid action to command the robot arm.
Point Clouds in 3D Pick-and-place Manipulation. Point clouds provide a flexible, geometric rep- resentation to encode object shapes and poses. In terms of pick-and-place manipulation, Simeonov et al. [23] used -En3a-iin point features to encode descriptor fields enabling sample efficient policy learning. Simeonov et al. [4] extended Diffusion Policy [24] to work with point cloud obser- vations and to learn multimodal actions. Pan al. [25] propose TAX-POSE which is closely related our method. Pan al. [25] used two segmented point clouds input and directly output new point cloud using weighted summary of target points together with residual predictions. Concur- rently, Fisner et al. [26] adopted TAX-POSE [25] with relative distance inferred by a kernel method. However, it was designed to output the new point cloud directly in one step without penalty for the generated results. Our method, instead, uses generative models to predict the movement of each point iteratively with a velocity model. Moreover, they are limited to single-task training, only work one-step pick-and-place settings, and thus cannot be applied to complex tasks without predefined prior actions. Recently, Shridhar et al. [1], Goyal al. [2], and Ke al. [3] showed impressive multi-task capabilities with transformer-base architectures. However, these methods require hun- dreds of expert demonstrations and cannot successfully learn high-precision tasks. In contrast, our

---

method leverages bi-equivariant symmetry and amortizes the action prediction across multiple tasks. As a result, it can solve high-precision tasks with few demonstrations.
Symmetry in Robot Learning. Robotic tasks defined 3D Euclidean space are invariant to trans- lations, rotations, and reflections which redefine the coordinate frame but do not otherwise alter the task. Recent advancements in equivariant modeling, such as those discussed by [27, 28, 29, 30], offer convenient approach to capturing these symmetries in robotics. Zhu et al. [31] and Huang et al. [32] utilized equivariant models to enforce pick symmetries for grasp learning. Yang et al. [33] proposed an equivariant policy for deformable and articulated object manipulation on top of pre-trained equivariant visual representation. Other works [23, 34, 8, 35, 9, 10, 11, 36, 37, 38, 39] leverage symmetries in pick and place and achieve high sample efficiency. However, they are lim- ited single-task pick-and-place equivariance. As result, they cannot be directly applied to the Mlug-Charger and Insert- Knife tasks without pre-place action. Our proposed method, however, can achieve bi-equivariance in the key-frame and multi-task setting. In addition, we employ equivariant action inference using an invariant point cloud generating process, which is different from previous methods.
3 Method
Problem Statement. Consider dataset D with samples of the form To, E) where Po nx* and P. Rm are point clouds that represent two segmented objects, Pab E (R++m)x33 represents the two objects at the desired configuration described by the lan- guage instruction L, and To E R4x4 and T3 E R4x4 are two rigid transformations in SE(3) rep- resented in homogeneous coordinates that transform P. and P. into the desired configuration, i.c., Pab = To Pa UT, Ps. As shown in Figure 1, for the pick, (Pa, Pb) indicates the gripper and the object pick (the flower). For the place, represents the placement (the mug) and the object arrange (the flower). In either the pick place setting our goal is model the policy function f: (Pa, Po, f) + which outputs the gripper movement E SE(3). We consider placing to include the pre-place action and the place action.
Imagination Policy. We factor action inference into two parts, point cloud generation (Figure 2ab) and transformation inference (Figure 2c). In the first part, we train generative model which, when conditioned on f, generates new coordinate for each point of Pa and P. to approximate Pab, i.e.,
fgen: (Pa, Po, E) + (Pa, Po) where Pa U P, Re Pat- In the second part, we estimate two transfor- mations F. from Pa to Pa. and Tú from P, to P, using singular value decomposition (SVD) [40]. Then, the pick action of the gripper can be calculated as Apick -ff and the pre-place and place action can be estimated as aplace = (fa)-
3.1 Pair Generation for Place
We first explain how the above method works in the place setting Splace: (Pa, Pb, + Oplace- The generative model fgen has two sequential parts, a point cloud feature encoder (Figure 2a) and conditional generator (Figure 2b). Then, we calculate the transformation Oplace from the generated points (Figure 2c). Finally, we prove a condition for when the full method bi-equivariant.
Encoding Point Feature. Given P. f and Pa = {pī )y we first compute feature at each point using two point cloud encoders &a and de- The encoder &a takes the XYZ coordinate and RGB color all points of Pa as input and outputs pointwise features { ) Similarly, { 1m which shares an architecture but has separate parameters.
Φb: Pb +
Generating Points. The combined point cloud Par is generated conditioned on the point features using modified version of Point Straight Flow [22]. This is generative flow model where, inference time, samples X taken by flowing over vector field parameterized by a neural network vg. Initial conditions are given by Xo x Uxx (2
F. (f and Fü {f using
'The pre- pre-place action is the prerequisite perform the place action. An example is shown in Figure 4c.
3

---

Point Cloud
To
Encodar
Fg 007
Point cloud Encoder
lffliws
Pao
Ta
(4)
(0)
(9)
Figure Architecture IMAGINATION POLICY. (a). Encoding the observed point features as Fa and Fb. (b). Conditional pair generation step r with the point fenture f ft language feature. (c). Estimating the rigid transformation (T, and Tú) from the observed point cloud to the generation using correspondence.
place scene random Gaussian z; illustrates
where _ E 1axx and X E Rmx3 which are sampled from scaled Gaussian. The network vo defines the vector field for the ODE:
dX, = no X,, Fa, Fb, () di, LE [0.
where X; is the intermediate point cloud states time t and ft is the encoded language feature of the language description e from CLIP [41]. To solve the ODE, we iteratively update Xi+Ar = X, + vo( (Z)) A.. for A steps. The model is trained by setting the optimal direction at any time as Pab Xo which provides the objective.
min E(Iluo Fo, fe,t) (Pab Xo))l²)
where = + - ()- Intuitively, this sets dX, to be the drift force needed to move X to Pas. Specifically, for single point p E Pa U P.. we sample noise z and a time step t to calculate the intermediate point
+ = (1-1)a0, T T ifp P. where a E {a.
The generator input is p² = ( Se, We then optimize 0 with respect to the loss function defined in Equation
Estimating the Action. Given two sets of points P. U Ps = (ph. P Up) ph. P,, and their corresponding target positions Pab = (Top) Topa: _Tap") (ToP) Top..... Topp) ) we can recover the rigid transformations T. and TB using SVD [40]. However, the output Pa U Po of the generator fgen is not constrained to be given by rigid transforms of the orig- inal two point clouds. Each point may move independently by transformation T such that Pa U Ps = (TIp), T?p2..... ... pa) U (T, P., Tī. Th" pb") We can still use SVD to es- timate the best fitting Te between (Pa. Pa) as well as T3 between (Po, Pb). Assuming Po represents the object to be arranged and Pa represents the placement, as shown in Figure 2, the pre-place or place action can calculated as aplace = (Ta)= 1%
Realizing Bi-equivariance. As noted in prior work [8, 9, 10], place actions that transform object B with respect to another object A are bi-equivariant. That is, independent transformations of object B with 9s E SE(3) and object A with Ia E SE(3) result in a change (a (d'lace = Ia Aplccee 1) to complete the rearrangement at the new configuration. Leveraging bi-equivariant symmetries can generalize learned place knowledge to different configurations and improve sample efficiency.Our placement model constrained to be bi-equivar,ant, with invariant generation during training.
Proposition Assuming
Gaussian noise Xo, the encoded point feature Fa and Fi are invariant rotations then Slace bi-equivariant
Splace (ga Pas 9b Pb) = paace (Pas Po)95
for all pairs of rotations (ga. gb) E SO(3) SO(3).
Proof. If X {=nlk=1 and F. U F, 1 notationiinvarian the intermediate point
states X, = LPab + (1 X) are rotation invariant with fixed Pab. Since all inputs to de are

---

Figure 3. Trajectory of the pick generation process ("grasp the bunana by crown"). Unlike the place generation, our pick generation conditioned the canonicalized gripper point cloud. The generated point clond at each timestep colored in orange.
invariant and the output always approaches Pab. we have
fgen (9a. Pa. 96 Pb) fgen (Pa, Pb)
(4)
With the same generated points Pab. the estimated transformation from rotated observation Pa
(gaPas JaPa: ... JaPa) to Pa is Taga' Similarly, the estimated transformation from g.Po to P..
Tog. Then, the new place action dplace can be calculated as aplace = (Taga -Tbb,
ga T -TT Iaaplace9h -1 which satisfies bi-equivariance.
3.2 Single Generation for Pick
Our pick network fpick: : (Pa, Pb) + apick has similar design to fplace. In this setting, Pa are the points of the gripper and Po are the points of the object to pick. The function fpick differs from fplace in that we only generate the new points for Po conditioned on Pa- Figure 3 illustrates the generation process of grasping the hanana by the crown.
Since the pose and the shape of the gripper are always known, we fix Pa canonical pose, sample only P from Gaussian distribution, and construct Xo P. U X3 We set the target Pab as the union of the canonicalized gripper with the point cloud P. posed it is held by the gripper, i.e., Pab = Pa U ToP.. We only use Ps to calculate the loss:
min E(Iluo(X, Fa, Fo. ( x #
(5)
After estimating Th from (Po, Po). the pick action is calculated as apick = T,
Proposition 2. Assuming on the pick target if the encoded point feature Fa and F& are rotation-invariare invariant: Fpick (Pa 9b Pb) = gb (a,, Pb).
Gaussian noise X.p- Spick is equivariant to rotations
Specifically, there is a rotation Jb acting on Po. the generated points Ps are the same as those without rotation. The estimated transformation from ga Ps to Pb is Tog," and the new pick action can be calculated as dpick = (I195 gof,,1 which realizes the desired equivariance property.
Experiments
Model Architecture Details. The generative models fpick and fplace share the same architecture. Each has two point cloud encoders and a generation network. We select PYCNN [42] as the back- bone of our point encoders, which output 64-dimension feature for each point. We use pre- trained CLIP-ViT32 [41] model as our language encoder and project the language embedding 32- dimension vector with linear layer. The time step is encoded as 32-dimension positional embedding. We also encode binary mask that indicates if the point belongs to P. or Ps as 32-dimension positional embedding. As a result, the generator input of point is 163-dimension vector. We adopt PSF [22] as our generator backbone. Both fpick and fplace are trained end-to-end with the MSE loss defined in Equation and Equation We use the Adam optimizer with an ini- tial learning rate of 10-4, Training takes hours to converge with 200k training steps on a single RTX-4090 graphic card. During inference, we randomly sample Xo from a Gaussian distribution and integrate over v4 with 1000 steps to generate Pat and calculate the action. Generating batch takes 20 seconds.
5

---

(4)
Figure 4, Illustration of the keyframe pipeline IMAGINATION POLICY on Kniffe-Kiife: (a) RGB-D image captured by the front camera and the segieenied point clouds. (b) pick generation, (e) preplace generation. and (d) place generation. The lup row shows the generated points with orange color and the bottomn row demonstrates the configurations of pick, preplace, and place with the calculated rigid transformations.
All point clouds P. and P. in our experiments are captured by RGB-D cameras instead of directly sampling from the ground truth mesh. We first center the point cloud and then downsample by selecting at most one point in each cell of 4mm voxel grid. We further randomly subsample or duplicate to get 2048 points for Pa and Po. To get SO(3) data augmentation to Pa and P. during training. i.e., an SO(3) rotation is sampled uniformly at each training step. This enforces symmetry properties from Proposition We found the results slightly outperform the equivariant point encoder of Vector Neurons [28]. as shown in Table 3. We hypothesize that VN's expressivity is not strong as that of PVCNN.
generation, we apply extensive
9bP) = fgen (Pa. Po). which leads to the desired
4.1 3D Key-frame Pick and Place
We conduct our primary experiments on six tasks shown in Figure 5 from RLbench [7] and compare it with three strong multi-task baselines [1, 2, 3].
3D Task Description. We choose the six difficult tasks from James al. [7] to test our proposed method. Phone-on-Base: The agent must pick up the phone and plug it onto the phone base cor- rectly. Stack- Wine: This task consists of grabbing the wine bottle and putting it on the wooden rack at one of three specified locations. Put-Plate: The agent is asked to pick up the plate and insert it between the red spokes in the colored dish rack. The colors of other spokes are randomly generated from the full set of 19 color instances. Put-Roll: This consists of grasping the toilet roll and slid- ing the roll onto its stand. This task requires high precision. Plug-Charger: The agent must pick up the charger and plug it into the power supply on the wall. Thus is also a sigh-precision task. Insert-Knife: This task requires picking up the knife from the chopping board and sliding it into its slot in the knife block. The different 3D tasks are shown graphically in Figure 5. Object poses randomly sampled at the beginning of each episode and the agent must generalize to novel poses.
Baselines. Our method is compared against three strong baselines: PerAct [1] is the state-of-the-art multi-task behavior cloning agent that tokenizes the voxel grids together with a language descrip- tion of the task and learns a language- conditiiond policy with Perceiver Transformer [43]. RVT [2] projects the 3D observation onto five orthographic images and uses the dense feature map of each image to generate 3D actions. 3D Diffuser Actor [3] is variation of Diffusion Policy [24] that denoises noisy actions conditioned on point cloud features. Comparison with this baseline tests the importance of point cloud generation since this baseline generates actions directly. RPDiff [4] con- sumes segmented Pa and P, and denoises the relative pose iteratively. To make a fair comparison, we adapt [4] to a multi-task policy. See Appendix 6.4 for our implementation details. NDFs [23] and its variation [34] are not included since they require per-object pretraining.
Settings. All methods are trained as multi-task models. There are four cameras (front, right shoul- der, left shoulder, hand) pointing toward the workspace. For our method, we formulate the action sequence as (pick, preplace, place), as shown in Figure 4. Specifically, our method generates the pick action with Spick. and the preplace and place action with fplace simultaneously. We use the ground truth mask to segment P. and P. for RPDiff [4] and our method, as shown in Figure 4a.
Training and Metrics We train our method with 1, 5, 10 demonstrations and train the base- lines with 10 demonstrations. All methods are evaluated on 25 unseen configurations and each

---

Figure 5. 3D pick-place tusks from RLBench 17]. From left to right the tasks Phone-on Baxe, Stack Wine, Put-Plate, Put Rofl, Plag- Charger, and Insers-Knife. The top row shows the initial secne and the bottom row shows the completion state.
Model
demos
phone-on-buse
stack-wine
put-plate
put-roll
plug-charger
insert-knife
IMAGINATION POLICY (ours)
4,00
2.67
1.33
2.78
IMAGINATION POLICY (ours)
5
78.67
97.33
0
1.39
24.00
38.67
IMAGINATION POLICY (ours)
10
90.67
97.33
34.67
23.61
26.67
42.67
RVT [2]
10
56.00
18.67
53.33
0
0
8.00
PerAct [i]
66.67
5.33
12.00
3D Diffusor Actor [3]
10
29.33
26.67
12,00
0
0
0
RPDiff [4]
10
62.67
32.00
5.33
0
2.67
Key-Frame Expert
98.67
100
74.6
56
72
90.6
demonstration episodes for training. Results are averaged over runs. Even with only 5 demos, our method can outperform existing baselines by a significant margin.
Table Performance comparisons benchmark,
(%)
when using
evaluation is averaged over 3 evaluation seeds. We report the mean success rate of each method in Table Since some tasks are very complex, to measure the effects of path planning, we also report the performance of the key-frame formulation used by our method with poses from the expert demonstrations (Key-Frame Expert) as an upper bound on performance.
Results. We report the results of all methods in Table 1. Several conclusions can be drawn from Table I: 1) IMAGINATION POLICY significantly outperforms all baselines trained with 10 demos on all the tasks except Pur-Plate. It can also achieve over 90% success rates in Phone-on-Base and Stack-Wine, 2) For tasks with a high-precision requirement, e.g., Plug-Charger, Insert-Knife and Put-Roll, IMAGINATION POLICY has relatively high success rate while all the baselines fail to learn good policy. 3) IMAGINATION POLICY achieves better sample efficiency and demonstrates few-shot learning performance. With one five demonstrations, it sometimes outperforms the base- lines trained with 10 demonstrations, e.g., IMAGINATION POLICY achieves 97. 22 success rate on Stack-Wine trained with 5 demos while the best baseline can only achieve 32%. We believe this sample efficiency is due to the models equivariance which allows exploit the symmetry inherent in the generation task. In the end, our method underperforms one baseline Pur-Plate. We hy- pothesize that the object in this task is symmetric and hard encode with distinguishable point features, which might result in wrong correspondences when estimating the rigid transformations. Since many complex manipulation tasks can be decomposed sequence of single pick and place, we illustrate that our method can address long- long-horizon tasks in Appendix 6.2.
4.2 Real Robot Experiment
We validated IMAGINATION POLICY on a physical robot. We trained a multi-tase task agent from scratch on tasks using total of just 30 demonstrations. There was no use of the simulated data or pre- training this experiment all demonstrations were performed on the real robot.
Settings. The experiment was performed on URS robot with Robotiq-85 end effector, as shown in Figure 6a. The workspace was a 48cm X 48em region on table. There were three RealSense 455 cameras mounted pointing toward the workspace. We split the workspace into two parts to place the object and the placement. The segmented point cloud was directly obtained by cropping the workspace accordingly. To collect the demonstrations, we released the UR5 brakes to push the arm physically and record data of the form (initial observation, pick pose, preplace pose, place pose). The combined point cloud Pas was constructed with segmented points and the poses. During testing, used Movelt path planner to execute the action sequentially.
7

---

(A) Workssaac Sertings
(b) Mug-Tree
(6) Plug-Flowa Flower
(4) Pour-Ball
Figure 6. Settings and of real-world experiments.
Task
demos
pick completions
# place complctions
# completions trials
Mug-Tree
10
15/15 (100%)
12/15 (80.0%)
12/15
80.0%
Flog-Flower
15/15 (100%)
14/15 (93.3%)
14/15
93.3%
Pour-Ball
14/15 (93.3%)
14/14 (100%)
14/15
93.3%
Table 2. Performance real-world experiments.
Tasks. We evaluate IMAGINATION POLICY on three pick and place tasks, shown in Figure 6bed.
Mug-Trel Tree: The robot needs to pick up the mug and place it on the mug holder. Plug-Flower: This task consists of picking up the flower and plugging into the mug. Pouring-Ball: The agent is asked to grasp the small blue cup and pour the ball into the big green cup.
Results. We collected 10 human demonstrations of each task. Our model trained for 200k SGD steps with the same settings as the simulated experiments. We evaluated 15 unseen configurations of each task. The results are reported in Table 2. Visualizations of the captured observation and the generated actions are shown in Appendix 6.5. Videos can be found in supplementary materials. Our failures are mainly caused by the distortion of observations and motion planning errors. For example, the handle of the green mug in Mug-Tree task might disappear due to sensor noise and calibration, which results in place failure.
Conclusion
In this work, we propose IMAGINATION POLICY, multi-task model for manipulation pick and place problems. It utilizes point cloud generation for key-frame manipulation policy learning by reasoning about the geometric configuration of the goal state. This process amortizes action predic- tion during generation by estimating the drift force of each point. We also analyze the key-frame equivariance of the task and implement it in the model by learning invarianninvariant point features. IMAGINATION POLICY demonstrates high sample efficiency and superior performance on six chal- lenging RLbench tasks against several strong baselines. Finally, we demonstrate that the method can effectively used to learn manipulation policies on a physical robot. We test our design choices using an ablation study on multimodal pick-part dataset in Appendix 6.
One limitation of the formulation in this paper is that it relies on segmented point clouds. We believe state-of-the-art segmentation models [44, 45] are sufficient to provide high-quality masks. Additionally, our generation process takes 20 seconds with 1000 steps to finish point cloud gen- eration. Fortunately, a large number works have studied a range of methods for improving the inference speed of diffusion models [46, 47, 48, 49, 50, 51, 22]. We leave applying these existing techniques future work. Moreover, this paper mainly focuses rigid-object manipulation. We add experiment of articulated object manipulation in Appendix 6.3. Our method might also work well for deformable objects. We will explore articulated objects and deformable objects in fu- ture work. Lastly, this paper assumes a fixed one-to-one correspondence between points in the object point clouds and generated point clouds. However, our pipeline of generation and pose estimation proposed here does not strictly require this. Specifically, one can generate point clouds without a correspondence and then train point cloud registration model to estimate the transformations.

---

Acknowledgments
This project were supported in part by NSF 1750649, NSF 2107256, NSF 2314182, NSF 2134178, NSF 2409351, and NASA 8ONSSC19K 1474. Dian Wang also funded by the JPMorgan Chase PhD fellowship. We would like to thank Jung Yeon Park and Nichols Crawford Taylor for their helpful discussions.
References
[1] M. Shridhar, I.. Manuelli, and D. Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. Conference on Robot Learning, pages 785-799, PMLR, 2023.
[2] A. Goyal, J. Xu, Y. Guo, V. Blukis, Y.W. Chao, and Fox. Rvt: Robotic view transformer for 3d object manipulation. In Conference on Robot Learning, pages 694-710. PMLR, 2023.
[3] T.-W. Ke, N. Gkanatsios, and K. Fragkiadaki. 3d diffuser actor: Policy diffusion with 3d scene representations. arXiv preprint arXiv:2402.10885, 2024.
[4] A. Simeonov, A. Goyal, L. Manuelli, L. Yen-Chen, Sarmiento, A. Rodriguez, Agrawal, and D. Fox. Shelving, stacking, hanging: Relational pose diffusion for multi-modal rearrange- ment. arXiv preprint arXiv:2307.04755 2023.
[5] S. James and A. J. Davison. Q-attention: Enabling efficient learning for rision-based robotic manipulation. IEEE Roboties and Automation Letters, 7(2):1612- 1619, 2022.
[6] P. Sundaresan, S. Belkhale, D. Sadigh, and J. Bohg. Kite: Keypoint-coondttooned policies for semantic manipulation. arXiv preprint arXiv:2306.16603 2023.
[7] S. James, Z. Ma, D. R. Arrojo, and A. J. Davison. Rlbench: The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, 5(2): 2020.
[8] H. Huang, D. Wang, R. Walters, and R. Platt. Equivariant Transporter Network. Proceedings of Robotics: Science and Systems, New York City, NY, USA, June 2022. doi: 10. 1560/RSSS XVII..0
[9] H. Huang, O. L. Howell, D. Wang, Zhu, R. Platt, and Walters. Fourier transporter: Bi- cquivariant robotic manipulation 3d. In The Twelfth International Conference on Learning Representations, 2024. URL https: ttattts://ppeeeeeeeetee //ppnreviiieeeeeee/oeeeeee nett/forum/dd-boawaaa0
[10] H. Ryu, H.-i. Lee, J.-H. Lee, and J. Choi. Equivariant descriptor fields: Se (3)-equivariant energy-based models for end-to-end visual robotic manipulation learning. arXiv preprint arXiv:2206.0832// 2022.
[11) H. Ryu, J. Kim, J. Chang, H. S. Ahn, J. Seo, T. Kim, J. Choi, and R. Horowitz. Diffusion-edfs: Bi-equivariant denoising generative modeling on se (3) for visual robotic manipulation. arXiv preprint arXiv:2309.0266
2023.
[12] A. Brock, T. Lim, J. M. Ritchie, and N. Weston. Generative and discriminative voxel modeling with convolutional neural networks. arXiv preprint arXiv: 160.:00660.. 2016.
[13] Kim, Yoo, Lee, and Hong. Setvae: Learning hierarchical composition for generative modeling of set-structured data. Proceedings of the IEEF/VV Conference on Computer Vision and Pattern Recognition, pages 15059-15068, 2021.
[14] Achlioptas, O. Diamanti, I. Mitliagkas, and L. Guibas. Learning representations and gen- erative models for 3d point clouds. In International conference machine learning, pages 40-49. PMLR, 2018.

---

[15] D. W. Shu, S. W. Park, and J. Kwon. 3d point cloud generative adversarial network based on tree structured graph convolutions. In Proceedings of the IEEE/CVF international conference on computer vision, pages 3859-3868, 2019.
[16] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan, Nor- malizing flows for probabilistic modeling and inference. Journal of Machine Learning Re- search, 22(57))1-644 2021.
[17] G. Yang. Huang, Z. Hao, M.-Y. Liu, S. Belongie, and B. Hariharan. Pointflow: 3d point cloud generation with continuous normalizing flows. In Proceedings of the IEEE/CVF inter- national conference computer vision, pages 4541-4550, 2019.
[18] L. Zhou, Y. Du, and J. Wu. 3d shape generation and completion through point-voxel diffusion. In Proceedings of the IEEE/CVE international conference on computer vision, pages 5826- 5835, 2021.
[19] S. Luo and W. Hu, Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2837-2845, 2021.
[20] A. Vahdat, F. Williams, Z. Gojeie, O. Litany, S. Fidler, K. Kreis, al. Lion: Latent point diffusion models for 3d shape generation. Advances in Neural Information Processing Systems, 1521011111, 2022.
[21] X. Liu, C. Gong, and Q Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint (xXiv:2209003000 2022.
[22] L. Wu, D. Wang, C. Gong, X. Liu, Y. Xiong, R. Ranjan, R. Krishnamoorthi, V. Chandra, and Q. Liu. Fast point cloud generation with straight flows. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9445-9454, 2023.
[23] A. Simeonov, Y. Du, A. Tagliasacchi, B. Tenenbaum, A. Rodriguez, P. Agrawal, and V. Sitz- mann. Neural descriptor fields: Se (3)-vuuivariant object representations for manipulation. In 2022 International Conference on Robotics and Automation (ICRA), pages 6394-6400. IEEE, 2022.
[24] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXhv:2303.04137, 2023.
[25] C. Pan, B. Okorn, H. Zhang, B. Eisner, and D. Held. Tax-pose: Task-specific cross-pose esti- mation for robot manipulation. In Conference on Robot Learning, pages 1783-1792. PMIR, 2023.
[26] B. Eisner, Y. Yang, T. Davchev, M. Vecerik, J. Scholz, and D. Held. Deep se (3)-equivariant geometric reasoning for precise placement tasks. arXiv preprint arXiv:2404. 13478, 2024.
[27] M. Weiler and G. Cesa. General B(2i-Equivariant Steerable CNNs. In Conference on Neural Information Processing Systems (NeurIPS), 2019.
[28] C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. J. Guibas. Vector neu- rons: A general framework for so (3)-equivariant networks. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 12200-12209, 2021.
[29] G. Cesa, L. Lang, and M. Weiler. A program to build E(N)-equivariant steerable CNNs. In International Conference on Learning Representations, 2022. URL httpss//ppeerrrri.rr net/forum?ine www.rrrnmiiddEeeeeeeeeaee
[30] Y.-L. Liao and T. Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint 2rXiv:2206.11990 2022.
10

---

[31] X. Zhu, D. Wang, O. Biza, G. Su, R. Walters, and R. Platt. Sample efficient grasp learning using equivariant models. Proceedings of Robotics: Science and Systems (RSS), 2022.
[32] H. Huang, D. Wang, X. Zhu, R. Walters, and R. Platt. Edge grasp network: A graph- se (3)-invariant approach to grasp detection. In 2023 IEEE International Conference Robotics and Automation (ICRA), pages 3882-3888. IEEE, 2023.
[33] Yang. C. Deng, J. Wu, R. Antonova, L. Guibas, and J. Bohg. Equivact: Sim (qu-equrvariant visuomotor policies beyond rigid object manipulation. arXiv preprint arXiv:2310.16050 2023.
[34] A. Simeonov, Y. Du, Y.-.. Lin, A. R. Garcia, L. P. Kaelbling, T. Lozano--éérre and P. Agrawal. Se (3)-equivariant relational rearrangement with neural descriptor fields. In Conference Robot Learning, pages 355-846. PMLR, 2023.
[35] H. Huang, D. Wang, A. Tangri, R. Walters, and R. Platt. Leveraging symmetries in pick and place. The International Journal of Robotics Research, page 02783649231225777, 2024.
[36] D. Wang, R. Walters, X. Zhu, and R. Platt. Equivariant learning in spatial action spaces. In Conference on Robot Learning, pages 1713-1723. PMLR, 2022.
[37] D. Wang, R. Walters, and R. Platt. -0u))-equivarian reinforcement learning. In International Conference on Learning Representations, 2022. URL https https:/ http:////peeeeeeeeeeeeeeeeeeeeee id=7F9cOhdvfk.
[38] D. Wang, M. Jia, X. Zhu, R. Walters, and R. Platt. On-robot learning with equivariant mod- els. 6th Anmual Conference on Robot Learning, 2022. URL https://openrerrie...t net/ forum?id ddkKuu6b2kkW
[39] D. Wang. J. Y. Park, N. Sortur, L. L. Wong, R. Walters, and R. Platt. The surprising ef- fectiveness of equivariant models in domains with latent symmetry. In International Con- ference on Learning Representations, 2023. URL /ttpps://ppenrvviii net/ fot/forrm?id- P4MUGRM4Acu.
[40] O. Sorkine- Sorkine-Homung and M. Rabinovich. Least-squares rigid motion using svd. Computing, (1):1-5, 2017.
[41] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervi- sion. In International conference on machine learning, pages $748-8763. PMIR, 2021.
[42] Z. Liu, H. Tang, Y. Lin, and S. Han. Point-voxel enn for efficient 3d deep learning, Advances in neural information processing systems, 32, 2019.
[43] A. Jaegle, S. Borgeaud, J.-B. Alayrac, C. Doersch, C. Ionescu, D. Ding, S. Koppula, D. Zoran, A. Brock, E. Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint arXiv:2101 1nn.21177.49999 14795, 2021.
[44] A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, et al. Segment anything. In Proceedings of the JEEF/CVF International Conference on Computer Vision, pages 4015-4026, 2023.
[45] L. Ke, M. Ye, M. Danelljan, Y.W. Tai, C.-K. Tang, F. Yu, et al. Segment anything in high quality, Advances in Neural Information Processing Systems, 36, 2024.
[46] E. Luhman and T. Luhman. Knowledge distillation iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 2021.
[47] T. Salimans and J. Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512. 2022.

---

[48] X. Liu, X. Zhang, J. Ma, J. Peng, et al. Instaflow: One step is enough for high-quality diffusion- based text-to-image generation. In The Twelfth International Conference on Learning Repre- sentations, 2023.
[49] H. Zheng, W. Nie, A. Vahdat, K. Azizzadenesheli, and A. Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning, pages 42390-42402. PMLR, 2023.
[50] Y. Song, P. Dhariwal, M. Chen, and I. Sutskever. Consistency models. arXiv preprint arXiv:2303.01469 2023.
2023.
[51] Luo, Y. Tan, L. Huang, Li, and Zhao. Latent consistency models: Synthesizing high- resolution images with few-step inference. arXiv preprint arXiv:2310.04378, 2023.
[52] B. Calli, A. Singh, A. Walsman, S. Srinivasa, P. Abbeel, and A. M. Dollar. The ycb object and model set: Towards common benchmarks for manipulation research. 2015 international conference on advanced robotics (ICAR), pages 510-517. IEEE, 2015.
[53] E. Coumans and Y. Bai. Pybullet, python module for physics simulation for games, robotics and machine learning. 2016.
[54] C. R. Qi, H. Su, K. Mo, and J. Guibas. Pointnet: Deep learning on point for 3d classification and segmentation. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 652-660,, 2017.
[55] Z. Wu, S. Song, Khosla, Yu, L. Zhang, X. Tang, and J. Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1912-1920, 2015.
12

---

6
Appendix
(a) banana crown
(b) mug handle
(c) spoca neck
(di fork hundle
(e) banana body
(0 mag body
(g) spoon tail
(b) lock head
Figure 7. Visualization of pick generation and place generation. Top row: multimodal pick-part training labels for different objects. Mid row; single generation for pick. Bottom row: pair generation for place. 'The generated point cloud is colored in orange. Note the model takes the rundomly rotated and downsampled cloud input.
Pick/Single Generation
Place/Pair Generation
min
error (9)
(cm)
(9)
(cm)
Imagination Policy
0.16
4.82
0.43
0.94
1.72
2.34
7.12
0.55
1.05
1.76
w/o downsample
0.47
7.77
39.55
0.43
0.96
1.77
0.29
9.28
28.94
0.65
1.19
1.85
w/o color augmentation
0.76
4.41
16.18
0.18
0.87
2.41
0.39
5.05
21.22
0.40
0.97
1.76
17.45
125.26
179.27
0.44
1.50
6.53
49.24
130.60
178.59
1.31
14.46
30.80
PointNet Encoder
0.42
3.30
10.09
0.26
0.88
1.56
0.96
4.82
14.31
0.33
0.98
1.74
Pretrained VN Encoder
0.75
5.24
34.06
0.26
0.80
1.56
0.92
6.01
20.68
0.44
1.04
2.13
Table 3. Ablation Results. We report minimum, mean, and maximum emor for single generation and pair generation over 100 runs with rundoml rotated and sampled isput.
6.1 Ablation Study
Multimodal Pick-part Dataset. To quantitatively measure the point cloud generation results and the equivariance IMAGINATION POLICY, we create a small pick-part dataset using four YCB objects [52] (banana, mug, spoon and fork). We load each object in the Pybullet simulator [53] and three cameras get the RGB-D images to extract the point cloud. Each object is assigned with two different expert grasps with corresponding language instructions, e.g., "grasp the mug by the handle", "grasp the mug by its body", as shown in Figure 7.
Training and Metrics. We trained a single pick generation model to generate all the objects con- ditioned on the canonicalized gripper points and language descriptions. We also train place gen- eration model generate both the gripper point cloud and the object point cloud. To evaluate the pick generation results, randomly rotate and randomly downsample the object point cloud (P,) to make starting pose unseen during training. We calculate the translation error and rotation error between the estimated grasp pose and the ground truth grasp pose. Note that rotating Pg results in change of the ground truth pick pose. To evaluate the place generation results, we randomly rotate and downsample the gripper (Pa) as well as the object (Pt) to make a scene unseen during training and calculate the translation error and rotation error between estimated transformation T,11, and the ground truth pose. Note that rotating either Pa or Ps changes the relative ground truth transfor- mation. We report the minimum, mean and maximum error over 100 runs in Table 3. We also show visualizations of the generated point clouds in orange in Figure 7.
Results. Table 3 includes 6 variations of our proposed methods. Several findings can be concluded from Table 3: (1) As shown in the first row, IMAGINATION POLICY can learn the multimodal dis- tribution and is equivariant. It realizes around 2 ~ 3" average rotation error and ccm translation error with different configurations of P. and Ps; (2) Without downsampling color information, the rotation enror slightly increases: (3) Without data augmentation in training, the performance de- creases dramatically since the model cannot learn rotation-invarian features. (4) Compared with
13

---

the results in the last two rows, the PVCNN-baae based point cloud encoder outperforms PointNet [54] and the pre-trained equívariant point cloud encoder from NDF [23]. Note that the pre-trained point cloud encoder consumes enormous 3D point clouds from ShapeNet [55] and makes use of Vector Neuron [28] which guaranteed to output the rotation invariant feature. We hypothesize that the architecture of Vector Neuron [28] and the standard representation limit its expressivity.
6.2 Task with Longer Horizon
(8)
(c)
(4)
f
Ilusstaaionn of Stack-three-Chiin From right: (a), initial observation, pick the bluc chair, (c). place the blue chair, (f), complete state.
green chair,
place the green chair,
Task
demos
finst pick place success rale second pick place success rule
overall success rase
oracle performance
hhre--Chairs
70.66% 4 2.61%
4.525%%
45.05% 3.71%
Table Performance Stack- Three-Chairs. Success over runs. For each test, the poses of the three chairs are randomly sampled with different seed from the training data.
25
10 domoostration episodes training. Results averaged
Many challenging robotic manipulation problems can be viewed through the lens single pick and place operation. We test IMAGINATION POLICY on the task with a longer horizon. Stack-three- Chairs requires picking the other two chairs and stacking them on top of the base (red) chairs follow- ing the RGB order. This a high-precision task requiring the agent correctly manage a sequence of pick and place. Even trained with 10 demos, first pick-place execution and maintain similar performance (68.00%) for the second pick-place execution. Detailed results are reported in Table 4. It demonstrates that IMAGINATION POLICY can address long-horizon tasks.
method can achieve 70.66% success rates in the
6.3 Task with Articulated Object
Figure 9. Illustration of Open-Opennneeee Meee-Meeewaee. From left to right: (a). initial observation, (h). pick the handle of the microwave, (e). open the dear the microwave, (d). final complete state, (e). segmentations of the door with handle (black color) and framne (red color).
Task Moen-Miiroowaee
demos training
success rate 69.33 %c 5.22%
oracle performance 90.66% +%% 2.61%
Table S. Performance on Open Microwave. Success rate (%) on 25 tests using 10 demonstration episodes for training. Results are averaged over runs. For each test, the pose the microwave randomly sampled with a different seed from the training data.
Articulated objects are special cases in manipulation since they are linked with several movable parts. We test IMAGINATION POLICY on Open-Microwave task to illustrate the potential of our method of addressing articulated object manipulation. Specifically, parts of the microwave, shown in Figure 9c, the door with the handle (black color) and the frame
segment the two movable
14

---

(red color). The task consists of grasping the handle of the microwave and opening the door. In our settings, the grasping is to infer the relative pose between the gripper and the door; the opening is to predict the relative pose between the door and the frame. Even trained with 10 demonstrations, our method can achieve 69.33% success rate, as reported in Table 5. We found that most failure cases are due the motion planning error and the collision between the door and the gripper when the gripper is closing. More complex articulated objects the relative poses between links and we leave it as the future work. Overall, it demonstrates that IMAGINATION POLICY can address articulated object manipulation.
also be manipulated by predicting
6.4 Baseline Details
The closest standard methods evaluated on RLBench [7] are PerAct [1], RVT [2], and 3D Diffuser Actor [3]. These are multi-task key-frame methods that address a problem setting very similar to ours. PerAct, RVT, and 3D Diffuser all create point clouds using RGBD data captured from four different camera views. This is exactly the pipeline in our method well. Specifically, PerAct transforms the raw RGBD input into a point cloud and then into a voxel map. RVT constructs the point cloud and then re ce-projects it onto orthographic images. 3D Diffuser Actor is also conditioned on the entire point cloud. The only difference concerning the input data between our method and these baselines is that we use per-object segmentation masks.
RPDiff [4] the baseline that consumed the same segmented point cloud (Pa and Pa) as ours. iteratively denoises the randomly sampled relative transformation poses conditioned on the current configuration of Pa and Pb. However, it only solve the single-step place problem trained single-task policy. To make fair comparison, model. Similar to our settings, we consider the pick problem as inferring the relative pose between the gripper (Pa) and the object to grasp (P)). The preplace action prediction can also be viewed as calculating the relative pose between the object (P,) and the placement (P)). We train the pick model and the place model separately, which is similar to ours. Since the original RPDiff learns a single-task single-step policy f: (Pa. Pb) + Tab. requires training 18 different models solve the six tasks in Table 1. We adapt it learning multi-task policy conditioned on the language embedding f: (Pa, Pe. fi) + Tab. Specifically, we used the same language embedding generated from CLIP [41]. To make RPDiff consume the language embedding, we map it to 128-dimension feature via a linear layer and concatenate it to 130-dimension time step embedding (the diffusion step). The model was trained and evaluated with the same settings in [4].
adapted it to multi-task key-frame prediction
(a)
(b)
(4)
Figure 10. Action inference Mug Tree with real data: (a) the observed real sensor point cloud and the inferred pick, preplace and place action from IMaciNation POLICY, (b) pick generation, (c) preplace generation, and (d) place generation. The top no shows the generated points with orange color the bottom demonstrates the conligurations pick, preplace, place transformation.. Please note that we used the point cloud from Franka- Emika Punda gripper to trinn the model and evalunted it with the Robotiq-85 gripper.
calculated rigid
15

---

6.5 Real-robot Experiments Pipeline
Figure 10 illustrates our pipeline key-frame action inference on Mug-Tree task with real-sensor data. The observed point cloud is shown in the first row of Figure 10a. The predicted pick, preplace and place action from IMAGINATION POLICY are plotted with RGB frames in the second row of Figure 10a. Specifically, Figure 10bed illustrate the pick generation, preplace generation, and place generation respectively.
For execution on robot, it requires a collision-free pick-and-place trajectory that connects the key- frame action. We use RRT-star as our motion planner and add the configuration of obstacles (the table, the mounting, and the cameras) to the planner to generate the trajectory.
6.6 Detailed Results on RLbench task
We report the results of our method and baselines RLbench tasks with 1.998 std emror in Table 6.
Model
demos
phone on huse
stack wine
put plate
pat-roll
plug churger
insert knife
IMAGINATION POLICY (ours)
4,00
+
1.33 2.61
+
2.72
IMAGINATION POLICY (ours)
10.45
24.00
IMAGINATION POLICY (ours)
10
90.67
+ 2.61
97.33
+ 2.61
34.67 + 10.45
23.61
5.44
26.67
13.82
42.67
+ 9.42
RVT
56.00
4,52
18.67
+ 2.61
53.33
+ 6.91
0
8.00
4.52
PerAct
66.67
11.39
5.33
- 2.62
12.00
+ 4.52
0
Diffusor 3D [3] RPDaI]4]
10
29.33
5.22
26.67
14.35
12.00
=
0
62.67
5.22
32.00
4.52
5.33
5.22
2.67
2.61
Discrete Expert
98.67
100
74.6
56
90.6
Table 6. Detailed performance comparisons on RL benchmark. Success rate (9) on 25 tests v.s. the number of demonsitation episodes (1,
10) training. Results averaged runs, with only demos, method outperform existing baselines significunt margin.
16


Reviews:
[{"review_id": "NJz0u4Wfvp", "reviewer": "Reviewer_1GT1", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 5, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The authors introduce a novel approach to generate pick and place actions. The proposed method receives as input two pointclouds and diffuse them to generate a desired relative configuration between them. Then, the rigid body transforms are recovered through SVD.", "summary_of_recommendation": "The method is very original, but the selected experiments are not properly selected to evaluate the benefit of their proposed method. Then, it is hard to validate if the idea is useful for solving robotics problems.", "technical_quality": 1, "clarity_of_presentation": 4, "potential_impact": 1, "robotics_focus": 3, "quality_of_limitations": 1, "questions_for_rebuttal": "- Given the originality and novelty of their approach, the authors should prove experimentally that the proposed approach to generate rigid bodies (pointcloud diffusion + SVD) is better (or at least to the level) than directly generating rigid body (SE(3) diffusion).", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DBOGoBo6On", "reviewer": "Reviewer_U3RP", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper proposes Imagination Policy, a method for manipulation policy learning that uses generative flow models to synthesize target points clouds of pick and place motions. For the place phase, the method takes point clouds of the picked object and the receptacle as input, encodes them into SO(3)-invariant point features, generates a joint point cloud containing the two objects in the target pose, and estimates the rigid transformations for both objects between initials poses before and after placing. For the pick phase, the method takes point clouds of the end-effector and the object to be picked up, and perform similar operations, except that the end-effector point cloud is fixed during generation. The method realizes encoder equivariance through extensive SO(3) data augmentation.", "summary_of_recommendation": "This paper introduces an interesting idea that suffers from limitations that prevent the proposed method from being used as general-purpose manipulation policies. The clarity of writing needs improvement and the authors need to perform more experiments to make a convincing argument. I recommend the authors spend more time to improve this submission.", "technical_quality": 3, "clarity_of_presentation": 4, "potential_impact": 2, "robotics_focus": 4, "quality_of_limitations": 2, "questions_for_rebuttal": "See above.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CxNwEcPRvK", "reviewer": "Reviewer_1BKy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 2, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This work proposes to use diffusion to generate target point cloud of two objects, and predicts tranformation of objects for manipulation.", "summary_of_recommendation": "The paper successfully predicts where objects should go based on initial point clouds and instructions, showing good results with extensive experiments. However, it is possible to use simpler method to solve the task, and current method doesn’t account for more complex situations like avoiding collisions.", "technical_quality": 4, "clarity_of_presentation": 4, "potential_impact": 2, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "1. Is it possible to directly predict two objects' SE(3) poses? Is it possible for a sample-based grasp model to predict the 6D pose?\n\n2. Is it possible to extend this work to consider collision or generate multiple targets?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Strengths:
+ Novel idea of using generative point-cloud models and flow models for manipulation.
+ The experimental setup includes evaluations on several simulated manipulation tasks.
+ The figures in the paper are informative. 

Weaknesses:
- The approach assumes segmented point-clouds of objects, whereas the baselines – RVT, PerAct, 3D Diffusor Actor – do not make this assumption. All reviewers highlighted that this difference makes the comparison unfair. 
- The formulation is limited to three-step-process tasks: pick, pre-place, and place. It is also limited to rigid-body transformations and does not reason about collisions.
- The writing can be greatly improved. The flow of the introduction is abrupt. Real-robot experiments are in the appendix. 
- Evaluations are missing comparisons to SE(3)-equivariant approaches such as Simeonov et al 2023, 2022.

Post Rebuttal:  
The authors addressed majority of the concerns and added another baseline with segmented input.
Decision: Accept