"""
=============================================================================
第 15 讲 · 代码 3：模型误差——复合误差与"优化器的诅咒"
=============================================================================

对应讲义：lectures/lecture15_model_based_rl.md 第 4–5 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
前两个文件的迷宫是确定性的、离散的，模型可以做到完全准确。现实中，学到的模型总有误差。两个后果：

1. **复合误差（compounding error）**：用模型往前"想象"多步时，每一步的小误差会累积、放大；
   而且一旦想象出的状态离开了训练数据覆盖的范围，模型的预测可能完全失控。
   —— A 部分：一个没有摩擦的单摆，用线性/三次多项式回归学习它的动力学，比较单步误差和多步误差。

2. **优化器的诅咒（optimizer's curse）**：规划 = 在模型中**挑选**看起来最好的策略。
   模型的误差有正有负，而被挑中的策略，往往恰好是那些"误差对它有利"的策略——
   所以规划出的策略在模型中的价值，系统性地**高于**它在真实环境中的价值。
   —— B 部分：一个会打滑的格子世界，用有限的数据估计模型，规划，然后在真实环境中检验。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture15
    python 03_model_error.py      （约 5 秒）
"""

import numpy as np

from maze import pad

np.set_printoptions(precision=4, suppress=True)


# =============================================================================
# A 部分：单摆与复合误差
# =============================================================================
DT = 0.05


def pendulum_step(theta, omega, u):
    """
    没有摩擦的单摆：θ'' = -sin(θ) + u（u 是施加的力矩）。用欧拉法离散化，时间步长 DT。
    【参数】都可以是数组（同时模拟很多条轨迹）。
    【返回】(θ', ω')
    """
    omega2 = omega + DT * (-np.sin(theta) + u)
    return theta + DT * omega2, omega2


def features(theta, omega, u, degree):
    """回归用的特征：θ 的 1～degree 次方、ω、u、常数项。degree = 1 就是线性模型。"""
    cols = [theta ** p for p in range(1, degree + 1)] + [omega, u, np.ones_like(theta)]
    return np.stack(cols, axis=-1)


def fit_model(n_data, degree, rng):
    """
    从 n_data 个随机的 (θ, ω, u) 采样真实的下一步（加一点测量噪声），用最小二乘拟合"状态的变化量"。
    训练数据只覆盖 θ ∈ [-2, 2]、ω ∈ [-2, 2]。
    【返回】一个函数 model(θ, ω, u) → (θ', ω')
    【为什么预测"变化量"而不是下一个状态本身？】变化量很小，更容易拟合；这是学习动力学模型的常用技巧。
    """
    th = rng.uniform(-2, 2, n_data)
    om = rng.uniform(-2, 2, n_data)
    u = rng.uniform(-1, 1, n_data)
    th2, om2 = pendulum_step(th, om, u)
    Y = np.stack([th2 - th, om2 - om], axis=-1) + rng.normal(0, 0.002, (n_data, 2))
    W, *_ = np.linalg.lstsq(features(th, om, u, degree), Y, rcond=None)

    def model(theta, omega, u_):
        d = features(theta, omega, u_, degree) @ W
        return theta + d[..., 0], omega + d[..., 1]
    return model


def rollout_errors(model, horizons, rng, n_traj=500, T=200):
    """
    从同样的初始状态、用同样的动作序列，分别在真实系统和模型中往前推 T 步（开环预测），
    记录第 k 步时 θ 的预测误差。

    【返回】列表：每个 k 对应 (误差的中位数, 发散的轨迹比例)
        "发散"：模型预测的 |θ| 超过 10 或变成了 nan（数值爆炸）。用中位数是因为少数发散的轨迹会让平均值失去意义。
    """
    th0 = rng.uniform(-1.5, 1.5, n_traj)
    om0 = rng.uniform(-1, 1, n_traj)
    U = rng.uniform(-1, 1, (n_traj, T))
    tt, ot = th0.copy(), om0.copy()
    tm, om = th0.copy(), om0.copy()
    out = []
    with np.errstate(all="ignore"):                  # 发散的轨迹会溢出，这里不需要看警告
        for k in range(1, T + 1):
            tt, ot = pendulum_step(tt, ot, U[:, k - 1])
            tm, om = model(tm, om, U[:, k - 1])
            if k in horizons:
                err = np.abs(tt - tm)
                diverged = ~np.isfinite(tm) | (np.abs(tm) > 10)
                out.append((float(np.median(np.where(np.isfinite(err), err, np.inf))), float(diverged.mean())))
    return out


