# 懸崖行走（Sutton & Barto 例 6.6）：比較 Q-learning 與 SARSA
# 4×12 格，從左下角走到右下角，中間的底排是懸崖（掉下去 −100 並回到起點），每走一步 −1
import numpy as np

H, W = 4, 12
START, GOAL = (0, 0), (0, 11)                                          # (列, 行)，第 0 列是最下面
MOVES = [(1, 0), (0, 1), (-1, 0), (0, -1)]                              # 上、右、下、左

def step(s, a):
    r, c = min(max(s[0] + MOVES[a][0], 0), H - 1), min(max(s[1] + MOVES[a][1], 0), W - 1)
    if r == 0 and 0 < c < W - 1:
        return START, -100.0                                            # 掉下懸崖
    return (r, c), -1.0

def run(method, rng, episodes=500, alpha=0.2, eps=0.1):
    Q = np.zeros((H, W, 4))
    pick = lambda s: rng.integers(4) if rng.random() < eps else int(Q[s].argmax())   # ε-貪婪
    returns = []
    for _ in range(episodes):
        s, G = START, 0.0
        a = pick(s)
        while s != GOAL:
            s2, r = step(s, a)
            a2 = pick(s2)
            if s2 == GOAL:
                target = r
            elif method == "Q-learning":
                target = r + Q[s2].max()                                # 用「最好的下一步」更新（離策略）
            else:
                target = r + Q[s2][a2]                                  # 用「實際要走的下一步」更新（同策略）
            Q[s][a] += alpha * (target - Q[s][a])
            s, a, G = s2, a2, G + r
        returns.append(G)
    return Q, returns

def greedy_path(Q):
    s, cells = START, {START}
    for _ in range(60):
        s, _ = step(s, int(Q[s].argmax()))
        cells.add(s)
        if s == GOAL:
            break
    return cells

for method in ("Q-learning", "SARSA"):
    avg = np.mean([np.mean(run(method, np.random.default_rng(k))[1][-100:]) for k in range(50)])
    Q, _ = run(method, np.random.default_rng(0))
    cells = greedy_path(Q)
    print(f"{method}：最後 100 回合的平均報酬（50 次平均）{avg:.1f}；學到的貪婪路徑：")
    for r in range(H - 1, -1, -1):
        print("   " + "".join("S" if (r, c) == START else "G" if (r, c) == GOAL else "*" if (r, c) in cells
                              else "x" if r == 0 else "." for c in range(W)))