# PPO（近端策略最佳化，Schulman 等 2017）學平衡倒單擺：演員–評論家網路、GAE 優勢估計、裁剪的目標函數
import json
import math
import random
from pathlib import Path

import torch
from torch import nn

def step(s, a):
    """經典 CartPole（同 dqn_cartpole.py）"""
    x, xd, th, thd = s
    f = 10.0 if a == 1 else -10.0
    c, sn = math.cos(th), math.sin(th)
    tmp = (f + 0.05 * thd * thd * sn) / 1.1
    thacc = (9.8 * sn - c * tmp) / (0.5 * (4 / 3 - 0.1 * c * c / 1.1))
    xacc = tmp - 0.05 * thacc * c / 1.1
    s = (x + 0.02 * xd, xd + 0.02 * xacc, th + 0.02 * thd, thd + 0.02 * thacc)
    return s, abs(s[0]) > 2.4 or abs(s[2]) > 12 * math.pi / 180

random.seed(0); torch.manual_seed(0)
actor = nn.Sequential(nn.Linear(4, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 2))    # 策略：輸出兩個動作的 logit
critic = nn.Sequential(nn.Linear(4, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))   # 價值函數 V(s)
opt = torch.optim.Adam(list(actor.parameters()) + list(critic.parameters()), lr=3e-4)
GAMMA, LAM, CLIP, N = 0.99, 0.95, 0.2, 2048

s, T, finished, curve = tuple(random.uniform(-0.05, 0.05) for _ in range(4)), 0, [], []
for it in range(1, 31):
    S, A, LOGP, R, D = [], [], [], [], []
    for _ in range(N):                                                  # 1. 用目前的策略收集 2048 步經驗
        with torch.no_grad():
            dist = torch.distributions.Categorical(logits=actor(torch.tensor(s)))
            a = int(dist.sample())
        s2, done = step(s, a)
        T += 1
        trunc = T >= 500
        S.append(s); A.append(a); LOGP.append(float(dist.log_prob(torch.tensor(a)))); R.append(1.0); D.append(done)
        s = s2
        if done or trunc:
            finished.append(T)
            s, T = tuple(random.uniform(-0.05, 0.05) for _ in range(4)), 0
    St, At, LOGPt = torch.tensor(S), torch.tensor(A), torch.tensor(LOGP)
    with torch.no_grad():                                               # 2. GAE：用 V 估計每一步的優勢
        V = critic(St).squeeze(1)
        V_next = torch.cat([V[1:], critic(torch.tensor([s])).view(1)])
        adv, gae = torch.zeros(N), 0.0
        for t in range(N - 1, -1, -1):
            nonterm = 0.0 if D[t] else 1.0
            delta = R[t] + GAMMA * V_next[t] * nonterm - V[t]
            gae = delta + GAMMA * LAM * nonterm * gae
            adv[t] = gae
        ret = adv + V
        adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    for epoch in range(10):                                             # 3. 同一批資料重複用 10 次，但限制策略不能變太多
        for idx in torch.randperm(N).split(64):
            dist = torch.distributions.Categorical(logits=actor(St[idx]))
            ratio = torch.exp(dist.log_prob(At[idx]) - LOGPt[idx])    # 新策略／舊策略的機率比
            clipped = torch.clamp(ratio, 1 - CLIP, 1 + CLIP)
            loss_pi = -torch.min(ratio * adv[idx], clipped * adv[idx]).mean()
            loss_v = (critic(St[idx]).squeeze(1) - ret[idx]).pow(2).mean()
            opt.zero_grad(); (loss_pi + 0.5 * loss_v).backward(); opt.step()
    recent = finished[-20:] or [T]
    curve.append([it * N, sum(recent) / len(recent)])
    if it % 5 == 0:
        print(f"{it * N:6d} 步：最近 20 回合平均撐 {curve[-1][1]:5.1f} 步")

def evaluate(n=20, seed=1000):
    out = []
    for k in range(n):
        rng = random.Random(seed + k)
        s, T = tuple(rng.uniform(-0.05, 0.05) for _ in range(4)), 0
        while T < 500:
            with torch.no_grad():
                s, done = step(s, int(actor(torch.tensor(s)).argmax()))
            T += 1
            if done:
                break
        out.append(T)
    return out
ev = evaluate()
print(f"訓練後取機率最大的動作測 20 回合：平均 {sum(ev) / 20:.1f} 步，撐滿 500 步的有 {ev.count(500)} 回合")
Path(__file__).with_name("ppo_cartpole.json").write_text(json.dumps({"curve": curve}), encoding="utf-8")