다섯 칸 복도에서
LLM 사후학습까지

python3 phase01_return.py          # 본문: Python 3.9 이상, 설치할 것 없음
python3 phase11_grpo.py

# 선택 실습 (extras/)
python3 -m pip install torch==2.8.0 "gymnasium[classic-control]==1.0.0"
python3 extras/autograd_check.py   # 손으로 유도한 gradient == PyTorch autograd
python3 -m pip install trl==0.24.0 transformers datasets
python3 extras/trl_grpo.py         # GPU 권장 (Colab T4 가능)

class Corridor(_Base):
    """5칸 복도. 상태 = 현재 칸 번호(int), 행동 = LEFT(0) 또는 RIGHT(1)."""

    def __init__(self, n=5, start=2, max_steps=20):
        self.n, self.start, self.max_steps = n, start, max_steps
        self.hole, self.goal = 0, n - 1           # 양 끝이 종료 칸
        if spaces is not None:                      # gymnasium이 있을 때만
            self.observation_space = spaces.Discrete(n)
            self.action_space = spaces.Discrete(2)

    def reset(self, seed=None, options=None):
        if spaces is not None:
            super().reset(seed=seed)
        self.pos, self.t = self.start, 0
        return self.pos, {}                         # (observation, info)

    def transition(self, s, a):
        """환경의 규칙(모델) P와 R: (s, a) → (s′, r, 종료 여부). Phase 2의 Bellman 계산이 쓴다."""
        assert a in (LEFT, RIGHT), a
        s2 = s + 1 if a == RIGHT else s - 1         # 전이 P: 결정적
        reward = {self.hole: -1.0, self.goal: 1.0}.get(s2, 0.0)
        return s2, reward, s2 in (self.hole, self.goal)

    def step(self, action):
        self.pos, reward, terminated = self.transition(self.pos, action)
        self.t += 1
        truncated = not terminated and self.t >= self.max_steps
        return self.pos, reward, terminated, truncated, {}

    def render(self):
        cells = ['H', '.', '.', '.', 'G']
        cells[self.pos] = '@'
        return ' '.join(cells)

def random_policy(obs, rng):
    """관측을 보지 않고 ←, →를 반반 고르는 정책."""
    return rng.choice([LEFT, RIGHT])


def run_episode(env, policy, seed=0):
    rng = random.Random(seed)               # 정책의 무작위성 (재현용 seed)
    obs, info = env.reset(seed=seed)
    trajectory = []                         # (s_t, a_t, r_{t+1}, s_{t+1}) 목록
    while True:
        action = policy(obs, rng)
        next_obs, reward, terminated, truncated, info = env.step(action)
        trajectory.append((obs, action, reward, next_obs))
        obs = next_obs
        if terminated or truncated:
            return trajectory, terminated

def step_api_demo():
    env = Corridor()
    obs, info = env.reset()
    assert (obs, info) == (2, {})
    # → 한 번: 3번 칸 도착, 아직 종료 아님
    assert env.step(RIGHT) == (3, 0.0, False, False, {})
    # → 한 번 더: 보물 칸 도착, 보상 +1, 종료
    assert env.step(RIGHT) == (4, 1.0, True, False, {})

    # ←→를 10번 반복하면 20걸음째 시간 제한에 걸린다 (truncated)
    env.reset()
    for i in range(20):
        obs, r, terminated, truncated, _ = env.step(LEFT if i % 2 == 0 else RIGHT)
    assert (obs, r, terminated, truncated) == (2, 0.0, False, True)

step() API 점검 통과
seed=0: 2→ 3→ 4 | 보상 [0.0, 1.0] | 종료
seed=1: 2← 1← 0 | 보상 [0.0, -1.0] | 종료
seed=2: 2← 1← 0 | 보상 [0.0, -1.0] | 종료

# pip install "gymnasium[box2d]" stable-baselines3
import gymnasium as gym
from stable_baselines3 import PPO

env = gym.make("LunarLander-v2")   # gymnasium 1.0 이상에서는 "LunarLander-v3"
model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=1_000_000)

def discounted_return(rewards, gamma=0.9):
    """식 (1.1): G_0 = R_1 + γR_2 + γ²R_3 + ...

    rewards[k]가 R_{k+1}이다. 즉 k번째 행동 직후에 받은 보상.
    """
    return sum(gamma ** k * r for k, r in enumerate(rewards))

def returns_backward(rewards, gamma=0.9):
    """식 (1.2): G_t = R_{t+1} + γ·G_{t+1},  종료 뒤 G_T = 0.

    뒤에서부터 한 번만 훑으면 모든 시점의 G_t가 나온다.
    """
    G, out = 0.0, []
    for r in reversed(rewards):
        G = r + gamma * G
        out.append(G)
    return out[::-1]                        # [G_0, G_1, ..., G_{T-1}]

def rewards_of(actions):
    """행동 목록을 복도에서 실행하고 보상 목록 [R_1, ..., R_T]를 돌려준다."""
    env = Corridor()
    env.reset()
    rewards = []
    for a in actions:
        _, r, terminated, truncated, _ = env.step(a)
        rewards.append(r)
    assert terminated, '예제 궤적은 종료 칸에서 끝나야 한다'
    return rewards


TRAJECTORIES = {                            # 이름: (행동 목록, 손계산 G_0)
    '→→':   ([RIGHT, RIGHT], 0.9),                  # 0 + 0.9·1
    '←→→→': ([LEFT, RIGHT, RIGHT, RIGHT], 0.729),   # 0.9³·1
    '←←':   ([LEFT, LEFT], -0.9),                   # 0 + 0.9·(−1)
}


def hand_check():
    for name, (actions, hand_G0) in TRAJECTORIES.items():
        rewards = rewards_of(actions)
        G0 = discounted_return(rewards)
        assert abs(G0 - hand_G0) < 1e-12, (name, G0, hand_G0)
        assert abs(returns_backward(rewards)[0] - G0) < 1e-12   # 두 식은 같은 값
        print(f'{name:5s} 보상 {rewards} → G_0 = {G0:.3f} (손계산 {hand_G0})')

    # 우회 궤적의 모든 시점 G_t: 표 1.2와 같은 값
    table = returns_backward(rewards_of(TRAJECTORIES['←→→→'][0]))
    assert all(abs(a - b) < 1e-12 for a, b in zip(table, [0.729, 0.81, 0.9, 1.0]))
    print('←→→→ 의 G_0..G_3 =', [round(g, 3) for g in table])

def gamma_effect():
    short = rewards_of([RIGHT, RIGHT])
    detour = rewards_of([LEFT, RIGHT, RIGHT, RIGHT])
    for gamma in (1.0, 0.9, 0.5):
        a, b = discounted_return(short, gamma), discounted_return(detour, gamma)
        print(f'γ={gamma}: 최단 {a:.3f}  우회 {b:.3f}  차이 {a - b:.3f}')

def estimate_J(p_right, episodes=20000, gamma=0.9, seed=0):
    """식 (1.3): J(π) = E[G_0] ≈ (1/N) Σ_i G_0^(i)   (몬테카를로 평균)

    정책 π: 매 칸에서 확률 p_right로 →, 나머지는 ←.
    """
    rng, env, total = random.Random(seed), Corridor(), 0.0
    for _ in range(episodes):
        env.reset()
        rewards, done = [], False
        while not done:
            a = RIGHT if rng.random() < p_right else LEFT
            _, r, terminated, truncated, _ = env.step(a)
            rewards.append(r)
            done = terminated or truncated
        total += discounted_return(rewards, gamma)
    return total / episodes

→→    보상 [0.0, 1.0] → G_0 = 0.900 (손계산 0.9)
←→→→  보상 [0.0, 0.0, 0.0, 1.0] → G_0 = 0.729 (손계산 0.729)
←←    보상 [0.0, -1.0] → G_0 = -0.900 (손계산 -0.9)
←→→→ 의 G_0..G_3 = [0.729, 0.81, 0.9, 1.0]
γ=1.0: 최단 1.000  우회 1.000  차이 0.000
γ=0.9: 최단 0.900  우회 0.729  차이 0.171
γ=0.5: 최단 0.500  우회 0.125  차이 0.375
P(→)=0.5: J ≈ +0.0009
P(→)=0.7: J ≈ +0.5469
P(→)=1.0: J ≈ +0.9000

def q_from_v(V, s, a, gamma=0.9):
    """식 (2.4): Q(s,a) = R(s,a) + γ·(1 − done)·V(s′)   (전이가 결정적인 복도)"""
    s2, r, done = ENV.transition(s, a)
    return r + gamma * (0.0 if done else V[s2])


def bellman_backup(V, s, pi, gamma=0.9):
    """식 (2.3): V(s) = Σ_a π(a|s) · Q(s,a)"""
    return sum(pi[a] * q_from_v(V, s, a, gamma) for a in (LEFT, RIGHT))

def evaluate_policy(p_right, gamma=0.9, sweeps=200):
    """반복 정책 평가: 모든 칸에 식 (2.3)을 동시에 적용하는 sweep을 되풀이한다."""
    pi = policy(p_right)
    V = {s: 0.0 for s in NONTERMINAL}       # 처음엔 아무것도 모른다: V = 0
    history = [dict(V)]
    for _ in range(sweeps):
        V = {s: bellman_backup(V, s, pi, gamma) for s in NONTERMINAL}  # 이전 V로 한꺼번에
        history.append(dict(V))
    return history

def exact_values(p_right, gamma=0.9):
    """식 (2.3)은 미지수 3개짜리 연립일차방정식이다. 손으로 푼 결과를 그대로 옮긴다.

    a = γ(1−p), b = γp 라 두면
      V1 = −(1−p) + b·V2
      V3 =  p     + a·V2
      V2 = a·V1 + b·V3  →  V2 = (−a(1−p) + b·p) / (1 − 2ab)
    """
    p = p_right
    a, b = gamma * (1 - p), gamma * p
    V2 = (-a * (1 - p) + b * p) / (1 - 2 * a * b)
    return {1: -(1 - p) + b * V2, 2: V2, 3: p + a * V2}

def mc_update(V, states, rewards, alpha=0.1, gamma=0.9):
    """식 (2.5) Monte Carlo: 에피소드가 끝난 뒤 V(S_t) ← V(S_t) + α[G_t − V(S_t)]"""
    G = 0.0
    for t in reversed(range(len(rewards))):     # Phase 1의 식 (1.2)로 G_t 계산
        G = rewards[t] + gamma * G
        V[states[t]] += alpha * (G - V[states[t]])
    return V


