Infinite-Horizon Markov Decision Processes

Finite-horizon stochastic dynamic programming obtains its boundary condition from the terminal reward. What value equation remains when the process has no terminal date? Discounting makes the infinite reward stream bounded and turns the Bellman equations into fixed-point problems.

The undiscounted expected total reward of policy $\boldsymbol{\pi} \in \Pi^{\mathrm{HR}}$ is

One drawback of this model is that we could easily encounter values that are $+\infty$ or $-\infty$, even in a setting as simple as a single-state MDP which loops back into itself and where the accrued reward is nonzero.

Therefore, it is often more convenient to work with an alternative formulation which guarantees the existence of a limit: the expected total discounted reward of policy $\boldsymbol{\pi} \in \Pi^{\mathrm{HR}}$ is defined to be:

for $0 \leq \gamma < 1$ and when $\max_{s \in \mathcal{S}} \max_{a \in \mathcal{A}_s}|r(s, a)| = R_{\max} < \infty$, in which case, $|v_\gamma^{\boldsymbol{\pi}}(s)| \leq (1-\gamma)^{-1} R_{\max}$.

Finally, another possibility for the infinite-horizon setting is the so-called average reward or gain of policy $\boldsymbol{\pi} \in \Pi^{\mathrm{HR}}$ defined as:

We won’t be working with this formulation in this course due to its inherent practical and theoretical complexities.

Extending the previous notion of optimality from finite-horizon models, a policy $\boldsymbol{\pi}^*$ is said to be discount optimal for a given $\gamma$ if:

Furthermore, the value of a discounted MDP $v_\gamma^*(s)$, is defined by:

More often, we refer to $v_\gamma$ by simply calling it the optimal value function.

As for the finite-horizon setting, the infinite horizon discounted model does not require history-dependent policies, since for any $\boldsymbol{\pi} \in \Pi^{HR}$ there exists a $\boldsymbol{\pi}^{\prime} \in \Pi^{MR}$ with identical total discounted reward:

Random Horizon Interpretation of Discounting¶

Can geometric termination give the discount factor a probabilistic meaning rather than treating it only as an algebraic device? The use of discounting can be motivated both from a modeling perspective and as a means to ensure that the total reward remains bounded. From the modeling perspective, we can view discounting as a way to weight more or less importance on the immediate rewards vs. the long-term consequences. There is also another interpretation which stems from that of a finite horizon model but with an uncertain end time. More precisely:

Let $v_\nu^{\boldsymbol{\pi}}(s)$ denote the expected total reward obtained by using policy $\boldsymbol{\pi}$ when the horizon length $\nu$ is random. We define it by:

Vector Representation in Markov Decision Processes¶

How do transition kernels, rewards, and policies become matrices and vectors that expose the Bellman equations as operator equations?

Let V be the set of bounded real-valued functions on a discrete state space S. This means any function $f \in V$ satisfies the condition:

where notation $\|f\|$ represents the sup-norm (or $\ell_\infty$-norm) of the function $f$.

When working with discrete state spaces, we can interpret elements of V as vectors and linear operators on V as matrices, allowing us to leverage tools from linear algebra. The sup-norm ($\ell_\infty$ norm) of matrix $\mathbf{H}$ is defined as:

where $\mathbf{H}_{s,j}$ represents the $(s, j)$-th component of the matrix $\mathbf{H}$.

For a Markovian decision rule $\pi \in \Pi^{MD}$, we define:

For a randomized decision rule $\pi \in \Pi^{MR}$, these definitions extend to:

In both cases, $\mathbf{r}_\pi$ denotes a reward vector in $\mathbb{R}^{|S|}$, with each component $\mathbf{r}_\pi(s)$ representing the reward associated with state $s$. Similarly, $\mathbf{P}_\pi$ is a transition probability matrix in $\mathbb{R}^{|S| \times |S|}$, capturing the transition probabilities under decision rule $\pi$.

For a nonstationary Markovian policy $\boldsymbol{\pi} = (\pi_1, \pi_2, \ldots) \in \Pi^{MR}$, the expected total discounted reward is given by:

Using vector notation, this can be expressed as:

This formulation leads to a recursive relationship:

where $\boldsymbol{\pi}^{\prime} = (\pi_2, \pi_3, \ldots)$.

For a stationary policy $\boldsymbol{\pi} = \mathrm{const}(\pi)$ with constant decision rule $\pi$, the total expected reward simplifies to:

This last expression is called a Neumann series expansion, and it’s guaranteed to exists under the assumptions of bounded reward and discount factor strictly less than one.