# =============================================================================
# B 部分：打滑的格子世界与优化器的诅咒
# =============================================================================
ROWS = COLS = 5
START, GOAL = (4, 0), (0, 4)
HOLES = {(1, 1), (2, 3)}
GAMMA = 0.95
N_S, N_A = ROWS * COLS, 4


def build_true_model(slip=0.2):
    """
    打滑的格子世界：选定方向后，以 1-slip 的概率走向这个方向，以 slip 的概率随机走向另外三个方向之一。
    到达终点 +1、掉进洞 -1，两者都是终止状态（用"吸收状态"表示：之后原地不动、奖励 0）。
    【返回】P[s, a, s']、R[s, a]（R 是"从 s 做 a 之后"的期望奖励）
    """
    moves = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    P = np.zeros((N_S, N_A, N_S))
    R = np.zeros((N_S, N_A))
    for s in range(N_S):
        r0, c0 = divmod(s, COLS)
        if (r0, c0) == GOAL or (r0, c0) in HOLES:
            P[s, :, s] = 1.0                           # 吸收状态
            continue
        for a in range(N_A):
            for b in range(N_A):
                p = (1 - slip) if a == b else slip / 3
                r1, c1 = r0 + moves[b][0], c0 + moves[b][1]
                if not (0 <= r1 < ROWS and 0 <= c1 < COLS):
                    r1, c1 = r0, c0
                s2 = r1 * COLS + c1
                P[s, a, s2] += p
                R[s, a] += p * (1.0 if (r1, c1) == GOAL else (-1.0 if (r1, c1) in HOLES else 0.0))
    return P, R


def value_iteration(P, R, iters=500):
    V = np.zeros(N_S)
    for _ in range(iters):
        V = (R + GAMMA * P @ V).max(axis=1)
    Q = R + GAMMA * P @ V
    return Q.argmax(axis=1), V


def evaluate(P, R, policy):
    """在模型 (P, R) 中精确计算确定性策略的价值（解线性方程组）。"""
    idx = np.arange(N_S)
    return np.linalg.solve(np.eye(N_S) - GAMMA * P[idx, policy], R[idx, policy])


def estimate_model(P, n, rng):
    """
    对每个 (s, a) 从真实环境中各采样 n 次下一个状态，数出经验转移概率 P̂。
    奖励是"到达哪个格子"的确定函数，所以 R̂ 由 P̂ 直接算出（模型的误差全部来自转移概率的估计）。
    """
    P_hat = np.zeros_like(P)
    for s in range(N_S):
        for a in range(N_A):
            nxt = rng.choice(N_S, size=n, p=P[s, a])
            P_hat[s, a] = np.bincount(nxt, minlength=N_S) / n
    reward_of = np.zeros(N_S)
    for s in range(N_S):
        rc = divmod(s, COLS)
        reward_of[s] = 1.0 if rc == GOAL else (-1.0 if rc in HOLES else 0.0)
    R_hat = P_hat @ reward_of
    for s in range(N_S):                               # 终止状态：保持吸收、奖励 0
        rc = divmod(s, COLS)
        if rc == GOAL or rc in HOLES:
            P_hat[s] = 0
            P_hat[s, :, s] = 1
            R_hat[s] = 0
    return P_hat, R_hat


