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
equal  contribution
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 05:58:13

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.

Main paper: 
 introduction: 
 Humans can look at a scene and imagine how it would look with the objects in it rearranged. For example, given a flower and a bottle on the table, we can imagine the flower placed in the bottle. Using this mental picture, we can then manipulate the objects to match the imagined scene. However, most robotic policy learning algorithms shortcut this process and map observations directly to actions (SE(3) poses or displacements) [1, 2, 3, 4, 5, 6] . These approaches lose important information 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 a robot arm to manipulate the objects. IMAGINATION POLICY consists of two generative processes, each of which uses the above method. As shown in Figure 1a , the pick generator generates the points of the object positioned relative to the gripper point cloud. The place generator generates a 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.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 learn high precision manipulation policies such as 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, such as how different parts of an object's surface should be displaced, which in turn facilitates precise 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 demonstrations. 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 is the first to apply the idea outside of the pick-place setting in the more general keyframe multitask setting. While IMAGINATION POLICY can still solve pick-place tasks, it can also solve more general manipulation tasks like Plug-Charger, Insert-Knife, and Open-Microwave.Our contributions in this paper are as follows. 1) We are the first to propose a generative point cloud model to estimate desired rigid body motion in a 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. 
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 a rigid action to command the robot arm.Point Clouds in 3D Pick-and-place Manipulation. Point clouds provide a flexible, geometric representation to encode object shapes and poses. In terms of pick-and-place manipulation, Simeonov et al. [23] used SE(3)-invariant point features to encode descriptor fields enabling sample efficient policy learning. Simeonov et al. [4] extended Diffusion Policy [24] to work with point cloud observations and to learn multimodal actions. Pan et al. [25] propose TAX-POSE which is closely related to our method. Pan et al. [25] used two segmented point clouds as input and directly output a new point cloud using a weighted summary of target points together with residual predictions. Concurrently, Eisner 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 in one-step pick-and-place settings, and thus cannot be applied to complex tasks without predefined prior actions. Recently, Shridhar et al. [1] , Goyal et al. [2] , and Ke et al. [3] showed impressive multi-task capabilities with transformer-based architectures. However, these methods require hundreds 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 in 3D Euclidean space are invariant to translations, 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 a 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 limited to single-task pick-and-place equivariance. As a result, they cannot be directly applied to the Plug-Charger and Insert-Knife tasks without a 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. 
method: 
 Problem Statement. Consider a dataset D with samples of the form (P a , P b , T a , T b , P ab , ℓ) where P a ∈ R n×3 and P b ∈ R m×3 are point clouds that represent two segmented objects, P ab ∈ R (n+m)×3 represents the two objects at the desired configuration described by the language instruction ℓ, and T a ∈ R 4×4 and T b ∈ R 4×4 are two rigid transformations in SE(3) represented in homogeneous coordinates that transform P a and P b into the desired configuration, i.e., P ab = T a • P a ∪ T b • P b . As shown in Figure 1 , for the pick, (P a , P b ) indicates the gripper and the object to pick (the flower). For the place, it represents the placement (the mug) and the object to arrange (the flower). In either the pick or place setting our goal is model the policy function f : (P a , P b , ℓ) → a which outputs the gripper movement a ∈ SE(3). We consider placing to include the pre-place action and the place action. 1Imagination Policy. We factor action inference into two parts, point cloud generation (Figure 2ab ) and transformation inference (Figure 2c ). In the first part, we train a generative model which, when conditioned on ℓ, generates a new coordinate for each point of P a and P b to approximate P ab , i.e., f gen : (P a , P b , ℓ) → ( Pa , Pb ) where Pa ∪ Pb ≈ P ab . In the second part, we estimate two transformations Ta from P a to Pa , and Tb from P b to Pb using singular value decomposition (SVD) [40] . Then, the pick action of the gripper can be calculated as a pick = ( Tb ) -1 Ta and the pre-place and place action can be estimated as a place = ( Ta ) -1 Tb .We first explain how the above method works in the place setting f place : (P a , P b , ℓ) → a place . The generative model f gen has two sequential parts, a point cloud feature encoder (Figure 2a ) and a conditional generator (Figure 2b ). Then, we calculate the transformation a place from the generated points (Figure 2c ). Finally, we prove a condition for when the full method is bi-equivariant.Encoding Point Feature. Given P a = {p i a } n i=1 and P b = {p j b } m j=1 , we first compute a feature at each point using two point cloud encoders ϕ a and ϕ b . The encoder ϕ a takes the XYZ coordinate and RGB color of all points of P a as input and outputs pointwise features {f i a } n i=1 . Similarly, ϕ b : P b → {f j b } m j=1 ,which shares an architecture but has separate parameters. Generating Points. The combined point cloud P ab is generated conditioned on the point features F a = {f i a } n i=1 and F b = {f j b } m j=1 using a modified version of Point Straight Flow [22] . This is a generative flow model where, at inference time, samples X 1 are taken by flowing over a vector field parameterized by a neural network v θ . Initial conditions are given by X where X Pa 0 ∈ R n×3 and X P b 0 ∈ R m×3 which are sampled from a scaled Gaussian. The network v θ defines the vector field for the ODE:0 = X Pa 0 ∪X P b 0 = {x k 0 } n+m k=1EQUATIONwhere X t is the intermediate point cloud states at time t and f ℓ is the encoded language feature of the language description ℓ from CLIP [41] . To solve the ODE, we iteratively update X t+∆t = X t + v θ (Z t )∆t for 1 ∆t steps. The model is trained by setting the optimal direction at any time t as P ab -X 0 which provides the objective,EQUATIONwhere X t = tP ab + (1 -t)X 0 . Intuitively, this sets dX t to be the drift force needed to move X t to P ab . Specifically, for a single point p k ∈ P a ∪ P b , we sample a noise x k 0 and a time step t to calculate the intermediate pointEQUATIONThe generator input is p k = (x k t , f k , f ℓ , t). We then optimize θ with respect to the loss function defined in Equation 2.P a ∪ P b = (p 1 a , p 2 a , • • • , p n a ) ∪ (p 1 b , p 2 b , • • • , p m b )and their corresponding target positionsP ab = (T a p 1 a , T a p 2 a , • • • , T a p n a ) ∪ (T b p 1 b , T b p 2 b , • • • , T b p m b ), we can recover the rigid transformations T a and T b using SVD [40] . However, the output Pa ∪ Pb of the generator f gen is not constrained to be given by rigid transforms of the original two point clouds. Each point may move independently by transformationT i α such that Pa ∪ Pb = (T 1 a p 1 a , T 2 a p 2 a , • • • , T n a p n a ) ∪ (T 1 b p 1 b , T 2 b p 2 b , • • • , T m b p m b ).We can still use SVD to estimate the best fitting Ta between (P a , Pa ) as well as Tb between (P b , Pb ). Assuming P b represents the object to be arranged and P a represents the placement, as shown in Figure 2 , the pre-place or place action can be calculated as a place = ( Ta ) -1 Tb .Realizing Bi-equivariance. As noted in prior work [8, 9, 10] , place actions that transform an object B with respect to another object A are bi-equivariant. That is, independent transformations of object B with g b ∈ SE(3) and object A with g a ∈ SE(3) result in a change (a ′ place = g a a place g -1 b ) 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 is constrained to be bi-equivariant, with invariant generation during training. Proposition 1. Assuming rotation-invariant Gaussian noise X 0 , if the encoded point feature F a and F b are invariant to rotations then f place is bi-equivariantf place (g a • P a , g b • P b ) = g a f place (P a , P b )g -1 b for all pairs of rotations (g a , g b ) ∈ SO(3) × SO(3). Proof. If X 0 = {x k 0 } n+m k=1 and F a ∪ F b = {f k } n+mk=i are rotation-invariant, the intermediate point states X t = tP ab + (1 -t)X 0 are rotation invariant with a fixed P ab . Since all inputs to v θ are invariant and the output always approaches P ab , we haveEQUATIONWith the same generated points P ab , the estimated transformation from rotated observationg a P a = (g a p 1 a , g a p 2 a , • • • , g a p n a ) to Pa is Ta g -1 a .Similarly, the estimated transformation fromg b P b to Pb is Tb g -1 b .Then, the new place action a ′ place can be calculated asa ′ place = ( Ta g -1 a ) -1 Tb g -1 b = g a T -1 a Tb g -1 b = g a a place g -1 b, which satisfies bi-equivariance .Our pick network f pick : (P a , P b ) → a pick has a similar design to f place . In this setting, P a are the points of the gripper and P b are the points of the object to pick. The function f pick differs from f place in that we only generate the new points for P b conditioned on P a . Figure 3 illustrates the generation process of grasping the banana by the crown.Since the pose and the shape of the gripper are always known, we fix P a in a canonical pose, sample only X P b 0 from a Gaussian distribution, and construct X 0 = P a ∪ X p b 0 . We set the target P ab as the union of the canonicalized gripper with the point cloud P b posed so it is held by the gripper, i.e., P ab = P a ∪ T b P b . We only use P b to calculate the loss: EQUATION 
