28  Reinforcement Learning and Optimization

Core idea: reinforcement learning is sequential decision-making, but many of its algorithms can be understood as optimization methods. Bellman equations are fixed-point optimization conditions, policy gradients are stochastic optimization algorithms, entropy-regularized methods use convex duality, and occupancy-measure formulations turn finite discounted MDPs into linear programs. This chapter organizes reinforcement learning through the language of optimization, geometry, and statistical approximation.

28.1 Learning goals

After reading this chapter, students should be able to:

• explain how value-based, policy-based, and model-based reinforcement learning are optimization problems;
• derive Bellman optimality equations from one-step optimization;
• formulate finite discounted MDPs as linear programs;
• derive the dual occupancy-measure formulation;
• connect policy gradients with stochastic gradient ascent;
• explain mirror descent, KL geometry, and entropy regularization on the policy simplex;
• derive the natural policy gradient from a local KL-constrained optimization problem;
• compare Euclidean gradients, mirror descent, and natural gradients;
• recognize the role of nonconvexity, sampling noise, and distribution shift in modern RL optimization;
• use AI tools responsibly to check derivations, dimensions, and assumptions.

28.2 25.1 Why optimization is central to reinforcement learning

In supervised learning, a common mathematical template is

\[ \min_{\theta} \frac{1}{n}\sum_{i=1}^n \ell(f_\theta(x_i),y_i) + \lambda R(\theta). \]

Reinforcement learning is more subtle because the data distribution depends on the policy being optimized. A policy changes the future states that will be observed, so the objective is not merely a loss over a fixed dataset.

For a discounted MDP, a parameterized policy \(\pi_\theta\) has objective

\[ J(\theta) = E_{\pi_\theta}\left[\sum_{t=0}^{\infty}\gamma^t R_{t+1}\right]. \]

The expectation depends on \(\theta\) in two ways:

1. \(\pi_\theta\) directly changes the actions selected;
2. the selected actions change the state distribution over time.

Thus reinforcement learning is optimization with feedback: the optimizer changes the data-generating process.

Three common RL optimization viewpoints are:

Viewpoint	Decision variable	Main equation or objective
value optimization	\(V\) or \(Q\)	Bellman optimality equation
policy optimization	\(\theta\) or \(\pi\)	maximize \(J(\pi)\)
occupancy optimization	\(d(s,a)\)	linear program over discounted visitation measures

Unifying theme. Bellman methods optimize by fixed-point iteration, policy-gradient methods optimize by stochastic gradients, and occupancy-measure methods optimize by constrained convex programming.

28.3 25.2 Bellman optimality as pointwise optimization

For a finite discounted MDP, the optimal value function satisfies

