Pearl
Lab

Chapter 07 · Trust-region policy optimization

TRPO

Abstract. Trust Region Policy Optimization (Schulman et al., 2015) stabilizes on-policy policy gradients by maximizing a surrogate for expected improvement subject to a hard constraint on how far the new policy may move from the old one in KL divergence. This chapter derives the surrogate from performance-difference lemmas, formulates the constrained problem, connects it to the natural policy gradient via the Fisher information matrix, and unpacks the practical solver: Hessian-vector products, conjugate gradients, and a backtracking line search. We close with the conceptual jump to PPO’s clipped objective. The laboratory visualizes surrogate improvement vs mean KL on a tabular softmax policy so the trust-region trade-off is concrete.

1. Overview

Chapter 06 ended with a warning: Euclidean steps in \(\theta\)-space are the wrong geometry for policies. TRPO is the paper that took that warning seriously and built a usable deep RL algorithm around a trust region in policy space.

TRPO sits between vanilla REINFORCE/A2C and PPO. It is more principled (and more expensive) than PPO; PPO is best understood as “TRPO’s idea, without conjugate gradients.” Reading TRPO carefully makes PPO’s clipping feel inevitable rather than arbitrary.

2. The instability problem (restated precisely)

On-policy methods collect data under \(\pi_{\theta_{\mathrm{old}}}\), compute a gradient estimate, update to \(\theta_{\mathrm{new}}\), then must discard that data. If the update collapses \(\pi\) (e.g. one action’s probability jumps to nearly 1 in many states), subsequent rollouts are useless and recovery is slow.

Two mismatches drive collapse:

  1. Parameter vs distribution. \(\|\Delta\theta\|_2\) small does not imply \(\mathrm{KL}(\pi_{\theta}\,\|\,\pi_{\theta+\Delta\theta})\) small.
  2. Surrogate vs true objective. The gradient is usually of a local approximation to \(J\), valid only near \(\theta_{\mathrm{old}}\). Large steps leave the region where the approximation holds.

TRPO addresses both: optimize a local surrogate, and constrain step size in KL, not in Euclidean \(\theta\).

3. Conservative policy improvement

