BUG: LQMarkov / solve_discrete_riccati_system averages over next-period states outside the inverse #982

Summary

solve_discrete_riccati_system (used by LQMarkov.stationary_values) iterates on the wrong Riccati recursion. It averages the quadratic correction term over next-period Markov states outside the matrix inverse. Because the control $u_t$ is chosen before $s_{t+1}$ is known, the expectation over $s_{t+1}$ belongs inside the inverse.

As a result, LQMarkov returns value matrices Ps that are not the cost of any policy, and decision rules Fs that are suboptimal. The error is zero only when each row of $\Pi$ has a single nonzero entry, or when $Q_i + \beta B_i' P_j B_i$ and $B_i' P_j A_i$ do not vary with $j$.

The problem is present in 0.10.1 and on current main (0.13.0, 70fa215).

The correct recursion

With $\bar P_i = \sum_j \Pi_{ij} P_j$, the Bellman equation

$$ x' P_i x + \rho_i = \min_u \Big{ x'R_i x + u' Q_i u + 2 u' N_i x + \beta \sum_j \Pi_{ij} \mathbb{E}\big[(A_i x + B_i u + C_i w)' P_j (A_i x + B_i u + C_i w) + \rho_j\big] \Big} $$

gives

$$ P_i = R_i + \beta A_i' \bar P_i A_i - (\beta B_i' \bar P_i A_i + N_i)' (Q_i + \beta B_i' \bar P_i B_i)^{-1} (\beta B_i' \bar P_i A_i + N_i), \qquad F_i = (Q_i + \beta B_i' \bar P_i B_i)^{-1} (\beta B_i' \bar P_i A_i + N_i). $$

What the code does

for j in range(m):
    sum1 += beta * Π[i, j] * As[i].T @ Ps[j] @ As[i]
    sum2 += Π[i, j] * \
            (beta * As[i].T @ Ps[j] @ Bs[i] + Ns[i].T) @ \
            solve(Qs[i] + beta * Bs[i].T @ Ps[j] @ Bs[i],
                  beta * Bs[i].T @ Ps[j] @ As[i] + Ns[i])
Ps1[i][:, :] = Rs[i] + sum1 - sum2

So sum2 is $\sum_j \Pi_{ij} (\beta B_i' P_j A_i + N_i)' (Q_i + \beta B_i' P_j B_i)^{-1} (\cdot)$, not the expression above.

The F formula in LQMarkov.stationary_values already uses $\bar P_i$ inside the inverse (_lqcontrol.py L536–545), which is correct. But it is evaluated at the wrong Ps.

Reproducer

The script below computes LQMarkov's solution and the solution of the correct recursion. It then evaluates the exact discounted cost of each policy by iterating on $V_i = R_i + F_i'Q_iF_i - F_i'N_i - N_i'F_i + \beta (A_i - B_iF_i)' \bar V_i (A_i - B_iF_i)$.

import numpy as np
import quantecon as qe

rng = np.random.default_rng(0)
n, k, m, β = 2, 1, 2, 0.95
Π = np.array([[0.5, 0.5],
              [0.3, 0.7]])
As = [np.array([[1.0, 0.1], [0.0, 0.9]]), np.array([[0.9, 0.3], [0.2, 1.0]])]
Bs = [np.array([[1.0], [0.0]]), np.array([[0.2], [1.0]])]
Rs = [np.eye(n), 2 * np.eye(n)]
Qs = [np.array([[0.1]]), np.array([[5.0]])]
Ns = [np.zeros((k, n)), np.zeros((k, n))]

lq = qe.LQMarkov(Π, Qs, Rs, As, Bs, Ns=Ns, beta=β)
Ps_lib, ds_lib, Fs_lib = lq.stationary_values()

def correct_riccati(tol=1e-12, max_iter=10_000):
    "Iterate P_i = R_i + β A_i' P̄_i A_i - (β B_i' P̄_i A_i + N_i)'(Q_i + β B_i' P̄_i B_i)^{-1}(...)."
    Ps = np.array([np.eye(n) for _ in range(m)])
    for _ in range(max_iter):
        P_bar = np.einsum('ij,jkl->ikl', Π, Ps)
        Fs = np.array([np.linalg.solve(Qs[i] + β * Bs[i].T @ P_bar[i] @ Bs[i],
                                       β * Bs[i].T @ P_bar[i] @ As[i] + Ns[i]) for i in range(m)])
        Ps_new = np.array([Rs[i] + β * As[i].T @ P_bar[i] @ As[i]
                           - (β * Bs[i].T @ P_bar[i] @ As[i] + Ns[i]).T @ Fs[i] for i in range(m)])
        if np.max(np.abs(Ps_new - Ps)) < tol:
            return Ps_new, Fs
        Ps = Ps_new

def policy_cost(Fs, tol=1e-12, max_iter=10_000):
    "Matrices V_i with x' V_i x = expected discounted cost of following u = -F_{s} x from state (x, i)."
    Vs = np.zeros((m, n, n))
    for _ in range(max_iter):
        V_bar = np.einsum('ij,jkl->ikl', Π, Vs)
        Vs_new = np.array([Rs[i] + Fs[i].T @ Qs[i] @ Fs[i] - Fs[i].T @ Ns[i] - Ns[i].T @ Fs[i]
                           + β * (As[i] - Bs[i] @ Fs[i]).T @ V_bar[i] @ (As[i] - Bs[i] @ Fs[i])
                           for i in range(m)])
        if np.max(np.abs(Vs_new - Vs)) < tol:
            return Vs_new
        Vs = Vs_new

Ps_cor, Fs_cor = correct_riccati()
V_lib, V_cor = policy_cost(Fs_lib), policy_cost(Fs_cor)

np.set_printoptions(precision=4, suppress=True)
print("quantecon", qe.__version__)
print("Fs (LQMarkov):        ", Fs_lib.reshape(m, -1))
print("Fs (correct):         ", Fs_cor.reshape(m, -1))
print("max |Ps_lib - V_lib| =", np.abs(Ps_lib - V_lib).max(),
      " (LQMarkov's P is not the cost of its own policy)")
print("max |Ps_cor - V_cor| =", np.abs(Ps_cor - V_cor).max())
for i in range(m):
    print(f"state {i}: eigenvalues of V_lib - V_cor =", np.linalg.eigvalsh(V_lib[i] - V_cor[i]))

Output (identical on 0.10.1 and main):

Fs (LQMarkov):         [[0.9644 0.3243]
 [0.2439 0.5083]]
Fs (correct):          [[0.9652 0.3295]
 [0.2486 0.5117]]
max |Ps_lib - V_lib| = 0.12200161184806113  (LQMarkov's P is not the cost of its own policy)
max |Ps_cor - V_cor| = 2.1316282072803006e-14
state 0: eigenvalues of V_lib - V_cor = [0.     0.0002]
state 1: eigenvalues of V_lib - V_cor = [0.     0.0006]

• LQMarkov's Ps differ from the true cost of its own Fs by 0.12. The correct recursion's Ps equal its policy's cost to machine precision.
• $V^{\text{lib}}_i - V^{\text{cor}}_i$ is positive semidefinite and nonzero, so LQMarkov's policy is strictly costlier than the optimal one.

Suggested fix

Inside the loop over i, form P_bar = sum_j Π[i, j] * Ps[j] and use the single-state Riccati update with P_bar. That is, replace sum1/sum2 with

P_bar = np.tensordot(Π[i], Ps, axes=1)
Ps1[i] = (Rs[i] + beta * As[i].T @ P_bar @ As[i]
          - (beta * Bs[i].T @ P_bar @ As[i] + Ns[i]).T
          @ solve(Qs[i] + beta * Bs[i].T @ P_bar @ Bs[i],
                  beta * Bs[i].T @ P_bar @ As[i] + Ns[i]))

The reference values in TestLQMarkov.test_stationary_scalar / test_stationary_mat (and the sequence tests that depend on Fs) were presumably generated by the current recursion, so they will need to be regenerated. A test that checks Ps against the policy-evaluation fixed point, as in the script above, would guard against a regression.

Context

Found by Tom Sargent while writing the new policy_regimes lecture. markov_jump_lq in lecture-python-advanced now states the correct equations and notes the discrepancy (QuantEcon/lecture-python-advanced.myst#392). Once this is fixed and released, those notes can be removed.

🤖 Generated with Claude Code