Provable Representation Learning for Imitation with Contrastive Fourier Features
Year: 2022

Abstract:
In imitation learning, it is common to learn a behavior policy to match an unknown target policy via max-likelihood training on a collected set of target demonstrations. In this work, we consider using offline experience datasets – potentially far from the target distribution – to learn low-dimensional state representations that provably accelerate the sample-efficiency of downstream imitation learning. A central challenge in this setting is that the unknown target policy itself may not exhibit low-dimensional behavior, and so there is a potential for the representation learning objective to alias states in which the target policy acts differently. Circumventing this challenge, we derive a representation learning objective that provides an upper bound on the performance difference between the target policy and a lowdimensional policy trained with max-likelihood, and this bound is tight regardless of whether the target policy itself exhibits low-dimensional structure. Moving to the practicality of our method, we show that our objective can be implemented as contrastive learning, in which the transition dynamics are approximated by either an implicit energy-based model or, in some special cases, an implicit linear model with representations given by random Fourier features. Experiments on both tabular environments and high-dimensional Atari games provide quantitative evidence for the practical benefits of our proposed objective.1

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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].

heading: 3 Background
text: 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.

heading: 4 Representation Learning with Performance Bounds
text: We now continue to our contributions, beginning by presenting performance difference bounds analogous to Lemma 1 but with respect to a specific representation φ. The bounds will necessarily depend on quantities which correspond to how “good” the representation is, and these quantities then form the representation learning objective for learning φ; ideally these quantities are independent of π∗, which is unknown.
Intuitively, we need φ to encapsulate the important aspects of the environment. To this end, we consider representation-based models PZ : Z ×A→ ∆(S) andRZ : Z ×A→ [−Rmax, Rmax] of the environment transitions and rewards. We define the error incurred by these models on the offline distribution as,
JR(RZ , φ)2 := E(s,a)∼doff [ (R(s, a)−RZ(φ(s), a))2 ] , (8)
JT(PZ , φ)2 := 1
2 E(s,a)∼doff [DKL(P(s, a)‖PZ(φ(s), a))] . (9)
We will elaborate on how these errors are translated to representation learning objectives in practice in Section 5, but for now we note that the expectation over doff already leads to a connection between theory and practice much closer than exists in other works, which often resort to supremums over state-action errors [30, 41, 1] or zero errors globally [17].

heading: 4.1 General Policies
text: We now relate the representation errors above to the performance difference, analogous to Lemma 1 but with JBC,φ. We emphasize that the performance difference below measures the difference in returns in the true MDPM, rather than, as is commonly seen in model-based RL [23], the difference in the latent MDP defined byRZ ,PZ . Theorem 2. Consider a representation function φ : S → Z and modelsRZ ,PZ as defined above. Denote the representation error as
R,T := |A|
1− γ JR(RZ , φ) + 2γ|A|Rmax (1− γ)2 JT(PZ , φ). (10)
Then the performance difference inM between π∗ and a latent policy πZ : Z :→ ∆(A) may be bounded as,
PerfDiff(πZ , π∗) ≤ (1 +Dχ2(dπ∗‖doff) 1 2 ) · R,T︸ ︷︷ ︸
offline pretraining
+C
√ 1
2 Ez∼dπ∗Z [DKL(π∗,Z(z)‖πZ(z))]︸ ︷︷ ︸ = const(π∗, φ) + JBC,φ(πZ)︸ ︷︷ ︸
downstream behavioral cloning
, (11)
where C = 2Rmax(1−γ)2 and d π∗ Z , π∗,Z are the marginalization of d π∗ , π∗ onto Z according to φ:
dπ∗Z (z) := Pr[z = φ(s) | s ∼ d π∗ ] ; π∗,Z(z) := E[π∗(s) | s ∼ dπ∗ , z = φ(s)]. (12)
Proof in Appendix B.2.
We thus have a guarantee showing that πZ can match the performance of π∗ regardless of the form of π∗ (, whether π∗ itself possesses a low-dimensional parameterization). Indeed, as long as φ is is learned well enough ( R,T → 0) and the space of candidate πZ is expressive enough, optimizing JBC,φ can achieve near-zero performance difference, since the latent policy optimizing JBC,φ is πZ = π∗,Z , and this setting of πZ zeros out the second term of the bound in (11). Note that, in general π∗ 6= π∗,Z ◦ φ, yet the performance difference of these two distinct policies is nevertheless zero when R,T = 0. The bound in Theorem 2 also clearly exhibits the trade-off due to the offline distribution, encapsulated by the Pearson χ2 divergence Dχ2(dπ∗‖doff). Remark (Representations agnostic to environment rewards). In some imitation learning settings, rewards are unobserved and so optimizing JR is infeasible. In these settings, one can consider setting RZ to an (unobserved) constant function Edoff [R(s, a)], ensuring JR(RZ , φ) ≤ Rmax. Furthermore, in environments where the reward does not depend on the action, ,R(s, a1) = R(s, a2)∀a1, a2 ∈ A, the bound in Theorem 2 can be modified to remove JR altogether (see Appendix B for details).

