Deep Reinforcement Learning Based on Local GNN for Goal-Conditioned Deformable Object Rearranging

There are two main approaches: Model-based approaches rely on physics to define the relationship between the configurations of the deformable object and manipulations. Establishing a data-driven dynamic model that can predict configurations of the deformable object is a new trend for accuracy. Wang et. [12] construct a latent space to represent configurations of the deformable object by CIGAN and establish an inverse model. After the data-driving model is established, the sampling-based method [13][14] is mostly used to get the manipulation related to configurations. However, if the initial configurations and goal configurations object varies considerably, interpolating intermediate configuration will bring an extra effort and cumulative error.

Data-driven approaches can be divided into two categories according to the source of the training data: imitation learning and reinforcement learning. In imitation learning, the manipulation task is often formulated as a supervised learning problem where the robot should imitate the observed behaviors. Nair et al. [15] adopt human demonstrations to solve the multi-step manipulation task. Reinforcement learning agent acquires rearranging skills through robot exploration. Matas et al. [16] train an end-to-end RL agent to complete the task of folding a towel and draping a piece of cloth over a hanger. Wu et al. [17] make the robot learn a pick-and-place policy for cloth folding and rope straightening by a conditional reinforcement learning method.

Designing a learning policy for a specific deformable object manipulation task has been widely investigated. However, the generalization of these algorithms is limited. Designing a general framework for deformable object manipulation has achieved some progress recently. Zeng et al. [18] proposed the transporter network that performs well on several rearranging tasks of rigid objects and Seita et al.[11] further improved the network and applied it to additional tasks of deformable objects. Lim et al.[19] made a more systematic classification of the rearrangement tasks, and achieved better results in sequence-conditioned task learning. Shridhar et al. [20] proposed CLIPort that can produce multi-task policies conditioned on natural language. However, it’s hard for these works to transfer to reality for 2 reasons. One reason is using CNN feature directly to represent the deformable object, which brings redundant information for the sparse feature of the deformable object. And the other reason is that end-to-end architecture also highly relies on quantities of data, which is hard to collect in reality. In our model design, the deformable object is represented as a graph, which provides rich semantics including keypoint positions and the topological relation between keypoints. And we use local GNN to process the graph features, which can reduce the size of the model. At the same time, our framework is consist of a keypoint detector (visual feature processing) and a local GNN (manipulation planning). All of these designs make our framework is also effective in reality.

The motivation of our framework is using a single complete graph whose nodes are the keypoints of the image to represent the deformable object with a large degree of freedom. Keypoint detection is the first step for graph construction. We borrowed the model in [21] and the design of our keypoint detector model is shown in Fig. 3. Given an image $I\in\mathbb{R}^{H\times W\times 3}$ of a deformable object, the feature CNN $f_{kp}$ extracts a feature map $f_{kp}(I)\in\mathbb{R}^{H^{{}^{\prime}}\times W^{{}^{\prime}}\times K}$ from $I$ firstly, we use $\Omega$ to denote the image domain ($H\times W$ lattice) and use $\Omega^{{}^{\prime}}$ to denote the feature domain ($H^{{}^{\prime}}\times W^{{}^{\prime}}$ lattice). $K$ heatmaps $\mathcal{H}_{c}(I;k),c\in\Omega^{{}^{\prime}}$ ($k=1,2,3,\dots,K$ ) are obtained in parallel as the channels of $f_{kp}(I)$ ($K$ is the number of keypoints), one heatmap for each keypoint.

Second, we can estimate coordinates $(\hat{x_{k}}^{{}^{\prime}},\hat{y_{k}}^{{}^{\prime}})$ of a keypoint by computing the (spatial) expected value of the lattice in each heatmap that has been normalized via spatial softmax:

$$(\hat{x_{k}}^{{}^{\prime}},\hat{y_{k}}^{{}^{\prime}})=\hat{c_{k}^{{}^{\prime}}}=\frac{\sum_{c\in\Omega^{{}^{\prime}}}c*\exp(\mathcal{H}_{c}(I;k))}{\sum_{c\in\Omega^{{}^{\prime}}}\exp(\mathcal{H}_{c}(I;k))}.$$

(1)

