E

cient Reinforcement Learning via
Representation Learning
Wen Sunffi

Joint work with Masatoshi Uehara (Cornell) & Xuezhou Zhang (Princeton)

Empirical RL for large-scale problems[AlphaGo, Silver et.al, 15][OpenAI Five, 18][OpenAI,19]Rich (nonlinear) function approximation + RL can work well w/ enough samples

Can we design provably ecient algorithms for

ffi

Rich Function Approx + RL ?

Can we design provably ecient algorithms forRich Function Approx + RL ?
RL

Environment w/complex highdim dataffi

𝒟DatasetCan we design provably ecient algorithms forRich Function Approx + RL ?
RL

Environment w/
complex highdim datausing ϕOur solution:
Representation Learning
Oracle:Datasetdffi

𝒟ϕ(s, a) ∈ ℝ

Episodic In nite Horizon Discounted MDPs
Policy: state to action

π(s) → a

Reward & Next Statefi

￼r(s, a), s′ ∼ P( ⋅ | s, a)Episodic In nite Horizon Discounted MDPs
Policy: state to action

π(s) → a

Reward & Next Stater(s, a), s′ ∼ P( ⋅ | s, a)
Objective:max J(π; P, r), where J(π; P, r) :=fi

￼

𝔼

π

[r(s0, a0) + γr(s1, a1) + γ r(s2, a2) + … | a ∼ π, P]
2Episodic In nite Horizon Discounted MDPs
Policy: state to action

π(s) → a

Reward & Next Stater(s, a), s′ ∼ P( ⋅ | s, a)
Objective:

max J(π; P, r), where J(π; P, r) :=π[r(s0, a0) + γr(s1, a1) + γ r(s2, a2) + … | a ∼ π, P]

fifi

￼𝔼

Assume xed initial state s02

Low-rank MDPs′(s, a)Transition matrixP∈ℝSA×S￼￼￼￼rank dhas

P ⋆(s′| s, a)s′(s, a)=⋆⋆ϕ (s, a)μ (s′)

Low-rank MDPs′(s, a)(s, a)P ⋆(s′| s, a)=Transition matrix

P∈ℝ

SA×Ss′⋆

μ (s′)⋆ϕ (s, a)hasrank d⋆

⋆￼

￼￼￼

￼￼

⋆

⊤ ⋆

∀s, a, s′, P (s′| s, a) = μ (s′) ϕ (s, a)￼∃μ , ϕ :⋆Low-rank MDPs′(s, a)(s, a)P ⋆(s′| s, a)=Transition matrix

P∈ℝSA×Ss′⋆

μ (s′)⋆ϕ (s, a)hasrank d⋆⋆∃μ , ϕ :⋆⋆⊤ ⋆∀s, a, s′, P (s′| s, a) = μ (s′) ϕ (s, a)

Low-rank MDP ≠ Linear MDPs (Jin et al, Yang & Wang)
￼￼￼￼￼

￼

￼

Linear MDP = low-rank + known ϕ⋆Low-rank MDPs′(s, a)(s, a)P ⋆(s′| s, a)=Transition matrix

P∈ℝ

SA×S

Representation Learnings′⋆

μ (s′)⋆ϕ (s, a)hasrank d⋆⋆∃μ , ϕ :⋆⋆⊤ ⋆∀s, a, s′, P (s′| s, a) = μ (s′) ϕ (s, a)

Low-rank MDP ≠ Linear MDPs (Jin et al, Yang & Wang)￼￼￼￼￼￼

￼Linear MDP = low-rank + known ϕ⋆Low-rank MDP is generale.g., Latent variable models where Z is the discrete latent space⋆

ϕ (s, a) ∈ Δ(Z)zzzsssaa[slide credit: Akshay Krishnumurthy]Low-rank MDP is generale.g., Latent variable models where Z is the discrete latent space⋆

ϕ (s, a) ∈ Δ(Z)zzzsssaa⋆⋆￼

