What Structural Conditions Permit Generalization
in Reinforcement Learning?
Joint work with
Simon Du, Sham Kakade, Jason Lee, Shachar Lovett,Gaurav Mahajan, Ruosong Wang

Solving Large-scale RL problems requires generalization[AlphaZero, Silver et.al, 17][OpenAI Five, 18][OpenAI,19]

Markov Decision Processes:a framework for RL• A policy:
π : States → Actions• Execute π to obtain a trajectory:s0, a0, r0, s1, a1, r1…sH−1, aH−1, rH−1
• Cumulative H-step reward:
πVH(s) =H−1

πrt s0 = s[∑]t=0π, QH(s, a) =H−1π[∑t=0rt s0 = s, a0 = aπ𝔼𝔼• Goal: Find a policy π that maximizes our value V (s0) from s0.
Episodic setting: We start at s0; act for H steps; repeat…

]

𝑑Provable Generalization in Supervised Learning (SL)Generalization is possible in the IID supervised learning setting!
To get ϵ-close to best in hypothesis class ℱ, we need # of samples that is:2• “Occam’s Razor” Bound (finite hypothesis class): need O(log | ℱ | /ϵ )Various Improvements:
2• VC dim: need only O(VC(ℱ)/ϵ )2• Classification: linearly separable + margin: O(margin)/ϵ )2• Linear Regression in dimensions: O(d/ϵ )• Deep Learning: the algorithm also determines the complexity control
The key idea in SL: data reuse
With a training set, we can simultaneously evaluate the loss of all hypotheses in our class!
4