def main():
    rng = np.random.default_rng(0)

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("A 部分：复合误差——单摆（θ'' = -sin θ + u），开环预测 k 步后 θ 的误差（500 条轨迹）")
    print("=" * 78)
    horizons = [1, 5, 10, 20, 50, 100, 200]
    print("   " + pad("模型", 28) + pad("", 12) + "".join(f"{'k=' + str(k):>9s}" for k in horizons))
    for label, n_data, degree in [("线性（300 个数据点）", 300, 1),
                                  ("三次多项式（300 个数据点）", 300, 3),
                                  ("三次多项式（3000 个数据点）", 3000, 3)]:
        model = fit_model(n_data, degree, rng)
        res = rollout_errors(model, horizons, np.random.default_rng(1))
        err_row = "".join(f"{e:9.4f}" if np.isfinite(e) else "     发散" for e, _ in res)
        div_row = "".join(f"{d * 100:8.1f}%" for _, d in res)
        print("   " + pad(label, 28) + pad("误差中位数", 12) + err_row)
        print("   " + pad("", 28) + pad("发散的比例", 12) + div_row)
    print("  观察：")
    print("  - 线性模型的单步误差只有千分之几弧度，但误差随 k 不断累积，200 步后已经是单步误差的上千倍——")
    print("    **单步准确，不代表多步准确**。单摆没有摩擦，误差不会被\"抹平\"，只会越积越多。")
    print("  - 三次多项式能表示 sin θ 的主要形状，单步误差小得多，多步预测在很长时间内都很准。")
    print("    但到了 k = 200，约 1% 的轨迹**发散**了：想象中的单摆转到了训练数据没有覆盖的角度（|θ| > 2），")
    print("    三次项在那里的外推完全不可信，一旦越界，误差就爆炸式地增长。")
    print("  - 更多的数据让单步更准，却不能阻止越界之后的发散——问题不在数据量，而在**分布偏移**：")
    print("    模型只在它见过的地方可信（和第 8、10 讲的分布偏移是同一件事）。")
    print("  - 这就是为什么现代的基于模型的方法（比如 MBPO）只用模型往前想象**很短**的几步。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("B 部分：优化器的诅咒——打滑的 5×5 格子世界（打滑概率 0.2，两个洞），用估计的模型规划")
    print("=" * 78)
    P, R = build_true_model()
    pi_star, V_star = value_iteration(P, R)
    s0 = START[0] * COLS + START[1]
    print(f"  真实的最优价值 V*(起点) = {V_star[s0]:.3f}")
    print("  每个 (s, a) 采样 n 次，估计模型 P̂，在 P̂ 上做价值迭代得到策略 π̂（200 次重复的平均）：")
    print("      n     π̂ 在模型中的价值    π̂ 的真实价值    差距（高估）   与最优的差距    对照：π* 在模型中的价值 - 真实价值")
    for n in [1, 2, 5, 10, 30, 100, 1000]:
        pred, true, ctrl = [], [], []
        for rep in range(200):
            P_hat, R_hat = estimate_model(P, n, np.random.default_rng(1000 * n + rep))
            pi_hat, V_hat = value_iteration(P_hat, R_hat)
            pred.append(V_hat[s0])
            true.append(evaluate(P, R, pi_hat)[s0])
            ctrl.append(evaluate(P_hat, R_hat, pi_star)[s0] - V_star[s0])
        pred, true = np.mean(pred), np.mean(true)
        print(f"   {n:5d}        {pred:7.3f}            {true:7.3f}        {pred - true:+7.3f}       {V_star[s0] - true:7.3f}"
              f"              {np.mean(ctrl):+7.3f}")
    print("  观察：")
    print("  - 规划出的策略在模型中\"看起来\"的价值，系统性地**高于**它的真实价值；数据越少，高估越严重。")
    print("  - 对照列：把**固定**的策略 π* 放到估计的模型中评估。n ≥ 5 时平均误差接近 0——")
    print("    模型本身的误差有正有负，并没有系统性地偏向乐观。n 极小（1、2）时它反而明显**偏低**：")
    print("    价值是转移概率的非线性函数，每个 (s, a) 只采样一次的模型把\"有时打滑\"变成了\"一定打滑\"，")
    print("    某些格子在模型里\"注定\"滑进洞——但这个偏差的方向与规划的高估**相反**。")
    print("    所以高估不是来自模型本身，而是来自**挑选**：规划在所有策略中挑出模型认为最好的那个，")
    print("    而它之所以\"看起来最好\"，部分原因恰恰是模型的误差对它有利（比如恰好低估了它掉进洞的概率）。")
    print("  - 这和第 4 讲 Q-learning 的最大化偏差、第 9 讲的奖励过度优化、第 10 讲离线 RL 的外推误差是同一个现象：")
    print("    **对一个有噪声的估计做优化，优化结果会系统性地利用噪声**。")
    print("  - 对策：更多（更好）的数据；在模型的不确定性上保持悲观（第 10 讲的 MOPO、MOReL 就是给模型的")
    print("    不确定性加惩罚）；或者只在模型可信的地方、只想象短短几步。")


if __name__ == "__main__":
    main()