Given s, a: z ∼ ϕ (s, a), s′ ∼ ν (z)[slide credit: Akshay Krishnumurthy]Provably ecient learning in low-rank mdp is plausible

Olive [JKALS,17]
Witness Rank [SJKAL,19]
BLin-UCB [DKLLMSW,21]
Mo le [MCKJA,21]

ffi

ff

ff

ff

ff

ff

ff

FLAMBE [AKKS,20]

Setting

Sample Complexity

Computation

Low Bellman rankd2A
ϵ 2(1 − γ)4

Ine icient

Low Witness rankd2A
ϵ 2(1 − γ)4

Ine icient

Bilinear Classd2A
ϵ 2(1 − γ)7

Ine icient

Low-nonnegative-rank MDP

d 6 A 13ϵ 2η 5(1 − γ)5Oracle-e icient

Low-rank MDP

d7A9
ϵ 10(1 − γ)22

Oracle-e icientProvably ecient learning in low-rank mdp is plausible

Olive [JKALS,17]
Witness Rank [SJKAL,19]
BLin-UCB [DKLLMSW,21]
Mo le [MCKJA,21]
FLAMBE [AKKS,20]

SettingSample Complexity

Computation

Low Bellman rankd2A
ϵ 2(1 − γ)4

Ine icient

Low Witness rankd2A
ϵ 2(1 − γ)4

Ine icient

Bilinear Classd2A
ϵ 2(1 − γ)7

Ine icient

Low-nonnegative-rank MDP

d 6 A 13ϵ 2η 5(1 − γ)5Oracle-e icient

Low-rank MDP

d7A9
ϵ 10(1 − γ)22

Oracle-e icient

ffiffiffff

ff

ff

ff

ff

FLAMBE is oracle-e cient and was state-of-art on low-rank MDPOur learning setting
1. Realizable hypothesis classes Γ, Φ⋆⋆μ ∈ Γ, ϕ ∈ ΦOur learning setting
1. Realizable hypothesis classes Γ, Φ⋆⋆μ ∈ Γ, ϕ ∈ Φ
2. Computation oracle:
Maximum Likelihood Estimation (MLE):(μ,̂ ϕ)̂ := arg max

￼

μ,ϕn∑
i=1ln (μ(s′i )ϕ(si, ai))

Our learning setting
1. Realizable hypothesis classes Γ, Φ⋆⋆μ ∈ Γ, ϕ ∈ Φ
2. Computation oracle:
Maximum Likelihood Estimation (MLE):(μ,̂ ϕ)̂ := arg max
μ,ϕn∑i=1ln (μ(s′i )ϕ(si, ai))

3. Learning Goal:￼Finding near-optimal policy w/ (tight) poly(A, d,1/(1 − γ), ln( | Φ | | Γ | ))Our algorithm: Rep-UCB(UCB-driven Representation Learning for online RL)
At iteration n:Our algorithm: Rep-UCB(UCB-driven Representation Learning for online RL)
At iteration n:nData generation from π :s ∼ d , a ∼ Uniform( ), s′ ∼ P ⋆( . | s, a)πn￼￼￼𝒜￼￼𝒜￼a′ ∼ Uniform( ), s′′ ∼ P ⋆( . | s′, a′)Our algorithm: Rep-UCB(UCB-driven Representation Learning for online RL)
At iteration n:n

Data generation from π :s ∼ d , a ∼ Uniform( ), s′ ∼ P ⋆( . | s, a)
πn

￼

￼

￼

￼

￼

￼

￼
𝒜

￼

￼

￼

𝒜

￼

𝒟𝒟￼
￼

𝒟𝒟a′ ∼ Uniform( ), s′′ ∼ P ⋆( . | s′, a′)Data Aggregation:n=′n =n−1 + {s, a, s′}′n−1 + {s′, a′, s′′}

Our algorithm: Rep-UCB(UCB-driven Representation Learning for online RL)At iteration n:nData generation from π :Data Aggregation:s ∼ d , a ∼ Uniform( ), s′ ∼ P ⋆( . | s, a)πna′ ∼ Uniform( ), s′′ ∼ P ⋆( . | s′, a′)n=′n =n−1 + {s, a, s′}′n−1 + {s′, a′, s′′}Representation / model
Learning (MLE)P̂ := (μ,̂ ϕ)̂ = arg max￼