heading: 4.2 Log-linear Policies
text: Theorem 2 establishes a connection between the performance difference and behavioral cloning over representations given by φ. The optimal latent policy for BC is π∗,Z , and this is the same policy which achieves minimal performance difference. Whether we can find πZ ≈ π∗,Z depends on how we parameterize our latent policy. If πZ is tabular or if πZ is represented as a sufficiently expressive neural network, then the approximation error is effectively zero. But what about in other cases?
In this subsection, we consider Z ⊂ Rd and log-linear policies of the form
πθ(z) = softmax(θ >z) := ( exp{θ>a z}∑ ã exp{θ>ã z} ) a∈A , (13)
where θ ∈ Rd×|A|. In general, π∗,Z cannot be expressed as πθ for some θ. Nevertheless, we can still derive strong bounds for this scenario, by considering factored linear [4] modelsRZ ,PZ : Theorem 3. Consider Z ⊂ Rd, a representation function φ : S → Z, and linear modelsRZ ,PZ: RZ(z, a) := r(a)>z ; PZ(s′|z, a) := ψ(s′, a)>z for some r : A→ Rd ; ψ : S ×A→ Rd.
Denote the representation error R,T as in Theorem 2. Then the performance difference between π∗ and a latent policy πθ(z) := softmax(θ>z) may be bounded as,
PerfDiff(πθ, π∗) ≤ (1 +Dχ2(dπ∗‖doff) 1 2 ) · R,T + C · ∥∥∥∥ ∂∂θJBC,φ(πθ) ∥∥∥∥
1
, (14)
where C = 11−γ ‖r‖∞ + γRmax (1−γ)2 ‖ψ‖∞.
Proof in Appendix B.3.
The statement of Theorem 3 makes it clear that realizability of π∗,Z is irrelevant for log-linear policies. It is enough to only have the gradient with respect to learned θ be close to zero, which is a guarantee of virtually all gradient-based algorithms. Thus, in these settings performing BC on top of learned representations is provably optimal regardless of both the form of π∗ and the form of π∗,Z . Remark (Kernel-based models). It is possible to extend the statement of Theorem 3 to generalized linear dynamics and reward models based on kernels by replacing the gradient in the bound with the functional gradient with respect to the kernel [16].