def td_update(V, s, r, s2, terminated, alpha=0.1, gamma=0.9):
    """식 (2.6) TD(0): 한 걸음마다 V(S_t) ← V(S_t) + α[R + γ(1−terminated)V(S_{t+1}) − V(S_t)]"""
    target = r + gamma * (0.0 if terminated else V[s2])   # truncated는 여기 들어가지 않는다
    delta = target - V[s]                                  # TD 오차 δ
    V[s] += alpha * delta
    return delta

def td_evaluate(p_right, episodes=10000, alpha=0.01, gamma=0.9, seed=0):
    """정책 π를 실제로 돌리면서 TD(0)만으로 V를 배운다. 환경 모델(transition)은 쓰지 않는다."""
    import random
    rng, env = random.Random(seed), Corridor()
    V = {s: 0.0 for s in NONTERMINAL}
    for _ in range(episodes):
        s, _ = env.reset()
        while True:
            a = RIGHT if rng.random() < p_right else LEFT
            s2, r, terminated, truncated, _ = env.step(a)
            td_update(V, s, r, s2, terminated, alpha, gamma)
            if terminated or truncated:
                break
            s = s2
    return V

def hand_check():
    # (1) sweep 표: 그림 2.2의 손계산
    h = evaluate_policy(0.7, sweeps=3)
    expect = [{1: 0, 2: 0, 3: 0}, {1: -0.3, 2: 0, 3: 0.7},
              {1: -0.3, 2: 0.36, 3: 0.7}, {1: -0.0732, 2: 0.36, 3: 0.7972}]
    for got, want in zip(h, expect):
        assert all(abs(got[s] - want[s]) < 1e-12 for s in NONTERMINAL), (got, want)

    # (2) 반복은 정확한 해로 수렴하고, Phase 1의 Monte Carlo 추정 J ≈ 0.547과 맞는다
    V, exact = evaluate_policy(0.7)[-1], exact_values(0.7)
    assert all(abs(V[s] - exact[s]) < 1e-12 for s in NONTERMINAL)
    assert abs(exact[2] - 0.36 / 0.6598) < 1e-12
    print('P(→)=0.7  V =', {s: round(v, 4) for s, v in exact.items()})
    print('P(→)=0.5  V =', {s: round(v, 4) for s, v in exact_values(0.5).items()})

    # (3) V로부터 Q: 표 2.1
    Q = {(s, a): q_from_v(exact, s, a) for s in NONTERMINAL for a in (LEFT, RIGHT)}
    assert abs(Q[2, RIGHT] - 0.9 * exact[3]) < 1e-12 and Q[1, LEFT] == -1.0
    assert abs(0.3 * Q[2, LEFT] + 0.7 * Q[2, RIGHT] - exact[2]) < 1e-12   # V = Σ π Q

    # (4) MC와 TD를 같은 에피소드 ←→→→ 에 한 번씩
    states, rewards, steps = run_actions([LEFT, RIGHT, RIGHT, RIGHT])
    V_mc = mc_update({s: 0.0 for s in NONTERMINAL}, states, rewards)
    V_td = {s: 0.0 for s in NONTERMINAL}
    deltas = [td_update(V_td, *step) for step in steps]
    assert deltas == [0.0, 0.0, 0.0, 1.0]
    assert V_td == {1: 0.0, 2: 0.0, 3: 0.1}
    # 뒤에서부터 갱신: t=3 V(3)=0.1, t=2 V(2)=0.09, t=1 V(1)=0.081, t=0 V(2)=0.09+0.1(0.729−0.09)
    want_mc = {1: 0.081, 2: 0.1539, 3: 0.1}
    assert all(abs(V_mc[s] - want_mc[s]) < 1e-12 for s in NONTERMINAL), V_mc
    print('MC 1회 후 V =', {s: round(v, 4) for s, v in V_mc.items()})
    print('TD 1회 후 V =', V_td, '· δ =', deltas)

    # (5) 두 번째 에피소드 →→ : TD는 정보를 한 칸 더 앞으로 전달한다
    _, _, steps = run_actions([RIGHT, RIGHT])
    deltas = [td_update(V_td, *step) for step in steps]
    assert abs(deltas[0] - 0.09) < 1e-12 and abs(V_td[2] - 0.009) < 1e-12
    assert abs(V_td[3] - 0.19) < 1e-12
    print('TD 2회 후 V =', {s: round(v, 4) for s, v in V_td.items()})

P(→)=0.7  V = {1: 0.0437, 2: 0.5456, 3: 0.8473}
P(→)=0.5  V = {1: -0.5, 2: 0.0, 3: 0.5}
MC 1회 후 V = {1: 0.081, 2: 0.1539, 3: 0.1}
TD 1회 후 V = {1: 0.0, 2: 0.0, 3: 0.1} · δ = [0.0, 0.0, 0.0, 1.0]
TD 2회 후 V = {1: 0.0, 2: 0.009, 3: 0.19}
TD 10000 에피소드 V = {1: 0.023, 2: 0.547, 3: 0.843} · 정확한 값 = {1: 0.044, 2: 0.546, 3: 0.847}

def q_learning_update(Q, s, a, r, s2, terminated, alpha=0.5, gamma=0.9):
    """식 (3.2): Q(s,a) ← Q(s,a) + α[ r + γ(1−terminated)·max_a′ Q(s′,a′) − Q(s,a) ]"""
    target = r + gamma * (0.0 if terminated else max(Q[s2]))
    td_error = target - Q[s][a]
    Q[s][a] += alpha * td_error
    return td_error

def epsilon_greedy(Q, s, epsilon, rng):
    """식 (3.3): 확률 ε로 무작위 행동(탐색), 나머지는 Q가 가장 큰 행동(활용)."""
    if rng.random() < epsilon:
        return rng.randrange(len(Q[s]))
    best = max(Q[s])
    return rng.choice([a for a, q in enumerate(Q[s]) if q == best])   # 동점이면 무작위

def train_q_learning(env, n_states, n_actions, episodes=300, alpha=0.5, gamma=0.9,
                     epsilon=0.2, seed=0):
    rng = random.Random(seed)
    Q = [[0.0] * n_actions for _ in range(n_states)]          # Q 테이블
    for ep in range(episodes):
        s, _ = env.reset(seed=seed + ep)
        while True:
            a = epsilon_greedy(Q, s, epsilon, rng)               # 행동 정책: ε-greedy
            s2, r, terminated, truncated, _ = env.step(a)
            q_learning_update(Q, s, a, r, s2, terminated, alpha, gamma)  # 목표 정책: greedy(max)
            if terminated or truncated:
                break
            s = s2
    return Q

def optimal_q(gamma=0.9, sweeps=100):
    """식 (3.1) Bellman 최적 방정식을 반복 적용(value iteration)해 Q*를 구한다. 정답지 용도."""
    env = Corridor()
    Q = {(s, a): 0.0 for s in NONTERMINAL for a in (LEFT, RIGHT)}
    for _ in range(sweeps):
        new = {}
        for s, a in Q:
            s2, r, done = env.transition(s, a)
            new[s, a] = r + gamma * (0.0 if done else max(Q[s2, LEFT], Q[s2, RIGHT]))
        Q = new
    return Q

def scripted_episode(Q, actions, alpha=0.5):
    env = Corridor()
    s, _ = env.reset()
    for a in actions:
        s2, r, terminated, _, _ = env.step(a)
        q_learning_update(Q, s, a, r, s2, terminated, alpha)
        s = s2


def hand_check():
    Qstar = optimal_q()
    want = {(1, LEFT): -1.0, (1, RIGHT): 0.81, (2, LEFT): 0.729,
            (2, RIGHT): 0.9, (3, LEFT): 0.81, (3, RIGHT): 1.0}
    assert all(abs(Qstar[k] - v) < 1e-12 for k, v in want.items()), Qstar

    Q = [[0.0, 0.0] for _ in range(5)]
    scripted_episode(Q, [RIGHT, RIGHT])             # 에피소드 A
    assert Q[2][RIGHT] == 0.0 and Q[3][RIGHT] == 0.5
    scripted_episode(Q, [RIGHT, RIGHT])             # 에피소드 B
    assert abs(Q[2][RIGHT] - 0.225) < 1e-12 and Q[3][RIGHT] == 0.75
    scripted_episode(Q, [LEFT, LEFT])               # 에피소드 C
    assert Q[2][LEFT] == 0.0 and Q[1][LEFT] == -0.5
    print('손계산 A→B→C 뒤 Q =', {s: Q[s] for s in NONTERMINAL})

손계산 A→B→C 뒤 Q = {1: [-0.5, 0.0], 2: [0.0, 0.225], 3: [0.0, 0.75]}
s=1: Q(s,←)=-0.992  Q(s,→)=+0.810   |  Q*(s,←)=-1.000  Q*(s,→)=+0.810
s=2: Q(s,←)=+0.729  Q(s,→)=+0.900   |  Q*(s,←)=+0.729  Q*(s,→)=+0.900
s=3: Q(s,←)=+0.810  Q(s,→)=+1.000   |  Q*(s,←)=+0.810  Q*(s,→)=+1.000
greedy 정책: ['→', '→', '→']

import gymnasium as gym
from phase03_qlearning import train_q_learning

env = gym.make("FrozenLake-v1", map_name="4x4", is_slippery=False)
Q = train_q_learning(env, n_states=16, n_actions=4, episodes=2000, gamma=0.95, epsilon=0.3)

env = gym.make("Taxi-v3")
Q = train_q_learning(env, n_states=500, n_actions=6, episodes=5000, gamma=0.95, epsilon=0.1)

class LinearQ:
    """one-hot 입력 + 선형층(편향 없음). 가중치 W[a][s]가 곧 Q 테이블의 칸 하나다."""

    def __init__(self, n_states=5, n_actions=2):
        self.W = [[0.0] * n_states for _ in range(n_actions)]

    def q(self, s):
        x = one_hot(s, len(self.W[0]))
        return [sum(w * xi for w, xi in zip(row, x)) for row in self.W]   # = [W[a][s] for a]

    def sgd_step(self, batch, target_net, lr, gamma=0.9):
        """식 (4.1)–(4.3)을 batch 평균으로 한 번 적용하고 loss를 돌려준다."""
        loss, grads = 0.0, {}
        for s, a, r, s2, done in batch:
            y = r + gamma * (0.0 if done else max(target_net.q(s2)))        # 식 (4.1)
            err = self.q(s)[a] - y                                          # Q − y
            loss += 0.5 * err ** 2 / len(batch)                             # 식 (4.2)
            grads[a, s] = grads.get((a, s), 0.0) + err / len(batch)         # 식 (4.3): ∂L/∂W[a][s]
        for (a, s), g in grads.items():
            self.W[a][s] -= lr * g
        return loss

