Reinforcement Learning PPO and SAC

Notation

MDP \((S,A,P,r,\gamma)\)

• State Space \(S\)
• Action Space \(A\)
• Transition Probability Function \(P(s'|s,a)\)
• Reward function \(r\)
• Discount factor \(\gamma\)
• Stochastic Policy \(\pi\)
• Objective \(J(\pi)=E_{\tau\sim \pi} [\sum_{t=0}^{T-1} \gamma ^t r(s_t,a_t)]\).

Bellman Function

• Value function \[ V^{\pi}(s)=E_{a\sim \pi,s'\sim P}[r(s,a)+\gamma V^{\pi}(s')] \]

• Q function \[ Q^{\pi}(s,a)=E_{s'\sim P,a'\sim \pi}[r(s,a)+\gamma Q^{\pi}(s',a')] \]

• Advantage \(A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s)\).

And that leads us to the trivial policy gradient of \[ \nabla_{\theta} J(\pi_{\theta})=E_{\pi_{\theta}}[\nabla_{\theta} A^{\pi_{\theta}}(s,a)\log \pi_{\theta}(a_t|s_t)] \]

PPO

PPO is a strong and robust baseline for on-policy RL, that does not usually needs too careful tuning. However, it is sensitive to reward design and poor on data efficiency due to its on-policy feature.

General Advantage Estimation

In the original advantage calculation, we have \(A_t=Q(s_{t},a_t)-V(s_t)=r(s_t,a_t)-V(s,t)+\gamma E[Q(s',a')]\).

A simple try to eliminate the need of \(Q\) is to introduce TD error \(\delta_t=r_t+\gamma V(s_{t+1})-V(s_t)\), which is an unbiased estimation of the advantage. However, in reality we use critic \(V_{\phi}(s_t)\) to approximate the true value \(V^{\pi}(s_t)\), which introduces high bias (it could also be explained by that \(\delta_t\) does not see future advantage). Thus \(\hat A^{TD}_t=\delta_t\) is of high bias.

Another try is the trivial REINFOCE algorithm where we use a Monto-Carlo estimation \(\hat A^{MC}_t=\sum \gamma^l r_{t+l}-V_t=\sum \gamma ^l \delta_{t+l}-V_{t}\). While significantly lowering down the bias, we face high variance of Monto-Carlo estimation.

PPO solves this by using a discount factor \(\lambda\) to seek a balance between the former two estimations, where \[ \hat A^{GAE}_t= \sum (\gamma \lambda)^l\delta_{t+l}-V_t \]

Importance Sampling & Clipping

Definition of importance Sampling: \(E_{x\sim p}[f(x)]=E_{x\sim q}[\frac{p(x)}{q(x)}f(x)]\). Thus we call \(\frac{p(x)}{q(x)}\) as the importance weight.

At some point of the learning process where we have a policy \(\pi_{old}\), and we want to obtain a new policy. The loss of the new policy \(\pi_{\theta}\) is \(L=E_{(s,a)\sim \pi_{\theta}} A(s,a)\). As only the old policy is available when rolling out trajectories, we could leverage importance sampling to fix this issue.

That is, \(L=E_{(s,a)\sim \pi_{\theta}} A(s,a)=E_{(s,a)\sim \pi_{old}}\frac{\pi_\theta(s,a)}{\pi_{old}(s,a)} A(s,a)\). And this is the unclipped surrogate loss of PPO. \(\nabla_{\theta} L(\theta)\) brings the same expression as the policy gradient mentioned above.

However, the importance weight could have high variance with unstable behavior. Therefore, TRPO remains the new policy to stay in a trust region where \(D_{KL}(\pi_{old}|\pi_{\theta})<\delta\). Following TRPO, PPO introduces clipped surrogate loss with \(L=E_t [\min(\rho A,clip(\rho,1-\epsilon,1+\epsilon)A)]\), where \(\rho\) is the important weight and \(\epsilon\) is a hyperparameter often set to \(0.2\).

Algorithm

When it comes to the final algorithm, we usually update several time for one simulation rollout.

• Calculate \(\delta,\hat A^{GAE}\) and \(\log \pi_{old}(s,a)\).
• Update \(k\) epochs:
  • For each minibatch:
    • Calculate the new \(\log \pi_{\theta}(s,a)\).
    • Compute the clipped policy surrogate loss.
    • Compute critic loss.
    • Compute the entropy loss \(H(\pi_{\theta}(\cdot |s_t))\).
    • Update actor and critic.

Q-Learning

Off-policy RL suffers from a distribution shift that we must use out-dated \((s,a,r,s')\) for value estimation. We maintain a replay buffer which we sampled from when training the policy.

The main idea of Q-learning is trying to bootstraps a \(Q\) function using the replay buffer. If the critic could precisely represent \(Q(s,a)\) then the actor could be easily learned by \(\pi^*(s)=\arg\max_a Q(s,a)\).

DDPG

From the Bellman equation, for a pair of \((s,a,r,s')\), we obtain \(Q^*(s,a)=r+\gamma Q^*(s',\pi^* (s'))\). Thus a trivial method to learn the Q function is to see it as a self-supervise regression problem: \[ L_{\theta}=||Q_{\theta}(s,a)-\operatorname{nograd}(r+\gamma Q_{\theta}(s',\pi_{\theta}(s')))||^2 \] As the Q-function is the dominating factor here in DDPG, we delay the update of actor. In other words, we update the critic multiple times before one actor update.

Though being classic in RL theory and proved convergence, the overestimation of \(Q\) and deterministic exploitation turns out to be harmful.

Overestimation: Once we overestimate \(Q\), the actor will be moving toward the overestimated state, causing further overestimation.

TD3

TD3 includes two techniques aiming to solve the problem of Q-Learning.

The first is twin critic. It learns two \(Q\) functions \(Q_1,Q_2\) in parrel and use \(\min(Q_1,Q_2)\) when calculating \(Q\) value. (Therefore, the calculation inside the no-grad becomes \(r+\gamma \min(Q_1,Q_2)\)).

The second is actor smoothing, where we set actor to be \(\pi(s')+\epsilon\) (\(\epsilon \sim N(0,c)\)) when calculating the no-grad part. That encourages exploration and smooth the Q-value peak caused by deterministic exploitation.

SAC

SAC encourages exploration by putting entropy into consideration, where \[ J(\pi)=E[\sum_t r_t+\alpha H(\pi(\cdot|s_t))] \] Then, the target of Q-learning (the no-grad part) becomes \(r+\gamma(Q(s',a')-\alpha \log \pi(a'|s'))\). That increase the exploration and randomness of the policy

Also, the actor optimize the policy by setting the loss \(E[\alpha \log \pi(a|s)-Q(s,a)]\).

The \(\alpha\) value is often calculated automatically, where we set loss \(J(\alpha)=E[-\alpha(\log \pi(a|s)+H)]\), which automatically tune \(\alpha\) when we need exploration or exploitation.