第 12 章 · 冰湖导航 MCTS 规划代码实验¶

# MCTS 先从起点看候选动作，再用模拟逐步判断哪个动作更值得走。
mcts_map = np.array([
    list("SFFF"),
    list("FHFF"),
    list("FFFH"),
    list("HFFG"),
])
n_rows, n_cols = mcts_map.shape
mcts_actions = {0: "左", 1: "下", 2: "右", 3: "上"}
mcts_arrows = {0: "←", 1: "↓", 2: "→", 3: "↑"}
action_delta = {0: (0, -1), 1: (1, 0), 2: (0, 1), 3: (-1, 0)}
side_actions = {0: (3, 1), 1: (0, 2), 2: (1, 3), 3: (2, 0)}
root_state = 0
goal_state = int(np.flatnonzero(mcts_map.reshape(-1) == "G")[0])

def state_to_rc(state):
    return divmod(int(state), n_cols)


def rc_to_state(row, col):
    return int(row * n_cols + col)


def state_name(state):
    row, col = state_to_rc(state)
    return f"{mcts_map[row, col]}({row},{col})"


def goal_distance(state):
    row, col = state_to_rc(state)
    goal_row, goal_col = state_to_rc(goal_state)
    return abs(row - goal_row) + abs(col - goal_col)


def move_state(state, action):
    row, col = state_to_rc(state)
    dr, dc = action_delta[action]
    next_row = min(max(row + dr, 0), n_rows - 1)
    next_col = min(max(col + dc, 0), n_cols - 1)
    return rc_to_state(next_row, next_col)


def mcts_step_model(state, action):
    tile = mcts_map[state_to_rc(state)]
    if tile in {"H", "G"}:
        return [(1.0, state, 0.0, True, "已结束")]

    candidates = [(action, 0.76, "按计划")] + [(a, 0.12, "滑向侧边") for a in side_actions[action]]
    outcomes = []
    old_dist = goal_distance(state)
    for actual_action, prob, cause in candidates:
        next_state = move_state(state, actual_action)
        next_tile = mcts_map[state_to_rc(next_state)]
        progress = 0.10 * (old_dist - goal_distance(next_state))
        reward = -0.02 + progress
        done = next_tile in {"H", "G"}
        if next_tile == "G":
            reward += 1.2
        elif next_tile == "H":
            reward -= 0.35
        outcomes.append((prob, next_state, reward, done, cause))
    return outcomes


def action_candidate_table(state):
    rows = []
    for action in mcts_actions:
        outcomes = mcts_step_model(state, action)
        expected_reward = sum(prob * reward for prob, _, reward, _, _ in outcomes)
        hole_risk = sum(prob for prob, next_state, _, _, _ in outcomes if mcts_map[state_to_rc(next_state)] == "H")
        goal_chance = sum(prob for prob, next_state, _, _, _ in outcomes if mcts_map[state_to_rc(next_state)] == "G")
        rows.append({
            "当前状态": state_name(state),
            "候选动作": mcts_actions[action],
            "一步期望得分": expected_reward,
            "掉洞概率": hole_risk,
            "到达概率": goal_chance,
            "可能到达": " / ".join(f"{prob:.2f}→{state_name(next_state)}" for prob, next_state, _, _, _ in outcomes),
        })
    return pd.DataFrame(rows)


def transition_detail_table(state):
    rows = []
    for action in mcts_actions:
        for prob, next_state, reward, done, cause in mcts_step_model(state, action):
            rows.append({
                "当前状态": state_name(state),
                "候选动作": mcts_actions[action],
                "实际结果": cause,
                "可能到达": state_name(next_state),
                "概率": prob,
                "即时得分": reward,
                "是否结束": done,
            })
    return pd.DataFrame(rows)