experiments: 
 Model Architecture Details. The generative models f pick and f place share the same architecture. Each has two point cloud encoders and a generation network. We select PVCNN [42] as the backbone of our point encoders, which output a 64-dimension feature for each point. We use a pretrained CLIP-ViT32 [41] model as our language encoder and project the language embedding to a 32-dimension vector with a linear layer. The time step t is encoded as a 32-dimension positional embedding. We also encode a binary mask that indicates if the point belongs to P a or P b as a 32-dimension positional embedding. As a result, the generator input of a point is a 163-dimension vector. We adopt PSF [22] as our generator backbone. Both f pick and f place are trained end-to-end with the MSE loss defined in Equation 2 and Equation 5. We use the Adam optimizer with an initial learning rate of 10 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 as strong as that of PVCNN.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 et al. [7] to test our proposed method. Phone-on-Base: The agent must pick up the phone and plug it onto the phone base correctly. 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 sliding 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 high-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 are 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 description of the task and learns a language-conditioned 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 a 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] consumes segmented P a and P b 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 shoulder, 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 f pick , and the preplace and place action with f place simultaneously. We use the ground truth mask to segment P a and P b for RPDiff [4] and our method, as shown in Figure 4a .Training and Metrics. We train our method with 1, 5, or 10 demonstrations and train the baselines with 10 demonstrations. All methods are evaluated on 25 unseen configurations and each evaluation is averaged over 3 evaluation seeds. We report the mean success rate of each method in Table 1 . 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 1 : 1) IMAGINATION POLICY significantly outperforms all baselines trained with 10 demos on all the tasks except Put-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 a relatively high success rate while all the baselines fail to learn a good policy. 3) IMAGINATION POLICY achieves better sample efficiency and demonstrates few-shot learning performance. With one or five demonstrations, it sometimes outperforms the baselines trained with 10 demonstrations, e.g., IMAGINATION POLICY achieves a 97.2% 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 it exploit the symmetry inherent in the generation task. In the end, our method underperforms one baseline in Put-Plate. We hypothesize that the object in this task is symmetric and is hard to encode with distinguishable point features, which might result in wrong correspondences when estimating the rigid transformations. Since many complex manipulation tasks can be decomposed as a sequence of single pick and place, we illustrate that our method can address long-horizon tasks in Appendix 6.2.We validated IMAGINATION POLICY on a physical robot. We trained a multi-task agent from scratch on 3 tasks using a total of just 30 demonstrations. There was no use of the simulated data or pretraining in this experiment -all demonstrations were performed on the real robot.Settings. The experiment was performed on a UR5 robot with a Robotiq-85 end effector, as shown in Figure 6a . The workspace was a 48cm × 48cm region on a 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 P ab was constructed with segmented points and the poses. During testing, we used MoveIt as our path planner to execute the action sequentially. Tasks. We evaluate IMAGINATION POLICY on three pick and place tasks, as shown in Figure 6bcd .Mug-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 it 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 was 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 a place failure. 
