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: This paper tackles the problem of learning representations for imitation learning, in the context of offline data (with demonstrations). Firstly, given some low dimensional paramterization of the state, the authors provide a performance bound on the learnt policy over the representation matching the ground truth policy, regardless of the structure of the ground truth, as well as bound on the sample efficiency. The authors show that it is possible to learn the model and representations using a contrastive learning objective, and the final representation is output by applying random fourier features to the features learnt via contrastive learning. The authors provide some toy experiments in tabular decision tree settings to verify this hypothesis as well as on offline Atari dataset.