class MLP:
    """입력 → 은닉층(ReLU) → 행동별 Q. 역전파를 손으로 쓴 2층 신경망."""

    def __init__(self, n_in, n_hidden, n_out, rng):
        self.p = {'W1': [[rng.gauss(0, 0.5) for _ in range(n_in)] for _ in range(n_hidden)],
                  'b1': [0.0] * n_hidden,
                  'W2': [[rng.gauss(0, 0.5) for _ in range(n_hidden)] for _ in range(n_out)],
                  'b2': [0.0] * n_out}

    def forward(self, x):
        W1, b1, W2, b2 = (self.p[k] for k in ('W1', 'b1', 'W2', 'b2'))
        z = [sum(w * xi for w, xi in zip(row, x)) + b for row, b in zip(W1, b1)]
        h = [max(0.0, v) for v in z]                                          # ReLU
        q = [sum(w * hi for w, hi in zip(row, h)) + b for row, b in zip(W2, b2)]
        return q, (x, z, h)

    def backward(self, cache, dq):
        """dq = ∂L/∂q 를 받아 모든 파라미터의 gradient를 연쇄법칙으로 계산한다."""
        x, z, h = cache
        W2 = self.p['W2']
        g = {'W2': [[d * hi for hi in h] for d in dq], 'b2': list(dq)}
        dh = [sum(dq[o] * W2[o][j] for o in range(len(dq))) for j in range(len(h))]
        dz = [d if zj > 0 else 0.0 for d, zj in zip(dh, z)]                  # ReLU의 미분
        g['W1'] = [[d * xi for xi in x] for d in dz]
        g['b1'] = dz
        return g

def dqn_update(net, target_net, batch, lr, gamma=0.9):
    """batch 하나로 식 (4.1)–(4.3)을 계산하고 SGD 한 걸음."""
    grads, loss = zeros_like(net.p), 0.0
    for s, a, r, s2, done in batch:
        q_next, _ = target_net.forward(one_hot(s2))
        y = r + gamma * (0.0 if done else max(q_next))           # 식 (4.1): target network로 y
        q, cache = net.forward(one_hot(s))
        err = q[a] - y
        loss += 0.5 * err ** 2 / len(batch)                      # 식 (4.2)
        dq = [err if i == a else 0.0 for i in range(len(q))]     # 고른 행동의 출력에만 gradient
        add_grads(grads, net.backward(cache, dq), 1 / len(batch))
    add_grads(net.p, grads, -lr)                                 # θ ← θ − lr·∇L
    return loss

def train_dqn(steps=3000, batch_size=16, lr=0.05, gamma=0.9, buffer_size=1000,
              learning_starts=200, target_every=50, seed=0):
    rng = random.Random(seed)
    net = MLP(5, 16, 2, rng)
    target_net = copy.deepcopy(net)                              # θ⁻: 주기적으로만 복사
    buffer = deque(maxlen=buffer_size)                           # replay buffer
    env = Corridor()
    s, _ = env.reset()
    for step in range(steps):
        epsilon = max(0.05, 1.0 - step / 1500)                   # 탐색을 점점 줄인다
        q, _ = net.forward(one_hot(s))
        a = rng.randrange(2) if rng.random() < epsilon else max((0, 1), key=lambda i: q[i])
        s2, r, terminated, truncated, _ = env.step(a)
        buffer.append((s, a, r, s2, terminated))                 # 경험 저장 (truncated는 저장 안 함)
        s = env.reset()[0] if terminated or truncated else s2
        if step >= learning_starts:
            dqn_update(net, target_net, rng.sample(buffer, batch_size), lr, gamma)
        if step % target_every == 0:
            target_net = copy.deepcopy(net)
    return net

def hand_check():
    # (1) LinearQ + target = online, batch 1, lr=α 이면 Phase 3의 Q-learning과 숫자가 같다
    net = LinearQ()
    for actions in ([RIGHT, RIGHT], [RIGHT, RIGHT], [LEFT, LEFT]):
        env = Corridor()
        s, _ = env.reset()
        for a in actions:
            s2, r, terminated, _, _ = env.step(a)
            net.sgd_step([(s, a, r, s2, terminated)], net, lr=0.5)
            s = s2
    assert net.W[RIGHT][3] == 0.75 and abs(net.W[RIGHT][2] - 0.225) < 1e-12
    assert net.W[LEFT][1] == -0.5                            # Phase 3 손계산 A→B→C와 동일

    # (2) target network를 고정하고 batch 2개로 한 걸음: 표 4.1
    target = copy.deepcopy(net)
    batch = [(2, RIGHT, 0.0, 3, False), (3, RIGHT, 1.0, 4, True)]
    loss = net.sgd_step(batch, target, lr=0.5)
    assert abs(loss - 0.06625) < 1e-12
    assert abs(net.W[RIGHT][2] - 0.3375) < 1e-12 and abs(net.W[RIGHT][3] - 0.8125) < 1e-12
    print(f'batch loss = {loss:.5f}, Q(2,→): 0.225 → {net.W[RIGHT][2]:.4f}, Q(3,→): 0.75 → {net.W[RIGHT][3]:.4f}')

    # (3) 손으로 쓴 역전파가 맞는지 수치미분(finite difference)으로 확인
    mlp = MLP(5, 4, 2, random.Random(1))
    x, a, y, eps = one_hot(2), RIGHT, 0.7, 1e-6
    def loss_fn():
        return 0.5 * (mlp.forward(x)[0][a] - y) ** 2
    q, cache = mlp.forward(x)
    g = mlp.backward(cache, [q[a] - y if i == a else 0.0 for i in range(2)])
    def nudge(k, i, j, d):
        if j is None:
            mlp.p[k][i] += d
        else:
            mlp.p[k][i][j] += d
    for k, i, j in (('W1', 0, 2), ('W1', 3, 2), ('W2', 1, 0), ('W2', 1, 3), ('b1', 1, None)):
        nudge(k, i, j, eps); up = loss_fn()
        nudge(k, i, j, -2 * eps); down = loss_fn()
        nudge(k, i, j, eps)
        numeric = (up - down) / (2 * eps)
        analytic = g[k][i] if j is None else g[k][i][j]
        assert abs(numeric - analytic) < 1e-6, (k, i, j, numeric, analytic)
    print('역전파 gradient == 수치미분: 확인')

batch loss = 0.06625, Q(2,→): 0.225 → 0.3375, Q(3,→): 0.75 → 0.8125
역전파 gradient == 수치미분: 확인
s=1: Q(s,←)=-1.000  Q(s,→)=+0.810   |  Q*(s,←)=-1.000  Q*(s,→)=+0.810
s=2: Q(s,←)=+0.729  Q(s,→)=+0.900   |  Q*(s,←)=+0.729  Q*(s,→)=+0.900
s=3: Q(s,←)=+0.810  Q(s,→)=+1.000   |  Q*(s,←)=+0.810  Q*(s,→)=+1.000

# -*- coding: utf-8 -*-
"""선택 실습 · Phase 4의 DQN을 PyTorch로 CartPole에 적용한다 (CleanRL dqn.py를 줄인 형태).

실행: pip install torch "gymnasium[classic-control]"  후  python extras/dqn_cartpole.py
CPU에서 1분 안쪽. 복도 버전과 바뀌는 것은 입력(4차원 실수)과 신경망 크기뿐이다.
"""
import random
import sys
from collections import deque

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn


def main(total_steps=30_000, seed=0):
    random.seed(seed)
    torch.manual_seed(seed)
    env = gym.make('CartPole-v1')
    q_net = nn.Sequential(nn.Linear(4, 120), nn.ReLU(), nn.Linear(120, 84), nn.ReLU(), nn.Linear(84, 2))
    target_net = nn.Sequential(nn.Linear(4, 120), nn.ReLU(), nn.Linear(120, 84), nn.ReLU(), nn.Linear(84, 2))
    target_net.load_state_dict(q_net.state_dict())
    optimizer = torch.optim.Adam(q_net.parameters(), lr=2.5e-4)
    buffer = deque(maxlen=10_000)                                   # replay buffer
    gamma, batch_size, returns, ep_return = 0.99, 128, [], 0.0
    obs, _ = env.reset(seed=seed)
    env.action_space.seed(seed)
    for step in range(total_steps):
        epsilon = max(0.05, 1.0 - step / (0.5 * total_steps))       # ε를 선형으로 감소
        if random.random() < epsilon:
            action = env.action_space.sample()
        else:
            with torch.no_grad():
                action = int(q_net(torch.tensor(obs, dtype=torch.float32)).argmax())
        next_obs, reward, terminated, truncated, _ = env.step(action)
        buffer.append((obs, action, reward, next_obs, terminated))
        ep_return += reward
        obs = next_obs
        if terminated or truncated:
            returns.append(ep_return)
            ep_return, (obs, _) = 0.0, env.reset()
        if step > 1_000 and step % 10 == 0:
            s, a, r, s2, d = map(lambda x: torch.as_tensor(np.array(x), dtype=torch.float32),
                                 zip(*random.sample(buffer, batch_size)))
            with torch.no_grad():
                y = r + gamma * (1 - d) * target_net(s2).max(dim=1).values          # 식 (4.1)
            q = q_net(s).gather(1, a.long().unsqueeze(1)).squeeze(1)
            loss = nn.functional.mse_loss(q, y)                                      # 식 (4.2)
            optimizer.zero_grad()
            loss.backward()                                                          # 식 (4.3): autograd
            optimizer.step()
        if step % 500 == 0:
            target_net.load_state_dict(q_net.state_dict())                           # θ⁻ ← θ
        if step % 5_000 == 0 and returns:
            print(f'step {step:6d}  최근 10 에피소드 평균 return {sum(returns[-10:]) / len(returns[-10:]):6.1f}')
    return sum(returns[-10:]) / 10


