Reinforcement Learning as Probabilistic Inference
report is made by
Pavel Temirchev
Deep RL
reading group
based on the research of Sergey Levine's team
Motivation
1) Sample Complexity!
2) Convergence to local optimas
The problems of standard RL:
Idea: encourage an agent to investigate all the promising strategies!
REMINDER: standard RL
Markov process:
a0
a_0
s0
s_0
a1
a_1
s1
s_1
a2
a_2
s2
s_2
p(τ)=p(s0)∏t=0Tp(at∣st)p(st+1∣st,at)
p(\tau) = p(s_0) \prod_{t=0}^T p(a_t|s_t) p(s_{t+1}|s_t, a_t)
Maximization problem:
π⋆=argmaxπ∑t=0TEst,at∼π[r(st,at)]
\pi^\star = \arg\max_\pi \sum_{t=0}^T \mathbb{E}_{s_t, a_t \sim \pi} [r(s_t, a_t)]
Qπ(st,at):=r(st,at)+∑t′=t+1TEst′,at′∼π[r(st′,at′)]
Q^\pi(s_t,a_t) := r(s_t,a_t) + \sum_{t'=t+1}^T \mathbb{E}_{s_{t'}, a_{t'} \sim \pi} [r(s_{t'}, a_{t'})]
Q-function:
Q⋆(st,at)=r(st,at)+Est+1V⋆(st+1)
Q^\star(s_t,a_t) = r(s_t,a_t) + \mathbb{E}_{s_{t+1}} V^\star(s_{t+1})
Bellman equality (optimal Q-function):
V⋆(st)=maxaQ⋆(st,a)
V^\star(s_t) = \max_a Q^\star(s_{t}, a)
τ=(s0,…,at,st,…,aT,sT)
\tau = (s_0, \dots, a_t, s_t, \dots, a_T, s_T)
Maximum Entropy RL
at∼N(⋅∣π⋆,σ2)
a_t \sim \mathcal{N}(\cdot| \pi^\star, \sigma^2)
Standard RL:
at∼expQ(st,at)
a_t \sim \exp{Q(s_t, a_t)}
Policy "proportional" to Q:
How to find such a policy?
minπKL(π(⋅∣s0)∣∣expQ(s0,⋅))=
\min_\pi\text{KL}\Big(\pi(\cdot|s_0)||\exp{Q(s_0, \cdot)}\Big) =
maxπEπ[Q(s0,a0)−logπ(a0∣s0)]=
\max_\pi \mathbb{E}_\pi \Big[ Q(s_0, a_0) - \log \pi(a_0|s_0) \Big] =
maxπEπ[∑tTr(st,at)+H(π(⋅∣st))]
\max_\pi \mathbb{E}_\pi \Big[ \sum_t^T r(s_t, a_t) {\color{pink}+ \mathcal{H} \big( \pi(\cdot|s_t) \big)}\Big]
Q(s0,a0)
Q(s_0, a_0)
RL as Probabilistic Inference
a0
a_0
s0
s_0
a1
a_1
s1
s_1
O0
\mathcal{O}_0
O1
\mathcal{O}_1
Optimality:
p(Ot=1∣st,at):=p(Ot∣st,at)=exp(r(st,at))
p(\mathcal{O}_t =1 |s_t, a_t) := p(\mathcal{O}_t |s_t, a_t) = \exp(r(s_t, a_t))
RL:
Which actions will lead as to the optimal future?
Probabilistic Inference:
Which actions were made given that the future is optimal?
p(at∣stOt:T)
p(a_t|s_t \mathcal{O}_{t:T})
Exact Probabilistic Inference
Let's find an optimal policy:
p(at∣st,Ot:T)=p(st∣Ot:T)p(st,at∣Ot:T)=
p(a_t|s_t, \mathcal{O}_{t:T}) = \frac{ \color{#00ff00} p(s_t, a_t|\mathcal{O}_{t:T})}{ \color{#ff0000} p(s_t|\mathcal{O}_{t:T})} =
=p(Ot:T)p(Ot:T∣st,at)p(at∣st)p(st)p(Ot:T∣st)p(st)p(Ot:T)
= \frac{p(\mathcal{O}_{t:T}|s_t, a_t) p(a_t|s_t) p(s_t)}{p(\mathcal{O}_{t:T})}\frac{p(\mathcal{O}_{t:T})}{p(\mathcal{O}_{t:T}|s_t) p(s_t)}
p(Ot:T)p(Ot:T∣st,at)p(at∣st)p(st)
\frac{p(\mathcal{O}_{t:T}|s_t, a_t) p(a_t|s_t) p(s_t)}{p(\mathcal{O}_{t:T})}
p(Ot:T∣st)p(st)p(Ot:T)
\frac{p(\mathcal{O}_{t:T})}{p(\mathcal{O}_{t:T}|s_t) p(s_t)}
where - prior policy
p(at∣st)
p(a_t|s_t)
if , then
p(at∣st)=∣A∣1
p(a_t|s_t) = \frac{1}{|\mathcal{A}|}
p(at∣st,Ot:T)∝p(Ot:T∣st)p(Ot:T∣st,at)
p(a_t|s_t, \mathcal{O}_{t:T}) \propto \frac{p(\mathcal{O}_{t:T}|s_t, a_t) }{p(\mathcal{O}_{t:T}|s_t)}
apply Bayes rule!
Let's introduce new notation:
αt(st,at):=p(Ot:T∣st,at)
\alpha_t(s_t, a_t) := p(\mathcal{O}_{t:T}|s_t, a_t)
βt(st):=p(Ot:T∣st)=∫αt(st,at)p(at∣st)dat
\beta_t(s_t) := p(\mathcal{O}_{t:T}|s_t) = \int \alpha_t(s_t, a_t) p(a_t|s_t)da_t
We can find all the and via Message Passing algorithm:
αt
\alpha_t
βt
\beta_t
For the timestep :
T
T
αT(sT,aT)=exp(r(sT,aT))
\alpha_T(s_T, a_T) = \exp(r(s_T, a_T))
βT(sT)=∫αT(sT,aT)p(aT∣sT)daT
\beta_T(s_T) = \int \alpha_T(s_T, a_T) p(a_T|s_T)da_T
Recursively:
αt(st,at)=∫βt+1(st+1)exp(r(st,at))p(st+1∣st,at)dst+1
\alpha_t(s_t, a_t) = \int \beta_{t+1}(s_{t+1}) \exp(r(s_t, a_t)) p(s_{t+1}|s_t, a_t)ds_{t+1}
βt(st)=∫αt(st,at)p(at∣st)dat
\beta_t(s_t) = \int \alpha_t(s_t, a_t) p(a_t|s_t)da_t
Exact Probabilistic Inference
Let's introduce new notation:
αt(st,at):=p(Ot:T∣st,at)
\alpha_t(s_t, a_t) := p(\mathcal{O}_{t:T}|s_t, a_t)
βt(st):=p(Ot:T∣st)=∫αt(st,at)p(at∣st)dat
\beta_t(s_t) := p(\mathcal{O}_{t:T}|s_t) = \int \alpha_t(s_t, a_t) p(a_t|s_t)da_t
We can find all the and via Message Passing algorithm:
αt
\alpha_t
βt
\beta_t
For the timestep :
T
T
αT(sT,aT)=exp(r(sT,aT))
\alpha_T(s_T, a_T) = \exp(r(s_T, a_T))
βT(sT)=∫αT(sT,aT)p(aT∣sT)daT
\beta_T(s_T) = \int \alpha_T(s_T, a_T) p(a_T|s_T)da_T
Recursively:
αt(st,at)=∫βt+1(st+1)exp(r(st,at))p(st+1∣st,at)dst+1
\alpha_t(s_t, a_t) = \int \beta_{t+1}(s_{t+1}) \exp(r(s_t, a_t)) p(s_{t+1}|s_t, a_t)ds_{t+1}
βt(st)=∫αt(st,at)p(at∣st)dat
\beta_t(s_t) = \int \alpha_t(s_t, a_t) p(a_t|s_t)da_t
Exact Probabilistic Inference
Soft Q and V functions
Qsoft(st,at):=logαt(st,at)
Q^{soft}(s_t, a_t) := \log\alpha_t(s_t, a_t)
Vsoft(st):=logβt(st)
V^{soft}(s_t) := \log\beta_t(s_t)
Recursively:
Vsoft(st)=logEp(at∣st)[expQsoft(st,at)]
V^{soft}(s_t) =\log \mathbb{E}_{p(a_t|s_t)} [\exp Q^{soft}(s_t, a_t)]
- soft maximum
Qsoft(st,at)=r(st,at)+logEp(st+1∣st,at)[expVsoft(st+1)]
Q^{soft}(s_t, a_t) = r(s_t, a_t) + \log \mathbb{E}_{p(s_{t+1}|s_t, a_t)} [\exp V^{soft}(s_{t+1})]
kinda Bellman equation
We can find analogues in the log-scale:
Soft and Hard Q and V functions
"Hard" Q and V functions:
V⋆(st)=maxatQ⋆(st,at)
V^\star(s_t) =\max_{a_t} Q^\star(s_t, a_t)
Q⋆(st,at)=r(st,at)+Ep(st+1∣st,at)V⋆(st+1)
Q^\star(s_t, a_t) = r(s_t, a_t) + \mathbb{E}_{p(s_{t+1}|s_t, a_t)} V^\star(s_{t+1})
Vsoft(st)=logEp(at∣st)[expQsoft(st,at)]
V^{soft}(s_t) =\log \mathbb{E}_{p(a_t|s_t)} [\exp Q^{soft}(s_t, a_t)]
Qsoft(st,at)=r(st,at)+logEp(st+1∣st,at)[expVsoft(st+1)]
Q^{soft}(s_t, a_t) = r(s_t, a_t) + \log \mathbb{E}_{p(s_{t+1}|s_t, a_t)} [\exp V^{soft}(s_{t+1})]
"Soft" analogues:
Q⋆(st,at)=r(st,at)+Ep(st+1∣st,at)maxat+1Q⋆(st+1,at+1)
Q^\star(s_t, a_t) = r(s_t, a_t) + \mathbb{E}_{p(s_{t+1}|s_t, a_t)} \max_{a_{t+1}} Q^\star(s_{t+1}, a_{t+1})
Qsoft(st,at)≈r(st,at)+maxst+1maxat+1Qsoft(st+1,at+1)
Q^{soft}(s_t, a_t) \approx r(s_t, a_t) + \max_{s_{t+1}} \max_{a_{t+1}} Q^{soft}(s_{t+1}, a_{t+1})
maxst+1
\max_{s_{t+1}}
What is being optimized?
p(τ∣O0:T)=p(O0:T)p(τ,O0:T)
p(\tau|\mathcal{O}_{0:T}) =\frac{p(\tau,\mathcal{O}_{0:T})} {p(\mathcal{O}_{0:T})}
Let's analyze an "exact variational inference" procedure:
true conditional
joint
evidence
argminqKL(q(τ)∣∣p(τ∣O0:T))=argminp(at∣st,Ot:T)KL(p(τ∣O0:T)∣∣p(τ∣O0:T))
\arg\min_q \text{KL}\big(q(\tau)||p(\tau|\mathcal{O}_{0:T})\big)
=
\arg\min_{p(a_t|s_t, \mathcal{O}_{t:T})}
\text{KL}\big(p(\tau|\mathcal{O}_{0:T})||p(\tau|\mathcal{O}_{0:T})\big)
=argminp(at∣st,Ot:T)KL(p(τ∣O0:T)∣∣p(τ,O0:T))
=
\arg\min_{p(a_t|s_t, \mathcal{O}_{t:T})}
\text{KL}\big(p(\tau|\mathcal{O}_{0:T})||p(\tau,\mathcal{O}_{0:T})\big)
What is being optimized?
where the joint ("exact") distribution is:
p(τ,O0:T)=p(s0)∏t=0Tp(at∣st)p(st+1∣st,at)exp(r(st,at))
p({\color{#ff0000}\tau},{\color{#00ff00}\mathcal{O}_{0:T}}) = {\color{#ff0000} p(s_0)\prod_{t=0}^Tp(a_t|s_t)p(s_{t+1}|s_t,a_t)} {\color{#00ff00}\exp\big(r(s_t, a_t)\big)}
and the variational one is:
p(τ∣O0:T)=p(s0∣O0:T)∏t=0Tp(at∣st,O0:T)p(st+1∣st,at,O0:T)
p(\tau|\mathcal{O}_{0:T}) = p(s_0|\mathcal{O}_{0:T})\prod_{t=0}^Tp(a_t|s_t,\mathcal{O}_{0:T})p(s_{t+1}|s_t,a_t,\mathcal{O}_{0:T})
we tried to find a policy which is optimal
only in an optimal environment!
We can fix this!
KL(p(τ∣O0:T)∣∣p(τ,O0:T))→minp(at∣st,Ot:T)
\text{KL}\big(p(\tau|\mathcal{O}_{0:T})||p(\tau,\mathcal{O}_{0:T})\big) \rightarrow \min_{p(a_t|s_t, \mathcal{O}_{t:T})}
p(τ,O0:T)=p(s0)∏t=0Tp(at∣st)p(st+1∣st,at)exp(r(st,at))
p(\tau,\mathcal{O}_{0:T}) = p(s_0)\prod_{t=0}^Tp(a_t|s_t)p(s_{t+1}|s_t,a_t) \exp\big(r(s_t, a_t)\big)
p(τ∣O0:T)=p(s0∣O0:T)∏t=0Tp(at∣st,O0:T)p(st+1∣st,at,O0:T)
p(\tau|\mathcal{O}_{0:T}) = p(s_0|{\color{#ff0000}\mathcal{O}_{0:T}})\prod_{t=0}^Tp(a_t|s_t,\mathcal{O}_{0:T})p(s_{t+1}|s_t,a_t,{\color{#ff0000}\mathcal{O}_{0:T}})
Variational Inference
The form of the - is our choice
q(τ)=p(s0)∏t=0Tπ(at∣st)p(st+1∣st,at)
q(\tau) = p(s_0)\prod_{t=0}^T\pi(a_t|s_t) p(s_{t+1}|s_t,a_t)
Minimization problem for VI
KL(q(τ)∣∣p(τ,O0:T))→minq
\text{KL}\big(q(\tau)||p(\tau,\mathcal{O}_{0:T})\big)
\rightarrow \min_q
q(τ)
q(\tau)
is a distribution over
ACHIEVABLE trajectories
q
q
fix the dynamics!
Then:
minqKL(q(τ)∣∣p(τ,O0:T))=−minqEqlogq(τ)p(τ,O0:T)=
\min_q \text{KL}\big(q(\tau)||p(\tau,\mathcal{O}_{0:T})\big) = - \min_q \mathbb{E}_q \log \frac{p(\tau,\mathcal{O}_{0:T})}{q(\tau)} =
Maximum Entropy RL Objective
=maxqEq[logp(s0)+∑t(logp(st+1∣st,at)+r(st,at))−
= \max_q \mathbb{E}_q \Big[ \log p(s_0)+\sum_{t} \big( \log p(s_{t+1}|s_t,a_t) + r(s_t, a_t) \big) -
−logp(s0)−∑t(logp(st+1∣st,at)−logπ(at∣st))]=
- \log p(s_0)-\sum_{t} \big( \log p(s_{t+1}|s_t,a_t) - \log \pi(a_t| s_t) \big) \Big]=
=maxπEπ∑t[r(st,at)+H(π(⋅∣st))]
= \max_\pi \mathbb{E}_\pi \sum_{t}\Big[ r(s_t, a_t) + \mathcal{H}\big( \pi(\cdot| s_t)\big) \Big]
Variational Inference
Returning to the Q and V functions
This objective can be rewritten as follows:
Vsoft(st)=log∫expQsoft(st,at)dat
V^{soft}(s_t) =\log \int \exp Q^{soft}(s_t, a_t) da_t
Qsoft(st,at)=r(st,at)+Ep(st+1∣st,at)Vsoft(st+1)
Q^{soft}(s_t, a_t) = r(s_t, a_t) + \mathbb{E}_{p(s_{t+1}|s_t, a_t)} V^{soft}(s_{t+1})
check it yourself!
π(at∣st)=exp(Vsoft(st))exp(Qsoft(st,at))
\pi(a_t|s_t) =\frac{\exp(Q^{soft}(s_t, a_t))}{\exp(V^{soft}(s_t))}
∑t=0TEst[−KL(π(at∣st)∣∣exp(Vsoft(st))exp(Qsoft(st,at)))+Vsoft(st)]→maxπ
\sum_{t=0}^T\mathbb{E}_{s_t} \Big[ -\text{KL}\Big(\pi(a_t|s_t)||\frac{\exp(Q^{soft}(s_t, a_t))}{\exp(V^{soft}(s_t))}\Big) + V^{soft}(s_t) \Big] \rightarrow \max_\pi
Then the optimal policy is:
where
- soft maximum
- normal Bellman equation
VI with function approximators
(neural nets)
- Maximum Entropy Policy Gradients
- Soft Q-learning
https://arxiv.org/abs/1702.08165 - Soft Actor-Critic
https://arxiv.org/abs/1801.01290
Maximum Entropy Policy Gradients
Eτ∼πθ∑t=0T[r(st,at)+H(πθ(⋅∣st))]→maxθ
\mathbb{E}_{\tau \sim \pi_\theta} \sum_{t=0}^T\Big[ r(s_t, a_t) + \mathcal{H}\big(\pi_\theta(\cdot|s_t)\big) \Big] \rightarrow \max_\theta
Directly maximize entropy-augmented objective
over policy parameters :
For gradients, use log-derivative trick:
∑t=0TE(st,at)∼qθ[∇θlogπθ(at∣st)∑t′=tT(r(st′,at′)−logπθ(at′∣st′)−b(st′))]
\sum_{t=0}^T\mathbb{E}_{(s_t,a_t) \sim q_\theta} \Big[ \nabla_\theta \log\pi_\theta(a_t|s_t) \sum_{t'=t}^T\Big( r(s_{t'}, a_{t'}) -\log\pi_\theta(a_{t'}|s_{t'}) - b(s_{t'}) \Big)\Big]
θ
\theta
- on-policy
- unimodal policies
Soft Q-learning
Train Q-network with parameters :
ϕ
\phi
E(st,at,st+1)∼D[Qϕsoft(st,at)−(r(st,at)+Vϕsoft(st+1))]2→minϕ
\mathbb{E}_{(s_t,a_t, s_{t+1}) \sim \mathcal{D}} \Big[ Q^{soft}_\phi(s_t, a_t) - \Big( r(s_t, a_t) + V^{soft}_\phi(s_{t+1})\Big) \Big]^2\rightarrow \min_\phi
use replay buffer
where
Vϕsoft(st)=log∫expQϕsoft(st,at)dat
V^{soft}_\phi(s_t) =\log \int \exp Q^{soft}_\phi(s_t, a_t) da_t
for continuous actions use
importance sampling
Policy is implicit
π(at∣st)=exp(Qϕsoft(st,at)−Vϕsoft(st))
\pi(a_t|s_t) = \exp\big(Q^{soft}_\phi(s_t, a_t) - V^{soft}_\phi(s_t)\big)
for samples use SVGD
or MCMC :D
Soft Q-learning
Exploration
Robustness
Multimodal Policy
Soft Actor-Critic
Train Q- and V-networks jointly with policy
E(st,at,st+1)∼D[Qϕsoft(st,at)−(r(st,at)+Vψsoft(st+1))]2→minϕ
\mathbb{E}_{(s_t,a_t, s_{t+1}) \sim \mathcal{D}} \Big[ Q^{soft}_\phi(s_t, a_t) - \Big( r(s_t, a_t) + V^{soft}_\psi(s_{t+1})\Big) \Big]^2\rightarrow \min_\phi
Q-network loss:
V-network loss:
V^soft(st)=Eat∼πθ[Qϕsoft(st,at)−logπθ(at∣st)]
\hat{V}^{soft}(s_t) = \mathbb{E}_{a_t \sim \pi_\theta} \Big[ Q^{soft}_\phi(s_t, a_t) - \log\pi_\theta(a_t|s_t) \Big]
Est∼D[V^soft(st)−Vψsoft(st)]2→minψ
\mathbb{E}_{s_t \sim \mathcal{D}} \Big[ \hat{V}^{soft}(s_t) - V^{soft}_\psi(s_{t}) \Big]^2 \rightarrow \min_\psi
Objective for the policy:
Est∼D,at∼πθ[Qϕsoft(st,at)−logπθ(at′∣st)]→maxθ
\mathbb{E}_{s_t \sim \mathcal{D}, \;a_t \sim \pi_\theta} \Big[ Q^{soft}_\phi(s_t, a_t) -\log\pi_\theta(a_{t'}|s_{t})\Big] \rightarrow \max_\theta
E(st,at,st+1)∼D[Qϕsoft(st,at)−(r(st,at)+Vψsoft(st+1))]2→minϕ
\mathbb{E}_{(s_t,a_t, s_{t+1}) \sim \mathcal{D}} \Big[ Q^{soft}_\phi(s_t, a_t) - \Big( r(s_t, a_t) + V^{soft}_\psi(s_{t+1})\Big) \Big]^2\rightarrow \min_\phi
Q-network loss:
V-network loss:
V^soft(st)=Eat∼πθ[Qϕsoft(st,at)−logπθ(at∣st)]
\hat{V}^{soft}(s_t) = \mathbb{E}_{a_t \sim \pi_\theta} \Big[ Q^{soft}_\phi(s_t, a_t) - \log\pi_\theta(a_t|s_t) \Big]
Est∼D[V^soft(st)−Vψsoft(st)]2→minψ
\mathbb{E}_{s_t \sim \mathcal{D}} \Big[ \hat{V}^{soft}(s_t) - V^{soft}_\psi(s_{t}) \Big]^2 \rightarrow \min_\psi
Objective for the policy:
Est∼D,at∼πθ[Qϕsoft(st,at)−logπθ(at′∣st)]→maxθ
\mathbb{E}_{s_t \sim \mathcal{D}, \;a_t \sim \pi_\theta} \Big[ Q^{soft}_\phi(s_t, a_t) -\log\pi_\theta(a_{t'}|s_{t})\Big] \rightarrow \max_\theta
Soft Actor-Critic
Soft Actor-Critic
Thank you for your attention!
and visit our seminars in RL Reading Group
telegram: https://t.me/theoreticalrl
REFERENCES:
Soft Q-learning:
https://arxiv.org/pdf/1702.08165.pdf
Soft Actor Critic:
https://arxiv.org/pdf/1801.01290.pdf
Big Review on Probabilistic Inference for RL:
https://arxiv.org/pdf/1805.00909.pdf
Implementation on TensorFlow:
https://github.com/rail-berkeley/softlearning
Implementation on Catalyst.RL:
https://github.com/catalyst-team/catalyst/tree/master/examples/rl_gym
Hierarchical policies (further reading):
Reinforcement Learning as Probabilistic Inference report is made by Pavel Temirchev Deep RL reading group based on the research of Sergey Levine's team