""" ============================================================================= 第 8 讲 · 代码 3：最大熵逆强化学习（MaxEnt IRL）——从行为反推奖励 ============================================================================= 对应讲义：lectures/lecture08_imitation_learning.md 第 4、5 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 行为克隆学的是"在每个状态做什么"（策略）。逆强化学习（IRL）换一个角度： **专家为什么这样做？——他在最大化什么奖励？** 一旦恢复了奖励，就可以用任何 RL / 规划算法求出策略。这有一个巨大的好处： **奖励比策略更容易迁移**。"水很危险、终点很好"这样的偏好，换一张地图依然成立； 而"在第 3 行第 2 列往右走"这样的策略，换一张地图就毫无意义了。 本文件实现 Ziebart et al. (2008) 的最大熵 IRL： - 假设专家是"玻尔兹曼理性"的：轨迹的概率 ∝ exp(轨迹的总奖励) - 奖励是特征的线性函数 R(s) = w · φ(s) - 用极大似然估计 w：梯度 = 专家的特征计数 - 当前奖励下的期望特征计数 ----------------------------------------------------------------------------- 【环境：两张 6×6 地图】 ----------------------------------------------------------------------------- 训练地图（专家在这里示范） 新地图（测试迁移，没有任何示范） 列 0 1 2 3 4 5 列 0 1 2 3 4 5 行0 . . . . . G 行0 G . . . . . 行1 . W W . . . 行1 . . . W W . 行2 . W W . s . 行2 . s . W W . 行3 . . . . s . 行3 . s . . . . 行4 s s . . . . 行4 . . . . s s 行5 A . . . . . 行5 . . . . . B G = 终点（好）、W = 水（很差）、s = 沙地（有点差）；A、B 为起点。 - 动作：上、下、左、右、不动（确定性移动，撞墙原地不动） - 回合长度 T = 20；每一步获得所在格子的奖励（到了终点可以"不动"，每一步都拿终点奖励） - 特征 φ(s) = [是否终点, 是否水, 是否沙地] - 真实奖励（IRL 看不到）：w* = [2.0, -3.0, -1.0] ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_maxent_irl.py （约 20 秒） """ import numpy as np np.set_printoptions(precision=3, suppress=True) MAP_TRAIN = [ ".....G", ".WW...", ".WW.s.", "....s.", "ss....", "A.....", ] MAP_NEW = [ "G.....", "...WW.", ".s.WW.", ".s....", "....ss", ".....B", ] MOVES = [(-1, 0), (1, 0), (0, -1), (0, 1), (0, 0)] # 上 下 左 右 不动 ARROWS = ["↑", "↓", "←", "→", "·"] N_ACTIONS = len(MOVES) T = 20 W_TRUE = np.array([2.0, -3.0, -1.0]) # ============================================================================= # 第 1 部分：环境 # ============================================================================= def build_env(grid): """ 把地图变成 (next_state, features, start)： next_state[s, a] ：确定性转移 features[s] ：[是否终点, 是否水, 是否沙地] start ：起点（A 或 B 所在格子） """ n_rows, n_cols = len(grid), len(grid[0]) n = n_rows * n_cols next_state = np.zeros((n, N_ACTIONS), dtype=int) features = np.zeros((n, 3)) start = None for r in range(n_rows): for c in range(n_cols): s = r * n_cols + c ch = grid[r][c] features[s] = [ch == "G", ch == "W", ch == "s"] if ch in "AB": start = s for a, (dr, dc) in enumerate(MOVES): nr, nc = r + dr, c + dc if not (0 <= nr < n_rows and 0 <= nc < n_cols): nr, nc = r, c next_state[s, a] = nr * n_cols + nc return next_state, features, start # ============================================================================= # 第 2 部分：规划（软 / 硬）与期望特征计数 # ============================================================================= def soft_value_iteration(reward, nxt): """ 有限视界的"软"价值迭代：把第 2 讲价值迭代中的 max 换成 log-sum-exp。 Q_t(s, a) = R(s) + V_{t+1}(s'(s, a)) （V_T = 0） V_t(s) = log Σ_a exp(Q_t(s, a)) （"软最大值"） π_t(a|s) = exp(Q_t(s, a) - V_t(s)) （玻尔兹曼策略） 【为什么是 log-sum-exp？】 最大熵模型假设轨迹的概率 ∝ exp(总奖励)。把所有以 (s, t) 开头的轨迹的 exp(奖励) 加起来， 它的对数恰好满足上面的递推。log-sum-exp 是 max 的平滑版本：某个 Q 远大于其他时， log Σ exp(Q) ≈ max Q；几个 Q 差不多时，它比 max 大一些（体现"有多条差不多好的路"）。 【返回】(policy, V0)：policy 形状 (T, N, A)；V0 形状 (N,)。 """ V = np.zeros(nxt.shape[0]) policy = np.zeros((T,) + nxt.shape) for t in range(T - 1, -1, -1): Q = reward[:, None] + V[nxt] m = Q.max(axis=1, keepdims=True) V = (m + np.log(np.exp(Q - m).sum(axis=1, keepdims=True)))[:, 0] # 数值稳定的 log-sum-exp policy[t] = np.exp(Q - V[:, None]) return policy, V def greedy_plan(reward, nxt): """ 普通的（硬 max）有限视界价值迭代（第 2 讲 8.7 节），返回确定性的非平稳策略（独热，形状 (T, N, A)）。 IRL 恢复奖励后，用它来求"在这个奖励下最好的策略"。 """ V = np.zeros(nxt.shape[0]) policy = np.zeros((T,) + nxt.shape) for t in range(T - 1, -1, -1): Q = reward[:, None] + V[nxt] best = Q.argmax(axis=1) policy[t, np.arange(nxt.shape[0]), best] = 1.0 V = Q.max(axis=1) return policy def expected_feature_counts(policy, start, nxt, feats): """ 在策略 policy 下，从 start 出发，T 步内特征的期望累计值 μ = E[Σ_t φ(s_t)]。 【前向递推】D_t(s) = 在时刻 t 处于状态 s 的概率： D_0 = e_start D_{t+1}(s') = Σ_{s,a} D_t(s) π_t(a|s) · 1[s'(s,a) = s'] μ = Σ_t Σ_s D_t(s) φ(s)。 """ n = nxt.shape[0] D = np.zeros(n) D[start] = 1.0 mu = np.zeros(feats.shape[1]) for t in range(T): mu += D @ feats flow = D[:, None] * policy[t] # 形状 (N, A)：从 s 选 a 的概率质量 D = np.bincount(nxt.ravel(), weights=flow.ravel(), minlength=n) # 【Python 小知识】np.bincount(下标, weights=权重)：把权重按下标累加。 # 这里把每个 (s,a) 的概率质量加到它的下一状态上。 return mu def prob_never_visit(policy, start, nxt, mask): """ 在策略 policy 下，从 start 出发的 T 步轨迹**一次都不进入** mask 标出的状态的概率。 【做法】和 expected_feature_counts 一样做前向递推，但每一步先把落在 mask 状态上的概率质量删掉， 最后剩下的总质量就是"从未进入"的概率。 """ n = nxt.shape[0] D = np.zeros(n) D[start] = 1.0 for t in range(T): D[mask] = 0.0 # 进入过 mask 的轨迹不再计入 flow = D[:, None] * policy[t] D = np.bincount(nxt.ravel(), weights=flow.ravel(), minlength=n) return float(D.sum()) def true_return(policy, env): """策略在**真实**奖励下的期望总奖励 = w* · μ（奖励是特征的线性函数）。""" nxt, feats, start = env return float(W_TRUE @ expected_feature_counts(policy, start, nxt, feats)) def sample_trajectory(policy, start, nxt, rng): """按（非平稳的）策略采样一条长度 T 的轨迹 [(s_0, a_0), ..., (s_{T-1}, a_{T-1})]。""" s, traj = start, [] for t in range(T): a = int(rng.choice(N_ACTIONS, p=policy[t, s])) traj.append((s, a)) s = nxt[s, a] return traj # ============================================================================= # 第 3 部分：最大熵 IRL # ============================================================================= def maxent_irl(demos, env, feats=None, lr=0.1, n_iters=300): """ 最大熵 IRL：求使专家轨迹的对数似然最大的奖励权重 w。 【似然】P(τ | w) = exp(Σ_t w·φ(s_t)) / Z(w)，Z(w) 是对所有可能轨迹求和的配分函数。 对数似然 L(w) = (1/N) Σ_i w·μ(τ_i) - log Z(w)，其中 log Z(w) 恰好是软价值迭代的 V_0(起点)。 【梯度】∇L(w) = μ_E - E_{τ~P(·|w)}[μ(τ)] = 专家的平均特征计数 - 当前模型的期望特征计数 直觉：如果专家比模型更常待在终点，就提高"终点"的权重；如果模型比专家更常下水， 就降低"水"的权重。收敛时两者相等——这叫**特征期望匹配**。 L(w) 是 w 的凹函数，所以梯度上升一定能找到全局最优。 【参数】 demos : 轨迹列表（都来自 env 的起点）。 feats : 特征矩阵（默认用 env 的特征；演示 4 会加一个常数特征）。 【返回】(w, 记录列表 [(迭代次数, w, 对数似然, 梯度范数), ...]) """ nxt, env_feats, start = env feats = env_feats if feats is None else feats mu_E = np.mean([feats[[s for s, _ in traj]].sum(axis=0) for traj in demos], axis=0) w = np.zeros(feats.shape[1]) log = [] for it in range(n_iters + 1): policy, V0 = soft_value_iteration(feats @ w, nxt) grad = mu_E - expected_feature_counts(policy, start, nxt, feats) if it % 50 == 0: log.append((it, w.copy(), w @ mu_E - V0[start], np.linalg.norm(grad))) w = w + lr * grad return w, log def maxent_irl_until_converged(mu_E, env, lr=0.1, tol=1e-5, max_iters=20000): """ 与 maxent_irl 相同的梯度上升，但直接从特征计数 μ_E 出发，一直迭代到梯度范数 < tol。 【为什么需要它？】 maxent_irl 固定只跑 300 轮。似然在某些方向上非常平（专家很少做的事，例如下水）， 那个方向上梯度很小、走得很慢，300 轮时还没收敛。演示 5 要比较"真正的极大似然解"，所以要跑到收敛。 【边界】 如果 μ_E 中某个特征的计数为 0（例如示范里一次都没下水），极大似然解**不存在**： 任何有限的 w 下，模型都有一点点概率下水，梯度永远指向"更负"，w 会一直减小。 这时函数会跑满 max_iters，并返回 converged = False。 【返回】(w, 迭代次数, converged) """ nxt, feats, start = env w = np.zeros(feats.shape[1]) for it in range(1, max_iters + 1): policy, _ = soft_value_iteration(feats @ w, nxt) grad = mu_E - expected_feature_counts(policy, start, nxt, feats) if np.linalg.norm(grad) < tol: return w, it, True w = w + lr * grad return w, max_iters, False def behavior_cloning_policy(demos, n_states): """ 对照组：表格型行为克隆。每个格子（按编号）取示范中最常见的动作；没见过的格子均匀随机。 注意它记住的是"第几号格子该做什么"——换了地图，格子编号的含义全变了。 """ counts = np.zeros((n_states, N_ACTIONS)) for traj in demos: for s, a in traj: counts[s, a] += 1 pi = np.full((n_states, N_ACTIONS), 1.0 / N_ACTIONS) seen = counts.sum(axis=1) > 0 pi[seen] = 0.0 pi[seen, counts[seen].argmax(axis=1)] = 1.0 return np.repeat(pi[None], T, axis=0) def render(grid, policy_t0): """画出 t = 0 时每个格子最可能的动作；水和终点用原字符标出。""" n_cols = len(grid[0]) lines = [] for r, row in enumerate(grid): cells = [] for c, ch in enumerate(row): s = r * n_cols + c cells.append(ch if ch in "GW" else ARROWS[int(policy_t0[s].argmax())]) lines.append(" " + " ".join(cells)) return "\n".join(lines) def main(): rng = np.random.default_rng(0) env_train = build_env(MAP_TRAIN) env_new = build_env(MAP_NEW) nxt, feats, start = env_train expert_policy, _ = soft_value_iteration(feats @ W_TRUE, nxt) # --------------------------------------------------------------------- print("=" * 78) print("演示 1：专家（在真实奖励下的最大熵策略）与示范数据") print("=" * 78) print(f" 真实奖励权重 w* = {W_TRUE}（终点、水、沙地）") print(" 专家在 t = 0 时最可能的动作：") print(render(MAP_TRAIN, expert_policy[0])) demos = [sample_trajectory(expert_policy, start, nxt, rng) for _ in range(100)] n_cols = len(MAP_TRAIN[0]) print(f"\n 收集 100 条示范（都从 A 出发）。一条示范轨迹（行, 列）：") print(" " + " → ".join(f"({s // n_cols},{s % n_cols})" for s, _ in demos[0][:12]) + " → ...") print(" （玻尔兹曼专家并不完美：大体上朝终点走、避开水和沙地，但路线带有随机性。）") # --------------------------------------------------------------------- print("\n" + "=" * 78) print("演示 2：最大熵 IRL 的学习过程") print("=" * 78) w_hat, log = maxent_irl(demos, env_train) print(" 迭代 w（终点, 水, 沙地） 对数似然 梯度范数（特征计数之差）") for it, w, ll, gn in log: print(f" {it:4d} {w} {ll:8.3f} {gn:.4f}") mu_E = np.mean([feats[[s for s, _ in traj]].sum(axis=0) for traj in demos], axis=0) policy_soft_hat, _ = soft_value_iteration(feats @ w_hat, nxt) print(f"\n 恢复的奖励权重 ŵ = {w_hat}，真实 w* = {W_TRUE}") print(f" 专家示范的平均特征计数 μ_E = {mu_E}") print(f" ŵ 下模型的期望特征计数 = {expected_feature_counts(policy_soft_hat, start, nxt, feats)}") print(" → 恢复的权重方向正确：终点好、水很差、沙地有点差；终点的权重几乎完全准确。") print(" 终点和沙地的计数已经匹配，水还没有（0.042 vs 0.03）：似然在\"水\"这个方向上非常平，") print(" 从第 50 轮到第 300 轮对数似然只变了 0.04，梯度上升走得很慢——300 轮时还没收敛。") print(" 水和沙地的惩罚都比真实值小。原因见演示 5：一是梯度上升还没收敛（只影响水），") print(" 二是这 100 条示范恰好比专家的平均水平更常踩沙地、更常下水。") # --------------------------------------------------------------------- print("\n" + "=" * 78) print("演示 3：迁移到一张新地图（没有任何示范）") print("=" * 78) n_states = nxt.shape[0] bc = behavior_cloning_policy(demos, n_states) policies = { "最优（用真实奖励规划）": lambda env: greedy_plan(env[1] @ W_TRUE, env[0]), "IRL（用恢复的奖励规划）": lambda env: greedy_plan(env[1] @ w_hat, env[0]), "行为克隆（照搬格子编号→动作）": lambda env: bc, "均匀随机": lambda env: np.full((T,) + env[0].shape, 1.0 / N_ACTIONS), } print(" 策略 训练地图（从 A） 新地图（从 B）") for name, make in policies.items(): print(f" {name:<28s} {true_return(make(env_train), env_train):8.2f} " f"{true_return(make(env_new), env_new):8.2f}") print("\n IRL 在新地图上规划出的策略（t = 0）：") print(render(MAP_NEW, greedy_plan(env_new[1] @ w_hat, env_new[0])[0])) print("\n 行为克隆在新地图上的\"策略\"（照搬训练地图的箭头）：") print(render(MAP_NEW, bc[0])) print(" → 在训练地图上，IRL 和行为克隆都达到了最优（数据足够多）。") print(" 换到新地图：行为克隆照搬\"这个格子该往哪走\"，从 B 出发一路向上、踩过沙地，") print(" 最后停在右上角不动——那正是**训练地图中终点的位置**！它学会的是\"去右上角\"，而不是\"去终点\"。") print(" IRL 恢复的是\"终点好、水差、沙地差\"，在新地图上重新规划，与最优完全一样好。") print(" **奖励比策略更容易迁移**——这是逆强化学习最重要的动机之一。") # --------------------------------------------------------------------- print("\n" + "=" * 78) print("演示 4：奖励的不可辨识性——加一个常数特征") print("=" * 78) feats_const = np.hstack([feats, np.ones((n_states, 1))]) w_c, _ = maxent_irl(demos, env_train, feats=feats_const) print(f" 特征改为 [终点, 水, 沙地, 常数 1]，恢复的 w = {w_c}") print(" → 常数特征的权重一直停在初始值 0，从未被更新。") print(" 原因：每条轨迹都恰好有 T 步，常数特征的计数永远是 T，专家和模型完全相同，梯度恒为 0。") print(" 给所有状态的奖励加同一个常数，所有轨迹的总奖励都加同一个数，轨迹的相对概率不变——") print(" 从行为中**原则上无法**区分这两种奖励。更一般地，许多不同的奖励都能解释同样的行为") print(" （讲义 4.2 节）。最大熵原则的作用，就是在这些候选中挑出一个\"不做多余假设\"的。") # --------------------------------------------------------------------- print("\n" + "=" * 78) print("演示 5：ŵ 为什么和 w* 不一样？——优化没收敛 + 有限样本") print("=" * 78) w_mle, n_it, _ = maxent_irl_until_converged(mu_E, env_train) pol_mle, V0_mle = soft_value_iteration(feats @ w_mle, nxt) print(f" (a) 同一批示范，梯度上升跑到收敛（{n_it} 轮）：ŵ = {w_mle}") print(f" 模型的期望特征计数 = {expected_feature_counts(pol_mle, start, nxt, feats)}，" f"对数似然 = {w_mle @ mu_E - V0_mle[start]:.4f}") print(" → 水的权重从 300 轮时的 -2.26 继续走到这里，三项计数完全匹配；对数似然却几乎没变——这就是\"平\"。") mu_true = expected_feature_counts(expert_policy, start, nxt, feats) print(f" (b) 专家真正的期望特征计数 E_w*[μ] = {mu_true}") print(f" 这 100 条示范的平均特征计数 μ_E = {mu_E}") print(" → 这批示范恰好比专家的平均水平更常踩沙地、更常下水。最大熵 IRL 忠实地匹配这批数据，") print(" 于是认为沙地和水\"没那么差\"。") w_inf, _, _ = maxent_irl_until_converged(mu_true, env_train) print(f" (c) 如果有无穷多条示范（直接用 E_w*[μ] 拟合）：ŵ = {w_inf}") print(" → 恢复了 w*。专家确实是\"线性奖励下的玻尔兹曼理性\"时，数据越多，最大熵 IRL 越准。") p_dry = prob_never_visit(expert_policy, start, nxt, feats[:, 1] > 0) print(f" (d) 专家的一条轨迹完全不下水的概率 = {p_dry:.4f}；100 条全都不下水的概率 = {p_dry ** 100:.3f}") ws, dry_mu = [], None # 继续使用 main 开头的随机数生成器 rng n_sets = 40 for _ in range(n_sets): reps = [sample_trajectory(expert_policy, start, nxt, rng) for _ in range(100)] mu_k = np.mean([feats[[s for s, _ in traj]].sum(axis=0) for traj in reps], axis=0) if mu_k[1] == 0: dry_mu = mu_k if dry_mu is None else dry_mu # 记下第一组"一次都没下水"的示范 continue ws.append(maxent_irl_until_converged(mu_k, env_train)[0]) ws = np.array(ws) print(f" 重新收集 {n_sets} 组、每组 100 条示范：{n_sets - len(ws)} 组一次都没下水；" f"其余 {len(ws)} 组各自拟合到收敛：") print(f" ŵ 的均值 = {ws.mean(axis=0)}，标准差 = {ws.std(axis=0)}") print(f" 水的权重最小 {ws[:, 1].min():.2f}，最大 {ws[:, 1].max():.2f}") print(" → 终点和沙地的估计平均下来接近真值；\"水\"的估计却散得很开。") print(" 专家很少做的事，只有极少的计数可以依据，对应的权重估计方差很大。") if dry_mu is not None: print(f" (e) 取一组一次都没下水的示范（μ_E = {dry_mu}）：") for iters in (300, 3000): w_d, _, ok = maxent_irl_until_converged(dry_mu, env_train, max_iters=iters) _, V0_d = soft_value_iteration(feats @ w_d, nxt) print(f" 跑 {iters:4d} 轮：w = {w_d}，对数似然 = {w_d @ dry_mu - V0_d[start]:.4f}，收敛：{ok}") print(" → 水的权重随迭代一直变负，对数似然一直缓慢上升：极大似然解不存在。") print(" 只要\"足够差\"，任何值都同样能解释这批数据——") print(" 这才是\"专家从不做的事，我们无法知道它有多糟\"的准确含义。") if __name__ == "__main__": main()