def draw_mcts_lake(ax, title, highlight_state=None, path_states=None):
    tile_color = {"S": "#dbeafe", "F": "#ecfeff", "H": "#fee2e2", "G": "#dcfce7"}
    path_states = set(path_states or [])
    for row in range(n_rows):
        for col in range(n_cols):
            state = rc_to_state(row, col)
            tile = mcts_map[row, col]
            ax.add_patch(plt.Rectangle((col - 0.5, row - 0.5), 1, 1, color=tile_color[tile], ec="#94a3b8"))
            ax.text(col, row, tile, ha="center", va="center", fontsize=15, fontweight="bold", color="#0f172a")
            if state in path_states:
                ax.add_patch(plt.Rectangle((col - 0.42, row - 0.42), 0.84, 0.84, fill=False, ec="#f97316", lw=2.4))
    if highlight_state is not None:
        row, col = state_to_rc(highlight_state)
        ax.scatter([col], [row], s=260, facecolors="none", edgecolors="#2563eb", linewidths=2.6)
    ax.set_title(title, loc="left", fontweight="bold")
    ax.set_xlim(-0.5, n_cols - 0.5)
    ax.set_ylim(n_rows - 0.5, -0.5)
    ax.set_xticks(range(n_cols))
    ax.set_yticks(range(n_rows))
    ax.grid(True, color="#cbd5e1", linewidth=0.8)

fig, ax = plt.subplots(figsize=(4.8, 4.4))
draw_mcts_lake(ax, "冰湖规划起点", highlight_state=root_state)
plt.tight_layout()
plt.show()

display(pd.DataFrame([
    {"得分来源": "每走一步", "作用": "鼓励更短路线", "数值": "-0.02"},
    {"得分来源": "靠近终点", "作用": "让中间状态也能比较好坏", "数值": "+0.10 × 距离缩短"},
    {"得分来源": "到达终点", "作用": "明确最终目标", "数值": "+1.20"},
    {"得分来源": "掉入洞中", "作用": "惩罚危险路线", "数值": "-0.35"},
]))
display(action_candidate_table(root_state).round(3))
display(transition_detail_table(root_state).round(3))

# 有限深度节点为 (位置, 剩余步数)，撞墙回到同一格也不会混用不同时间的价值。
MCTS_HORIZON = 18
MCTS_GAMMA = 0.99
MCTS_BUDGET = 1000
model_outcomes = {
    (state, action): mcts_step_model(state, action)
    for state in range(n_rows * n_cols) for action in mcts_actions
}
rollout_cumulative = {}
for state in range(n_rows * n_cols):
    scores = np.array([
        sum(prob * reward for prob, _, reward, _, _ in model_outcomes[state, action])
        for action in mcts_actions
    ])
    weights = np.exp((scores - scores.max()) * 5.0)
    rollout_cumulative[state] = np.cumsum(weights / weights.sum())
    rollout_cumulative[state][-1] = 1.0


def sample_model_step(state, action, rng):
    draw = rng.random()
    cumulative = 0.0
    outcomes = model_outcomes[state, action]
    for prob, next_state, reward, done, _ in outcomes:
        cumulative += prob
        if draw < cumulative:
            return next_state, reward, done
    return outcomes[-1][1:4]


def rollout(state, remaining, rng, gamma):
    total, discount, steps = 0.0, 1.0, 0
    while steps < remaining and mcts_map[state_to_rc(state)] not in {"H", "G"}:
        # 仅用于模拟后半段的启发式策略，不是最终执行动作的选择规则。
        action = int(np.searchsorted(rollout_cumulative[state], rng.random()))
        state, reward, done = sample_model_step(state, action, rng)
        total += discount * reward
        discount *= gamma
        steps += 1
        if done:
            break
    return total, steps


def backpropagate(path, tail_return, gamma, n_state, n_action, w_action):
    value = tail_return
    for key, action, reward in reversed(path):
        value = reward + gamma * value
        n_state[key] += 1
        n_action[key, action] += 1
        w_action[key, action] += value
    return value


