Trust-Region Policy Optimization

Let $\pi_\theta$ denote a stochastic policy with parameters $\theta$, and let $\eta(\pi) = \mathbb{E}_{\tau \sim \pi}\left[\sum_{t=0}^{\infty} \gamma^t r_t\right]$ be its expected discounted return. Given a current policy $\pi_{\theta_{\mathrm{old}}}$, a classical identity (used throughout policy-based methods) expresses the performance of any other policy $\pi$ in terms of the advantage function $A_{\pi_{\mathrm{old}}}(s,a)$ of the old policy:

where $\rho_\pi(s) = \sum_t \gamma^t P(s_t = s \mid \pi)$ is the discounted state-visitation distribution under $\pi$. This identity is exact but useless computationally, because it requires sampling states from the new policy $\pi$, which we have not yet computed. TRPO's first move is to replace $\rho_\pi$ with $\rho_{\pi_{\mathrm{old}}}$, giving the surrogate objective:

This substitution is exactly what makes the objective usable: it can be estimated from trajectories generated by the old policy, using The Advantage Function estimates already computed, and (with an importance-sampling ratio over actions) differentiated with respect to $\theta$. The substitution is also the source of error — for policies that are far from $\pi_{\mathrm{old}}$, $\rho_\pi$ and $\rho_{\pi_{\mathrm{old}}}$ diverge, and $L$ no longer tracks $\eta$ faithfully. The key theoretical contribution of TRPO is a bound quantifying exactly how much $L$ can diverge from $\eta$, expressed via a KL-divergence penalty between $\pi_{\mathrm{old}}$ and $\pi$. This transforms an intuitive idea — "don't move the policy too far" — into a precise, provable stopping condition. The practical algorithm then solves, at each iteration,

where $\overline{D}_{\mathrm{KL}}$ is the KL-divergence averaged over states visited by $\pi_{\theta_{\mathrm{old}}}$, and $\delta$ is a trust-region radius. This is a trust-region method in the classical optimization sense: instead of trusting a linear/quadratic local model everywhere, we only trust it within a region — here, a region in policy-distribution space rather than raw parameter space, which matters because a small parameter change can produce an enormous change in action probabilities (or vice versa).

The monotonic-improvement bound

Define $D_{\mathrm{KL}}^{\max}(\pi_{\mathrm{old}}, \pi) = \max_s D_{\mathrm{KL}}\big(\pi_{\mathrm{old}}(\cdot\mid s) \,\|\, \pi(\cdot\mid s)\big)$. The central theorem states that for $\alpha = D_{\mathrm{KL}}^{\max}(\pi_{\mathrm{old}}, \pi)$ and a constant $\epsilon = \max_{s,a} |A_{\pi_{\mathrm{old}}}(s,a)|$,

This says: the true return of the new policy is at least the surrogate value minus a penalty that grows with how far, in the worst-case state, the new policy has drifted from the old one. Two immediate corollaries follow. First, at $\pi = \pi_{\mathrm{old}}$, both sides equal $\eta(\pi_{\mathrm{old}})$ and the KL term is zero, so the bound is tight at the current policy — there is no slack to exploit trivially. Second, define

Then $M(\pi_{\mathrm{old}}) = \eta(\pi_{\mathrm{old}})$ and $\eta(\pi) \ge M(\pi)$ for all $\pi$. Any policy update that increases $M$ therefore cannot decrease $\eta$:

This is a minorization-maximization (MM) algorithm: $M$ is a lower bound (a minorant) on $\eta$ that touches $\eta$ at the current point, and maximizing $M$ each iteration guarantees the true objective is non-decreasing. This is the same template that underlies expectation-maximization algorithms, and it is what gives TRPO its name: theoretically maximizing $M$ instead of $L$ alone yields guaranteed monotonic improvement.

From theory to practice