The calculation of coordinates is fully-differentiable and can limit the information flow to improve model performance. And we can get the coordinates in image domain $\Omega$ by a direct mapping:

$$\hat{c_{k}}=(\hat{x_{k}},\hat{y_{k}})=(\frac{H\hat{x_{k}}^{{}^{\prime}}}{H^{{}^{\prime}}},\frac{W\hat{y_{k}}^{{}^{\prime}}}{W^{{}^{\prime}}}).$$

(2)

Finally, we can get a gaussian heatmap by a gaussian-like function centred at $\hat{c_{k}}$ with a small fixed standard deviation $\sigma$:

$$\Phi_{c}(I;k)=\exp(-\frac{1}{2\sigma^{2}}\left\|c-\hat{c_{k}}\right\|^{2})$$

(3)

The end result is a gaussian heatmap $\Phi(I)=\sum_{k=1}^{K}\Phi(I,k)$, $\Phi(I)\in\mathbb{R}^{H\times W}$, where the location of $K$ maxima is the estimated keypoint coordinates. Our training goal is to make the gaussian heatmap of the estimated keypoint coordinates $C^{{}^{\prime}}=\{c_{i}^{{}^{\prime}}=(x_{i}^{{}^{\prime}},y_{i}^{{}^{\prime}})\}_{i=1}^{K}$ and that of the true keypoint coordinates $C=\{c_{i}=(x_{i},y_{i})\}_{i=1}^{K}$ as close as possible instead of making the keypoint coordinates close directly. It is easy to recover the keypoint coordinates exactly from these heatmaps, this gaussian heatmap representation is equivalent to the 2D coordinates. In addition, regressing the heatmap will make the model easier to train compared to regressing coordinates directly because of a higher feature dimension.

The result of keypoint detection is shown in Fig. 4. Our keypoint detector perform well in three kinds of deformable object (rope, rope ring, and cloth) involved in our dataset.

After getting the current coordinates $C_{t}$ and the goal coordinates $C_{g}$ of keypoints, we can get a single complete graphs, whose nodes are $C_{t}$ and $C_{g}$.

To develop an optimal manipulation policy for the rearranging task, we design a DQN, which aims to generate a manipulation that can narrow the gap between current and goal configurations. The network architecture is shown in Fig. 5.

The classic setting of DQN can be represented as a Markov decision process (MDP) define as a tuple $(S,A,T,r,\gamma,\omega)$, where $S$ is the set of full observable states (representation graphs), $A$ is the set of actions (pick-and-place actions), $T:S\times A\times S$ is the state transition probability function (related to dynamics of the deformable object). The reward function $r(s,a)$ specifies the reward received when the agent takes action $a$ at state $s$, $\gamma$ is the discount factor and $\omega$ denotes all of the parameters of the Q-value network.

Given a current state $s_{t}$, where $t$ is the time instant, the robot could obtain the reward $r_{t+1}=r(s_{t},a_{t+1})$, and the accumulated reward is denoted as $R_{t+1}=\sum_{i=t+1}^{\infty}\gamma^{i-t-1}r_{i}$. For rearranging policy $\pi$, we use $Q^{\pi}(s_{t})=\mathbb{E}[R_{t+1}|s_{t},\pi]$ to denote the action-value mapping and $Q*$ is the estimated one. We resort to the Bellman equation to calculate the optimal action policy.

Considering that $S\in\mathbb{R}^{L\times N}$ ($L=2K$, $N$ is the dimension of node representation vector in the graph) and $A\in\mathbb{R}^{L}$ (we use the position of keypoints in current and goal images as candidate positions for picking and placing), the action-value mapping is a sequence to sequence mapping. So we utilize a GNN (Graph Neural Network) as the function approximator to estimate the action-value mapping $Q(\omega)\rightarrow Q*$. Parameters $w$ can be iteratively estimated by minimizing the following temporal difference objective:

$$\hat{w}=\mathop{\arg\min}\limits_{\omega}\mathbb{E}[(r_{t}+\gamma\mathop{\max}\limits_{a_{t+1}}Q(s_{t},a_{t+1};\omega)-Q(s_{t-1},a_{t};\omega))^{2}]$$

(4)

