Deformable Linear Object Prediction Using Locally Linear Latent Dynamics

First, we introduce some fundamental concepts.

Rope Dynamics. We consider nonlinear rope dynamics

where $t\in\{0,1,...,T\}$ is the time step, $x(t)\in\mathcal{X}\subseteq\mathbb{R}^{m}$ is the state, $u(t)\in\mathcal{U}\subseteq\mathbb{R}^{n}$ is the control input. Since we cannot get the rope state easily, we denote the rope image as the state $x(t)$. In our case, $u=(x,y,l,\theta)$, where $x$ and $y$ are positions in the image space, moving length $l$ and moving angle $\theta$ are in the world space.

Latent Dynamics. The rope dynamics are non-linear, very complicated to learn, and hard to generalize, while linear dynamics are generally simpler to learn and can help make state predictions more easily. Thus, we consider a linear dynamics model

in the latent space. Here, $g(t)\in\mathcal{G}\subseteq\mathbb{R}^{p}$ is the latent state, $a(t)\in\mathcal{A}\subseteq\mathbb{R}^{q}$ is the latent action, $K(t)\in\mathcal{K}\subseteq\mathbb{R}^{p\times p}$ is the state matrix, and $L(t)\in\mathcal{L}\subseteq\mathbb{R}^{p\times q}$ is the control matrix. The state matrix $K(t)$ is a transition matrix from the current latent state $g(t)$ to the next latent state $g(t+1)$. The control matrix $L(t)$ learns the impact of the latent action $a(t)$ on the state transition.

For globally linear latent dynamics, $K(t)=K$ and $L(t)=L$ for all $t\in\{0,1,...,T\}$, where $K$ and $L$ are fixed matrices. For locally linear latent dynamics, $K(t)$ depends on the state $x(t)$ and $L(t)$ depends on both state $x(t)$ and action $u(t)$.

Encoder Models. The state encoder model is a map $\phi_{e}:\mathcal{X}\to\mathcal{G}$, the action encoder model is a map $\varphi_{e}:\mathcal{U}\to\mathcal{A}$. We then have

Encoder models encode the state $x$ into the latent state $g$ and the action $u$ into the latent action $a$. For the state $x$ (i.e., raw image), it is high-dimensional and redundant, so we want to reduce the state size and use the smaller latent state as a representation. For the action $u$ (i.e., grasping positions, gripper moving length, and gripper moving angle), the raw action usually has a very non-linear effect on the state, so we lift the actions to a higher dimension and make them easier to linearize in the latent space.

Decoder Model. The state decoder model is a map $\phi_{d}:\mathcal{G}\to\mathcal{X}$, the action decoder model is a map $\varphi_{d}:\mathcal{A}\to\mathcal{U}$. We then have

Paired with encoder models, decoder models decode the latent state $g$ into the reconstructed state $\hat{x}$ and the latent action $a$ into the reconstructed action $\hat{u}$.

Dynamics Model. The dynamics model contains one map $\psi_{s}:\mathcal{X}\to\mathcal{K}$ for the state matrix and the other map $\psi_{c}:\mathcal{X},\mathcal{U}\to\mathcal{L}$ for the control matrix. We then have

We consider the prediction framework in Figure 1. It consists of the state autoencoder model, action autoencoder model, and dynamics model.

When predicting, given a new state $x(t)$ and a new action $u(t)$, we can get latent state $g(t)$, latent action $a(t)$, state matrix $K(t)$, control matrix $L(t)$. Using locally linear latent dynamics, we can have predicted next latent state $g^{pred}(t+1)=K(t)g(t)+L(t)a(t)$. By decoding the predicted next latent state $g^{pred}(t+1)$, we will have the predicted next state $x^{pred}(t+1)$.

Loss Function. We consider minimizing the losses for state autoencoder model, action autoencoder model, and dynamics model. The loss of state autoencoder model is

The state loss $\mathcal{L}_{state}$ makes sure the reconstructed state $\hat{x}$ is similar to ground truth state $x$. The loss of action autoencoder model is

