"""
第 13 章配套代码：REINFORCE（Monte-Carlo policy gradient, Williams 1992）。
纯 numpy，环境是一个 1xN 走廊：从最左走到最右得 +1。
策略是 softmax(theta)，用策略梯度定理 ∇J = E[∇logπ(a|s) * G] 直接对参数上升。
对比"有 baseline 降方差"前后的学习曲线。
Runnable with: numpy only.  python3 13_reinforce_policy_gradient.py
"""
import numpy as np

N = 7                 # 走廊长度，状态 0..N-1，目标 N-1
ACTIONS = [-1, +1]    # 左 / 右
GAMMA = 0.99


def softmax(z):
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()


def run_episode(theta, rng, max_steps=50):
    """用当前策略走一局，返回轨迹 (states, actions, rewards)。"""
    s = 0
    S, A, R = [], [], []
    for _ in range(max_steps):
        probs = softmax(theta[s])          # 每个状态一组动作 logits
        a = rng.choice(2, p=probs)
        s_next = min(N - 1, max(0, s + ACTIONS[a]))
        r = 1.0 if s_next == N - 1 else -0.05
        S.append(s); A.append(a); R.append(r)
        s = s_next
        if s == N - 1:
            break
    return S, A, R


def reinforce(episodes=1500, lr=0.1, use_baseline=True, seed=0):
    rng = np.random.default_rng(seed)
    theta = np.zeros((N, 2))               # 策略参数: 每状态2个动作 logit
    baseline = np.zeros(N)                 # 状态价值的滑动估计（baseline）
    curve = []
    for ep in range(episodes):
        S, A, R = run_episode(theta, rng)
        # 计算每步的回报 G_t（从 t 到结束的折扣累计）
        G = np.zeros(len(R))
        g = 0.0
        for t in reversed(range(len(R))):
            g = R[t] + GAMMA * g
            G[t] = g
        # 策略梯度上升: theta[s,a] += lr * (G_t - b(s)) * ∇logπ
        for t, (s, a) in enumerate(zip(S, A)):
            adv = G[t] - (baseline[s] if use_baseline else 0.0)
            probs = softmax(theta[s])
            grad = -probs
            grad[a] += 1.0                 # ∇logπ(a|s) = e_a - softmax
            theta[s] += lr * adv * grad
            if use_baseline:
                baseline[s] += 0.05 * (G[t] - baseline[s])
        curve.append(sum(R))
    return theta, curve


def report():
    print("=== REINFORCE (policy gradient) on 1x%d 走廊 ===" % N)
    for use_b in (False, True):
        theta, curve = reinforce(use_baseline=use_b)
        first = np.mean(curve[:100]); last = np.mean(curve[-100:])
        greedy = "".join("→" if softmax(theta[s])[1] > 0.5 else "←" for s in range(N - 1))
        tag = "有 baseline" if use_b else "无 baseline"
        print(f"\n[{tag}] 平均回报 前100局={first:+.3f} -> 末100局={last:+.3f}")
        print(f"       学到的贪婪策略（状态0..{N-2}）: {greedy}  (全 → = 一路向右最优)")
    print("\n结论: 策略梯度直接优化 softmax 策略；baseline 把 G_t 减去 b(s) 降低方差、更快收敛。")


if __name__ == "__main__":
    report()