Maximizing $M(\pi)$ directly is impractical: the penalty coefficient $C$ is enormous in typical problems (because $\gamma$ close to $1$ makes $(1-\gamma)^2$ tiny), producing minuscule, glacially slow steps. TRPO makes two practical relaxations: 1. Replace the max-KL penalty with a hard constraint on the average KL-divergence, $\overline{D}_{\mathrm{KL}}^{\rho_{\mathrm{old}}}(\theta_{\mathrm{old}}, \theta) \le \delta$, chosen empirically (e.g., $\delta \approx 0.01$) rather than derived from $C$. 2. Solve the constrained problem approximately using local Taylor expansions: linearize the objective, quadratically approximate the constraint. Writing $g = \nabla_\theta L_{\theta_{\mathrm{old}}}(\theta)\big|_{\theta_{\mathrm{old}}}$ for the policy gradient and $F$ for the Fisher information matrix of the policy (the Hessian of the KL constraint at $\theta_{\mathrm{old}}$), the local problem becomes

whose solution (via Lagrange multipliers) is the natural gradient direction

Forming and inverting $F$ explicitly is infeasible for networks with thousands to millions of parameters ($O(n^3)$ for a naive inverse, $O(n^2)$ just to store $F$). TRPO instead computes the step direction $x \approx F^{-1}g$ using the conjugate gradient (CG) method, which only needs the ability to compute Fisher-vector products $F v$ for arbitrary vectors $v$ — computable in $O(n)$ time via automatic differentiation (a Pearlmutter-style Hessian-vector-product trick applied to the KL divergence), without ever materializing $F$. Because the linear/quadratic approximations are only locally valid, the raw natural-gradient step can violate the true KL constraint or even decrease the true surrogate objective. TRPO corrects this with a backtracking line search: it tries the full step, and if either the KL constraint is violated or the surrogate objective did not improve, it shrinks the step geometrically until both conditions hold. ``` PSEUDOCODE — Trust-Region Policy Optimization (single iteration) Inputs: current policy parameters θ_old, trust-region size δ, CG iterations K, backtracking coefficient α ∈ (0,1), max backtracking steps L 1. Collect a batch of trajectories using π_θold, interacting with the environment 2. Estimate advantages Â(s,a) for all visited (s,a) # e.g., via GAE, using a value baseline 3. Form the surrogate objective estimate: L(θ) = (1/N) Σ_i [ π_θ(a_i | s_i) / π_θold(a_i | s_i) ] Â_i 4. Compute policy gradient g = ∇_θ L(θ) at θ = θ_old 5. Use Conjugate Gradient to approximately solve F x = g (so x ≈ F⁻¹ g) without forming F explicitly: function Fisher_vector_product(v): kl = mean_s KL[ π_θold(·|s) ‖ π_θ(·|s) ] # evaluated at θ = θ_old grad_kl = ∇_θ kl gv = grad_kl · v # scalar return ∇_θ (gv) # this equals F v (+ damping·v) x ← ConjugateGradient(Fisher_vector_product, g, iterations=K) 6. Compute the maximal step length under the quadratic KL model: β ← sqrt( 2δ / (xᵀ F x) ) # xᵀFx estimated via one more Fisher-vector product full_step ← β x 7. Backtracking line search (ensures true improvement + true constraint satisfaction): for j = 0, 1, ..., L-1: θ_new ← θ_old + α^j full_step if KL(θ_old, θ_new) ≤ δ and L(θ_new) ≥ L(θ_old) and L(θ_new) is finite: accept θ_new; break if no step accepted after L tries: θ_new ← θ_old # reject the update this iteration 8. (Separately) Fit the value function baseline V_φ by regression on observed returns Output: θ ← θ_new ``` The CG solve in step 5 finds $x$ minimizing $\tfrac12 x^\top F x - g^\top x$, which is mathematically equivalent to solving $Fx = g$; running $K \ll n$ iterations (typically $K \approx 10$–$20$) gives a good approximate direction far cheaper than exact inversion. The line search in step 7 is essential: it is what upgrades an approximate* natural-gradient step into one that verifiably respects the trust-region constraint and improves the surrogate, restoring the guarantees that the Taylor approximations alone would forfeit.

