TarGF: Learning Target Gradient Field for Object Rearrangement

The object arrangement has been previously studied by researchers from the robotics community Schuster et al. (2012); Abdo et al. (2015, 2016) and the graphics community Fisher et al. (2012); Yu et al. (2011); Yeh et al. (2012); Handa et al. (2016); Qi et al. (2018). However, most of them focus on manually designing rules/ energy functions to find a goal Schuster et al. (2012); Abdo et al. (2015, 2016) or synthesising a scene configuration Yu et al. (2011); Yeh et al. (2012); Handa et al. (2016); Qi et al. (2018) that satisfies the preference of the user. The most recent work Kapelyukh and Johns (2021) tries to learn a GNN to output an arrangement tailored to human preferences. However, they neglect the physical constraints during the process. 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 Batra et al. (2020); Weihs et al. (2021) mitigate the goal specification problem by providing a goal scene to the agent, making the reward shaping feasible. Concurrent works Kant et al. (2022); Sarch et al. (2022) 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 Netanyahu et al. (2022) 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 fundamental problem in 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.

Without reward, previous works explore algorithms to learn a control policy from expert demonstrations via imitation learning Pomerleau (1991); Ross and Bagnell (2010); Ross et al. (2011); Ho and Ermon (2016) or inverse reinforcement learning Fu et al. (2017); Liu et al. (2020); Kostrikov et al. (2018); Ziebart et al. (2008). Imitation learning (IL) Pomerleau (1991); Ross and Bagnell (2010); Ross et al. (2011) aims to directly optimise the policy by cloning the behaviour from the expert trajectories. The inverse reinforcement learning (IRL) Fu et al. (2017); Liu et al. (2020); Kostrikov et al. (2018); Ziebart et al. (2008) tries to learn a reward function from data for reinforcement learning. Considering the difficulty in collecting expert trajectories, there are some example-based methods proposed to learn policy only with a set of successful examples Fu et al. (2018, 2017); Reddy et al. (2019); Eysenbach et al. (2021). The most recent work is RCE Eysenbach et al. (2021) which directly learns a recursive classifier from successful outcomes and interactions for policy optimisation. We also try to find an example-based control method for rearrangement with a set of examples. Our method can learn the target gradient fields from the examples, without any additional efforts in reward engineering. The target gradient fields can provide meaningful reward signals and action guidance for reinforcement learning and traditional planners. To the best of our knowledge, such usage of gradients fields is not explored in previous works.

We argue that estimating the gradients of the log-target-likelihood ${\nabla }_s\log p_{tar}\left(s\right)$ can tackle the first two difficulties mentioned in Sec. 3: For the inaccessibility problem, we can approximate the reward increment between two adjacent states by the first-order Taylor expansion using the gradient field $\log p_{tar}\left(s_{t+1}\right)-\log p_{tar}\left(s_t\right)\approx ⟨{\nabla }_s\log p_{tar}\left(s_t\right),s_{t+1}-s_t⟩$. As shown in Sec. 4.4, this helps to derive a surrogate objective of Eq. 1.

For the sparsity problem, we can leverage the gradient field to help with exploration since the gradient ${\nabla }_s\log p_{tar}\left(s\right)$ indicates the direction that points to the higher density region of the target distribution $p_{tar}$. To estimate the gradients of the log-target-likelihood ${\nabla }_s\log p_{tar}\left(s\right)$, noted as the target gradient field (TarGF) ${\Phi }_{tar}$, we adopt the deep score-matching technique Song et al. (2020) that has achieved impressive results in many research areas recentlyShi et al. (2021); Meng et al. (2021); Luo and Hu (2021).

To address the adaptation problem, we further demonstrate two ways to incorporate the trained target gradient field ${\Phi }_{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 actions with a distributed path planner. 2) For the model-free setting (reinforcement learning), we leverage the target gradient to estimate reward and provide suggested actions in residual policy learning Silver et al. (2018).

In the followings, we will first introduce how to train the target gradient field ${\Phi }_{tar}$ via score-matching in Sec. 4.2 and then introduce our rearrangement framework under the model-based and model-free settings in Sec. 4.3 and Sec. 4.4 respectively.