if __name__ == '__main__':
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 30_000)

step   5000  최근 10 에피소드 평균 return   19.6
step  10000  최근 10 에피소드 평균 return   94.4
step  15000  최근 10 에피소드 평균 return   73.2
step  20000  최근 10 에피소드 평균 return  115.5
step  25000  최근 10 에피소드 평균 return  201.3

def softmax(logits):
    m = max(logits)
    e = [math.exp(x - m) for x in logits]
    return [x / sum(e) for x in e]


def grad_log_softmax(logits, a):
    """식 (5.2): ∂ log π(a|s) / ∂θ_s = onehot(a) − π(·|s)"""
    pi = softmax(logits)
    return [(1.0 if i == a else 0.0) - p for i, p in enumerate(pi)]


def new_theta():
    """칸마다 로짓 [θ_←, θ_→]. 0이면 π = [0.5, 0.5]."""
    return {s: [0.0, 0.0] for s in NONTERMINAL}


def p_right(theta, s):
    return softmax(theta[s])[RIGHT]

def returns_backward(rewards, gamma=0.9):
    G, out = 0.0, []
    for r in reversed(rewards):                     # Phase 1의 식 (1.2)
        G = r + gamma * G
        out.append(G)
    return out[::-1]


def reinforce_update(theta, states, actions, rewards, alpha=0.1, gamma=0.9, baseline=None):
    """식 (5.3)/(5.4): θ_s ← θ_s + α · (G_t − b(s_t)) · ∇ log π(a_t|s_t)

    에피소드 하나의 gradient를 업데이트 전 θ로 모두 계산한 뒤 한꺼번에 더한다.
    """
    G = returns_backward(rewards, gamma)
    steps = []
    for t, (s, a) in enumerate(zip(states, actions)):
        adv = G[t] - (baseline[s] if baseline else 0.0)
        steps.append((s, adv, grad_log_softmax(theta[s], a)))
    for s, adv, g in steps:
        theta[s] = [th + alpha * adv * gi for th, gi in zip(theta[s], g)]
    return G

def train_reinforce(episodes=2000, alpha=0.1, seed=0):
    rng, theta = random.Random(seed), new_theta()
    for _ in range(episodes):
        reinforce_update(theta, *rollout(theta, rng), alpha=alpha)
    return theta

def gradient_samples(baseline, state=3, episodes=20000, seed=0):
    """π = 반반에서 state를 방문할 때마다 (G_t − b)·∂log π/∂θ_state[→] 값을 하나씩 모은다."""
    rng, theta, out = random.Random(seed), new_theta(), []
    for _ in range(episodes):
        states, actions, rewards = rollout(theta, rng)
        G = returns_backward(rewards)
        for t, s in enumerate(states):
            if s == state:
                b = baseline[s] if baseline else 0.0
                out.append((G[t] - b) * grad_log_softmax(theta[s], actions[t])[RIGHT])
    return out


