"""
=============================================================================
第 10 讲 · 代码 3：离线控制——外推误差与悲观原则
=============================================================================

对应讲义：lectures/lecture10_offline_rl.md 第 4、5 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
01、02 文件做的是离线**评估**。现在做离线**控制**：只用一批固定的历史数据，学出一个好策略。

最直接的想法：Q-learning 是异策略的（第 4 讲），那就在历史数据上直接跑 Q-learning。
但这会出大问题：Q-learning 的目标里有 max_a' Q(s', a')——它会去"挑"那些
**数据中从来没出现过的动作**。这些动作的 Q 值从来没被更新过，只是一个任意的初始值
（在深度 RL 中，就是神经网络随意外推出来的值）。如果这个任意值碰巧偏高，
智能体就会被它吸引——而且没有任何机会通过试错来纠正（离线！）。这叫**外推误差**。

解决办法的核心思想是**悲观**：对数据中很少出现或没有出现的动作，不要轻信它们的价值。

-----------------------------------------------------------------------------
【环境：5×5 网格，有岩浆】
-----------------------------------------------------------------------------
        列 0 1 2 3 4
    行0    . . . . G        G：终点，+10，结束
    行1    . L L L .        L：岩浆，-10，结束
    行2    . . . L .        每一步：-1
    行3    L L . L .        动作：上下左右（10% 的概率随机滑向一个方向）
    行4    S . . . .        γ = 0.95

    最优路线：从 S 沿最下面一行向右，再沿最右一列向上到 G。

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

import numpy as np

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

GRID = [
    "....G",
    ".LLL.",
    "...L.",
    "LL.L.",
    "S....",
]
N_ROWS, N_COLS = 5, 5
N_CELLS = N_ROWS * N_COLS
END = N_CELLS                     # 吸收状态（第 2 讲 10.1 节）
N_STATES = N_CELLS + 1
MOVES = [(-1, 0), (1, 0), (0, -1), (0, 1)]
ARROWS = ["↑", "↓", "←", "→"]
N_ACTIONS = 4
GAMMA = 0.95
SLIP = 0.1
START = 4 * N_COLS + 0


# =============================================================================
# 第 1 部分：环境模型（只用于生成数据和评估，学习算法看不到）
# =============================================================================
def build_model():
    """返回 (P, R)：P 形状 (N_STATES, 4, N_STATES)，R 形状 (N_STATES, 4)。终点和岩浆进入 END。"""
    P = np.zeros((N_STATES, N_ACTIONS, N_STATES))
    R = np.zeros((N_STATES, N_ACTIONS))
    for s in range(N_CELLS):
        r, c = divmod(s, N_COLS)
        ch = GRID[r][c]
        for a in range(N_ACTIONS):
            if ch in "GL":
                P[s, a, END] = 1.0
                R[s, a] = 10.0 if ch == "G" else -10.0
                continue
            R[s, a] = -1.0
            for actual in range(N_ACTIONS):
                prob = (1 - SLIP) + SLIP / 4 if actual == a else SLIP / 4
                dr, dc = MOVES[actual]
                nr, nc = r + dr, c + dc
                if not (0 <= nr < N_ROWS and 0 <= nc < N_COLS):
                    nr, nc = r, c
                P[s, a, nr * N_COLS + nc] += prob
    P[END, :, END] = 1.0
    return P, R


P_TRUE, R_TRUE = build_model()


def value_iteration(P, R, iters=500):
    """第 2 讲的价值迭代，返回 Q*。"""
    Q = np.zeros((N_STATES, N_ACTIONS))
    for _ in range(iters):
        Q = R + GAMMA * P @ Q.max(axis=1)
    return Q


def policy_value(policy):
    """精确计算确定性策略 policy（长度 N_STATES 的动作数组）从起点出发的真实价值。"""
    idx = np.arange(N_STATES)
    P_pi = P_TRUE[idx, policy]
    R_pi = R_TRUE[idx, policy]
    V = np.linalg.solve(np.eye(N_STATES) - GAMMA * P_pi, R_pi)
    return float(V[START])


# =============================================================================
# 第 2 部分：收集离线数据
# =============================================================================
def behavior(kind, rng):
    """
    返回一个行为策略（函数：状态 → 动作）。
        "expert"：最优策略，但以 0.1 的概率随机动作（示范很好，但覆盖很窄）
        "random"：完全随机（示范很差；配合随机起点，覆盖很广）
    """
    Q_star = value_iteration(P_TRUE, R_TRUE)
    greedy = Q_star.argmax(axis=1)
    if kind == "expert":
        return lambda s: int(rng.integers(N_ACTIONS)) if rng.random() < 0.1 else int(greedy[s])
    return lambda s: int(rng.integers(N_ACTIONS))


NORMAL_CELLS = [s for s in range(N_CELLS) if GRID[s // N_COLS][s % N_COLS] not in "GL"]


def collect(kind, n_episodes, rng, max_steps=50):
    """
    按行为策略收集 n_episodes 个回合，返回转移列表 [(s, a, r, s'), ...]。
        - "expert" 回合从起点 S 出发；
        - "random" 回合从一个随机的普通格子出发（否则随机游走几乎总是在起点附近掉进岩浆，
          数据根本覆盖不到通往终点的路）；
        - kind = "mixed" 时，一半回合用 expert、一半用 random。
    """
    data = []
    for k in range(n_episodes):
        this_kind = kind if kind != "mixed" else ("expert" if k % 2 == 0 else "random")
        pol = behavior(this_kind, rng)
        s = START if this_kind == "expert" else int(rng.choice(NORMAL_CELLS))
        for _ in range(max_steps):
            a = pol(s)
            s_next = int(rng.choice(N_STATES, p=P_TRUE[s, a]))
            data.append((s, a, R_TRUE[s, a], s_next))
            s = s_next
            if s == END:
                break
    return data


def empirical_model(data):
    """
    从数据"数"出经验模型：N(s,a)、P̂(s'|s,a)、R̂(s,a)。
    所有离线算法都只通过它接触数据（表格情况下，在数据上反复做 Q-learning 等价于在经验模型上做价值迭代）。
    """
    counts = np.zeros((N_STATES, N_ACTIONS, N_STATES))
    r_sum = np.zeros((N_STATES, N_ACTIONS))
    for s, a, r, s2 in data:
        counts[s, a, s2] += 1
        r_sum[s, a] += r
    N = counts.sum(axis=2)
    P_hat = np.divide(counts, N[:, :, None], out=np.zeros_like(counts), where=N[:, :, None] > 0)
    R_hat = np.divide(r_sum, N, out=np.zeros_like(r_sum), where=N > 0)
    return N, P_hat, R_hat


# =============================================================================
# 第 3 部分：离线算法
# =============================================================================
def behavior_cloning(N):
    """行为克隆：每个状态选数据中最常见的动作；没见过的状态随便选（这里选 0）。"""
    return N.argmax(axis=1)


def naive_offline_q(N, P_hat, R_hat, q_init=0.0, iters=500):
    """
    朴素的离线 Q 学习（拟合 Q 迭代）：在经验模型上做价值迭代，只更新数据中出现过的 (s,a)。

        对 N(s,a) > 0：  Q(s,a) ← R̂(s,a) + γ Σ_s' P̂(s'|s,a) max_a' Q(s', a')
        对 N(s,a) = 0：  Q(s,a) 保持初始值 q_init（从来没有数据来更新它）

    【外推误差从哪里来？】
        目标中的 max_a' Q(s', a') 会考虑 s' 的**所有**动作，包括没见过的。
        本例中每一步的奖励都是 -1，见过的动作的 Q 值大多是负的；
        没见过的动作的 Q 值停在 0——看起来反而是"最好的"！
        于是 max 选中它，把这个虚假的高估值传播给所有前驱状态；最终的贪心策略也会选它。
        在深度 RL 中，q_init 的角色由"神经网络在没见过的输入上随意外推出来的值"扮演。
    """
    Q = np.full((N_STATES, N_ACTIONS), float(q_init))
    Q[END] = 0.0
    seen = N > 0
    for _ in range(iters):
        target = R_hat + GAMMA * P_hat @ Q.max(axis=1)
        Q = np.where(seen, target, Q)
    return Q


def pessimistic_lcb(N, P_hat, R_hat, c=5.0, iters=500):
    """
    悲观方法 1：置信下界（LCB）。对每个 (s,a) 减去一个与数据量有关的"不确定性惩罚"：

        Q(s,a) ← R̂(s,a) - c / sqrt(N(s,a)) + γ Σ_s' P̂(s'|s,a) max_a' Q(s', a')
        N(s,a) = 0 时，Q(s,a) = 可能的最低值（完全不相信）

    【直觉】数据越少，估计越不可信，就越要往坏处打算。这是第 11 讲 UCB"乐观"原则的反面：
        在线学习时，可以通过尝试来纠正错误，所以对不确定的动作保持**乐观**（鼓励探索）；
        离线学习时，无法尝试，任何错误都没法纠正，所以对不确定的动作保持**悲观**。
    """
    V_min = -10.0 / (1 - GAMMA)                      # 理论上的最低可能价值
    penalty = np.where(N > 0, c / np.sqrt(np.maximum(N, 1)), 0.0)
    Q = np.full((N_STATES, N_ACTIONS), V_min)
    Q[END] = 0.0
    seen = N > 0
    for _ in range(iters):
        target = R_hat - penalty + GAMMA * P_hat @ Q.max(axis=1)
        Q = np.where(seen, target, V_min)
        Q[END] = 0.0
    return Q


def support_constrained_q(N, P_hat, R_hat, min_count=1, iters=500):
    """
    悲观方法 2：限制在数据的支撑集内（BCQ 等"策略约束"方法的表格版）。

        目标中的 max 只在"数据中出现过至少 min_count 次的动作"里取；
        最终策略也只从这些动作中选。

    【直觉】干脆不考虑数据里没有的动作——"没见过的就不做"。
    """
    allowed = N >= min_count
    Q = np.zeros((N_STATES, N_ACTIONS))
    for _ in range(iters):
        Qmask = np.where(allowed, Q, -np.inf)
        v_next = np.where(allowed.any(axis=1), Qmask.max(axis=1), 0.0)
        Q = np.where(allowed, R_hat + GAMMA * P_hat @ v_next, -np.inf)
        Q[END] = 0.0
    return Q


def cql_tabular(N, P_hat, R_hat, alpha=1.0, lr=0.1, iters=3000):
    """
    悲观方法 3：保守 Q 学习（CQL，Kumar et al. 2020）的表格版。损失函数：

        L(Q) = 1/2 Σ_{(s,a)∈数据} (Q(s,a) - y(s,a))²                    ← 普通的 TD 误差
             + α Σ_{s∈数据} [ log Σ_a' exp Q(s,a') - Q(s, a_数据) ]     ← CQL 正则项

    正则项的作用：log-sum-exp 是 Q(s,·) 的"软最大值"，最小化它会把 s 处**所有**动作的 Q 往下压，
    而减去数据中动作的 Q 又把**数据中出现的动作**往上拉。净效果：没见过的动作被压低，
    见过的动作相对被抬高——学到的 Q 是真实 Q 的一个（大致的）下界。

    【实现】用经验模型上的"期望梯度"做梯度下降（N 是计数，除以总数得到数据分布）：
        TD 部分的梯度：    N(s,a) · (Q(s,a) - y(s,a))
        CQL 部分的梯度：   α · [ N(s) · softmax(Q(s,·)) - N(s,·) ]
    """
    total = N.sum()
    Ns = N.sum(axis=1, keepdims=True)
    Q = np.zeros((N_STATES, N_ACTIONS))
    for _ in range(iters):
        y = R_hat + GAMMA * P_hat @ Q.max(axis=1)
        grad_td = N * (Q - y)
        z = Q - Q.max(axis=1, keepdims=True)
        soft = np.exp(z) / np.exp(z).sum(axis=1, keepdims=True)
        grad_cql = alpha * (Ns * soft - N)
        Q -= lr * (grad_td + grad_cql) / total * N_STATES
        Q[END] = 0.0
    return Q


# =============================================================================
# 第 4 部分：演示
# =============================================================================
def render(policy, N=None):
    """画出策略；数据中从没出现过的 (状态, 所选动作) 用 ? 标出。"""
    lines = []
    for r in range(N_ROWS):
        cells = []
        for c in range(N_COLS):
            s = r * N_COLS + c
            ch = GRID[r][c]
            if ch in "GL":
                cells.append(ch)
            elif N is not None and N[s, policy[s]] == 0:
                cells.append("?")
            else:
                cells.append(ARROWS[policy[s]])
        lines.append("    " + " ".join(cells))
    return "\n".join(lines)


def run_all(N, P_hat, R_hat):
    """运行所有算法，返回 {名字: (贪心策略, 起点处的 max Q 估计)}。"""
    out = {}
    out["行为克隆"] = (behavior_cloning(N), float("nan"))
    for name, Q in [("朴素离线 Q 学习", naive_offline_q(N, P_hat, R_hat)),
                    ("悲观：LCB 惩罚", pessimistic_lcb(N, P_hat, R_hat)),
                    ("悲观：支撑集约束", support_constrained_q(N, P_hat, R_hat)),
                    ("悲观：CQL", cql_tabular(N, P_hat, R_hat))]:
        out[name] = (Q.argmax(axis=1), float(Q[START].max()))
    return out


def main():
    Q_star = value_iteration(P_TRUE, R_TRUE)
    V_opt = policy_value(Q_star.argmax(axis=1))
    print(f"最优策略从起点出发的真实价值：{V_opt:.2f}")
    print(render(Q_star.argmax(axis=1)))

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 1：外推误差——在\"接近专家\"的数据上直接跑 Q 学习（20 个回合）")
    print("=" * 78)
    rng = np.random.default_rng(0)
    data = collect("expert", 20, rng)
    N, P_hat, R_hat = empirical_model(data)
    seen_pairs = int((N[:N_CELLS] > 0).sum())
    print(f"  数据：{len(data)} 条转移，覆盖了 {seen_pairs} / {N_CELLS * N_ACTIONS} 个 (状态, 动作) 对")
    Q_naive = naive_offline_q(N, P_hat, R_hat)
    pol = Q_naive.argmax(axis=1)
    print(f"  朴素离线 Q 学习：起点处估计的价值 max_a Q = {Q_naive[START].max():.2f}，"
          f"其贪心策略的真实价值 = {policy_value(pol):.2f}")
    print("  学到的贪心策略（? = 选了一个数据中从没出现过的动作）：")
    print(render(pol, N))
    print("  → 估计的价值（0 左右）远高于真实价值——严重高估。")
    print("    原因：没见过的动作的 Q 值停在初始值 0，而见过的动作的 Q 值大多是负的（每步 -1），")
    print("    于是 max 总是选中\"没见过的\"动作。策略里满是问号：它专门挑数据里没有的动作去做，")
    print("    其中一些会把它带进岩浆或让它原地打转。离线学习无法通过试错发现这些错误。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：各种方法在三种数据上的表现（10 次独立收集的平均）")
    print("=" * 78)
    print(f"  （表中为贪心策略的真实价值；最优为 {V_opt:.2f}；从起点直接走进岩浆的价值约为 -10.5）")
    names = ["行为克隆", "朴素离线 Q 学习", "悲观：LCB 惩罚", "悲观：支撑集约束", "悲观：CQL"]
    print("   数据类型                  " + "".join(f"{n:>13s}" for n in names))
    for kind, n_ep, label in [("expert", 20, "接近专家（20 回合）"), ("random", 200, "完全随机（200 回合）"),
                              ("mixed", 50, "一半一半（50 回合）")]:
        vals = {n: [] for n in names}
        for rep in range(10):
            data = collect(kind, n_ep, np.random.default_rng(100 + rep))
            N, P_hat, R_hat = empirical_model(data)
            for n, (pol, _) in run_all(N, P_hat, R_hat).items():
                vals[n].append(policy_value(pol))
        print(f"   {label}   " + "".join(f"{np.mean(vals[n]):13.2f}" for n in names))
    print("  观察：")
    print("  - 接近专家的数据：行为克隆很好（照着专家做就行）；朴素 Q 学习很差（外推误差）；")
    print("    LCB 和 CQL 与行为克隆一样好——它们只在数据覆盖的范围内做决策。")
    print("    支撑集约束稍差：它只要求\"见过至少一次\"，专家偶然做过一次的随机动作也被当作可选项。")
    print("  - 完全随机的数据：行为克隆很差（模仿的是一个随机策略）；朴素 Q 学习反而达到了最优——")
    print("    数据覆盖了几乎所有 (状态, 动作)，外推误差不再是问题。LCB、CQL 略差一些：")
    print("    悲观有代价，它会低估那些数据较少的好动作。")
    print("  - 一半一半：LCB 和 CQL 最好（达到最优），朴素 Q 学习依然被外推误差拖累。")
    print("  - 总结：没有一种方法在所有数据上都最好，但悲观方法（尤其 LCB、CQL）是最**稳健**的——")
    print("    在不知道数据质量的情况下，它们很少犯大错。这是离线 RL 研究中的一个核心结论。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：悲观的方法学到了什么？（演示 1 的同一批接近专家的数据）")
    print("=" * 78)
    data = collect("expert", 20, np.random.default_rng(0))
    N, P_hat, R_hat = empirical_model(data)
    for name, (pol, qmax) in run_all(N, P_hat, R_hat).items():
        if name == "行为克隆":
            continue
        print(f"  【{name}】起点处估计的价值 {qmax:7.2f}，真实价值 {policy_value(pol):7.2f}")
        print(render(pol, N))
    print("  → 在数据覆盖的路线上（下面一行和右边一列），悲观方法的策略没有问号，")
    print("    估计的价值也不再虚高（LCB 甚至明显低估——这是有意为之的保守）。")
    print("    左上角的问号出现在数据中从没到过的状态：那里所有动作都没见过，选哪个都是\"?\"，")
    print("    但悲观的策略根本不会走到那里，所以无关紧要。")
    print("    代价是：它们不会去尝试数据中没有的、可能更好的动作——离线 RL 的上限受限于数据的覆盖范围。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 4：悲观过了头——数据太少时（完全随机，只有 20 个回合）")
    print("=" * 78)
    data = collect("random", 20, np.random.default_rng(100))
    N, P_hat, R_hat = empirical_model(data)
    Q_lcb = pessimistic_lcb(N, P_hat, R_hat)
    pol = Q_lcb.argmax(axis=1)
    print(f"  LCB 策略的真实价值：{policy_value(pol):.2f}；它在起点选择的动作：{ARROWS[pol[START]]}")
    print(render(pol, N))
    print("  → 起点正上方就是岩浆，LCB 选择了\"直接走进岩浆\"（价值约 -10.5）！")
    print("    数据太少，通往终点的每一步都只有一两个样本，不确定性惩罚层层累积，")
    print("    \"走一条长路去终点\"的悲观估计比\"立刻结束\"还要差——一个确定的坏结果，")
    print("    胜过一个不确定的好结果。悲观的强度（这里的 c）需要根据数据量仔细权衡。")


if __name__ == "__main__":
    main()