Per-iteration time complexity. Let $n$ be the number of policy parameters and $N$ the number of sampled transitions per batch. Sampling trajectories and computing advantages costs $O(N)$ forward/backward passes through the policy and value networks. Each Fisher-vector product costs two backpropagation passes through the policy network, $O(n)$ (proportional to network size, same order as an ordinary gradient computation), and CG runs for $K$ iterations, giving $O(Kn)$ for the search-direction computation — versus $O(n^3)$ for naive exact Fisher-matrix inversion and $O(n^2)$ just to store $F$. Since $K$ is a small constant (commonly 10–20) independent of $n$, TRPO's per-iteration cost is effectively linear in the number of parameters, the same asymptotic order as a plain policy-gradient step, at the price of a constant-factor overhead of roughly $2K{+}O(1)$ backward passes plus the line search's extra forward evaluations (each cheap, $O(N)$, since it only re-evaluates the objective and KL, not full gradients).

Space complexity is $O(n + N)$: parameters, their gradients, and CG's working vectors ($O(n)$), plus the stored batch of transitions ($O(N)$). No $O(n^2)$ matrix is ever stored, which is what makes TRPO usable for policies with millions of parameters.

Correctness argument (why monotonic improvement holds). The proof rests on three pieces fitting together: 1. Exact identity. The performance-difference identity $\eta(\pi) = \eta(\pi_{\mathrm{old}}) + \mathbb{E}_{s\sim\rho_\pi, a\sim\pi}[A_{\pi_{\mathrm{old}}}(s,a)]$ is an algebraic consequence of telescoping the discounted-return difference over time steps and using the definition of the advantage function; it holds for any two policies, with no approximation. 2. Bounding the distribution-shift error. Swapping $\rho_\pi$ for $\rho_{\mathrm{old}}$ introduces an error bounded by a total-variation argument: the difference in expected advantage under the two state distributions is controlled by how much probability mass shifts between them, which itself is bounded by the (max) KL-divergence between the policies via Pinsker's-inequality-type reasoning, propagated through the discounted horizon. This chain of bounds yields exactly the constant $C = 4\epsilon\gamma/(1-\gamma)^2$ in the theorem, showing the surrogate-true gap is controlled — not merely correlated — by KL-divergence. 3. MM argument. Given the bound $\eta(\pi) \ge M(\pi) := L_{\mathrm{old}}(\pi) - C\,D_{\mathrm{KL}}^{\max}(\mathrm{old}, \pi)$ with equality at $\pi = \pi_{\mathrm{old}}$, any $\pi$ with $M(\pi) \ge M(\pi_{\mathrm{old}}) = \eta(\pi_{\mathrm{old}})$ satisfies $\eta(\pi) \ge \eta(\pi_{\mathrm{old}})$. This is the textbook argument for MM (minorize-maximize) algorithms: improving a global lower bound that touches the objective at the current point cannot decrease the true objective.

Honesty about the practical algorithm's guarantees. The theorem uses $D_{\mathrm{KL}}^{\max}$ and a specific, typically very conservative constant $C$; the practical algorithm uses an average-KL constraint with an empirically chosen $\delta$, and solves the constrained problem only approximately via linearization plus CG plus line search. Strictly speaking, this means the textbook monotonic-improvement guarantee is not exactly preserved by the implemented algorithm — it is a well-motivated approximation. What is preserved, mechanically, by the line search is a weaker but still meaningful guarantee: each accepted step provably satisfies the average-KL constraint and provably does not decrease the estimated surrogate objective on the current batch. Empirically this tracks true performance closely when $\delta$ is small and batches are large enough to keep the advantage/objective estimates low-variance, but it is a heuristic bridge, not a formally airtight guarantee, over the sampling noise and the max-vs-average KL substitution.

Consider a small discrete-action policy $\pi_\theta$ with two scalar parameters $\theta = (\theta_1, \theta_2)$, controlling action logits in a 2-state, 2-action toy MDP. Suppose one iteration collects $N=4$ transitions and computes the following (illustrative, hand-computable) quantities at $\theta_{\mathrm{old}} = (0.0, 0.0)$:

Suppose the policy gradient at $\theta_{\mathrm{old}}$ works out to $g = (0.55, 0.20)$ (obtained by differentiating the importance-weighted surrogate, which at $\theta = \theta_{\mathrm{old}}$ reduces to the ordinary score-function policy gradient weighted by $\hat A$), and suppose the Fisher information matrix — here small enough to write explicitly for illustration, though TRPO would never form it this way at scale — is

Step — natural gradient via CG. Instead of inverting $F$, TRPO's conjugate-gradient routine solves $Fx = g$ iteratively. Starting from $x_0 = 0$, residual $r_0 = g$, direction $p_0 = r_0$:

• Iteration 1: $\alpha_0 = \frac{r_0^\top r_0}{p_0^\top F p_0} = \frac{0.55^2+0.20^2}{(0.55,0.20)F(0.55,0.20)^\top} = \frac{0.3425}{0.55(0.55{+}0.02){+}0.20(0.055{+}0.10)} \approx \frac{0.3425}{0.344} \approx 0.996$

giving $x_1 \approx (0.548, 0.199)$, quite close to the true $F^{-1}g$ already after one step (since $F$ is close to diagonal here). Continuing CG for a second iteration would essentially converge to the exact solution

Step — step-size scaling. With trust-region radius $\delta = 0.01$, compute $x^\top F x \approx 0.551(0.551{+}0.029) + 0.290(0.055{+}0.145) \approx 0.550$. Then

so the proposed full step is $\Delta\theta = \beta\, x^\star \approx (0.105, 0.055)$.

Step — line search. TRPO evaluates $\theta_{\mathrm{new}} = \theta_{\mathrm{old}} + \Delta\theta = (0.105, 0.055)$: it recomputes the exact average KL divergence between $\pi_{\theta_{\mathrm{old}}}$ and $\pi_{\theta_{\mathrm{new}}}$ over the sampled states (not the quadratic approximation) and the exact surrogate objective $L(\theta_{\mathrm{new}})$. Suppose this full step yields $\overline{D}_{\mathrm{KL}} = 0.013 > \delta$ — the quadratic model under-estimated the true divergence, a common occurrence for larger steps. TRPO backtracks: it tries $0.5 \times \Delta\theta = (0.053, 0.028)$, recomputes KL and $L$; suppose now $\overline{D}_{\mathrm{KL}} = 0.006 \le \delta$ and $L(\theta_{\mathrm{new}}) = 0.31 > L(\theta_{\mathrm{old}}) = 0.28$. Both conditions hold, so this half-step is accepted as $\theta_{\mathrm{new}}$. This trace illustrates the two-layer safety net explicitly: the quadratic KL model determines an initial candidate step, and the exact KL and objective, checked in the line search, are what actually gate acceptance — protecting against the inevitable mismatch between the local quadratic approximation and the true, nonlinear KL-divergence surface.

TRPO was introduced by John Schulman and collaborators at UC Berkeley (advised by Pieter Abbeel), building on decades-earlier ideas about natural gradients in the parameter space of probability distributions, conservative policy iteration, and trust-region methods from classical numerical optimization.

TRPO's direct descendants are ubiquitous in modern deep reinforcement learning. Proximal Policy Optimization, which simplifies TRPO's machinery into a first-order clipped objective, became the default policy-optimization algorithm behind many robotics locomotion controllers, game-playing agents, and — notably — the reinforcement-learning-from-human-feedback (RLHF) stage used to align large language models, where a KL-constrained-style objective against a reference policy is central for exactly the same reason TRPO cares about it: preventing a single bad update from catastrophically shifting model behavior. Natural-gradient and Fisher-vector-product techniques developed for TRPO's conjugate-gradient solver also influenced later second-order and preconditioned optimization methods in deep learning more broadly. Trust-region reasoning — bound the divergence between successive policies, not just the parameter step — remains the conceptual backbone for stable policy improvement across continuous-control robotics, autonomous-vehicle planning research, and multi-agent training pipelines that require careful control over policy drift.

Ask the tutor: How does the clipped surrogate objective in Proximal Policy Optimization approximate the same trust-region guarantee without conjugate gradients or a line search?