def mean_std(xs):
    m = sum(xs) / len(xs)
    return m, (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5

def sigmoid(x):
    return 1 / (1 + math.exp(-x))


def hand_check():
    assert grad_log_softmax([0.0, 0.0], RIGHT) == [-0.5, 0.5]

    # 성공 에피소드 →→ (α = 1): 표 5.1 첫 줄
    th = new_theta()
    G = reinforce_update(th, *run_scripted([RIGHT, RIGHT]), alpha=1.0)
    assert th[2] == [-0.45, 0.45] and abs(p_right(th, 2) - sigmoid(0.9)) < 1e-12
    assert abs(p_right(th, 3) - sigmoid(1.0)) < 1e-12

    # 실패 에피소드 ←← : 음의 return이 ←를 밀어내서 결과적으로 → 확률이 같은 만큼 오른다
    th = new_theta()
    reinforce_update(th, *run_scripted([LEFT, LEFT]), alpha=1.0)
    assert abs(p_right(th, 2) - sigmoid(0.9)) < 1e-12 and abs(p_right(th, 1) - sigmoid(1.0)) < 1e-12

    # 우회 에피소드 ←→→→ : 나쁜 행동(2에서 ←)도 G=0.729 > 0 이라 강화된다
    th = new_theta()
    reinforce_update(th, *run_scripted([LEFT, RIGHT, RIGHT, RIGHT]), alpha=1.0)
    assert abs(th[2][RIGHT] - (0.9 - 0.729) / 2) < 1e-12
    assert abs(p_right(th, 2) - sigmoid(0.171)) < 1e-12
    print('←→→→ 한 번 뒤 P(→|s) =', {s: round(p_right(th, s), 4) for s in NONTERMINAL})

    # 같은 에피소드에 baseline b(s) = V^π(s) (π = 반반, Phase 2): 상태마다 크기가 달라진다
    th = new_theta()
    V = exact_values(0.5)                            # {1: −0.5, 2: 0, 3: 0.5}
    reinforce_update(th, *run_scripted([LEFT, RIGHT, RIGHT, RIGHT]), alpha=1.0, baseline=V)
    assert abs(p_right(th, 1) - sigmoid(1.31)) < 1e-12 and abs(p_right(th, 3) - sigmoid(0.5)) < 1e-12
    print('baseline 사용  P(→|s) =', {s: round(p_right(th, s), 4) for s in NONTERMINAL})

←→→→ 한 번 뒤 P(→|s) = {1: 0.6921, 2: 0.5426, 3: 0.7311}
baseline 사용  P(→|s) = {1: 0.7875, 2: 0.5426, 3: 0.6225}
3번 칸 방문당 θ₃[→] gradient 추정: baseline 없음 평균 +0.246 표준편차 0.354 | baseline V^π 평균 +0.246 표준편차 0.247
REINFORCE 2000 에피소드 뒤 P(→|s) = {1: 0.933, 2: 0.987, 3: 0.989}

def gae(rewards, values, last_value, terminated, gamma=0.9, lam=0.95):
    """식 (6.1) δ_t = r_{t+1} + γV(s_{t+1}) − V(s_t)
       식 (6.2) A_t = δ_t + γλ·A_{t+1}            (뒤에서부터, Phase 1의 식 (1.2)와 같은 모양)

    values[t] = V(s_t). 에피소드가 terminated면 마지막 다음 상태의 가치는 0,
    truncated면 last_value = V(s_T)로 미래를 추정해 채운다 (bootstrapping).
    """
    T = len(rewards)
    next_values = values[1:] + [0.0 if terminated else last_value]
    deltas = [rewards[t] + gamma * next_values[t] - values[t] for t in range(T)]
    adv, A = [0.0] * T, 0.0
    for t in reversed(range(T)):
        A = deltas[t] + gamma * lam * A
        adv[t] = A
    returns = [a + v for a, v in zip(adv, values)]          # critic의 학습 목표
    return deltas, adv, returns

def train_a2c(iterations=300, episodes_per_batch=8, actor_lr=1.0, critic_lr=0.05,
              gamma=0.9, lam=0.95, seed=0):
    """A2C 방식: 에피소드 몇 개를 모아 GAE로 advantage를 만들고 actor와 critic을 함께 갱신."""
    rng, env = random.Random(seed), Corridor()
    theta, V = new_theta(), {s: 0.0 for s in NONTERMINAL}
    for _ in range(iterations):
        batch = []
        for _ in range(episodes_per_batch):
            s, _ = env.reset()
            states, actions, rewards = [], [], []
            while True:
                a = RIGHT if rng.random() < p_right(theta, s) else LEFT
                s2, r, terminated, truncated, _ = env.step(a)
                states.append(s); actions.append(a); rewards.append(r)
                if terminated or truncated:
                    break
                s = s2
            last = 0.0 if terminated else V[s2]
            _, adv, rets = gae(rewards, values_of(V, states), last, terminated, gamma, lam)
            batch += list(zip(states, actions, adv, rets))
        n = len(batch)
        grads = {s: [0.0, 0.0] for s in NONTERMINAL}
        for s, a, A, ret in batch:
            g = grad_log_softmax(theta[s], a)
            grads[s] = [x + A * gi / n for x, gi in zip(grads[s], g)]   # actor: E[A ∇log π]
            V[s] += critic_lr * (ret - V[s])                          # critic: V(s) → GAE return
        for s in NONTERMINAL:
            theta[s] = [x + actor_lr * gi for x, gi in zip(theta[s], grads[s])]
    return theta, V

def hand_check():
    V = exact_values(0.5)                                    # {1: −0.5, 2: 0, 3: 0.5}
    states, actions, rewards = run_scripted([LEFT, RIGHT, RIGHT, RIGHT])
    vals = values_of(V, states)
    close = lambda xs, ys: all(abs(x - y) < 1e-12 for x, y in zip(xs, ys))

    deltas, A1, _ = gae(rewards, vals, 0.0, True, lam=1.0)
    assert close(deltas, [-0.45, 0.5, 0.45, 0.5])
    assert close(A1, [0.729, 1.31, 0.9, 0.5])                # λ=1: G_t − V(s_t) (Monte Carlo)
    _, A0, _ = gae(rewards, vals, 0.0, True, lam=0.0)
    assert close(A0, deltas)                                  # λ=0: δ_t (TD)
    _, Ah, _ = gae(rewards, vals, 0.0, True, lam=0.5)
    assert close(Ah, [-0.0883125, 0.80375, 0.675, 0.5])
    print('δ =', deltas)
    print('A(λ=1) =', [round(a, 4) for a in A1], '| A(λ=0.5) =', [round(a, 4) for a in Ah])

    # 첫 걸음(2에서 ←)만 actor–critic으로 갱신 (α=1): REINFORCE와 반대 방향
    th = new_theta()
    g = grad_log_softmax(th[2], LEFT)
    th[2] = [x + 1.0 * deltas[0] * gi for x, gi in zip(th[2], g)]
    assert abs(p_right(th, 2) - sigmoid(0.45)) < 1e-12
    print(f'2에서 ← 를 δ = −0.45 로 평가 → P(→|2) = {p_right(th, 2):.4f} (REINFORCE는 이 걸음의 ←를 강화했다)')

δ = [-0.45, 0.5, 0.45, 0.5]
A(λ=1) = [0.729, 1.31, 0.9, 0.5] | A(λ=0.5) = [-0.0883, 0.8037, 0.675, 0.5]
2에서 ← 를 δ = −0.45 로 평가 → P(→|2) = 0.6106 (REINFORCE는 이 걸음의 ←를 강화했다)
A2C 뒤 P(→|s) = {1: 0.924, 2: 0.989, 3: 0.979} | critic V = {1: 0.68, 2: 0.885, 3: 0.991}

def ppo_clip_term(ratio, adv, eps=0.2):
    """식 (7.2): min(ρA, clip(ρ, 1−ε, 1+ε)A) 와, 이 항을 log π로 미분한 계수.

    ∂ρ/∂θ = ρ·∇log π 이므로 clip되지 않은 쪽이 선택되면 gradient 계수는 ρA,
    clip된 쪽이 선택되면 상수이므로 0이다.
    """
    unclipped = ratio * adv
    clipped = max(1 - eps, min(1 + eps, ratio)) * adv
    if clipped < unclipped:                         # 더 작은(보수적인) 쪽이 clip된 항
        return clipped, 0.0
    return unclipped, ratio * adv

def epoch_demo(epochs=20, lr=0.2, eps=0.2, clip=True):
    """2번 칸에서 모은 고정 batch [(→, A=+1), (←, A=−1)]를 epoch마다 다시 학습한다."""
    batch = [(RIGHT, +1.0), (LEFT, -1.0)]
    old = softmax([0.0, 0.0])                       # π_old: 데이터를 모을 때의 정책, 고정
    theta, history = [0.0, 0.0], [0.5]
    for _ in range(epochs):
        pi, grad = softmax(theta), [0.0, 0.0]
        for a, A in batch:
            ratio = pi[a] / old[a]                                    # 식 (7.1)
            coef = ppo_clip_term(ratio, A, eps)[1] if clip else ratio * A
            g = grad_log_softmax(theta, a)
            grad = [x + coef * gi / len(batch) for x, gi in zip(grad, g)]
        theta = [x + lr * gi for x, gi in zip(theta, grad)]          # 경사 상승
        history.append(softmax(theta)[RIGHT])
    return history

def collect(theta, V, rng, episodes, gamma, lam):
    env, batch = Corridor(), []
    for _ in range(episodes):
        s, _ = env.reset()
        states, actions, rewards, logps = [], [], [], []
        while True:
            pi = softmax(theta[s])
            a = RIGHT if rng.random() < pi[RIGHT] else LEFT
            s2, r, terminated, truncated, _ = env.step(a)
            states.append(s); actions.append(a); rewards.append(r); logps.append(math.log(pi[a]))
            if terminated or truncated:
                break
            s = s2
        _, adv, rets = gae(rewards, values_of(V, states), 0.0 if terminated else V[s2], terminated, gamma, lam)
        batch += list(zip(states, actions, logps, adv, rets))
    return batch


def train_ppo(iterations=60, episodes=8, epochs=4, lr=0.5, critic_lr=0.1, eps=0.2,
              gamma=0.9, lam=0.95, seed=0):
    rng = random.Random(seed)
    theta, V = new_theta(), {s: 0.0 for s in NONTERMINAL}
    log = []
    for it in range(iterations):
        batch = collect(theta, V, rng, episodes, gamma, lam)          # π_old로 rollout
        clipped = 0
        for _ in range(epochs):                                         # 같은 batch를 K번 재사용
            grads = {s: [0.0, 0.0] for s in NONTERMINAL}
            for s, a, old_logp, A, ret in batch:
                ratio = math.exp(math.log(softmax(theta[s])[a]) - old_logp)
                _, coef = ppo_clip_term(ratio, A, eps)
                clipped += coef == 0.0 and A != 0.0
                g = grad_log_softmax(theta[s], a)
                grads[s] = [x + coef * gi / len(batch) for x, gi in zip(grads[s], g)]
            for s in NONTERMINAL:
                theta[s] = [x + lr * gi for x, gi in zip(theta[s], grads[s])]
        for s, _, _, _, ret in batch:                                   # critic: V → GAE return
            V[s] += critic_lr * (ret - V[s])
        log.append((it, clipped / (epochs * len(batch))))
    return theta, V, log

def hand_check():
    assert ppo_clip_term(1.4, 0.9) == (1.2 * 0.9, 0.0)              # 좋은 행동, 이미 많이 올림 → 멈춤
    assert ppo_clip_term(0.6, -0.5) == (0.8 * -0.5, 0.0)            # 나쁜 행동, 이미 많이 내림 → 멈춤
    obj, coef = ppo_clip_term(1.1, 0.9)
    assert abs(obj - 0.99) < 1e-12 and abs(coef - 0.99) < 1e-12     # 범위 안: 평범한 policy gradient
    assert ppo_clip_term(0.6, 0.5) == (0.3, 0.3)                    # 좋은 행동을 실수로 내렸으면 되돌린다
    assert ppo_clip_term(1.5, -1.0) == (-1.5, -1.5)                 # 나쁜 행동을 실수로 올렸으면 되돌린다

    with_clip, no_clip = epoch_demo(clip=True), epoch_demo(clip=False)
    assert with_clip[2] < 0.6 < with_clip[3]                      # 3번째 epoch에 경계 0.6 = 1.2 × 0.5를 넘고
    assert all(p == with_clip[3] for p in with_clip[3:])           # 그 뒤로는 gradient가 0이라 멈춘다
    assert with_clip[:4] == no_clip[:4] and no_clip[-1] > 0.92     # clip이 없으면 계속 올라간다
    print('epoch별 P(→|2), clip 있음:', [round(p, 3) for p in with_clip[:6]], '…', round(with_clip[-1], 3))
    print('epoch별 P(→|2), clip 없음:', [round(p, 3) for p in no_clip[:6]], '…', round(no_clip[-1], 3))

epoch별 P(→|2), clip 있음: [0.5, 0.55, 0.598, 0.643, 0.643, 0.643] … 0.643
epoch별 P(→|2), clip 없음: [0.5, 0.55, 0.598, 0.643, 0.684, 0.72] … 0.921
PPO 60회 반복 뒤 P(→|s) = {1: 0.912, 2: 0.977, 3: 0.947} | critic V = {1: 0.435, 2: 0.891, 3: 0.992}
clip된 샘플 비율 (처음 3회, 마지막 3회): [0.0, 0.03, 0.08] [0.0, 0.0, 0.0]

# -*- coding: utf-8 -*-
"""선택 실습 · Phase 7의 PPO를 PyTorch로 CartPole에 적용한다 (CleanRL ppo.py를 줄인 형태).

실행: pip install torch "gymnasium[classic-control]"  후  python extras/ppo_cartpole.py
HF Unit 8 Part 1은 같은 구조로 LunarLander를 학습한다. ENV_ID만 바꿔 보자.
"""
import sys

import gymnasium as gym
import torch
import torch.nn as nn
from torch.distributions import Categorical

ENV_ID = 'CartPole-v1'


def main(iterations=60, seed=0):
    torch.manual_seed(seed)
    env = gym.make(ENV_ID)
    n_obs, n_act = env.observation_space.shape[0], env.action_space.n
    actor = nn.Sequential(nn.Linear(n_obs, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, n_act))
    critic = nn.Sequential(nn.Linear(n_obs, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))
    opt = torch.optim.Adam(list(actor.parameters()) + list(critic.parameters()), lr=2.5e-4)
    T, gamma, lam, eps, epochs, minibatch = 1024, 0.99, 0.95, 0.2, 4, 256
    obs, _ = env.reset(seed=seed)
    returns, ep_return = [], 0.0
    for it in range(iterations):
        # ① rollout: π_old로 T걸음 모으기
        buf = {k: [] for k in ('obs', 'act', 'logp', 'rew', 'done', 'val')}
        for _ in range(T):
            o = torch.tensor(obs, dtype=torch.float32)
            with torch.no_grad():
                dist = Categorical(logits=actor(o))
                a = dist.sample()
                v = critic(o).squeeze()
            obs, r, terminated, truncated, _ = env.step(int(a))
            for k, x in zip(buf, (o, a, dist.log_prob(a), r, float(terminated), v)):
                buf[k].append(x)
            ep_return += r
            if terminated or truncated:
                returns.append(ep_return)
                ep_return, (obs, _) = 0.0, env.reset()
        # ② GAE (식 6.1–6.2). truncated는 이 축약판에서 종료처럼 다룬다.
        with torch.no_grad():
            next_v = critic(torch.tensor(obs, dtype=torch.float32)).squeeze()
        adv, A = [0.0] * T, 0.0
        for t in reversed(range(T)):
            nv = next_v if t == T - 1 else buf['val'][t + 1]
            nonterminal = 1.0 - buf['done'][t]
            delta = buf['rew'][t] + gamma * nv * nonterminal - buf['val'][t]
            A = delta + gamma * lam * nonterminal * A
            adv[t] = A
        b_obs, b_act = torch.stack(buf['obs']), torch.stack(buf['act'])
        b_logp, b_val = torch.stack(buf['logp']), torch.stack(buf['val'])
        b_adv = torch.tensor(adv)
        b_ret = b_adv + b_val
        # ③ 같은 batch를 epochs번, minibatch로 나눠 재사용 (식 7.1–7.3)
        for _ in range(epochs):
            for idx in torch.randperm(T).split(minibatch):
                dist = Categorical(logits=actor(b_obs[idx]))
                ratio = (dist.log_prob(b_act[idx]) - b_logp[idx]).exp()
                a = b_adv[idx]
                a = (a - a.mean()) / (a.std() + 1e-8)                    # advantage 정규화
                pg_loss = -torch.min(ratio * a, ratio.clamp(1 - eps, 1 + eps) * a).mean()
                v_loss = 0.5 * (critic(b_obs[idx]).squeeze() - b_ret[idx]).pow(2).mean()
                loss = pg_loss + 0.5 * v_loss - 0.01 * dist.entropy().mean()
                opt.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(list(actor.parameters()) + list(critic.parameters()), 0.5)
                opt.step()
        if it % 10 == 0 and returns:
            print(f'iter {it:3d}  최근 10 에피소드 평균 return {sum(returns[-10:]) / len(returns[-10:]):6.1f}')
    return sum(returns[-10:]) / 10


if __name__ == '__main__':
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 60)

iter   0  최근 10 에피소드 평균 return   20.1
iter  10  최근 10 에피소드 평균 return   79.8
iter  20  최근 10 에피소드 평균 return  152.8
iter  30  최근 10 에피소드 평균 return  212.8
iter  40  최근 10 에피소드 평균 return  208.4
iter  50  최근 10 에피소드 평균 return  254.0

