Advanced Policy Gradients: TRPO

Created on July 06, 2026

2026   ·   rl   policy-gradients   trpo   trust-region   fisher-information   kl-divergence   total-variations   ·   reinforcement-learning

In the previous post we explained the vanilla policy gradient (VPG) and actor-critic (A2C) methods. In this post, we will study some advanced policy gradient methods. Shown below is the reward maximization objective which is the average value of total rewards on trajectories parametrized by $\theta$.

\[\begin{align*} J(\theta) = \mathbb{E}_{\tau \sim p(\tau|\theta)} [\sum_t r(s_t,a_t)] \end{align*}\]

$p(\tau\vert\theta)$ is the trajectory probability derived in the previous post. It depends on initial state, policy and transition probabilities. In VPG and A2C, we use gradient based methods (SGD/ADAM) to improve upon the policy. We compute $\nabla_\theta J(\theta)$ and update the policy parameters:

\[\begin{align*} \nabla_\theta J(\theta) = \sum_{t=0}^{T} \mathbb{E}_{s_t, a_t} \left[ \nabla_\theta \log \pi_\theta(a_t|s_t) A^\pi(s_t, a_t) \right] \end{align*}\] \[\begin{align*} \theta_{new} \leftarrow \theta_{old} + \alpha \nabla_\theta J(\theta_{old}) \end{align*}\]

The gradient ascent should in theory improve the objective at every update. However, in practice, the gradient estimate is noisy and based upon sampling, therefore, it is not guaranteed that every update step we take improves the objective. An interesting example discussed in Kakade2002 paper is the case when reward is sparse i.e. reward comes at the end of the episode. If we take random steps, the path to the reward state is exponential in terms of the number of states in our MDP. Thus, any advantage estimate will have 0 value (since we won’t see any reward in meaningful number of steps). The estimated gradient as a result will be zero and the policy will get stuck.

One can note that gradient is an expectation over state $s_t$ i.e. it depends on the state visitation probability ($p(s_t)$)

\[\begin{align*} \nabla_\theta J(\theta) = \sum_{t=0}^{T} p(s_t) p(a_t|s_t) \nabla_\theta \log \pi_\theta(a_t|s_t) A^\pi(s_t, a_t) \end{align*}\]

States which have low visitation probability, may not be sampled in our estimate despite having possibly a high value. The situation worsens as policy is updated, gradient update further lowers the already low visitation probability because it may find the current state has good reward. There may be better reward waiting at another state but it was never visited and as policy gets updated chances of its visitation reduce further. Policy gradient methods seem to exploit more and explore less. This is also a classic local optima trap, gradient estimates vanish and policy is stuck. Policy initialization will also play a significant role in final converged policy.

Guaranteeing Improvement in every update

We reformulate the reward maximization problem to maximize change in reward as the policy is updated ($J(\theta_{new}) - J(\theta_{old})$) rather than total reward. $\theta_{new}$ represents the updated parameters of the new policy and it is the variable we want to optimize. $\theta_{old}$ represents the parameters of the old policy before the update. We will show later that by introducing appropriate constraints to this problem we can guarantee improvement in every update step. Note, $\theta_{old}$ is fixed, only $\theta_{new}$ is the variable in this problem.

\[\begin{align*} J(\theta_{new}) - J(\theta_{old}) &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t r(s_t,a_t)] - \mathbb{E}_{\tau \sim p(\tau|\theta_{old})} [\sum_t \gamma^t r(s_t,a_t)]\\ &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t r(s_t,a_t)] - \mathbb{E}_{s_0\sim p(s_0)} [V^{\theta_{old}}(s_0)]\\ &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t r(s_t,a_t)] - \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [V^{\theta_{old}}(s_0)] \end{align*}\]

The reason we can make the last change to $p(\tau\vert\theta_{new})$ is because rest of the states of transition don’t impact the inner expectation. Instead of taking expectation over the initial state marginal distribution, we take expectation over the entire joint distribution of trajectory. We can show that taking expectation over marginal distribution is same as taking expectation over the joint distribution.