\[ V^*(s) = \max_{a\in A(s)} \left\{ r(s,a)+\gamma\sum_{s'}P(s'\mid s,a)V^*(s') \right\}. \]

The right-hand side is a pointwise optimization problem at each state. Define the optimal Bellman operator

\[ (TV)(s) = \max_a \left\{ r(s,a)+\gamma\sum_{s'}P(s'\mid s,a)V(s') \right\}. \]

Then \(V^*\) is the fixed point of \(T\):

\[ V^*=TV. \]

More precisely, since \(T\) is a \(\gamma\)-contraction in the sup norm, the iteration

\[ V_{k+1}=TV_k \]

converges to \(V^*\) for any initial vector \(V_0\).

This is optimization without gradients. The maximization over actions is local, and the fixed-point iteration propagates this local optimization through time.

28.3.1 Interactive: optimization landscape for a one-state policy

A one-state, two-action discounted MDP already gives a policy objective as a function of the action probability \(p=\pi(a_1)\). The shape may be simple, but the example helps students see that a policy is an optimization variable.

28.4 25.3 Exact finite MDP example

Consider a two-state MDP with two actions. We can evaluate every deterministic stationary policy and compare the resulting values.

import numpy as np
from itertools import product

S = 2
A = 2
gamma = 0.9

# P[s, a, s_next]
P = np.array([
    [[0.80, 0.20], [0.25, 0.75]],
    [[0.40, 0.60], [0.10, 0.90]],
])

r = np.array([
    [1.0, 0.2],
    [0.0, 1.5],
])

mu0 = np.array([1.0, 0.0])

def evaluate_policy(policy):
    P_pi = np.zeros((S, S))
    r_pi = np.zeros(S)
    for s in range(S):
        a = policy[s]
        P_pi[s] = P[s, a]
        r_pi[s] = r[s, a]
    V = np.linalg.solve(np.eye(S) - gamma * P_pi, r_pi)
    return V, mu0 @ V

rows = []
for policy in product(range(A), repeat=S):
    V, objective = evaluate_policy(policy)
    rows.append((policy, V, objective))

for policy, V, objective in rows:
    print(f"policy={policy}, V={np.round(V, 4)}, J={objective:.4f}")

policy=(0, 0), V=[7.1875 5.625 ], J=7.1875
policy=(0, 1), V=[12.4324 13.7838], J=12.4324
policy=(1, 0), V=[0.8106 0.6344], J=0.8106
policy=(1, 1), V=[12.1445 13.6474], J=12.1445

This example uses brute-force optimization over deterministic policies. Brute force is not practical for large state spaces because the number of deterministic stationary policies is

\[ |A|^{|S|}. \]

The combinatorial growth motivates dynamic programming and gradient-based methods.

28.5 25.4 Policy gradient as stochastic optimization

Let \(\pi_\theta\) be a differentiable policy. The policy-gradient theorem states that, under standard regularity assumptions,

\[ \nabla_\theta J(\theta) = \frac{1}{1-\gamma} E_{S\sim d_{\pi_\theta}, A\sim \pi_\theta} \left[ \nabla_\theta \log \pi_\theta(A\mid S) Q^{\pi_\theta}(S,A) \right], \]

where \(d_{\pi_\theta}\) is the normalized discounted state-occupancy distribution.

A stochastic gradient ascent update is

\[ \theta_{k+1} = \theta_k+\n\alpha_k \widehat{\nabla_\theta J(\theta_k)}. \]

The gradient estimator is noisy because it is computed from sampled trajectories. This places policy-gradient methods within stochastic approximation.

For a trajectory \(\tau=(S_0,A_0,R_1,S_1,A_1,R_2,\ldots)\), the REINFORCE estimator has the form

\[ \widehat{g} = \sum_{t\ge 0} G_t\nabla_\theta \log \pi_\theta(A_t\mid S_t). \]

A baseline \(b(S_t)\) can reduce variance:

\[ \widehat{g} = \sum_{t\ge 0} (G_t-b(S_t))\nabla_\theta \log \pi_\theta(A_t\mid S_t). \]

The baseline does not change the expected gradient if it does not depend on \(A_t\).

28.5.1 Interactive: gradient ascent paths

This figure compares small and large step sizes for gradient ascent on a simple policy objective. The purpose is not to claim that RL objectives are always one-dimensional, but to visualize the optimization tradeoff between slow learning and unstable steps.

28.6 25.5 Geometry of the policy simplex

For a finite action space, a randomized policy at a state is a probability vector

\[ \pi(\cdot\mid s)\in \Delta_A, \]

where

\[ \Delta_A = \left\{p\in R^A: p_a\ge 0,\ \sum_a p_a=1\right\}. \]

A naive Euclidean update

\[ \pi_{new}(a\mid s) = \pi(a\mid s)+\alpha g(a) \]

may leave the simplex. One may project back, but projection can behave poorly near the boundary.

A more natural update uses KL divergence. Given a gradient-like advantage vector \(A(s,a)\), define

\[ \pi_{new}(\cdot\mid s) = \arg\max_{p\in \Delta_A} \left\{ \sum_a p(a)A(s,a) - \frac{1}{\eta}D_{KL}(p\|\pi(\cdot\mid s)) \right\}. \]

The solution is the exponentiated-gradient or mirror-descent update

\[ \pi_{new}(a\mid s) = \frac{\pi(a\mid s)\exp(\eta A(s,a))} {\sum_b \pi(b\mid s)\exp(\eta A(s,b))}. \]

This update stays inside the simplex and moves multiplicatively rather than additively.

28.6.1 Interactive: mirror descent on the policy simplex

The update moves a policy distribution toward actions with positive advantage while keeping probabilities nonnegative and summing to one.

import numpy as np

pi = np.array([0.30, 0.40, 0.30])
advantage = np.array([0.8, -0.2, 0.4])
eta = 1.5

weights = pi * np.exp(eta * advantage)
pi_new = weights / weights.sum()

print("old policy:", np.round(pi, 4))
print("advantage:", advantage)
print("new policy:", np.round(pi_new, 4))
print("sum:", pi_new.sum())

old policy: [0.3 0.4 0.3]
advantage: [ 0.8 -0.2  0.4]
new policy: [0.5416 0.1611 0.2972]
sum: 0.9999999999999999

28.7 25.6 Bregman divergences and mirror descent

A Bregman divergence generated by a differentiable strictly convex function \(\psi\) is

\[ D_\psi(x,y) = \psi(x)-\psi(y)-\langle \nabla\psi(y),x-y\rangle. \]

Mirror descent solves local optimization problems of the form

\[ x_{k+1} = \arg\min_{x\in C} \left\{ \langle g_k,x\rangle + \frac{1}{\eta_k}D_\psi(x,x_k) \right\}. \]

For policy maximization, one changes the sign:

\[ \pi_{k+1} = \arg\max_{\pi\in \Pi} \left\{ \langle \widehat{A}_k,\pi\rangle - \frac{1}{\eta_k}D_\psi(\pi,\pi_k) \right\}. \]

When \(\psi(p)=\sum_a p(a)\log p(a)\), the induced Bregman divergence is KL divergence. Thus entropy geometry gives multiplicative policy updates.

Optimization lesson. Euclidean geometry is natural for unconstrained vectors. KL geometry is natural for probability distributions.

28.8 25.7 Natural policy gradients

The ordinary gradient depends on the parameterization. A small Euclidean parameter change may create a large policy change, or a large parameter change may create a small policy change.

Natural gradient methods measure policy change by KL divergence. Consider the local problem

\[ \max_{\Delta\theta} \quad \nabla J(\theta)^T\Delta\theta \qquad \text{subject to} \qquad \frac{1}{2}\Delta\theta^T F(\theta)\Delta\theta\le \epsilon, \]

where \(F(\theta)\) is the Fisher information matrix. The solution direction is

\[ \Delta\theta \propto F(\theta)^{-1}\nabla J(\theta). \]

Thus the natural gradient update is

\[ \theta_{k+1} = \theta_k+ \alpha_k F(\theta_k)^{-1}\nabla J(\theta_k). \]

For a policy model,

\[ F(\theta) = E_{S\sim d_\pi,A\sim \pi} \left[ \nabla_\theta\log\pi_\theta(A\mid S) \nabla_\theta\log\pi_\theta(A\mid S)^T \right]. \]

This is the Riemannian metric induced by the statistical model. Natural gradient methods are closely related to trust-region policy optimization.

28.8.1 Interactive: Euclidean gradient versus natural gradient

The same objective can have curved level sets. The natural gradient rescales directions according to local geometry.

import numpy as np

# Two-action softmax policy with scalar logit theta:
# pi(action 1) = sigmoid(theta), pi(action 0) = 1 - sigmoid(theta).

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

theta = 0.3
p = sigmoid(theta)

# Suppose the policy gradient with respect to theta is estimated as follows.
grad = 0.18

# Fisher information for Bernoulli logistic parameterization.
F = p * (1 - p)

ordinary_step = grad
natural_step = grad / F

print("p(action 1):", round(p, 4))
print("Fisher information:", round(F, 4))
print("ordinary gradient direction:", round(ordinary_step, 4))
print("natural gradient direction:", round(natural_step, 4))

p(action 1): 0.5744
Fisher information: 0.2445
ordinary gradient direction: 0.18
natural gradient direction: 0.7363

28.9 25.8 Linear programming formulation of discounted MDPs

For a finite discounted MDP, the optimal value function is the solution of a linear program. The primal value-function formulation is

\[ \min_V \sum_s \mu(s)V(s) \]

subject to the Bellman inequalities

\[ V(s) \ge r(s,a)+\gamma\sum_{s'}P(s'\mid s,a)V(s') \quad \text{for all }s,a. \]

Here \(\mu\) is any probability distribution with positive mass on relevant states. At optimum, \(V=V^*\).

Why does this work? The inequalities say

\[ V\ge TV. \]

The optimal value function \(V^*\) is the smallest vector satisfying these inequalities. Minimizing a positive weighted sum selects this smallest feasible upper bound.

28.9.1 Interactive: Bellman inequalities as feasible half-spaces

For a tiny MDP, the Bellman inequalities form a feasible region in value space. The optimal value is the smallest feasible vector in the appropriate partial order.

28.10 25.9 Occupancy measures and the dual linear program

The dual viewpoint optimizes over discounted state-action visitation measures. For a policy \(\pi\), define the discounted occupancy measure

\[ d_\pi(s,a) = (1-\gamma)\sum_{t=0}^{\infty}\gamma^t P_\pi(S_t=s,A_t=a). \]

This is a probability distribution over state-action pairs when the initial state distribution is fixed.

The dual linear program is

\[ \max_d \sum_{s,a}d(s,a)r(s,a) \]

subject to

\[ \sum_a d(s,a) = (1-\gamma)\mu(s) + \gamma\sum_{s',a'}d(s',a')P(s\mid s',a') \quad \text{for all }s, \]

and

\[ d(s,a) \ge 0. \]

The constraints are flow-conservation equations. Discounted probability mass enters from the initial distribution and flows through transitions.

Given a feasible occupancy measure, a policy can be recovered by

\[ \pi(a\mid s) = \frac{d(s,a)}{\sum_b d(s,b)} \]

when \(\sum_b d(s,b)>0\).

28.10.1 Interactive: discounted occupancy as flow

This example visualizes how discounted visitation mass is distributed across state-action pairs.

import numpy as np

S = 2
A = 2
gamma = 0.9
mu = np.array([1.0, 0.0])

P = np.array([
    [[0.80, 0.20], [0.25, 0.75]],
    [[0.40, 0.60], [0.10, 0.90]],
])

pi = np.array([
    [0.70, 0.30],
    [0.20, 0.80],
])

# Build policy-induced transition matrix.
P_pi = np.zeros((S, S))
for s in range(S):
    for a in range(A):
        P_pi[s] += pi[s, a] * P[s, a]

# Normalized discounted state distribution d_state.
d_state = (1 - gamma) * mu @ np.linalg.inv(np.eye(S) - gamma * P_pi)

d_sa = np.zeros((S, A))
for s in range(S):
    d_sa[s] = d_state[s] * pi[s]

print("discounted state distribution:", np.round(d_state, 4))
print("state-action occupancy:")
print(np.round(d_sa, 4))
print("total mass:", d_sa.sum())

discounted state distribution: [0.4262 0.5738]
state-action occupancy:
[[0.2983 0.1279]
 [0.1148 0.459 ]]
total mass: 1.0000000000000004

28.11 25.10 Entropy regularization and convex duality

Entropy regularization replaces a hard greedy policy by a smooth optimization problem. At a fixed state, consider

\[ \max_{p\in\Delta_A} \left\{ \sum_a p(a)q(a) + \alpha H(p) \right\}, \]

where

\[ H(p)=-\sum_a p(a)\log p(a). \]

The solution is

\[ p^*(a) = \frac{\exp(q(a)/\alpha)}{\sum_b\exp(q(b)/\alpha)}. \]

The optimal value is

\[ \alpha\log\sum_a\exp(q(a)/\alpha). \]

Thus entropy turns the maximum into the smooth log-sum-exp function. As \(\alpha\downarrow 0\), the log-sum-exp approaches the maximum:

\[ \alpha\log\sum_a\exp(q(a)/\alpha) \to \max_a q(a). \]

28.11.1 Interactive: entropy regularization smooths greedy choice

As the temperature increases, the policy becomes more random. As it decreases, the policy approaches a greedy action.

28.12 25.11 Constrained reinforcement learning

Many applied problems require constraints. Examples include safety, budget, fairness, risk, or resource constraints. A constrained discounted MDP can be written as

\[ \max_\pi J_r(\pi) \quad \text{subject to} \quad J_c(\pi)\le C, \]

where \(J_r\) is expected reward and \(J_c\) is expected cost.

The Lagrangian is

\[ L(\pi,\lambda) = J_r(\pi)-\lambda(J_c(\pi)-C), \qquad \lambda\ge 0. \]

A primal-dual update has the schematic form

\[ \pi_{k+1} \approx \arg\max_\pi L(\pi,\lambda_k), \]

\[ \lambda_{k+1} = \left[\lambda_k+\beta_k(J_c(\pi_k)-C)\right]_+. \]

The multiplier increases when the constraint is violated and decreases when the constraint is comfortably satisfied.

28.13 25.12 Online learning and regret

In online learning, an agent repeatedly chooses actions and receives losses. The regret after \(T\) rounds is

\[ R_T = \sum_{t=1}^T \ell_t(a_t) - \min_a\sum_{t=1}^T \ell_t(a). \]

In bandits, the loss or reward is observed only for the selected action. In full-information online learning, the entire loss vector is observed.

Exponentiated-gradient methods update action probabilities by

\[ p_{t+1}(a) = \frac{p_t(a)\exp(-\eta \ell_t(a))} {\sum_b p_t(b)\exp(-\eta \ell_t(b))}. \]

This is mirror descent with negative entropy. It is mathematically close to entropy-regularized policy improvement.

The online-learning viewpoint helps explain exploration bonuses, optimism, posterior sampling, and adversarial robustness.

28.14 25.13 Nonconvexity and modern RL optimization

Tabular discounted MDPs have strong structure. The value-function LP is convex, and dynamic programming has global convergence. However, modern deep RL is usually nonconvex because policies and value functions are represented by neural networks.

Several difficulties arise:

Difficulty	Optimization issue
nonlinear function approximation	nonconvex loss surfaces
bootstrapping	moving targets
off-policy data	distribution mismatch
max operator	overestimation bias
exploration	objective is not fully observed
long horizon	high-variance gradients

A useful habit is to distinguish three kinds of error:

\[ \text{total error} = \text{optimization error} + \text{statistical error} + \text{approximation error}. \]

Optimization error comes from not solving the chosen empirical problem exactly. Statistical error comes from finite data. Approximation error comes from the fact that the chosen function class may not contain the true value function or optimal policy.

28.15 25.14 Python example: projected policy gradient on a bandit

The following example compares Euclidean projected gradient ascent with exponentiated-gradient ascent for a three-action bandit.

import numpy as np

q = np.array([1.0, 0.2, 0.7])
eta = 0.4
T = 20

def project_simplex(v):
    # Euclidean projection onto the probability simplex.
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u)
    rho = np.nonzero(u * np.arange(1, len(v) + 1) > (cssv - 1))[0][-1]
    theta = (cssv[rho] - 1) / (rho + 1)
    return np.maximum(v - theta, 0)

p_euclid = np.ones(3) / 3
p_mirror = np.ones(3) / 3

history = []
for k in range(T):
    p_euclid = project_simplex(p_euclid + eta * q)
    weights = p_mirror * np.exp(eta * q)
    p_mirror = weights / weights.sum()
    history.append((k + 1, p_euclid.copy(), p_mirror.copy()))

print("final projected-gradient policy:", np.round(p_euclid, 4))
print("final mirror-descent policy:", np.round(p_mirror, 4))
print("best action:", int(np.argmax(q)))

final projected-gradient policy: [1. 0. 0.]
final mirror-descent policy: [0.9154 0.0015 0.083 ]
best action: 0

The Euclidean projection quickly hits the boundary. Mirror descent approaches the best action multiplicatively and often gives smoother probability paths.

28.16 25.15 AI-assisted learning components

28.16.1 AI prompt: identify the optimization variable

Give an AI tool a paragraph describing an RL method and ask:

What is the optimization variable? Is the method optimizing over values, policies, model parameters, or occupancy measures? What constraints are present?

Then verify the answer yourself. Many RL descriptions hide the optimization variable behind algorithmic language.

28.16.2 AI prompt: check a Lagrangian derivation

Ask:

For the constrained problem \(\max_\pi J_r(\pi)\) subject to \(J_c(\pi)\le C\), derive the Lagrangian and the dual update. Check the sign of the multiplier update.

The sign is important. If the cost violates the constraint, the multiplier should increase.

28.16.3 AI prompt: compare Euclidean and KL geometry

Ask:

Explain why KL divergence is often more natural than Euclidean distance for policy optimization on the probability simplex. Give a two-action example.

Then check whether the response distinguishes parameter distance from distribution distance.

28.16.4 AI prompt: debug mirror descent code

Give the following pseudocode to an AI tool:

weights = exp(eta * advantage)
new_policy = weights / sum(weights)

Ask what is missing. The answer should note that exponentiated-gradient mirror descent around an old policy uses

\[ \pi_{new}(a)\propto \pi_{old}(a)\exp(\eta A(a)). \]

Without the old policy factor, the update is just softmax over advantages.

28.17 25.16 Summary

This chapter connected reinforcement learning with optimization.

• Bellman optimality is pointwise optimization plus fixed-point computation.
• Policy gradient is stochastic gradient ascent on expected return.
• Mirror descent gives natural updates on the policy simplex.
• KL divergence and entropy are central because policies are probability distributions.
• Natural gradients correct Euclidean gradients using Fisher information.
• Finite discounted MDPs admit primal and dual linear programming formulations.
• Occupancy measures provide a convex view of state-action visitation.
• Constrained RL can be studied using Lagrangians and primal-dual updates.
• Deep RL adds nonconvexity, sampling noise, distribution shift, and approximation error.

28.18 Exercises

28.18.1 Conceptual exercises

1. Explain why the data distribution in RL depends on the policy being optimized.
2. Compare value iteration, policy gradient, and occupancy-measure optimization.
3. Why is KL divergence natural for policy optimization?
4. Explain the difference between optimization error, statistical error, and approximation error.
5. Why can a Euclidean policy update leave the probability simplex?

28.18.2 Mathematical exercises

1. Prove that the Bellman inequalities \(V\ge TV\) imply \(V\ge V^*\).
2. Derive the exponentiated-gradient update \[ \pi_{new}(a) = \frac{\pi(a)\exp(\eta A(a))} {\sum_b\pi(b)\exp(\eta A(b))}. \]
3. Derive the natural-gradient direction from the KL-constrained local optimization problem.
4. Show that the discounted occupancy measure satisfies the flow-conservation equations.
5. Derive the entropy-regularized maximizer \[ p^*(a)=\frac{\exp(q(a)/\alpha)}{\sum_b\exp(q(b)/\alpha)}. \]

28.18.3 Computational exercises

1. Implement projected gradient ascent and mirror descent for a five-action bandit.
2. For a small MDP, compute \(V^*\) using value iteration and compare it with the solution of the Bellman inequality LP.
3. Estimate an occupancy measure from simulated trajectories and compare it with the matrix formula.
4. Implement natural policy gradient for a two-action logistic policy.
5. Simulate a constrained bandit and implement a primal-dual multiplier update.

28.18.4 AI-assisted exercises

1. Ask an AI tool to explain why entropy regularization produces a softmax policy. Then write your own derivation.
2. Give an AI tool a policy-gradient update and ask it to identify whether it is Euclidean, natural-gradient, or mirror-descent style.
3. Ask an AI tool to generate a small constrained MDP. Then verify whether the proposed constraints are mathematically well-defined.
4. Ask an AI tool to compare the primal and dual LP formulations for discounted MDPs. Check whether it correctly defines occupancy measures.
5. Ask an AI tool to debug a natural-gradient implementation where the Fisher matrix is singular. Propose a regularization fix.

28.19 Instructor notes

This chapter is a bridge between earlier algorithmic chapters and later statistical-learning/offline-RL chapters. For MA Applied Math students, emphasize fixed points, convex duality, Bregman divergence, and constrained optimization. For MS Statistics students, emphasize sampling noise, stochastic gradients, occupancy distributions, and finite-sample estimation. The chapter can be taught after policy-gradient and entropy-regularized RL, or earlier as a unifying perspective before advanced topics.