From 9acea2482fcff51b33798d47ef728d38cc52a4f9 Mon Sep 17 00:00:00 2001 From: Matt McKay Date: Fri, 16 Oct 2026 17:22:34 +1100 Subject: [PATCH] FIX: correct the Markov-jump Riccati recursion in LQMarkov (#982) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit solve_discrete_riccati_system, which LQMarkov.stationary_values uses, averaged the quadratic correction term over next-period Markov states outside the inverse. The control is chosen before s_{t+1} is known, so the expectation belongs inside. With P_bar_i = sum_j Π[i, j] P_j, each P_i now gets the single-state Riccati update evaluated at P_bar_i. The returned Ps are now the cost of following the returned Fs, and the Fs are optimal. The left factor of the correction term is formed directly, not as the transpose of the right factor. In the transposed form, rounding asymmetry in Ps is not damped by the closed loop. The iteration then diverges on well-posed problems with an unstable A, including the single-regime case. The TestLQMarkov reference values came from the old recursion and are regenerated. Two tests are added: Ps must equal the exact cost of following Fs, and with one regime or identical regimes LQMarkov must match LQ. Closes #982. Co-Authored-By: Claude Opus 5.5 --- quantecon/_matrix_eqn.py | 42 ++++++++----- quantecon/tests/test_lqcontrol.py | 99 +++++++++++++++++++++++++------ 2 files changed, 108 insertions(+), 33 deletions(-) diff --git a/quantecon/_matrix_eqn.py b/quantecon/_matrix_eqn.py index 4d2a431d4..1d3a30011 100644 --- a/quantecon/_matrix_eqn.py +++ b/quantecon/_matrix_eqn.py @@ -227,11 +227,25 @@ def solve_discrete_riccati(A, B, Q, R, N=None, tolerance=1e-10, max_iter=500, def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta, tolerance=1e-10, max_iter=1000): - """ + r""" Solves the stacked system of algebraic matrix Riccati equations in the Markov Jump linear quadratic control problems, by iterating Ps matrices until convergence. + The system consists of the m equations + + .. math:: + + 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) + + for :math:`i = 1, \ldots, m`, where + :math:`\bar{P}_i = \sum_{j=1}^m \Pi_{ij} P_j`. The expectation over + next period's Markov state enters inside the inverse because the + control is chosen before that state is realized. + Parameters ---------- Π : array_like(float, ndim=2) @@ -265,7 +279,7 @@ def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta, Returns ------- - Ps : array_like(float, ndim=2) + Ps : array_like(float, ndim=3) The fixed point of the stacked system of algebraic matrix Riccati equations, consists of m n x n P(s) matrices @@ -287,9 +301,6 @@ def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta, else: fail_msg = "Convergence failed after {} iterations." - # == Prepare array for iteration == # - sum1, sum2 = np.empty((n, n)), np.empty((n, n)) - # == Main loop == # iteration = 0 while error > tolerance: @@ -299,18 +310,17 @@ def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta, else: error = 0 + # P_bar[i] = sum_j Π[i, j] Ps[j] + P_bar = np.einsum('ij,jkl->ikl', Π, Ps) for i in range(m): - # Initialize arrays - sum1[:, :] = 0. - sum2[:, :] = 0. - 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 + BPA_N = beta * Bs[i].T @ P_bar[i] @ As[i] + Ns[i] + # Formed directly: as BPA_N.T, the left factor would let + # rounding asymmetry in Ps grow from one sweep to the next + APB_N = beta * As[i].T @ P_bar[i] @ Bs[i] + Ns[i].T + Ps1[i][:, :] = \ + Rs[i] + beta * As[i].T @ P_bar[i] @ As[i] - \ + APB_N @ solve(Qs[i] + beta * Bs[i].T @ P_bar[i] @ Bs[i], + BPA_N) error += np.max(np.abs(Ps1[i] - Ps[i])) Ps[:, :, :] = Ps1[:, :, :] diff --git a/quantecon/tests/test_lqcontrol.py b/quantecon/tests/test_lqcontrol.py index 5e250a329..fe500df56 100644 --- a/quantecon/tests/test_lqcontrol.py +++ b/quantecon/tests/test_lqcontrol.py @@ -95,6 +95,35 @@ def test_stationary_mat(self): assert_allclose(val_func_lq, val_func_answer, atol=1e-3) +def _lq_markov_policy_cost(lq, Fs): + """ + Return the matrices V(s) such that, when it is finite, the expected + discounted cost of following u = -F(s) x from (x, s) in `lq` is + x' V(s) x plus a term that does not depend on x. They solve the stacked + linear system + + V(i) = R_F(i) + beta A_F(i)' [sum_j Π[i, j] V(j)] A_F(i) + + where A_F(i) = A(i) - B(i) F(i) and + R_F(i) = R(i) + F(i)' Q(i) F(i) - F(i)' N(i) - N(i)' F(i). + + """ + m, n = lq.m, lq.n + M = np.identity(m * n**2) + b = np.empty(m * n**2) + for i in range(m): + A_F = lq.As[i] - lq.Bs[i] @ Fs[i] + rows = slice(i * n**2, (i+1) * n**2) + for j in range(m): + # vec(A_F' V A_F) = kron(A_F', A_F') vec(V) + M[rows, j * n**2:(j+1) * n**2] -= \ + lq.beta * lq.Π[i, j] * np.kron(A_F.T, A_F.T) + R_F = lq.Rs[i] + Fs[i].T @ lq.Qs[i] @ Fs[i] - \ + Fs[i].T @ lq.Ns[i] - lq.Ns[i].T @ Fs[i] + b[rows] = R_F.ravel() + return np.linalg.solve(M, b).reshape(m, n, n) + + class TestLQMarkov: def setup_method(self): @@ -154,8 +183,8 @@ def test_scalar_sequences_with_seed(self): lq_markov_scalar = self.lq_markov_scalar x0 = 2 - expected_x_seq = np.array([[2., 1.15977567, 0.6725398]]) - expected_u_seq = np.array([[1.28044866, 0.7425166]]) + expected_x_seq = np.array([[2., 1.15843266, 0.67098311]]) + expected_u_seq = np.array([[1.28313468, 0.74321256]]) expected_w_seq = np.array([[1.3486939, 0.55721062, 0.53423587]]) expected_state = np.array([1, 1, 1]) @@ -172,11 +201,11 @@ def test_stationary_scalar(self): lq_markov_scalar = self.lq_markov_scalar - P_answer = np.array([[[1.51741465]], - [[1.07334181]]]) + P_answer = np.array([[[1.52007965]], + [[1.07741061]]]) d_answer = np.array([0., 0.]) - F_answer = np.array([[[-0.54697435]], - [[-0.64022433]]]) + F_answer = np.array([[[-0.54745226]], + [[-0.64156734]]]) Ps, ds, Fs = lq_markov_scalar.stationary_values() @@ -189,10 +218,10 @@ def test_mat_sequences(self): lq_markov_mat = self.lq_markov_mat1 x0 = np.array([[1000, 1, 25]]) - expected_x_seq = np.array([[1.00000000e+03, 1.01372101e+03], + expected_x_seq = np.array([[1.00000000e+03, 1.01342472e+03], [1.00000000e+00, 1.00000000e+00], [2.50000000e+01, 2.61845443e+01]]) - expected_u_seq = np.array([[1013.72101253]]) + expected_u_seq = np.array([[1013.42471946]]) expected_w_seq = np.array([[0.41782708, 1.18454431]]) expected_state = np.array([1, 1]) @@ -207,15 +236,15 @@ def test_mat_sequences(self): def test_stationary_mat(self): lq_markov_mat = self.lq_markov_mat1 - d_answer = np.array([16.2474886, 16.31935939]) - P_answer = np.array([[[4.5144056e-02, 1.8627227e+01, 1.9348906e-01], - [1.8627227e+01, 7.9343733e+03, 8.0055130e+01], - [1.9348906e-01, 8.0055130e+01, 8.2991316e-01]], - [[5.3606323e-02, 2.0136657e+01, 2.1763323e-01], - [2.0136657e+01, 7.8100167e+03, 8.1960509e+01], - [2.1763323e-01, 8.1960509e+01, 8.8413147e-01]]]) - F_answer = np.array([[[-0.98437714, 19.2051657, -0.83142157]], - [[-1.01434303, 21.58480004, -0.83851124]]]) + d_answer = np.array([16.31235346, 16.38442572]) + P_answer = np.array([[[4.5345640e-02, 1.8660614e+01, 1.9430712e-01], + [1.8660614e+01, 7.9402383e+03, 8.0190872e+01], + [1.9430712e-01, 8.0190872e+01, 8.3325640e-01]], + [[5.3840146e-02, 2.0171229e+01, 2.1853447e-01], + [2.0171229e+01, 7.8154682e+03, 8.2093814e+01], + [2.1853447e-01, 8.2093814e+01, 8.8762671e-01]]]) + F_answer = np.array([[[-0.98417975, 19.23774624, -0.83061122]], + [[-1.01410488, 21.6197524, -0.83758363]]]) Ps, ds, Fs = lq_markov_mat.stationary_values() @@ -223,6 +252,42 @@ def test_stationary_mat(self): assert_allclose(P_answer, Ps, atol=1e-6) assert_allclose(d_answer, ds, atol=1e-6) + def test_stationary_values_are_policy_costs(self): + # x' P(s) x must be the x-dependent part of the discounted cost of + # following u = -F(s) x from (x, s). This fails if the expectation + # over next period's Markov state is taken outside the inverse in + # the Riccati recursion (#982). + Π = np.array([[0.5, 0.5], + [0.3, 0.7]]) + As = [[[1., .1], [0., .9]], [[.9, .3], [.2, 1.]]] + Bs = [[[1.], [0.]], [[.2], [1.]]] + Rs = [np.eye(2), 2 * np.eye(2)] + lq_markov_mat2 = LQMarkov(Π, [.1, 5.], Rs, As, Bs, beta=.95) + + for lq in (self.lq_markov_scalar, self.lq_markov_mat1, + lq_markov_mat2): + Ps, ds, Fs = lq.stationary_values() + + assert_allclose(Ps, _lq_markov_policy_cost(lq, Fs), rtol=1e-6) + + def test_identical_regimes_match_lq(self): + # With one regime, or identical regimes, LQMarkov solves a standard + # LQ problem. A is unstable, so this also fails if the Riccati + # iteration lets rounding asymmetry in Ps grow. + A = np.array([[1.5, .3], [.1, 1.5]]) + B = np.array([[1.], [1.]]) + P, F, d = LQ(1., np.eye(2), A, B, beta=.99).stationary_values() + + for Π in ([[1.]], [[.9, .1], [.2, .8]]): + m = len(Π) + lq = LQMarkov(Π, [1.] * m, [np.eye(2)] * m, [A] * m, [B] * m, + beta=.99) + Ps, ds, Fs = lq.stationary_values() + + for i in range(m): + assert_allclose(Ps[i], P, rtol=1e-8) + assert_allclose(Fs[i], F, rtol=1e-8) + def test_beta_one_without_shocks_one_state(self): lq = LQMarkov([[1.]], [1.], [1.], [.5], [0.], beta=1)