\[\begin{align*} \mathbb{E}_{x \sim p(x)} [f(x)] &= \int f(x) p(x) \, dx \\ &= \int f(x) \left( \int p(x, y) \, dy \right) dx \\ &= \int_x \int_y f(x) p(x, y) \, dy \, dx \\ &= \mathbb{E}_{x, y \sim p(x, y)} [f(x)] \end{align*}\] \[\begin{align*} J(\theta_{new}) - J(\theta_{old}) &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t r(s_t,a_t)] - \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [V^{\theta_{old}}(s_0)]\\ &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t r(s_t,a_t)] - \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_{t=0} \gamma^t V^{\theta_{old}}(s_t)-\gamma^{t+1} V^{\theta_{old}}(s_{t+1})]\\ &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t( r(s_t,a_t) + \gamma V^{\theta_{old}}(s_{t+1})-V^{\theta_{old}}(s_t))]\\ &= \mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t A^{\theta_{old}}(s_t,a_t)] \end{align*}\]

where second equality is a telescopic sum and assumes infinite horizon episodes.

Maximizing the objective $J(\theta_{new})$ is same as maximizing above function since $\theta_{old}$ is a parameter from the previous iteration and is known. The only place $\theta_{new}$ appears in the objective function is in the expectation. The catch is computing the expectation over $\theta_{new}$ is not easy in practice because we are maximizing over a distribution from which we need to sample to form the estimates of the objective. We will expand further:

\[\begin{align*} &\max_{\theta_{new}}\mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t A^{\theta_{old}}(s_t,a_t)]\\ &\max_{\theta_{new}} \sum_t \mathbb{E}_{s_t,a_t \sim p(s_t,a_t)}[\gamma^t A^{\theta_{old}}(s_t,a_t)]\\ &=\max_{\theta_{new}}\sum_t \mathbb{E}_{s_t\sim p(s_t|\theta_{new})}[\mathbb{E}_{a_t\sim \pi^{\theta_{new}}(a_t|s_t)}[\gamma^t A^{\theta_{old}}(s_t,a_t)]]\\ &= \max_{\theta_{new}}\sum_t \mathbb{E}_{s_t\sim p(s_t|\theta_{new})}[\mathbb{E}_{a_t\sim \pi^{\theta_{old}}(a_t|s_t)}[\frac{\pi^{\theta_{new}}(a_t|s_t)}{\pi^{\theta_{old}}(a_t|s_t)}\gamma^t A^{\theta_{old}}(s_t,a_t)]] \end{align*}\]

Using importance sampling we can get rid of sampling from $\theta_{new}$ dependent distribution for inner action sampling but for outer samples we still need to sample from state marginal distribution.

We can show that the approximation of replacing $p(s\vert\theta_{new})$ with $p(s\vert\theta_{old})$ will work if policies $\pi$ are also close to each other i.e. if $\pi^{\theta_{new}}(a_t\vert s_t)$ is close to $\pi^{\theta_{old}}(a_t\vert s_t)$. In that case we can write the optimization problem above as:

\[\begin{align*} &\max_{\theta_{new}}\mathbb{E}_{\tau \sim p(\tau|\theta_{new})} [\sum_t \gamma^t A^{\theta_{old}}(s_t,a_t)]\\ &\approx \max_{\theta_{new}}\sum_t \mathbb{E}_{s_t\sim p(s_t|\theta_{old})}[\mathbb{E}_{a_t\sim \pi^{\theta_{old}}(a_t|s_t)}[\frac{\pi^{\theta_{new}}(a_t|s_t)}{\pi^{\theta_{old}}(a_t|s_t)}\gamma^t A^{\theta_{old}}(s_t,a_t)]]\\ & s.t. \ \ \pi^{\theta_{new}}(a_t|s_t) \text{ is close to } \pi^{\theta_{old}}(a_t|s_t) \end{align*}\]

The reason above approximation works is provided in the TRPO paper. They show:

\[\begin{align} \label{eq:trpo_bound} J(\theta_{new}) - J(\theta_{old}) &\geq \sum_t \mathbb{E}_{s_t\sim p(s_t|\theta_{old})}[\mathbb{E}_{a_t\sim \pi^{\theta_{old}}(a_t|s_t)}[\frac{\pi^{\theta_{new}}(a_t|s_t)}{\pi^{\theta_{old}}(a_t|s_t)}\gamma^t A^{\theta_{old}}(s_t,a_t)]] - cf(TV(\pi_{new}, \pi_{old})^2) \end{align}\]

This implies we can maximize the lower bound on right hand side. TV is the total variation distance and it’s a measure of distance between two distributions. Further, the second term: $cf(TV(\pi_{new}, \pi_{old})^2)$ can approximately be moved to constraint. Let’s understand total variation first.

Total Variation Distance

Total variation distance (TVD) is a measure of similarity between two distributions. Let $P$ and $Q$ be two distributions, for any event $a$, you can assign probability to $a$ through $P$ and $Q$ as $P(a)$ and $Q(a)$ respectively. The absolute difference, $\vert P(a) - Q(a)\vert$ gives how much the distributions differ for event $a$. If two distributions are similar, for most events they will assign similar scores. TVD gives the worst case difference in probabilities across all events $a$:

\[\begin{align*} TV(P,Q) = \max_{a} |P(a)-Q(a)| \end{align*}\]

If two distributions are identical across all events, $TV(P,Q) = 0$. If two distributions are defined such that for each outcome where $P$ assigns non-zero probability, $Q$ assigns zero probability and vice-versa, then we can create an event as union of these outcomes such that $P(a) = 1$ and $Q(a) = 0$. Such a distribution will have the highest TVD i.e. $TV(P,Q)=1$. TVD is a measure of distance between two distributions and can also be written as:

\[\begin{align*} TV(P,Q) = \frac{1}{2}\sum_x |P(x) - Q(x)| \leq 1 \end{align*}\]

Let $X\sim P$ and $Y \sim Q$ be two random variables. Let $J(X,Y)$ be joint distribution of $X$ and $Y$. We can have multiple joint distributions s.t. $P(X) = \sum_y J(X,Y)$ and $Q(Y) = \sum_x J(X,Y)$ are the marginal distributions. It can be shown that:

\[\begin{align*} TV(P,Q) \leq J(X\neq Y) \end{align*}\]

Also, there exists some clever joint distribution where equality can be obtained i.e.

\[\begin{align*} TV(P,Q) = J^*(X\neq Y) = \min_J J(X\neq Y) \end{align*}\]
TVD Example

What does $TV(P,Q)<\epsilon$ imply? $TV(P,Q) = J^*(X\neq Y) \leq \epsilon$ means, there exists some joint distribution in which both $X$ and $Y$ take the same values for all events with probability greater than $1-\epsilon$.

For example: let P be an unbiased coin, Q be a biased coin with Q(heads)=0.7 i.e. $TV(P,Q)=0.2$. This implies under some joint distribution both coins will take the same values with $Pr \geq 0.8$. If you assume both coins are independent $Pr(X=Y) = 0.5$, so the distribution we are looking for is not independent. We can solve a system of equations and get to this optimal distribution. Such a distribution looks something like this:

\[\begin{align*} J(H,H) = 0.5, J(T,T) = 0.3, J(H,T) = 0, J(T,H) = 0.2 \end{align*}\]

Note, $P(X=H) = 0.5$ and $Q(Y=H) = 0.7$ from the distribution above, while $P(X=Y) = 0.8$. This property is useful in proving why we can replace \(\mathbb{E}_{s_t\sim p(s_t \|\theta^{new})}\) with $\mathbb{E}_{s_t\sim p(s_t|\theta^{old})}$ which in turn is useful in computing the bound in eq $\eqref{eq:trpo_bound}$.

Claim. $p(s_t\vert\theta^{new})$ is close to $p(s_t\vert\theta^{old})$ if $\pi^{new}(a_t\vert s_t)$ is close to $\pi^{old}(a_t\vert s_t)$