# SFT를 마친 모델이라고 가정한 출발점이자 reference 정책 π_ref.
SFT_PROBS = {
    ('3×4=', ()):     _probs({'1': 0.4, '7': 0.4, '2': 0.1, '6': 0.05, '<eos>': 0.05}),
    ('3×4=', ('1',)): _probs({'2': 0.5, '<eos>': 0.3, '1': 0.1, '6': 0.05, '7': 0.05}),
    ('2×3=', ()):     _probs({'6': 0.3, '7': 0.3, '1': 0.2, '2': 0.1, '<eos>': 0.1}),
}
DEFAULT_PROBS = _probs({'<eos>': 0.8})   # 그 밖의 prefix: 대개 바로 끝낸다 (나머지 토큰 0.05씩)

class TinyLM:
    def __init__(self):
        self.logits = {}                                   # prefix → 로짓 벡터 (필요할 때 생성)

    def _row(self, prompt, prefix):
        key = (prompt, tuple(prefix))
        if key not in self.logits:
            self.logits[key] = [math.log(p) for p in SFT_PROBS.get(key, DEFAULT_PROBS)]
        return self.logits[key]

    def dist(self, prompt, prefix):
        """π(· | x, y<t): 다음 토큰 분포"""
        return softmax(self._row(prompt, prefix))

    def sample(self, prompt, rng):
        tokens = []
        while len(tokens) < MAX_NEW_TOKENS:
            p = self.dist(prompt, tokens)
            tok = rng.choices(VOCAB, weights=p)[0]
            tokens.append(tok)
            if tok == '<eos>':
                break
        return tokens

    def token_logps(self, prompt, tokens):
        """식 (8.2)의 각 항: log π(y_t | x, y<t)"""
        return [math.log(self.dist(prompt, tokens[:t])[VOCAB.index(tok)]) for t, tok in enumerate(tokens)]

    def grad_token(self, prompt, tokens, t):
        """∂ log π(y_t | x, y<t) / ∂(그 prefix의 로짓) = onehot(y_t) − π  (Phase 5의 식 (5.2)와 같다)"""
        p = self.dist(prompt, tokens[:t])
        a = VOCAB.index(tokens[t])
        return (prompt, tuple(tokens[:t])), [(1.0 if i == a else 0.0) - pi for i, pi in enumerate(p)]

    def apply(self, grads, lr):
        """grads: {prefix key: 로짓 gradient}. 경사 상승 한 걸음."""
        for key, g in grads.items():
            row = self._row(*key)
            for i, gi in enumerate(g):
                row[i] += lr * gi

    def copy(self):
        other = TinyLM()
        other.logits = {k: list(v) for k, v in self.logits.items()}
        return other

def verify(prompt, tokens):
    """검증기: <eos>로 끝났고 숫자가 정답과 같으면 1, 아니면 0."""
    finished = tokens and tokens[-1] == '<eos>'
    return 1.0 if finished and ''.join(tokens[:-1]) == ANSWERS[prompt] else 0.0

def train_reinforce_kl(beta, iterations=150, group=16, lr=0.5, seed=0):
    """식 (8.4): 시퀀스 보상 R(x,y) − β·log π(y|x)/π_ref(y|x) 를 REINFORCE로 최대화한다."""
    rng, lm = random.Random(seed), TinyLM()
    ref = lm.copy()                                         # π_ref: 학습 시작 시점(SFT) 모델, 고정
    for _ in range(iterations):
        grads = {}
        for prompt in ANSWERS:
            samples = [lm.sample(prompt, rng) for _ in range(group)]
            scores = []
            for y in samples:
                log_ratio = sum(lm.token_logps(prompt, y)) - sum(ref.token_logps(prompt, y))
                scores.append(verify(prompt, y) - beta * log_ratio)
            baseline = sum(scores) / len(scores)              # Phase 5의 baseline: 같은 질문의 평균
            for y, score in zip(samples, scores):
                for t in range(len(y)):
                    key, g = lm.grad_token(prompt, y, t)
                    add_grad(grads, key, g, (score - baseline) / (group * len(ANSWERS)))
        lm.apply(grads, lr)
    return lm, ref

def kl_optimal_p_correct(prompt, beta):
    """식 (8.5): max E[R] − β·KL 의 해는 π*(y) ∝ π_ref(y)·exp(R(y)/β).

    R이 0/1뿐이면 정답 전체의 확률 p_ref가 exp(1/β)배로 커진 뒤 다시 정규화된다.
    """
    p_ref = p_correct(TinyLM(), prompt)
    w = p_ref * math.exp(1 / beta)
    return w / (w + 1 - p_ref)

def hand_check():
    lm = TinyLM()
    y = ['1', '2', '<eos>']
    logps = lm.token_logps('3×4=', y)                       # 식 (8.2)의 각 항
    assert all(abs(a - math.log(b)) < 1e-12 for a, b in zip(logps, [0.4, 0.5, 0.8]))
    assert abs(math.exp(sum(logps)) - 0.16) < 1e-12          # 식 (8.1): 0.4 × 0.5 × 0.8
    assert abs(p_correct(lm, '3×4=') - 0.16) < 1e-12 and abs(p_correct(lm, '2×3=') - 0.24) < 1e-12
    print('log π(12⏎ | 3×4=) =', [round(x, 4) for x in logps], '합', round(sum(logps), 4))

    # 정답 응답 하나로 시퀀스 REINFORCE 한 걸음 (R=1, baseline 0, α=1): 식 (8.3)
    grads = {}
    for t in range(len(y)):
        key, g = lm.grad_token('3×4=', y, t)
        add_grad(grads, key, g, 1.0 * verify('3×4=', y))
    lm.apply(grads, lr=1.0)
    p0 = lm.dist('3×4=', [])[VOCAB.index('1')]
    print(f'한 번 갱신 뒤: π(1|3×4=) 0.400 → {p0:.4f}, P(12⏎) 0.160 → {p_correct(lm, "3×4="):.4f}')
    assert p0 > 0.4 and p_correct(lm, '3×4=') > 0.16

log π(12⏎ | 3×4=) = [-0.9163, -0.6931, -0.2231] 합 -1.8326
한 번 갱신 뒤: π(1|3×4=) 0.400 → 0.6163, P(12⏎) 0.160 → 0.3451
β=0.0: P(정답) = {'3×4=': 0.908, '2×3=': 0.938}  KL(π‖π_ref) = {'3×4=': 1.42, '2×3=': 1.157}
β=0.5: P(정답) = {'3×4=': 0.541, '2×3=': 0.675}  KL(π‖π_ref) = {'3×4=': 0.39, '2×3=': 0.428}
β=0.5의 이론적 최적 P(정답) = {'3×4=': 0.585, '2×3=': 0.7}

# 사람이 두 응답을 비교해 고른 결과: (프롬프트, 더 나은 응답 y_w, 덜 나은 응답 y_l)
PREFERENCES = [
    ('3×4=', '12⏎', '7⏎'),
    ('3×4=', '12⏎', '1⏎'),
    ('3×4=', '1⏎', '7⏎'),      # "1"은 틀렸지만 정답의 첫 자리라 "7"보다 낫다고 판단
    ('2×3=', '6⏎', '7⏎'),
    ('2×3=', '6⏎', '1⏎'),
]

def train_reward_model(prefs, epochs=200, lr=0.5):
    """식 (9.1)–(9.2): Bradley–Terry. P(y_w ≻ y_l) = σ(r(y_w) − r(y_l)),  L = −log σ(r_w − r_l)

    여기서는 응답마다 점수 하나를 갖는 표가 보상 모델이다. 처음 보는 응답은 0점.
    """
    r = {}
    for _ in range(epochs):
        for x, yw, yl in prefs:
            rw, rl = r.get((x, yw), 0.0), r.get((x, yl), 0.0)
            g = 1 - sigmoid(rw - rl)                    # −∂L/∂r_w = ∂L/∂r_l 의 크기
            r[x, yw] = rw + lr * g
            r[x, yl] = rl - lr * g
    return r

def rlhf_ppo(rm, beta=0.1, iterations=80, group=16, epochs=4, lr=0.5, critic_lr=0.2,
             eps=0.2, lam=0.95, seed=0):
    rng, lm = random.Random(seed), TinyLM()
    ref, V = lm.copy(), {}                              # π_ref(고정), critic V(prefix)
    for _ in range(iterations):
        batch = []
        for x in ANSWERS:                               # ① rollout: π_old로 응답 생성
            for _ in range(group):
                y = lm.sample(x, rng)
                old = lm.token_logps(x, y)
                refl = ref.token_logps(x, y)
                # ② 토큰별 보상 (식 9.3): 매 토큰 −β·(log π_old − log π_ref), 마지막 토큰에 RM 점수
                rewards = [-beta * (o - r) for o, r in zip(old, refl)]
                rewards[-1] += rm.get((x, show(y)), 0.0)
                # ③ critic 값과 GAE (γ = 1): 응답이 끝나면 다음 가치는 0
                keys = [(x, tuple(y[:t])) for t in range(len(y))]
                vals = [V.get(k, 0.0) for k in keys]
                adv, A = [0.0] * len(y), 0.0
                for t in reversed(range(len(y))):
                    nxt = vals[t + 1] if t + 1 < len(y) else 0.0
                    A = rewards[t] + nxt - vals[t] + lam * A
                    adv[t] = A
                batch.append((x, y, old, adv, [a + v for a, v in zip(adv, vals)], keys))
        n_tokens = sum(len(b[1]) for b in batch)
        for _ in range(epochs):                         # ④ PPO: 같은 batch를 여러 번
            grads = {}
            for x, y, old, adv, _, _ in batch:
                new = lm.token_logps(x, y)
                for t in range(len(y)):
                    _, coef = ppo_clip_term(math.exp(new[t] - old[t]), adv[t], eps)
                    key, g = lm.grad_token(x, y, t)
                    add_grad(grads, key, g, coef / n_tokens)
            lm.apply(grads, lr)
        for x, y, _, _, rets, keys in batch:            # ⑤ critic: V(prefix) → return
            for k, ret in zip(keys, rets):
                V[k] = V.get(k, 0.0) + critic_lr * (ret - V.get(k, 0.0))
    return lm, ref