Note that for any induced matrix norm $\|\cdot\|$ (i.e., a norm satisfying $\|\mathbf{H}\mathbf{v}\| \leq \|\mathbf{H}\| \cdot \|\mathbf{v}\|$ for all vectors $\mathbf{v}$) and any matrix $\mathbf{H}$, the spectral radius is bounded by:

This inequality provides a practical way to verify the convergence condition $\rho(\mathbf{H}) < 1$ by checking the simpler condition $\|\mathbf{H}\| < 1$ rather than trying to compute the eigenvalues directly.

We can now verify that $(\mathbf{I} - \gamma \mathbf{P}_\pi)$ is invertible and the Neumann series converges.

1. Norm of the transition matrix: Since $\mathbf{P}_\pi$ is a stochastic matrix (each row sums to 1 and all entries are non-negative), its $\ell_\infty$-norm is:

   $$\|\mathbf{P}_\pi\| = \max_{s \in S} \sum_{j \in S} [\mathbf{P}_\pi]_{s,j} = \max_{s \in S} 1 = 1.$$

   (27)

2. Norm of the scaled matrix: Using the homogeneity property of norms, we have:

   $$\|\gamma \mathbf{P}_\pi\| = |\gamma| \cdot \|\mathbf{P}_\pi\| = |\gamma| \cdot 1 = |\gamma|.$$

   (28)

3. Bounding the spectral radius: Since the spectral radius is bounded by the matrix norm:

   $$\rho(\gamma \mathbf{P}_\pi) \leq \|\gamma \mathbf{P}_\pi\| = |\gamma|.$$

   (29)

4. Verifying convergence: Since $0 \leq \gamma < 1$ by assumption, we have:

   $$\rho(\gamma \mathbf{P}_\pi) \leq |\gamma| < 1.$$

   (30)

   This strict inequality guarantees that $(\mathbf{I} - \gamma \mathbf{P}_\pi)$ is invertible and the Neumann series converges.

Therefore, the Neumann series expansion converges and yields:

Consequently, for a stationary policy, $\mathbf{v}_\gamma^{\pi}$ can be determined as the solution to the linear equation:

which can be rearranged to:

We can also characterize $\mathbf{v}_\gamma^{\pi}$ as the solution to an operator equation. More specifically, define the transformation $\mathrm{L}_\pi$ by

for any $\mathbf{v} \in V$. Intuitively, $\mathrm{L}_\pi$ takes a value function $\mathbf{v}$ as input and returns a new value function that combines immediate rewards ($\mathbf{r}_\pi$) with discounted future values ($\gamma \mathbf{P}_\pi\mathbf{v}$).

Therefore, we view $\mathrm{L}_\pi$ as an operator mapping elements of $V$ to $V$: i.e., $\mathrm{L}_\pi: V \rightarrow V$. The fact that the value function of a policy is the solution to a fixed-point equation can then be expressed with the statement:

This is a fixed-point equation: the value function $\mathbf{v}_\gamma^{\pi}$ is a fixed point of the operator $\mathrm{L}_\pi$.

Solving Operator Equations¶

Which iterative and Newton-like methods solve a fixed point when direct matrix inversion is unavailable or inappropriate?

The operator equation we encountered in MDPs, $\mathbf{v}_\gamma^{\pi} = \mathrm{L}_\pi \mathbf{v}_\gamma^{\pi}$, is a specific instance of a more general class of problems known as operator equations. These equations appear in various fields of mathematics and applied sciences, ranging from differential equations to functional analysis.

Operator equations can take several forms, each with its own characteristics and solution methods:

1. Fixed Point Form: $x = \mathrm{T}(x)$, where $\mathrm{T}: X \rightarrow X$. Common in fixed-point problems, such as our MDP equation, we seek a fixed point $x^*$ such that $x^* = \mathrm{T}(x^*)$.

2. General Operator Equation: $\mathrm{T}(x) = y$, where $\mathrm{T}: X \rightarrow Y$. Here, $X$ and $Y$ can be different spaces. We seek an $x \in X$ that satisfies the equation for a given $y \in Y$.

3. Nonlinear Equation: $\mathrm{T}(x) = 0$, where $\mathrm{T}: X \rightarrow Y$. A special case of the general operator equation where we seek roots or zeros of the operator.

4. Variational Inequality: Find $x^* \in K$ such that $\langle \mathrm{T}(x^*), x - x^* \rangle \geq 0$ for all $x \in K$. Here, $K$ is a closed convex subset of $X$, and $\mathrm{T}: K \rightarrow X^*$ (the dual space of $X$). These problems often arise in optimization, game theory, and partial differential equations.