The action loss $\mathcal{L}_{action}$ makes sure the reconstructed action $\hat{u}$ is similar to ground truth action $u$. The loss of dynamics model is

The dynamics loss $\mathcal{L}_{dyn}$ sets a locally linear constraint on latent dynamics. The loss of prediction model is

The prediction loss $\mathcal{L}_{pred}$ makes sure the predicted state $x^{pred}$ is similar to ground truth state $x$. The expansions of Equation (4) and Equation (5) are in Appendix V-C. The overall training loss is

where $\lambda_{1}$, $\lambda_{2}$, and $\lambda_{3}$ are coefficients.

Training Prediction Framework. Figure 2 shows 4 steps on how to train the prediction framework $\mathcal{M}$.

Training Autoencoders. In Figure 2 Step 1, we use Equation (2) and Equation (3) to train the autoencoders to get both latent states $g$ and latent actions $a$. When state loss $\mathcal{L}_{state}<\epsilon_{state}$ and action loss $\mathcal{L}_{action}<\epsilon_{action}$, we save latent state $g$ and latent action $a$ in each time step. Here, $\epsilon_{state}$ and $\epsilon_{action}$ are very small values. We make sure state $x$ and action $u$ are similar to reconstructed state $\hat{x}$ and reconstructed action $\hat{u}$, respectively. We also do early stopping to prevent the autoencoders from overfitting to the training set.

Training Dynamics Model. The state matrix $K$ depends on state $x$ and the control matrix $L$ depends on both state $x$ and action $u$. In Figure 2 Step 2, we train dynamics model to minimize the dynamics loss $\mathcal{L}_{dyn}$ in Equation (4) and prediction loss $\mathcal{L}_{pred}$ in Equation (5).

There are two methods for model training: (i) training autoencoder models first and then training dynamics model, and (ii) training autoencoder models and dynamics model together. We discuss the model training order and our preference in Section IV ablation study.

To wrap up, after training the autoencoders and the dynamics model using the training set, we start to make predictions for new states and actions in the test set. The detailed architectures and the hyperparameters for the autoencoders and the dynamics model are in the Appendix V-D and Appendix V-A.

We consider using cross-entropy method (CEM) [31] to select optimal action given trained prediction framework $\mathcal{M}$. CEM is an optimization algorithm. It works by repeating two phases: (i) providing a probability distribution and sampling from this distribution and (ii) minimizing the cross-entropy between the current distribution and a target distribution for better sampling in the next iteration. Our Algorithm 1 is an extension of CEM. For each time step, given trained prediction framework $\mathcal{M}$, initial image $x_{init}$, and goal image $x_{goal}$, we first initialize a uniform distribution $Q_{pre}$. Then, we sample $m$ action sequences from distribution $Q_{pre}$. Each action sequence has length $h$, meaning that there are $h$ actions in this sequence. We use the prediction framework $\mathcal{M}$ to predict the future states $x^{pred}$ given initial image $x_{init}$ and each action sequence. Assuming we can get a predicted goal image $x_{goal}^{pred}$ after taking $h$ actions, we then calculate the Binary Cross Entropy (BCE) loss between the predicted goal images $x_{goal}^{pred}$ and goal image $x_{goal}$ we want. By sorting $m$ BCE losses, we choose $n~{}(n<m)$ action sequences that have lowest losses. We apply multivariate normal distribution $Q_{post}$ to fit $n$ samples. If the KL divergence between $Q_{pre}$ and $Q_{post}$ is not less than a specified small value, we will repeat the same process until the distribution $Q$ converges. After distribution convergence, we execute the optimal action with the lowest loss.

We use Rope Manipulation Dataset from paper [26]. This dataset is challenging because (i) the rope is deformable and we do not know the rope dynamics beforehand, (ii) the images are not taken straight down and there is a less-than-ninety-degree angle between the camera and the table, and (iii) the consecutive actions are not applied on the same grasping position. These challenges combine to make it difficult to predict the rope state accurately at each time step.

State. Different from [26] which uses the whole image as an input, we merely want to get the rope-related information in the state. Thus, we do color space segmentation and add a black mask (Figure 2, Step 1) to replace the noisy background. We resize the image from $240\times 240$ to $50\times 50$ and apply data augmentation consisting of translation, vertical flips, and horizontal flips.

