All Lessons

Constraining Policy Updates via Kullback–Leibler Divergence

This lesson explores the fundamental role of KL divergence in stabilizing reinforcement learning algorithms. We will dissect how measuring the distance between probability distributions prevents catastrophic policy collapses during optimization.

AI Narration Press play to listen
0  / 7 paragraphs
Click any paragraph to jump · Scroll freely without breaking narration

At the heart of modern policy optimization lies a simple yet profound intuition: we want to improve our agent's behavior, but we must not change it too drastically in a single step. Imagine navigating a foggy mountain; taking giant leaps might accidentally send you off a cliff, whereas small, measured steps ensure you stay on the path. In machine learning, the 'distance' between the old policy and the new policy is quantified by the Kullback–Leibler (KL) divergence. This metric tells us how much information is lost when we approximate the new distribution with the old one, serving as a trust region for our updates.

Mathematically, the KL divergence between two discrete probability distributions, $P$ (the old policy) and $Q$ (the new policy), is defined as the expected logarithmic difference between the probabilities. The formula is expressed as $$D_{KL}(P || Q) = \sum_x P(x) \log \left( \frac{P(x)}{Q(x)} \right)$$. It is crucial to note that this measure is asymmetric, meaning $D_{KL}(P || Q) \\≠ D_{KL}(Q || P)$. In the context of policy gradients, we typically measure how much the new policy $\pi_\theta$ diverges from the old policy $\pi_{\theta_{old}}$ under the state distribution visited by the old policy.

When updating a policy parameterized by $\theta$, our primary objective is to maximize the expected return $J(\theta)$. However, unconstrained gradient ascent can lead to large updates that degrade performance, a phenomenon often called 'policy collapse.' To mitigate this, we frame the update as a constrained optimization problem. We seek to maximize the surrogate objective function subject to a constraint on the KL divergence: $$\text{maximize}_\theta \quad \hat{J}(\theta) \quad \text{subject to} \quad D_{KL}(\pi_{\theta_{old}} || \pi_\theta) \le \delta$$. Here, $\delta$ is a hyperparameter controlling the maximum allowable step size in the probability space.

Solving this constrained problem directly is computationally expensive, so algorithms like Trust Region Policy Optimization (TRPO) use approximations. By employing a second-order Taylor expansion, we approximate the objective function as linear and the KL constraint as quadratic. The KL divergence can be approximated locally using the Fisher Information Matrix $F$, such that $$D_{KL}(\pi_{\theta_{old}} || \pi_\theta) \approx \frac{1}{2} (\theta - \theta_{old})^T F (\theta - \theta_{old})$$. This transforms the problem into a quadratically constrained linear program, which has a closed-form solution involving the inverse of the Fisher matrix.

In practice, computing the inverse of the Fisher Information Matrix is often intractable for high-dimensional neural networks. Consequently, modern approaches like Proximal Policy Optimization (PPO) relax the hard constraint into a penalty term within the loss function. Instead of a rigid boundary, we add a KL penalty coefficient $\beta$ to the objective: $$L^{CLIP}(\theta) = \hat{J}(\theta) - \beta \\· D_{KL}(\pi_{\theta_{old}} || \pi_\theta)$$. This allows the algorithm to adaptively adjust the step size; if the KL divergence grows too large, the penalty term dominates the gradient, naturally pushing the update back toward the old policy.

The significance of KL divergence extends beyond mere stabilization; it ensures that the data distribution used for training remains relevant. Reinforcement learning relies on importance sampling to estimate gradients using data collected from an older policy. If the new policy diverges too far from the behavior policy that generated the data, the importance sampling ratios become unstable, leading to high variance estimates. By constraining the KL divergence, we ensure that the off-policy data remains a valid proxy for on-policy performance, maintaining the statistical integrity of the learning process.

To summarize, the Kullback–Leibler divergence acts as the guardian of stability in policy gradient methods. It translates the abstract concept of 'trust' into a rigorous mathematical boundary, preventing the optimizer from making reckless updates that could destroy learned behaviors. Whether implemented as a hard constraint in TRPO or a soft penalty in PPO, monitoring $D_{KL}$ is essential for achieving consistent and robust convergence in complex environments. Without this constraint, deep reinforcement learning would likely remain too unstable for practical application.