def rlhf_optimum(rm, prompt, beta):
    """식 (8.5)와 같은 원리: π*(y) ∝ π_ref(y)·exp(r(y)/β). 응답을 모두 나열해 정확히 계산."""
    ref = TinyLM()
    w = {show(y): q * math.exp(rm.get((prompt, show(y)), 0.0) / beta) for y, q in all_responses(ref, prompt)}
    z = sum(w.values())
    return {y: v / z for y, v in w.items()}

def hand_check():
    one_step = train_reward_model(PREFERENCES[:1], epochs=1, lr=1.0)
    assert one_step == {('3×4=', '12⏎'): 0.5, ('3×4=', '7⏎'): -0.5}   # 처음 loss = log 2, 기울기 ∓0.5
    assert abs(-math.log(sigmoid(0)) - math.log(2)) < 1e-12
    print('첫 비교 한 걸음: r(12⏎) = +0.5, r(7⏎) = −0.5, 선호 확률 σ(1) =', round(sigmoid(1.0), 4))

첫 비교 한 걸음: r(12⏎) = +0.5, r(7⏎) = −0.5, 선호 확률 σ(1) = 0.7311
보상 모델 점수: {'3×4=12⏎': 4.63, '2×3=6⏎': 3.8, '3×4=1⏎': 0.0, '2×3=1⏎': -1.89, '2×3=7⏎': -1.91, '3×4=7⏎': -4.64}
3×4=  P(정답): 0.160 → 0.980  (이론적 최적 1.000)   KL(π‖π_ref) = 1.714
2×3=  P(정답): 0.240 → 0.985  (이론적 최적 1.000)   KL(π‖π_ref) = 1.344

def dpo_loss_and_coef(lm, ref, x, yw, yl, beta=0.1):
    """식 (10.3): L = −log σ(β·m),  m = [log π(y_w) − log π_ref(y_w)] − [log π(y_l) − log π_ref(y_l)]

    −∂L/∂m = β·σ(−β·m) 이므로, 경사 상승 방향은 β·σ(−βm)·(∇log π(y_w) − ∇log π(y_l)).
    """
    m = (seq_logp(lm, x, yw) - seq_logp(ref, x, yw)) - (seq_logp(lm, x, yl) - seq_logp(ref, x, yl))
    return -math.log(sigmoid(beta * m)), beta * sigmoid(-beta * m), m


def train_dpo(prefs, beta=0.1, epochs=300, lr=1.0):
    lm = TinyLM()
    ref = lm.copy()
    for _ in range(epochs):
        grads = {}
        for x, yw, yl in prefs:
            _, coef, _ = dpo_loss_and_coef(lm, ref, x, yw, yl, beta)
            add_seq_grad(grads, lm, x, yw, +coef / len(prefs))    # 좋은 응답의 확률 ↑
            add_seq_grad(grads, lm, x, yl, -coef / len(prefs))    # 나쁜 응답의 확률 ↓
        lm.apply(grads, lr)
    return lm, ref

KTO_LABELS = [('3×4=', '12⏎', True), ('3×4=', '7⏎', False), ('3×4=', '1⏎', False),
              ('2×3=', '6⏎', True), ('2×3=', '7⏎', False), ('2×3=', '1⏎', False)]


def train_kto(labels, beta=0.1, epochs=300, lr=1.0):
    """식 (10.4): r = log π(y)/π_ref(y),  z = KL(π‖π_ref) (기준점, gradient 없음)
       좋음:  L = 1 − σ(β(r − z))    나쁨:  L = 1 − σ(β(z − r))
    """
    lm = TinyLM()
    ref = lm.copy()
    for _ in range(epochs):
        grads = {}
        z = {x: max(0.0, kl(lm, ref, x)) for x in ANSWERS}         # detach된 기준점
        for x, y, good in labels:
            r = seq_logp(lm, x, y) - seq_logp(ref, x, y)
            v = sigmoid(beta * (r - z[x]) if good else beta * (z[x] - r))
            coef = beta * v * (1 - v) * (1 if good else -1)        # −∂L/∂r
            add_seq_grad(grads, lm, x, y, coef / len(labels))
        lm.apply(grads, lr)
    return lm, ref

def hand_check():
    lm = TinyLM()
    ref = lm.copy()
    loss, coef, m = dpo_loss_and_coef(lm, ref, '3×4=', '12⏎', '7⏎')
    assert m == 0.0 and abs(loss - math.log(2)) < 1e-12 and abs(coef - 0.05) < 1e-12
    print(f'처음: m = 0, DPO loss = log 2 = {loss:.4f}, gradient 계수 β·σ(0) = {coef}')

처음: m = 0, DPO loss = log 2 = 0.6931, gradient 계수 β·σ(0) = 0.05
DPO 뒤:
  3×4=  P(정답) 0.160 → 0.805   KL = 1.287
  2×3=  P(정답) 0.240 → 0.896   KL = 1.201
  암묵적 보상 β·log π/π_ref: {'3×4=12⏎': 0.162, '3×4=1⏎': -0.106, '3×4=7⏎': -0.912}
KTO 뒤:
  3×4=  P(정답) 0.160 → 0.485   KL = 0.406
  2×3=  P(정답) 0.240 → 0.711   KL = 0.675

# -*- coding: utf-8 -*-
"""선택 실습 · Phase 10의 DPO를 TRL로 실제 LLM에 적용한다.

실행 (GPU 권장, Colab T4 가능):
    pip install trl==0.24.0 transformers datasets
    python extras/trl_dpo.py                           # Qwen/Qwen2.5-0.5B-Instruct
    python extras/trl_dpo.py <모델 이름> <max_steps>

Phase 9의 PREFERENCES와 같은 (prompt, chosen, rejected) 형식이다. 보상 모델도 rollout도 없다.
"""
import random
import sys

from datasets import Dataset
from trl import DPOConfig, DPOTrainer

MODEL = sys.argv[1] if len(sys.argv) > 1 else 'Qwen/Qwen2.5-0.5B-Instruct'
MAX_STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 100


def make_pairs(n=512, seed=0):
    rng = random.Random(seed)
    rows = []
    for _ in range(n):
        a, b = rng.randint(2, 19), rng.randint(2, 19)
        wrong = a * b + rng.choice([-a, -b, b, a, 1, -1])            # 그럴듯한 오답
        rows.append({'prompt': [{'role': 'user', 'content': f'{a}×{b}는 얼마인가요? 숫자 하나로만 답하세요.'}],
                     'chosen': [{'role': 'assistant', 'content': str(a * b)}],
                     'rejected': [{'role': 'assistant', 'content': str(wrong)}]})
    return Dataset.from_list(rows)


args = DPOConfig(
    output_dir='dpo-out',
    beta=0.1,                          # 식 (10.3)의 β
    per_device_train_batch_size=8,
    learning_rate=5e-7,
    max_steps=MAX_STEPS,
    logging_steps=1,
    report_to='none',
)
trainer = DPOTrainer(model=MODEL, args=args, train_dataset=make_pairs())   # ref_model은 자동으로 복사
trainer.train()

{'loss': 0.6931, 'grad_norm': 0.1568, 'rewards/chosen': 0.0, 'rewards/rejected': 0.0, 'rewards/margins': 0.0, 'logps/chosen': -65.59, 'logps/rejected': -67.09, ...}
{'loss': 0.6931, 'grad_norm': 0.1384, 'rewards/chosen': -1.3e-06, 'rewards/rejected': -9.5e-07, 'rewards/margins': -3.3e-07, ...}

def group_advantages(rewards, mode='std', eps=1e-4):
    """식 (11.1): A_i = (r_i − mean(r)) / (std(r) + eps)   — 같은 질문의 G개 응답끼리 비교.

    mode='mean' 이면 표준편차로 나누지 않는다 (Dr. GRPO).
    mode='rloo' 이면 나를 뺀 나머지의 평균을 뺀다 (RLOO, 식 12.1).
    std는 모집단 표준편차(n으로 나눔)다. TRL은 표본 표준편차(n−1)를 쓴다.
    """
    G = len(rewards)
    mean = sum(rewards) / G
    if mode == 'rloo':
        return [r - (sum(rewards) - r) / (G - 1) for r in rewards]
    if mode == 'mean':
        return [r - mean for r in rewards]
    std = (sum((r - mean) ** 2 for r in rewards) / G) ** 0.5
    return [(r - mean) / (std + eps) for r in rewards]

DEFAULT = dict(adv='std', agg='seq_mean', eps_low=0.2, eps_high=0.2, ratio='token', beta=0.04, dynamic=False)


def grpo_loss_grads(lm, ref, group, cfg=DEFAULT):
    """식 (11.2): 그룹 하나의 목적함수를 로짓 gradient로 바꿔 돌려준다.

    group: [(prompt, tokens, old_logps, advantage)]
    cfg['agg']    'seq_mean'  응답별 토큰 평균 → 그룹 평균           (GRPO 원형)
                  'token_mean' 그룹 전체 토큰 평균                    (DAPO)
                  'const'      토큰 합 / (G · 최대 생성 길이)          (Dr. GRPO)
    cfg['ratio']  'token' 토큰마다 ratio / 'sequence' 응답 하나에 ratio (GSPO)
    """
    G, grads = len(group), {}
    total_tokens = sum(len(y) for _, y, _, _ in group)
    for x, y, old, A in group:
        new = lm.token_logps(x, y)
        refl = ref.token_logps(x, y)
        if cfg['ratio'] == 'sequence':                      # GSPO: exp(평균 log ratio) 하나를 공유
            seq_ratio = math.exp(sum(n - o for n, o in zip(new, old)) / len(y))
        for t in range(len(y)):
            ratio = seq_ratio if cfg['ratio'] == 'sequence' else math.exp(new[t] - old[t])
            lo, hi = 1 - cfg['eps_low'], 1 + cfg['eps_high']
            clipped = (A > 0 and ratio > hi) or (A < 0 and ratio < lo)
            # Phase 7의 ppo_clip_term과 같은 규칙. GSPO에서도 ∂s/∂log π_t = s/|y| 이고,
            # seq_mean 가중치 1/(G|y|)가 그 1/|y|를 담당하므로 계수는 똑같이 ratio·A다.
            coef = 0.0 if clipped else ratio * A
            # KL 벌점 k3 = π_ref/π − log(π_ref/π) − 1 의 gradient 계수: −β(1 − π_ref/π)
            coef += -cfg['beta'] * (1 - math.exp(refl[t] - new[t]))
            if cfg['agg'] == 'seq_mean':
                w = 1 / (G * len(y))
            elif cfg['agg'] == 'token_mean':
                w = 1 / total_tokens
            else:
                w = 1 / (G * MAX_NEW_TOKENS)
            key, g = lm.grad_token(x, y, t)
            add_grad(grads, key, g, coef * w)
    return grads