def action_statistics(key, n_state, n_action, w_action, exploration=1.4):
    rows = []
    for action in mcts_actions:
        visits = n_action.get((key, action), 0)
        mean = w_action.get((key, action), 0.0) / visits if visits else np.nan
        bonus = exploration * math.sqrt(math.log(max(1, n_state.get(key, 0))) / visits) if visits else np.inf
        rows.append({"action_id": action, "action": mcts_actions[action], "visits": visits,
                     "mean_value": mean, "explore": bonus, "UCT": mean + bonus if visits else np.inf})
    # UCT 只决定下一次模拟探索谁；实际执行按访问次数，平手时按平均回报。
    return pd.DataFrame(rows).sort_values(
        ["visits", "mean_value", "action_id"], ascending=[False, False, True]
    ).reset_index(drop=True)


def mcts_plan(start_state, remaining, simulations=1000, seed=12, gamma=0.99, exploration=1.4, trace_count=0):
    if remaining < 0 or simulations < 1 or not 0 <= gamma <= 1 or exploration < 0:
        raise ValueError("步数非负、模拟次数为正，折扣在 [0,1] 内，探索系数非负。")
    rng = np.random.default_rng(seed)
    n_state, n_action, w_action = defaultdict(int), defaultdict(int), defaultdict(float)
    traces, snapshots = [], []
    root_key = (int(start_state), remaining)
    result = {"root_key": root_key, "N_state": n_state, "N_action": n_action,
              "W_action": w_action, "traces": traces, "snapshots": snapshots,
              "action": None, "table": pd.DataFrame()}
    if remaining == 0 or mcts_map[state_to_rc(start_state)] in {"H", "G"}:
        return result

    for simulation in range(1, simulations + 1):
        state, left, path = int(start_state), remaining, []
        tail_return, rollout_steps, expanded = 0.0, 0, False
        while left > 0:
            key = (state, left)
            unvisited = [a for a in mcts_actions if n_action[key, a] == 0]
            if unvisited:
                action = int(rng.choice(unvisited))
                expanded = True
            else:
                # Selection：沿已访问动作按 UCT 向下走，随机转移按环境概率抽样。
                log_visits = math.log(n_state[key])
                action = max(mcts_actions, key=lambda a:
                    w_action[key, a] / n_action[key, a]
                    + exploration * math.sqrt(log_visits / n_action[key, a]))
            next_state, reward, done = sample_model_step(state, action, rng)
            path.append((key, action, reward))
            state, left = next_state, left - 1
            if done:
                break
            if expanded:
                # Expansion 后只做剩余步数的 rollout，本次不再扩展第二个新动作。
                tail_return, rollout_steps = rollout(state, left, rng, gamma)
                break

        assert len(path) + rollout_steps <= remaining
        assert len({key for key, _, _ in path}) == len(path)
        root_return = backpropagate(path, tail_return, gamma, n_state, n_action, w_action)
        if simulation <= trace_count:
            traces.append({
                "模拟轮次": simulation,
                "选择与扩展路径": " → ".join(f"{state_name(k[0])}[剩{k[1]}步]/{mcts_actions[a]}" for k, a, _ in path),
                "扩展新动作数": int(expanded), "后续模拟步数": rollout_steps,
                "总步数": len(path) + rollout_steps, "根节点回报": root_return,
            })
        if trace_count and (simulation % max(1, simulations // 5) == 0 or simulation == simulations):
            row = {"模拟次数": simulation}
            row.update({f"{name}访问次数": n_action[root_key, a] for a, name in mcts_actions.items()})
            snapshots.append(row)

    table = action_statistics(root_key, n_state, n_action, w_action, exploration)
    result.update(action=int(table.loc[0, "action_id"]), table=table)
    assert n_state[root_key] == simulations
    for key, count in n_state.items():
        assert sum(n_action[key, a] for a in mcts_actions) == count
    return result


def show_action_table(plan):
    display(plan["table"].rename(columns={
        "action_id": "动作编号", "action": "动作", "visits": "访问次数",
        "mean_value": "平均后续回报", "explore": "探索项", "UCT": "UCT（仅供模拟选择）",
    }).round(4))


root_plan = mcts_plan(root_state, MCTS_HORIZON, MCTS_BUDGET, seed=12, trace_count=8)
mcts_root_df = root_plan["table"]
display(pd.DataFrame(root_plan["traces"]).round(4))
display(pd.DataFrame(root_plan["snapshots"]))
show_action_table(root_plan)
print("起点推荐动作:", mcts_actions[root_plan["action"]], "（按访问次数排序）")

# 执行与搜索使用独立随机数。真实抽样保留掉洞、撞墙等结果，不筛选安全结果。
environment_rng = np.random.default_rng(2026)
state, remaining = root_state, MCTS_HORIZON
episode_states, episode_rows, episode_plans = [state], [], []
episode_return, discount = 0.0, 1.0
for step in range(MCTS_HORIZON):
    plan = root_plan if step == 0 else mcts_plan(state, remaining, MCTS_BUDGET, seed=12 + step)
    episode_plans.append(plan)
    action = plan["action"]
    next_state, reward, done = sample_model_step(state, action, environment_rng)
    episode_return += discount * reward
    episode_rows.append({
        "步次": step + 1, "当前状态": state_name(state), "剩余步数": remaining,
        "重新规划的动作": mcts_actions[action], "动作访问次数": int(plan["table"].loc[0, "visits"]),
        "估计后续回报": plan["table"].loc[0, "mean_value"], "实际下一状态": state_name(next_state),
        "即时奖励": reward, "累计折扣回报": episode_return, "是否结束": done,
    })
    state, remaining, discount = next_state, remaining - 1, discount * MCTS_GAMMA
    episode_states.append(state)
    if done:
        break

display(pd.DataFrame(episode_rows).round(4))
next_plan = episode_plans[1] if len(episode_plans) > 1 else None
if next_plan is not None:
    print("实际下一状态重新搜索:", state_name(next_plan["root_key"][0]), "剩余步数:", next_plan["root_key"][1])
    show_action_table(next_plan)
middle_plan = episode_plans[len(episode_plans) // 2]
print("中间状态重新搜索:", state_name(middle_plan["root_key"][0]), "剩余步数:", middle_plan["root_key"][1])
show_action_table(middle_plan)
print("本次抽样到达终点:", state == goal_state, "折扣回报:", round(episode_return, 4))

fig, axes = plt.subplots(1, 2, figsize=(10.8, 4.7))
draw_mcts_lake(axes[0], "逐步重新规划的实际路线", highlight_state=state, path_states=episode_states)
for before, after in zip(episode_states, episode_states[1:]):
    if before != after:
        r0, c0 = state_to_rc(before)
        r1, c1 = state_to_rc(after)
        axes[0].annotate("", xy=(c1, r1), xytext=(c0, r0), arrowprops={"arrowstyle": "->", "lw": 1.8, "color": "#2563eb"})
if next_plan is not None:
    next_table = next_plan["table"]
    bars = axes[1].bar(next_table["action"], next_table["mean_value"], color="#2563eb")
    for bar, visits in zip(bars, next_table["visits"]):
        axes[1].annotate(f"n={visits}", (bar.get_x() + bar.get_width() / 2, bar.get_height()),
                         ha="center", va="bottom", xytext=(0, 5), textcoords="offset points")
    axes[1].set_title(f"下一状态 {state_name(next_plan['root_key'][0])}，剩 {next_plan['root_key'][1]} 步", loc="left", fontweight="bold")
    axes[1].set_ylabel("重新搜索的平均后续回报（按访问次数排序）")
    axes[1].margins(y=0.25)
    axes[1].axhline(0, color="#94a3b8", linewidth=0.9)
    axes[1].grid(True, axis="y", color="#e2e8f0", linewidth=0.8)
else:
    axes[1].text(0.5, 0.5, "已到终止状态，不再规划", ha="center", va="center")
    axes[1].set_axis_off()
plt.tight_layout()
plt.show()

# 固定剩余步数切片，不把不同规划时长的价值混成一张图。
fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.7))
x = np.arange(len(mcts_root_df))
axes[0].bar(x - 0.18, mcts_root_df["mean_value"], width=0.36, color="#2563eb", label="平均回报")
axes[0].bar(x + 0.18, mcts_root_df["UCT"], width=0.36, color="#f97316", label="UCT 总分")
axes[0].set_xticks(x, mcts_root_df["action"])
axes[0].set_title("起点：UCT 用于探索，执行按访问次数", loc="left", fontweight="bold")
axes[0].set_ylabel("评分")
axes[0].axhline(0, color="#94a3b8", linewidth=0.9)
axes[0].grid(True, axis="y", color="#e2e8f0", linewidth=0.8)
axes[0].legend()

state_count = n_rows * n_cols
slice_remaining = MCTS_HORIZON - 2
value_grid = np.full(state_count, np.nan)
policy_grid = np.full(state_count, -1)
for state in range(state_count):
    key = (state, slice_remaining)
    if root_plan["N_state"].get(key, 0):
        table = action_statistics(key, root_plan["N_state"], root_plan["N_action"], root_plan["W_action"])
        policy_grid[state] = int(table.loc[0, "action_id"])
        value_grid[state] = float(table.loc[0, "mean_value"])
value_grid = value_grid.reshape(n_rows, n_cols)
policy_grid = policy_grid.reshape(n_rows, n_cols)

im = axes[1].imshow(np.ma.masked_invalid(value_grid), cmap="RdYlGn", vmin=-0.5, vmax=1.5)
for state in range(state_count):
    r, c = state_to_rc(state)
    tile = mcts_map[r, c]
    arrow = "" if tile in {"H", "G"} or policy_grid[r, c] < 0 else mcts_arrows[int(policy_grid[r, c])]
    label = "终止" if tile in {"H", "G"} else ("未访问" if np.isnan(value_grid[r, c]) else f"{value_grid[r, c]:.2f}")
    axes[1].text(c, r, f"{tile}\n{label}\n{arrow}", ha="center", va="center", color="#0f172a", fontweight="bold")
axes[1].set_title(f"剩 {slice_remaining} 步：最多访问动作的平均回报", loc="left", fontweight="bold")
axes[1].set_xticks(range(n_cols))
axes[1].set_yticks(range(n_rows))
fig.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04)
plt.tight_layout()
plt.show()