heading: 4.3 Sample Efficiency
text: The previous theorems show that we can reduce imitation learning to (1) representation learning on an offline dataset, and (2) behavioral cloning on target demonstrations with the learned representation. How does this compare to performing BC on the target demonstrations directly? Intuitively, representation learning should help when the learned representation is “good” (, R,T is small) and the complexity of the representation space Z is low (, |Z| is small or Z ⊂ Rd for small d). In this subsection, we formalize this intuition for a simple setting. We consider finite S,Z. For the representation φ, we assume access to an oracle φM := OPT φ(DoffM ) which yields an error R,T(φM ). For BC we consider learning a tabular πZ on the finite demonstration set Dπ∗N . We have the following theorem, which characterizes the expected performance difference when using representation learning. Theorem 4. Consider the setting described above. Let φM := OPT φ(DoffM ) and πN,Z be the policy resulting from BC with respect to φM . Then we have,
EDπ∗N [PerfDiff(πN,Z , π∗)] ≤ (1 +Dχ2(d off‖dπ∗) 12 ) · R,T(φM ) + C · √ |Z||A| N , (15)
where C is as in Theorem 2.
Proof in Appendix C.1.
Note that application of vanilla BC to this setting would achieve a similar bound but with R,T = 0 and |Z| = |S|. Thus, an improvement from representation learning is expected when R,T(φM ) and |Z| are small.

heading: 4.4 Comparison to Bisimulation
text: The form of our representation learning objectives – learning φ to be predictive of rewards and next state dynamics – recalls similar ideas in the bisimulation literature [17, 20, 41, 13]. However, a key difference is that in bisimulation the divergence over next state dynamics is measured in the latent representation space; , a divergence between φ◦P(s, a) and f(φ(s), a) for some “latent space model” f , whereas our proposed representation error is between P(s, a) and PZ(φ(s), a). We find that this difference is crucial, and in fact there exist no theoretical guarantees for bisimulation similar to those in Theorems 2 and 3. Indeed, one can construct a simple example where the use of latent space models leads to a complete failure, see Figure 1.

heading: 5 Learning the Representations in Practice
text: The bounds presented in the previous section suggest that a good representation φ should be learned to minimize JR, JT in (8) and (9). To this end we propose to learn φ in conjunction with auxiliary models of rewardRZ and dynamics PZ . The offline representation learning objective is given by,
Jrep(RZ ,PZ , φ) := 1
2 E(s,a,r,s′)∼doff [αR · (r −RZ(φ(s), a))2 − αT · logPZ(s′|φ(s), a)], (16)
where αR, αT are appropriately chosen hyperparameters; in our implementation we choose αR = 1, αT = (1 − γ)−1 to roughly match the coefficients of the bounds in Theorems 2 and 3. Once one chooses parameterizations of RZ ,PZ , this objective may be optimized using any stochastic sample-based solver, SGD.

heading: 5.1 Contrastive Learning
text: One may recover a contrastive learning objective by parameterizing PZ as an energy-based model. Namely, consider parameterizing PZ as
PZ(s′|z, a) ∝ ρ(s′) exp { −||z − g(s′, a)||2/2 } , (17)
where ρ is a fixed (untrainable) distribution over S (typically set to the distribution of s′ in doff ) and g is a learnable function S ×A→ Z (, a neural network). Then Edoff [− logPZ(s′|φ(s), a)] yields a contrastive loss:
Edoff [− logPZ(s′|φ(s), a)] = 1
2 Edoff [||φ(s)−g(s′, a)||2]+logEs̃′∼ρ[exp{−||φ(s)−g(s̃′, a)||2/2}].
Similar contrastive learning objectives have appeared in related works [10, 40], and so our theoretical bounds can be used to explain these previous empirical successes.