Action. The original action $u=(x,y,l,\theta,\mathbbm{1})$ includes $x$ and $y$ positions in the image space, moving length $l$ and moving angle $\theta$ in the world space, and an indicator $\mathbbm{1}$ that decides whether this action validly moves the rope or not. Situations for the invalid actions include the robotic gripper not contacting the rope and not moving the rope according to the command. We re-calculate the action $x$ and $y$ positions according to the change of image size and filter the invalid actions when training and testing.

One-step Prediction. In Figure 2 Step 3, we use Equation (1) to get predicted next latent state $g_{pred}$ based on the current state matrix $K$, the current latent state $g$, the current control matrix $L$, and the current latent action $a$. In Figure 2 Step 4, we use the same decoder in the state autoencoder and decode the predicted next latent state $g_{pred}$ into predicted next state $x_{pred}$.

Multi-step Prediction. In order to test our framework’s prediction horizon, we use the trained prediction model for multi-step prediction. Given state $x(0)$, action $u(0)$, and trained model $\mathcal{M}$, we can predict next state $x^{pred}(1)$. Then using predicted state $x^{pred}(1)$ and ground truth action $u(1)$, we predict the next state $x^{pred}(2)$. By iterating the same process, we get ten predicted states shown in Figure 3.

Comparison. We compare our results with Ground Truth and results using Embedded to Control (E2C) [8]. E2C is a method for learning and control of non-linear dynamics from images. Unlike E2C using optimal control formulation in the latent space, we use the dynamics model to get the state matrix and the control matrix for state prediction in the latent space. To make a fair comparison, we use the same neural network architecture for autoencoders and the same latent size for the state. When running the E2C on the Rope Manipulation Dataset, we find the reconstruction and prediction results are blurry, and we cannot observe a clear rope shape from the image. Thus, we add a prediction loss $\mathcal{L}_{pred}$ and a state loss $\mathcal{L}_{state}$ on the original loss of variational bound and KL divergence. We name this new approach E2C++. Figure 3 (a) and (b) show that our method can predict the future states well in ten steps, and the rope shape is similar to the ground truth. For Figure 3 (c), there is a clear distinction between GT and our method at time step 8 because the previous states (time step 1 to 7) gradually accumulate the state prediction error. E2C++ can only predict a blurry rope shape within two-time steps. The remaining state predictions using E2C++ do not make sense.

Qualitative Results. We use the sampling-based optimization method, shown in Algorithm 1, to plan rope manipulation. The objective is to generate an optimal action sequence. Similar to the standard cross-entropy method (CEM), our algorithm stops repeating in the inner loop when the probability distribution $Q$ converges. However, unlike repeating two phases for the standard CEM, our algorithm needs to do more steps of sampling action sequences, predicting the goal state, calculating and sorting the loss, and fitting a new distribution. In the experiments, we first sample the optimal action by applying CEM. Then we use the trained prediction framework to predict the next state $x^{pred}(t+1)$ given the current state $x(t)$ and the sampled optimal action.

Figure 4 shows the experimental results. We overlay sampled action on the rope and compare it with the ground truth action. By comparing the Ground Truth (GT) action with sampled action using CEM, we can find sampled action is close to GT action concerning grasping positions, moving length, and moving angle. However, sometimes there are many possible actions to move the rope from one state to the next state. For example, Figure 4 (c) shows the different grasping positions for sampled action, but the rope ends up with a similar predicted next state as the GT state.

Quantitative Results. We calculate the Mean Square Error (MSE) between sampled actions and ground truth actions. We also apply MSE between the predicted next state and ground truth next state. After running 100 sampled one-step actions using our method and E2C++, we get the mean and standard deviation in Table I. By comparing Action Score, we can observe that our score is lower. It means the sampled action is closer to the ground truth action when using our prediction model. From State Score, we can find our method has a lower score. It means our method is better at predicting the next state after applying the sampled action. Detailed methods to get the quantitative results are mentioned in Appendix V-B.