Sample E cient RL inthe Tabular Case(no generalization here)• Thm: In the episodic setting, poly(S, A, H,1/ϵ)
samples su ce to nd an ϵ-opt policy with the3E algo. [Kearns & Singh ‘98]
also: [Brafman& Tennenholtz ‘02; K. '03]
Key idea: optimism + dynamic programming• Regret guarantees with model based algos:fiffi

•

ffi

•

[Auer+ ‘09]
Provable Q-learning (+bonus):[Strehl+ (2006)], [Szita & Szepesvari ‘10],[Jin+ ‘18]
(asymptotically) optimal reget: Reg(#episodes) =[Azar+ ’17],[Dann+’17]

HSA ⋅ #episodes

I: Provable Generalization in RL
Q1: Can we find an ϵ-opt policy with no S dependence?
• How can we reuse data to estimate the value of all policies in a policy class ℱ?
Idea: Trajectory tree algodataset collection: uniformly at random choose actions for all H steps in an episode.
estimation: uses importance sampling to evaluate every f ∈ ℱ• Thm:[Kearns, Mansour, & Ng ‘00]
H
2To find an ϵ-best in class policy, the trajectory tree algo uses O(A log( | ℱ | )/ϵ ) samples• Only log( | ℱ | ) dependence on hypothesis class size.
• There are VC analogues as well.
H• Can we avoid the 2 dependence to find an ϵ-best-in-class policy?

Agnostically, NO!
HProof: Consider a binary tree with 2 -policies and a sparse reward at a leaf node. 6

II: Provable Generalization in RL• Q2: Can we nd an ϵ-opt policy with no S, A dependence andpoly(H,1/ϵ, "complexity measure") samples?

fi•With various stronger assumptions, yes.
• Linear Bellman Completion: [Munos, ’05, Zanette+ ‘19]
• Linear MDPs: [Wang & Yang’18]; [Jin+ ’19] (the transition matrix is low rank)• Linear Quadratic Regulators (LQR): standard control theory model• FLAMBE / Feature Selection: [Agarwal, K., Krishnamurthy, Sun ’20]
• Linear Mixture MDPs: [Modi+’20, Ayoub+ ’20]
• Block MDPs [Du+ ’19]
• Factored MDPs [Sun+ ’19]
• Kernelized Nonlinear Regulator [K.+ ’20]

This talk:
Structural conditions
Under which Generalization is possible

Warm up
The linear Bellman Complete model:
dDef: Given feature map ϕ(s, a) ∈ ℝ , Bellman operator has linear closuredd

Given any w ∈ ℝ , there exists a θ ∈ ℝ , such that:⊤∀s, a : θ ϕ(s, a) = r(s, a) +:= T(w)s′∼P(s,a) [max wa′⊤

Examples:￼￼￼￼Linear MDPs, Linear Quadratic Regulator𝔼]ϕ(s′, a′)Warm up
The linear Bellman Complete model:
Generalization is possible here:fi∃ an algorithm, nding ϵ-near optimal policy only needspoly(H, d,1/ϵ) many samples

Warm up
The linear Bellman Complete model:
what’s the structure here that permits generalization?
We can rewrite the average Bellman error (averaged over any roll-in π)
in a bilinear form:⋆⊤

Given a Q candidate w ϕ(s, a), we have:ϕ(s′, a′)]⊤⊤= s,a∼π [w ϕ(s, a) − T(w) ϕ(s, a)] = ⟨w − T(w), s,a∼πϕ(s, a)⟩ϕ(s, a) − r(s, a) −￼𝔼￼

𝔼𝔼𝔼￼

￼

s,a∼π [w⊤s′∼P(s,a) max wa′⊤

Warm up
The linear Bellman Complete model:
what’s the unique structure here that permits generalization?
We can also estimate the value of the bilinear form:⊤⊤

De ne discrepancy ℓ(s, a, s′, w) = w ϕ(s, a) − r(s, a) − max w ϕ(s′, a′)a′

We have:ℓ(s,a,s′,w)=w−T(w),⟨s,a∼π￼￼￼

Note we have data reuse: given data from π, we can evaluate all w𝔼￼￼fi𝔼ϕ(s,a)⟩s,a∼πWarm up
The linear Bellman Complete model:
In summary, it has a bilinear structure:
For any roll-in policy π, and any w, we have: (1)s,a∼π [wϕ(s, a) − r(s, a) −s′∼P(s,a) max w
a′ϕ(s′, a′) = ⟨w − T(w),]⊤AND (2) there exists a discrepancy function ℓ, s.t.,ℓ(s,a,s′,w)=w−T(w),⟨s,a∼πϕ(s,a)⟩π𝔼𝔼￼￼￼￼

Note that the analytical form of bilinear structure is unknown𝔼￼𝔼𝔼⊤πϕ(s, a)⟩• Realizable Hypothesis class: {f ∈ ℱ},

with associated state-action value, (greedy) value and policy: Qf (s, a), Vf (s), πf• can be model based or model-free class.

Def: A (ℱ, ℓ) forms an (implicit) Bilinear class class if there areWh ∈ ℱ ↦ ℋ, & Xh ∈ ℱ ↦ ℋ (ℋ being some Hilbert space):

Note: Wh & Xh are implicit—no need to known themff𝔼• Bilinear regret: on-policy di erence between claimed reward and true reward⋆sh,ah∼πf[Qf (sh, ah) − r(sh, ah) − Vf (sh+1)] ≤ ⟨Wh( f ) − Wh( f ), Xh( f )⟩• Data reuse: there is discrepancy function ℓf (s, a, s′, g) & policy πest s.t.
⋆ℓ(s,a,s,g)=W(g)−W(f),X(f)
,∀g
∈
ℱ[]
⟨⟩
sh∼πf ,ah∼πest f h h h+1hhh
￼𝔼BiLinear Classes: structural properties to enablegeneralization in RL

Back to Linear Bellman Complete:⊤(1) Bilinear regret: for any f(s, a) := w ϕ(s, a), we have:sh,ah∼πf [w⊤ϕ(sh, ah) − r(sh, ah) −ϕ(s′, a′) = (w − T(w)) − (w − T(w )),] ⟨⋆Wh( f )⋆Wh( f ⋆)AND (2) data-reuse: there exists a discrepancy function ℓ, s.t.,ℓ(s,a,s′,w′)=(w′−T(w′))−(w−T(w)),πf⟨⋆⋆ϕ(s,a)s,a∼πf⟩
𝔼𝔼

￼

￼

￼

￼

￼

￼

𝔼

￼

￼

Note Wh( f ) is unknown as the Bellman backup T(w) is unknown

𝔼𝔼s′∼P(s,a) max wa′⊤⟩

πϕ(s, a)Xh( f )

The Algorithm: BiLin-UCBFor t = 0 → T:

• Find the “optimistic” ft ∈ ℱ:
2arg max Vf (s0), s.t., σh ( f ) ≤ R, ∀hf∈ℱ• Sample m trajectories πft and create a batch dataset:
D = {(sh, ah, sh+1) ∈ trajectories}• Update the cumulative discrepancy function σh( ⋅ ), ∀h𝔼

22σh ( ⋅ ) ← σh ( ⋅ ) + (∑(sh,ah,sh+1)∈Dℓft(sh, ah, sh+1, ⋅ )/ | D | )Note here we roughly have: σh(g) ≈t∑(
i=02sh,ah∼πfiℓfi(sh, ah, sh+1, g))2

Theorem 2: Generalization in RL• Theorem: [Du, Kakade., Lee, Lovett, Mahajan, S, Wang ’21]
⋆Assume ℱ is a bilinear class and the class is realizable, i.e. f ∈ ℱ.
32Using γT ⋅ poly(H) ⋅ log(1/δ)/ϵ trajectories, the BiLin-UCB algorithm
returns an ϵ-opt policy (with prob. ≥ 1 − δ).
T−11⊤
γT is the max. info. gain γT := max ln det I + ∑ Xh( ft)Xh( ft)•
λ t=0h,f0…fT−1∈ℱ()• γT ≈ d log T for Xh in d-dimensions• The proof is “elementary” using the elliptical potential function.
[Dani, Hayes, K. ’08]

Proof Sketch⋆⋆

1. Optimism: V (s0) ≤ Vft(s0); this can be veri ed by showing f is always a feasible solution
2. Using optimism, we can upper bound per-episode regret Bilinear form (“simulation” lemma):πft⋆πft

V (s0) − V (s0) ≤ Vft(s0) − V (s0) ≤πft

⋆

H−1

∑
h=0

⋆

Wh( ft) − Wh( f ), Xh( ft)

3. If πft is really sub-optimal, i.e., V (s0) − V (s0) ≥ ϵ, then ft has large bilinear regret:∃h, s.t., Wh( ft) − Wh( f ), Xh( ft) ≥ ϵ/H⋆

4. Recall in the alg, we have a constraint σh( ft) ≤ R, i.e.,t−1

Wh( ft) − Wh( f ), Xh( fi)) ≤ R(∑⋆2⋆ ⊤⋆⇒ (Wh( ft) − Wh( f )) Σt;h(Wh( ft) − Wh( f )) ≤ R + λfi

𝔼i=0

∑(
t−1i=0(sh,ah∼πfi [ℓfi(sh, ah, sh+1, ft)])

Σh;t :=t−1

∑
i=0⊤)

Xh( fi)Xh( fi) + λI2≤R

Proof Sketchπft⋆

3. If πft is really sub-optimal, i.e., V (s0) − V (s0) ≥ ϵ, then ft has large bilinear value:∃h, s.t., Wh( ft) − Wh( f ), Xh( ft) ≥ ϵ/H⋆

4. Recall in the alg, we have a constraint σ( ft) ≤ R, i.e.,t−1∑(
τ=0⋆ ⊤sh,ah∼πfτℓfτ(sh, ah, sh+1, ft))⋆(Wh( ft) − Wh( f )) Σt;h(Wh( ft) − Wh( f )) ≤ R + λ(Σh;t :=t−1∑
i=02≤R)Xh( fi)Xh( fi)⊤ + λI

5. Finally, combine the two results and using Cauchy-Schwartz, we have:
ϵ/H ≤

⋆

Wh( ft) − Wh( f )⇒Σh;tXh( ft)Xh( ft)≥−1Σh;t≤−1Σh;tR+λXh( ft)Σ−1h;tϵH R+λ2𝔼

In d-dimensional setting, such event cannot happen more than Õ (d/ϵ ) many times….

• Theorem: [Du, Kakade., Lee, Lovett, Mahajan, S, Wang ’21]

The following models are bilinear classes for some discrepancy function ℓ( ⋅ )• Linear Bellman Completion: [Munos, ’05, Zanette+ ‘19]
• Linear MDPs: [Wang & Yang’18]; [Jin+ ’19] (the transition matrix is low rank)• Linear Quadratic Regulators (LQR): standard control theory model• FLAMBE / Feature Selection: [Agarwal, K., Krishnamurthy, Sun ’20]
• Linear Mixture MDPs: [Modi+’20, Ayoub+ ’20]
• Block MDPs [Du+ ’19]
• Factored MDPs [Sun+ ’19]
• Kernelized Nonlinear Regulator [K.+ ’20]⋆⋆
• Linear Q & V• Reactive PSR/POMDP, Generalized linear MDP, and more…..

• (almost) all “named” models (with provable generalization) are bilinear classes⋆⋆
two exceptions: deterministic linear Q ; Q -state-action aggregation• Bilinear classes generalize the: Bellman rank [Jiang+ ‘17]; Witness rank [Sun+ ’19]

The feature selection problem:⋆⊤ ⋆⋆⋆Low-rank MDP: P(s′| s, a) = μ (s′) ϕ (s, a), μ & ϕ unknown⋆dFunction approximation for feature: ϕ ∈ Ψ ⊂ S × A ↦ ℝ⋆

Function approximation for Q ::= {w ϕ(s, a) : w ∈ BallW, ϕ ∈ Ψ}⊤

Q: can we nd ϵ optimal policy w/ samples poly(d, ln(Ψ), H,1/ϵ)?

fi￼

Yes & it’s in the Bilinear class!

￼𝒬

Another Example: Feature Selection for Low-rank MDPAnother Example: Feature Selection for Low-rank MDP
The feature selection problem:⊤1. Claim: on-policy Bellman error of f := w ϕ has Bilinear form:
sh,ah∼πf [Qf (sh, ah) − rh −sh−1,ah−1∼πf sh∼P(⋅|sh−1,ah−1),ah∼πf (⋅|sh) [Qf (sh, ah) − rh −(sh−1, ah−1),μ (sh) ah∼πf Qf (sh, ah) − rh −[∫sϕ (sh−1, ah−1) μ (sh)h￼s′∼P(⋅|s,a) max Qf (sh+1, a′)]
a′
⟩

￼￼

￼

￼

𝔼

￼
𝔼𝔼𝔼

￼

Xh( f )s′∼P(⋅|s,a) max Qf (sh+1, a′)]a′⋆h𝔼𝔼𝔼𝒬𝔼𝔼𝔼⟨

sh−1,ah−1∼πf ϕ⋆𝔼∫sah∼πf [Qf (sh, ah) − rh −⊤ ⋆￼

sh−1,ah−1∼πf⋆￼=s′∼P(⋅|s,a) max Qf (sh+1, a′)]a′￼==s′∼P(⋅|s,a) max Qf (sh+1, a′)]
a′

￼

:

￼∀f ∈Wh( f )Another Example: Feature Selection for Low-rank MDP
The feature selection problem:
2. Claim: The bilinear regret is estimable using some ℓ
De ne ℓ(s, a, s′, g) =1{a = πg(s)}1/A

Qg(s, a) − r(s, a) − max Qg(s′, a′)()a′ℓ(s,a,s′,g)=W(g)−W(f),X(f)⟨⟩sh∼πf ah∼U(A)h h h+1hhh

￼

￼￼

𝔼fi𝔼￼

￼

⋆

⋆

Another Example: Linear Q & V⋆

What’s the role of linear function approximation in RL?
We have linear MDP, Linear Bellman complete…
⋆The most natural one should just be Q being linear-realizable:⋆⋆ ⊤Q (s, a) = (w ) ϕ(s, a), under known feature ϕHowever generalization is impossible here:⋆• Theorem [Weisz, Amortila, Szepesvári ‘21]: There exists an MDP w/ linear Q ,d Hs.t any online RL algorithm requires Ω(min(2 ,2 )) samples to output a near
optimal policy

⋆Additional Example: Linear Q & V⋆

What’s the role of linear function approximation in RL?
⋆⋆ ⊤However, if we further assume V (s) = (θ ) ψ(s), then we will be ok:• Theorem [Du, Kakade., Lee, Lovett, Mahajan, S, Wang ’21]⋆⋆
For any MDP with both Q and V being realizable (under some known

features, e.g., RKHS), there exists an algorithm that learns with # of samples:
3d poly(H)2ϵAdditional Example: Linear Q & V⋆:= {w ϕ(s, a) : w ∈ Ballw}, V ∈⋆⊤Function classes for Q ∈⋆⋆Linear Q & V has Bilinear Structure⋆:= {θ ψ(s) : θ ∈ BallB}⊤

Key step: pre-process function class{Q, V} := {(w, θ) : ∀s, max w ϕ(s, a) = θ ψ(s)}⊤⊤a(1) on-policy Bellman error of any f := (w, θ) has bilinear form:sh,ah∼πf [w⊤ϕ(sh, ah) − rh −s′∼P(sh,ah)θ

ψ(s′)] = ⟨Wh([w, θ]) − Wh([w , θ ]), Xh([w, θ])⟩⊤⋆⊤⋆⊤(2) there exists ℓ(s, a, s′, [w, θ]) = w ϕ(s, a) − r(s, a) − θ ψ(s′), s.t.,⋆, θ ]), Xh( f )⟩, ∀[w′, θ′]⋆￼￼￼￼￼￼￼𝔼￼￼𝒬𝒱￼sh,ah∼πf [ℓ(sh, ah, s′h+1, [w′, θ′])] = ⟨Wh([w′, θ′]) − Wh([w￼𝔼𝔼⋆Final Example: Linear Mixture Model (model-based)
The linear Mixture Model:Function class: ℱ = {P : P(s′| s, a; θ) = θ ϕ(s, a, s′), θ ∈ BallW}⊤Why this model is ever interesting?
iImagine we have d simulators, P (s′| s, a), i ∈ [d],we assume the ground truth is the linear mixture of d simulators:⋆P (s′| s, a) =d∑￼￼￼￼￼

i=1

⋆i

θ [i]P (s′| s, a)The linear Mixture Model:

Function class: ℱ = {P : P(s′| s, a; θ) = θ ϕ(s, a, s′), θ ∈ BallW}⊤(Notation: f ∈ ℱ is a potential transition, and f⋆⋆:= P )

Claim 1: For any f ∈ ℱ, the on-policy Bellman error of Qf under πf has bilinear form:￼￼￼￼πf [Qf (sh, ah) − r(sh, ah) −𝔼𝔼Final Example: Linear Mixture Model (model-based)s′∼f ⋆(sh,ah)Vf (s′)] = ⟨Wh( f ) − Wh( f⋆), Xh( f )⟩The linear Mixture Model:Function class: ℱ = {P : P(s′| s, a; θ) = θ ϕ(s, a, s′), θ ∈ BallW}⊤(Notation: f ∈ ℱ is a potential transition, and f⋆⋆:= P )

Claim 2: For any f ∈ ℱ, there exists discrepancy ℓf (s, a, s′, g) to measure bilinear form:ℓf (s, a, s′, g) =s′∼g(s,a) [Vf (s′)] − Vf (s′)

s.t., πf ℓf (sh, ah, s′h+1, g) = ⟨Wh(g) − Wh( f ), Xh( f )⟩￼￼￼𝔼￼￼￼￼⋆￼

𝔼Final Example: Linear Mixture Model