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: The authors have considered the problem of imitation learning using only visual observations, and they've explored the paradigms of using methods based on adversarial learning and optimal transport to solve this problem. In particular, the authors propose two new methods--P-SIL and P-DAC--presumably as a way to explore design choices and tradeoffs in this space. Some experiments designed to demonstrate the superiority of P-SIL and P-DAC are presented.