Proof. We want to show if action distribution is close then state marginal distributions are also close, i.e. if $TV(\pi^{new}(a_t\vert s_t), \pi^{old}(a_t\vert s_t)) < \epsilon$ then $TV(p(s_t\vert\theta^{new}),p(s_t\vert\theta^{old}))<c$

Intuitively, this makes sense, since if action distribution is close, both trajectories as they unroll end up taking same actions with high probability. For computing marginal distribution note, it only depends on action distribution and transition probabilities. Since transition probability is policy independent, as long as actions are the same, marginal distribution remains the same.

If total variation is less than $\epsilon$, then we know there is some joint distribution where both policies actions will be same with high probability i.e. $Pr(X = Y) \geq 1 - \epsilon$. Let’s focus on step $t$. We will compute $p(s_t\vert\theta^{new})$:

\[\begin{align*} p(s_t|\theta^{new}) &= p(s_t|\theta^{new}, \text{all t past actions for both policies are same})\\ &. P(\text{all t past actions for both policies are same})\\ & + p(s_t|\theta^{new}, \text{all t past actions for both policies are not same})\\ &.P(\text{all t past actions for both policies are not same}) \end{align*}\]

At each state probability of taking the same action is $\geq 1-\epsilon$. If we make a mistake in any of the $t$ steps we can’t claim $p(s_t\vert\theta^{new})$ is same as $p(s_t\vert\theta^{old})$ rather the distribution under new policy will be some other distribution we can call $p_{mistake}$