conclusion: 
 In this work, we propose IMAGINATION POLICY, a 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 prediction 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 rotation-invariant point features. IMAGINATION POLICY demonstrates high sample efficiency and superior performance on six challenging RLbench tasks against several strong baselines. Finally, we demonstrate that the method can effectively be used to learn manipulation policies on a physical robot. We test our design choices using an ablation study on a multimodal pick-part dataset in Appendix 6.1.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 generation. Fortunately, a large number of 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 to future work. Moreover, this paper mainly focuses on rigid-object manipulation. We add one 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 future 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 a point cloud registration model to estimate the transformations. Table 3 . Ablation Results. We report the minimum, mean, and maximum error for single generation and pair generation over 100 runs with randomly rotated and sampled input. 
ablation study: 
 Multimodal Pick-part Dataset. To quantitatively measure the point cloud generation results and the equivariance of 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 use three cameras to 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 conditioned on the canonicalized gripper points and language descriptions. We also train a place generation model to generate both the gripper point cloud and the object point cloud. To evaluate the pick generation results, we randomly rotate and randomly downsample the object point cloud (P b ) to make a 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 P b results in a change of the ground truth pick pose. To evaluate the place generation results, we randomly rotate and downsample the gripper (P a ) as well as the object (P b ) to make a scene unseen during training and calculate the translation error and rotation error between estimated transformation T -1 a Tb and the ground truth pose. Note that rotating either P a or P b changes the relative ground truth transformation. 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 distribution and is equivariant. It realizes around 2• ∼ 3• average rotation error and 1cm translation error with different configurations of P a and P b ; (2) Without downsampling or color information, the rotation error slightly increases; (3) Without data augmentation in training, the performance decreases dramatically since the model cannot learn rotation-invariant features. (4) Compared with the results in the last two rows, the PVCNN-based point cloud encoder outperforms PointNet [54] and the pre-trained equivariant 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 is guaranteed to output the rotation invariant feature. We hypothesize that the architecture of Vector Neuron [28] and the standard representation limit its expressivity. Many challenging robotic manipulation problems can be viewed through the lens of a 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 following the RGB order. This is a high-precision task requiring the agent to correctly manage a sequence of pick and place. Even trained with 10 demos, our method can achieve 70.66% success rates in the first pick-place execution and maintain a 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. 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, we segment the two movable parts of the microwave, as shown in Figure 9e , the door with the handle (black color) and the frame (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 to the motion planning error and the collision between the door and the gripper when the gripper is closing. More complex articulated objects can also be manipulated by predicting the relative poses between links and we leave it as the future work. Overall, it demonstrates that IMAGINATION POLICY can address articulated object manipulation.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 as 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-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] is the baseline that consumed the same segmented point cloud (P a and P b ) as ours. It iteratively denoises the randomly sampled relative transformation poses conditioned on the current configuration of P a and P b . However, it can only solve the single-step place problem trained as a single-task policy. To make a fair comparison, we adapted it to a multi-task key-frame prediction model. Similar to our settings, we consider the pick problem as inferring the relative pose between the gripper (P a ) and the object to grasp (P b ). The preplace action prediction can also be viewed as calculating the relative pose between the object (P b ) and the placement (P b ). 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 : (P a , P b ) → T ab , it requires training 18 different models to solve the six tasks in Table 1 . We adapt it to learning a multi-task policy conditioned on the language embedding f : (P a , P b , f ℓ ) → T ab . Specifically, we used the same language embedding generated from CLIP [41] . To make RPDiff consume the language embedding, we map it to a 128-dimension feature via a linear layer and concatenate it to a 130-dimension time step embedding (the diffusion step). The model was trained and evaluated with the same settings in [4] . For execution on a robot, it requires a collision-free pick-and-place trajectory that connects the keyframe 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.We report the results of our method and baselines on RLbench tasks with ±1.98 std error in Table 6 .Model  
 
 Appendix: The pre-place action is the prerequisite to perform the place action. An example is shown in Figure4c. 