The target score network ${\Phi }_{tar}:S\to S$ 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 ${\nabla }_s\log p_{tar}\left(s\right)$.

To train the target score network, we adopt the Denoising Score-Matching (DSM) objective proposed by Vincent (2011), which can guarantee a reasonable estimation of the ${\nabla }_s\log p_{tar}\left(s\right)$. In training, we pre-specify a noise distribution $q_{\sigma }\left(\tilde{s}|s\right)$ on the collected target state $s$. Then, we match the output of the target gradient field with the score of the perturbed target distribution $q_{\sigma }\left(\tilde{s}\right)=\int q_{\sigma }\left(\tilde{s}|s\right)p_{tar}\left(s\right)ds$:

The DSM objective guarantees the optimal score network satisfies ${\Phi }_{tar}^{\ast }\left(s\right)={\nabla }_s{q}_{\sigma }\left(s\right)$ almost surely. When $\sigma$ is small enough, we have ${\nabla }_s{q}_{\sigma }\left(s\right)\approx {\nabla }_s\log p_{tar}\left(s\right)$, so that ${\Phi }_{tar}^{\ast }\left(s\right)\approx {\nabla }_s\log p_{tar}\left(s\right)$.

In practice, we adopt a variant of DSM proposed by Song et al. (2020), which conducts DSM under different noise scales simultaneously. We implement the target score network as a graph-based network ${\Phi }_{tar}(s,t)$ that is conditioned on a noise level $t\in (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\in R^{f_s}$ and static properties (e.g.category, bounding-box) $p^i\in R^{f_p}$ as the node feature $[s^i,p^i]\in R^{f_s+f_p}$. The input graph is pre-processed by linear initial feature extraction layers and then passed through two Edge-Convolution layers. After message-passing, the output on each node serves as the component of the target gradient on the i-th object’s state ${\Phi }_{tar}^i\left(s\right)\in R^{f_s}$.

The target gradient $g={\Phi }_{tar}\left(s\right)$ can be translated into a gradient-based action $a^g=G\left(g\right)$, 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 ${\Phi }_{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 $a^g=G\left(g\right)=\frac{v_{max}}{||g|{|}_2}\cdot g$ into the action space, 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 of the rearrangement task.

However, the gradient-based action is imperfect(e.g.leading to collisions between the objects), as the target gradient has no access to the information of environment dynamics. This may be inefficient and unsafe due to the severe collisions between objects.

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 ${\varphi }_{\theta }$ takes the gradient-based action $a^g$ as reference velocities and then outputs the collision-free velocities $a={\varphi }_{\theta }\left(a^g\right)$ 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.

We further propose a framework based on the model-free reinforcement learning (RL) algorithm to accomplish the rearrangement tasks in the model-free scenario, where the agent needs to adapt the dynamics via interacting with the environment.

To tackle the inaccessibility problem mentioned in Sec. 3, we first derive an equivalent $J^{\ast }\left(\pi \right)$ of the original objective in Eq. 1 and further obtain a surrogate objective $\hat{J}\left(\pi \right)$ to approximate this equivalent:

Further, we derive a surrogate objective $\hat{J}\left(\pi \right)$ to approximate the equivalent objective $J^{\ast }\left(\pi \right)$. 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 delta log-likelihood using the Taylor expansion:

In this way, we can optimise the surrogate objective $\hat{J}\left(\pi \right)$ by assigning $\sum _{1\le k\le t}⟨{\Phi }_{tar}\left(s_t\right),s_k-s_{k-1}⟩$ as the immediate reward. In practice, we simply assign the last term of the cumulative summation as the immediate reward $r_t=⟨{\Phi }_{tar}\left(s_t\right),s_t-s_{t-1}⟩$, and this version is shown to be the most effective.

To tackle the sparsity problem mentioned in Sec. 3, our idea is to build our policy ${\pi }_{\theta }$ upon the gradient-based action $a_t^g$ mentioned in Sec. 4.3, as $a_t^g$ is the best direction to increase the likelihood:

This approach is similar to the residual policy learning proposed by Silver et al. (2018). The $a_t^g$ serves as an initial policy, and the agent’s policy network ${\pi }_{\theta }$ takes the target gradient $g_t$ and the current state $s_t$ as input and outputs $\pi (s_t,g_t)$ to ‘polish’ the $a_t^g$. The agent finally outputs ‘polished’ action $a_t=\pi (s_t,g_t)+a_t^g$. 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_t^g$).

We implement 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_t^{tar}+\lambda \ast r_t^i$, where $r_t^{tar}=\sum _{1\le k\le t}⟨{\Phi }_{tar}\left(s_t\right),s_k-s_{k-1}⟩$ is a centralised target likelihood reward, $\lambda >0$ is a hyperparameter and $r_t^i={\sum }_{j\ne i}{c}_{i,j}$ is a decentralised collision-aware reward 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$.

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 real data with implicit target distribution.

Ball Rearrangement includes three environments with increasing difficulty of rearrangement: Circling, Clustering, and the hybrid of the first two, Circling + Clustering. There are 21 balls with the same geometry bounded in a square in each environment. In Clustering and Circling+Clustering, the balls are divided equally into three colours. The goal of each environment is as follows: Circling requires all balls to form into 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; 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., friction coefficient and mass. 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 program to measure the similarity between the state and the target distribution.

Room Rearrangement is built on a more realistic dataset, 3D-FRONT Fu et al. (2021). 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 dataset and shuffle the room with a Brownian Process at the beginning to get an initial state. The agent can assign velocities for each object in the room at each time step. The agent has to learn implicit knowledge learned from the target examples to tidy up the room. More details are in the Appendix A.

For both baseline comparison and ablation studies, we collect a fixed number of trajectories $T=\left\{{\tau }_i\right\}$ starting from the same initial states for each of 5 different random seeds. For each random seed, we collect 100 and $83\times 8$ 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:

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\to 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_{\tau \sim T}\left[F_{proxy}\right(s_t\left)\right]$ and the confidence interval over 5 random seeds. We do not report PL on room rearrangement as it is hard to program the human preferences.

Coverage Score (CS): measures the diversity and fidelity of the rearrangement results (i.e.the terminal states $S_T=\left\{s_T\right\}$ of the trajectories) by calculating the Minimal-Matching-Distance(MMD) Achlioptas et al. (2018) between $S_T$ and a fixed set of examples $S_{gt}$ from $p_{tar}$: $\sum _{s_{gt}\in S_{gt}}\underset{s_T\in S_T}{min}||s_{gt}-s_T||$.

Averaged Collision Number (ACN) reflects the efficiency of the rearrangement process since object collision will lead to object blocking and deceleration. ACN is calculated as ${\sum }_{\tau \in T}{\sum }_{s_t\in \tau }{c}_t$, where $c_t$ denotes the total collision number at time step $t$. More details are in the Appendix D.

We compare our framework (Ours(SAC) and Ours(ORCA)) 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.

Leaning-based Baselines: GAIL: A classical inverse RL method that trains a discriminator as the reward function. RCE: The SOTA example-based reinforcement learning method. 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 to reach the goal using goal-conditioned 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.

Ball Rearrangement: We present the quantitative results on ball rearrangement tasks in the first three columns of Fig. 5. As seen from 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 across all three tasks, while 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 goal may be far away from the initial state. 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. Besides, 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). Therefore, the balls may have unexpected behaviour during the process.

The classifier-based reward learning method, GAIL, RCE, and SQIL, basically fails in most tasks, as shown in Fig. 3 and Fig. 5. 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: The dynamics of the room rearrangement task are too complex to apply explicit planning-based methods, so we only compare our model-free 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.

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-bar 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 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 collisions which hurts the efficiency.

Generalisation: We further evaluate the generalisation of our methods to different numbers of objects. To be specific, we train all the policies in the ball environment with $3\times 7=21$ balls. During the testing phase, we directly transfer each policy to environments with increased ball number, i.e., $3\times 8$, $3\times 9$, and $3\times 10$, without any fine-tuning. As shown in Fig. 6, we report the averaged PL increment from initial to target states of each policy in each configuration. The results indicate that our framework generalises well to different numbers of balls with the least performance drop compared with baselines in all three ball rearrangement tasks.