\[\begin{align*} p(s_t|\theta^{new}) &= (1-\epsilon)^t p(s_t|\theta^{old}) + (1-(1-\epsilon)^t) p_{mistake}(s_t) \\ p(s_t|\theta^{new}) - p(s_t|\theta^{old}) &= -(1-(1-\epsilon)^t)p(s_t|\theta^{old})+ (1-(1-\epsilon)^t) p_{mistake}(s_t)\\ &= (1-(1-\epsilon)^t) (p_{mistake}(s_t) - p(s_t|\theta^{old})) \end{align*}\] \[\begin{align*} TV(p(s_t|\theta'),p(s_t|\theta)) &= (1-(1-\epsilon)^t) TV(p_{mistake}(s_t),p(s_t|\theta))\\ &\leq (1-(1-\epsilon)^t) \leq \epsilon t \end{align*}\]

where last inequality follows from $(1-\epsilon)^t > 1 - \epsilon t$ if $\epsilon \in [0,1]$. $\blacksquare$

So, we have shown if action is $\epsilon$ close in TV terms, marginals are close in TV terms for small epsilon but bound becomes looser as $t$ increases.

Total Variation to KL divergence

Enforcing TV bounds can be hard given non differentiability. So, a more convenient bound is by relating TV with KL divergence via inequality:

\[\begin{align*} TV(P,Q) \leq \sqrt{0.5*KL(P,Q)} \end{align*}\]

We solve:

\[\begin{align*} & \max_{\theta^{new}}\sum_t \mathbb{E}_{s_t\sim p(s_t|\theta)}[\mathbb{E}_{a_t\sim \pi^{\theta}(a_t|s_t)}[\frac{\pi^{\theta^{new}}(a_t|s_t)}{\pi^{\theta}(a_t|s_t)}\gamma^t A^\theta(s_t,a_t)]]\\ & s.t. KL(\pi^{\theta^{new}}(a_t|s_t), \pi^{\theta}(a_t|s_t))<\epsilon \end{align*}\]

Solving the constraint optimization problem

The next part of this blog pertains to solving above optimization problem. There are several ways to attack the problem. Below is a penalty based method where infeasible points penalize the objective function. This particular formulation is based on optimizing the Lagrangian:

\[\begin{align*} L(\theta', \lambda) = \sum_t \mathbb{E}_{s_t\sim p(s_t|\theta)}[\mathbb{E}_{a_t\sim \pi^{\theta}(a_t|s_t)}[\frac{\pi^{\theta'}(a_t|s_t)}{\pi^{\theta}(a_t|s_t)}\gamma^t A^\theta(s_t,a_t)]] - \lambda_k ( KL(\pi^{\theta'}(a_t|s_t), \pi^{\theta}(a_t|s_t)) - \epsilon) \end{align*}\]

We solve the problem iterating between two steps. In first step we maximize the objective + penalty term and in the second step we adapt the lagrange multiplier based on constraint satisfaction:

  1. At given $\lambda_k$, $\theta’_{k+1} \leftarrow \max L(\theta’, \lambda_k)$
  2. $\lambda_{k+1} \leftarrow \lambda_t + \alpha(KL(\pi^{\theta’_{k+1}}(a_t\vert s_t), \pi^{\theta}(a_t\vert s_t)) - \epsilon)$

So in step 2 we decrease the penalty if the constraint is satisfied and increase the penalty if it is not satisfied.

The problem with this approach is you still need to solve a massive optimization problem in step 1 and we have to do that iteratively. Next, we discuss an alternative way to solve this problem via Taylor series expansion of both objective and constraint. This approach was proposed as natural gradients in Kakade2001. Before going into natural gradients, we will first understand fisher information, then we will understand KL divergence, we will see how both are related. Then we will come back to the optimization problem and rewrite solution in form of fisher information and interpret it.

Fisher Information

The usual parameter estimation process is to first make a choice of distribution that will model given data. This statistical model has parameters $\theta$ and a functional form (PMF or PDF) that gives assignment to any outcome of the experiment (as $p(x;\theta)$). We then estimate $\theta$ based on the likelihood of observations coming from the distribution. Fisher information gives a measure of the ability of modeling choice $p$ to estimate $\theta$. Consider an RV $X$ which is once modeled with gaussian distribution with mean $\theta$ and another time with gaussian distribution with mean $\theta$ and variance $25$. Which one will be able to estimate the mean better? One feels variance = 1 will do better, since its less dispersed. Fisher information is denoted by $I_X(\theta)$. We should ideally subscript on modeling choice i.e. $I_{\sigma=1}(\theta)$ or $I_{\sigma=5}(\theta)$. We use the $I_X(\theta)$ notation, since for each distribution a different RV is assumed.

Different distributions can be based on same underlying parameter. For instance, let RV $X$ be defined as $1$ if coin tosses head and $0$ otherwise. The underlying parameter $\theta$ is true probability of coin landing head. We can compute $I_x(\theta)$ as information in single toss of coin. We can also toss the coin 100 times and let’s say RV $D$ represents a sequence of coin tosses. We can intuitively say $I_D(\theta) > I_x(\theta)$.

Note, fisher information is not dependent on the actual samples collected, it’s based on the choice of distribution function used for modeling. A single toss, whether head or tail will have less Fisher information than $N$ tosses. Similarly, sample collected from high variance gaussian will have less FI. Fisher information is a function of $\theta$ i.e. it depends on the true value of parameter we are estimating as well. We will formally define Fisher information and relate it to RL optimization problem.

Score function

$P(X;\theta)$ is the PDF/PMF of a distribution parameterized by $\theta$. $P(X=x;\theta)$ gives probability random variable takes value of $x$. What happens to this value when we change true parameter value $\theta$? For some $x$ probability may increase, for some it may reduce. The slope of change at a particular point is the score at that point. The assignment of probability to $x$ is completely dependent on the model we chose.

\[\begin{align*} S(X;\theta) &= \frac{d\log P(X;\theta)}{d\theta} \end{align*}\]

For multivariate $\theta$, score is a vector

\[\begin{align*} S(X;\theta) &= \nabla_\theta \log P(X;\theta) \end{align*}\]

Notice, score function is also a RV (since its function of $X$). We can compute expected value and variance of score function. It is easy to show the $E[S]=0$ by plugging in the definition and taking expectation.

High variance of score means that as $\theta$ changes the probability assignment of each sample changes wildly. As you change $\hat{\theta}$, your estimate value, to fit the sample, you will get high likelihood at some estimate and at most other values of $\theta$ likelihood will be low. Think of this as turning the knob of a radio, high variance implies at some $\theta$ data will suddenly fit well and changing it even slightly may lead to considerably lower likelihood.

The variance of the score function is the fisher information $I(\theta)$.

\[\begin{align*} I_X(\theta) &= Var(S(X;\theta))\\ &= E[S^2] - E[S]^2\\ &= E[S^2] \end{align*}\]

For multivariate case variance is the co-variance matrix called fisher information matrix (FIM):

\[\begin{align} \label{eq:fim} I_X(\theta) &= Var(S(X;\theta))\\ &= E[SS^T]\\ &= E[\nabla_\theta \log P(X;\theta). \nabla_\theta \log P(X;\theta)^T] \end{align}\]

Note how FIM is average of rank 1 matrices.

Another interesting property of FI is it can also be defined as the expected value of hessian of log-likelihood.

\[\begin{align*} E[\nabla_\theta^2 \log P(X;\theta)]=-E[\nabla_\theta \log P(X;\theta). \nabla_\theta \log P(X;\theta)^T] = -I_X(\theta) \end{align*}\]

It can be shown further that FI is related to the variance of an estimate of $\theta$. Consider you create an unbiased estimator $\hat{\theta} = g(X)$. As it is unbiased estimate, $E[\hat{\theta}] = \theta$, where expectation is over RV $X$. We take derivative wrt $\theta$ on both side

\[\begin{align*} \frac{d\int \hat \theta p(x;\theta) dx}{d\theta} &= 1 \\ \frac{\int \hat \theta dp(x;\theta) dx}{d\theta} &= 1 \\ \int p(x) \hat \theta \frac{d\log p(x;\theta) dx}{d\theta} &= 1 \\ E[\hat\theta S(X;\theta)] &= 1 \end{align*}\]

Further, we can say $Cov(\hat\theta, S(X;\theta)) = E[\hat\theta S(X;\theta)] - E[\hat\theta]\underbrace{E[ S(X;\theta)]}_{=0} = 1$.

Note, $Cov(X,Y)^2 \leq Var(x) Var(Y)$ (follows from correlation definition which is always less than 1). Using above we can see:

\[\begin{align*} Var(S).Var(\hat \theta) &\geq 1 \\ Var(\hat \theta) \geq \frac{1}{I_X(\theta)} \end{align*}\]

This shows variance of unbiased estimator is bounded on the low end by inverse of fisher information. The lower the fisher information in an estimate, the more the variance. This is also the Cramer Rao Bound. We can also apply this further to write mean squared error in estimate of parameters as sum of $MSE = bias(\hat \theta)^2 + Var(\hat \theta)$. For an unbiased error, we show we can’t do better than the inverse of fisher information. So, lower information means higher error estimate.

All this is cool, how do we relate fisher information and KL divergence? We will compute KL divergence between two distributions which are slightly perturbed i.e $KL(\pi^{\theta’}\Vert\pi^\theta)$, this is a function of $\theta’$ and $\theta$:

\[\begin{align*} KL(\pi^{\theta'}||\pi^\theta) &= {KL(\pi^{\theta}||\pi^\theta)} + \nabla_{\theta'}KL(\pi^{\theta'}||\pi^\theta)|_{\theta} (\theta' - \theta) + 0.5 (\theta' - \theta)^T\nabla^2_{\theta'}KL(\pi^{\theta'}||\pi^\theta)|_{\theta} (\theta' - \theta) \\ &= \underbrace{KL(\pi^{\theta}||\pi^\theta)}_{=0} + \nabla_{\theta'}KL(\pi^{\theta'}||\pi^\theta)|_{\theta} (\theta' - \theta) + 0.5 (\theta' - \theta)^T\nabla^2_{\theta'}KL(\pi^{\theta'}||\pi^\theta)|_{\theta} (\theta' - \theta) \end{align*}\]

Now lets compute the first gradient, which turns out to be $0$:

\[\begin{align*} KL(\pi^{\theta'}||\pi^\theta) &= \sum_a \pi^{\theta'}(\log \pi^{\theta'} - \log \pi^{\theta})\\ \nabla_{\theta'}KL(\pi^{\theta'}||\pi^\theta) &= \sum_a d\pi^{\theta'}(\log \pi^{\theta'} - \log \pi^{\theta}) + d\pi^{\theta'} \\ \nabla_{\theta'}KL(\pi^{\theta'}||\pi^\theta)|_{\theta} &= d\sum_a \pi^{\theta} = 0 \end{align*}\]

Second gradient:

\[\begin{align*} \nabla_{\theta'}KL(\pi^{\theta'}||\pi^\theta) &= \sum_a \nabla_{\theta'}\pi^{\theta'}(\log \pi^{\theta'} - \log \pi^{\theta}) + \nabla_{\theta'}\pi^{\theta'} \\ \nabla^2_{\theta'}KL(\pi^{\theta'}||\pi^\theta)&= \sum_a (\log \pi^{\theta'} - \log \pi^{\theta}+1) {\nabla_{\theta'}^2\pi^{\theta'}}_{d\times d}+ \sum_a \nabla_{\theta'}\pi^{\theta'}_{d \times 1}(\frac{ {\nabla_{\theta'}\pi^{\theta'}}^T}{\pi^{\theta'}})\\ &= \sum_a (\log \pi^{\theta'} - \log \pi^{\theta}+1) {\nabla_{\theta'}^2\pi^{\theta'}}_{d\times d} + \sum_a \pi^{\theta'}\nabla_{\theta'}\log\pi^{\theta'}(\frac{ \pi^{\theta'} \nabla_{\theta'}\log \pi^{\theta'}}{\pi^{\theta'}})\\ \nabla^2_{\theta'}KL(\pi^{\theta'}||\pi^\theta) |_{\theta}&= \underbrace{\nabla_{\theta'}^2\sum_a {\pi^{\theta'}}}_{=0} + I_{a|s}(\theta')|_{\theta}\\ &= I_{a|s}(\theta) \end{align*}\]

This makes sense, since the more fisher information you have, the more the distribution will change as the underlying parameter changes, causing the policies to drift further apart and hence higher KL divergence. We finally have KL in terms of fisher information as:

\[\begin{align*} KL(\pi^{\theta'}||\pi^\theta) &\approx \frac{1}{2} (\theta' - \theta)^T I_{a|s}(\theta)(\theta' - \theta) \end{align*}\]

Back to optimization problem

Coming back to the problem, restating here:

\[\begin{align*} & \max_{\theta'}\sum_t \mathbb{E}_{s_t\sim p(s_t|\theta)}[\mathbb{E}_{a_t\sim \pi^{\theta}(a_t|s_t)}[\frac{\pi^{\theta'}(a_t|s_t)}{\pi^{\theta}(a_t|s_t)}\gamma^t A^\theta(s_t,a_t)]]\\ & s.t. KL(\pi^{\theta'}(a_t|s_t), \pi^{\theta}(a_t|s_t))<\epsilon \ \ \forall s_t \end{align*}\]

We do two more changes now, we linearize the objective by taking first order taylor approximation and we replace KL constraint with fisher information.

Linearlize Objective

For easy writing, let’s write the objective in terms of the variable $\theta’$. Note $\theta$ is considered a constant here.

\[\begin{align*} L(\theta') &= \sum_t \mathbb{E}_{s_t\sim p(s_t|\theta)}[\mathbb{E}_{a_t\sim \pi^{\theta}(a_t|s_t)}[\frac{\pi^{\theta'}(a_t|s_t)}{\pi^{\theta}(a_t|s_t)}\gamma^t A^\theta(s_t,a_t)]] \end{align*}\]

By taylor’s approximation around $\theta$:

\[\begin{align*} L(\theta') &\approx L(\theta) + \nabla_{\theta'}L(\theta')|_{\theta}(\theta' - \theta) \end{align*}\]

Since we are maximizing under variable $\theta’$, we can drop the first term from the approximation to get equivalent problem as:

\[\begin{align*} & \max_{\theta'} \nabla_{\theta'}L(\theta')|_{\theta}(\theta' - \theta)\\ & s.t. \frac{1}{2} (\theta' - \theta)^T I_{a|s}(\theta)(\theta' - \theta)<\epsilon \ \ \forall s \end{align*}\]

The number of constraints is large in above formulation, one for each state in state-space. One approximation done in the TRPO paper is to replace state specific constraint with a single average constraint:

\[\begin{align*} & KL(\pi^{\theta'}(a_t|s_t), \pi^{\theta}(a_t|s_t))<\epsilon \ \ \forall s_t \\ & \approx E_{s_t\sim p(s_t|\theta)}[KL(\pi^{\theta'}(a_t|s_t), \pi^{\theta}(a_t|s_t)]\\ & = \frac{1}{2} (\theta' - \theta)^T \bar{I}_{a|s}(\theta)(\theta' - \theta)<\epsilon \end{align*}\]

Let’s also compute $\nabla_{\theta’}L(\theta’)$. Only part where $\theta’$ appears is in policy. We can take derivative and replace with log to get:

\[\begin{align*} \nabla_{\theta'}L(\theta')|_{\theta} &= \sum_t \mathbb{E}_{s_t\sim p(s_t|\theta)}[\mathbb{E}_{a_t\sim \pi^{\theta}(a_t|s_t)}[\nabla_{\theta'} \log \pi^{\theta'}(a_t|s_t)|_{\theta}\gamma^t A^\theta(s_t,a_t)]] \end{align*}\]

Notice this is the standard policy gradient. Only difference is now we also have constraint on fisher information. Note by definition fisher information is symmetric positive semi-definite. Using the lagrangian method and writing KKT conditions shows we can write this as solving $I^{-1} \nabla_\theta L(\theta)$

\[\begin{align*} L(\theta',\lambda) &= \nabla_{\theta'}L(\theta')|_{\theta}(\theta' - \theta) -\lambda ( \frac{1}{2} (\theta' - \theta)^T \bar{I}_{a|s}(\theta)(\theta' - \theta) -\epsilon)\\ \nabla_{\theta'} L(\theta',\lambda) &= \nabla_{\theta'}L(\theta')|_{\theta} - \lambda \bar{I}_{a|s} (\theta' - \theta) = 0 \\ \theta' &= \theta + \frac{1}{\lambda} \bar{I}_{a|s}^{-1} \nabla_{\theta'}L(\theta')|_{\theta} \end{align*}\]

This gives the update equation. Equality to zero follows from the KKT conditions. Since $I$ is huge, $d\times d$, where $d$ is the number of parameters in the model, computing its inverse is $O(d^3)$, hence infeasible. Note $I$ is symmetric positive semi definite, since it is expectation over rank 1 positive semi-definite matrices (see eq $\eqref{eq:fim}$). In such a case we can use conjugate gradient method for solving a system of linear equations, rather than taking inverse. The theory of CG I need to read, found a good link for later: paper

The above KKT based solution, still needs to honor the quadratic constraint. We can solve this iteratively as the penalty method mentioned earlier. We first compute the search direction as natural gradient $s = I^{-1}L$, then we take a step in this direction i.e. $\theta’ = \theta + \beta s$. $\beta$ is computed as $\sqrt(\frac{\epsilon}{s^TI(\theta)s})$ i.e. maximize the step size till constraint is satisfied. Next, beta is reduced till we see improvement in objective ($J(\theta’)-J(\theta)$). Below is the summary of steps taken:

  1. Choose lambda to maximize the step size so that fisher information constraint is still satisfied
  2. The obtained step size may lead us to go outside taylor approximation zone, so reduce $\lambda$ till we get original (non linearized) objective i.e. $J(\theta’)-J(\theta)$ improved and constraint is also satisfied.

This is the TRPO algorithm for policy optimization. There are some more subtleties on practical implementation which we skip in this post. In further posts, we will go in depth into other advanced RL methods like PPO and GRPO.