Locally Latent Dynamics v.s. Globally Latent Dynamics. For locally latent dynamics, state matrix $K(t)$ and control matrix $L(t)$ are different at each time step. However, $K$ and $L$ are constant matrices for globally latent dynamics. If $K$ and $L$ are constant matrices, it will set a more strict constraint in the latent space. It is difficult to train a good-performance model with globally linear latent dynamics since we need to guarantee the dynamics are linear for all latent states. In our experiments, we use locally linear latent dynamics.

Order for Training Models. We have autoencoder models for state and action and the dynamics model. The order to train models is important to get good performance. We use two methods: (i) training autoencoder models first and then training dynamics model, (ii) training autoencoder models and dynamics model together. We demonstrate that method (i) can provide better prediction and control results, as shown in Figure 3, while method (ii) reconstructs blurry rope images, and we cannot find a clear rope shape. The reason is that we can get a good latent state estimation when training autoencoder models. After having a good latent state estimation, we then train the dynamics model to get good system dynamics. For method (ii), training all models together to optimize all losses is more difficult to get a good prediction result. It is because we need to find the optimal neural network weights to get accurate latent state, latent action, state matrix, and control matrix at the same time.

Latent State and Latent Action Size. Since the dimension of state $x$ is $50\times 50$, we need to shrink its size to a smaller dimension for latent state $g$. Different from the state, we use the autoencoder model to expand the action dimension from $(x,y,l,\theta)$ to a larger size. We do experiments to find how small we can get for the latent size to achieve good performance. Experiments show that the latent state size cannot be less than 50, and latent action size cannot be less than 10. When the latent state and latent action size are too small, there is no enough information about the real state and action. Besides, their sizes do not need to be too large (e.g., latent state size is 300+ and latent action size is 100+). One reason is that it increases more time to train since we have state matrix $K$ and control matrix $L$ as well. Another reason is that the larger size does not improve the prediction performance much. Thus, we choose 80 and 80 for latent state size and latent action size, respectively.

Whether Using Action Autoencoder. In the prediction framework, we use the autoencoder to lift the action space to a higher dimensional space so it can work linearly with the latent state in the latent space. We also try not to use the autoencoder in the prediction framework, and we cannot achieve good prediction results due to the non-linearity of the impact of the four-element action.

Our experimental results are based on the hyper-parameters in Table II.

We calculate the Mean Square Error (MSE) between the sampled actions and the ground truth actions.

where $u_{\text{sample}}$ is the sampled action and $u_{\text{GT}}$ is the ground truth action. $n=4$ since there are four elements $(x,y,l,\theta)$ in the action $u$.

We also apply MSE between predicted next state and ground truth next state.

where $x_{\text{pred}}$ is the predicted next state and $x_{\text{GT}}$ is the ground truth state. $n=2500$ since the state is a $50\times 50$ image. To get the MSE for the state, we first do pixel-wise subtraction between $x_{\text{sample}}$ and $x_{\text{GT}}$. Then we square each pixel value in the image matrix. After that, we sum up all the pixel values and divide it by the image size $n$.

After applying different sampled one-step actions on different states and running it 100 times based on our method and E2C++, we get the mean and standard deviation for the State Score and the Action Score.

The expansion of the dynamic loss is

The expansion of the prediction loss is

Our experimental results are based on the following neural networks. Figure 5 includes the input, output, and details of each layer for both the state encoder and the state decoder. Figure 6 introduces the action encoder and the action decoder. In the action decoder, the Multiplication and Addition layers make sure that each element (position $x$, position $y$, moving length $l$, moving angle $\theta$) in the action are within certain range. We multiply tensor $[50,50,0.14,2\pi]$ to the output from the Sigmoid layer and then add another tensor $[0,0,0.01,0]$. The position $x\in[0,50]$, the position $y\in[0,50]$, the moving length $l\in[0.01,0.15]$, and the moving angle $\theta\in[0,2\pi]$. Figure 7 includes two mappings. One mapping is from the state $x$ to the state matrix $K$. The other mapping is from the state $x$ and the action $u$ to the control matrix $L$.