Advanced Policy Optimization

Reusing On-Policy Data

\nabla_\theta J(\theta) =\mathbb E_{ s\sim d^{\pi_\theta},\, a\sim\pi_\theta(\cdot\mid s) } \left[ \nabla_\theta\log\pi_\theta(a\mid s) A^{\pi_\theta}(s,a) \right].

The estimator is on-policy in two ways: states follow the current policy's occupancy $d^{\pi_\theta}$ , and actions follow $\pi_\theta(\cdot\mid s)$ .

Suppose a batch was collected by $\pi_{\theta_k}$ and one gradient step produces

\theta_{k+1}=\theta_k+\alpha\widehat g_k.

The old state-action pairs did not come from $\pi_{\theta_{k+1}}$ . Reusing them naively substitutes the wrong sampling distribution into the next gradient estimate.

This is expensive because reliable policy gradients need large trajectory batches. We want an outer loop that collects data and an inner loop that performs several optimization steps:

\text{collect once} \longrightarrow \text{estimate advantages once} \longrightarrow \text{optimize several times}.

Every inner step makes the batch more off-policy. Advanced policy optimization asks how to reuse it without letting the candidate policy drift so far that the old data becomes misleading.

Importance Sampling

Let $p$ be a target distribution and $q$ a behavior distribution. If $q(x)>0$ wherever $p(x)>0$ , then

\boxed{ \mathbb E_{x\sim p}[f(x)] =\mathbb E_{x\sim q} \left[\frac{p(x)}{q(x)}f(x)\right] }.

The likelihood ratio corrects for sampling from the wrong distribution.

For a trajectory,

p_\theta(\tau) =p(s_0)\prod_{t=0}^{H-1} \pi_\theta(a_t\mid s_t) P(s_{t+1}\mid s_t,a_t).

Comparing a target policy $\pi_\theta$ with a behavior policy $\pi_\beta$ , the initial-state and transition factors cancel:

\boxed{ \frac{p_\theta(\tau)}{p_\beta(\tau)} =\prod_{t=0}^{H-1} \frac{\pi_\theta(a_t\mid s_t)} {\pi_\beta(a_t\mid s_t)} }.

Define the per-step ratio

\rho_t(\theta) =\frac{\pi_\theta(a_t\mid s_t)} {\pi_\beta(a_t\mid s_t)}.

Then the exact off-policy objective is

J(\theta) =\mathbb E_{\tau\sim p_\beta} \left[ \left(\prod_{t=0}^{H-1}\rho_t(\theta)\right) R(\tau) \right].

The exact off-policy gradient

Applying the likelihood-ratio identity gives

\boxed{ \nabla_\theta J(\theta) =\mathbb E_{\tau\sim p_\beta} \left[ \left(\prod_{j=0}^{H-1}\rho_j(\theta)\right) \left( \sum_{t=0}^{H-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) \right) R(\tau) \right] }.

See derivation

Write the target expectation under behavior trajectories:

J(\theta) =\mathbb E_{\tau\sim p_\beta} \left[ \frac{p_\theta(\tau)}{p_\beta(\tau)}R(\tau) \right].

Differentiate:

\begin{aligned} \nabla_\theta J(\theta) &=\mathbb E_{\tau\sim p_\beta} \left[ \frac{\nabla_\theta p_\theta(\tau)}{p_\beta(\tau)} R(\tau) \right]\\ &=\mathbb E_{\tau\sim p_\beta} \left[ \frac{p_\theta(\tau)}{p_\beta(\tau)} \nabla_\theta\log p_\theta(\tau) R(\tau) \right]. \end{aligned}

Only the policy factors depend on $\theta$ :

\nabla_\theta\log p_\theta(\tau) =\sum_t\nabla_\theta\log\pi_\theta(a_t\mid s_t).

Substituting the trajectory ratio gives the result.

Causality lets the score at time $t$ use reward-to-go and only the prefix ratio:

\boxed{ \nabla_\theta J(\theta) =\mathbb E_{\tau\sim p_\beta} \left[ \sum_{t=0}^{H-1} \left(\prod_{j=0}^{t}\rho_j(\theta)\right) \nabla_\theta\log\pi_\theta(a_t\mid s_t) G_t \right] }.