Decision: 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.

Reviews: review: title: title: Pointcloud Diffusion as a Pick-Place action generator
paper_summary: paper_summary: 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.
main_review: main_review: The work is very original. The common approach to solve pick and place problems is by learning a model that generates directly the desired rigid transform on the robot (for picking tasks) or on the object (for placing tasks). Instead, the authors propose a novel approach that instead generates the desired point cloud and then, estimate the rigid body by SVD. 

**Strenghts**
- The presented idea is very novel and different with respect to the common approaches. If the authors are able to validate experimentally their proposed approach, the method would be of high interest for the robotics community.

**Weaknesses**
- The selected experimetal evaluation is not properly selected. The authors compare their method against PerAct, RVT and 3D Diffuser Actor. Nevertheless, these methods are usually trained asuming raw RGBD data, while their method assume segmented full pointclouds of the interested objects. Thus, the comparison is unfair and unproper to validate the benefit of their method.
- Instead, the authors could compare the benefit of generating a point cloud and infering the rigid body against works that are more similar to theirs. For example, the authors could compare the placing against [1], that consider two pointclouds and instead applies diffusion in the SE(3) space. For the picking, authors could compare against [2], that consider also a target object pointcloud and diffuses the SE(3) pose of the gripper pose. This methods fit better as baselines as they consider the same information available for generating their actions and they only differ in the proposed method in the way the rigid body is estimated.

