V ARIATIONAL M ODEL -B ASED I MITATION L EARNING IN H IGH -D IMENSIONAL O BSERVATION S PACES
A BSTRACT
We consider the problem setting of imitation learning where the agent is provided a ﬁxed dataset of demonstrations.
While the agent can interact with the environment for exploration, it is oblivious to the reward function used by the demonstrator.
This setting is representative of many applications in robotics where task demonstrations may be straightforward while reward shaping or conveying stylistic aspects of human motion may be difﬁcult.
For this setting, we develop a variational model-based imitation learning algorithm (VMIL) that is capable of learning policies from visual observations.
Through experiments, we ﬁnd that VMIL is more sample efﬁcient compared to prior algorithms in several challenging visionbased locomotion and manipulation tasks, including a high-dimensional in-hand dexterous manipulation task.
All results including videos can be found online at https://sites.google.com/view/vim .
I NTRODUCTION
The ability of reinforcement learning (RL) agents to autonomously learn by interacting with environment presents a promising approach for learning diverse skills.
However, reward speciﬁcation has remained a major challenge in deployment of RL in practical settings (Amodei et al., 2016; Everitt & Hutter, 2019; Rajeswaran* et al., 2018).
The ability to imitate humans or other expert trajectories enables us to avoid the reward speciﬁcation problem while also circumventing hard exploration challenges in RL.
It also presents a more natural way to teach robots various tasks and skills.
In this work, we aim to develop a new algorithm that can learn from limited expert demonstration data and scale to high-dimensional observation and action spaces.
Behaviour cloning (BC) is a classic algorithm to imitate expert demonstrations (Pomerleau, 1988).
It uses supervised learning to greedily match the expert behaviour at demonstrated expert states, without regard for the marginal state distribution.
However, due to environment stochasticity or compounding actor errors over time, the agent will drift from the expert state distribution towards out-of-distribution states and ultimately fail to mimic the demonstrations (Ross et al., 2011).
This problem is especially prevalent in highdimensional domains or domains with limited data.
While a wide policy coverage by the expert (Spencer et al., 2021) or the ability to repeatedly and interactively query the expert policy (Ross et al., 2011) can circumvent difﬁculties to an extent, such conditions are seldom met in practical applications.
An alternate line of work based on inverse RL (IRL) (Finn et al., 2016b; Fu et al., 2018a) and adversarial imitation learning (Ho & Ermon, 2016) aim to not only match actions at demonstrated states, but also the long term visitation distribution (Ghasemipour et al., 2019).
These approaches explicitly train a GAN-based classiﬁer (Goodfellow et al., 2014; Finn et al., 2016a) to estimate the expert state-action distribution, which serves as a reward for training an RL agent.
While these methods have achieved substantial improvement over behaviour cloning without additional expert supervision, they are difﬁcult to deploy in realistic scenarios.
This is mainly due to three reasons: (1) the need to optimize an RL objective that requires on-policy data collection and contributes towards sample complexity; (2) non-stationarity of the reward function that changes as the RL agent learns; and (3) high-dimensionality of the observation space that requires representation learning.
To overcome the aforementioned difﬁculties, we develop a new model-based adversarial imitation learning algorithm that trains a variational dynamics model.
Our approach has several advantages: (1) model learning serves as a rich auxiliary task for learning visual representations that support efﬁcient learning of control policies in a stable fashion, (2) we optimize the correct (on-policy) imitation learning objective, (3) model-based policy optimization allows us to train an agent onpolicy without the sample complexity required by these algorithms.
On four simulated environments with image observations and continuous actions, including a vision-based dexterous manipulation task, we ﬁnd that our algorithm is signiﬁcantly more performant than a prior state-of-the-art modelfree method.
R ELATED W ORK
We begin by reviewing the literature on imitation learning and recent advances in image-based RL.
Imitation Learning.
Ho & Ermon (2016) develop GAIL - a practical imitation learning algorithm with a ﬁxed amount of demonstration data, without the ability to further query the expert.
The key component is to train a binary classiﬁer D ψ : S ×
A → [0 , 1] that discriminates between expert and policy rollouts similar to a GAN.
The agent’s policy π is then trained using reinforcement learning to maximize E τ ∼ π
[ (cid:80) t log D ψ ( s, a )] − λ H ( π ) , where H ( π ) is an entropy regularizer.
Fu et al.
(2018a) develop a similar algorithm, but like GAIL it still requires on-policy data.
Several off-policy algorithms have been developed (Reddy et al., 2020; Blond ´ e & Kalousis, 2019; Kostrikov et al., 2019) to tackle sample inefﬁciency, largely by replacing an on-policy RL algorithm (TRPO, PPO) with an off-policy one.
However, this introduces a distribution mismatch as it replaces the expectation under the policy state-action marginal distribution with the expectation under the distribution of the replay buffer.
While this tends to work in practice, Kostrikov et al.
(2020a) show that the distribution mismatch can signiﬁcantly degrade performance if the buffer distribution is signiﬁcantly off from the policy distribution.
Alternatively Finn et al.
(2016b) use a similar approach to reward learning in combination with a learned locally linear dynamics model, which makes policy optimization easier.
In this work we extend this approach to high-dimensional POMDPs with arbitrary dynamics.
Learning From Images.
Reinforcement learning from images is an inherently difﬁcult task, since the agent needs to learn meaningful visual representations from reward signal alone.
A recent line of research (Gelada et al.
; Hafner et al., 2019; Lee et al., 2020; Hafner et al., 2020) train a variational model of the image-based environment as an auxiliary task.
Lee et al.
(2020) use the model for representation learning purposes only and train an off-policy SAC-based algorithm on top of the latent representation.
On the other hand, Hafner et al.
(2020) use rollouts from the learned model to train an agent in an on-policy fashion within the latent space of the model.
We base our algorithm on this approach, as it allows us to train an agent on-policy and target the correct objective of imitation learning, while still maintaining the sample efﬁciency of off-policy approaches.
Separately
Kostrikov et al.
(2020b) introduce DrQ, a sample efﬁcient approach for off-policy Q-learning from images in a completely model-free way by introducing image-augmentation for both the target and actor Q-networks.
The authors report state of the art performance on both asymptotic returns and sample efﬁciency.
We later empirically compare our method to SQIL (Reddy et al.
(2020)), an established off-policy imitation learning algorithm combined with DrQ image augmentation.
M ODEL -B ASED I MITATION L EARNING IN POMDP S
In this section, we consider learning in image-based domains as a POMDP.
3. M ODELING VIA CONTROL AS INFERENCE IN POMDP S
Following the control as inference approach (Levine, 2018; Lee et al., 2020), we jointly model the POMDP dynamics and policy learning.
We choose choose a joint variational distribution that factors into state inference, latent dynamics, and a policy term.
However, unlike prior works (Lee et al., 2020), we base the policy distribution entirely on the inferred latent state.
This later allows us to optimize the policy in a model-based fashion under the learned latent model.
Our joint variational distribution is as follows:
[Equation]
1: T | x 1: τ +1 , a 1: τ ) =
τ (cid:89) t =
q ( z t +1 | x t +1 , z t , a t )
T − 1 (cid:89) t = τ + p ( z t +1 | z t , a t )
T − 1 (cid:89) t = τ + π ( a t | z t )
Using this variational distribution, the evidence lower bound becomes:
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
All the distributions are represented as diagonal Gaussians with the mean and variance parameterized by the output of neural networks.
3. P RACTICAL A LGORITHM AND T RAINING
We optimize the two different components of Equation 1 in an alternating fashion as shown in
Figure 1.
We utilize the latent dynamics model by Hafner et al.
(2019), with the exception that we do not include a reward predictor.
Following Fu et al.
(2018b) we can consider log p ( O t = 1 |
z t , a t ) as the probability that the state-action pair belongs to an expert ”event” distribution.
We train a discriminator on top of the latent representation using the following objective:
[Equation]
[Equation]
as we ﬁnd that including the next state helps with alleviating model exploitation.
We also add
Gaussian input noise (M¨uller et al., 2019) as it helps improve convergence speed and stability.
At this point, we optimize the policy objective as a reinforcement learning problem with a reward function r ( z t , a t , z t +1 )
=
D ψ ( z t , a t , z t +1 ) .
We can use several algorithms to optimize this component, but we would like to (1) avoid reinforcement learning with a non-stationary reward and (2) use an on-policy training algorithm in order to target the correct objective.
Because of these goals we opt for a model-based value expansion approach based on Hafner et al.
E XPERIMENTS
In our experiments we want to answer several questions: (1)
can
VMIL successfully scale to environments with image observations, (2) how does VMIL compare to state of the art model-free approaches, (3) how does VMIL perform in different data regimes, (4) can VMIL solve realistic manipulation tasks, (5) can VMIL solve environments with complex interaction and physics?
Tasks.
Our test environments are shown in Figure 3.2.
We ﬁrst evaluate our method on the Cheetah and Walker tasks from the DeepMind Control Suite (Tassa et al., 2018).
We also consider a
Sawyer peg-insertion task, where the arm and peg positions are randomized, which creates a difﬁcult conﬁguration for imitation learning when the conﬁguration is out of the expert distribution.
Our ﬁnal environment is the Baoding balls task from Nagabandi et al.
This is a challenging task for policy learning, even in the state-based case.
All observations are from images, while the
Baoding balls task additionally includes robot proprioception.
The DeepMind Control Suite tasks have between 1 and 10 demonstration rollouts, generated using SAC training from state.
The Sawyer task uses 50 demos generated using an analytic controller.
Finally the Baoding balls uses 25 demos generated using the method from Nagabandi et al.
Results.
Experiment results are shown in Figure 3.
To answer questions (1) and (2) above we compare VMIL to SQIL with DrQ image augmentations, as introduced in Section 2, on the Cheetah and Walker tasks.
While the model-free baseline initially outperforms our method, it’s performance has high variance between random seeds and exhibits signiﬁcant instability.
We hypothesize that this is due to (1) sparse rewards associated with imitation learning, which prevent the development of reliable image features and (2) the static nature of the reward in the SQIL method, as only expert demonstrations receive positive reward.
The addition of a discriminator-based training, similar
Kostrikov et al.
(2019); Blond ´ e & Kalousis (2019), might help alleviate these issues, however that would introduce the problem of representation learning using non-stationary reward signal.
To answer question (3), we vary the number of expert trajectories available to our method, results shown in the ﬁrst two plots of Figure 4.
We ﬁnd that performance is comparable and we still get stable convergence with only a single expert rollout.
To answer question (4), we evaluate our method on a
Sawyer peg-insertion tasks, where it reaches over 50% success rate, while the model-free baseline fails to complete the task.
Towards the end of training we see meaningful degradation in performance.
We hypothesize that this is due to the evolution of the latent space representation in the course of training.
A potential solution would be annealing the amount of input noise to the discriminator or exploring alternative stabilization techniques, (Blond ´ e et al.
To answer question (5) we deploy our method on the Baoding Balls task from Nagabandi et al.
This is a complex task, which is difﬁcult to solve with policy training, even in the state-based domain.
We manage to reach up to 14% success rate, while the model-free baseline fails to solve the task, however we encounter similar degradation in performance, as in the peg-insertion task.
C ONCLUSION
In this work we presented VMIL, a model-based imitation learning algorithm that works from highdimensional observations, such as images.
Our method achieves better asymptotic returns and is more stable and sample efﬁcient than state-of-the art off-policy model-free approaches in the same domain.
The improved sample efﬁciency and safety of the learning procedure make real-world imitation learning feasible
and we plan to evaluate our algorithm on a real robot setting.
Although we observe some degradation in performance on environments with out-of-distribution states and more complex dynamics, we believe that by utilizing additional stabilization and reward shaping techniques we can fully solve those cases as well.
Finally the use of task-agnostic dynamics models opens many doors to future research, such as ofﬂine, multi-task, continual and meta-imitation learning.