Trust Region Policy Optimization (TRPO)

While basic policy gradient methods update the policy by taking steps in the direction of the estimated gradient, choosing the right step size (learning rate) is challenging. Too small a step leads to slow progress, while too large a step can drastically change the policy, potentially leading to a catastrophic drop in performance from which the agent might never recover. We saw how Actor-Critic methods use a critic to get better gradient estimates, but the step size issue remains.

Trust Region Policy Optimization (TRPO) addresses this stability problem directly. Instead of just following the gradient, TRPO aims to maximize the policy's expected performance improvement in each step, but critically, it adds a constraint: the new policy must not diverge too far from the old policy. This constrained region around the current policy is the "trust region," implying we only trust our performance estimate within this local area.

The Core Idea: Constrained Policy Updates

Imagine you are adjusting the parameters $\theta$ of your policy $\pi_{\theta}$. Standard policy gradient methods tell you which direction $g$ (the gradient) increases performance locally. They then update $\theta_{k+1} = \theta_k + \alpha g$, where $\alpha$ is the learning rate. TRPO takes a different approach. It seeks to find the best possible parameters $\theta$ that maximize an objective function representing performance improvement, while ensuring that the "distance" between the new policy $\pi_{\theta}$ and the old policy $\pi_{\theta_k}$ remains small.

This "distance" is typically measured using the Kullback-Leibler (KL) divergence, $D_{KL}(\pi_{\theta_k} || \pi_{\theta})$, which quantifies how much the probability distribution defined by the new policy differs from the old one.

The TRPO Optimization Problem

TRPO aims to maximize a surrogate objective function, which approximates the expected advantage of the new policy $\pi_{\theta}$ relative to the old policy $\pi_{\theta_k}$. This objective uses importance sampling to correct for the fact that we are using data collected with $\pi_{\theta_k}$ to evaluate $\pi_{\theta}$:

$$L_{\theta_k}(\theta) = \mathbb{E}_{s \sim \rho_{\theta_k}, a \sim \pi_{\theta_k}} \left[ \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)} A^{\pi_{\theta_k}}(s,a) \right]$$

Here, $\rho_{\theta_k}$ is the distribution of states encountered when running policy $\pi_{\theta_k}$, and $A^{\pi_{\theta_k}}(s,a)$ is the advantage estimate (e.g., from GAE) under the old policy.

The key innovation in TRPO is constraining the update step. We want the average KL divergence between the old and new policies, over the states visited by the old policy, to be less than some small constant $\delta$:

$$\mathbb{E}_{s \sim \rho_{\theta_k}} [D_{KL}(\pi_{\theta_k}(\cdot|s) || \pi_{\theta}(\cdot|s))] \le \delta$$

So, the full TRPO optimization problem at each iteration $k$ is:

$$\theta_{k+1} = \arg \max_{\theta} L_{\theta_k}(\theta)$$ $$\text{subject to } \bar{D}_{KL}(\pi_{\theta_k} || \pi_{\theta}) \le \delta$$

where $\bar{D}_{KL}$ denotes the average KL divergence mentioned above. The hyperparameter $\delta$ defines the size of the trust region.

Practical Implementation via Approximations

Solving this constrained optimization problem exactly is difficult. TRPO makes it tractable using approximations:

1. Objective Approximation: The objective function $L_{\theta_k}(\theta)$ is approximated by its first-order Taylor expansion around $\theta_k$: $L_{\theta_k}(\theta) \approx L_{\theta_k}(\theta_k) + g^T (\theta - \theta_k)$, where $g = \nabla_{\theta} L_{\theta_k}(\theta)|_{\theta=\theta_k}$ is the policy gradient evaluated at the current parameters $\theta_k$. Since $L_{\theta_k}(\theta_k)$ is constant, maximizing this is equivalent to maximizing $g^T (\theta - \theta_k)$.

2. Constraint Approximation: The KL divergence constraint is approximated by its second-order Taylor expansion around $\theta_k$: $\bar{D}_{KL}(\pi_{\theta_k} || \pi_{\theta}) \approx \frac{1}{2} (\theta - \theta_k)^T F (\theta - \theta_k)$, where $F$ is the Fisher Information Matrix (FIM). The FIM acts like a metric on the policy parameter space, measuring how much the output distribution of the policy changes for small changes in parameters. It's defined as $F = \mathbb{E}_{\pi_{\theta_k}}[ \nabla_{\theta} \log \pi_{\theta}(a|s) \nabla_{\theta} \log \pi_{\theta}(a|s)^T ]|_{\theta=\theta_k}$.

Substituting these approximations, the problem becomes finding the update step $\Delta \theta = \theta - \theta_k$ that solves:

$$\max_{\Delta \theta} g^T \Delta \theta$$ $$\text{subject to } \frac{1}{2} \Delta \theta^T F \Delta \theta \le \delta$$

This is a standard quadratic program with a quadratic constraint.

The Conjugate Gradient Method

Calculating, storing, and inverting the FIM $F$ directly is computationally expensive, especially for large neural networks. TRPO cleverly avoids this by using the Conjugate Gradient (CG) algorithm. CG is an iterative method for solving linear systems of equations of the form $Fx = g$. Importantly, CG only requires the ability to compute matrix-vector products, specifically $Fv$ for arbitrary vectors $v$. This FIM-vector product can be computed efficiently without explicitly forming $F$, using techniques similar to calculating the gradient itself (often involving automatic differentiation).

CG approximately finds the optimal step direction $\Delta \theta \propto F^{-1} g$. The magnitude of the step is then scaled to satisfy the KL constraint $\frac{1}{2} \Delta \theta^T F \Delta \theta = \delta$.

Line Search

Because the Taylor approximations are only accurate locally, the step calculated via CG might not actually satisfy the KL constraint or improve the true (unapproximated) objective function. Therefore, TRPO performs a line search: it tries exponentially shrinking step sizes along the direction found by CG $(\alpha^j \Delta \theta$ for $j=0, 1, 2, ...$, where $\alpha \in (0, 1)$ is a backtracking coefficient) until it finds a step that both satisfies the average KL constraint and yields a positive value for the surrogate objective $L_{\theta_k}$.

Conceptual Flow of a TRPO Update

The following diagram illustrates the sequence of operations within a single TRPO policy update:

This diagram shows the typical workflow for one update step in TRPO. It involves data collection, advantage estimation, gradient calculation, using the conjugate gradient method with Fisher-vector products to find a search direction, and finally a line search to ensure constraints are met before updating the policy.

Advantages and Disadvantages of TRPO

Advantages:

• Stability: Generally provides more stable and monotonic performance improvements compared to standard policy gradient methods.
• Theoretical Foundation: Based on optimizing a well-defined objective subject to a meaningful constraint related to policy change.
• Strong Performance: Often achieves good results on complex continuous control benchmarks.

Disadvantages:

• Computational Complexity: Significantly more computationally expensive per update than first-order methods like A2C or vanilla policy gradients due to the CG iterations and the need for second-order information (FIM-vector products).
• Implementation Complexity: More complex to implement correctly compared to simpler policy gradient variants.
• Compatibility Issues: The derivation and approximations used in TRPO don't easily accommodate certain neural network techniques like parameter sharing between actor and critic, dropout, or batch normalization, as these can interfere with the accurate computation of the KL divergence and its quadratic approximation.

TRPO represented a significant step forward in developing reliable policy optimization algorithms. However, its complexity motivated the search for simpler methods that could achieve similar stability guarantees. This led directly to the development of Proximal Policy Optimization (PPO), which we will discuss next.