By formulating the temporal difference error as the following objective:

$$L(w)=(r_{t}+\gamma\mathop{\max}\limits_{a_{t+1}}Q(s_{t},a_{t+1};\omega)-Q(s_{t-1},a_{t};\omega))^{2}$$

(5)

We transform the optimization problem into a standard regression estimation problem. We will elaborate the state space, action space, and reward design as follows.

We first conduct experiments to evaluate the performance of our framework on multiple rearranging tasks. The robot is given random goal configurations by only visual input, and the robot needs to rearrange the deformable object to the goal configurations without any sub-goal input. We define that the robot completes a rearranging task within 30 picking and placing actions as a success, and the rest is a failure. The situation of completing a rearranging task is that the average pixel distance between the corresponding keypoints in the current and goal image is less than 10. We have chosen this threshold through lots of experiments. The two configurations that satisfy this situation are similar enough. We evaluate the model on the new 100 testing goal-conditioned rearranging tasks for each type of task. The result of the evaluation experiment is shown in TABLE II. The tasks involved in the circle shape (the geometry is very strict) are hard to complete for complex dynamics, the success rates have reached 75% (rearranging into a circle) and 70% (move the rope ring) even on these tasks. The success rate is also acceptable on the 2-D deformable object rearranging task. Compared with the previous framework for multi-task deformable object rearranging skill learning (the tasks involved are different) [11][19], our success rate is almost equivalent, which can verify the performance of our proposed model.

However, the training of our model is much faster and the size of our model is much smaller. Because we don’t rely on the CNN-oriented image feature adopted in previous work, our graph feature is better at handling the sparse information of the deformable object. The calculation of attention on short sequences (32 keypoints in our model setting) is more efficient than multiple convolution calculation on the whole image, where the most area without the deformable object is redundant. Fig. 8 shows some examples of goal-conditioned deformable object rearranging tasks involved in our dataset. All of these tasks are completed in 30 manipulations.

Instead of directly performing global attention calculation on the sequence of keypoints to generate the distribution of the Q value, we have adopted a local GNN. An experiment has also been conducted to validate the rationality of this design. We compared the reward obtained of our local GNN and the origin GNN in the rearranging task training. The number of layers for both models is set to be the same. As Fig. 9 shown, the reward obtained by local GNN is larger than the reward obtained by origin GNN. The improvements convincingly demonstrated our model design’s positive contribution to agent performance. At the same time, the computational cost of local GNN is half that of origin GNN because the attention calculation is only performed on points from the same image (self-attention) or points from the different image (cross-attention) at each layer, which can further improve the efficiency of our model.

In our framework, the processing of visual features (keypoint detector) and the planning of manipulation (local GNN) are naturally separated, which enables the transfer of our framework from simulation to reality. Our model can learn multiple goal-conditioned deformable object rearranging skills from a large quantity of data in simulation and these skills can be used in reality with only keypoint detector fine-tune. We test our proposed method in a physical environment with the setup shown in Fig. 11. The rope is placed on the platform, and a UR5 robotic manipulator with a suction cup is placed in front of the platform for picking and placing. Images are captured with a Realsence camera, which is fixed on the top of the platform. It’s should be noticed that we use a suction cup as the end-effector and we add 6 markers to attach to the rope for sucking. These markers have litter influence on the dynamic of the rope but can make our experiment easier to conduct. It’s easy to replace the sucker with the gripper by considering orientation in the picking and placing action. The orientation can be the tangent direction at the picking and placing point.

Affected by the Covid-19 epidemic, we only tested the goal-conditioned rope rearranging task. We collect 500 real images of the rope and fine-tune the keypoint detector to reduce errors caused by image style differences. The number of keypoints is set as 6, which is different from the simulation. Our model relies on GNN, which makes the number of keypoints can be not fixed. We can choose a different number of keypoints for deformable objects with different physical properties in reality, which further ensures the transferability of the model. The results are shown in Fig. 10. $I_{t}$ denotes the image at $t$ instant, $I_{g}$ denotes the goal image of the task and $a_{t}$ denotes the picking and placing action taken at $t$ instant. The experimental results demonstrate that the skills our framework learned in the simulation are also effective in reality.