𝒟

￼

￼￼

￼

￼

￼

￼
𝒜

￼

￼

𝒜

￼

𝔼

￼

￼

￼

𝒟𝒟￼

𝒟𝒟

𝒟μ,ϕ⊤

ln(μ(s′)ϕ(s,a))+′nnOur algorithm: Rep-UCB(UCB-driven Representation Learning for online RL)
At iteration n:nData generation from π :Data Aggregation:s ∼ d , a ∼ Uniform( ), s′ ∼ P ⋆( . | s, a)
πnn=a′ ∼ Uniform( ), s′′ ∼ P ⋆( . | s′, a′)′n =n−1 + {s, a, s′}′n−1 + {s′, a′, s′′}(Linear bandit style) bonus under ϕ:̂
−1 ̂̂b(s, a) = c ϕ(s, a)Σ ϕ(s, a)

Σ=∑￼

𝒟￼

￼￼￼￼

￼

￼
𝒜

￼

￼

𝒜￼

𝔼￼￼￼

𝒟𝒟￼𝒟𝒟𝒟𝒟s,a∈̂ a)ϕ(s,̂ a)ϕ(s,⊤nRepresentation / model
Learning (MLE)P̂ := (μ,̂ ϕ)̂ = arg maxμ,ϕ⊤ln(μ(s′)ϕ(s,a))+′nnOur algorithm: Rep-UCB(UCB-driven Representation Learning for online RL)
At iteration n:nData generation from π :Data Aggregation:s ∼ d , a ∼ Uniform( ), s′ ∼ P ⋆( . | s, a)πnn=a′ ∼ Uniform( ), s′′ ∼ P ⋆( . | s′, a′)π−1 ̂̂
b(s, a) = c ϕ(s, a)Σ ϕ(s, a)= max J(π; P,̂ r + b)πΣ=∑￼𝒟￼￼￼￼￼￼￼𝒜￼￼𝒜￼𝔼￼￼￼

𝒟𝒟￼

𝒟𝒟𝒟s,a∈𝒟′n−1 + {s′, a′, s′′}(Linear bandit style) bonus under ϕ:̂

Plan w/ reward + bonus
n+1′n =n−1 + {s, a, s′}̂ a)ϕ(s,̂ a)ϕ(s,⊤nRepresentation / model
Learning (MLE)P̂ := (μ,̂ ϕ)̂ = arg max
μ,ϕ⊤ln(μ(s′)ϕ(s,a))+′nn

PAC-Bound of Rep-UCB in low-rank MDPAssume trajectory-reward is normalized in [0,1]. W/ highprobability, it nds an ϵ near optimal policy, with # of samples:Xxld A||||Õ⋅lnΓΦ()
25ϵ(1−γ)
()

fi

42

PAC-Bound of Rep-UCB in low-rank MDPAssume trajectory-reward is normalized in [0,1]. W/ highprobability, it nds an ϵ near optimal policy, with # of samples:
Xxl

d A||||Õ⋅lnΓΦ
()
25ϵ(1−γ)()42For reference, prior SOTAFLAMBE has the following bound:d A||||Õ⋅lnΓΦ()1022( ϵ (1 − γ))fi7 9

Applying our new techniques to O

ine RLO ine RL: we only have a static= {s, a, s′}, where⋆(s, a) ∼ πb, s′ ∼ P ( . | s, a)￼￼𝒟dataset

ffl

ffl

ffl

[Image from BAIR blog post: https://bair.berkeley.edu/blog/2020/12/07/o ine/]Applying our new techniques to O

ine RLO ine RL: we only have a static= {s, a, s′}, where⋆(s, a) ∼ πb, s′ ∼ P ( . | s, a)dataset￼￼𝒟𝒟

Goal: learn to nd some high quality policy solely fromffl

ffl

fiffl

[Image from BAIR blog post: https://bair.berkeley.edu/blog/2020/12/07/o ine/]Coverage condition of the oine data

A comparator policy π is covered by o ine data if the relative condition number is bounded:x ( s,a∼dπϕ (s, a)ϕ (s, a) ) x
Cπ* := max<∞x x ⊤ ( s,a∼d πbϕ ⋆(s, a)ϕ ⋆(s, a)⊤) x
⊤⋆

⋆

⊤

ffl

ffl

𝔼

𝔼

Note coverage is wrt true representation only!

Coverage condition of the oine data

A comparator policy π is covered by o ine data if the relative condition number is bounded:x ( s,a∼dπϕ (s, a)ϕ (s, a) ) x
Cπ* := max<∞x x ⊤ ( s,a∼d πbϕ ⋆(s, a)ϕ ⋆(s, a)⊤) x⊤⋆⋆⊤

Note coverage is wrt true representation only!

Goal is to learn robustly, i.e., as long as there is a high quality policy that isπb

ffl

ffl

𝔼𝔼

covered by d , we want to compete against it!

The Rep-LCB Algorithm(Low con dence bound driven oine RL)

1. Representation / model
Learning (MLE) underP̂ := (μ,̂ ϕ)̂ = arg max

￼fflfi𝔼𝒟𝒟μ,ϕ⊤ln(μ(s′) ϕ(s, a))The Rep-LCB Algorithm(Low con dence bound driven o

2. Penalty w/ ϕ̂−1 ̂̂
b(s, a) = c ϕ(s, a)Σ ϕ(s, a)Σ=∑￼fflfi𝔼𝒟𝒟𝒟s,a∈̂ a)ϕ(s,̂ a)ϕ(s,⊤nine RL)

1. Representation / model
Learning (MLE) underP̂ := (μ,̂ ϕ)̂ = arg max
μ,ϕ⊤ln(μ(s′) ϕ(s, a))The Rep-LCB Algorithm(Low con dence bound driven o

2. Penalty w/ ϕ̂

3. Conservative Plan

−1 ̂̂b(s, a) = c ϕ(s, a)Σ ϕ(s, a)

w/ reward + penalty

π ̂ = max J(π; P,̂ r − b)Σ=π∑￼fflfi𝔼𝒟𝒟𝒟s,a∈̂ a)ϕ(s,̂ a)ϕ(s,⊤nine RL)

1. Representation / model
Learning (MLE) underP̂ := (μ,̂ ϕ)̂ = arg max
μ,ϕ⊤ln(μ(s′) ϕ(s, a))The guarantee of Rep-LCB
Assume the behavior policy πb(a | s) ≥ w, ∀s; W/ high probability, for ALL
comparator policy π* (include history-dependent ones):dJ(π*; r) − J(π;̂ r) ≤ Õ1.5( (1 − γ)
2wCπ*⋅ ln( | Φ | | Γ | )n)The guarantee of Rep-LCB
Assume the behavior policy πb(a | s) ≥ w, ∀s; W/ high probability, for ALL
comparator policy π* (include history-dependent ones):dJ(π*; r) − J(π;̂ r) ≤ Õ1.5( (1 − γ)2wCπ*⋅ ln( | Φ | | Γ | )n)The guarantee of Rep-LCB
Assume the behavior policy πb(a | s) ≥ w, ∀s; W/ high probability, for ALL
comparator policy π* (include history-dependent ones):dJ(π*; r) − J(π;̂ r) ≤ Õ1.5( (1 − γ)2wCπ*⋅ ln( | Φ | | Γ | )n)(prior work CPPO [Uehara & Sun, 21] can achieve similar guarantee,
but is a version-space alg)Summary
1. Improved online Representation Learning algorithm for low-rank MDP:
Oracle-e cient + tight sample complexity
2. New oine RL algorithm for low-rank MDP:

Partial coverage + Oracle-e cient

ffi

ffl

ffi

Rep-UCB / LCB: https://arxiv.org/pdf/2110.04652.pdf