[1] Simeonov, Anthony, et al. "Shelving, Stacking, Hanging: Relational Pose Diffusion for Multi-modal Rearrangement." 7th Annual Conference on Robot Learning.

[2] Urain, Julen, et al. "Se (3)-diffusionfields: Learning smooth cost functions for joint grasp and motion optimization through diffusion." 2023 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2023.
questions: questions: - 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).
limitations: limitations: 1
review_summary: review_summary: 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.
title: title: Official Review of Paper 456 by Reviewer U3RP
paper_summary: paper_summary: 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.
main_review: main_review: Strengths:

* The idea to use flow models for predicting target point clouds is interesting and novel.
* The accompanying figures and videos are useful in illustrating the ideas introduced in the paper.

Issues:

* The flow of the introduction section feels unnatural. The first two paragraphs motivate the need for local geometric information in robot manipulation policies and introduce Imagination Policy method, but the third paragraph jumps abruptly to discussing about learning from multimodal robot data and dives into related works in that aspect. I suggest editing the intro to improve flow and clarity.
* The authors propose a method that works only with tasks that can be formulated as a three-step-process: pick, pre-place, place. Can the method solve long-horizon tasks that require multiple steps to complete? Can the method deal with manipulation of articulated objects? If the proposed method can solve these tasks, it would be great if the authors can add simulated and real robot experiments that involve these kinds of tasks. If not, the authors should explicitly acknowledge these limitations both in the introduction and the dedicated limitation section. I would be careful with the wording in "state-of-the-art performance across various tasks... compared with several strong baselines" (L10).
* Unlike the baselines that the authors compared with in the paper (RVT, PerAct), the proposed method assumes knowledge of which objects in the scene are ones to be picked up and placed on. This makes the proposed method less general-purpose than the competing baselines.
* The problem formulation in the paper resembles [1,2] more than the baselines mentioned in Table 1 of the paper. Could the authors add a comparison with one of these works in Table 1?
* The authors should move their real robot experiments from the appendix to the main paper. In addition to showing qualitative experiments, it would be great it the authors can add quantitative comparison to at least one baseline and showcase the generalizability of the proposed method when input object appearances and poses change.

Minor corrections:

* Consider citing [2,3] related to equivariant robot manipulation.
* I found that the large amount of symbols in the method section make it difficult to read through and understand. The authors should consider improving or adding figures in the paper and connect the symbols in the text with the figures.

Citations:

1. Simeonov et al. SE(3)-equivariant Object Representations for Manipulation. ICRA 2022.
2. Simeonov et al. SE(3)-Equivariant Relational Rearrangement with Neural Descriptor Fields. CoRL 2022.
3. Yang et al. EquivAct: SIM(3)-Equivariant Visuomotor Policies beyond Rigid Object Manipulation. ICRA 2023.
questions: questions: See above.
limitations: limitations: 2
review_summary: review_summary: 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.
title: title: Good pipeline but handles simpler problem
paper_summary: paper_summary: This work proposes to use diffusion to generate target point cloud of two objects, and predicts tranformation of objects for manipulation.
main_review: main_review: This paper proposes to predict the target point cloud from the initial point cloud and language instruction. The method validates broad manipulation tasks and achieves good results.

Although the task benchmark is extensive, the subproblem this work solving only requires rigid transformation. It is possible to use a 6D grasp pipeline to predict the transformation, for example, by sampling 6D poses, and evaluating its quality according to the target point cloud [1].

This work only considers segmented point clouds and does not consider collision-free valid targets.

This pipeline has the potential to generalize to more challenging tasks, for example, dealing with collision [2] or soft body [3].

I think this work requires additional efforts to demonstrate its ability.

[1] A billion ways to grasp: An evaluation of grasp sampling schemes on a dense, physics-based grasp data set

[2] SE(3)-DiffusionFields: Learning smooth cost functions for joint grasp and motion optimization through diffusion

[3] Any-point Trajectory Modeling for Policy Learning
questions: questions: 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?

2. Is it possible to extend this work to consider collision or generate multiple targets?
limitations: limitations: 3
review_summary: review_summary: 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.