Geometry-aware RL for Manipulation of Varying Shapes and Deformable Objects
Authors: Tai Hoang
Huy Le
Philipp Becker
Vien Anh Ngo
Gerhard Neumann
Venue: ICLR 2025 Oral (2025)
Keywords: Robotic Manipulation
Equivariance
Graph Neural Networks
Reinforcement Learning
Deformable Objects
Primary area: reinforcement learning
TL;DR: Geometry-aware RL with heterogeneous SE(3) equivariant back-bone policy for robotic manipulation

Abstract: Manipulating objects with varying geometries and deformable objects is a major challenge in robotics. Tasks such as insertion with different objects or cloth hanging require precise control and effective modelling of complex dynamics. In this work, we frame this problem through the lens of a heterogeneous graph that comprises smaller sub-graphs, such as actuators and objects, accompanied by different edge types describing their interactions. This graph representation serves as a unified structure for both rigid and deformable objects tasks, and can be extended further to tasks comprising multiple actuators. To evaluate this setup, we present a novel and challenging reinforcement learning benchmark, including rigid insertion of diverse objects, as well as rope and cloth manipulation with multiple end-effectors. These tasks present a large search space, as both the initial and target configurations are uniformly sampled in 3D space. To address this issue, we propose a novel graph-based policy model, dubbed Heterogeneous Equivariant Policy (HEPi), utilizing $SE(3)$ equivariant message passing networks as the main backbone to exploit the geometric symmetry. In addition, by modeling explicit heterogeneity, HEPi can outperform Transformer-based and non-heterogeneous equivariant policies in terms of average returns, sample efficiency, and generalization to unseen objects. Our project page is available at https://thobotics.github.io/hepi.

Published as a conference paper at ICLR 2025
GEOMETRY-AWARE RL FOR MANIPULATION OF VARY- ING SHAPES AND DEFORMABLE OBJECTS
Tai Hoang" Huy Le Philipp Becker', Ngo Anh Vien', Gerhard Neumann'
Autonomous Learning Robots, Karlsruhe Institute of Technology 2 Bosch Center for Artificial Intelligence
ABSTRACT
Manipulating objects with varying geometries and deformable objects is a major challenge in robotics. Tasks such insertion with different objects cloth hang- ing require precise control and effective modelling of complex dynamics. In this work, we frame this problem through the lens of heterogeneous graph that com- prises smaller sub-graphs, such as actuators and objects, accompanied by different edge types describing their interactions. This graph representation serves as a uni- fied structure for both rigid and deformable objects tasks, and can be extended further to tasks comprising multiple actuators. To evaluate this setup, we present a novel and challenging reinforcement learning benchmark, including rigid in- sertion of diverse objects, as well as rope and cloth manipulation with multiple end-effectors. These tasks present large search space, as both the initial and tar- get configurations are uniformly sampled in 3D space. To address this issue, we propose novel graph-based policy model, dubbed Heterogeneous Equivariant Policy (HEPi), utilizing (E)(3 equivariant message passing networks as the main backbone exploit the geometric symmetry. In addition, by modeling explicit heterogeneity, HEPi can outperform Transformer- based and non-heterogeneous equivariant policies in terms of average returns, sample efficiency, and generaliza- tion to unseen objects. Our project page is available here.
INTRODUCTION
Geometric structure plays crucial role in robotic manipulation. For instance, in insertion tasks, robot must precisely align objects with their corresponding target placements. Understanding the geometries therefore essential in such tasks, as each pair requires unique alignment (Zeng al., 2020; Tang et al., 2024). Similarly, for deformable objects like cloths, whose shapes change over time, successfully completing the task requires policy that can capture these dynamic geometric changes (Lin et al., 2021: Antonova et al., 2021; Shi et al., 2024). In both scenarios, these geometric structures can be naturally represented as graphs, a widely adopted framework (Wang et al., 2018; Huang et al., 2020; Ryu et al., 2023; Shi et al., 2024). In this paper, we frame manipulation problems as heterogeneous graphs. Taking the Mloth-Hagging task as an example, depicted in Figure 1, the cloth and the actuators by a set of directed inter-edges. Each node is associated with a geometric vector representing its 3D coordinates. Yet, this representation results in high-dimensional observation and action spaces, which makes learning policies that generalize seamlessly to novel orientations, poses, and unseen geometries challenging.
robot learning
represented as two distinct node sets, connected
To address this issue, recent works (Zeng et al., 2020; Huang et al., 2022: 2024; Ryu et al., 2023) in- troduced equivariance in the S3) space as an inductive bias. Using Equivariant Message Passing Networks (EMPNss) they learn policies that generalize to different poses by leveraging the geometric structure of the scene. These works learn by imitation, and most only consider simple pick-and- place tasks, where the model only has to produce a desired end-effector pose to be reached by a controller. An exception is the recent Equibot (Yang et al., 2024), where the policy outputs velocity vectors rather than static end-effector poses, enabling success in more complex and dynamic tasks like cloth folding and wrapping by imitation. This work investigates how to transfer these ideas from imitation
"Corresonndence ww..hhnnn@kitddee

---

Published as a conference paper at ICLR 2025
Heterogeneous Equivariant Policy (HEPi)
EMPN
Xn
Σ
EMPN
EMPN
xm
x1 Velocity vector
Actuator @:
Objest @:
Target
Intra-edge - Inter-edge
Figure 1: Left: Cloth-Hanging task represented by a heterogeneous graph that comprises two disjoint node sets, objects, and actuators, connected through directed, fully-connected inter-edges. Intra-edges occur within each set (both objects and actuators) to capture relationships within clus- ters. Information aggregated from objects to actuators via inter-edges. The target distance is ab- sorbed into the feature representation rather than treated as separate node type. Right: Overview of Heterogeneous Equivariant Policy (HEPi), consisting of multiple Equivariant Message Passing Networks (EMPNs) process the graph, and the outputs are aggregated to generate the final action.
to reinforcement learning. Unlike supervised imitation, training policies with reinforcement learn- ing presents additional challenges, particularly due the need for high-frequency data collection and efficient adaptation to new experiences. Large policy networks struggle these settings, as they unable to quickly adapt to changing data (Andrychowicz et al., 2021). To address these is- sues, we design a lightweight heterogenous equivariant architecture, amenable to efficient on on-policy reinforcement learning. The architecture's equivariance allows generalizing between poses and its heterogeneity enables us to include and exploit knowledge about the scene as well as the unactuated and actuated objects in it. For training, we find that naively using Proximal Policy Optimization (PPO) (Schulman et al., 2017), can result in suboptimal performance, and we propose to employ more principled trust region approach from Otto al. (2021) to achieve stable convergence.
To evaluate our approach and future advancements in this direction, we propose novel suite of seven tasks, realized using NIVIDA IsaacLab (Mittal al., 2023) to utilize its GPU-based simula- tion engine. They are designed to highlight the role of geometric structure in manipulation tasks, with a progressive increase in difficulty, from simple rigid-body manipulation with diverse objects to more challenging tasks involving multiple actuators and deformable objects. Our experimen- tal results demonstrate that the proposed Heterogeneous Equivariant Policy (HEPi) outperforms both Transformer-basee and pure EMPN baselines, particularly in complex 3D manipulation tasks. HEPi's s integration of equivariance and explicit heterogeneity modelling improves performance in terms of average returns, sample efficiency, and generalization to unseen objects.
To summarize, our contributions are i) novel benchmark comprising rigid insertion of varying ge- ometries and deformable objects manipulation that is particularly well-suited for geometry aware reinforcement learning research; ii) HEPi, a graph-based policy that expressive and computa- tionally efficient while being constrained to be (euu)vvrvvarr, perfectly suitable for solving complex 3D manipulation tasks under reinforcement learning settings; iii) a theoretical justification and extensive empirical analysis for design choices.
BACKGROUND
Message Passing Neural Networks (MPNN) Consider a graph 9 = (V,E), where represents the nodes and E the edges. In standard Graph Neural Network (GNN) (Battaglia et al., 2018), each node E upddates its feature representation by aggregating information from its neighbors N(v). This process is formalized
1 (k+1) =
(())
+
1
(F())
p(R)
euv)
(1)
uVN(v)
2

---