The prefix corrects both the current action and the state distribution created by all preceding actions.

Why exact correction is impractical

Even modest deviations multiply across time:

\log\prod_{j=0}^{t}\rho_j =\sum_{j=0}^{t}\log\rho_j.

A few trajectories receive enormous weights while most receive almost zero. The estimator remains unbiased but has unusably high variance.

For normalized weights $w_i$ , the effective sample size

\mathrm{ESS} =\frac{\left(\sum_iw_i\right)^2}{\sum_iw_i^2}

collapses when a few samples dominate.

The practical surrogate retains only the current action ratio:

\boxed{ L^{\mathrm{IS}}(\theta') =\mathbb E_{ s\sim d^{\pi_\beta},\, a\sim\pi_\beta } \left[ \frac{\pi_{\theta'}(a\mid s)}{\pi_\beta(a\mid s)} A^{\pi_\beta}(s,a) \right] }.

This deliberately ignores state-distribution mismatch. It is useful only while the candidate policy remains close to the behavior policy.

Policy Improvement from Old Advantages

The right justification for several old-batch updates is approximate policy iteration. We evaluate the old policy, then fit a new policy that prefers its positive-advantage actions.

For policies $\pi$ and $\pi'$ , the performance-difference lemma states

\boxed{ J(\pi')-J(\pi) =\mathbb E_{\tau\sim\pi'} \left[ \sum_{t=0}^{\infty}\gamma^tA^\pi(s_t,a_t) \right] }.

See derivation

Start with

J(\pi')-J(\pi) =J(\pi')- \mathbb E_{\tau\sim\pi'}[V^\pi(s_0)].

The value term can be written as a telescoping series:

-V^\pi(s_0) =\sum_{t=0}^{\infty}\gamma^t \left[ \gamma V^\pi(s_{t+1})-V^\pi(s_t) \right].

Therefore,

\begin{aligned} J(\pi')-J(\pi) &=\mathbb E_{\tau\sim\pi'} \left[ \sum_{t=0}^{\infty}\gamma^t \left( r_t+\gamma V^\pi(s_{t+1})-V^\pi(s_t) \right) \right]. \end{aligned}

Taking the conditional expectation over $s_{t+1}$ ,

r(s_t,a_t)+\gamma\mathbb E[V^\pi(s_{t+1})\mid s_t,a_t] -V^\pi(s_t) =A^\pi(s_t,a_t).

Substitution proves the identity.

Using the normalized discounted occupancy

d^\pi(s) =(1-\gamma)\sum_{t=0}^{\infty} \gamma^t\Pr(s_t=s\mid\pi),

we can write

J(\pi')-J(\pi) =\frac1{1-\gamma} \mathbb E_{ s\sim d^{\pi'},\, a\sim\pi' }[A^\pi(s,a)].

At each fixed state, change the action distribution back to $\pi$ :

\mathbb E_{a\sim\pi'}[A^\pi(s,a)] =\mathbb E_{a\sim\pi} \left[ \frac{\pi'(a\mid s)}{\pi(a\mid s)}A^\pi(s,a) \right].

The only remaining new-policy quantity is its state distribution. Replacing $d^{\pi'}$ with $d^\pi$ yields the tractable surrogate

\boxed{ L_\pi(\pi') =\frac1{1-\gamma} \mathbb E_{ s\sim d^\pi,\, a\sim\pi } \left[ \frac{\pi'(a\mid s)}{\pi(a\mid s)}A^\pi(s,a) \right] }.

The old advantage is exact. The approximation is freezing the state visitation distribution.