heading: 5.2 Linear Models with Contrastive Fourier Features
text: While the connection between temporal contrastive learning and approximate dynamics models has appeared in previous works [30], it is not immediately clear how one should learn the approximate linear dynamics required by Theorem 3. In this section, we show how the same contrastive learning objective can be used to learn approximate linear models, thus illuminating a new connection between contrastive learning and near-optimal sequential decision making; see Appendix A for pseudocode.
We propose to learn a dynamics model P(s′|s, a) ∝ ρ(s′) exp{−||f(s) − g(s′, a)||2/2} for some functions f, g (, neural networks), which admits a similar contrastive learning objective as mentioned above. Note that this parameterization does not involve φ. To recover φ, we may leverage random Fourier features from the kernel literature [33]. Namely, for k-dimensional vectors x, y we can approximate exp{−||x − y||2/2} ≈ 2dϕ(x)
>ϕ(y), where ϕ(x) := cos(Wx + b) for W a d × k matrix with entries sampled from a standard Gaussian and b a vector with entries sampled uniformly from [0, 2π]. We can therefore approximate P as follows, using E(s, a) := Es̃′∼ρ[exp{−||f(s)− g(s̃′, a)||2/2}]:
P(s′|s, a) = ρ(s ′)
E(s, a) exp{−‖f(s)− g(s′, a)‖2/2} ≈ 2ρ(s
′)
d · E(s, a) ϕ(f(s))>ϕ(g(s′, a)). (18)
Finally, we recover φ : S → R|A|d, ψ : S ×A→ R|A|d as φ(s) := [ϕ(f(s))/E(s, a)]a∈A ; ψ(s ′, a) := [1a=ã · 2ρ(s′)ϕ(s′, ã)/d]ã∈A , (19)
ensuring P(s′|s, a) ≈ φ(s)>ψ(s′, a) (, PZ(s′|z, a) = ψ(s′, a)>z), as required for Theorem 3.3 Notably, during learning explicit computation of ψ(s′, a) is never needed, while E(s, a) is straightforward to estimate from data when ρ(s′) = doff(s′),4 thus making the whole learning procedure – both the offline representation learning and downstream behavioral cloning – practical and easy to implement with deep neural network function approximators for f, g,RZ ; in fact, one can interpret φ as simply adding an additional untrainable neural network layer on top of f , and so these derivations may partially explain why previous works in supervised learning found benefit from adding a non-linear layer on top of representations [15].

heading: 6 Experiments
text: We now empirically verify the performance benefits of the proposed representation learning objective in both tabular and Atari game environments. See environment details in Appendix D.

heading: 6.1 Tree Environments with Low-Rank Transitions
text: For tabular evaluation, we construct a decision tree environment whose transitions exhibit low-rank structures. To achieve this, we first construct a “canonical” three-level binary tree where each node is associated with a stochastic reward for taking left or right branch and a stochastic transition matrix indicating the probability of landing at either child node. We then duplicate this canonical tree in the state space so that an agent walks down the duplicated tree while the MDP transitions are determined by the canonical tree, thus it is possible to achieve R,T = 0 with |Z| = 8. We collect the offline data using a uniform random policy and the target demonstrations using a return-optimal policy.
We learn contrastive Fourier features as described in Section 5.2 using tabular f, g. We then fix these representations and train a log-linear policy on the target demonstrations. For the baseline, we learn a vanilla BC policy with tabular parametrization directly on target demonstrations.
We also experiment with representations given by singular value decomposition (SVD) of the empirical transition matrix, which is another form of learning factored linear dynamics. Figure 2 shows the performance achieved by the learned policy with and without representation learning. Representation learning consistently yields significant performance gains, especially with few target demonstrations. SVD performs similar to contrastive Fourier features when the offline data is abundant with respect to the state space size, but degrades as the offline data size reduces or the state space grows.

