BUG: LQMarkov.compute_sequence uses next period's A, B and C in the law of motion

The bug is in the simulation loop of LQMarkov.compute_sequence in quantecon/_lqcontrol.py (around L619-L633), which indexes As, Bs and Cs by the destination period's regime. Start by running the reproducer in the issue to confirm the mismatch with stationary_values. Done means the loop uses the regime current at the start of each step, the regression test test_compute_sequence_switching_regimes in quantecon/tests/test_lqcontrol.py passes, and the 15 existing tests still pass.

Summary

LQMarkov documents the law of motion x_{t+1} = A(s_t) x_t + B(s_t) u_t + C(s_t) w_{t+1}, and stationary_values derives F, P and d for that timing. compute_sequence instead steps from x_t to x_{t+1} with A(s_{t+1}), B(s_{t+1}) and C(s_{t+1}), while still setting u_t = -F(s_t) x_t. Whenever the regime switches between states whose A, B or C differ, the returned paths are therefore not the optimal paths of the documented model, and their realised cost does not match the value from stationary_values (see the reproducer). No existing test can catch this: both seeded sequence tests draw constant regime paths ([1, 1, 1] and [1, 1]), and lq_markov_mat1 has the same A, B and C in both regimes.

What the code does

The class docstring states the law of motion with s_t (L358-L362). stationary_values builds F(s) from A(s), B(s) and sum_j Π[s, j] P_j, and it builds d from tr(P_j C(s) C(s)') (L532-L553). The simulation in compute_sequence uses the next regime instead:

https://github.com/QuantEcon/QuantEcon.py/blob/70fa215530518c0b504c76ba8458c8cfaafd52c9/quantecon/_lqcontrol.py#L619-L633

        Cw_path = np.empty((self.n, T+1))
        for i in range(T+1):
            Cw_path[:, i] = Cs[state[i]] @ w_path[:, i]

        # == Use policy sequence to generate states and controls == #
        x_path[:, 0] = x0.flatten()
        u_path[:, 0] = - (Fs[state[0]] @ x0).flatten()
        for t in range(1, T):
            Ax = As[state[t]] @ x_path[:, t-1]
            Bu = Bs[state[t]] @ u_path[:, t-1]
            x_path[:, t] = Ax + Bu + Cw_path[:, t]
            u_path[:, t] = - (Fs[state[t]] @ x_path[:, t])
        Ax = As[state[T]] @ x_path[:, T-1]
        Bu = Bs[state[T]] @ u_path[:, T-1]
        x_path[:, T] = Ax + Bu + Cw_path[:, T]

The loop index t is the destination period. So the step from x_{t-1} to x_t uses A(s_t), B(s_t) and C(s_t), even though u_{t-1} = -F(s_{t-1}) x_{t-1} was chosen for regime s_{t-1}.

Reproducer

import numpy as np
import quantecon as qe

# Two regimes that alternate every period; A and B differ by regime, C = 0.
# Π is a permutation, so the #982 averaging error does not arise here.
Π = [[0., 1.],
     [1., 0.]]
A, B = np.array([1.2, 0.5]), np.array([1.0, 0.3])
beta, x0, T = 0.9, 1.0, 600

lq = qe.LQMarkov(Π, Qs=[1., 1.], Rs=[1., 1.], As=A, Bs=B, beta=beta)
Ps, _, Fs = lq.stationary_values()
P, F = Ps[:, 0, 0], Fs[:, 0, 0]


def cost(x, u):  # sum_t beta^t (x_t^2 + u_t^2), since Q = R = 1
    return np.sum(beta**np.arange(T) * (x[:T]**2 + u**2))


for seed in (0, 1):  # compute_sequence draws s_0 at random
    x, u, _, s = lq.compute_sequence(x0, ts_length=T, random_state=seed)
    x, u = x[0], u[0]

    # Documented law x_{t+1} = A(s_t) x_t + B(s_t) u_t, same rule u_t = -F(s_t) x_t
    x_doc = np.empty(T + 1)
    x_doc[0] = x0
    for t in range(T):
        x_doc[t+1] = (A[s[t]] - B[s[t]] * F[s[t]]) * x_doc[t]
    u_doc = -F[s[:T]] * x_doc[:T]

    resid = x[1:] - (A[s[1:]] * x[:T] + B[s[1:]] * u)
    print(f"s_0..s_3 = {s[:4]}")
    print("  x_0..x_3, compute_sequence:", x[:4])
    print("  x_0..x_3, documented law  :", x_doc[:4])
    print("  max|x_{t+1} - A(s_{t+1}) x_t - B(s_{t+1}) u_t|:", np.abs(resid).max())
    print(f"  x0' P(s_0) x0                   = {P[s[0]] * x0**2:.10f}")
    print(f"  realised cost, documented law   = {cost(x_doc, u_doc):.10f}")
    print(f"  realised cost, compute_sequence = {cost(x, u):.10f}")

Output on main (70fa215):

s_0..s_3 = [0 1 0 1]
  x_0..x_3, compute_sequence: [1.         0.30240861 0.29905565 0.090437  ]
  x_0..x_3, documented law  : [1.         0.54136203 0.23639858 0.12797722]
  max|x_{t+1} - A(s_{t+1}) x_t - B(s_{t+1}) u_t|: 0.0
  x0' P(s_0) x0                   = 1.7903655698
  realised cost, documented law   = 1.7903655699
  realised cost, compute_sequence = 1.6384709616
s_0..s_3 = [1 0 1 0]
  x_0..x_3, compute_sequence: [1.         0.98891248 0.29905565 0.29573986]
  x_0..x_3, documented law  : [1.         0.43667374 0.23639858 0.10322905]
  max|x_{t+1} - A(s_{t+1}) x_t - B(s_{t+1}) u_t|: 0.0
  x0' P(s_0) x0                   = 1.3518125360
  realised cost, documented law   = 1.3518125360
  realised cost, compute_sequence = 2.4866632394

compute_sequence's path satisfies x_{t+1} = A(s_{t+1}) x_t + B(s_{t+1}) u_t exactly and leaves the documented path at x_1. Its realised cost misses x0' P(s_0) x0 (1.638 vs 1.790, and 2.487 vs 1.352). The documented-law path matches it to within 2e-11. With B = (1, -1) instead, the optimal F stabilises the documented law (|x_40| ≈ 1.7e-19), but compute_sequence's path grows (|x_40| ≈ 2.9e4).

Suggested fix

Step each period with the regime that is current at the start of the step. Replacing L619-L633 with a single forward loop does this and removes Cw_path:

        # == Use policy sequence to generate states and controls == #
        x_path[:, 0] = x0.flatten()
        for t in range(T):
            s = state[t]
            u_path[:, t] = -(Fs[s] @ x_path[:, t])
            x_path[:, t+1] = (As[s] @ x_path[:, t] + Bs[s] @ u_path[:, t]
                              + Cs[s] @ w_path[:, t+1])

A smaller diff that keeps the current structure also works:

• index As and Bs by state[t-1] inside the loop and by state[T-1] after it;
• set Cw_path[:, i] = Cs[state[i-1]] @ w_path[:, i] for i in range(1, T+1).

The two versions gave bit-identical output in every case I ran.

• Outputs. The random draws do not change (states are drawn first, then w), so state and w_path are the same for a given seed. x_path and u_path change only on paths that switch between regimes whose A, B or C differ.
• Existing tests. Both seeded tests have constant regime paths. With either version applied to main, their x, u, w and state are bit-identical, and quantecon/tests/test_lqcontrol.py passes unchanged (15 passed).
• Lectures. In lecture-python-advanced, the LQMarkov.compute_sequence calls in markov_jump_lq and tax_smoothing_1/2/3 all use the same A, B and C in every regime, so those paths do not change. I checked this bit-for-bit for ex1_a in markov_jump_lq, whose regime alternates every period.

Here is a regression test for TestLQMarkov in quantecon/tests/test_lqcontrol.py, using a switching chain. It fails on main and passes with either version. Because Π is a permutation, #982 does not affect it:

    def test_compute_sequence_switching_regimes(self):
        # Regimes alternate every period, and A and B differ across them
        beta, T = 0.9, 600
        lq = LQMarkov([[0., 1.], [1., 0.]], [1., 1.], [1., 1.],
                      [1.2, 0.5], [1., 0.3], beta=beta)
        Ps, ds, Fs = lq.stationary_values()
        x, u, w, state = lq.compute_sequence(1., ts_length=T,
                                             random_state=0)
        x, u, s = x[0], u[0], state[:T]
        A, B = lq.As[:, 0, 0], lq.Bs[:, 0, 0]

        # x_{t+1} = A(s_t) x_t + B(s_t) u_t
        assert_allclose(x[1:], A[s] * x[:T] + B[s] * u)
        # Realised discounted cost equals x0' P(s0) x0
        cost = np.sum(beta**np.arange(T) * (x[:T]**2 + u**2))
        assert_allclose(cost, Ps[state[0], 0, 0])

While there: the docstring summary (L561-L563) says u_0,..., u_T, but u_path has T columns (u_0, ..., u_{T-1}). LQ.compute_sequence has the same wording (L260-L262) and the same T-column u_path.

Context

Found while verifying #984, the fix for #982. The state[t] indexing goes back to the commit that added LQMarkov (5fea2de, merged in #489).

🤖 Generated with Claude Code