Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

TRPO and Natural Policy Gradient

4 minute read

Published:

TL;DR: TRPO (Trust Region Policy Optimisation) places a hard KL-divergence constraint on policy updates, guaranteeing that performance never decreases. The natural policy gradient steers updates in the direction of steepest ascent in policy space (not parameter space), using the Fisher information matrix as a Riemannian metric. Together, these ideas provide principled, stable policy optimisation with monotonic improvement guarantees.
TRPO trust region
TRPO trust region policy optimisation (Schulman et al., 2015)

The Problem with Vanilla Gradient Ascent

Gradient ascent in parameter space has an implicit flaw: the same step size in parameter space can produce wildly different changes in policy behaviour depending on the local curvature. A large step might leap past the trust region into a region where the policy is catastrophically bad, and recovery is difficult.

The core question TRPO asks is: how do we define a “safe” step size that is invariant to the parameterisation of the policy?

The Natural Gradient

In Euclidean gradient ascent, the update is:

θ ← θ + α ∇_θ J(θ)

The natural gradient (Amari 1998) replaces the Euclidean metric with the Fisher information matrix \(F(\theta)\), which measures the local curvature of the KL divergence between policy distributions:

F(θ) = E_π [ ∇_θ log π_θ(a|s) ∇_θ log π_θ(a|s)^T ]

The natural gradient update is:

θ ← θ + α F(θ)^{-1} ∇_θ J(θ)

This steepest ascent direction in distribution space is invariant to reparameterisation: regardless of how we represent the policy (softmax, log-linear, neural network), the natural gradient points in the same direction in policy space.

Key Insight: The Fisher information matrix is the Riemannian metric on the manifold of probability distributions. Computing the natural gradient is equivalent to asking: "which parameter change produces the maximum improvement in objective per unit of KL divergence from the current policy?" This is the right question to ask for policy optimisation.

TRPO: Formalising the Trust Region

TRPO (Schulman et al. 2015) formalises this as a constrained optimisation:

max_θ E_t [ (π_θ(a_t|s_t) / π_{θ_old}(a_t|s_t)) · Â_t ] subject to: E_t [ KL( π_{θ_old}(·|s_t) || π_θ(·|s_t) ) ] ≤ δ

The constraint keeps the new policy within a trust region of radius \(\delta\) in KL space. The monotonic improvement theorem guarantees that, within this constraint, the true objective cannot decrease.

Solving the Constrained Problem

Directly inverting \(F\) for large neural networks is intractable. TRPO avoids this via:

  1. Conjugate gradient: efficiently compute \(F^{-1} g\) without forming \(F\) explicitly, using matrix-vector products \(F v = \nabla_\theta (\nabla_\theta \log \pi \cdot v)\).
  2. Line search: after finding the natural gradient direction, perform a backtracking line search to find the largest step that satisfies the KL constraint and actually improves the surrogate objective.

This makes TRPO second-order in concept but avoids full Hessian computation.

TRPO vs PPO

TRPO’s theoretical guarantees are attractive, but the conjugate gradient solve and line search make it 10–100× more expensive per update than first-order methods. PPO (Schulman et al. 2017) approximates the same constraint with a clipping heuristic, achieving comparable empirical stability at a fraction of the cost.

For most practical applications, PPO is preferred. TRPO remains valuable when theoretical guarantees matter — for instance, in safety-critical settings — and its analysis underpins the theoretical foundations of PPO.

ACKTR: Natural Gradient with Kronecker Factoring

ACKTR (Wu et al. 2017) approximates the Fisher matrix using Kronecker-factored structure (K-FAC), enabling efficient second-order updates without conjugate gradient. It achieves the stability of TRPO with computational cost closer to first-order methods, and outperforms A2C and TRPO on several Atari and MuJoCo benchmarks.

References

  • Schulman, J., Levine, S., Moritz, P., Jordan, M., & Abbeel, P. (2015). Trust Region Policy Optimization. ICML. arXiv:1502.05477.
  • Amari, S. (1998). Natural Gradient Works Efficiently in Learning. Neural Computation, 10(2), 251–276.
  • Kakade, S. (2002). A Natural Policy Gradient. NeurIPS.
  • Wu, Y., Mansimov, E., Liao, S., Grosse, R., & Ba, J. (2017). Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation. NeurIPS. arXiv:1708.05144.
  • Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347.

You May Also Enjoy

GAPE: Remember to Forget — Gated Adaptive Positional Encoding

4 minute read

Published:

GAPE is a drop-in RoPE augmentation that adds content-aware attention logit biases: a query-gate suppresses irrelevant distant context while a key-gate preserves salient distant tokens. Provably sharper attention and improved long-context robustness — no architecture changes needed.

PolyNSD: Polynomial Neural Sheaf Diffusion

4 minute read

Published:

PolyNSD replaces the NSD propagation operator with a degree-K Chebyshev polynomial in the normalised sheaf Laplacian, achieving SOTA on homo- and heterophilic benchmarks with only diagonal restriction maps and dramatically lower memory usage.

TDA in Materials Science: Topology of Structure and Phase

3 minute read

Published:

Persistent homology characterises the multi-scale structure of materials — from pore geometry in catalysts to glass transition in amorphous materials. H₀ captures connectivity, H₁ captures channels and rings, H₂ captures enclosed voids. These topological descriptors predict mechanical properties, adsorption capacity, and phase transitions more accurately than geometric averages.

TDA in Drug Discovery: Molecular Topology

3 minute read

Published:

Persistent homology provides multi-scale topological fingerprints of molecular structures — capturing rings, cavities, and protein pockets that traditional cheminformatics descriptors miss. Applications include binding site detection, ADMET property prediction, and protein-ligand interaction modelling.