def train_grpo(cfg=DEFAULT, iterations=60, G=8, epochs=2, lr=1.0, seed=0):
    rng, lm = random.Random(seed), TinyLM()
    ref = lm.copy()
    for _ in range(iterations):
        groups = []
        for x in ANSWERS:
            ys = [lm.sample(x, rng) for _ in range(G)]                 # ① 질문마다 G개 생성
            rs = [verify(x, y) for y in ys]                            # ② 검증기 보상 0/1
            if cfg['dynamic'] and len(set(rs)) == 1:
                continue                                               # DAPO: 신호 없는 그룹 제외
            As = group_advantages(rs, cfg['adv'])                      # ③ 그룹 안에서 상대 평가
            groups.append([(x, y, lm.token_logps(x, y), A) for y, A in zip(ys, As)])
        for _ in range(epochs):                                        # ④ 같은 그룹으로 몇 번 갱신
            for group in groups:
                lm.apply(grpo_loss_grads(lm, ref, group, cfg), lr / len(ANSWERS))
    return lm, ref

def hand_check():
    close = lambda xs, ys: all(abs(a - b) < 1e-3 for a, b in zip(xs, ys))
    assert close(group_advantages([1, 0, 0, 1]), [1, -1, -1, 1])
    assert group_advantages([1, 1, 1, 1]) == [0.0, 0.0, 0.0, 0.0]        # 모두 같으면 신호 없음
    assert close(group_advantages([1, 0, 0, 0]), [1.732, -0.577, -0.577, -0.577])
    print('A([1,0,0,1]) =', [round(a, 3) for a in group_advantages([1, 0, 0, 1])],
          '| A([1,0,0,0]) =', [round(a, 3) for a in group_advantages([1, 0, 0, 0])])
    lm = TinyLM()
    ys = [['1', '2', '<eos>'], ['7', '<eos>'], ['1', '<eos>'], ['1', '2', '<eos>']]
    rs = [verify('3×4=', y) for y in ys]
    assert rs == [1.0, 0.0, 0.0, 1.0]
    print('그룹:', [show(y) for y in ys], '보상:', rs)

A([1,0,0,1]) = [1.0, -1.0, -1.0, 1.0] | A([1,0,0,0]) = [1.732, -0.577, -0.577, -0.577]
그룹: ['12⏎', '7⏎', '1⏎', '12⏎'] 보상: [1.0, 0.0, 0.0, 1.0]
3×4=  P(정답) 0.160 → 0.933   KL(π‖π_ref) = 1.508
2×3=  P(정답) 0.240 → 0.965   KL(π‖π_ref) = 1.254

# -*- coding: utf-8 -*-
"""선택 실습 · Phase 11의 GRPO를 TRL로 실제 LLM에 적용한다.

실행 (GPU 권장, Colab T4 가능):
    pip install trl==0.24.0 transformers datasets
    python extras/trl_grpo.py                          # Qwen/Qwen2.5-0.5B-Instruct, 200 step
    python extras/trl_grpo.py <모델 이름> <max_steps>    # 예: 작은 테스트 모델로 2 step만

TinyLM과 바뀌는 것은 모델과 데이터 크기뿐이다. verifier는 tinylm.verify와 같은 역할을 한다.
"""
import random
import re
import sys

from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer

MODEL = sys.argv[1] if len(sys.argv) > 1 else 'Qwen/Qwen2.5-0.5B-Instruct'
MAX_STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 200


def make_dataset(n=512, seed=0):
    rng = random.Random(seed)
    rows = []
    for _ in range(n):
        a, b = rng.randint(2, 19), rng.randint(2, 19)
        rows.append({'prompt': [{'role': 'user', 'content': f'{a}×{b}는 얼마인가요? 숫자 하나로만 답하세요.'}],
                     'answer': str(a * b)})
    return Dataset.from_list(rows)


def verifier(completions, answer, **kwargs):
    """검증기 보상 (식 11.1의 r_i): 응답의 첫 정수가 정답이면 1, 아니면 0."""
    rewards = []
    for completion, gold in zip(completions, answer):
        text = completion[0]['content'] if isinstance(completion, list) else completion
        m = re.search(r'-?\d+', text)
        rewards.append(1.0 if m and m.group() == gold else 0.0)
    return rewards


args = GRPOConfig(
    output_dir='grpo-out',
    num_generations=8,                 # G: 질문 하나당 응답 수
    per_device_train_batch_size=8,     # 한 step에 질문 1개 × 8개 응답
    max_completion_length=32,
    learning_rate=1e-6,
    beta=0.04,                         # KL 계수 (TRL 0.24 기본값은 0.0)
    epsilon=0.2,                       # clip 폭. DAPO 식으로는 epsilon_high=0.28
    loss_type='grpo',                  # Phase 12: 'grpo' | 'dapo' | 'dr_grpo' | 'bnpo'
    scale_rewards='group',             # Phase 12: 'none'이면 Dr. GRPO 식 advantage
    importance_sampling_level='token', # Phase 12: 'sequence'이면 GSPO
    max_steps=MAX_STEPS,
    logging_steps=1,
    report_to='none',
)
trainer = GRPOTrainer(model=MODEL, reward_funcs=verifier, args=args, train_dataset=make_dataset())
trainer.train()

{'loss': 0.0, 'grad_norm': 0.0, 'completions/mean_length': 32.0, 'completions/clipped_ratio': 1.0,
 'rewards/verifier/mean': 0.0, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.0, ...}
{'loss': 0.0, 'grad_norm': 0.0, 'completions/mean_length': 32.0, 'completions/clipped_ratio': 1.0,
 'rewards/verifier/mean': 0.0, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.0, ...}

VARIANTS = {
    'GRPO':     dict(DEFAULT),
    'RLOO':     dict(DEFAULT, adv='rloo', beta=0.0),
    'Dr. GRPO': dict(DEFAULT, adv='mean', agg='const', beta=0.0),
    'DAPO':     dict(DEFAULT, agg='token_mean', eps_high=0.28, dynamic=True, beta=0.0),
    'GSPO':     dict(DEFAULT, ratio='sequence', beta=0.0),
}

def token_weights(advantages, lengths, agg):
    """토큰 하나가 loss에서 받는 가중치 (ratio = 1인 첫 갱신 기준). 응답 i의 모든 토큰이 같은 값."""
    G, total = len(lengths), sum(lengths)
    if agg == 'seq_mean':
        return [A / (G * L) for A, L in zip(advantages, lengths)]
    if agg == 'token_mean':
        return [A / total for A in advantages]
    return [A / (G * MAX_NEW_TOKENS) for A in advantages]       # 'const'

def hand_check():
    close = lambda xs, ys: all(abs(a - b) < 1e-3 for a, b in zip(xs, ys))
    rewards, lengths = [1, 0, 0, 1], [3, 2, 2, 3]              # 12⏎, 7⏎, 1⏎, 12⏎
    A_std = group_advantages(rewards, 'std')
    A_mean = group_advantages(rewards, 'mean')
    A_rloo = group_advantages(rewards, 'rloo')
    assert close(A_mean, [0.5, -0.5, -0.5, 0.5]) and close(A_rloo, [0.667, -0.667, -0.667, 0.667])
    table = {
        'GRPO (std, 응답별 평균)': token_weights(A_std, lengths, 'seq_mean'),
        'DAPO (std, 전체 토큰 평균)': token_weights(A_std, lengths, 'token_mean'),
        'Dr. GRPO (평균만, 상수 분모)': token_weights(A_mean, lengths, 'const'),
        'RLOO (나 뺀 평균, 응답별 평균)': token_weights(A_rloo, lengths, 'seq_mean'),
    }
    assert close(table['GRPO (std, 응답별 평균)'], [0.0833, -0.125, -0.125, 0.0833])
    assert close(table['DAPO (std, 전체 토큰 평균)'], [0.1, -0.1, -0.1, 0.1])
    assert close(table['Dr. GRPO (평균만, 상수 분모)'], [0.0417, -0.0417, -0.0417, 0.0417])
    for name, w in table.items():
        print(f'{name:22s} 토큰당 가중치 {[round(x, 4) for x in w]}')

    # DAPO clip-higher: A>0 인 토큰의 ratio가 1.25일 때
    from phase07_ppo import ppo_clip_term
    assert ppo_clip_term(1.25, 1.0, eps=0.2)[1] == 0.0                 # GRPO: 이미 멈춤
    hi = 1 + 0.28
    assert 1.25 < hi                                                    # DAPO: 상한 1.28 → 아직 학습

    # GSPO: 토큰 log ratio [0.10, −0.05, 0.02] → 응답 ratio는 기하평균
    logr = [0.10, -0.05, 0.02]
    s = math.exp(sum(logr) / len(logr))
    assert abs(s - 1.02360) < 1e-4
    print('GSPO 응답 ratio =', round(s, 4), '| 토큰 ratio =', [round(math.exp(x), 4) for x in logr])

GRPO (std, 응답별 평균)     토큰당 가중치 [0.0833, -0.125, -0.125, 0.0833]
DAPO (std, 전체 토큰 평균)   토큰당 가중치 [0.1, -0.1, -0.1, 0.1]
Dr. GRPO (평균만, 상수 분모)  토큰당 가중치 [0.0417, -0.0417, -0.0417, 0.0417]
RLOO (나 뺀 평균, 응답별 평균)  토큰당 가중치 [0.0556, -0.0833, -0.0833, 0.0556]
GSPO 응답 ratio = 1.0236 | 토큰 ratio = [1.1052, 0.9512, 1.0202]
같은 seed, 30회 반복 뒤 (장난감 문제라 순위는 의미가 없다):
  GRPO     평균 P(정답) 0.868   평균 KL 1.128
  RLOO     평균 P(정답) 0.734   평균 KL 0.778
  Dr. GRPO 평균 P(정답) 0.600   평균 KL 0.441
  DAPO     평균 P(정답) 0.879   평균 KL 1.167
  GSPO     평균 P(정답) 0.899   평균 KL 1.215