Successive Approximation Method¶

For equations in fixed point form, a common numerical solution method is successive approximation, also known as fixed-point iteration:

The convergence of successive approximation depends on the properties of the operator $\mathrm{T}$. In the simplest and most common setting, we assume $\mathrm{T}$ is a contraction mapping. The Banach Fixed-Point Theorem then guarantees that $\mathrm{T}$ has a unique fixed point, and the successive approximation method will converge to this fixed point from any starting point. Specifically, $\mathrm{T}$ is a contraction if there exists a constant $q \in [0,1)$ such that for all $x,y \in X$:

where $d$ is the metric on $X$. In this case, the rate of convergence is linear, with error bound:

However, the contraction mapping condition is not the only one that can lead to convergence. For instance, if $\mathrm{T}$ is nonexpansive (i.e., Lipschitz continuous with Lipschitz constant 1) and $X$ is a Banach space with certain geometrical properties (e.g., uniformly convex), then under additional conditions (e.g., $\mathrm{T}$ has at least one fixed point), the successive approximation method can still converge, albeit potentially more slowly than in the contraction case.

In practice, when dealing with specific problems like MDPs or differential equations, the properties of the operator often naturally align with one of these convergence conditions. For example, in discounted MDPs, the Bellman operator is a contraction in the supremum norm, which guarantees the convergence of value iteration.

Newton-Kantorovich Method¶

The Newton-Kantorovich method is a generalization of Newton’s method from finite dimensional vector spaces to infinite dimensional function spaces: rather than iterating in the space of vectors, we are iterating in the space of functions.

Newton’s method is often written as the familiar update:

which makes it look as though the essence of the method is “take a derivative and invert it.” But the real workhorse behind Newton’s method (both in finite and infinite dimensions) is linearization.

At each step, the idea is to replace the nonlinear operator $F:X \to Y$ by a local surrogate model of the form

where $L$ is a linear map capturing how small perturbations in the input propagate to changes in the output. This is a Taylor-like expansion in Banach spaces: the role of the derivative is precisely to provide the correct notion of such a linear operator.

To find a root of $F$, we impose the condition that the surrogate vanishes at the next iterate:

Solving this linear equation gives the increment $h$. In finite dimensions, $L$ is the Jacobian matrix; in Banach spaces, it must be the Fréchet derivative.

But what exactly is a Fréchet derivative in infinite dimensions? To understand this, we need to generalize the concept of derivative from finite-dimensional calculus. In infinite-dimensional spaces, there are several notions of differentiability, each with different strengths and requirements:

1. Gâteaux (Directional) Derivative

We say that the Gâteaux derivative of $F$ at $x$ in a specific direction $h$ is:

This quantity measures how the function $F$ changes along the ray $x + th$. While this limit may exist for each direction $h$ separately, it doesn’t guarantee that the derivative is linear in $h$. This is a key limitation: the Gâteaux derivative can exist in all directions but still fail to provide a good linear approximation.

2. Hadamard Directional Derivative

Rather than considering a single direction of perturbation, we now consider a bundle of perturbations around $h$. We ask how the function changes as we approach the target direction from nearby directions. We say that $F$ has a Hadamard directional derivative if:

This is a stronger condition than Gâteaux differentiability because it requires the limit to be uniform over nearby directions. However, it still doesn’t guarantee linearity in $h$.

3. Fréchet Derivative

The strongest and most natural notion: $F$ is Fréchet differentiable at $x$ if there exists a bounded linear operator $L$ such that:

This definition directly addresses the inadequacy of the previous notions. Unlike Gâteaux and Hadamard derivatives, the Fréchet derivative explicitly requires the existence of a linear operator $L$ that provides a good approximation. Key properties:

• $L$ must be linear in $h$ (unlike the directional derivatives above)

• The approximation error is $o(\|h\|)$, uniform in all directions

• This is the “true” derivative: it generalizes the Jacobian matrix to infinite dimensions

• Notation: $L = F'(x)$ or $DF(x)$

Relationship:

In the context of the Newton-Kantorovich method, we work with an operator $F: X \to Y$ where both $X$ and $Y$ are Banach spaces. The Fréchet derivative $F'(x)$ is the best linear approximation of $F$ near $x$, and it’s exactly this linear operator $L$ that we use in our linearization $F(x+h) \approx F(x) + F'(x)h$.

Now apart from those mathematical technicalities, Newton-Kantorovich has in essence the same structure as that of the original Newton’s method. That is, it applies the following sequence of steps:

1. Linearize the Operator: Given an approximation $x_n$, we consider the Fréchet derivative of $F$, denoted by $F'(x_n)$. This derivative is a linear operator that provides a local approximation of $F$ near $x_n$.

2. Set Up the Newton Step: The method then solves the linearized equation for a correction $h_n$:

   $$F'(x_n) h_n = -F(x_n).$$

   (47)

   This equation represents a linear system where $h_n$ is chosen so that the linearized operator $F(x_n) + F'(x_n)h_n$ equals zero.

3. Update the Solution: The new approximation $x_{n+1}$ is then given by:

   $$x_{n+1} = x_n + h_n.$$

   (48)

   This correction step refines $x_n$, bringing it closer to the true solution.

4. Repeat Until Convergence: We repeat the linearization and update steps until the solution $x_n$ converges to the desired tolerance, which can be verified by checking that $\|F(x_n)\|$ is sufficiently small, or by monitoring the norm $\|x_{n+1} - x_n\|$.

The convergence of Newton-Kantorovich does not hinge on $F$ being a contraction over the entire domain (as it could be the case for successive approximation). The convergence properties of the Newton-Kantorovich method are as follows:

1. Local Convergence: Under mild conditions (e.g., $F$ is Fréchet differentiable and $F'(x)$ is invertible near the solution), the method converges locally. This means that if the initial guess is sufficiently close to the true solution, the method will converge.

2. Global Convergence: Global convergence is not guaranteed in general. However, under stronger conditions (e.g., $F$ is analytic and satisfies certain bounds), the method can converge globally.

3. Rate of Convergence: When the method converges, it typically exhibits quadratic convergence. This means that the error at each step is proportional to the square of the error at the previous step:

   $$\|x_{n+1} - x^*\| \leq C\|x_n - x^*\|^2$$

   (49)

   where $x^*$ is the true solution and $C$ is some constant. This quadratic convergence is significantly faster than the linear convergence typically seen in methods like successive approximation.

Optimality Equations for Infinite-Horizon MDPs¶

How does actionwise maximization turn fixed-policy evaluation into the Bellman optimality equation?

Recall that in the finite-horizon setting, the optimality equations are:

where $v_n(s)$ is the value function at time step $n$ for state $s$, $A_s$ is the set of actions available in state $s$, $r(s, a)$ is the reward function, $\gamma$ is the discount factor, and $p(j | s, a)$ is the transition probability from state $s$ to state $j$ given action $a$.

Intuitively, we would expect that by taking the limit of $n$ to infinity, we might get the nonlinear equations:

which are called the optimality equations or Bellman equations for infinite-horizon MDPs.

We can adopt an operator-theoretic perspective by defining operators on the space $V$ of bounded real-valued functions on the state space $S$. For a deterministic Markov rule $\pi \in \Pi^{MD}$, define the policy-evaluation operator:

The Bellman optimality operator is then:

where $\Pi^{MD}$ is the set of Markov deterministic decision rules, $\mathbf{r}_\pi$ is the reward vector under decision rule $\pi$, and $\mathbf{P}_\pi$ is the transition probability matrix under decision rule $\pi$.

Note that while we write $\max_{\pi \in \Pi^{MD}}$, we do not implement the above operator by enumerating all decision rules. Rather, the fact that we compare policies based on their value functions in a componentwise fashion means that maximizing over the space of Markovian deterministic rules reduces to the following update in component form:

For convenience, we define the greedy selector $\mathrm{Greedy}(v) \in \Pi^{MD}$ that extracts an optimal decision rule from a value function:

In Puterman’s terminology, such a greedy selector is called $v$-improving (or conserving when it achieves the maximum). This operator will be useful for expressing algorithms succinctly:

• Value iteration: $v_{k+1} = \Bellman v_k$, then extract $\pi = \mathrm{Greedy}(v^*)$

• Policy iteration: $\pi_{k+1} = \mathrm{Greedy}(v^{\pi_k})$ with $v^{\pi_k}$ solving $v = \mathrm{L}_{\pi_k}v$

The equivalence between these two forms can be shown mathematically, as demonstrated in the following proposition and proof.

Algorithms for Solving the Optimality Equations¶

What computational trade-off separates value iteration, policy evaluation, and policy improvement?

The optimality equations are operator equations. Therefore, we can apply general numerical methods to solve them. Applying the successive approximation method to the Bellman optimality equation yields a method known as “value iteration” in dynamic programming. A direct application of the blueprint for successive approximation yields the following algorithm:

The termination criterion in this algorithm is based on a specific bound that provides guarantees on the quality of the solution. This is in contrast to supervised learning, where we often use arbitrary termination criteria based on computational budget or early stopping when the learning curve flattens. This is because establishing implementable generalization bounds in supervised learning is challenging.

However, in the dynamic programming context, we can derive various bounds that can be implemented in practice. These bounds help us terminate our procedure with a guarantee on the precision of our value function and, correspondingly, on the optimality of the resulting policy.

Bellman contraction laboratory¶

Use the controls below to change the discount factor, transition persistence, reward asymmetry, and starting value. Before moving $\gamma$, predict how it will change the slope of the error envelope. The middle panel compares the observed error and Bellman residual with the contraction bound; the text below reports the final greedy policy.

Exact Scheduling MDP for Inference Serving¶

How can a large request-level system be reduced to a finite MDP, and which predictive distinctions disappear in that reduction?

The inference examples have so far treated the scheduling rule as fixed and the GPU clock as the action. A different decision interface fixes the clock and asks which phase should receive the next unit of service. Prefill admits new requests into decode and consumes cache; decode advances requests already producing output tokens. Serving either phase delays the other.

An exact request-level Markov state would contain every prompt length, generated token count, cache allocation, and waiting time. For computation, these quantities are aggregated into the finite state

where $p$ counts waiting prefill jobs, $d$ counts active decode jobs, and $a$ is the oldest prefill-age bin. The actions are

An action is masked when its phase is empty. Prefill is also masked at $d=6$, which represents the cache limit in this abstraction. A Bernoulli arrival probability for each 0.1-second decision period is calibrated from the load-normalized version of the same five-minute Azure trace used in the modeling chapter.

The phase rates and powers come from one measured NVIDIA L4 run of Qwen/Qwen2.5-7B-Instruct served by vLLM 0.28.0. The reduced MDP uses the 1,125 MHz requested clock level; its batch-balanced median realized graphics clock was 939.375 MHz. The measured prefill rate and the trace’s mean prompt length determine the Bernoulli probability that a prefill action completes one aggregate prompt. A successful completion moves that job into decode. The measured decode rate and mean output length similarly determine one aggregate expected completion budget per decode action. That budget is shared symmetrically across the active jobs, so adding jobs does not multiply the model’s expected completion capacity. Queue counts are capped at six, and arrivals beyond that cap are recorded as drops. These choices define a complete transition matrix $P_{ss'}^u$ on 245 states.

The one-step cost assigns separate penalties to congestion, old prompt work, decode stalls, dropped requests, and energy:

For prefill, decode, and idle, respectively, the measured profile gives $E(u)=(6.427,6.274,2.339)$ joules per decision period. Each value is the measured phase-power summary at the requested 1,125 MHz level multiplied by 0.1 seconds, rather than a direct request-level energy measurement. Only the ratio $E(u)/E_{\max}$ enters the stage cost.

With $\gamma=0.99$, cost-minimizing value iteration applies

Iteration stops when $\lVert V_{n+1}-V_n\rVert_\infty<10^{-10}$. The final Bellman residual is checked independently and must be below 10^-8.

Policy slices compare the optimal phase decision across the two queue lengths and the age of the oldest prompt. A replay applies the resulting policy to fixed evaluation episodes and compares it with the decode-priority rule.

from pathlib import Path
import sys

from IPython.display import HTML, display

code_dir = Path.cwd() / "code"
if str(code_dir) not in sys.path:
    sys.path.insert(0, str(code_dir))

from inference_replay import render_serving_replay

display(HTML(render_serving_replay(
    Path("artifacts/inference_serving/textbook_results.json"),
    view="scheduling",
)))

The reduced scheduling MDP has the following transition parameters and value-iteration certificate. The Bellman residual applies only to this finite model.

import pandas as pd

dp_row = pd.read_csv("artifacts/inference_serving/metrics_dp.csv").iloc[0]
pd.DataFrame(
    {
        "value": [
            f"{dp_row['bellman_residual']:.3e}",
            f"{int(dp_row['iterations']):,}",
            dp_row["profile_status"],
            f"{dp_row['arrival_probability']:.6f}",
            f"{dp_row['prefill_completion_probability']:.6f}",
            f"{dp_row['decode_completion_probability']:.6f}",
            f"{dp_row['prefill_energy_j']:.3f}",
            f"{dp_row['decode_energy_j']:.3f}",
            f"{dp_row['idle_energy_j']:.3f}",
        ]
    },
    index=[
        "Bellman residual",
        "value-iteration sweeps",
        "profile provenance",
        "arrival probability",
        "prefill completion probability",
        "aggregate decode completion budget",
        "prefill energy (J)",
        "decode energy (J)",
        "idle energy (J)",
    ],
).rename_axis("quantity")

Download the dynamic-programming certificate (CSV)

Value iteration required 2,454 sweeps on this high-discount problem. Its independently recomputed Bellman residual is $9.823\times10^{-11}$, below both the 10^-8 acceptance threshold and the 10^-10 stopping tolerance. The certificate concerns the supplied 245-state transition matrix. The optimal policy serves decode whenever $d>0$, serves prefill when $d=0$ and $p>0$, and idles only in the empty state. The same rule appears in all five age slices. For this calibration, the decode-stall penalty and aggregate completion budget make the transparent decode-priority heuristic exactly optimal within the reduced model. A richer state or a different cost can produce a switching boundary instead.

The result is exact for the stated finite MDP, not for vLLM. Aggregating request ages and lengths removes distinctions that can affect head-of-line waiting and cache release. The calibrated transition kernel is stationary, the clock is fixed, temperature is absent from the state, and the action set excludes mixed prefill-decode batches. Measurements from one L4 deployment calibrate the phase rates and powers; they do not establish scheduling performance across L4 systems. The request-level replay from the earlier chapters remains a model-audit tool, not part of the optimality proof.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def solve_scheduling_mdp(
    mdp: SchedulingMDP | None = None,
    *,
    tolerance: float = 1e-10,
    maximum_iterations: int = 100_000,
) -> DPSolution:
    """Solve the reduced discounted scheduling model by value iteration."""

    model = make_scheduling_mdp() if mdp is None else mdp
    model.validate()
    if tolerance <= 0.0 or maximum_iterations <= 0:
        raise ValueError("tolerance and maximum_iterations must be positive")
    value = np.zeros(model.states.shape[0], dtype=float)
    iterations = 0
    for iterations in range(1, maximum_iterations + 1):
        continuation = np.einsum("sak,k->sa", model.transitions, value)
        q_values = model.stage_cost + model.gamma * continuation
        q_values = np.where(model.valid_actions, q_values, np.inf)
        updated = np.min(q_values, axis=1)
        if np.max(np.abs(updated - value)) <= tolerance:
            value = updated
            break
        value = updated
    continuation = np.einsum("sak,k->sa", model.transitions, value)
    q_values = np.where(
        model.valid_actions,
        model.stage_cost + model.gamma * continuation,
        np.inf,
    )
    bellman_update = np.min(q_values, axis=1)
    residual = float(np.max(np.abs(bellman_update - value)))
    policy = np.argmin(q_values, axis=1).astype(int)
    labels = tuple(SCHEDULING_ACTIONS[int(action)] for action in policy)
    return DPSolution(
        value=value,
        policy=policy,
        policy_labels=labels,
        iterations=iterations,
        bellman_residual=residual,
        mdp=model,
    )

Newton-Kantorovich Applied to Bellman Optimality¶

We now apply the Newton-Kantorovich framework to the Bellman optimality equation. Let

The problem is to find $v$ such that $\Bellman v = v$, or equivalently $\mathrm{B}(v) := \Bellman v - v = 0$. The operator $\Bellman$ is piecewise affine, hence not globally differentiable, but it is directionally differentiable everywhere in the Hadamard sense and Fréchet differentiable at points where the maximizer is unique.

We consider three complementary perspectives for understanding and computing its derivative.

Perspective 1: Max of Affine Maps¶

In tabular form, for finite state and action spaces, the Bellman operator can be written as a pointwise maximum of affine maps:

where $P_a \in \mathbb{R}^{|S| \times |S|}$ is the transition matrix associated with action $a$. Each $Q_a v := r^a + \gamma P_a v$ is affine in $v$. The operator $\Bellman$ therefore computes the upper envelope of a finite set of affine functions at each state.

At any $v$, let the active set at state $s$ be

Then the Hadamard directional derivative exists and is given by

If the active set is a singleton, this expression becomes linear in $h$, and $\Bellman$ is Fréchet differentiable at $v$, with

where $\pi_v(s) := a^*(s)$ is the greedy policy at $v$.

Perspective 2: Envelope Theorem¶

Consider now a value function approximated as a linear combination of basis functions:

At a node $s_i$, define the parametric maximization

Define

so that $v_i(c) = \max_a F_i(a, c)$. Since $F_i$ is linear in $c$, we can apply the envelope theorem (Danskin’s theorem): if the optimizer $a_i^*(c)$ is unique or selected measurably, then

We do not need to differentiate the optimizer $a_i^*(c)$ itself. The result extends to the subdifferential case when ties occur, where the Jacobian becomes set-valued.

This result is useful when solving the collocation equation $\Phi c = v(c)$. Newton’s method requires the Jacobian $v'(c)$, and this expression allows us to compute it without involving any derivatives of the optimal action.

Perspective 3: The Implicit Function Theorem¶

The third perspective applies the implicit function theorem to understand when the Bellman operator is differentiable despite containing a max operator. The maximization problem defines an implicit relationship between the value function and the optimal action, and the implicit function theorem tells us when this relationship is smooth enough to differentiate through.

The Bellman operator is defined as

The difficulty is that the max operator encodes a discrete selection: which action achieves the maximum. To apply the implicit function theorem, we reformulate this as follows. For each action $a$, define the action-value function:

The optimal action at $v$ satisfies the optimality condition:

Now suppose that at a particular $v$, action $a^*(s)$ is a strict local maximizer in the sense that there exists $\delta > 0$ such that

This strict inequality is the regularity condition needed for the implicit function theorem. It ensures that the optimal action is unique at $v$ and remains so in a neighborhood of $v$.

To see why, consider any perturbation $v + h$ with $\|h\|$ small. Since $Q_a$ is linear in $v$, we have:

The perturbation term is bounded: $|\gamma \sum_j p(j \mid s,a) h(j)| \leq \gamma \|h\|$. Therefore, for $\|h\| < \delta/\gamma$, the strict gap ensures that

Thus $a^*(s)$ remains the unique maximizer throughout the neighborhood $\{v + h : \|h\| < \delta/\gamma\}$.

The implicit function theorem now applies: in this neighborhood, the mapping $v \mapsto a^*(s; v)$ is constant (and hence smooth), taking the value $a^*(s)$. This allows us to write

as an explicit formula that holds throughout the neighborhood. Since $Q_{a^*(s)}(\cdot, s)$ is an affine (hence smooth) function of $v$, we can differentiate it:

More precisely, for any perturbation $h$:

This is the Fréchet derivative:

where $\pi_v(s) = a^*(s)$ is the greedy policy.

The role of the implicit function theorem: It guarantees that when the maximizer is unique with a strict gap (the regularity condition), the argmax function $v \mapsto a^*(s; v)$ is locally constant, which removes the non-differentiability of the max operator. Without this regularity condition (specifically, at points where multiple actions tie for optimality), the implicit function theorem does not apply, and the operator is not Fréchet differentiable. The active set perspective (Perspective 1) and the envelope theorem (Perspective 2) provide the tools to handle these non-smooth points.

Connection to Policy Iteration¶

We return to the Newton-Kantorovich step:

Suppose $\Bellman'(v_n) = \gamma P_{\pi_{v_n}}$ for the greedy policy $\pi_{v_n}$. Then

which is exactly policy evaluation for $\pi_{v_n}$. Recomputing the greedy policy from $v_{n+1}$ yields the next iterate.

Thus, policy iteration is Newton-Kantorovich applied to the Bellman optimality equation. At points of nondifferentiability (when ties occur), the operator is still semismooth, and policy iteration corresponds to a semismooth Newton method. The envelope theorem is what justifies the simplification of the Jacobian to $\gamma P_{\pi_v}$, bypassing the need to differentiate through the optimizer. This completes the equivalence.

The Semismooth Newton Perspective¶

The three perspectives we developed above (the active set view, the envelope theorem, and the implicit function theorem) all point toward a deeper framework for understanding Newton-type methods on non-smooth operators. This framework, known as semismooth Newton methods, was developed precisely to handle operators like the Bellman operator that are piecewise smooth but not globally differentiable. The connection between policy iteration and semismooth Newton methods has been rigorously developed in recent work Gargiani et al. (2022).

The classical Newton-Kantorovich method assumes the operator is Fréchet differentiable everywhere. The derivative exists, is unique, and varies continuously with the base point. But the Bellman operator $\Bellman$ violates this assumption at any value function where multiple actions tie for optimality at some state. At such points, the implicit function theorem fails, and there is no unique Fréchet derivative.

Semismooth Newton methods address this by replacing the notion of a single Jacobian with a generalized derivative that captures the behavior of the operator near non-smooth points. The most commonly used generalized derivative is the Clarke subdifferential, which we can think of as the convex hull of all possible “candidate Jacobians” that arise from limits approaching the non-smooth point from different directions.

For the Bellman residual $\mathrm{B}(v) = \Bellman v - v$, the Clarke subdifferential at a point $v$ can be characterized explicitly using our first perspective. Recall that at each state $s$, we defined the active set $\mathcal{A}^*(s; v) = \arg\max_a Q_a(v, s)$. When this set contains multiple actions, the operator is not Fréchet differentiable. However, it remains directionally differentiable in all directions, and the Clarke subdifferential consists of all matrices of the form

In words, the generalized Jacobian is the set of all matrices $I - \gamma P_\pi$ where $\pi$ is any policy that selects an action from the active set at each state. When the maximizer is unique everywhere, this set reduces to a singleton, and we recover the classical Fréchet derivative. When ties occur, the set has multiple elements: precisely the convex combinations mentioned in Perspective 1.

The semismooth Newton method for solving $\mathrm{B}(v) = 0$ proceeds by selecting an element $J_k \in \partial \mathrm{B}(v_k)$ at each iteration and solving

What this tells us is that any choice from the Clarke subdifferential yields a valid Newton-like update. In the context of the Bellman equation, choosing $J_k = I - \gamma P_{\pi_k}$ where $\pi_k$ is any greedy policy corresponds exactly to the policy evaluation step in policy iteration. The freedom in selecting which action to choose when ties occur translates to the freedom in selecting which element of the subdifferential to use.

Under appropriate regularity conditions (specifically, when the residual function is BD-regular or CD-regular), the semismooth Newton method converges locally at a quadratic rate Gargiani et al. (2022). This means that near the solution, the error decreases quadratically:

This theoretical result explains an empirical observation that has long been noted in practice: policy iteration typically converges in very few iterations, often just a handful, even when the state and action spaces are enormous and the space of possible policies is exponentially large.

The semismooth Newton framework also suggests a spectrum of methods interpolating between value iteration and policy iteration. Value iteration can be interpreted as a Newton-like method where we choose $J_k = I$ at every iteration, ignoring the dependence of $\Bellman$ on $v$ entirely. This choice guarantees global convergence through the contraction property but sacrifices the quadratic local convergence rate. Policy iteration, at the other extreme, uses the full generalized Jacobian $J_k = I - \gamma P_{\pi_k}$, achieving quadratic convergence but at the cost of solving a linear system at each iteration.

Between these extremes lie methods that use approximate Jacobians. One natural variant is to choose $J_k = \alpha I$ for some scalar $\alpha > 1$. This leads to the update

This is known as $\alpha$-value iteration or successive over-relaxation when $\alpha > 1$. For appropriate choices of $\alpha$, this method retains global convergence while achieving better local rates than standard value iteration, and it requires only pointwise operations rather than solving a linear system. The Newton perspective thus unifies existing algorithms and generates new ones by systematically exploring different approximations to the generalized Jacobian.

The connection to semismooth Newton methods places policy iteration within a broader mathematical framework that extends far beyond dynamic programming. Semismooth Newton methods are used in optimization (for complementarity problems and variational inequalities), in PDE-constrained optimization (for problems with control constraints), and in economics (for equilibrium problems). The Bellman equation, viewed through this lens, is simply one instance of a piecewise smooth equation, and the tools developed for such equations apply directly.

Policy Iteration¶

While we derived policy iteration-like steps from the Newton-Kantorovich method, it’s worth examining policy iteration as a standalone algorithm, as it has been traditionally presented in the field of dynamic programming.

The policy iteration algorithm for discounted Markov decision problems is as follows:

As opposed to value iteration, this algorithm produces a sequence of both deterministic Markovian decision rules $\{\pi_n\}$ and value functions $\{\mathbf{v}^n\}$. We recognize in this algorithm the linearization step of the Newton-Kantorovich procedure, which takes place here in the policy evaluation step 3 where we solve the linear system $(\mathbf{I}-\gamma \mathbf{P}_{\pi_n}) \mathbf{v} = \mathbf{r}_{\pi_n}$. In practice, this linear system could be solved either using direct methods (eg. Gaussian elimination), using simple iterative methods such as the successive approximation method for policy evaluation, or more sophisticated methods such as GMRES.

Summary and Outlook¶

Discounting converts an infinite stream of rewards into a bounded value and makes the Bellman maps contractions in the sup norm. Value iteration applies successive approximation; policy iteration alternates a fixed-policy linear solve with greedy improvement. The inference-scheduling example also shows the price of state reduction: a smaller MDP averages over distinctions that can still affect future completions.

The hard maximum in the Bellman operator selects one action and is nondifferentiable at ties. Can the decision rule remain stochastic while the operator becomes smooth? Smooth and regularized dynamic programming answers by placing a convex regularizer on the action distribution.

Self-checks¶