# REINFORCE（策略梯度）：只用 4 個權值的線性策略，學平衡倒單擺；並和由模型算出的 LQR 增益比較
import math
import numpy as np
from scipy.linalg import expm, solve_discrete_are

def step(s, f):
    """經典 CartPole：台車 1 kg、擺桿 0.1 kg（半長 0.5 m）、每步 0.02 s；推力 f = ±10 N"""
    x, xd, th, thd = s
    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 = np.array([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

SCALE = np.array([1, 1, 10, 1])                                        # 擺角放大 10 倍，讓四個輸入大小相近

def train(seed, episodes=200, lr=0.03, gamma=0.99):
    rng = np.random.default_rng(seed)
    w, baseline, lens = np.zeros(4), 0.0, []
    for ep in range(episodes):
        s, S, A = rng.uniform(-0.05, 0.05, 4), [], []
        for t in range(500):
            p = 1 / (1 + math.exp(-np.clip(w @ (s * SCALE), -30, 30)))  # 往右推的機率 π(右 | s)
            a = int(rng.random() < p)
            S.append(s * SCALE); A.append(a)
            s, done = step(s, 10.0 if a else -10.0)
            if done:
                break
        G = np.zeros(len(A)); acc = 0.0
        for t in range(len(A) - 1, -1, -1):                            # 每一步之後的折扣報酬總和
            acc = 1 + gamma * acc; G[t] = acc
        S, A = np.array(S), np.array(A)
        P = 1 / (1 + np.exp(-np.clip(S @ w, -30, 30)))
        w += lr * ((A - P) * (G - baseline)) @ S / len(A)              # ∇log π = (a − p)·φ(s)，乘上（報酬 − 基準線）
        baseline = 0.9 * baseline + 0.1 * G.mean()
        lens.append(len(A))
    return w, lens

for seed in (1, 2, 3):
    w, lens = train(seed)
    first = next((i + 1 for i, n in enumerate(lens) if n == 500), None)
    print(f"種子 {seed}：第 {first} 回合首次撐滿 500 步，最後 20 回合平均 {np.mean(lens[-20:]):.0f} 步，「往右推」logit 的係數 {np.round(w * SCALE, 2)}")

# 對照：在直立點線性化，用 LQR 算回授增益（推力 u = −Kx）
den = 0.5 * (4 / 3 - 0.1 / 1.1)
a_th, b_th = 9.8 / den, -1 / (1.1 * den)
Ac = np.array([[0, 1, 0, 0], [0, 0, -0.05 * a_th / 1.1, 0], [0, 0, 0, 1], [0, 0, a_th, 0]])
Bc = np.array([[0], [1 / 1.1 - 0.05 * b_th / 1.1], [0], [b_th]])
M = expm(np.block([[Ac, Bc], [np.zeros((1, 5))]]) * 0.02)
A, B = M[:4, :4], M[:4, 4:]
P = solve_discrete_are(A, B, np.diag([1, 1, 10, 1]), np.array([[0.1]]))
K = np.linalg.solve(0.1 + B.T @ P @ B, B.T @ P @ A)[0]
print("LQR：往右推的程度 −Kx 的係數", np.round(-K, 2), "（依序為 x、ẋ、θ、θ˙）")