Provable Representation Learning for Imitation with Contrastive Fourier Features

1 Introduction . In the field of sequential decision making one aims to learn a behavior policy to act in an environment to optimize some criteria . The well-known field of reinforcement learning ( RL ) corresponds to one aspect of sequential decision making , where the aim is to learn how to act in the environment to maximize cumulative returns via trial-and-error experience [ 38 ] . In this work , we focus on imitation learning , where the aim is to learn how to act in the environment to match the behavior of some unknown target policy [ 25 ] . This focus puts us closer to the supervised learning regime , and , indeed , a common approach to imitation learning – known as behavioral cloning ( BC ) – is to perform max-likelihood training on a collected a set of target demonstrations composed of state-action pairs sampled from the target policy [ 31 , 34 ] . Since the learned behavior policy produces predictions ( actions ) conditioned on observations ( states ) , the amount of demonstrations needed to accurately match the target policy typically scales with the state dimension , and this can limit the applicability of imitation learning to settings where collecting large amounts of demonstrations is expensive , , in health [ 18 ] and robotics [ 24 ] applications . The limited availability of target demonstrations stands in contrast to the recent proliferation of large offline datasets for sequential decision making [ 28 , 19 , 7 , 21 ] . These datasets may exhibit behavior far from the target policy and so are not directly relevant to imitation learning via max likelihood 1Find experimental code at https : //github.com/google-research/google-research/tree/ master/rl_repr . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . training . Nevertheless , the offline datasets provide information about the unknown environment , presenting samples of environment reward and transition dynamics . It is therefore natural to wonder , is it possible to use such offline datasets to improve the sample efficiency of imitation learning ? Recent empirical work suggests that this is possible [ 40 , 10 ] , by using the offline datasets to learn a low-dimensional state representation via unsupervised training objectives . While these empirical successes are clear , the theoretical foundation for these results is less obvious . The main challenge in providing theoretical guarantees for such techniques is that of aliasing . Namely , even if environment rewards or dynamics exhibit a low-dimensional structure , the target policy and its demonstrations may not . If the target policy acts differently in states which the representation learning objective maps to the same low-dimensional representation , the downstream behavioral cloning objective may end up learning a policy which “ averages ” between these different states in unpredictable ways . In this work , we aim to bridge the gap between practical objectives and theoretical understanding . We derive an offline objective that learns low-dimensional representations of the environment dynamics and , if available , rewards . We show that minimizing this objective in conjunction with a downstream behavioral cloning objective corresponds to minimizing an upper bound on the performance difference between the learned low-dimensional BC policy and the unknown and possibly high-dimensional target policy . The form of our bound immediately makes clear that , as long as the learned policy is sufficiently expressive on top of the low-dimensional representations , the implicit “ averaging ” occurring in the BC objective due to any aliasing is irrelevant , and a learned policy can match the target regardless of whether the target policy itself is low-dimensional . Extending our results to policies with limited expressivity , we consider the commonly used parameterization of setting the learned policy to be log-linear with respect to the representations ( , a softmax of a linear transformation ) . In this setting , we show that it is enough to use the same offline representation learning objective , but with linearly parameterized dynamics and rewards , and this again leads to an upper bound showing that the downstream BC policy can match the target policy regardless of whether the target is low-dimensional or log-linear itself . We compare the form of our representation learning objective to “ latent space model ” approaches based on bisimulation principles , popular in the RL literature [ 20 , 41 , 22 , 13 ] , and show that these objectives are , in contrast , very liable to aliasing issues even in simple scenarios , explaining their poor performance in recent empirical studies [ 40 ] . We continue to the practicality of our own objective , and show that it can be implemented as a contrastive learning objective that implicitly learns an energy based model , which , in many common cases , corresponds to a linear model with respect to representations given by random Fourier features [ 33 ] . We evaluate our objective in both tabular synthetic domains and high-dimensional Atari game environments [ 11 ] . We find that our representation learning objective effectively leverages offline datasets to dramatically improve performance of behavioral cloning . 2 Related Work . Representation learning in sequential decision making has traditionally focused on learning representations for improved RL rather than imitation . While some works have proposed learning action representations [ 8 , 30 ] , our work focuses on state representation learning , which is more common in the literature , and whose aim is generally to distill aspects of the observation relevant to control from those relevant only to measurement [ 3 ] ; see [ 27 ] for a review . Of these approaches , bisimulation is the most theoretically mature [ 17 , 13 ] , and several recent works apply bisimulation principles to derive practical representation learning objectives [ 20 , 41 , 6 ] . However , the existing theoretical results for bisimulation fall short of the guarantees we provide . For one , many of the bisimulation results rely on defining a representation error which holds globally on all states and actions [ 1 ] . On the other hand , theoretical bisimulation results that define a representation error in terms of an expectation are inapplicable to imitation learning , as they only provide guarantees bounding the performance difference between policies that are “ close ” ( , Lipschitz ) in the representation space and say nothing regarding whether an arbitrary target policy in the true MDP can be represented in the latent MDP [ 20 ] . In Section 4.4 , we will show that these shortcomings of bisimulation fundamentally limit its applicability to an imitation learning setting . In contrast to RL , there are comparatively fewer theoretical works on representation learning for imitation learning . One previous line of research in this vein is given by [ 9 ] , which considers learning a state representation using a dataset of multiple demonstrations from multiple target policies . Accordingly , this approach requires that each target policy admits a low-dimensional representation . In contrast , our own work makes no assumption on the form of the target policy , and , in fact , this is one of the central challenges of representation learning in this setting . As imitation learning is close to supervised learning , it is an interesting avenue for future work to extend our results to more common supervised learning domains . We emphasize that our own contrastive objectives are distinct from typical approaches in image domains [ 15 ] and popular in image-based RL [ 26 , 37 , 36 ] , which use prior knowledge to generate pairs of similar images ( , via random cropping ) . We avoid any such prior knowledge of the task , and our losses are closer to temporal contrastive learning , more common in NLP [ 29 ] . 3 Background . We begin by introducing the notation and concepts we will build upon in the later sections . MDP Notation We consider the standard MDP framework [ 32 ] , in which the environment is given by a tupleM : = 〈S , A , R , P , µ , γ〉 , where S is the state space , A is the action space , R : S ×A→ [ −Rmax , Rmax ] is the reward function , P : S ×A→ ∆ ( S ) is the transition function,2 µ ∈ ∆ ( S ) is the initial state distribution , and γ ∈ [ 0 , 1 ) is the discount factor . In this work , we restrict our attention to finite action spaces , |A| ∈ N. A stationary policy in this MDP is a function π : S → ∆ ( A ) . A policy acts in the environment by starting at an initial state s0 ∼ µ and then at time t ≥ 0 sampling an action at ∼ π ( st ) . The environment then produces a reward rt = R ( st , at ) and stochastically transitions to a state st+1 ∼ P ( st , at ) . While we consider deterministic rewards , all of our results readily generalize to stochastic rewards , in which case the same bounds hold forR ( s , a ) denoting the expected value of the reward at ( s , a ) . The performance associated with a policy π is its expected future discounted reward when acting in the manner described above : JPerf ( π ) : = E [ ∞∑ t=0 γtR ( st , at ) ∣∣∣∣∣ π , M ] . ( 1 ) The visitation distribution of π is the state distribution induced by the sequential process : dπ ( s ) : = ( 1− γ ) ∞∑ t=0 γt · Pr [ st = s|π , M ] . ( 2 ) Behavioral Cloning ( BC ) In imitation learning , one wishes to recover an unknown target policy π∗ with access to demonstrations of π∗ acting in the environment . More formally , the demonstrations are given by a dataset Dπ∗N = { ( si , ai ) } Ni=1 where si ∼ dπ∗ , ai ∼ π∗ ( si ) . A popular approach to imitation learning is behavioral cloning ( BC ) , which suggests to learn a policy π to approximate π∗ via max-likelihood optimization . That is , one wishes to use the N samples to approximately minimize the objective JBC ( π ) : = E ( s , a ) ∼ ( dπ∗ , π∗ ) [ − log π ( a|s ) ] . ( 3 ) In this work , we consider using a state representation function to simplify this objective . Namely , we consider a function φ : S → Z . Given this representation , one no longer learns a policy π : S → ∆ ( A ) , but rather a policy πZ : Z → ∆ ( A ) . The BC loss with representation φ becomes JBC , φ ( πZ ) : = E ( s , a ) ∼ ( dπ∗ , π∗ ) [ − log πZ ( a|φ ( s ) ) ] . ( 4 ) A smaller representation space Z can help reduce the hypothesis space for πZ compared to π , and this in turn reduces the number of demonstrations N needed to achieve small error in JBC , φ . However , whether a small error in JBC , φ translates to a small error in JBC depends on the nature of φ , and thus how to determine a good φ is a central challenge . Offline Data In this work , we consider learning φ via offline objectives . We assume access to a dataset of transition tuples DoffM = { ( si , ai , ri , s′i ) } Mi=1 sampled independently according to si ∼ doff , ai ∼ UnifA , ri = R ( si , ai ) , s′i ∼ P ( si , ai ) , ( 5 ) 2We use ∆ ( X ) to denote the simplex over a set X . where doff is some unknown offline state distribution . We assume that the support of doff includes the support of dπ∗ ; , dπ∗ ( s ) > 0⇒ doff ( s ) > 0 . The uniform sampling of actions inDoff follows similar settings in related work [ 5 ] and in principle can be replaced with any distribution uniformly bounded from below by η > 0 and scaling our derived bounds by 1|A|η . At times we will abuse notation and write samples of these sub-tuples as ( s , a ) ∼ doff or ( s , a , r , s′ ) ∼ doff . Learning Goal Similar to related work [ 9 ] , we will measure the discrepancy between a candidate π and the target π∗ via the performance difference : PerfDiff ( π , π∗ ) : = |JPerf ( π ) − JPerf ( π∗ ) | . ( 6 ) At times we will also use the notation PerfDiff ( πZ , π∗ ) , and this is understood to mean PerfDiff ( πZ ◦ φ , π∗ ) . While we focus on the performance difference , all of our results may be easily modified to alternative evaluation metrics based on distributional divergences , , DTV ( dπ∗‖dπ ) or DKL ( dπ∗‖dπ ) , where DTV is the total variation ( TV ) divergence and DKL is the Kullback Leibler ( KL ) divergence . Also note that we make no assumption that π∗ is an optimal or near-optimal policy inM . In the case of vanilla behavioral cloning , we have the following relationship between JBC and the performance difference , which is a variant of Theorem 2.1 in [ 34 ] . Lemma 1 . For any π , π∗ , the performance difference may be bounded as PerfDiff ( π , π∗ ) ≤ Rmax ( 1− γ ) 2 √ 2Edπ∗ [ DKL ( π∗ ( s ) ‖π ( s ) ) ] = Rmax ( 1− γ ) 2 √ const ( π∗ ) + 2JBC ( π ) . ( 7 ) See Appendices B and C for all proofs . Remark ( Quadratic dependence on horizon ) . Notice that the guarantee above for vanilla BC includes a quadratic dependence on horizon in the form of ( 1− γ ) −2 , and this quadratic dependence is maintained in all our subsequent bounds . While there exists a number of imitation learning works that aim to reduce this dependence , the specific problem our paper focuses on – aliasing in the context of learning state representations – is an orthogonal problem to quadratic dependence on horizon . Indeed , if some representation maps two very different raw observations to the same latent state , no downstream imitation learning algorithm ( regardless of sample complexity ) will be able to learn a good policy . Still , extending our representation learning bounds to more sophisticated algorithms with potentially smaller dependence on horizon , like DAgger [ 35 ] , is a promising direction for future work .

Summary: The authors study the impact of representation learning when applied to behavioral cloning (BC), and present a series of interesting theoretical results that establish its utility in this context. In particular, the authors show that there is a direct relationship between the quality of the state representation and the performance of the policy derived using BC over that representation. The authors then present a practical representation learning algorithm designed using the insight gained from their theoretical results and show that using it in combination with BC results in an improvement over comparison approaches.