Published as a conference paper at ICLR 2025
where f. is the feature vector of node at iteration eur is the edge feature between nodes a and v, & and vp are often deep neural networks, and + represents an aggregation function like summation, mean max.
SE(3) Equivariance and Invariance In this work, we focus on geometric graphs g = (V. & x),
where each node is associated with coordinates Dv E X R and steerable geometric features f, A feature is considered steerable if it transforms consistently under the action of group 9 E G through group representation p. For instance, a vector v R under rotation R E SO(3) transforms as v' = Rr. In this case, is finiteddimensional vector, and its transformation is described by an invertible matrix representation with determinant 1, p(g) R.
Equivariance and invariance are two fundamental concepts in this context, can be formalized as follows, using the notation from Brandstetter et al. (2022):
Definition 2.1. Let G be a group with representations px and p² A function X y is equivariant if
p (g)[[())] - f(00 [(0x (g)[2)), EgEG
and invariant if
((=)=
G
In simpler terms, equivariance guarantees that applying a transformation g to the input space and then applying the function f produces the same result as applying first and then transforming the output space V. On the other hand, invariance implies that the function f remains unchanged when the input undergoes transformation in X.
Symmetries in MDPs A Markov Decision Process (MDP) is defined by the tuple (S, where the set of states, the actions, P(s'|s |,, a) the transition probability, R(s, a) the reward function, and [0, 1] the discount factor. The goal is to find policy и S - A that maximizes the expected discounted reward E, D R(s,,
In MDPs with symmetries, both the transition distribution P(s'|s, a) and policy distribution (a]s) are invariant under group transformations E left-regular representation Ly and Kg for state and action, respectively, resulting in the following conditions:
[(].)(''|[_g()), Kz(a)) a)) T. m(ass),
and similarly for the reward function: R(Lg(s), ([Lg(s),Kjaa)) Kyaa))= R(s, a). This allows leveraging symme- tries to reduce the complexity learning. potentially improving sample efficiency and generaliza- tion, it results in group-structured MDP homomorphism (Van der Pol et al., 2020).
3 METHODOLOGGY
Problem Statement We aim to solve robotic manipulation problems using an on-policy actor- critic reinforcement learning approach. To address the symmetries present the state and action spaces of the Markov Decision Process (MDP), we leverage equivariant policies, ensuring that transformations applied the state space are consistently reflected in the action space. To handle the complexities of robotic manipulation, where actuators and objects play distinct roles, we propose the Heterogeneous Equivariant Policy (HEPi), which comprises three key components:
Equivariant MPN backbone: An efficient and expressive EMPN capable of exploiting en- vironment symmetries, thereby significantly reducing the search space complexity.
Heterogeneous graph design and update rules: A graph structure with distinct actuator and object nodes, with tailored message- message-passing rules to handle the system's 's heterogeneity.
Employing a principled trust-region method to stabilize training in complex, high- dimensional environments.
3.1 EQUIVARIANT MPN BACKBONE
An Equivariant Message Passing Network (EMPN) can be constructed (Brandstetter et al., 2022; Duval et al., 2023; Bekkers al., 2024) by enforcing equivariance in the functions and in
3

---

Published as a conference paper at ICLR 2025
Equation 1, ensuring that their inputs and outputs steerable and transform consistently under the group G, as described in Definition 2. 1. Constructing such functions for high-dimensional steerable features is challenging and typically requires spherical harmonics embeddings (Duval et al., 2023), where matrix-vector multiplications are carried out using Clebsch-Gordar tensor products, followed by steerable activation functions (Brandstetter al., 2022; Bekkers al., 2024). While these op- erations guarantee equivariance, they also introduce high computational complexity, making them impractical for reinforcement learning settings. To mitigate this issue, Bekkers al. (2024) intro- duced the PONITA framework, an efficient equivariant message- message-passing approach. We use it as our EMPN backbone and refer to it as EMPN throughout the rest of the paper for consistency.
Cav) = k( x., ) f is defined as linear function with steerable feature f., and summation used the aggregation function, + E. Here, is a convolution kernel, and (z,, x.) represents the relative position between nodes and v. This leads to the convolutional messagg-passinngg update rule, f. EuEN(w) k(2,, x.)
Consider the message function in Equation 1, where v/(f, (f) f(k))
On a regular grid, e.g., an image, each relative position (æ - x.,) has a corresponding weight W.., and is stored in one single matrix W. However, this does not apply to non-uniform grids, e.g. point clouds, k( Ze) in this case can be parameterized by neural network, resulting in the formulation fy = Jxx ko (E F. (ud Under this general formulation, Bekkers al. (2024) showed that can made (3) equivariant by "lifting" the input domain from TR to X R x S². Specifically, for every position E IR.3 an associated orientation S is introduced. This allows features in EMPN to be embedded in both spatial and orientation spaces, leading to the following convolutional form:
f. - == ko (o([, o.), (px, of)))fud
Furthermore, improve computational efficiency, the kernel function is factorized
Ki (op(( ow), (pv: o.)]) K ki (o, o.)) kg (o) (p. - -P.), |lo, 1 (p.. - -p)ll)
where ((1) handles spatial interactions based the relative position (p.- p.) and the perpendicular component ||o, 1 (P.. P.)lI. (2(2) manages orientation-based interactions via dot products o. Ou- and (3) performs channel-wise mixing across features. This formulation preserves the universal approximation property of equivariant functions while being significantly more computationally efficient and not requiring specialized network structures. Furthermore, can be sampled on uniform grid over S², making EMPN only approximately equivariant. However, in practice, it achieves strong performance even with a limited number of samples, as discussed Appendix D.
3.2 HETEROGENEOUS EQUIVARIANT POLICY
In robotic manipulation tasks, actuators and objects play fundamentally distinct roles. The graph defined as G = (V,&), where V = Vact Vanj represents disjoint node sets for actuators and objects. Our approach captures these roles by first processing local information within the object and actuator clusters and then aggregating it globally to the actuators via directed, fully-connected inter-edges, shown in Figure 1. This design distinguishes object-to-object, actuator-to tottor--t-cctuttrr. and peject-to-actuatee interactions, allowing the system to separate local processing from global information exchange.
The updates for both object and actuator nodes can be expressed as
pobj, new
Dobj
pobj
k(aaaii ..hhi Ocbj--ots) rahi
Vobji
(e))
eact, new
Puut local
k-xeett m0e andtaa pact
1 E Vuct.
WE (w))
pact, final
pact, new
+ Oact-global
pact
ko x Jobj-act ff new
V E Vast-
4

---

Published as a conference paper at ICLR 2025
Here, gobj, new represents the updated object features after local object-to-object interactions, f.rt. refers to the updated actuator features after local actuator-to-actualee to-actutoo interactions, and fou. final is the final feature for actuator nodes after aggregating information from both the objects and its actua- tor neighbors. Here each of the kemels k(,,"' otjj. ). k(,;'' '5 (at and k(., "; act has it own learnable parameters, allowing them specialize the learning process for each interaction type.
Moreover, each node v E V encodes its node_typp type as a one-hot scalar-vector, along with normal- ized position vectors Pv and velocities V.. For object nodes, the feature vector also includes the relative distance to the target, du target, embedding target information directly without the need for an additional target node. For actuator nodes, the output consists of both scalar e and vector Vout- where the final output vector is computed as Vour = c. v. This setup ensures flexibility for diverse tasks while maintaining consistency with the geometrie properties of the system.
Value Function We employ DeepSets (Zaheer et al., 2017) with the same input structure as the policy to preserve permutation invariance of the node features, while keeping the architecture both simple and computationally efficient, similar to the prior work Simm et al. (2021). The value func- tion is computed as ((s) = MLPPuter (Zwev MLPiiner(Ss)) where 8y represents the feature of node These node features may differ from those used in the policy network to capture task-specific ob- servations. For example, in the Cloth- Hoag-nangg task, the value function considers features from all nodes, while the policy network focuses only the hole boundary nodes. Full details of the input features used for the value function are provided in the Appendix B.
TrgstnRegion Projection Layers Standard on-policy reinforcement learning approaches such as Proximal Policy Optimization (PPO) (Schulman et al., 2017), learn a policy by optimizing the sur- rogate objective
0 arg max E (s.a)
A T (s, (s) Tez (als)
To(als)
s.t. DKL (To (To.) S 5.
where A (s, a)) is the advantage function. Here. DKL (Tell is the KL-divergence between the new policy Te and the old policy Ta., constrained by 5 to ensure stable updates and prevent overly large policy changes. PPO approximates this trust region by clipping the importance sampling ratio to limit updates. This, however, requires careful hyperparameters turning to make it work stably, as pointed out by Andrychowicz al. (2021). We will show later in the Results Section 4.2, this is also applied to graph-based policy. On the other hand, Trust Region Projection Layers (TRPL) (Otto al., 2021) adopt a more principled approach. TRPL projects policy parameters onto trust region boundaries using a differentiable convex optimization, ensuring stability by projecting both the mean and variance of the Gaussian policy to satisfy trust region constraints.
In this paper, we adopt TRPL to ensure stable policy updates, and we will show in the Result Sec- tion 4.2, TRPL consistently outperforms PPO, with little hyperparameter tuning required, especially in tasks requiring complex exploration, as also observed in the prior works (Otto et al., 2023; Li et al., 2023; Celik et al., 2024).
3.3 THEORETICAL JUSTIHCCAION
HEPi is inspired by adding global Virtual Nodes (VNg) to Message Passing Neural Networks, (MPNNs). For example, Southern et al. (2024) proposed MPNN + VNG which seperates local and global updates. Here, the local updates are equivalent to object node updates, and the global ones correspond to our actuator node update. Based on this interpretation, we show that locally connecting actuator nodes to only k-nearest object nodes not capture relevant relation between object and actuator nodes, while treating the actuator nodes as VN that connects to all object nodes can. We name the graph network with locally connected actuators and object nodes as MPNN VNLocal- We show that for HEPi any two actuator and object nodes can exchange information, while this is not the case for MPNN VNLocll-
Proposition 3.1. For MPNN + VNiocd, the Jacobian a is independent of u. whenever object node u and actuator node are separated by more than 2 hops. In contrast, HEPi with node connections and updates as described in Section 3.2 can exchange information between any actuator and object nodes after single layer.
The proof is provided in Appendix A. This result implies that HEPi's connection design allows the actuators to receive relevant information to predict actions w.r.t changes at object nodes. In contrast,
5

---

Published as a conference paper at ICLR 2025
Rigid Insertion
Rigid Insertion with Agents
0=
t-=10
Rope Shaping
Cloth Hanging
1E0
Q
t=0
Figure 2: Illustration our diverse and challenging manipulation tasks, involving both rigid and deformable objects. These tasks require precise control under complex geometric constraints, co- ordination between multiple actuators, and handling intricate interactions between objects and actuators. The variety of tasks highlights the need for policies that can understand the geometric structure in large observation and action spaces.
for MPNN VNrocal the actuators could fail to predict relevant actions changes object node u. We provide experimental ablations across various values of & to clearly emphasize this distinction Appendix D.
EXPERIMENTS
In this section, outline the experimental setup and present the results comparing the proposed НEPi against other baselines.
4.1 EXPERIMENTAL SETUP
Task Design Our task design, illustrated in Figure 2, emphasizes testing the role of geometric structure and information exchange between objects and actuators in robotic manipulation. To focus on this, we abstract away the specifics the robot body and consider only end-effector control. We introduce two categories of tasks: rigid manipulation diverse geometries and deformable object manipulation, all implemented in NVIDIA IsaacLab (Mittal et al., 2023) to leverage its GPU-based parallelization capabilities'.
The rigid manipulation tasks are inspired by Transporter Net (Zeng et al., 2020). Rigid-Sliding mimics using suction gripper slide an object across 2D plane to a target position and orienta- tion, with 10 distinct objects, and randomized initial and target poses. Next, Rigid- removes the physical connection between the actuator and the object, allowing the actuator in the z-y plane to push the object to a desired target position and orientation. Rigid- Insertinn similar the assembly kit task in Transporter Net, extends this to 3D, requiring precise alignment and insertion of objects into holes, using 8 different objects. Additionally, we introduce novel Rigid-Insertion. Tworionn------ntts task, where two linear actuators work together to control an object, guiding it to target randomly positioned in the upper hemisphere of the S
freely
For deformable object manipulation, we first adopt the Rope-Closing task from Laezza al. (2021), where two actuators manipulate a deformable rope to wrap around a cylindrical object in 2D plane, with randomized initial configurations. We then introduce a novel task, Rope-Shaping, which increases complexity by requiring the rope to form specific shape (a "W" from the LASA dataset (Khansari-Zades Billard, 201 I)) to a desired orientation. Finally, we introduce Cloth-Hanging where four actuators control the corners of cloth to hang it onto hanger, with randomized starting positions and orientations 3D space.
These tasks present range of manipulation challenges, emphasizing the role of geometric structure and requiring complex exploration strategies to coordinate the agents in completing the tasks. Full task details, including reward definitions, provided in Appendix B.
video showcasing the tasks be found in the supplementary material.

---

Published as a conference paper at ICLR 2025
HEPi (Ours
EMPN
Transformer
rigid-sliding-2D
rigid-insertion-2DD
rigid-insertion- two-agents-3D
rigid- Soliddpushnng-2DD
-1.00
-1550
1.00
IGM Raturn
1.25
2.00
-150
1.50
2.50
-1.75
-1775
-2.00
-2.00
-2.00
3.00
8
12
16
20
12
20
12
18
30
rope-closing-2D
rope-shaping-2D
cloth-hanging-30
22000
-2.00
-3.00
-3.00
-1660
-4.00
10
Env Steps (1M)
Env Steps (1M)
Env Steps (1M)
Figure 3: Evaluation curves for our seven manipulation tasks, comparing HEPi (ours), EMPN, and Transformer baselines. Results are averaged over 10 seeds, using IQM with 95% confidence intervals. HEPi consistently outperforms EMPN and Transformer in tasks requiring complex ex- ploration and heterogeneity handling, such as rigidd-meertionnnoogggnn---- rigid- Sggid-pushing-21 and cloth-hanging-3D.
Baselines We compare HEPi against two primary baselines: policies based Transformer (Vaswani et al., 2017) and a naive EMPN. Transformers serve as a strong baseline to evaluate in our setting it can be seen as fully-connected GNN (Battaglia et al., 2018), and have achieved state- of-the-art performance in other graph-base based reinforcement learning problems (Kurin et al., 2021; Trabucco al., 2022; Gupta al., 2022; Hong al., 2022). In addition, for the Cloth-Hanginng task, we evaluate two additional baselines, Heterogeneous GNN (HeteroGNN) and a naive GNN to highlight the effectiveness of incorporating equivariant constraints in large 3D space.
Our experimental setup aims to answer the following key questions: (1) Can explicitly modeling het- erogeneity between actuators and objects, combined with SE))) equivariance, improve performance in complex 3D tasks? (2) How well does HEPi generalize when dealing with different geometries, resolutions in rigid tasks, and varying sample spaces the complex 3D cloth-hangee task? (3) Attention mechanisms are often employed in GNNs to capture heterogeneity, do they offer the same benefits as explicitly modeling heterogeneity in HEPi? (4) Does using trust-region methods HEPi stabilize the training process more effectively than naive PPO?
4.2 RESULTS AND DISCUSSIONS
In the main evaluations, generate 1000 scenes per task (sampled according to Appendix B) and compute the undiscounted return over 10 seeds, and report the average using Interquartile Mean (IQM) (Agarwal et al., 2021) with 95% confident interval. Figure shows the evaluation curves for the seven manipulation tasks. Overall, both EMPN and HEPi outperform the Transformer in terms of sample complexity, owing their ability to exploit symmetry.
For final performance on rigid tasks, firstly in rigid-slding-2 2D and iiiddiinsertion-2D4: tasks, HEPi and Transformer policies perform comparably, suggesting that the limited task complexity does fully leverage the benefits of equivariant constraints. However, when the search space grows larger, as in the case of rigid-ing rrriddinsertionn- waeoonttoo-geeee two-agents-3D Transformer struggles to find a policy that general- izes to all poses. EMPN, on the other hand, gets stuck in local optima due to its lack of expressive- ness, especially in tasks requiring more exploration, such as Sggid-uushig-222 and iegiiddissetiinn-ae wwweerttoootwwooggereee amo-agents-33 D. In contrast, HEPi's explicit handling of heterogeneity allows for more effective exploration, leading to better overall performance.
rigid-insenn iigid-insertion- 2D+z
In rope-closing and rope-shaping, simple deformable object tasks, the Transformer exhibits poor generalization, likely due to the complexity introduced by non-rigid constraints and random orienta- tions. HEPi and EMPN perform similarly on the 2D tasks, but as tasks scale up 3D environments, such as cloth Sotth-hanging-333 3D. HEPi shows a significant advantage, outperforming both baselines. This highlights the importance of explicitly capturing heterogeneity and task geometry in manipulation.

---

Published as a conference paper at ICLR 2025
HEPS (Ours)
EMPN
Transformer
HeterrGNN
GNN
-1.00
-1.00
-1.00
IQM Return
1.25
11.25
-1.25
11.00
1150
-1.50
-1.75
11.75
-1.75
2
3
Env Steps (1M)
Env Steps (1M)
Env Steps (1M)
(a) Original 3D Space
(b) 2D Half Circle
(c) 2D Quarter Circle
Figure 4: Performance of different models on the Cloth Haaging task across varying sample spaces. Overall, performance improves as the sample space decreases. In terms of final performance, het- erogeneous models outperform homogeneous baselines most cases, demonstrating the benefits of explicit heterogeneity modeling. Additionally, applying equívariant constraints is critical for achiev- ing superior performance in 3D tasks. More results can be found in Appendix D.
Low-resolution objects (- 20 nodes)
High-resolution objects (- 1200 nodes)
0.15
HEPI (Ours)
Transformer
-2.35
rigid-slid giddgS
1100
Postiio Noise
0.1
-1.89
IOM Return
-1.0
-1.25-
0.05
-1.5
1.50-
1.75
0,0
-2.00
2
0.0
0.05
0.1
0.15
0.05
0.1
0.15
Velocity Nolse (std)
Velocity Nolse (std)
Figure 5: Left: Analysis of noise sensitivity and scalability high-resolution objects in the Rigid- Pushing task. Heatmaps show average returns under varying levels of artificial Gaussian noise in position and velocity inputs for both low-resolution and high-resolution objects. A single HEPi agent, trained on low-resolution with additive Gaussian Noise (a 0.01), was used for all evalu- ations. Right: Generalization performance on Rigid-Sliding and Rigid-Insertion tasks. Models are trained on one object (plus), two objects (plus, star), and three objects (plus, star, pentagon) and tested on the remaining unseen objects. Overall, HEPi generalizes well to unseen objects, performs consistently across resolutions, and handles noise effectively, making it suitable for real-world tasks.
Generalizability Figure 4 compares the performance of different models the Cloth- Hooth-Hnggig task across varying sample spaces. As expected, smaller sample spaces simplify the exploration and improve performance. Heterogeneous models (HeteroGNN and HEPi) consistently achieve higher final returns than their homogencous counterparts, demonstrating greater expressiveness. However, HeteroGNN requires more samples, whereas HEPi's use of EMPN significantly improves sample efficiency by leveraging equivariant constraints in large 3D spaces.
Next, evaluate the robustness of HEPi noisy inputs and its ability handle high-resolution object meshes on the Rigid-Pushing task. GNNs naturally capture locality through message pass- ing, allowing them scale effectively to higher-resolution graphs without retraining (Li al., 2020; Freymuth et al., 2023). During training, we added Gaussian noise (a 0.01) to normalized posi- tions and velocities to encourage diverse node representations, a common GNN regularization tech- nique (Godwin et al., 2022). The best HEPi agent was then evaluated with varying Gaussian noise levels applied to pre-normalized inputs (environment noise) and across low-resolution (~ nodes) and high-resolution (~ 1200 nodes) object meshes. As shown in Figure (left), HEPi maintains high performance across resolutions with only mild degradation at higher noise levels, demonstrating its scalability and robustness noisy and diverse object representations.
Finally, we evaluate the generalization these models to unseen objects on two rigid tasks: rigid- sliding and rigid-insertion. Both tasks trained subsets of objects-oon one (plus), (plus, star),
8

---

Published as a conference paper at ICLR 2025
HEPi (Ours)
HEP Attention
EMPN Attention
EMPN
rigid-sliding-2D
niiidiinsertion224z
rigid-insertion two-agents-3D
rigid-pushing-2D
-1.00
-1.00
L0M Return
-1 1.000-
-1.60
-1.20
1
-1.80-
-1.40
1.40
rope-closing-2D
rope-shaping-22
clotn-hanging-38
Computational overhead
-1.00
I0M Return
11500
-2.00
2.00
-1.75
-2.00
-2.00
-2.40
-1.40
Figure 7: Performance comparison on tasks with and without attention mechanisms over 10 seeds. Adding attention significantly increases the training time but does not improve performance, as shown on the right. The computational overhead is measured as the ratio of training time per iteration (over seven tasks) relative to HEPi
and three (plus, star, pentagon) and tested on the remaining objects. Figure (right) shows that HEPi generalizes better than the Transformer baseline, benefiting from the ability of graph-based models to exploit object topology. In contrast, Transformers lack structure-awar embeddings and struggle with graph-structured inputs, as noted in prior work (Hong et al., 2022).
Attention Attention mechanisms are widely used in graph neural networks to capture heterogene- ity. this experiment, we examine the impact of adding attention as an aggregation function in Equation 1 to both homogeneous and heterogeneous graph networks, framing the popular Graph Attention Network (GAT) framework (Veličković et al., 2018). However, as shown in Figure 7, attention does not provide any noticeable benefit across the tasks.
and it ultimately complicates the learning process in on-policy reinforce- ment learning, making the optimization landscape more difficult to traverse. Additionally, adding attention significantly increases training time without improving performance, e.g., for HEPi it al- most doubled.
While attention helps capture some heterogeneity, particularly tasks like wtt-iiiistttitttoooeee
Training Stability To investigate the impact of the TRPL method, we compare it against PPO. For a fair comparison, we perform a grid search over the cl clippeps parameter in PPO (Appendix E.3). Overall, as depicted in Figure in tasks requiring high exploration such as cloth-hanging-333 PPO struggles to maintain conservative updates, often resulting in unstable performance. However, in tasks with lower-dimensional action space, such as 2D environments, well-tuned PPO performs comparably to TRPL in terms of final average return, though being less sample efficient. This suggests that while PPO can tuned to perform adequately in simpler action spaces, TRPL provides more stability and robustness, particularly in complex 3D environments that demand more effective exploration control.
RELATED WORK
Equivariant Policies for Robotic Manipulation In imitation learning for robotic manipulation, equivariance has been widely applied to reduce the effort of collecting human demonstrations (Zeng al., 2020; Huang et al., 2022: 2024; Ryu et al., 2023; Yang al., 2024). Most prior work fo- cuses on simple pick-and-place tasks, where the policy outputs the 3D pose of the end-effector. By leveraging equivariance, these policies can generalize across different object poses, significantly reducing the number of required demonstrations (e.g., only 5 to 10 demonstrations in (Ryu et al., 2023)). However, 3D pose actions are insufficient for tasks that require higher dexterity or involve deformable objects. To address this limitation, Equibot (Yang al., 2024) designed equivariant policy that outputs velocity vectors, achieving success in more complex tasks such as cloth folding and object wrapping.

---

Published as a conference paper at ICLR 2025
HERS with TRPL.
HEPi with PPO
Transformer with TRPL.
Transformer with PPO
rigid-insertion-2D4*
rope-shaping-2D
cloth-hanging-
-1.00
-1.00
-2.00
IOM Return
-1.25
-1.50
-1.50
-3.00
-2.00
-1.77
-2.00
-4.00
-2.50
12 16
2
10
5
Env Steps (1M)
Env Steps (1M)
Env Steps (1M)
Figure 8: Performance comparison between HEPi and Transformer models with TRPL and PPO over 10 seeds. TRPL shows stable performance across all tasks, while PPO struggles in tasks requir- ing high exploration, especially in 3D environments like cloth-hanging-32. In tasks with lower- dimensional action space (e.g., 2D tasks), both methods perform comparably when carefully tuned.
There has been limited work on exploiting equivariant policies in reinforcement learning, and ex- isting approaches have largely focused 2D spaces (Wang et al., 2022; Nguyen et al., 2023). In this work, we extend the study of equivariant policies to 3D space within a reinforcement learning setting, which, to the best of our knowledge, has not yet been explored for robotic manipulation.
RL with GNNs Graph-based representations in reinforcement learning have shown great success across diverse domains, including molecular design (Simm et al., 2021), adaptive mesh refinement (Freymuth et al., 2023), and multi-agent systems like traffic light control (van der Pol et al., 2022). Among these, the most closely related work ours morphology reinforcement learning (Wang et al., 2018; Huang et al., 2020; Trabucco et al., 2022; Hong et al., 2022; Gupta et al., 2022), where the robot's kinematic structure represented as a graph, allowing actuators to be controlled through message passing. This approach enables policies to generalize across different robot topologies, particularly in locomotion tasks (Gupta et al., 2022). Chen al. (2023) extend these ideas to handle 3D environments, but only with sub-equivariant policies for rotations around the gravity axis.
However, these works primarily exploit the locality of graph structures, framing the problem as multi-agent reinforcement learning graphs (Jiang et al., 2020), where each node can make de- cisions influenced by its neighbors. contrast, our work focuses on the underactuated problem where only small subset of nodes (actuators) controls much larger set of object nodes, framing a heterogeneous graph. The interactions between the nodes are therefore much more complex to capture using homogeneous GNN models.
CONCLUSION
We have demonstrated that robotic manipulation problems can be effectively represented as hetero- geneous graphs, comprising two sub-graphs to capture the geometric structure the environment. Building on this, introduced HEPi, a graph-based policy featuring multiple equivariant message- passing networks as its backbone. These networks are constrained to be equivariant under SF((3 transformations, which significantly improves sample efficiency. Furthermore, HEPi explicitly mod- els heterogeneity by assigning distinct network parameters for each interaction type, reducing mes- sage mixing and improving expressiveness. This approach has proven less prone converging on sub-optimal solutions. To assess the effectiveness of our approach, we developed a new reinforce- ment learning benchmark focused on manipulating objects with diverse geometries and deformable materials. Our results show that HEPi outperforms both the state-of-the-at Transformer and its non-heterogeneous, non-equivariant counterparts.
Limitation In our current setup, we abstract away the robot body, focusing solely on end-effector movements. Future work could explore incorporating a more structured representation of actuator nodes, potentially leveraging the robot's full morphology. Moreover, although approach does not require full object meshes, we assume that the keypoint coordinates are readily available as our main observation. This limitation could be addressed by integrating tate--o-ther of-tee-----hrtttt- computer vision techniques to extract keypoints from cameras (1 (Tumanyan et al., 2024; Hou et al., 2024), using these object nodes, thus increasing its applicability in real-world scenarios.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTI
This work is part of the M4 subproject of the DFG Al Research Unit 5339, which focuses on design- ing distributed policy optimization methodologies to accelerate manufacturing process maturation. We gratefully acknowledge financial support from the German Research Foundation (DFG) and computational resources provided by bwHPC and the HoreKa supercomputer, funded by the Min- istry of Science, Research, and the Arts Baden- Württemberg and the German Federal Ministry of Education and Research.
We thank Niklas Freymuth, Philipp Dahlinger, Onur Celik, Aleksandar Taranovic, and Mayank Mittal for their insightful discussions during both the initial and rebuttal phases of this paper, as well as for their technical support. We also appreciate Reviewer 6nhr for suggesting the Rigid-Pushing task.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in Neural Informa- tion Processing Systems, 34, 2021.
Marcin Andrychowicz, Anton Raichuk, Piotr Stañczyk, Manu Orsini, Sertan Girgin, Raphaël Marinier, Leonard Hussenot, Matthieu Geist, Olivier Pietquin, Marcin Michalski, Sylvain Gelly, and Olivier Bachem. What matters for on-policy deep actor-critic methods? large-scale study. In International Conference on Learning Representations, 2021.
Rika Antonova, Peiyang Shi, Hang Yin, Zehang Weng, and Danica Kragic Jensfelt. Dynamic envi- ronments with deformable objects. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021. URL https: openreview. jMCBA.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzale Vinícius Flores Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Çaglar Gülçehre, H. Francis Song, Andrew J. Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey R. Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matthew Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. CoRR, abs/1 http//806.012... 2018.
Erik Bekkers, Sharvaree Vadgama, Rob Hesselink, Putri A Van der Linden, and David W. Romero. Fast, expressive $ mathrm (SE) equivariant networks through weight-sharing in position- orientation space. In The Twelfth International Conference on Learning Representations, 2024.
Albert Bou, Matteo Bettini, Sehastian Dittert, Vikash Kumar, Shagun Sodhani, Xiaomeng Yang. Gianni De Fabritiis, and Vincent Moens. Torchrl: A data-driven decision-making library for pytorch, 2023.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling, Geomet- ric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations, 2022.
Chen Cai, Truong Son Hy, Rose Yu, and Yusu Wang. On the connection between MPNN and graph transformer. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Muchine Learning Research, pp. 3408-3430. PMLR, 23-29 Jul 2023.
Onur Celik, Aleksandar Taranovic, and Gerhard Neumann. Acquiring diverse skills using curricu- lum reinforcement learning with mixture experts. In Forty-first International Conference on Machine Learning. 2024. URL https: //ppnreveeeeeeeeeeeeeeeeeeeeeeee http:a/eeeeeeetttooooooee-------
Runfa Chen, Jiaqi Han, Fuchun Sun, and Wenbing Huang. Subequivariant graph reinforcement learning in 3d environment. In International Conference Machine Learning. PMLR,, 2023.
11

---

Published as a conference paper at ICLR 2025
Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Liò, and Michael Bronstein. On over-squashing in message passing neural networks: the impact of width, depth, and topology. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
Alexandre Duval, Simon V. Mathis, Chaitanya K. Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D. Malliaros, Taco Cohen, Pietro Lio, Yoshua Bengio, and Michael Bronstein. A hitchhiker's guide to geometric gnns for 3d atomic systems, 2023.
Matthias Fey and Jan Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop Representation Learning Graphs and Manifolds, 2019.
Niklas Freymuth, Philipp Dahlinger, Tobias Daniel Würth, Simon Reisch, Luise Kärger, and Ger- hard Neumann. Swarm reinforcement learning for adaptive mesh refinement. In Thirly-seventh Conference on Neural Information Processing Systems, 2023. URL net cest/Eorumiics ? r2 qe 4uc.
Jonathan Godwin, Michael Schaarschmidt, Alexander Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia. Simple GNN regularisa- tion for 3d molecular property prediction and beyond. In International Conference on Learning Representations, 2022. URL https: http///ppprreeeeeeeeee net
Agrim Gupta, Linxi Fan, Surya Ganguli, and Li Fei-Fei. Metamorph: Learning universal con- trollers with transformers. In International Conference Learning Representations, 2022. URL
https: /pwp/eeeeeeeeeeeeee/eeeeeoeeeeee .et/ forum? ittp:///weee.eee/ooooooooooo--e (oooummidd-pmmttt_Gr7
Sunghoo Hong, Deunsol Yoon, and Kee-Eung Kim. Structure-aware transformer policy for inho-
mogeneous multi-task reinforcement learning. In International Conference on Learning Repre- sentations, 2022.
Chengkai Hou, Zhengrong Xue, Bingyang Zhou, Jinghan Ke, Lin Shao, and Huazhe Xu. Key-grid: Unsupervised 3d keypoints detection using grid heatmap features. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: / //penreii.//pereiee...
Haojie Huang, Dian Wang. Robin Walters, and Robert Platt. Equivariant transporter network. In Proceedings of Roboties: Science and Systems, New York City, NY, USA, June 2022. doi: 10. 15607/RSS.20211 202.XX111.00...
Haojie Huang, Owen Lewis Howell, Dian Wang, Xupeng Zhu, Robert Platt, and Robin Walters. Fourier transporter: Bi-equivariant robotic manipulation in 3d. In The Twelfth International Con- ference on Learning Representations, 2024.
Wenlong Huang, Igor Mordatch, and Deepak Pathak. One policy control them all: Shared modular policies for agnot-ig control. In ICML, 2020.
Jiechuan Jiang, Chen Dun, Tiejun Huang, and Zongqing Lu. Graph convolutional reinforcement learning. In International Conference on Learning Representations, 2020. URL https://
openreview forumfooummooo www.://ooooodidd///dooooooo di=HkkxdkkSDDD
S. Mohammad Khansari- and Aude Billard. Learning stable nonlinear dynamical systems with gaussian mixture models. IEEE Transactions Robotics, 27(5):943-957, 2011. doi: 10.
1109/RO.2201 1.215941..
Vitaly Kurin, Maximilian Igl, Tim Rocktäschel, Wendelin Boehmer, and Shimon Whiteson. My body is a cage: the role of morphology in graph-based incompatible control. In International Conference on Learning Representations, 2021.
Rita Laezza, Robert Gieselmann, Florian T. Pokorny, and Yiannis Karayiannidis. Reform: A robot learning sandbox for deformable linear object manipulation. In 2021 IEEE International Confer- ence on Robotics and Automation (ICRA), pp. 4717-4723, 2021. doi: 10.1109/94448500..2202
9561766.
12

---

Published as a conference paper at ICLR 2025
Ge Li, Hongyi Zhou, Dominik Roth, Serge Thilges, Fabian Otto, Rudolf Lioutikov, and Gerhard Neumann. Open the black box: Step- based policy updates for temporally correlated episodic reinforcement learning. In The Twelfth International Conference Learning Representations, 2023.
temporallycoorelltee
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. In Advances in Neural Information Processing Systems, volume 33, pp- 6755 6766, 2020.
Xingyu Lin, Yufei Wang, Jake Olkin, and David Held. Softgym: Benchmarking deep rein- forcement learning for deformable object manipulation. In Jens Kober, Fabio Ramos, and Claire Tomlin (eds.), Proceedings of the 2020 Conference on Robot Learning, volume 155 of Proceedings of Machine Learning Research, pp. 432 448. PMLR, 16-18 Nov 2021. URL m.ls.press/ 1ww.5551in2a..tnmn..
Mayank Mittal, Calvin Yu, Qinxi Yu, Jingzhou Liu, Nikita Rudin, David Hoeller, Jia Lin Yuan, Ritvik Singh, Yunrong Guo, Hammad Mazhar, Ajay Mandlekar, Buck Babich, Gavriel State, Marco Hutter, and Animesh Garg. Orbit: unified simulation framework for interactive robot learning environments. IEEE Robotics and Automation Letters, 86))3740-37447, 2023. doi: 10. N0.1109/LRA.222332700044
Hai Huu Nguyen, Andrea Baisero, David Klee, Dian Wang, Robert Platt, and Christopher Am- ato. Equivariant reinforcement learning under partial observability. 7th Annual Conference on Robot Learning, 2023.
Fabian Otto, Philipp Becker, Vien Anh Ngo, Hanna Carolin Maria Ziesche, and Gerhard Neu- mann. Differentiable trust region layers for deep reinforcement learning. International Confer- Learning Representations, 2021. URL https: /wwp//ppwewww...///// net ffrrmorrdiii qY2D-AO1Vn.
/
Fabian Otto, Onur Celik, Hongyi Zhou, Hanna Ziesche, Vien Anh Ngo, and Gerhard Neumann. Deep black-box reinforcement learning with movement primitives. In Karen Liu, Dana Kulic, and Jeff Ichnowski (eds.), Proceedings of The 6th Conference on Robot Learning, volume 205 of Proceedings of Machine Learning Research, pp. 1244-1265. PMLR, 14-18 Dec 2023. URL ht htttss://proceedi.e mmr...mm........ www.mmrrrsessss. press www.rppeesssuvve vvtp//225/////// 20t00//otto2aaahhtte
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Viti,,, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W W, 2017.
Eran Rosenbluth, Jan Tönshoff, Martin Ritzert, Berke Kisin, and Martin Grohe Distinguished in uniform: Self-attention vs. virtual nodes. In The Twelfth International Conference on Learning Representations, 2024.
Hyunwoo Ryu, Jiwoo Kim, Junwoo Chang, Hyun Seok Ahn, Joohwan Seo, Taehan Kim, Jongeun Choi, and Roberto Horowitz. Diffusion-edfs: Bi-equivariant denoising generative modeling on se(3) for visual robotic manipulation. arXiv preprint arXiv:2309.0248 02685, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, hbs/1707.063477 2017.
Haochen Shi, Huazhe Xu, Zhiao Huang, Yunzhu Li, and Jiajun Wu. Robocraft: Learning to see, simulate, and shape elasto-plastic objects in 3d with graph networks. The International Journal of Robotics Research, (3(4):533-549 2024. doi: 10. 10.1177/078564922119900
Gregor N. C. Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hernández-LLoato Symmetry- aware actor-critic for 3d molecular design. In International Conference on Learning Representa- tions, 2021. URL ht http:: openreview n forumm iid jiyK E1xYN.
Joshua Southern, Francesco Di Giovanni, Michael Bronstein, and Johannes F Lutzeyer. Under- standing virtual nodes: Oversmoothing, oversquashing, and node heterogeneity. arXiv preprint arXiv:2405.1 13526, 2024.
13

---

Published as a conference paper at ICLR 2025
Bingjie Tang. Iretiayo Akinola, Jie Xu, Bowen Wen, Ankur Handa, Karl Van Wyk, Dieter Fox, Gaurav Sukhatme, Fabio Ramos, and Yashraj Narang. Automate: Specialist and generalist assembly policies over diverse geometries. In Robotics: Science and Systems, 2024.
Brandon Trabucco, Mariano Phielipp, and Glen Berseth. AnyMorph: Learning transferable po- lices by inferring agent morphology. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 21677- 21691. PMIR, 17-23 Jul 2022.
Narek Tumanyan, Assaf Singer, Shai Bagon, and Tali Dekel. Dino-tracker: Taming dino for self- supervised point tracking in single video, 2024.
Elise Van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehock, and Max Welling. Mdp homo- morphic networks: Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems, 33:4199 199-4210.. 2020.
Elise van der Pol, Herke van Hoof, Frans A Oliehoek, and Max Welling. Multi-agent MDP ho- momorphic networks. In International Conference on Learning Representations, 2022. URL
net htw.://iwww...t/oooordidddoooo
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018.
Dian Wang, Robin Walters, and Robert Platt. (Quu)eequivarian reinforcement learning. In Inter- national Conference Learning Representations, 2022.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet:
Learning structured policy with graph neural networks. In International Conference on Learning Representations, 2018. URL
.et/ wwwe/foouuie forum?id= Sitammmmdd-sss---- SHMCDD
Jingyun Yang. Ziang Cao, Congyue Deng, Rika Antonova, Shuran Song, and Jeannette Bohg. Equi- bot: SIM(3)-equivariant diffusion policy for generalizable and data efficient learning. In 8th Annual Conference on Robot Learning. 2024.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Andy Zeng. Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong, Ivan Krasin, Dan Duong, Vikas Sindhwani, and Johnny Lee. Transporter networks: Rearranging the visual world for robotic manipulation. Conference on Robot Learning (CoRL., 2020.
14

---

Published as a conference paper at ICLR 2025
A PROOFS OF PROPOSITION 3.1
Proof. We write simplified updates of HEPi and MPNN VNiocal as follows,
HEPi:
pobi.(+11) = -pchi,10) to
W Σ
k(:: )
E Vanj.
fact, (t+11 poct.()) + Wlocal,(()
Σ
k(,," (aci-ac) Pf
Vaca
(3)
wEN(o)a
(lc,(1+1) pact, new. (I+1) + p0
+ a
w Σ
kes: Abbj-ace rybj.
v E Vact-
Mc Vitn
where a is activation function, and ((( is the weight at layer L. Note that the object nodes are updated through L layers.
MPNN + Vitcaa::
(1bi (()),.()
w
k( Oonj-obj) ) (obi.(1)
Vabj.
(()(())a
pact. (1+1)
(1)
(1)
k ma p (p)ct())
v. E Vuct
(4)
WEN(s
pact (1++.
fort new.(L+)))))))
+ ())t.() to
f
W
k(+, oooj-aat) p.hji.r. (E)
E Vact-
As seen, the main difference between HEPi and MPNN VNiocal is at treating the actuator nodes as VN nodes as in MPNN + VNG or normal graph nodes with k-NN connections.
For HEPi, the actuator nodes updated through every object node as in the third equation in Eq. 3. Explicitly, we compute its Jacobian w.r.t object nodes as
(1+1)
L pact a'(2' (=(2/ocall (W))))lloo
(1
Σ
1e Taus Pact V (.get,())
(5)
o' (2. (Kr Fui Oubj-act)
with 20 - Σ Eneva k(,, ar Lubi. and ze Jocal,(1) (wppalll))))))
EweN(r) k(s Ouet-ce ) (())
be the evaluation of the argument of the function a. This shows that any object node exchange information with the actuator nodes after a single layer of the object-actuator update.
For MPNN VNiocat. if an actuator node and object node each other, the message from node u sent to will arrive either through another actuator node (via actuator-actuator updates) through node where the k-NN connections of those actuator nodes overlap (depicted in Figure 9). However, in both cases, the Jacobian at the actuator node becomes independent of the feature the object node u, i.c.. can only receive a homogeneous value from the VN (overlapping node or other actuator node) (Southern et al., 2024). Consequently, the policy could fail to predict relevant actions response to changes the object node u.
than 2-hops distant from
Related Work for Oversquashing in Graph Neural Networks Transformers (Vaswani et al., 2017) can be viewed as fully connected GNNs under the Message Passing Neural Network (MPNN) framework (Battaglia al., 2018), since self-attention be seen mechanism aggregate
15

---

Published as a conference paper at ICLR 2025
(a) Overlapping node.
(b) No overlapping node.
Figure 9: Demonstration of graph with overlapping and non-overlapping nodes, Actuator nodes in red, object nodes are either white or orange.
messages from its neighbors. While GNNs offer efficient local information processing with linear complexity O(VV + Fi)) they struggle with over-squashing, limiting their ability to propagate information across distant nodes (Di Giovanni et al., 2023). In contrast, Transformers, by being fully connected, allow every node exchange information with all others, making them well-suited for tasks requiring global information aggregation, though at higher quadratic complexity ((VV)²).
Recent studies have shown that introducing virtual nodes (VN) in GNNs can mitigate the over- squashing issue by facilitating long-range information exchange while retaining the lower complex- ity of GNNs (Di Giovanni al., 2023; Cai al., 2023; Rosenbluth al., 2024). Our HEPi builds on this insight, treating actuators as virtual nodes to enable efficient global information aggregation from object nodes, as shown in Figure
B TASKS DETAILS
Here, provide detailed specifications for each of the seven manipulation tasks introduced in the main paper.
B. RIGID-SLIDING
The goal of the Rigid-Sliding task is to control object using suction gripper and slide it on 2D plane desired target position and orientation. The agent controls the object's linear velocity and angular velocity in the yaw direction.
Figure 10: Example trajectory of Rigid Sliding task.
Input and Output
The input space for each node includes:
Gripper nodes: node_type position Pa. velocity V.. angular velocity wa.
Object nodes: node type, position Po: distance to target durget-
The output consists of the gripper's s linear velocity v, and vector from which the angular velocity wa is derived. Specifically, the vector is decomposed into its parallel and tangential components
16

---

Published as a conference paper at ICLR 2025
with respect to the position vector r, where V|
|m|
vir
and T1 = VII The angular velocity
then computed as wa |||||
rxv
Sample Space
Initial pose: (x, ys yuw) E [-11, x [-
Target pose: yyaw E m].
Reward Function The reward consists of multiple sub-rewards:
Distance to goal:
Rgoul lIp Pyoal
where Po is the object position and Pgoot is the target position.
Rotation distance:
Reotation quat_diff(rr... yoool)
where H, and goo! are the object and goal orientations in quaternion.
Object velocity penalty:
Vobject Vangular + Ulinear where Vangular and Vlinear are the angular and linear velocities of the object.
Action rate penalty:
Aactions
Va: ai-1
where a, and a,-1 are the actions at the current and previous time steps.
The total time- deme-dependent reward with T 100 is:
- 0.8 Pogga 0.4 R Rouationn
0. Vobject
0..002 actions:
1< - 2,
Riot
4.0.ggool Protton
0.1 Vatject 0.002. nntton.nctionss:
t t2 T 2.
B.2 RIGID-PUSHING
The goal of the Rigid-Pushing task is control an object using rod and push it 2D plane to desired target position and orientation. The agent controls the object's 's linear velocity v.
Figure 11: Example trajectory of Rigid Pushing task.
Input and Output
The input space for each node includes:
Gripper nodes: node type. position Pa. velocity Va. angular velocity wa.
Object nodes: node type, position Pa. distance target charget and object's linear veloc-
ity V..
The output consists of the actuator's linear velocity Va.
Sample Space
Initial pose: (x, y, 8ym)) E [-0.5, 0.5)3 x [-n,
Target pose: 0yaw E [[-T, T].
17

---

Published as a conference paper at ICLR 2025
Reward Function
The reward consists of multiple sub-rewards:
Distance to goal:
R Rgaul
lp= Pyoad
where is the object position and Pooot is the target position.
Rotation distance:
Reotation quat- difff " gooal
where F, and Egoul are the object and goal orientations in quaternion.
Distance to object:
Robject IPo Pactuator
encouraging the rod (actuator) to stay close to the object during pushing,
The total time-dependent reward with T 100 is:
0.8 0.08 0.22Rojjott:
<< T 5,
R
8.0 0.8 Roottion 0.22
t t2 T - 5.
B.3 RIGID-INSERTION
The Rigid-Insee Insertion task extends Rigid-Sliding to 3D, where the agent must control both linear and angular velocities to move object along the z-axis and insert it into hole. The state space includes (x, y, 2, 0), and precise alignment is required near the hole before sliding and rotating the object into place.
Figure 12: Example trajectory of Rigid Insertion task.
Input and Output The input space for each node includes:
Gripper nodes: node.type, position Pa. velocity v.. angular velocity wa.
Object nodes:
node_type, position Po: distance to target durget-
The output consists of the gripper's linear velocity e, and vector from which the angular velocity
wi is derived. Specifically, the vector is decomposed into its parallel and tangential components
with respect to the position vector r, where Vi
v.r and v = =v- -II- The angular velocity |r|
is then computed Wa rXv1 ||
Sample Space
Initial pose: (x., y, 2, (yaw) [-1,1 x [0,0.5] [-n,m].
Target pose: 0yaw E [-T,
Reward Function The total reward consists of the following sub-rewards:
Distance to goal:
Ryoul Ilp. Pgoat
where P. is the object's position, and Pgoat is the target position.
18

---

Published as a conference paper at ICLR 2025
Rotation distance:
Rotation = quat.diff((ro Tgoal)
where r, and yooul the object and goal orientations in quatermion.
Distance along a-axis:
Rgual. P Pgoul,=ll where Po,s and Pgoat, represent the positions along the z-axis.
The total time-dependent reward with T 100 is defined as:
goa8 2.0 0.4 T<T-2 4.0 gotatonn - 0. 0.4
Riox
t T
B.4
The Rigid-Insertion- Sonnd--neertion-w--geeeeee Two-ttents task extends Rigid-Insertion into 3D with two linear actuators con- trolling the object. The object's 's initial position and the target are sampled from the upper hemi- sphere. Control limited to linear velocities along two axes, simplifying the task for stability in physical systems.
Figure 13: Example trajectory Rigid Insertion with Two Agents task.
Input and Output The input space for each node includes:
Gripper nodes: node hode_type, position Рa. velocity V.-
Object nodes: node_type, position P.. distance to target duaget-
The output consists of the gripper's velocities; however, only the linear velocity
Sample Space
Initial pose: translate in [0.25, 0.75], y 1-0.75,0.75).. [0.5, 1.25], and rotate around its own axis between [-x, m.
Target pose: (Opich. yyaw) E [-x/2,0] [-,, These samples lie in an upper- hemisphere when rotating unit-vector [l, 0, 0J" an initial pose of the target placement, as shown in Figure 14.
Reward Function The total reward consists of the following components:
Distance to goal:
Rgeal = |lp - Pooat
where P. is the object's position, and Pyoal is the target position.
Rotation distance:
Reotation quat quaf.diff(r.. g
where r, and goo! are the object and goal orientations in quaternion.
The time-dependent reward with T 100 defined as:
-0.8/gaat 0.08 Rotationn
1< -
Ris
-4.0 0.6
t t2 T 22
19

---

Published as a conference paper at ICLR 2025
Figure 14: Sample space Rigid- Sogid-Insertion-Tee Twoortion-TTo-Agettts task.
B.5 ROPE-CLOSING
In Rope-Closing two actuators manipulate the endpoints of deformable rope in 2D plane, with the goal of wrapping the rope around cylindrical object and closing the loop. The rope is segmented into 40 links, with each node representing the pose of link.
Figure 15: Example trajectory of Rope Closing task.
Input and Output
The input space for each node includes:
Gripper nodes: node position Pa. velocity Va.
Object nodes: node_type, position P.. velocity distance to target diaret-
The output consists of the gripper's linear velocity Ue.
Sample Space this task, the rope starts in a straight configuration, and only the mid-point position is sampled, with the rope constrained move accordingly.
Initial rope mid-point: Oyaw E 1- π /4]
Target cylinder: (x, y, yum) 0.0. 0.5] x [- n, m.
Reward Function
The total reward consists of:
Rope closing reward:
Reope closing
|| (gripppe Mgripper,))
(Iriipper, Mgripper, )
Center wrapping reward:
Rwrapping (Thanger: Yhanger) )
Crope,z Czope.p) y
where Crope is the center of the rope:
Thansos
Crope
Puode;
Tepodes
20

---

Published as a conference paper at ICLR 2025
Link velocity penalty:
Vinks
|link,
Pllinks
Action rate penalty:
Aactions
Vaz
ai-1
The total time dimeddependent reward with T 200 is defined as:
{
0.01 Vinks 0.001 1 Actions:
20,
Rexe
-2.0 closing
0.8 R Hrapping
0.01 Vinks - 0.001Agationss:
2 20.
B.6 ROAE-NGAPING
This task requires the rope to form specific shape (e.g., "W") to a desired orientation by control- ling the actuators. The rope is segmented into 80 links, with each node representing the pose of link.
Figure 16: Example trajectory Rope Shaping task.
Input and Output
The input space for each node includes:
Gripper nodes: node St ype, position Pa. velocity Va-
Object nodes: node_type. position Po. velocity Vo. distance to target durget-
The output consists of the gripper's linear velocity Ua.
Sample Space this task, the rope starts in a straight configuration, and only the mid-point position is sampled, with the rope constrained move accordingly.
Initial rope mid-point: 6yam E 1- - /T/
Target orientation: 0yaw 1
Reward Function To allow translation invariance, as the task focuses solely on shaping rather
than the object's position, the shape descriptor is computed using both local and global geometric features of the rope. Given the positions N points along the rope, P = [Pi..... .NNl first compute vectors between adjacent points, Pi+1 1 P.. and their normalized form v, -vv v The angles between consecutive segments are then 0: arccos( V++1).
Next, we compute the global direction vector, PN P1. and normalize it as g g/ g/lgg|| The angles between each segment and the global direction are fgllball arccos( g).
Finally, we compute the relative positions of the points with respect to the midpoint (p1 + p2)/2, and their distances d, = [pp. m||. The shape descriptor is formed by concatenating these angles and distances:
Dshape 10...... 0N Igiobal.1..... Ogtobal.N 11 m. PN m, .i....,ddyy dx].
With the shape descriptor defined, the task's 'ss reward function encourages the current rope configura- tion to match the target shape, while also penalizing rapid changes in actions. The specific reward components are:
21

---

Published as a conference paper at ICLR 2025
Shape matching reward:
Rahape
Durrent
D Parget
|l
where Deurrent and Durgot represent the current and target shape descriptors.
Action rate penalty:
Auciion
0i-1
where a, and a,-1 represent the actions at consecutive time steps.
The total time-dependent reward with T 200 is defined as:
shape 0.0001 actions:
<< < 10,
Ras
Rshape 0.0001Anctionss
t 82 22--10 10.
B.7 CLOTH- HANGING
In the Cloth-Hanging task, four actuators are attached to the corners of cloth with hole. The goal to hang the cloth onto beam by aligning the hole with beam. The cloth is represented of particles and simulated using multi-body mass-spring-damper system. For the policy, only the hole boundary particles are used as the object representation (knn points around the hole centroid, with knn k 10). However, for the value function, all particle points are used, required by the reward function defined below.
set
Figure 17: Example trajectory of Cloth Hanging task.
Input and Output
The input space for each node includes:
Gripper nodes:
node_type, position Pa. velocity V.-
Object nodes: node_type, position Pa, velocity Vo. distance to target hanger darget. dis- tance to initial shape dinitial- For the value function, instead of distances, the absolute coor- dinates of the target and initial node are used as features, while keeping the other features the same. We observed that providing the absolute coordinates of the target and initial points to each node makes learning easier. This may be due to the ability of MLP in DeepSets to infer similarity via dot products between feature channels.
The output consists of the grippers' linear velocity Va.
Sample Space In this task, the cloth always starts a straight configuration, and the mid-point position of the cloth is sampled, with the cloth constrained to move accordingly:
Hole location: each hole location (x,y) is sampled within predefined range of offsets from the cloth's center, with a fixed radius. total, we generate 20 unique hole locations.
Initial cloth mid-point: Optch [-T, m).
Target hanger: (z, E [-0.5, 0.5] and (f 6puch- 6yaw) //2 [-ar, .]. This results in sample that lies in the upper hemisphere with quarter part of the lower hemisphere when rotating a unit-vector [0, 1, 0)" an initial pose of the target hanger, shown in Figure 18.
22

---

Published as a conference paper at ICLR 2025
Figure 18: Sample space of the Cloth-Hanging task.
Reward Function
The total reward consists of the following sub-rewards:
Hole-hanger alignment reward:
Chole-hanger Chhole
Changer || + 0.1 cOs(Qna))))
where Chole and Changer represent the centroids of the cloth's 's hole and the hanger, and aaiggn measures their alignment. The first term is similar the previous work from Antonova et al. (2021). Point velocity penalty:
N
V points
Upoint,
N
where Upuint, is the velocity of the i-th point on the cloth, and N is the total number of points.
Cloth distortion penalty:
1 M Leurrett Amitial
D distortion
M
Linitial:
where leurrent, and /initial, are the current and initial lengths of the i-th edge of the cloth, and M is the total number of edges.
Action rate penalty:
Aaction
V(a 1
where a and 0:-1 represent the actions consecutive time steps.
The total time- deeendepndennt reward with T 100 is defined as:
-0.8Rpole-bangg hanger 0.2V ppoins distoiistotionn 0.00224scioon t<
Ris
4.0. hanger 0.2V points
1.0DDistortinnn distortion 0.00224
1.0.
t> T
C RIGID-BODY TASK OBJECTS
Table presents the number of nodes used for each of the different geometrical shapes considered in the rigid manipulation experiments.
D FURTHER EXPERIMENTS
Ablation smaller sample space Here, show detail version of the ablation in Fig- ure addition to different sample spaces in the main paper, we evaluate all the methods on the
scenarios when drawing 0yaw (-1/8, π
and with fixed-angle setting at 0aa 0yaw 0.
23

---

Published as a conference paper at ICLR 2025
ATE E
Figure 19: Overview of all objects used in the rigid manipulation tasks. From left to right: Triangle, Heart, Diamond, Pentagon, Hexagon, Plus, Star, A-shape, T-shape, and E-shape. In sliding tasks, all objects are used, while the A-shupe, and E-shape are excluded in the insertion tasks due to its complex shape.
Table I: Number of nodes for different geometrical shapes.
Shape
Low Resolution #Nodes
High Resolution #Nodes
Triangle
6
1128
Diamond
8
736
Pentagon
10
1032
Hexagon
12
1120
T-shape
16
1152
Star
20
1068
Plus
24
1224
A-shape
23
1660
E-shape
24
1972
Heart
25
1170
As shown in Figure 20, performance generally increases the orientation range narrows. Interest- ingly, HeteroGNN performs better than its non-heterogeneous GNN in most cases, showing its high expressiveness though being less sample efficient. This phenomenon can attributed the shared networks among all the edge types of the naive GNN model. However, once employing EMPN the backbone, HEPi consistently outperforms all the baselines in terms of both the performance and sample effeciency. This proves equivariant constraint plays crucial role to reduce the problem complexity in large 3D space.
Ablation on K-NN for obj-to-act edges Ablation in Figure 21 investigates the update of MPNN grgid-insertion-tweee
+ VN with varying k-nearest neighbors on tasks rigid-insertion, rope-shaping. For the two insertion tasks, the maximum node size is 25; therefore, only vary from to 10. On the other hand, rope-shaping task has 80 nodes, so i is picked from {1, 10, 20, 50). As shown, when & is small, there no-overlapping nodes, the actuator nodes miss the infor- mation from distant nodes. Meanwhile, when k is bigger, the number of overlapping nodes increase, and therefore, these nodes be reached after one layer of message passing, thus resulting in higher expected return.
and
Ablation on Number of Message-Passing Steps ing the number of message- mhessage-passing steps (#MPs) in both HEPi and naive EMPN model. Our goal to determine whether increasing the number of message--psssngg steps improves policy learning.
this ablation, we examine the impact of vary-
As shown in Figure 22, increasing the number of message-passing steps does not yield significant improvements. Moreover, Figure 23 specifically demonstrates that increasing number of message passing lead to oversquashing, where the amount of information exponentially decays by the number hops, and hence lead to performance drop. HEPi, on the other hand, efficiently reduces the information transmission time only hop. These results suggest that, in our design, fewer
24

---

Published as a conference paper at ICLR 2025
HEPS (Ours)
EMPN
Transformer
HeterrGNN
GNN
-1.00
-1.00
-1.00
IQM Return
1.25
11.25
-1.25
11.00
1150
-1.50
-1.75
-1.75
-1.75
2
2
3
Env Steps (1M)
Env Steps (1M)
Env Steps (1M)
(a)
(b)
(c)
-1.00
11.00
1QM Return
11.55
1.25
-1550
1.50
-1.75
11.7
Env Steps (1M)
Env Steps (1M)
(d)
(c)
Figure 20: Performance of difffrrent models on the Cloth- Haaning task across various sample spaces. Assuming the global scene located at [0, 1, 0 then from left to right, we generate sample by rotating by (a) 6 E TT /4 /2) yaw E (-T, т, x), (b) 8yaw
1- /2), and (c) 8yaw E (-T/4, x/4). Meanwhile, the bottom row shows results for (d) 8yaw E (-T/8, и r8), and (e) the fixed orientation at Orl 0, 0yaw = As the sample space decreases, performance improves across all models, with HEPi consistently outperforming the baselines. The additional plot with fixed orientation on the bottom are averaged over seeds while the others with 10 seeds.
rigid-insertion-ZDzz
rigid-insertion- two-agents-3D
rope-shaping-20
-1.00
-1.00
-2.00
-1.25
1OM Return
11.25
-1.50
-1.50
+3.00
-1.75
-1.75
k-10
k=
-2.00
-2.00
-4.00
12
Env Steps (1M)
Env Steps (1M)
Env Steps (1M)
Figure 21: Ablation different k-nearest neighbors choices for obj-to-uct edges in MPNN VNLocal updates (in Section 3.3), evaluated across multiple tasks: higid-insertion rigid-insertion two-agents, and rope-shaping. Results are averaged over 8 seeds.
message-passing steps suffice to capture the necessary information flow, reinforcing the efficiency of our graph design.
Ablation on Orientation Discretization (ori_dim) In this ablation, we explore the impact of varying the orientation discretization dimension (ori-dim) in the Equivariant Message Passing Network (EMPN). The ori_dim. controls how finely we sample orientations from the S sphere, where higher or 1. iidiim increases the number of samples and better approximates full equivari- ance. In our main experiments, we used a default ori rii of 16. Here, we vary ori .ddim across 8, 16, and 24.
As shown in Figure 24, increasing ori dim generally improves performance in 3D tasks, as a finer discretization better captures orientation changes. However, this comes at the cost of increased
25

---

Published as a conference paper at ICLR 2025
#MP
#MP==3
#MP
rigid-insertion-2D+z
rigid-insertion- two-agents-3D
cloth-hanging-3D
-1.00
-1.00
-1.00
1QM Return
-1.20
-1.20
-1220
-1.40
-1.40
-1.40
НЕРi
EMPN
НEPi
EMPN
НЕPİ
EMPN
Figure 22: Ablation on the number of pnessage-passinge steps (#MPs) for HEPi and EMPN models. For HEPi, #MPs=m refers to object-to-object message-passing layers. Across all tasks, in- creasing the number message-pasee paassiggg steps beyond certain point does not improve performance, as the proposed graph design already efficiently transmits information from observations to actions. Results are averaged over 5 seeds.
HEPS
kwl
km3
kw5
Iw
rigid-insertion-2D--
rigid-insertion-2D+*
rigid-insertion-2D+
-1.00
-1.00
-1.00
JOM Retum
-1.25
-1.25
-1.25
-1.50
-1550
-1.50
-1.75
-1.75
-1.7
-2.00
-2.00
-2.00
8 12
16
0
12 16
20
12 16 20
Env Steps (1M)
Env Steps (1M)
Env Steps (1M)
(a)
(b)
(c) 3
Figure 23: Ablation on different k-nearest neighbors for obj-to-act edges in MPNN + VN (in
Section 3.3) updates, evaluated the Rigid-Insertinn task with varying message passing steps {1, 2. for object nodes. Increasing the number message passing steps degrades performance due to oversquashing. Results averaged over 5 seeds.
training time due to the higher computational demand. Notably, in simpler 2D tasks, increasing ori-dim does not significantly affect performance, focus 3D tasks for this ablation.
Ablation K-Nearest Neighbor Graph (KNN_k) In this ablation, evaluate the effect of
varying the number of nearest neighbors (XNN_K) used to connect object nodes in our graph. Instead of relying on mesh edges, we use a K-nearest neighbor (KNN) graph ensure a more generic representation, a common practice in PointCloud-based representations.
As shown Figure 25, our default setting of KNN_K=3 performs comparably to higher values of KNN_K. Increasing the number of nearest neighbors does not provide additional benefits and can even degrade performance slightly, seen in the aopeshhainng-2D task. This is likely due the overcapacity of messages being passed through the network, which introduces unnecessary com- plexity in the message aggregation process.
E EVALUATION DETAILS
E. IMPLEMENTATION DETAILS
All experiments conducted machine equipped with an NVIDIA A100 NVIDIA H100 GPU. We utilized the TorchRL framework (Bou et al., 2023) for the implementation of PPO and TRPL. algorithms, and PyG (PyTorch Geometric) (Fey & Lenssen, 2019) for han-
26

---

Published as a conference paper at ICLR 2025
0=24
O=16
0=g
rigid-insertion-2D4z
rigid-insertion- two-agents-30
cloth-hanging-3D
-1.00
-1.00
-1.00
1QM Return
-1.20
-1.20
-1200
L
-1.40
-1.40
-1.40
НЕРİ
EMPN
НEPi
EMPN
EMPN
Figure 24: Ablation on the orientation discretization dimension bori. Increasing ori_dim improves performance in 3D tasks, such as rigid-insertion and cloth-hunging, by better approximat- ing full equivariance. However, higher or i_dim also increases training time. Results are averaged
over 5 seeds.
nigid-sliding-DD
igiidinneetion-2D4z
rigid-insertion- two-agents-3D
rope-shaping-20
-1.00
-1.80
1QM Return
-1.00
-1.00
-2.00
-1.20
11.20
11.20
-2.20-
-1.40
1.40
-140
-2.00
k
km3
k=5
km3
ke
ker
Km3
kw
ke5
km3
Figure 25: Ablation on the number of nearest neighbors (KNN_k) used in the KNN graph on HEPi. Increasing KNN_K beyond does improve performance and may even reduce it, particularly in tasks like rope-shaping-211 2D, due to message overcapacity. Results are averaged over 5 seeds.
dling the graph-based structure. The Transformer architecture was implemented using the torch. nn. TransformerEncooee and torch. nn. www..Trasoorrrerccceeeeeee pack- ages from PyTorch (Paszke al., 2017).
E.2 COMPUTATIONAL TIME
We report the computational time for each model on all tasks here. Table 2 reports the total training time and Table demonstrates the size of the input graph, which the main factor contributing to the total training time.
Table 2: Total training time for each task (in hours). (*) We note that the fast training speed of Transformer might be attributed to the internal optimization implementation of PyTorch. (**) In Rope-Shaping task, we report the training time on NVIDIA H100 GPU.
HEPi
EMPN
Transformer'
Rigid-Sliding
2h 10m
2h 56m
1h 155m
Rigid-Pushing
2h 58m
3h 58m
1h 52m
Rigid-Insertion
1h 56m
2h 35m
lh 14m
Rigid-Ing
Inseetioo-TTw-AAAe
Ih 3m
Ih 20m
36m
1h 14m
Ih 40m
51m
Hope-Shaping
2h 58m
4h 57m
Ih 56m
Cloth Cloth-Hanging
2h 40m
2h 36m
2h 21m
27

---

Published as a conference paper at ICLR 2025
Table 3: Maximum number of nodes and type of connection for each task.
Shape
Maximum #nodes
Graph Connections
Rigid-Sliding
25
knn=3
Rigid-Insertion
25
knn=3
Rigid-Pushing
25
knn=3
Mhgic-Inseriio-TToAAgentt
25
knn=3
Rope-Closing
40
knn=3
Rope-Shaping
80
knn=3
Cloth Haaning
10
complete
E.3 GRID SEARCH FOR PPO
To fairly compare PPO with TRPL, we perform a grid search for PPO over 5 seeds and select the best-performing configuration based on the maximum return. We then run the chosen configuration on 5 additional seeds and report the results in Figure 8. Specifically, we tune the clip eps param- eter, which controls how much the new policy allowed deviate from the old policy, with values {0.1, 0.2, 0.3, 0.5), and explore both with and without annealing (annea 1 _cl ccliippps=Tre or False). The clip-eps bounds the probability ratio between the new and old policies to (1 (1- + e), preventing large updates and ensuring stable learning. Lower values of cl clip_eps result in more conservative updates, while higher values allow more flexibility in policy updates. Annealing progressively decreases cl clip-eps time, tightening the constraint training pro- gresses.
E.4 HYPERPARAMETERS
We presents the hyperparameters used across all policy models (HEPi, EMPN, and Transformer) for all the tasks in Table 4.
Table Hyperparameters used for all tasks, In EMPN, the number of layers (*) corresponds the number of message-passian steps.
HEPi
EMPN
Transformer
contextual std
true
true
true
latent dim.
64
64
64
activation
GELU
GELU
ReLU
dropout
false
false
false
num layers num heads
n.a.
num messages (obj-to-obj)
n.a.
n.a.
num messages (obj-to-act) num messages (act-to-act)
I
n.a.
n.a.
1
n.a.
n.a.
ponita orientation dim. ponita degree
16
16
n.a.
2
2
n.a.
ponita spatial hidden dim. ponita fiber hidden dim. ponita widening factor
[64, 64]
[64, 64]
n.a.
[64, 64)
[64, 64]
n.a.
n.a.
The following tables, Table 5 and Table provide details on environment settings, data collection parameters, and training hyperparameters.
28

---

Published as a conference paper at ICLR 2025
Table 5: Hyperparameters for Rigid Environments
Rigid-Sliding
Rigid-Insertion /Rigid-Pushing
Rigid-Insertion Two-Agents
Environments
time steps
100
100
100
warmup steps
episode length (in sec.)
4
4
4
decimation
simulation dt
0.01
0.01
0.01
Data Collection
frames per batch total frames
100k
100k
100k
20M
20M / 30M
6M
Input Graph
obj-to-obj edges
knn=3
knn=3
knn=3
act-to-act edges
n.a.
n.a.
complete
Training
epochs
5
5
5
mini-batch size
1000
1000
1000
learning rate (actor)
3e-4
3e-4
3e-4
learning rate (critic)
3e-4
3e-4
3e-4
critic coeff. entropy coeff.
0.5
0.5
0.5
0.005
0.005
0.005
clip gradient norm
false
false
false
Projection
trust region coeff. mean bound covariant bound
4.0
1.0
0.05
0.05
0.05
0.001
0.0025
0.0025
Critic (DeepSets)
num inner layers
2
num outer layers
2
2
2
hidden dim.
64
64
64
activation
ReL.U
ReLU
ReLUU
layer norm
true
true
true
29

---

Published as a conference paper at ICLR 2025
Table 6: Hyperparameters for Deformable Environments
Rope-Closing
Rope-Shaping
Cloth-Hanging
Environments
time steps
200
200
100
warmup steps
10
10
episode length (in sec.)
4
4
2
decimation
2
2
simulation dt
0.01
0.01
0.01
Data Collection
frames per batch total frames
40k
40k
10k
4M
10M
5M
Input Graph
obj-to-obj edges
knn=3 complete
knn=3 complete
complete complete
act-to-act edges
Training
epochs
5
mini-batch size
200
200
200
learning rate (actor)
3e-4
3e-4
3e-4
learning rate (critic)
3e-4
3e-4
3e-4
critic coeff. entropy coeff.
0.5
0.5
0.5
0.005
0.005
0.005
clip gradient norm Projection trust region coeff. mean bound covariant bound
false
false
false
4.0
1.0
4.0
0.05
0.05
0.05
0.001
0.0025
0.001
Critic (DeepSets)
num inner layers num outer layers hidden dim. activation
2
2
2
2
64
64
ReLU
ReLU
ReLU
layer norm
true
true
true
30


Reviews:
[{"review_id": "iyzOhlzuy6", "reviewer": "Reviewer_qL9N", "summary": "The paper proposes a novel setting for representing robotic manipulation problems as heterogeneous graph learning problems. The authors introduce a graph-based policy model, *Heterogeneous Equivariant Policy (HEPi)*, featuring multiple *SE(3)* equivariant message-passing networks (EMPNs) to model smaller sub-graphs like actuators and objects. HEPi explicitly models heterogeneity by assigning distinct network parameters for each interaction type to reduce message mixing and improve expressiveness. HEPi is claimed to be the first study of equivariant policies on 3D space within a reinforcement learning setting for robotic manipulation.\n\nThe authors theoretically prove that, for HEPi any two actuator and object nodes can exchange information while the graph network with locally connected actuators and object nodes can not. This justifies their design of actuator nodes as global virtual nodes to connect all object nodes\n\nThe authors test the proposed heterogeneous graph representation with reinforcement learning for 6 rigid and deformable object tasks, including Rigid-Sliding, Rigid-Insertion, Rigid-Insertion-Two-Agents, Rope-Closing, Rope-Shaping, Cloth-Hanging. Empirical results show that the proposed approach is more sample efficient and less likely to converge on sub-optimal solutions compared with the state-of-the-art Transformer and non-heterogeneous EMPN methods.", "questions": "1. Can HEPi generalize to more complicated geometries in the real world other than simple ropes or shapes like hearts and stars, which might bring much computational complexity to graph learning?\n2. How sensitive is HEPi to the perturbation of the object coordinates from the observation?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and presents its ideas clearly with solid formulation.\n2. The proposed approach of modeling robotic manipulation problems as heterogeneous graph learning problems is well-motivated and clever to unify the structure for both rigid and deformable object tasks using sub-graphs for actuators and objects.\n3. The adaptation of *SE(3)* equivariant message passing networks is reasonable and suitable to exploit the geometric symmetry for improving the sample efficiency in the large 3D search space of configurations. As far as I know, HEPi is one of the first studies of equivariant policies on 3D space within a reinforcement learning setting for robotic manipulation.\n4. The empirical experiments are comprehensive and sound to support the arguments. Most results are averaged over 10 seeds using\ninterquartile mean with 95% confidence intervals, which is also statistically robust.", "weaknesses": "1. The experiments are limited to simple geometric shapes like ropes, triangles, and stars, which can be easily modeled using a few nodes (< 100 in all the experiments). However, the target objects in most 3D manipulation tasks are more complex and sophisticated, e.g., cabinets or dresses, which require significantly more nodes to model. I doubt whether HEPi can still generalize to all those complicated geometries, which brings much computational complexity to graph learning.\n2. The current framework assumes that the object coordinates are readily available in the observation, which is basically a simulation setting. In real applications without coordinates, such coordinates can only be extracted from other CV models in the wild. As a graph policy model, HEPi is likely to suffer from the cumulated error in the observation and fail the tasks compared to non graph learning methods.\n3. Some typos:\n- L17: objects -> object.\n- L182: “lifts”, the left quotation mark.\n- L344: There should be a blank space before HEPi.\n\nNevertheless, the overall idea is novel, and empirical experiments are solid enough to support the proposed setting. I recommend that the paper should be accepted.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pneCWOhSoM", "reviewer": "Reviewer_Z2z8", "summary": "The authors propose a novel SE(3) equivariant RL method called HEPi, and introduce a new benchmark for future geometry-aware RL evaluations.  The paper use EMPNs as algorithm backbone to allow model generalize between poses and its heterogeneity. The paper use graph as representation and do thorough mathmetical analyze.", "questions": "Thank you for the detailed appendix, which has addressed my questions about the specifics. If the authors can supplement the experiments as suggested, I would be willing to increase my score.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "+ Clever Design: This paper includes several ingenious designs. \n   + The use of EMPNs The use of SE(3)-equivariant networks ensures that the model naturally possesses SE(3) generalization, thereby reducing the search space and complexity. \n   + The use of TRPL the introduction of TRPL in place of traditional PPO makes hyperparameter tuning easier, addressing a significant challenge in RL training.\n\n+ Detailed Appendix Experimental Description: One of the major issues in RL is poor reproducibility. However, in this paper, the authors provide a detailed appendix that lists experimental information, including the reward function and hyperparameters. This makes the paper highly reproducible.\n\n+ Benchmark Design: This paper provides a detailed demonstration of the proposed benchmark tasks in the video and clearly defines these tasks in the appendix. As a result, these tasks and environments can be easily adopted by the research community.\n\nI am inclined to accept this paper.", "weaknesses": "+ Real-World Application: As the authors note, this paper uses ground-truth model inputs and does not incorporate a physical robot, which means it cannot be directly applied to real-world scenarios. However, given the significant contributions of this paper in terms of benchmarking and methodology, I do not consider this a critical issue. Nonetheless, I still recommend that the authors include some specific experiments to evaluate this aspect. I suggest author can simply add visual capture pipeline. The details are as follows:\n    + For task like rigid insertion and rigid sliding, you can simply add camera to take picture then use a common pose estimation module to estimate the position. The potential errors in pose estimation can be used to test the robustness of the proposed method to input errors.\n    + For tasks involving fabrics or ropes, you can directly use a camera to capture point clouds and then construct a graph using the point cloud data.\n+ Geometry aware: This paper does not seem to have any special designs focused on geometry, although graphs do play a role. However, the primary focus of the paper appears to be on SE(3)-equivariant designs.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1BUz9mc5hj", "reviewer": "Reviewer_ozCf", "summary": "This paper addresses the task of manipulating diverse rigid and deformable objects with geometry-aware rl policy.\n\nA heterogeneous graph is proposed to represent rigid and deformable object manipulation tasks. To leverage geometric symmetry for better task performance and sample efficiency, a heterogeneous equivariant policy that utilizes SE(3) equivariant message passing networks is proposed. \n\nEvaluation is carried out on a self-curated rl benchmark, including rigid insertion of diverse objects, as well as rope and cloth manipulation with multiple end-effectors. The proposed method outperforms baseline methods in terms of average returns, sample efficiency, and generalizability.", "questions": "Could the authors expand more on line 51-52 about attention making the optimization landscape more difficult to traverse?\n\nWould the policy be able to generalize to objects with different physical properties?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed heterogeneous graph representation is a general representation for both rigid and deformable object manipulation tasks, and captures the geometric structure of the object. HEPi is a novel formulation of graph-based equivariant policy build on top of the heterogeneous graph representation that enables rl for manipulation of diverse shapes and deformable objects.\n\nThe paper has demonstrated thorough evaluations and ablations in simulation, with convincing results proving that the proposed method has better performance, sample-efficiency, and generalizability. Detailed discussions are provided for the results, providing interesting insights.", "weaknesses": "The object model that contains vertices are required for the object node representation, which is the main observation for the policy. This might not be easy to get in the real world due to a lot of self-occlusions during deformable object manipulation.\n\nThe paper demonstrates thorough evaluations across various simulation benchmarks, but lacks real-world evaluations to make the method fully convincing in terms of it's usefulness to real-world robot applications.\n\ntypo at line 74: equivariacne->equivariance.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i3zGT3wWoi", "reviewer": "Reviewer_6nhr", "summary": "The paper trains policies for one or two end effectors to manipulate rigid and deformable objects (ropes and cloths). It models the actuator and object of interest with a heterogeneous graph and advocates for a heterogeneous equivariant policy (HEPi). The paper benchmarks its method on tasks including manipulating rigid and deformable objects in IsaacLab and reports its superior performance over typical baselines.", "questions": "Please see the weaknesses above.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "In general, I find the story in this paper intuitive and convincing. While the two core building blocks (heterogeneous graphs and equivariant policies) have been explored in prior arts, this paper presents a proper combination of both in the setting of rigid/deformable object manipulation. I feel the proposed method has the potential to be extended to manipulate fluid and rigid/soft/fluid-coupled systems, and it would be a pleasant surprise if the paper could demonstrate a fluid scene similar to the “Pour Water” task in Lin et al 2020. Of course, I understand this is outside the scope of the current submission.", "weaknesses": "I generally agree with the limitations listed in the paper. I think the paper can be improved in the following ways:\n\n1. While the technical method and the story look promising, the benchmark scenes are still relatively simple because the setup only models a point-like end effector, ignoring kinematic and dynamic constraints in practical robotic hand/arm structures. Therefore, transferring the result to a real-world robot is not straightforward and probably requires more algorithmic development.\n\n2. From a dynamic perspective, the 2D rigid-sliding task is too trivial to be a valuable benchmark for evaluating this paper and its baselines. This is already reflected in the result: “HEPi and Transformer policies perform comparably, suggesting that the limited task complexity does not fully leverage the benefits of equivariant constraints. ” I suggest this experiment remove the suction gripper and try pushing the rigid object to its target position and orientation with one/two end effectors and their contact with the object, similar to the setup in http://sain.csail.mit.edu/.\n\n3. Deformable objects can be classified by their dimensions: codimension-2 (e.g., ropes), codimension-1 (e.g., cloths), and codimension-0 (e.g., rubber balls). The paper seems to consider the first two categories only. Handling full-dimension soft objects in the proposed method may be non-trivial because it requires observations of object coordinates and cannot be easily obtained by “integrating state-of-the-art computer vision techniques to extract key points from cameras.” I suggest the paper clarify this point by being more precise with the concept of deformable objects it deals with.", "quality": null, "clarity": null, "significance": null, "originality": null, "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)