Kakade & Langford’s conservative policy iteration and related performance difference lemmas bound how much \(J\) can change when you switch policies. Informally, for policies \(\pi\) and \(\pi'\),

(1) \[ J(\pi') - J(\pi) = \mathbb{E}_{s\sim d^{\pi'},\, a\sim\pi'} \big[ A^{\pi}(s,a) \big]. \]

Exact, but awkward: the expectation uses the new occupancy \(d^{\pi'}\), which you do not have until you deploy \(\pi'\). Replace \(d^{\pi'}\) by \(d^{\pi}\) and you get a first-order surrogate that matches the true improvement to linear order in the policy change:

(2) \[ L_\pi(\pi') \;=\; J(\pi) + \mathbb{E}_{s\sim d^{\pi},\, a\sim\pi'} \big[ A^{\pi}(s,a) \big]. \]

Then \(L_\pi(\pi)=J(\pi)\) and \(\nabla_{\pi'} L_\pi|_{\pi'=\pi} = \nabla J\) (in the appropriate sense). Maximizing \(L_\pi\) is therefore a sensible local proxy for maximizing \(J\) — provided \(\pi'\) stays near \(\pi\) so that \(d^{\pi'}\approx d^{\pi}\) and higher-order terms remain small.

Theorem 3.1 · Improvement lower bound (schematic)

Under mild conditions there exists \(C>0\) such that \[ J(\pi') \;\ge\; L_\pi(\pi') - C\cdot \mathbb{E}_{s\sim d^{\pi}} \big[ \mathrm{KL}\big(\pi'(\cdot\mid s)\,\|\,\pi(\cdot\mid s)\big) \big]. \] Maximizing \(L_\pi\) while keeping mean KL small therefore lower-bounds true improvement — the theoretical heart of TRPO.

4. Importance-sampled surrogate

Sampling \(a\sim\pi'\) in (2) is inconvenient when you only have data from \(\pi\). Rewrite with importance weights:

(3) \[ L_\pi(\pi') = J(\pi) + \mathbb{E}_{s\sim d^{\pi},\, a\sim\pi} \left[ \frac{\pi'(a\mid s)}{\pi(a\mid s)} A^{\pi}(s,a) \right]. \]

In parameterized form, with \(\theta_{\mathrm{old}}\) fixed during an update and \(\theta\) the free variable,

(4) \[ L_{\theta_{\mathrm{old}}}(\theta) = \mathbb{E}_{t} \left[ \frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} \hat A_t \right]. \]

(Constant \(J(\pi_{\theta_{\mathrm{old}}})\) dropped — it does not affect \(\arg\max_\theta\).) The ratio \(r_t(\theta)=\pi_\theta(a_t\mid s_t)/\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)\) is the same object PPO will later clip.

Remark · On-policy data

The expectation in (4) is estimated from rollouts of \(\pi_{\theta_{\mathrm{old}}}\). After the update you throw the batch away. TRPO is firmly on-policy — unlike DQN or SAC.

5. The KL trust region

Definition 5.1 · Mean KL constraint

TRPO constrains the average forward KL under the state distribution of the old policy: \[ \bar D_{\mathrm{KL}} \big(\theta_{\mathrm{old}},\theta\big) = \mathbb{E}_{s\sim d^{\pi_{\theta_{\mathrm{old}}}}} \Big[ \mathrm{KL}\big( \pi_{\theta_{\mathrm{old}}}(\cdot\mid s) \,\big\|\, \pi_{\theta}(\cdot\mid s) \big) \Big] \;\le\; \delta. \] Typical \(\delta\) in deep RL: on the order of \(10^{-2}\) (e.g. \(0.01\)).

Why KL rather than total variation or parameter norms?

  • Appears naturally in the improvement lower bound (Theorem 3.1).
  • Local quadratic approximation of KL is the Fisher information metric — connects to natural gradients.
  • Invariant to reparameterization of \(\theta\) (to first order) — Euclidean \(\|\Delta\theta\|\) is not.
Analysis · Forward vs reverse KL

TRPO uses \(\mathrm{KL}(\pi_{\mathrm{old}}\,\|\,\pi_{\mathrm{new}})\) — mode-covering / zero-avoiding relative to the old policy: \(\pi_{\mathrm{new}}\) is discouraged from putting tiny mass where \(\pi_{\mathrm{old}}\) had mass. The opposite direction would behave differently; do not swap them casually when reading code.

6. The TRPO optimization problem

(5) \[ \begin{aligned} \theta_{k+1} &= \arg\max_{\theta} \; L_{\theta_k}(\theta)\\ &\quad\text{s.t.}\quad \bar D_{\mathrm{KL}}(\theta_k,\theta) \le \delta. \end{aligned} \]

This is a constrained nonlinear program in high dimension. TRPO solves a local approximation:

(6) \[ \begin{aligned} \maximize_{x}\quad & g^\top x \\ \text{s.t.}\quad & \tfrac12 x^\top F x \le \delta, \end{aligned} \]

where \(x=\theta-\theta_k\), \(g=\nabla_\theta L_{\theta_k}(\theta)|_{\theta_k}\) is the policy gradient (same direction as \(\nabla J\) at \(\theta_k\)), and \(F\) is the Fisher information matrix — the Hessian of mean KL at \(x=0\):

(7) \[ F = \nabla^2_\theta \bar D_{\mathrm{KL}}(\theta_k,\theta) \big|_{\theta=\theta_k} = \mathbb{E}_{s,a\sim\pi_{\theta_k}} \big[ \nabla_\theta\log\pi_{\theta_k}(a\mid s)\, \nabla_\theta\log\pi_{\theta_k}(a\mid s)^\top \big]. \]

7. Natural gradient connection

The solution of the quadratic program (6) is analytic:

(8) \[ x = \sqrt{\frac{2\delta}{g^\top F^{-1} g}} \; F^{-1} g. \]

The direction \(F^{-1}g\) is exactly the natural policy gradient from Chapter 06. The scalar rescales the step to sit on the KL ball boundary \(\tfrac12 x^\top F x = \delta\) (approximately). TRPO = natural gradient step with an explicitly enforced trust-region radius, plus a line search to repair approximation error.

Geometry

\(F\) defines a Riemannian metric on parameter space whose geodesics locally track KL. Steepest ascent of \(J\) in that metric is \(F^{-1}\nabla J\). Euclidean ascent ignores the metric — hence brittle steps.

8. Computing \(F^{-1}g\) without forming \(F\)

For a deep net, \(F\) is \(|\theta|\times|\theta|\) — impossible to materialize or invert. TRPO never builds \(F\). It only needs Hessian-vector products \(Fv\), estimated by autodiff:

(9) \[ F v = \nabla_\theta \Big( \big(\nabla_\theta \bar D_{\mathrm{KL}}\big)^\top v \Big) \Big|_{\theta_k}. \]

Conjugate gradient (CG) solves \(F x = g\) using a sequence of such products — typically \(10\)–\(15\) iterations — yielding \(x\approx F^{-1}g\) well enough for the step direction. Then rescale as in (8).

Damping: practice often uses \(F+\varepsilon I\) for numerical stability (Levenberg–Marquardt style).

Definition 8.1 · Fisher-vector product

Any routine that, given vector \(v\), returns an estimate of \(Fv\) without instantiating \(F\). This is the computational core that PPO later avoids entirely.

9. Backtracking line search

The quadratic model (6) can disagree with the true surrogate and true KL. TRPO therefore proposes \(x\) from (8), then tries \(\theta \leftarrow \theta_k + \beta^j x\) for \(j=0,1,2,\ldots\) (e.g. \(\beta=0.8\)) until both hold:

  1. Surrogate improves: \(L_{\theta_k}(\theta) > L_{\theta_k}(\theta_k)\) (or a small positive threshold).
  2. Empirical mean KL \(\le \delta\).

If no step satisfies the constraints, skip the update. This is why TRPO is conservative in the everyday sense as well as the Kakade–Langford sense.

10. Full algorithm (practical TRPO)

For iteration k = 1, 2, …:
  Collect trajectories with π_{θ_k}
  Compute advantages Â_t (typically GAE)
  Estimate g = ∇_θ L_{θ_k}(θ) |_{θ_k}
  Use CG with Fisher-vector products to get x ≈ F⁻¹ g
  Rescale x so that ½ xᵀ F x ≈ δ
  Line search along x for KL ≤ δ and L improvement
  θ_{k+1} ← accepted point
  Fit value function V_φ on returns / TD targets (separate regression)

Value function fitting is usually an ordinary supervised problem (MSE on returns or GAE-style targets) — not constrained by the KL trust region. Some implementations use a separate trust region or conjugacy for \(V\) as well; the original paper focuses the KL constraint on the policy.

11. Advantages in TRPO practice

TRPO’s theory is written with \(A^{\pi}\); practice uses GAE(\(\gamma,\lambda\)) from Chapter 06: \(\hat A_t=\sum_l(\gamma\lambda)^l\delta_{t+l}\). Advantage normalization (subtract mean, divide by std within batch) is common for variance control.

Multiple SGD epochs over the same batch are not free in TRPO the way they later are (somewhat) in PPO — the surrogate and KL are defined w.r.t. \(\theta_k\), and the hard constraint is enforced once per iteration via CG + line search, not via many small clipped steps.

12. Theory notes & caveats

  • The monotonic improvement guarantee holds for the idealized problem (maximize \(L\) under KL); approximations (sampling, CG truncation, neural nets, GAE bias) void the literal guarantee. Empirically TRPO is still far more stable than vanilla PG.
  • Forward KL constraint does not by itself prevent all bad updates; line search on the empirical surrogate matters.
  • TRPO scales poorly to very large policies because CG needs many backprops per update. That cost is a primary industrial reason PPO won.
  • Compatible with discrete and continuous policies (softmax / Gaussian), as long as \(\log\pi\) and KL (or its Hessian-vector product) are tractable.
Analysis · TRPO vs DQN stability stories

DQN stabilizes value bootstrapping (replay + target net). TRPO stabilizes policy steps (KL ball + surrogate). Different objects, same moral: unconstrained deep TD / deep PG updates are too aggressive; insert a lag or a trust region.

13. Bridge to PPO

PPO keeps the importance-ratio surrogate (4) but replaces the hard KL constraint with a clipped objective:

(10) \[ L^{\mathrm{CLIP}}(\theta) = \mathbb{E}_t \Big[ \min\big( r_t(\theta)\hat A_t,\; \mathrm{clip}(r_t(\theta), 1-\varepsilon, 1+\varepsilon)\hat A_t \big) \Big]. \]

Clipping removes incentive to push \(r_t\) far outside \([1-\varepsilon,1+\varepsilon]\), mimicking a trust region without Fisher-vector products or CG. PPO also often adds an adaptive KL penalty variant. Next chapter derives (10) carefully and contrasts it with TRPO step by step.

14. Laboratory · trust region on a softmax policy

Tabular softmax policy on a 4×4 goal grid (same spirit as Chapter 06). Each “TRPO-style step” estimates advantages by Monte Carlo, forms the surrogate gradient in logit space, takes a natural-ish step (diagonal Fisher approximation — full CG is overkill for a table), then backtracks until mean KL \(\le\delta\) and surrogate improves. Compare to an unconstrained Euclidean step of similar raw magnitude — watch KL explode and returns dip.

Surrogate vs KL ball

Line search · δ constraint
Trust region
Update type

Policy (greedy)

Collect a batch first, then take a trust-region or vanilla step.

Batch eps0
Mean |Â|
Surrogate Δ
Mean KL
Line steps
Avg G₀

Solver ledger

Collect batch → TRPO-style step. Ledger shows proposed KL, backtracking, and accepted ΔL.

15. Worked calculations

Quadratic trust step. Suppose 1-D \(\theta\), \(g=2\), \(F=4\), \(\delta=0.01\). Natural direction \(F^{-1}g=0.5\). Scale \(\eta=\sqrt{2\delta/(g F^{-1} g)}=\sqrt{0.02/(2\cdot 0.5)}=\sqrt{0.02}=0.1414\). Step \(x=\eta\cdot 0.5\approx 0.0707\). Check: \(\tfrac12 F x^2=\tfrac12\cdot 4\cdot(0.0707)^2\approx 0.01=\delta\).

Importance ratio. Old \(\pi_{\mathrm{old}}(a|s)=0.25\), new \(0.40\), \(\hat A=2\). Contribution to surrogate: \(rA=(0.40/0.25)\cdot 2=3.2\). If \(\hat A=-2\), contribution \(-3.2\) — TRPO still uses the linear surrogate; PPO would clip the ratio when it tries to move too far against a negative advantage.

  1. Roll out π_old; compute  (GAE or MC).
  2. Form L(θ)=E[(π_θ/π_old) Â]; compute g=∇L at θ_old.
  3. Estimate Fv via KL Hessian-vector products; CG ⇒ x≈F⁻¹g.
  4. Rescale x onto ½ xᵀFx = δ.
  5. Backtrack βʲx until KL≤δ and L improves.
  6. Refit V; repeat. PPO will replace steps 3–5 by clipping r(θ).

16. What follows

PPO keeps TRPO’s surrogate ratio and advantage machinery but substitutes clipping (and optional KL penalties) for conjugate gradients. After PPO, the course returns to off-policy continuous control (DDPG/TD3) and entropy-regularized soft actor–critic (SAC).

See also

  • Schulman et al., 2015 — Trust Region Policy Optimization.
  • Kakade & Langford — Approximately optimal approximate RL (CPI).
  • Schulman et al. — High-dimensional continuous control using GAE.
  • Chapter 06 — policy gradient theorem, natural gradient motivation.

Previous

← Policy Gradients

Up next

PPO

Clipped surrogates — TRPO’s idea without CG.

Continue →