heading: 6.2 Atari 2600 with Deep Neural Networks
text: We now study the practical benefit of the proposed contrastive learning objective to imitation learning on 60 Atari 2600 games [11], taking for the offline dataset the DQN Replay Dataset [7], which for
3While we use random Fourier features with the radial basis kernel, it is clear that a similar derivation can be used with other approximate featurization schemes and other kernels.
4In our experiments, we ignore the scaling E(s, a) altogether and simply set φ(s) := ϕ(f(s)). We found this simplification to have little effect on results.
each game provides 50M steps collected during DQN training. For the target demonstrations, we take 10k single-step transitions from the last 1M steps of each dataset, corresponding to the data collected near the end of the DQN training. For the offline data, we use all 50M transitions of each dataset.
For our learning agents, we extend the implementations found in Dopamine [14]. We use the standard Atari CNN architecture to embed image-based inputs to vectors of dimension 256. In the case of vanilla BC, we pass this embedding to a log-linear policy and learn the whole network end-to-end with behavioral cloning. For contrastive Fourier features, we use separate CNNs to parameterize g, f in the objective in Section 5.2, and the representation φ := ϕ ◦ f is given by the random Fourier feature procedure described in Section 5.2. A log-linear policy is then trained on top of this representation, but without passing any BC gradients through φ. This corresponds to the setting of Theorem 3. We also experiment with the setting of Theorem 2; in this case the setup is same as for contrastive Fourier features, only that we define φ := f (, PZ is an energy-based dynamics model) and we parameterize πZ as a more expressive single-hidden-layer softmax policy on top of φ.
We compare contrastive learning with Fourier features and energy-based models to two latent space models, DeepMDP [20] and Deep Bisimulation for Control (DBC) [41], in Figure 3. Both linear (Fourier features) and energy-based parametrization of contrastive learning achieve dramatic performance gains (> 40%) on over half of the games. DeepMDP and DBC, on the other hand, achieve little improvement over vanilla BC when presented as a separate loss from behavioral cloning. Enabling end-to-end learning of the latent space models as an auxiliary loss to behavioral cloning leads to better performance, but DeepMDP and DBC still underperform contrastive learning. See Appendix D for further ablations.

heading: 7 Conclusion
text: We have derived an offline representation learning objective which, when combined with BC, provably minimizes an upper bound on the performance difference from the target policy. We further showed that the proposed objective can be implemented as contrastive learning with an optional projection to Fourier features. Interesting avenues for future work include (1) extending our theory to multi-step contrastive learning, popular in practice [40, 10], (2) deriving similar results for policy learning in offline and online RL settings, and (3) reducing the effect of offline distribution shifts. We also note that our use of contrastive Fourier features for learning a linear dynamics model may be of independent interest, especially considering that a number of theoretical RL works rely on such an approximation (, [4, 39]), while to our knowledge no previous work has demonstrated a practical
and scalable learning algorithm for linear dynamics approximation. Determining if the technique of learning contrastive Fourier features works well for these settings offers another interesting direction to explore.
Limitations One of the main limitations inherent in our theoretical derivations is the dependence on distribution shift, in the form of 1 +Dχ2(doff‖dπ∗) 1 2 in all bounds. Arguably, some dependence on the offline distribution is unavoidable: Certainly, if a specific state does not appear in the offline distribution, then there is no way to learn a good representation of it (without further assumptions on the MDP). In practice one potential remedy is to make sure that the offline dataset sufficiently covers the whole state space; the inequality 1 +Dχ2(p‖q) ≤ ‖p/q‖2∞ will ensure this limits the dependence on distribution shift.
Acknowledgments and Disclosure of Funding
We thank Bo Dai, Rishabh Agarwal, Mohammad Norouzi, Pablo Castro, Marlos Machado, Marc Bellemare, and the rest of the Google Brain team for fruitful discussions and valuable feedback.

Review summaries:
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.

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. 

This paper proposes an approach for learning a low-dimensional state representation that favors the training procedure of a policy with imitation learning. One of the main challenges is that the representation learning strategy is run on offline data, and is possibly not informed about the expert policy to imitate. The paper presents a theoretical framework to bound the performance of the student policy in function of the expressiveness of the representation and the distance to the distance between expert and student on low-dimensional states. The theoretical findings are then put into practice by deriving an algorithm to learn representations via contrastive learning and Fourier features. The algorithm is then evaluated on both toy tabular examples and on more complex high-dimensional tasks (Atari).