Several inner updates are reasonable while the old advantages remain informative and $d^{\pi'}$ remains close to $d^\pi$ .

Controlling State-Distribution Shift

At a state, total-variation distance between two policies is

D_{\mathrm{TV}} \left(\pi'(\cdot\mid s),\pi(\cdot\mid s)\right) =\frac12\sum_a |\pi'(a\mid s)-\pi(a\mid s)|.

Assume

\max_sD_{\mathrm{TV}} \left(\pi'(\cdot\mid s),\pi(\cdot\mid s)\right) \leq\alpha.

The policies can be coupled to choose the same action at each shared state with probability at least $1-\alpha$ . The probability of no disagreement for $t$ decisions is at least $(1-\alpha)^t$ , so

D_{\mathrm{TV}}(p_t^{\pi'},p_t^\pi) \leq1-(1-\alpha)^t \leq t\alpha.

Thus small action-distribution changes control state-distribution changes over a limited horizon.

For a bounded function $f$ ,

\left| \mathbb E_{p'}[f]-\mathbb E_p[f] \right| \leq2D_{\mathrm{TV}}(p',p)\|f\|_\infty.

Combining these ideas gives the structural bound

\boxed{ \text{true improvement} \geq \text{old-data surrogate improvement} -\text{penalty for policy change} }.

For the infinite-horizon discounted case, the penalty contains

\sum_{t=0}^{\infty}t\gamma^t =\frac{\gamma}{(1-\gamma)^2},

so long effective horizons make distribution shift especially costly.

Pinsker's inequality connects total variation to KL divergence:

\boxed{ D_{\mathrm{TV}}(p,q) \leq\sqrt{\frac12D_{\mathrm{KL}}(p\|q)} }.

We can therefore control the approximation error by keeping the candidate policy close in KL.

KL Trust Regions

Euclidean parameter distance is a poor measure of policy change. A tiny parameter step can drastically alter action probabilities, while a large step in an insensitive direction may barely change them.

KL divergence measures change in the represented distributions:

D_{\mathrm{KL}}(p\|q) =\mathbb E_{x\sim p} \left[\log\frac{p(x)}{q(x)}\right].

It is nonnegative but asymmetric:

D_{\mathrm{KL}}(p\|q) \neq D_{\mathrm{KL}}(q\|p).

We usually average statewise KL under old-policy states:

\boxed{ \overline D_{\mathrm{KL}}(\theta',\theta) =\mathbb E_{s\sim d^{\pi_\theta}} \left[ D_{\mathrm{KL}} \left( \pi_\theta(\cdot\mid s) \,\|\, \pi_{\theta'}(\cdot\mid s) \right) \right] }.

The trust-region problem is

\boxed{ \begin{aligned} \max_{\theta'}\quad &\mathbb E_t \left[ \frac{\pi_{\theta'}(a_t\mid s_t)} {\pi_\theta(a_t\mid s_t)} \widehat A_t \right]\\ \text{subject to}\quad &\overline D_{\mathrm{KL}}(\theta',\theta)\leq\delta. \end{aligned} }.

A simpler penalized objective is

L_{\mathrm{pen}}(\theta') =L_\theta(\theta') -\beta\overline D_{\mathrm{KL}}(\theta',\theta).

Large $\beta$ makes updates conservative. Small $\beta$ permits more movement. Because the useful scale changes during training, we can adapt $\beta$ :

\boxed{ \beta\leftarrow \left[ \beta+\eta_\beta \left(\overline D_{\mathrm{KL}}-\delta\right) \right]_+ }.

If measured KL is too large, policy movement becomes more expensive. If it is too small, the penalty relaxes.

Proximal Policy Optimization

PPO approximately prevents the updated policy from moving too far from the policy that generated the batch.

Let

\rho_t(\theta) =\frac{\pi_\theta(a_t\mid s_t)} {\pi_{\mathrm{old}}(a_t\mid s_t)} =\exp \left( \log\pi_\theta(a_t\mid s_t) -\log\pi_{\mathrm{old}}(a_t\mid s_t) \right).

The raw importance-weighted surrogate is

L^{\mathrm{IS}}(\theta) =\mathbb E_t[\rho_t(\theta)\widehat A_t].

For $\widehat A_t>0$ , it rewards increasing $\rho_t$ without limit. For $\widehat A_t<0$ , it rewards decreasing $\rho_t$ without limit. Repeated updates can therefore move far beyond the region where the old advantage is reliable.

The clipped objective

PPO uses the pessimistic lower envelope

\boxed{ L^{\mathrm{CLIP}}(\theta) =\mathbb E_t \left[ \min\left( \rho_t(\theta)\widehat A_t, \operatorname{clip} (\rho_t(\theta),1-\epsilon,1+\epsilon) \widehat A_t \right) \right] }.

When $A>0$ ,

\ell(\rho,A) = \begin{cases} \rho A,&\rho\leq1+\epsilon,\\ (1+\epsilon)A,&\rho>1+\epsilon. \end{cases}

The sample stops rewarding further increases beyond $1+\epsilon$ . Moving in the harmful direction still receives the unclipped penalty.

When $A<0$ ,

\ell(\rho,A) = \begin{cases} (1-\epsilon)A,&\rho<1-\epsilon,\\ \rho A,&\rho\geq1-\epsilon. \end{cases}

The sample stops rewarding further probability reduction below $1-\epsilon$ , while increasing the bad action remains penalized.

This is why PPO uses

\min(\rho A,\operatorname{clip}(\rho)A)

rather than simply replacing $\rho$ by its clipped value.

Clipping does not enforce

1-\epsilon\leq\rho_t\leq1+\epsilon

for every sample. It removes the incentive for excessive movement in one direction; it does not pull an out-of-range ratio back. Shared network parameters can also change unsampled probabilities.

At the start of the update, $\rho_t=1$ , so clipped and unclipped objectives have the same gradient. After policy movement, clipping introduces deliberate bias in exchange for lower variance and safer reuse.

Advantages and value targets

PPO is usually paired with Generalized Advantage Estimation:

\delta_t^V =r_t+\gamma(1-d_t)V_\phi(s_{t+1})-V_\phi(s_t),

\widehat A_t =\sum_{l=0}^{H-t-1} (\gamma\lambda)^l\delta_{t+l}^V.

The matching value target is

\widehat V_t^{\mathrm{target}} =\widehat A_t+V_\phi(s_t).

Old log-probabilities, advantages, and value targets are computed once and frozen throughout the inner epochs. Recomputing the denominator with the updated policy would collapse the ratio toward one and destroy the surrogate's meaning.

The complete loss

A common minimization loss is

\boxed{ \mathcal J(\theta,\phi) =-L^{\mathrm{CLIP}}(\theta) +c_VL_V(\phi) -c_H\mathbb E_t[\mathcal H(\pi_\theta(\cdot\mid s_t))] }.

where

L_V(\phi) =\mathbb E_t \left[ (V_\phi(s_t)-\widehat V_t^{\mathrm{target}})^2 \right].

The entropy bonus discourages premature collapse of the action distribution.

\textbf{Clipped PPO}

\begin{aligned} &\text{Initialize actor }\pi_\theta\text{ and critic }V_\phi.\\[0.25em] &\text{Repeat:}\\[0.25em] &\qquad \theta_{\mathrm{old}}\leftarrow\theta,\\[0.25em] &\qquad \text{collect trajectories with }\pi_{\theta_{\mathrm{old}}},\\[0.25em] &\qquad \text{store }\log\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t),\\[0.25em] &\qquad \text{compute and freeze }\widehat A_t \text{ and }\widehat V_t^{\mathrm{target}},\\[0.25em] &\qquad \text{for several epochs:}\\[0.25em] &\qquad\quad \text{shuffle the batch into minibatches},\\[0.25em] &\qquad\quad \rho_t\leftarrow \exp(\log\pi_\theta(a_t\mid s_t)-\log\pi_{\mathrm{old}}(a_t\mid s_t)),\\[0.25em] &\qquad\quad \text{optimize the clipped actor objective, value loss, and entropy},\\[0.25em] &\qquad\quad \text{measure KL and stop early if it becomes too large}.\\[0.25em] &\qquad \text{collect a fresh batch with the updated policy}. \end{aligned}

An empirical old-to-new KL estimate is

\widehat D_{\mathrm{KL}} \approx\frac1B\sum_t \left[ \log\pi_{\mathrm{old}}(a_t\mid s_t) -\log\pi_\theta(a_t\mid s_t) \right].

Clipping approximates a trust region; monitoring KL and stopping inner optimization early remain important safeguards.

Natural Policy Gradients

The ordinary gradient is the steepest-ascent direction under a Euclidean parameter constraint:

\max_\delta g^\top\delta \quad\text{subject to}\quad \delta^\top\delta\leq\varepsilon,

where $g=\nabla_\theta J(\theta)$ . Natural gradients instead measure local distance by policy KL.

For $\theta'=\theta+\delta$ , the local KL expansion is

\boxed{ D_{\mathrm{KL}} (p_\theta\|p_{\theta+\delta}) \approx\frac12\delta^\top F(\theta)\delta },

where

\boxed{ F(\theta) =\mathbb E_{x\sim p_\theta} \left[ \nabla_\theta\log p_\theta(x) \nabla_\theta\log p_\theta(x)^\top \right] }

is the Fisher information matrix.

See derivation

Taylor-expand the candidate log-density:

\log p_{\theta+\delta}(x) \approx\log p_\theta(x) +\delta^\top\nabla_\theta\log p_\theta(x) +\frac12\delta^\top \nabla_\theta^2\log p_\theta(x)\delta.

Substitute into KL:

\begin{aligned} D_{\mathrm{KL}}(p_\theta\|p_{\theta+\delta}) &\approx -\delta^\top \mathbb E[\nabla_\theta\log p_\theta(x)]\\ &\quad- \frac12\delta^\top \mathbb E[\nabla_\theta^2\log p_\theta(x)] \delta. \end{aligned}

\mathbb E_{p_\theta}[\nabla_\theta\log p_\theta(x)]=0.

The Fisher identity is

-\mathbb E[\nabla_\theta^2\log p_\theta(x)] =\mathbb E \left[ \nabla_\theta\log p_\theta(x) \nabla_\theta\log p_\theta(x)^\top \right].

Substitution yields the quadratic approximation.

For policies, we average the statewise Fisher under current-policy states:

F(\theta) =\mathbb E_{ s\sim d^{\pi_\theta},\, a\sim\pi_\theta } \left[ \nabla_\theta\log\pi_\theta(a\mid s) \nabla_\theta\log\pi_\theta(a\mid s)^\top \right].

Linearize the objective and quadratically approximate KL:

\begin{aligned} \max_\delta\quad &g^\top\delta\\ \text{subject to}\quad &\frac12\delta^\top F\delta\leq\varepsilon. \end{aligned}

The solution is

\boxed{ \delta^* =\sqrt{ \frac{2\varepsilon}{g^\top F^{-1}g} } F^{-1}g }.

See derivation

Form the Lagrangian

\mathcal L(\delta,\lambda) =g^\top\delta -\lambda \left(\frac12\delta^\top F\delta-\varepsilon\right).

Stationarity gives

g-\lambda F\delta=0,

\delta=\frac1\lambda F^{-1}g.

The optimum uses the full KL budget:

\frac1{2\lambda^2}g^\top F^{-1}g=\varepsilon.

Solving for $\lambda$ and substituting yields the stated step.

The direction

\widetilde\nabla_\theta J =F^{-1}\nabla_\theta J

is the natural policy gradient. Directions that strongly change action probabilities have large Fisher curvature and are shrunk. Directions that barely change the policy can move farther in parameter space.

Trust Region Policy Optimization

TRPO approximates the natural-gradient trust-region step without constructing or inverting the Fisher matrix.

At the current policy, estimate the surrogate gradient $g$ and solve

Fx=g.

For a network with millions of parameters, explicitly forming $F$ is impossible. Automatic differentiation can instead compute Fisher-vector products

Fv =\nabla_{\theta'} \left[ \nabla_{\theta'}\widehat D_{\mathrm{KL}}(\theta')^\top v \right]_{\theta'=\theta}

without materializing the matrix.

Conjugate gradients approximately solves the damped system

\boxed{ (F+\eta I)x=g }.

Damping improves conditioning and prevents extreme movement along poorly estimated low-curvature directions.

Scale the direction to the quadratic KL radius:

\delta_{\mathrm{quad}} =\sqrt{ \frac{2\delta_{\mathrm{KL}}}{x^\top Fx} }x.

This step is based on local approximations and a finite-sample Fisher estimate. TRPO therefore performs backtracking line search. It tries

\theta_{\mathrm{candidate}} =\theta+\alpha\delta_{\mathrm{quad}}, \qquad \alpha\in\{1,\xi,\xi^2,\ldots\},

and accepts the first candidate that improves the measured surrogate while satisfying the empirical KL limit.

\textbf{TRPO}

\begin{aligned} &\text{Collect trajectories with }\pi_\theta.\\[0.25em] &\text{Estimate fixed advantages and the surrogate gradient }g.\\[0.25em] &\text{Define }v\mapsto(F+\eta I)v\text{ using Fisher-vector products}.\\[0.25em] &\text{Use conjugate gradients to approximately solve }(F+\eta I)x=g.\\[0.25em] &\text{Scale }x\text{ to the quadratic KL radius}.\\[0.25em] &\text{Backtrack along the step until:}\\[0.25em] &\qquad \text{the surrogate improves, and}\\[0.25em] &\qquad \widehat D_{\mathrm{KL}}\leq\delta_{\mathrm{KL}}.\\[0.25em] &\text{Accept the candidate and collect a new batch}. \end{aligned}

TRPO more faithfully enforces a KL trust region. PPO replaces the second-order solve and line search with ordinary minibatch optimization and approximate clipping.

Property	TRPO	PPO
Policy control	Explicit empirical KL check	Pessimistic ratio clipping
Optimizer	Conjugate gradients and line search	SGD or Adam
Batch reuse	One large constrained step	Several minibatch epochs
Constraint fidelity	Higher	Approximate
Implementation complexity	High	Moderate

Practical Diagnosis

PPO is usually the practical default when standard minibatch optimization, accelerator efficiency, and simple distributed rollout collection matter. TRPO is useful when tighter control over policy movement justifies curvature estimation and line search.

The main PPO diagnostics are:

\widehat D_{\mathrm{KL}}(\pi_{\mathrm{old}},\pi_\theta),

\operatorname{clipfrac} =\frac1B\sum_t \mathbb 1[|\rho_t-1|>\epsilon],

policy entropy, value loss, and explained variance

\operatorname{EV} =1- \frac{ \operatorname{Var}(\widehat V^{\mathrm{target}}-V_\phi) }{ \operatorname{Var}(\widehat V^{\mathrm{target}}) }.

Symptom	Likely cause	Useful response
KL rises sharply within an epoch	Learning rate or epoch count is too large	Stop early and reduce inner optimization
Most samples are clipped	Batch is already too off-policy	Use fewer epochs or smaller updates
KL is tiny and clip fraction is near zero	Updates are overly conservative	Increase learning rate or reuse cautiously
Entropy collapses early	Policy becomes deterministic too quickly	Increase entropy regularization or reduce actor updates
Ratios are unstable	Old log-probabilities or policy scale are wrong	Store the denominator once and inspect log-ratios
Policy metrics look healthy but returns fail	Critic advantages are unreliable	Inspect value loss and explained variance
Value loss decreases while policy degrades	Critic is overfitting a stale batch	Reduce critic epochs and collect fresh data

Before trusting a PPO run, verify that old log-probabilities are stored and frozen, advantages and value targets are computed once, terminal and truncation bootstrapping are correct, ratios use log-probability differences, and inner optimization stops when KL becomes excessive.

The full progression is

\underbrace{ J(\pi')-J(\pi) =\frac1{1-\gamma} \mathbb E_{d^{\pi'},\pi'}[A^\pi] }_{\text{exact policy improvement}}

\Downarrow\quad\text{correct actions and freeze states}

\underbrace{ L_\pi(\pi') =\frac1{1-\gamma} \mathbb E_{d^\pi,\pi} \left[\frac{\pi'}{\pi}A^\pi\right] }_{\text{old-data surrogate}}

\Downarrow\quad\text{control policy movement}

\underbrace{ \max_{\pi'}L_\pi(\pi') \quad\text{subject to}\quad \overline D_{\mathrm{KL}}(\pi,\pi')\leq\delta }_{\text{trust-region update}}.

PPO approximates this constraint with clipped ratios or an adaptive KL penalty. Natural gradients use the Fisher geometry of policy space. TRPO combines that geometry with conjugate gradients and line search. All three aim to extract more learning from each batch without trusting old advantages beyond their local region.