Imitation Learning from Pixel Observations for Continuous Control

1 INTRODUCTION . Learning continuous control policies directly from pixel observations is an important problem due to its potential impact on fields like robotics , autonomous driving and video games . These domains have rich resources of data available of humans performing expert-level demonstrations that our software agents do not leverage as they are often trained from scratch without any knowledge of how humans think about these problems . However , using unlabeled video data is challenging as it i ) requires distilling a representation of the world into the policy of an agent , and ii ) we do not know the underlying actions and reasoning process of the expert . This renders common algorithms like canonical behavioral cloning ( Pomerleau , 1988 ; 1991 ) useless in the no-action setting . Recently the community has advanced our understanding in learning visual representations and learning to imitate demonstrations provided as proprioceptive states . Visual representation learning has been crucial in recent advancements for sample-efficient reinforcement learning ( RL ) directly from pixels in continuous spaces , e.g . with reconstruction ( Finn et al. , 2016 ; Yarats et al. , 2019 ) , contrastive learning ( Srinivas et al. , 2020 ; Stooke et al. , 2020 ) , unsupervised pre-training ( Liu & Abbeel , 2021 ; Yarats et al. , 2021b ; Seo et al. , 2021 ) , world models ( Hafner et al. , 2018 ; 2019 ; 2020 ) and data augmentation ( Yarats et al. , 2021c ; Raileanu et al. , 2020 ; Laskin et al. , 2020 ) . These approaches require known reward signals from the environment , which are not always available or well-defined . When expert demonstrations are available , imitation learning ( IL ) and inverse RL ( IRL ) methods overcome the issue of not having a reward signal and seek to recover the expert agent ( Ng & Russell , 2000 ) . These methods are very effective when low-dimensional proprioceptive states and actions are available and typically consist in learning based on the mismatch between the expert and agent ’ s state ( -action ) distributions . In this paper , we combine the budding areas of model-free image-based reinforcement learning and state-action imitation learning to control non-trivial continuous dynamical systems from pixel demonstrations . We extend two leading proprioceptive-state approaches to comparing the pixel trajectories of the learner and expert – see fig . 1 for an illustration . The first approach , pixel sinkhorn imitation learning , P-SIL , extends optimal transport approaches for imitation learning to the image 1Code at https : //anonymous.4open.science/r/ImitateFromPixelsICLR-B0EF setting ( Papagiannis & Li , 2020 ; Dadashi et al. , 2021 ) . A key component here is how we learn and compare latent representations using the cosine distance on a target encoder that is updated with the RL encoder ’ s weights . Our second approach , pixel discriminator actor critic , P-DAC , is a GAIL-based method for pixels ( Ho & Ermon , 2016 ; Kostrikov et al. , 2019 ; Torabi et al. , 2018 ) . We propose several key modifications , including data augmentation and using the RL encoder for representations , that enable P-DAC to scale to non-trivial control tasks without needing access to expert actions . The imitation rewards from these methods are optimized with DrQ-v2 as an underlying image-based RL backbone ( Yarats et al. , 2021a ) . We demonstrate the versatility and strong performance of both approaches on DeepMind control suite tasks , and find that both P-DAC and P-SIL are able to recover expert performance while outperforming canonical extensions of corresponding state-based IL methods , such as SIL ( Papagiannis & Li , 2020 ) and DAC ( Kostrikov et al. , 2019 ; Ho & Ermon , 2016 ) in terms of performance and sample efficiency . Our extensive ablation study reveals insights into the key design decisions for pixel imitation learning . Finally , we also provide an easy-to-use implementation of these methods with unified DrQ-v2 backbone to make it simple for practitioners to build upon in this under-explored area . 2 BACKGROUND AND RELATED WORK . 2.1 REINFORCEMENT LEARNING ( RL ) WITH PIXEL OBSERVATIONS . Pixel reinforcement learning can be instantiated as an infinite-horizon Markov decision process ( MDP ) ( Bellman , 1957 ; Sutton & Barto , 2018 ) , where the agent ’ s state is approximated by a stack of consecutive RGB frames ( Mnih et al. , 2015 ) . The MDP is of the form ( O , A , P , R , γ , d0 ) where O is the observation space , A is the action space , P : O × A → ∆ ( O ) is the transition function ( ∆ ( O ) is a probability distribution over O ) , R : O ×A → R is the reward function , γ ∈ [ 0 , 1 ) is the discount factor and d0 is the initial state distribution . The RL problem consists of finding an optimal policy π : O → ∆ ( A ) that maximizes the expected long-term reward Eπ [ ∑∞ t=0 γ tR ( ot , at ) ] , where o0 ∼ d0 , at ∼ π ( ·|ot ) and ot+1 ∼ P ( ·|ot , at ) . 2.2 IMITATION LEARNING . In imitation learning , agents do not have access to the environment reward R. Instead , they are provided with a dataset of multiple expert trajectories which the agent aims to imitate , where each trajectory is of the form oe = ( oe0 , . . . , o e T ) ∈ OT . More formally , agent trajectories oa = ( oa0 , . . . , o a T ) under the policy π have to be close to expert trajectories o e under some metric between trajectories . We note these trajectories are empirical proxies for the occupancy distributions under the learner and expert policies ( Ho & Ermon , 2016 ; Dadashi et al. , 2021 ) . This paper follows the line of research in imitation learning via agent-expert density matching , initially studied by Ng & Russell ( 2000 ) ; Boularias et al . ( 2011 ) ; Englert et al . ( 2013 ) . Here , imitation-reward signals are derived from trajectories obtained by an expert policy ; the agent is then trained using RL methods , which use these imitation rewards as learning signals . Generative Adversarial Learning for IL P-DAC builds on the budding area of state ( -action ) based generative adversarial imitation learning , originally initiated by Ho & Ermon ( 2016 ) with the GAIL method , which uses adversarial learning combined with RL . A significant body of work built upon GAIL , notably AIRL and WGAIL ( Fu et al. , 2018 ; Xiao et al. , 2019 ; Kostrikov et al. , 2020 ) , which propose alternative losses relying on other probabilistic divergences . Also , Torabi et al . ( 2018 ) extend GAIL to the no-action setting , a setting which we consider in this paper as well . DAC ( Kostrikov et al. , 2019 ) significantly improves upon GAIL by turning it into an offline algorithm , and adding a gradient penalty . Our P-DAC approach directly builds upon DAC , extending it to the visual setting . A key difference , which allowed us to stabilize visual DAC , was to store imitation rewards in the replay buffer and to not recompute them with up-to-date discriminators and encoders . Also , while DAC parameterizes the discriminator directly on the ( proprioceptive ) observation space , we compose the discriminator with the RL image encoder to leverage representations learned in the RL pipeline . GAIL-based methods in visual spaces are nascent , and contain works leveraging auxiliary losses ( Cetin & Celiktutan , 2021 ) . Toyer et al . ( 2020 ) propose a robust visual IL benchmark in which they evaluate mainly adversarial approaches in the discrete action setting . Finally , Rafailov et al . ( 2021 ) concurrently propose a model-based extension of GAIL , which does adversarial learning in the latent space of a world model . As a baseline , they use a direct extension of DAC to visual spaces by parameterizing a discriminator in image space , which does not perform well on nearly all tasks they consider , while our P-DAC extension achieves expert performance on all considered tasks ( including tasks considered by Rafailov et al . ( 2021 ) , such as walker walk and cheetah run ) . Optimal Transport for IL P-SIL extends to the visual case recent SOTA works on imitation learning via optimal transport ( OT ) , namely Dadashi et al . ( 2021 ) and Papagiannis & Li ( 2020 ) . They both leverage optimal transport matching to define rewards . The former uses a greedy approximation of the Wasserstein while the latter leverages the entropic Wasserstein with cosine cost . However , the latter also learns an embedding for states adversarially , while we leverage the encoder of the RL algorithm , avoiding the need for minimax optimization , and allowing us to avoid any computational slowdown . Finally , most of the adversarial algorithms in the GAIL line ( Ho & Ermon , 2016 ; Fu et al. , 2018 ) can also be interpreted from the standpoint of the minimization of an OT functional between the state ( -action ) occupancy distribution of the agent and expert ( Xiao et al. , 2019 ) . App . A reviews relevant background on optimal transport . 3 RECIPES FOR VISUAL IMITATION LEARNING . We now describe our IL approaches , which consist of alternative ways of defining imitation rewards that an RL backbone algorithm can learn from , and significantly improve upon canonical extensions of proprioceptive-state methods . In this paper , we leverage DrQ-v2 as the underlying RL algorithm ( Yarats et al. , 2021a ) . The data augmentation strategy throughout the paper ( including for baselines ) hence consists of random shifts with padding and a random crop to restore the original image dimension , followed by bilinear interpolation . For all methods , we gather episodes under the current policy , evaluate reward trajectories , and update the replay buffer with such episodes , replacing environment rewards by imitation rewards . We summarize the overall process in algorithm 1 . We also illustrate our two proposed methods in fig . 1 . 3.1 PIXEL SINKHORN IMITATION LEARNING ( P-SIL ) Our first approach extends imitation learning algorithms based on optimal transport to a setting where only pixel observations are available . We define imitation-reward signals via the negative entropic Wasserstein distance between embedded agent and experts ’ image trajectories . We interpret an image trajectory o = ( o1 , . . . , oT ) as a discrete probability measure of the form µo = 1 T ∑T t=1 δot , where images ot ∈ RC×H×W are atoms weighted uniformly over time . Optimal transport distances directly on image observations provide weak signal given that i ) the metric between individual images does not take into account spatial relationships between individual pixels , and ii ) the sample complexity of OT grows exponentially with the number of dimensions ( the number of samples in each trajectory necessary for a good estimate of the OT plan π ) ( Genevay et al. , 2019 ) . To alleviate these challenges , we embed image trajectories using DrQ-v2 ’ s encoder : oa , φ = [ fφ ( o a 1 ) , . . . , fφ ( o a T ) ] , oe , φ = [ fφ ( o e 1 ) , . . . , fφ ( o e T ) ] . ( 1 ) In order to be agnostic to the scale of the encoded states , we consider the cosine distance as metric dc between encoded visual observations , similarly to Papagiannis & Li ( 2020 ) , Ct , t′ = dc ( o a , φ t , o e , φ t′ ) = 1− 〈oa , φt , o e , φ t′ 〉 ||oa , φt || ||o e , φ t′ || . ( 2 ) We estimate the entropic Wasserstein distance with the cosine cost between embedded trajectories W2 ( µoa , φ , µoe , φ ) = min ψ∈Ψ T∑ t , t′=1 Ct , t′ψt , t′ −H ( ψ ) , H ( ψ ) = − T∑ t , t′=1 ψt , t′ logψt , t′ , ( 3 ) where Ψ = { ψ ∈ RT×T : ψ1 = ψT1 = 1T 1 } is the set of coupling matrices , in order to obtain an optimal alignment ψ ? . Finally , we extract rewards for each of the agent ’ s states as r ( oa , φt ) = − T∑ t′=1 Ct , t′ψ ? t , t′ . ( 4 ) If we are provided with multiple expert trajectories oe1 , . . . , oeN , we infer the nearest-neighbor expert by computing the embedded entropic Wasserstein distance with cosine cost between each expert trajectory and the agent ’ s rollout . We then set rewards with the alignment to the closest expert : e ? = arg min n∈ { 1 , ... , N } W2 ( µoa , φ , µoen , φ ) r ( o a , φ t ) = − T∑ t′=1 dc ( o a , φ t , o e ? , φ t′ ) ψ ? t , t′ . ( 5 ) The rewards computed in ( 5 ) are non-stationary as the encoder fφ is updated to optimize the critic loss . To increase the stability of P-SIL , we use a target network fφ′ , which is updated every Tupdate environment steps with the weights of DrQ-v2 ’ s encoder . We use fφ′ to estimate rewards via ( 5 ) . Also , we rescale P-SIL rewards so that the first episode is normalized to have a total return of −10 to account for rapid changes in encoder representations during initial training . An important benefit of P-SIL is that computing imitation rewards does not add any computational burden as the cost of ( 5 ) is negligible , and we use the encoder of the RL agent . This is in contrast to adversarial approaches ( Papagiannis & Li , 2020 ; Ho & Ermon , 2016 ) , which require training a discriminator by solving an inner-loop maximization problem and then inferring rewards using the discriminator . We thus preserve the attractive computational efficiency of the base algorithm DrQ-v2 . In summary , we compute imitation rewards as the entropic Wasserstein discrepancy between the agent ’ s and expert ’ s encoded states for the single-demonstration case . For multiple expert demonstrations , we use the transport cost from the closest expert demonstration in the OT sense . 3.2 PIXEL DISCRIMINATOR ACTOR CRITIC ( P-DAC ) We now propose and study P-DAC as a GAIL-based method for pixels that builds on the work by Ho & Ermon ( 2016 ) ; Torabi et al . ( 2018 ) ; Kostrikov et al . ( 2019 ) . The key differences with previous works are that i ) we compose the discriminator with the RL encoder to leverage its representations , ii ) we apply data augmentation to discriminator inputs and iii ) we store imitation rewards in the buffer rather than recomputing them at each iteration . We now describe the main components of P-DAC in more detail . In contrast to P-SIL , it is necessary to train a discriminator , along with DrQ-v2 ’ s other networks ( actor , critic and encoder ) . To train the discriminator , we sample batches oe1 : N and o a 1 : N of size N from both the expert buffer and the policy buffer , respectively , apply data augmentation , and encode augmented observations using DrQ-v2 ’ s encoder , which then yields : oa , φ = [ fφ ◦ aug ( oa1 ) , . . . , fφ ◦ aug ( oaN ) ] , oe , φ = [ fφ ◦ aug ( oe1 ) , . . . , fφ ◦ aug ( oeN ) ] . ( 6 ) Per common practice ( Kostrikov et al. , 2019 ) , we omit importance sampling when estimating the policy state occupancy via replay buffer . We then maximize max D N∑ n=1 logD ( oa , φn ) + N∑ n=1 log ( 1−D ( oe , φn ) ) − Eōφ [ ∥∥∇D ( ōφ ) ∥∥− 1 ] 2 , ( 7 ) with respect to the discriminator ’ s weights and add a gradient penalty as recommended by Kostrikov et al . ( 2019 ) . Here , ōφs are sampled along straight lines between agent and expert ’ s encoded observations . To compute rewards at the end of episodic rollouts , we embed the observation trajectory with DrQ-v2 ’ s encoder and then apply the discriminator to the encoded observations to obtain rewards r ( oat ) = log ( σ ◦D ( oa , φt ) ) − log ( 1− σ ◦D ( oa , φt ) ) , ( 8 ) where σ is the sigmoid function . We then add the episode with imitation rewards computed in ( 8 ) to the replay buffer . This differs from the original GAIL implementation ( Ho & Ermon , 2016 ) , where training is online , and from its follow-up offline extension DAC ( Kostrikov et al. , 2019 ) that recomputes rewards at each training iteration . We will show in the experiment section that data augmentation is a key component to scaling to challenging environments . Finally , in contrast with previous works that parameterize the discriminator directly in observation space ( Ho & Ermon , 2016 ; Torabi et al. , 2018 ; Kostrikov et al. , 2019 ) , we compose it with the RL encoder , which allows us to leverage its representations .

Summary: This paper introduces two methods for imitation learning from pixel observations based on adversarial learning and optimal transport. Representation from the RL encoder are adopted for calculation imitation rewards. P-SIL (Pixel Sinkhorn Imitation Learning) learns and compares latent representations of the encoded expert and agent behaviors; and P-SAC (Pixel Discriminator Actor Critic) augments the representations of the agents and the experts. The key component of this work is the RL encoder used in both P-SIL and P-SAC.