Trajectory VAE for multi-modal imitation

1 INTRODUCTION . Imitation learning has become successful in a wide range of sequential decision making problems , in which the goal is to mimic expert behaviour given demonstrations ( Ziebart et al. , 2008 ; Wang et al. , 2017 ; Li et al. , 2017 ; D ’ Este et al. , 2003 ) . Compared with reinforcement learning , imitation learning does not require access to a reward function – a key advantage in domains where rewards are not naturally or easily obtained . Instead , the agent learns a behavioural policy implicitly through demonstrated trajectories . Expert demonstrations are typically assumed to be provided by a human demonstrator and generally can vary from person to person , e.g. , according to their personality , experience and skill at the task . Therefore , when capturing demonstrations from multiple humans , observed behaviours may be distinctly different due to multi-modal structure caused by differences between demonstrators . Variations like these , which are very common in video games where players often cluster into distinct play styles , are typically not modelled explicitly as the structure of these differences is not known a priori but instead emerge over time as part of the changing meta-game . In this paper , we propose Trajectory Variational Autoencoder ( T-VAE ) a deep generative model that learns a structured representation of the latent features of human demonstrations that result in diverse behaviour , enabling the imitation of different types of emergent behaviour . In particular , we use a Variational Autoencoder ( VAE ) to maximise the Evidence Lower Bound ( ELBO ) of the log likelihood of the expert demonstrations on the trajectory level where the policy is directly learned from optimising the ELBO . Not only can our model reconstruct expert demonstrations , but we empirically demonstrate it learns a meaningful latent representation of distinct emergent variances in the observed trajectories . 2 RELATED WORK . Popular imitation learning methods include behavior cloning ( BC ) ( Pomerleau , 1991 ) , which is a supervised learning method that learns a policy from expert demonstration of state-action pairs . However , this approach assumes independent observations which is not the case for sequential decision making problems , as future observations depend on previous actions . It has been shown that BC can not generalise well to unseen observations ( Ross & Bagnell , 2010 ) . Ross et al . ( 2011 ) proposed a new iterative algorithm , which trains a stationary deterministic policy with no regret learning in an online setting to overcome this issue . Torabi et al . ( 2018 ) also improve behaviour cloning with a two-phase approach where the agent first learns an inverse dynamics model via interacting with the environment in a self-supervised fashion , and then use the model to infer missing actions given expert demonstrations . An alternative approach is Apprenticeship Learning ( AC ) ( Abbeel & Ng , 2004 ) , which uses inverse reinforcement learning to infer a reward function from expert trajectories . However , it suffers from expensive computation due to the requirement of repeatedly performing reinforcement learning from tabula-rasa to convergence . Whilst each of these methods has had successful applications , none are able to capture multi-modal structure in the demonstration data representative of underlying emergent differences in playstyle . More recently , the learning of a latent space for imitation learning has been studied in the literature . Generative Adversarial Imitation Learning ( GAIL ) ( Ho & Ermon , 2016 ) learns a latent space of demonstrations with a Generative Adverserial Network ( GAN ) ( Goodfellow et al. , 2014 ) like approach which is inherently mode-seeking and does not explicitly model multi-modal structure in the demonstrations . This limitation was addressed by ( Li et al. , 2017 ) , who built on the GAIL framework to infer a latent structure of expert demonstrations enabling imitation of diverse behaviours . Similarly , ( Wang et al. , 2017 ) combined a VAE with a GAN architecture to imitate diverse behaviours . However , these methods require interacting with the environment and rollouts of the policy whilst learning . For comparison we note our method does not need access to the environment simulator during training and is computationally cheaper , as the policy is learned simply by gradient descent using a fixed dataset of trajectories . Additionally , whilst the aim in GAIL is to keep the agent behaviour close to the expert ’ s state distribution , our model can serve as an alternative approach to capturing state sequence structure . In work more closely related to our approach , ( Co-Reyes et al. , 2018 ) have also proposed a Variational Auto encoder ( VAE ) ( Kingma & Welling , 2013 ) that embeds the expert demonstration on the trajectory level which showed promising results . However their approach only encodes the trajectories of the states whereas ours encodes both the state and action trajectories , which also allows us to learn the policy directly from the probabilistic model rather than adding a penalty term to the ELBO . Rabinowitz et al . ( 2018 ) also learns an interpretable representation of the latent space in a hierarchical way , but their focus is more on representing the mental states of other agents and is different from our goal of imitating diverse emergent behaviours . 3 METHODS . 3.1 PRELIMINARIES . Let the tuple ( S , A , P , r , I ) denote the infinite-horizon Markov Decision Processes ( MDP ) with : S the state space , A the action space , P the transition probability distribution , r the reward function and I the distribution of the initial state s0 . Let πE : S × A −→ [ 0 , 1 ] denote the expert policy which we do not know , under which expert trajectories τ of states and actions are generated from , i.e. , s0 ∼ I , at ∼ πE ( at|st ) , st+1 ∼ P ( st+1|at , st ) . The goal of imitation learning is to learn a policy π that best explains the trajectories without knowledge of the reward signal r . 3.2 TRAJECTORY VAE ( T-VAE ) . Given N demonstrated trajectories { τ ( i ) } Ni=1 of states and actions , where each τ ( i ) = { ( s ( i ) t , a ( i ) t ) } Ti t=1 , where Ti is the length of trajectory τ ( i ) . The marginal likelihood of the set of trajectories is composed of a sum over the marginal likelihoods of each individual trajectory log pθ ( τ ( 1 ) , · · · , τ ( N ) ) = ∑N i=1 log pθ ( τ ( i ) ) . We use a latent variable model and assume the prior distribution of latent z is standard Gaussian . Rather than a VAE which is applied on the data point level , we use a VAE on the trajectory level ( which consists of a time sequence of data points ) , as shown in Figure 1a . We call our model Trajectory VAE ( T-VAE ) . ( a ) Graphical model of T-VAE ( b ) Computation graph of T-VAE 3.2.1 ENCODER NETWORK . We encode whole trajectories into the latent space in order to embed useful features of different behaviours and extract distinguishing features which differ from trajectory to trajectory . Note that the latent z is therefore a single variable rather than a sequence that depends on t. In order to utilise all information , we encode both the states and actions , i.e. , qφ ( z|τ ( i ) ) = qφ ( z| { ( s ( i ) t , a ( i ) t ) } Ti t=1 ) . We assume that the approximate posterior qφ ( z|τ ( i ) ) has a Gaussian distribution , whose mean and log variance parameters are constructed as follows ( and illustrated in the top half of Figure 1b ) : 1 ) concatenate states with actions at each time step [ st , at ] ; 2 ) feed the sequence into a bidirectional LSTM ; 3 ) mean-pool over the outputs along the time horizon ( T ( i ) ) ; 4 ) pass through two separate fully connected layers . 3.2.2 DECODER NETWORK . The decoder p ( τ ( i ) |z ) can be decomposed as pθ ( τ ( i ) |z ) = pθSD ( { s ( i ) t } Tt=1|z ) pθPD ( { a ( i ) t } Tt=1|z , { s ( i ) t } Tt=1 ) ( 1 ) where we call pθSD ( { s ( i ) t } Tt=1|z ) the state decoder and pθPD ( { a ( i ) t } Tt=1|z , { s ( i ) t } Tt=1 ) the policy decoder . Instead of having a separate policy decoder and control it to be consistent with the state decoder ( as proposed by ( Co-Reyes et al. , 2018 ) ) , T-VAE models the policy and state decoder jointly and enables consistency inherently . Note that the state decoder does not depend on the policy and therefore during training , we do not need to interact with the environment nor conduct rollouts of the policy , making the learning process simpler and computationally relatively cheap . For the state decoder , we assume Gaussian distribution on each s ( i ) t where the variance is fixed and the mean ŝ ( i ) t is computed recursively . At time t < = T ( i ) , we concatenate ŝ ( i ) t with the latent z to form [ ŝ ( i ) t , z ] which is fed into a LSTM cell , the output is then fed into a fully connected layer to produce ŝ ( i ) t+1 . The fully connected layers guarantee that the dimensionality is preserved ( ŝt has the same dimension as st ) . For the action decoder , we assume Gaussian distribution over at for continuous actions and Multinomial/Bernoulli distributions for discrete actions . The variational parameter to be learned â ( i ) t+1 is therefore the mean and the logits vector in the two cases respectively . Similarly as the state decoder , â ( i ) t+1 is generated recursively from [ â ( i ) t , ŝ ( i ) t , z ] which is fed into a LSTM followed by a fully connected layer . Continuous actions are output at this stage , or an additional softmax/sigmoid activation function is applied to the output to generate discrete actions . If the action space consists of a mixture of continuous and discrete actions , we assume the actions are independent conditional on the states and latent variable , and the policy decoder can be factored as the product . An illustration of the entire model can be found in figure 1b with the bottom half representative of the decoder network . 3.2.3 VARIATIONAL BOUND . The marginal likelihood for each trajectory can be written as log pθ ( τ ( i ) ) = DKL ( qφ ( z|τ ( i ) ) ||pθ ( z|τ ( i ) ) ) + L ( θ , φ ; τ ( i ) ) ( 2 ) where DKL represents the KL divergence between the approximate posterior and the true posterior , andL ( θ , φ ; τ ( i ) ) is the variational lower bound of the marginal likelihood of τ ( i ) which is decomposed into 3 terms : the L2 reconstruction loss for the state decoder , the L2 or cross entropy/sigmoid reconstruction loss for the policy decoder and a KL divergence between the posterior and prior distribution of the latent variable z . Formally : L ( θ , φ ; τ ( i ) ) = Eqφ ( z|τ ( i ) ) [ − log qφ ( z|τ ( i ) ) + log pθ ( τ ( i ) , z ) ] ( 3 ) = Eqφ ( z|τ ( i ) ) [ log pθ ( τ ( i ) |z ) ] −DKL ( qφ ( z|τ ( i ) ) |pθ ( z ) ) ( 4 ) = Eqφ ( z|τ ( i ) ) [ log pθ ( { s ( i ) t } Tt=1|z ) + log pθ ( { a ( i ) t } Tt=1| { s ( i ) t } Tt=1 , z ) ] −DKL ( qφ ( z|τ ( i ) ) |pθ ( z ) ) ( 5 ) As log pθ ( τ ( i ) ) ≥ L ( θ , φ ; τ ( i ) ) , the encoder and decoder network parameters can then be optimised with stochastic gradient descent .

Summary: This paper presents an approach to multi-modal imitation learning by using a variational auto-encoder to embed demonstrated trajectories into a structured latent space that captures the multi-modal structure. This is done through a stochastic neural network with a bi-directional LSTM and mean pooling architecture that predicts the mean and log-variance of the latent state. This is followed by a state and action/policy decoder (both LSTMs) that recursively generate trajectories from latent space samples. The entire model is trained by optimising the ELBO on a set of pre-specified expert demonstrations. At test time, samples are generated from the latent space and recursively decoded to generate state and action trajectories. The method is tested on three low-dimensional continuous control tasks and is able to learn structured latent spaces capturing the modes in the training data as well as generating good trajectory reconstructions.