Model-Based Robot Imitation with Future Image Similarity

Abstract

We present a visual imitation learning framework that enables learning of robot action policies solely based on expert samples without any robot trials. Robot exploration and on-policy trials in a real-world environment could often be expensive/dangerous. We present a new approach to address this problem by learning a future scene prediction model solely from a collection of expert trajectories consisting of unlabeled example videos and actions, and by enabling action selection using future image similarity. In this approach, the robot learns to visually imagine the consequences of taking an action, and obtains the policy by evaluating how similar the predicted future image is to an expert sample. We develop an action-conditioned convolutional autoencoder, and present how we take advantage of future images for zero-online-trial imitation learning. We conduct experiments in simulated and real-life environments using a ground mobility robot with and without obstacles in reaching target objects. We explicitly compare our models to multiple baseline methods requiring only offline samples. The results confirm that our proposed methods perform superior to previous methods, including 1.5 $\times$ and 2.5 $\times$ higher success rate in two different tasks than behavioral cloning.

Appendix

1.1 Implementation Details

We implemented the CNN models using the PyTorch library. The encoder/decoder networks followed the architecture of DCGAN (Radford et al. 2015), using their discriminator as our encoder CNN and their generator as our decoder CNN. Specifically, the encoder has 6 convolutional layers with a $3\times 3$ kernel and stride of 2. The network layers have 64, 128, 256, 512, 512, 128 channels. Our input images are resized to $64\times 64$, resulting in a feature map of size $128\times 3\times 3$. For the linear-representation model shown in Fig. 6a, we reshape this to be a vector of size $128\cdot 3\cdot 3$ then use a fully-connected layer to reduce the dimensionality to 4096. Our action network has two layers to increase the dimensionality to 64 then 256.

In the convolutional-representation model used in Fig. 6b, we leave the representation as-is. Our actions are 3-dimensional vectors for robot pose ($x,y,\theta$), which are used as input to the action network. The action network has two layers that produces a 576-dimensional vector which we reshape to a spatial tensor of size $64\times 3\times 3$. We concatenate this tensor along the channel axis of the convolutional representation, which is then used as input to the decoder. The convolutional future prediction model contains 5 convolution layers with a $3\times 3$ kernel and a stride of 1. The layers contain 256, 512, 512, 256, 128 channels.

Our decoder contains 6 deconvolutional layers for upsampling. All have a $3\times 3$ kernel and a stride of 2. In the deconvolutional layer, a stride of 2 doubles in output size. The layers contain 512, 512, 256, 128, 64, 3 channels. The last layer is followed by a $\tanh$ activation function. All other layers in all networks were followed by batch normalization and used the LeakyReLU activation function with the negative slope set to 0.2. We minimize our loss function with gradient descent using the Kingma and Ba (2014) solver and learning rate set to 0.001.

The LSTMs are implemented similar to Denton and Fergus (2018). $LSTM_{\phi }$ and $LSTM_{\psi }$ are both single layer LSTMs with 256 cells in each layer. Each network has a linear embedding layer and a fully connected output layer. At inference, the output of $LSTM_{\psi }$ is concatenated to $z_I$ and $z_a$, and fed to the decoder. The output dimensionalities of the LSTM networks are g = 128 and $\mu _{\phi } = \mu _{\psi } = 64$.

1.2 Training Information

Our training curves for the image predictor model and the critic are shown in Fig. 15. For the image predictor of both datasets, we set the learning rate = 0.001 and batch size = 60. The $\beta$ multiplier for the KL loss was set to 0.0001 in our experiments. The learning rate of the value function was set to 5E−6. The weights of the image predictor were held constant when training the value function.

Fig. 15

Training curves for the image predictor for the simulated and lab datasets with learning rate = 0.001 and batch size = 60, and the value function with learning rate = 5E−6