TarGF: Learning Target Gradient Field for Object Rearrangement
Abstract
Object Rearrangement is to move objects from an initial state to a goal state.
Here, we focus on a more practical setting in object rearrangement, i.e ., rearranging objects from shuffled layouts to a normative target distribution without explicit goal specification.
However, it remains challenging for AI agents, as it is hard to describe the target distribution (goal specification) for reward engineering or collect expert trajectories as demonstrations.
Hence, it is infeasible to directly employ reinforcement learning or imitation learning algorithms to address the task.
This paper aims to search for a policy only with a set of examples from a target distribution instead of a handcrafted reward function.
We employ the score-matching objective to train a Tar get G radient F ield (TarGF) , indicating a direction on each object to increase the likelihood of the target distribution.
For object rearrangement, the TarGF can be used in two ways: 1) For model-based planning, we can cast the target gradient into a reference control and output actions with a distributed path planner; 2) For model-free reinforcement learning, the TarGF is not only used for estimating the likelihood-change as a reward but also provides suggested actions in residual policy learning.
Experimental results in ball rearrangement and room rearrangement demonstrate that our method significantly outperforms the state-of-the-art methods in the quality of the terminal state, the efficiency of the control process, and scalability. The code and demo videos are on https://sites.google.com/view/targf .
Introduction
[Figure]
❌
✅ ✅
[Figure]
As shown in Fig.
1, we consider object rearrangement without explicit goal specification [1, 2], where the agent is required to manipulate a set of objects from an initial layout to a normative distribution.
This task is taken for granted by humans [3] and widely exists in our daily life, such as tidying up a table [4], placing furniture [5] , and sorting parcels [6].
However, compared with conventional robot tasks [7, 8, 9, 10, 11], there are two critical challenges in our setting: 1) Hard to clearly define the target state (goal), as the target distribution is diverse and the patterns are difficult to describe in the program, e.g ., how to evaluate the tidiness quantitatively?
As a result, it is difficult to manually design a reward/objective function, which is essential to modern control methods, e.g ., deep reinforcement learning (DRL) [12].
In the physical world, the agents also need to find an executable path to reach the target state efficiently and safely, i.e ., finding a short yet collision-free path to reach the target distribution.
To this end, the agents are required to jointly learn to evaluate the quality of the state (reward) and move objects to increase the quality efficiently.
In other words, the agents should answer the two questions: What kind of layouts does the user prefer , and How to manipulate objects to meet the target distribution ?
Thus, it is necessary to search for a policy without dependence on the explicit goal state and shaped reward.
Imitation learning (IL) can learn a reward function [13, 14, 15] or learn a policy [16, 17, 18] via supervised learning from expert trajectories.
However, collecting a large number of expert trajectories is expensive.
Recently, example-based reinforcement learning (RL) [19, 20] has tried to reduce the reliance on expert trajectories, i.e ., learning a policy only with examples of successful outcomes.
However, it is still difficult to learn a classifier to provide accurate reward signals for learning in the case of high-dimensional state and action space, as there are many objects to control in object rearrangements.
In this work, we separate the task into two stages: 1) learning a score function to estimate the target-likelihood-change between adjacent states, and 2) learning/building a policy to rearrange objects with the score function.
Inspired by the recent advances in score-based generative models [21, 22, 23, 24, 25], we employ the score-matching objective to learn a target gradient field (TarGF) by perturbing the well-arranged examples.
The target gradient field can be viewed as object-wise guidance that provides a pseudo velocity for each object, pointing to regions with a higher density of the target distribution, e.g ., tidier configurations.
To rearrange objects in the physical world, we further demonstrate two ways to leverage the target gradient field in control.
We integrate the gradient-based guidance with a model-based path planner to rearrange objects without collision.
We further derive a reward function for reinforcement learning by estimating the likelihood change via the target gradient.
With the TarGF-based reward and a TarGF-guided residual policy network [26], we can train an effective RL-based policy from interactions.
In experiments, we evaluate the effectiveness of our methods on two typical rearrangement tasks: Ball Rearrangement and Room Rearrangement .
First, we introduce three metrics to analyse the quality of the terminal state and the efficiency of the control process.
In Ball Rearrangement , we demonstrate that both of our control methods can efficiently reach a terminal configuration of high quality and diversity.
We also observe that our methods can find a path with fewer collisions and shorter lengths to converge compared with learning-based baselines [14, 20] and planning-based baselines [27].
In
Room Rearrangement , we further evaluate the feasibility of our methods in complex scenes where there are multiple heterogeneous objects.
Our contributions are summarised as follows:
• We introduce a score-based framework that learns a target gradient field (TarGF) by estimating the denoising direction of the perturbed states of well-arranged target examples for object rearrangement.
• We point out two usages of the target gradient field in object rearrangement: 1) providing guidance in moving direction; 2) deriving a reward for learning.
• We conduct experiments to demonstrate that both traditional planning and RL methods can benefit from the target gradient field and significantly outperform the state-of-the-art methods in sample diversity, realism and safety.
Related Works 2.1 Object and Scene (Re-)Arrangements
The object rearrangement has been previously studied by researchers from the robotics community [28, 2, 29] and the graphics community [30, 31, 32, 33, 34].
However, most of them focus on manually designing rules/ energy functions to find a goal [28, 2, 29] or synthesising a scene configuration [31, 32, 33, 34] that satisfies the preference of the user.
The most recent work [4] tries to learn a GNN to output an arrangement tailored to human preferences.
However, they neglect the physical process of rearrangement.
Hence, the accessibility and transition cost from the initial and target states are not guaranteed.
Considering the physical interaction, the recent works on scene rearrangement [35, 36] mitigate the goal specification problem by providing a goal scene to the agent, making the reward shaping feasible.
Concurrent works [37, 38] also notice the necessity of automatic goal inference for tying rooms and exploit the commonsense knowledge from Large Language Model (LLM) or memex graph to infer rearrangements goals when the goal is unspecified.
Another concurrent work [39] aims at discovering generalisable spatial goal representations via graph-based active reward learning for 2D object rearrangement.
In this paper, we focus on a more practical setting of rearrangement: how to estimate the similarity (i.e., the target likelihood) between the current state and the example sets and manipulate objects to maximise it .
Here, only a set of positive target examples are required to learn the gradient field, rather than the prior knowledge about specific scenarios.
2.2 Learning without Reward Engineering
Without hand-engineered reward, previous works explore algorithms to learn a control policy from expert demonstrations via behavioural cloning [16, 17, 18, 13] or inverse reinforcement learning [40, 15, 41, 42].
Imitation learning (IL) [16, 17, 18] aims to directly learn a policy by cloning the behaviour from the expert trajectories.
The inverse reinforcement learning (IRL) [40, 15, 41, 42] tries to learn a reward function from data for subsequent reinforcement learning.
Considering the difficulty in collecting expert trajectories, there are some example-based methods proposed to learn a policy only with a set of successful examples [19, 14, 20].
The most recent work is RCE [20] which directly learns a recursive classifier from successful outcomes and interactions for policy optimisation.
We also try to develop an example-based control method for rearrangement with a set of examples.
Our method can learn the target gradient field from the examples, without any additional efforts in reward engineering.
The target gradient field can provide meaningful reward signals and action guidance for reinforcement learning and traditional path planners.
To the best of our knowledge, such usage of gradient field is not explored in previous works.
Problem Statement
We aim to learn a policy that moves objects to reach states close to a target distribution without relying on reward engineering.
Similar to the example-based control [20], the grounded target distribution p tar ( s ) is unknown to the agent, and the agent is only given a set of target examples S ∗ =
{ s ∗ } where s ∗ ∼ p tar ( s ) .
In practice, these examples can be provided by the user without accessing the dynamics of the robot.
The agent starts from an initial state s 0 ∼ p 0 ( s 0 ) , where N objects are randomly placed in the environment.
At time step t , the agent takes action π ( a t | s t ) imposed on the objects to reach a next state with dynamics p ( s t +1 | s t , a t ) .
Here, the goal of the agent is to search for a policy π ∗ ( a t | s t ) that maximises the discounted log-target-likelihood of the future states:
[Equation]
where γ ∈ (0 , 1] denotes the discount factor.
Notably, we assume that p tar ( s ) > 0 everywhere since we can perturb the original data with a tiny noise ( e.g ., N (0 , 0 . 00001) ) to ensure the perturbed density is always positive.
[21] also used this trick to tackle the manifold hypothesis issue.
This objective reveals three challenges: 1) Inaccessibility problem: The grounded target distribution p tar in Eq.
1 is inaccessible.
Thus, we need to learn a function to approximate the log-target-likelihood for policy search.
2) Sparsity problem: The log-target-likelihood is sparse in the state space due to the Manifold Hypothesis mentioned in [21].
Hence, even if we have access to the target distribution p tar , it is still difficult to explore the high-density region.
3) Adaptation problem: The dynamics p ( s t +1 | s t , a t ) is also inaccessible.
To maximise the cumulative sum in Eq.
1, the agent is required to adapt to the dynamics to efficiently increase the target likelihood.
To address these problems, we partition the task into 1) learning to estimate the log-target-likelihood (reward) of the state and 2) learning/building a policy to rearrange objects to adapt to the dynamics.
Method 4.1 Motivation
Estimating the gradient of the log-target-likelihood ∇ s log p tar ( s ) can tackle the first two problems mentioned in Sec.
For the inaccessibility problem, we can approximate the likelihood increment between two adjacent states by the first-order Taylor expansion log p tar ( s t +1 ) − log p tar ( s t )
≈ ⟨∇ s log p tar ( s t ) , s t +1 − s t ⟩ .
This helps to derive a surrogate objective (Eq. 4) of Eq.
𝑎 𝑡𝑔
Env
Policy Network 𝜋 𝜃 𝑂𝑅𝐶𝐴Planner 𝜙 𝑠 𝑡+
[Figure]
∗
Perturbed
[Figure]
∗
Perturbation 𝒩(0,𝜎 2 𝑰)
Score Matching ℒ2(𝛿,𝑔) (b) Control Objects with TarGF (a) Learning Target Gradient Field (TarGF)
𝑠
𝑡 𝑎 𝑡 = 𝜋 𝜃 𝑠 𝑡 ,𝒈 𝒕 +𝒢(𝒈 𝒕 ) 𝑎 𝑡
RL-based Control
Model-based Planning 𝑔 = Φ ෩𝑠 ∗ 𝛿 = ෡𝑠 ∗ −𝑠 ∗ 𝜎
[Figure] [Figure]
Grad2Action 𝒢
[Figure]
𝑟 𝑡 (𝑠 𝑡 ,𝒈 𝒕 )
[Figure]
𝑡𝑔 𝑔 𝑡 +
The target examples are first perturbed by Gaussian noise. The target score network is forced to match the denoising direction of the perturbation. (b) Our framework is based on the trained target score network. In the model-free setting, the target score network provides exploration guidance and reward estimation for RL.
In the model-based setting, the TarGF provides reference velocities for a model-based planner based on ORCA.
The planner then outputs collision-free velocities for objects. For the sparsity problem, we can leverage the gradient field to help with exploration since the gradient ∇ s log p tar ( s ) indicates the fastest direction to increase the target likelihood. To estimate ∇ s log p tar ( s ) , we employ the score-matching [22], a novel generative model that achieved impressive results in many research areas recently[23, 25, 24], to train a target score network Φ tar : S → S .
To address the adaptation problem, we further demonstrate two approaches to incorporate the trained target gradient field Φ tar with control algorithms for object rearrangement: 1) For the model-based setting, our framework casts the target gradient into a reference control and outputs an action with a distributed path planner.
For the model-free setting (reinforcement learning), we leverage the target gradient to estimate reward and provide suggested action in residual policy learning [26].
In the followings, we will first introduce how to train the target gradient field Φ tar via score-matching in Sec.
and then introduce our rearrangement framework under the model-based and model-free settings in Sec.
4.3 and Sec.
4.4 respectively.
4.2 Learning the Target Gradient Field from Examples
The target score network Φ tar is the core module of our framework, which aims at estimating the score function ( i.e .the gradient of log-density) of the target distribution ∇ s log p tar ( s ) .
To train the target score network, we adopt the Denoising Score-Matching (DSM) objective proposed by [43], which can guarantee a reasonable estimation of the ∇ s log p tar ( s ) .
In training, we prespecify a noise distribution q σ ( (cid:101) s | s ) = N ( (cid:101) s ; s , σ 2 I ) .
Then, DSM matches the output of the target score network with the score of the perturbed target distribution q σ ( (cid:101) s )
= (cid:82) q σ (
[Equation]
s | s ) p tar
[Equation]
s ) d s
[Equation]
The DSM objective guarantees the optimal score network satisfies Φ ∗ tar ( s ) = ∇ s q σ ( s ) almost surely.
When σ is small enough, we have ∇ s q σ ( s )
≈ ∇
s log p tar ( s ) , so that Φ ∗ tar ( s )
≈ ∇
s log p tar ( s ) .
In practice, we adopt a variant of DSM proposed by [22], which conducts DSM under different noise scales simultaneously. In this way, we can efficiently try different noise scales after only one training.
The target score network is implemented as a graph-based network Φ tar ( s , t ) which is conditioned on a noise level t ∈ (0 , 1] .
The input state s =
[ s 1 , s 2 , ..., s N ] is constructed as a fully-connected graph that takes each object’s state( e.g .position, orientation)
s i ∈ R f s and static properties ( e.g .category, bounding-box) p i ∈ R f p as the node feature [ s i , p i ] ∈
R f s + f p .
The input graph is pre-processed by linear feature extraction layers and then passed through two Edge-Convolution layers.
After message-passing, the output on the i-th node serves as the component of the target gradient on the i-th object’s state Φ itar ( s ) ∈ R f s .
We defer structure details to Appendix.
B.
4.3 Model-based Planning with the Target Gradient Field
The target gradient g =
Φ tar ( s ) can be translated into a gradient-based action a g = G ( g ) , where G denotes the gradient-to-action translation.
A gradient over the state space can be viewed as velocities imposed on each object.
For instance, if the state of each object is a 2-dimensional position s i =
[ x, y ] , then the target gradient component on the i-th object can be viewed as linear velocities on two axes Φ tar ( s )
i =
[ v x , v y ] .
When the action space is also velocities, we can construct a g by simply projecting the target gradient into the action space a g = G ( g ) = v max
|| g
|| 2 · g , where the v max denotes the speed limit.
Following the gradient-based action, objects will move towards directions that may increase the likelihood of the next state, which meets the need for the rearrangement task.
However, the gradient-based action will potentially lead objects to collide with each other since the target gradient cannot infer the environment constraints from only target examples.
To adapt the gradient-based action to the environment dynamics, we incorporate the target gradient field with a model-based off-the-shelf collision-avoidance algorithm ORCA .
The ORCA planner ϕ takes the gradient-based action a g as reference velocities and then outputs the collision-free velocities a = ϕ ( a g ) with the least modification to the reference velocities.
In this way, the adapted action can lead objects to move towards a higher likelihood region efficiently and safely.
For more details of the
ORCA planner, we defer to Appendix B.2.
4.4 Learning Policy with the Target Gradient Field
The model-based approach to ‘correct’ the target gradient assumes the objects are of circular shape. This limits the scope of this method when objects are of non-circular shape ( e.g ., furniture). To this end, we further propose a model-free approach based on reinforcement learning (RL) for object rearrangement, where the agent needs to adapt the dynamics via online interactions.
To tackle the inaccessibility problem mentioned in Sec. 3, we first derive an equivalent J ∗ ( π ) of the original objective in Eq. 1 by subtracting a constant from the original objective:
[Equation]
Further, we derive a surrogate objective ˆ J ( π ) to approximate the J ∗ ( π ) .
We notice that in most physical simulated tasks, the distance between two adjacent states || s t − s t − [-
-]1 || 2 is quite small, which inspires us to approximate the log-likelihood-change of the adjacent states using the Taylor expansion:
[Equation]
In this way, we can optimise the surrogate objective ˆ J ( π ) by assigning (cid:80)
[Equation]
as the immediate reward.
In practice, we simply assign the last term of the cumulative summation as the immediate reward r t = ⟨ Φ tar ( s t ) , s t − s t − 1 ⟩ , and this version is shown to be the most effective.
To further tackle the sparsity problem, our idea is to build a residual policy π θ upon the gradient-based action a gt mentioned in Sec.
4.3, as a gt helps to explore the regions with a high target likelihood:
[Equation]
This approach is similar to the residual policy learning proposed by [26].
The a gt serves as an initial policy, and the residual policy network π θ takes the target gradient g t and the current state s t as input and outputs π ( s t , g t ) to ‘correct’ the a gt .
The agent finally outputs a ‘corrected’ action a t =
π ( s t , g t ) + a gt .
In this way, the agent can not only benefit from the efficient exploration but also have less training burden as it ‘stands on the giant’s shoulder’
( i.e .the gradient-based action a gt ).
We employ the multi-agent Soft-Actor-Critic(SAC) as the reinforcement learning algorithm backbone.
Each object is regarded as an agent, and each agent can communicate with all the other agents.
At each time step t , the reward for the i-th agent is r t =
r tart + λ ∗ r it , where r tart = (cid:80) 1 ≤ k ≤ t ⟨ Φ tar ( s t ) , s k − s k − 1 ⟩ is a centralised target likelihood reward, λ > 0 is a hyperparameter and r it = (cid:80)
j ̸
= i
c i,j is a decentralised collision penalty where c i,j is a collision penalty.
Specifically, when a collision is detected between object i and j ,
c i,j = − 1 , otherwise c i,j = 0 .
Experiment Setups 5.1 Environment
We design two object rearrangement tasks without explicit goal specification for evaluation: Ball Rearrangement and Room Rearrangement , where the former uses controlled environments for better numerical analysis and the latter is built on a real dataset with implicit target distribution.
Ball Rearrangement includes three environments with increasing difficulty of target distribution:
Circling , Clustering , and the hybrid of the first two, Circling + Clustering .
There are 21 balls bounded in a square in each task. All the balls have the same geometry.
In Clustering and Circling+Clustering , the balls are divided equally into three colours.
The goals of the three tasks are as follows:
Circling requires all balls to form a circle.
The centre of the circle can be anywhere in the environment;
Clustering requires all balls to form into three clusters by colour.
The joint centre of each cluster has two choices ( i.e ., red-green-blue or red-blue-green clockwise); Circling + Clustering requires all balls to form into a circle where balls of the same type are adjacent.
To slightly increase the complexity of the physical dynamics, we enlarge the lateral friction coefficient of each ball and make each ball with different physical properties, e.g ., different friction coefficients and masses.
The target examples of each task are sampled from an explicit process.
Thus, we can define a ‘pseudo likelihood’ on a given state according to the sampling process to measure the similarity between the state and the target distribution.
[Figure] [Figure] Room Rearrangement is built on a more realistic dataset, 3D-FRONT [44].
After cleaning the data, we get a dataset with 839 rooms, 756 for training( e.g .target examples), and 83 for testing.
In each episode, we sample a room from the test split and shuffle the room via a Brownian Process at the beginning to get an initial state.
The agent can assign angular and linear velocities for each object in the room at each time step.
To tidy up the room, the agent has to learn implicit knowledge learned from the target examples.
More details are in the Appendix A.
5.2 Evaluation
For both baseline comparison and ablation studies, we collect a fixed number of trajectories T
= { τ i } starting from the same initial states for each of 5 different random seeds.
For each random seed, we collect 100 and 83 × 8 (8 initial states for each room) trajectories for ball and room rearrangement respectively.
Then we calculate the following metrics on the trajectories, where the PL is only reported in ball rearrangement. More details of evaluation are in the Appendix D.
Pseudo-Likelihood (PL) measures the efficiency of the rearrangement process and the sample quality of the results.
We specify a pseudo-likelihood function F proxy : S → R + for each environment that measures the similarity between a given state and the target examples.
For each time step t , we report the mean PL over trajectories E τ ∼ T
[ F proxy ( s t )] and the confidence interval over 5 random seeds.
We do not report PL on room rearrangement as it is hard to program human preferences.
Coverage Score (CS): measures the diversity and fidelity of the rearrangement results ( i.e .the terminal states S T = { s T } of the trajectories) by calculating the Minimal-Matching-Distance (MMD) [45] between S T and a fixed set of examples S gt from p tar : (cid:80) s gt ∈ S gt min s
T ∈ S T || s gt − s T
|| .
If the rearrangement results of a method miss some modes in the ground truth example set, the coverage score will increase so as to hurt the performance.
Averaged Collision Number (ACN) reflects the safety and efficiency of the rearrangement process since object collision will lead to object blocking and deceleration.
ACN is calculated as (cid:80) τ ∈ T (cid:80) s t ∈ τ c t , where c t denotes the total collision number at time step t .
[Figure]
PL
Circling
[Figure]
Circling
[Figure]
Oracle
[Figure]
O u r s ( O
R C A )
G o a l ( O
R C A )
O u r s w / o R L
O u r s ( S
A C )
G
[Figure]
l ( S
A C )
O u r s w /
[Figure]
i d u
[Figure]
l
[Figure]
L
[Figure]
r s ( O
R C A )
G o a l ( O
R C A )
G o a l ( S
A C )
O u r s w / o R e s i
d u a l O u r s ( S A C )
O u r s w / o R L
G A I L ( S A C ) S Q I L R C E O u r s ( O R C A )
G o a l ( S A C )
G o a l ( O R C A ) O u r s ( S A C )
O u r s w / o R L
G A I L ( S A C ) S Q I L O u r s w / o R e s i
d u a l R C E
O u r s ( S
A C )
G A I L ( S
A C )
G o a l ( S
A C )
R C E
S Q I L
Room Arrangement
O u r s w / o R e s i
d u a l
O u r s ( O
R C A )
G o a l ( S A
C )
G o a l ( O
R C A )
O u r s ( S
A C )
G A I L ( S
A C )
O u r s w / o R L R C E S Q I L 0.
[Figure] [Figure]
O u r s ( O
R C A )
O u r s w / o R e s i
d u a l
O u r s ( S
A C )
G o a l ( S A
C )
G o a l ( O
R C A )
O u r s w / o R L
G A I L ( S
A C ) R C E S Q I L
O u r s ( O
R C A )
O u r s w / o R e s i
d u a l
O u r s ( S
A C )
G o a l ( O
R C A )
G o a l ( S A
C )
O u r s w / o R L
G A I L ( S
A C ) S Q I L R C E
O u r s ( S
A C )
G o a l ( S A
C )
G A I L ( S
A C )
S Q I L
R C E
5.3 Baselines
We compare our framework, i.e ., Ours(ORCA) and Ours(SAC) , with planning-based baselines under the model-based setting and with learning-based baselines under the model-free setting. Besides, we design heuristic baselines serving as the upper bounds for ball rearrangement tasks. All the reinforcement learning-based methods share the same network design, capacity and hyperparameters.
More details of baselines are in Appendix. C. Leaning-based Baselines: GAIL: A classical inverse RL method that trains a discriminator as the reward function. RCE:
A strong example-based reinforcement learning method proposed recently.
SQIL: A competitive example-based imitation learning method.
Goal-RL: An ‘open-loop’ method that first generates a goal at the beginning of the episode via a VAE trained from the target examples and then trained a goal-conditioned policy to reach the proposed goal via RL. Planning-based Baselines: Goal(ORCA): This method first generates a goal via the same VAE as in Goal(SAC) . Then the agent assigns an action pointing to the goal as reference control and adjusts the preference velocity using ORCA, similar to Sec. 4.3.
Heuristic-based Baseline: Oracle: We slightly perturb the samples from the target distribution and take these samples as the rearrangement results.
The rearrangement results of this method will be so close to the target distribution yet different from the target examples. Hence, we take this method as a reference for normalising the PL curve and CS bars.  Result Analysis 6.1 Baseline Comparison Ball Rearrangement: We present the quantitative results on ball rearrangement tasks in the first three columns of Fig.
As shown on the top row of Fig.
5, the PL curves of our model-free framework ( Ours(SAC) ) and model-based framework ( Ours(ORCA) ) are significantly higher and steeper than the baselines in all three tasks, which shows the effectiveness of our method in arranging objects towards the target distribution efficiently. In the middle row of Fig. 5, Ours(ORCA) achieves the best performance in CS across all three tasks. Goal(SAC) outperforms Ours(SAC) on Clustering and Circling+Clustering since the goal is proposed by a generative model that can naturally generate diverse and realistic goals.
From the qualitative results shown in Fig.
3, our methods also significantly outperform the baselines in realism.
Our methods also achieve the lowest ACN in most cases, except that Goal(SAC) is slightly better than Ours(SAC) in Circling .
The Goal(SAC) is the strongest baseline yet loses to our framework in efficiency since the generated goals are far away from the initial state. We validate this problem by comparing the averaged path length of the Goal(SAC) with ours in Appendix. E.3.
Results show that our methods converge to rearrangement resulting in a shorter path.
Besides, the Goal(SAC) suffers from invalid goal proposals.
As shown in Fig.
3, the Goal(SAC) and Goal(ORCA) can rearrange objects in the right trend, but the objects may block each other’s way.
In the Appendix E.4, we show the VAE used in Goal(SAC) is not guaranteed to generate a legal state( e.g .the balls overlap in goal).
The classifier-based reward learning method, GAIL , RCE , and SQIL , basically fails in most tasks, as shown in Fig.
3 and Fig.
This is due to the over-exploitation problem that commonly exists in classifier-based reward learning methods.
We defer the experiments to Appendix E.2.
Room Rearrangement: It is infeasible to employ the ORCA planner in room rearrangement since all the furniture is non-circular. Thus, we only compare Ours(SAC) framework with learning-based baselines.
As shown in Fig.
4 and the last column of Fig.
5, our method significantly outperforms other baselines in sample realism of rearrangement results and the safety of the rearrangement process, which proves our method can also work on a more complex task. As all the methods are evaluated on the unseen rooms, results on room rearrangement also demonstrate the generalisation ability of our method.
6.2 Ablation Studies and Analysis
We conduct ablation studies on ball rearrangement tasks to investigate: 1) The effectiveness of the exploration guidance provided from the target gradient field and 2)
The necessity of combining the target gradient field with control algorithms to adapt to the environment dynamics.
To this end, we construct an ablated version named Ours w/o Residual that drops the residual policy learning and another one named Ours w/o RL that only takes gradient-based action a g as policy.
Ours w/o Residual: The CS bars in Fig.
5 shows that Ours(SAC) significantly outperforms the Ours w/o Residual in CS, which means the agent is easy to over-exploit few modes with high density.
We also show qualitative results in Appendix E.1 that without the residual policy learning, the objects are usually arranged into a single pattern.
Ours w/o RL: The PL curves in Fig.
5 shows that without the control algorithm serving as the backbone, the rearrangement efficiency of Ours w/o RL is significantly lower than Ours(SAC) and Ours(ORCA) .
Meanwhile, the ACN of Ours w/o RL is significantly larger than Ours(SAC) and Ours(ORCA) , which indicates the gradient-based action may cause severe collisions without adapting to the environment dynamics.
Scalability: We further evaluate the scalability of our methods to be zero-shot transferred to different numbers of objects.
To be specific, we train all the policies in the ball environment with 3 × 7 =balls.
During the testing phase, we directly transfer each policy to environments with increased ball numbers, i.e., 3 × 8 , 3 × 9 , and 3 × 10 , without any fine-tuning.
In Fig.
6.2, we report the averaged PL increment from the initial to target states of each method.
Results show that our method still outperforms baselines under different numbers of balls while achieving a slight performance drop. [Figure]
Conclusion
In this study, we first analyse the challenges of object rearrangement without explicit goal specification and then incorporate a target gradient field with either model-based planning or model-free learning approaches to achieve this task.
Experiments demonstrate our effectiveness by comparisons with learning-based and planning-based baselines.
Limitation and Future Work.
This work only considers the planar state-based rearrangement with simplified dynamics where the agent can directly control the velocities of objects.
Future works may extend our framework to more realistic scenarios: 1) Visual Observation: We may explore how to conduct efficient and accurate score-based reward estimation from image-based examples and leverage the recent progress on visual state representation learning [46, 47, 48] to develop the visual gradient-based action.
Complex Dynamics: When the agent can only move one object at a time, we may explore hierarchical methods where high-level policy chooses which object to move and low-level policy leverage the gradient-based action for motion planning.
When we can only impose forces on objects instead of velocities, we may explore a hierarchical policy where a high-level policy outputs target velocities and a low-level employ a velocity controller, such as a PID controller, to mitigate the velocity errors.
Ethics Statement.
Our method has the potential to build home-assistant robots.
We evaluate our method in simulated environments, which may introduce data bias.
However, similar studies also have such general concerns.
We do not see any possible major harm in our study.
Acknowledgment
This work was supported by the National Natural Science Foundation of China -Youth Science Fund (No.62006006): Learning Visual Prediction of Interactive Physical Scenes using Unlabelled Videos. Fangwei Zhong was supported by China National Post- doctoral Program for Innovative Talents (Grant No. BX2021008). We also thank Tianhao Wu and Yali Du for their insightful discussions.
The three tasks in ball rearrangement only differ in state representation, pseudo-likelihood function and target distribution:
Circling:
The state is represented as a full-connected graph where each node contains the twodimensional position of each ball s i =
[ x, y ] .
The target examples are drawn from an explicit process with an intractable density function: We first uniformly sample a legal centre o = [ x o , y o ] that allows all balls to place in a circle around this centre without overlap.
Then we uniformly sample a radius r o from the legal radius range based on the sampled centre. After sampling the centre and radius, we uniformly place all balls in a circle centred in o with radius r o .
The pseudo likelihood function is defined as F proxy ( s ) = exp − ( σ θ + σ r ) , where σ θ and σ r denote the standard deviation of the angle between two adjacent balls and the distances from each ball to the centre of gravity of all balls, respectively.
Intuitively, if a set of balls are arranged into a circle, then the σ r and σ θ should be close to zero, achieving higher pseudo-likelihood.
Clustering : The state is represented as a full-connected graph where each node contains the twodimensional position of each ball s i = [ x, y ] and the one-dimensional category feature of each ball c i ∈ { 0 , 1 , 2 } . We generate the target examples in two stages: First, we sample the positions of each ball from a Gaussian Mixture Model(GMM) p GMM : R K ∗ 2 → R + with two modes:
[Equation]
(cid:89)
[Equation]
N ((0 . 18 sin( [Equation] π
[Equation] 18 cos(03 π
[Equation] [Equation] [Equation]
( s [Equation] ) (cid:89)
[Equation] N ((
[Equation] 18 sin( [Equation]
3 π ) , 0 . [Equation] (13 π )) , [Equation]
[Equation] )( s i ) (cid:89) [Equation] K
{+N ((0 . 
 sin(23 π ) , 0 . 18 cos(23 π [Equation]+}
where K = 21 denotes the total number of balls.
In the second stage, we slightly adjust the positions sampled from the GMM to eliminate overlaps between these positions by stepping the physical simulation.
Since the results mainly depend on the GMM, we take the density function of GMM as the pseudo likelihood function F proxy ( s ) = p GMM ( s ) Circling+Clustering : The state representation is the same as that of Clustering .
To generate the target examples, we first generate a circle using the same process in Circling . Then, starting with any ball in the circle, we colour the ball three colours, in turn, K 3 for each. The order is randomly selected from red-yellow-blue and red-blue-yellow. This sampling is also explicit yet with an intractable density similar to Circling .
The pseudo-likelihood function is defined as F proxy ( s ) = exp [Equation] [Equation] · exp
[Equation] [Equation] R + σ G + σ B ) − σ C where σ R denotes the standard deviation of the angle between two adjacent red balls, and σ G and σ B and σ C denotes the standard deviation of the positions of red, green and blue centres. Intuitively, the first term exp − ( σ θ + σ r ) measures the pseudo likelihood of balls forming a circle and the next term exp − ( σ R + σ G + σ B )+ σ C measures the pseudo likelihood of balls being clustered into three piles.
The common settings shared by each task are summarised as follows:
Horizon: Each training episode contains steps. Initial Distribution: We first uniformly sample rough locations for each ball, and then we eliminate overlaps between these positions by stepping the physical simulation. Dynamics: The floor and wall are all absolutely smooth planes.
All the balls are bounded in an 0.3m x 0.3m area, with the radius being 0.025m. We set the friction coefficients of all balls to 100.0 since we observe that setting a small( e.g .not larger than 1.0) friction coefficient does not significantly affect the dynamics. Besides, to increase the complexity of the dynamics, we set the masses and restitution coefficients of all green and blue balls to 0.1 and 0.99, respectively. The masses and restitution coefficients of all red balls are set to 10 and 0.1, respectively. We observe that under these dynamics, the collision may significantly harm the efficiency of the rearrangement process.
Hence, the agent has to adapt to the dynamics for more efficient object rearrangement. Target Examples: We collect 100,000 examples for each task as target examples . A.2 Room Rearrangement Dataset and Simulator: We clean the 3D-Front dataset [44] to obtain bedrooms that consist of four walls and three to eight objects.
We augment each room by flipping two times and rotating four times to get eight augmentations.
We then import these rooms into iGibson [49] to run the physical simulation.
The 756x8 rooms are used for target examples that are used for training of the target score network, the classifier-based baselines and the VAE in goal-conditioned baselines.
The other 83x rooms are used to initialise the room in the test phase, so the room rearrangement task is performed on the test dataset.
The rooms in the training dataset are only used for learning prior knowledge to arrange the room. State and Action Spaces: The state consists of a aspect ratio r a ∈ R + and an object state s o ∈ R K ×where K denotes the number of objects. The aspect ratio r a = tanh ( b x b y ) where b x and b y denotes the horizontal and vertical wall bounds. The object state is the concatenation of sub-states of all the objects s o =
[ s 1 , s 2 , ... s i , ..., s K ] where the sub-state of the i-th object s i ∈ R 6 is consists of 2-D position, 1-D orientation, 2-D bounding bound and a 1-D category label. The action is also a concatenation of sub-actions of all the objects a o =
[ a 1 , a 2 , ... a i , ..., a K ] . For the i-th object, the action a i ∈ R 3 consists of a 2-D linear and a 1-D angular velocity. The whole action space is normalised into a 3 × K dimensional unit-box [ − 1 , 1] 3 × K by the velocity bounds. Horizon: Each training episode contains 250 steps.
Initial Distribution: To guarantee the initial state is accessible to the high-density region of the target distribution, we sample an initial state in two stages: First, we sample a room from the 83x rooms in the test dataset. Then we perturb this room by 1000 Brownian steps. Dynamics: For more efficient environment reset and physical simulation, we build a ‘proxy simulator’ based on PyBullet [50] to replace the original iGibson simulator.
We use iGibson to load and save the metadata of each room.
Then we reload these rooms in the proxy simulator, where each object is replaced by a box with the same geometry. We set the friction coefficient of all the objects in the room to zero, as the dynamics of the room are complex enough. B Details of Our Method B.1 Training the Target gradient Field The complete training objective is the SDE-based score-matching objective proposed by [22]:
[Equation]
[Equation]
(7) where p 0 t ( s ( t ) | s (0))
=
N ( s ( t ); s (0) , 1 2log σ ( σ 2 t − 1) I ) and σ = 25 is a hyper-parameter.
Using this objective, we can obtain the estimated score w.r.t.
different levels of the noise-perturbed target distribution p ttar ( s ( t )) = (cid:82) p 0 t ( s ( t ) | s (0)) p tar ( s (0)) d s (0) simultaneously.
In this way, we can assign different noise-level t for different tasks conveniently.
The choice of t will be described in the sub-section.
B.
We set τ = 0 . 1 and the simulation duration of each timestep ∆ t = 0 . For each agent(object),
ORCA only considers the 2-nearest agents as neighbours, since in our experiments ORCA often has no solution when the number of neighbours is larger than
[Equation]
In all ball rearrangement tasks, we choose t = 0 .
1 as the initial noise scale for the target score network.
The noise level decays to 0 within an episode.
B.3 Complete Derivation of Surrogate Objective
[Equation]
For all the ball rearrangement tasks, we choose t = 0 . 1 for the target score network when outputting the gradient-based action a gt = G ( Φ tar ( s t , 0 . 1)) and t = 0 .
01 when estimating the immediate reward r t = ⟨ Φ tar ( s t , 0 . 01) , s t +1 − s t ⟩ .
At each time step t , each object also receives a collision penalty c it .
So the total reward for the i-th object at timestep t is r it =
⟨ Φ tar ( s t , 0 . 01) , s t +1 − s t ⟩ + λ ∗ c it where λ denotes a hyper-parameter to balance the immediate reward and the collision penalty.
We choose λ = 5 for Clustering and Circling+Clustering , λ = 3 for Circling and λ = 0 .
2 for room rearrangement.
We conduct reward normalisation for both immediate reward and the collision penalty, which maintains a running mean µ t and a standard deviation σ t of a given reward sequence and returns a z-score z t =
r t − µ t σ t as the normalised reward.
For room rearrangement, we choose t = 0 .
01 for both outputting the gradient-based action and estimating the immediate reward.
C Details of Baselines C.1 Goal-based Baselines
These baselines refer to the Goal-SAC and Goal-ORCA in experiments.
Goal Proposal: This type of baseline first train a VAE on the target examples and then leverages the trained VAE for the goal proposal.
The VAE is implemented as a GNN, and the model capacity is similar to our target score network for a fair comparison.
We choose λ kl = 0 . 01 for Circling and Clustering and λ kl = 0 . 02 for Circling+Clustering .
Execution: At the beginning of each episode, the agent first proposes a goal for this episode using the VAE and then reaches the goal via a control algorithm.
In Goal-ORCA , the agent reaches the goal by a planning-based method: The agent first assigns velocities for each ball that points to the corresponding goal.
Then these velocities are updated by the ORCA planner ϕ to be collision-free.
In
Goal-SAC , the agent trains a multi-agent goal-conditioned policy via goal-conditioned RL to reach the goal: Similar to Ours-SAC , the reward of the i-th object at timestep t is r it
=
||
s t − s goal || 1 + λ ∗ c
it where s goal denotes the goal proposal, λ is a hyper-parameter and c it =
(cid:80) j ̸ = i col i,j is the total number of collisions between the i-th object and the others. Here col i,j = 1 when We choose λ
=for ball rearrangement and λ = 0 .
2 for room rearrangement.
C.2 Classifier-based Baselines
These baselines refer to the RCE , SQIL and GAIL in experiments.
RCE and SQIL are implemented based on the codes [20] released by RCE’s authors.
We only modify γ = 0 .
95 , the training steps decrease to 0.5 million ( i.e ., the same number of training steps as other methods) and the model architecture.
The architecture of actor and critic networks is implemented the same as ours( i.e ., the same feature extraction layers and edge-convolutional layers).
GAIL is actually a modification of our learning-based framework:
Keeping the RL agent the same( i.e ., multi-agent SAC), GAIL ’s reward is given by a discriminator.
The discriminator takes only one state instead of two adjacent states as input since the ‘ground-truth’ reward( i.e ., likelihood) is defined on the current state.
At each training step, we update the discriminator by distinguishing between the agents’ and the expert’s states(for one step) and then update the RL policy under the reward given by the discriminator(for one step).
The agent also receives a collision reward during training similar to
Ours-SAC and Goal-SAC : r it =
D ( s t +1 ) + λ ∗ c
it where D denotes the classifier trained by GAIL .
We do not conduct reward normalisation for GAIL as the learned reward is unstable.
We choose λ
=for room rearrangement, Circling and Circling+Clustering , λ = 100 for Clustering .
D Details of Evaluations D.1 Ball rearrangement
We collect 100 trajectories for each task starting from the same set of initial states.
To calculate the coverage score, we sample fixed sets examples from the target distribution serving as S gt for
Circling , Clustering , and Circling+Clustering , respectively.
We sample 20 examples for for Circling and Circling+Clustering and 50 examples for Clustering .
Since the balls in the same category can actually be viewed as a two-dimensional point cloud, we measure the distance between two states by summing the CDs between each pile of balls by category.
D.2 Room rearrangement
For each room in 83 test rooms, we collect eight initial states.
Then we collect 83x8 trajectories starting from these initial states for each method.
The coverage score is calculated by averaging the coverage score in each room condition since the state dimension differs in different rooms.
For each room in 83 test rooms, we calculate the coverage score between the eight ground truth states and eight rearrangement results and then the averaged coverage score over the 83 rooms is taken as the final coverage score for a method.
We measure the distance between two states by calculating the average L2 distance between the positions of the corresponding objects.
E Additional Results E.1 Single-mode Problem of Ours w/o Residual
In Figure.
7 we show qualitative results of Ours w/o Residual .
Apparently, the balls are arranged into a single pattern in all three ball rearrangement tasks, while the examples from the target distribution are diverse.
In the most difficult task Circling + Clustering , the agent cannot even reach a terminal state with a high likelihood. This result indicates that Ours w/o Residual failed to explore the high-density region of target distribution without residual learning. [Figure] E.2 Effectiveness of Reward Learning
In each ball rearrangement task, we collect 100 trajectories, each of which is run by a hybrid policy.
The hybrid policy takes the first 50 steps using the Ours(ORCA) and the next 50 steps using random actions. To evaluate the effectiveness of our method on reward learning, we compare the estimated reward curve of Ours(SAC) and GAIL with the pseudo likelihood(PL) curve of the trajectories. As shown in 8, the reward curve of Ours(SAC) best fits the trend of the pseudo-likelihood curve, which shows the effectiveness of our reward estimation method. [Figure]
In figure. E.3, we report the comparative results of our framework and goal-based baselines( e.g ., Goal(SAC) , Goal(ORCA) ) on a new metric named absolute state change(ASC) .
The ASC measures the sum of the absolute paths of all small balls in the rearrangement process.
[Equation]
As shown in figure. E.3, Ours(SAC) and Ours(ORCA) are significantly better than Goal(SAC) and
Goal(ORCA) , respectively. This result explains why our method’s likelihood curves are better than the goal-based baselines’: The proposed goal is far away from the initial state, which harms the efficiency of goal-based approaches. E.4 Visualisations of Goals Proposed by the VAE
We demonstrate the visualisations of goal proposals by the VAE used in goal-based baselines.
Typically, the proposed goals indeed form a reasonable shape that is similar to the target examples.
However, it is hard to generate a fully legal goal since the VAE is not accessible to the dynamics of the environment or has enough data to infer the physical constraints of the environment. As shown in [Table] [Table]
Fig.
9, there exist many overlaps between balls in generated goals, which causes the balls to have conflicting goals and thus harms the efficiency of goal-based baselines
[Table]
[Figure]
and Goal(SAC) .
E.5 Six-modes Clustering
Task Settings: This task is a six-modes extension of Clustering , where the centres of clusters are located in the following six patterns.
Defining the joint centres’ positions as a latent variable C
= ( C r , C g , C b ) where C r , C g and C b denotes centres of red, green and blue balls respectively and the above six modes as { c i } 1 ≤ i ≤ 6 , the
C obeys a categorical distribution p ( C = c i )
= 16 .
The target distribution of six-modes clustering is a Gaussian Mixture Model:
[Equation]
Notably, the ‘mean mode’ of above six modes is the origin, i.e ., 16 (cid:80) 6 i =1 c i =
If the policy arranges the balls into this mean pattern, then the balls should be centred around (0 , 0)
( i.e ., all positions of balls obey N (0 , 0 .
05 I ) ).
As shown in Figure 10 (a), we illustrate an example for each mode.
[Table]
Implementation Details: We evaluate Ours(SAC) and Ours(ORCA) on this task.
The performances are normalised by the Oracle.
Results: We report the pseudo-likelihood curves of our methods.
As shown in Figure 10 (b), the rearrangement results are close to ground truth examples from the target distribution.
We further compute the latent distributions for rearrangement results of Ours(SAC) using the Bayesian
Theorem:
[Equation]
p ( C = c i | s ) indicates which mode a state belongs to.
To demonstrate our rearrangement results are close to one of the six modes instead of concentrating on a ‘mean’ pattern, we evaluate the average entropy of the latent distributions
E [ H ( p ( C | s ))] .
As shown in Figure E.5, the average entropy of our methods is even lower than ground truths’.
This indicates each state in our rearrangement results distinctly belonging to one of the categories.
We also report the average latent distribution E [ p ( C | s )] .
As shown in Figure E.5, the averaged latent distribution (overall rearrangement results) of our methods achieve comparable orders of magnitude in different modes.
This shows the rearrangement of our methods can cover all the mode centres.
E.6 Move One-ball at a Time
Task Settings: This task is an extension of Clustering where the policy can only move one ball at a time.
We increase the horizon of each episode from 100 to 300.
At each time step, the agent can choose one ball to take a velocity-based action for 0.1 seconds.
Implementation Details: We design a bi-level approach based on our method as shown in Figure 11 (b): Every 20-time steps, the high-level planner outputs an object index i t with the largest target gradient’s component:
[Equation]
where g it ∈ R 2 denotes the component of the target gradient on the i-th object.
In the followingsteps, the ORCA planner computes the target velocity according to g it and masks all the other object’s velocities to zero.
We compare our method with a goal-based baseline where the agent first generates goals for each object via the VAE used in Goal(ORCA).
Then the high-level planner chooses the object with the farthest distance to the goal as i t .
The low-level planner of this baseline is the same as ours.
Results: As shown in Fig.
11 (a), (c), our method achieves more appealing results, better efficiency and performance compared with the goal-based baseline.
E.7 Force-based Dynamics
Task Settings: This task is an extension of Circling + Clustering where the agent can only impose forces on the objects instead of velocities.
At each time step, the agent can assign a two-dimensional force f i ∈ R 2 on each object.
Implementation Details: Similar to the ‘one at a time’ experiment, we design a bi-level approach to tackle this task, as Fig.
12 (b) illustrates.
The high-level policy outputs a target velocity everysteps.
In the following 8 steps after the target velocities are outputted, the low-level PID controller receives the target velocity and outputs the force-based action to minimise the velocity error.
We set
{+K P = 10 . 0 , K I = 0 . 0 , K D = 0 . 
+}
for PID controller.
This policy is compared with Ours(SAC) in the main paper.
Results: As shown in Figure 12.
(a) and (c), this force-based policy achieves comparable performance with Ours(SAC) yet suffers from a slight efficiency drop due to the control error of PID.
E.8 Image-based Reward Learning
[Figure]
Task Settings:
In this task, the target score network is trained on image-based target examples.
We simply render the state-based target example set used in Ours(SAC) to 64x64x3 images and
Implementation Details: Our target score network of
Ours(Image) is trained on the image-based examples set.
The policy network and gradient-based action are state- based since we focus on image-based reward learning instead of visual-based policy learning.
This approach is compared with Ous(SAC) and Goal(SAC) in the main paper.
Results: Results in Figure 13 show that Ours(Image) achieves slightly better performance than
Goal(SAC) yet is lower than Ours(SAC) due to the increment of the dimension.
This indicates that our reward learning method is still effective in image-based settings.