# 在同一组随机环境中比较两种策略的成功率与回报。
def evaluate_policy(use_mcts, episodes=96, simulations=400):
    rows = []
    for episode in range(episodes):
        # 两种策略共享每局的环境种子，规划随机数不会消耗真实转移的随机数。
        env_rng = np.random.default_rng(10000 + episode)
        action_rng = np.random.default_rng(20000 + episode)
        state, total, discount = root_state, 0.0, 1.0
        for step in range(MCTS_HORIZON):
            if use_mcts:
                plan = mcts_plan(state, MCTS_HORIZON - step, simulations,
                                 seed=30000 + episode * MCTS_HORIZON + step)
                action = plan["action"]
            else:
                action = int(action_rng.integers(4))
            state, reward, done = sample_model_step(state, action, env_rng)
            total += discount * reward
            discount *= MCTS_GAMMA
            if done:
                break
        rows.append({"成功": state == goal_state, "折扣回报": total, "步数": step + 1})
    return pd.DataFrame(rows)


random_evaluation = evaluate_policy(False)
mcts_evaluation = evaluate_policy(True)
evaluation_summary = pd.DataFrame([
    {"策略": name, "局数": len(data), "成功局数": int(data["成功"].sum()),
     "成功率": data["成功"].mean(), "平均折扣回报": data["折扣回报"].mean(), "平均步数": data["步数"].mean()}
    for name, data in [("均匀随机", random_evaluation), ("每步 MCTS（400 次模拟）", mcts_evaluation)]
])
display(evaluation_summary.round(4))
print("成功率是固定地图和种子下的有限样本结果，不保证每局成功；距离奖励也不等于成功概率。")