좋은 답변은
어떻게 학습되는가?

def robot_values(p_bridge=0.4, p_careful=0.75, gamma=0.9):
    q_careful, q_rush = 8.0, 0.5*10 + 0.5*(-6)
    v1 = p_careful*q_careful + (1-p_careful)*q_rush
    q_safe, q_bridge = 2.0, -1.0 + gamma*v1
    v0 = p_bridge*q_bridge + (1-p_bridge)*q_safe
    return dict(V0=v0, V1=v1, Q_safe=q_safe, Q_bridge=q_bridge,
                A_bridge=q_bridge-v0, A_safe=q_safe-v0,
                A_careful=q_careful-v1, A_rush=q_rush-v1)

def discounted_returns(rewards, gamma=0.9):
    out, carry = [], 0.0
    for r in reversed(rewards):
        carry = r + gamma*carry
        out.append(carry)
    return list(reversed(out))

import torch

def reinforce_walkthrough(action=1, use_baseline=False):
    # action=1: 왼쪽 8점, action=0: 오른쪽 2점.
    # rollout에서 이 행동이 관측되었다고 가정한 한 번의 재현 예제.
    z = torch.tensor(0.0, requires_grad=True)
    optimizer = torch.optim.SGD([z], lr=0.1)
    policy = torch.distributions.Bernoulli(logits=z)
    # 실제 수집 시: observed_action = policy.sample()
    observed_action = torch.tensor(float(action))
    reward = 8.0 if action == 1 else 2.0
    # 한 번 행동하고 종료했으므로 return G = reward.
    G = torch.tensor(reward)
    p_before = torch.sigmoid(z).detach()
    # 교육용으로 환경의 두 보상을 알아 정확한 V를 계산한다.
    # 실제 환경에서는 별도 추정값을 사용할 수 있다.
    baseline = (2.0 + 6.0 * p_before) if use_baseline else 0.0
    signal = (G - baseline).detach()
    loss = -signal * policy.log_prob(observed_action)
    optimizer.zero_grad()
    loss.backward()
    loss_gradient = z.grad.item()
    optimizer.step()
    return p_before.item(), loss_gradient, torch.sigmoid(z).item()

# 왼쪽 기본형: (0.5, -4.0, 약 0.598688)
# 오른쪽 기본형: (0.5, +1.0, 약 0.475021)
# 평균 5를 빼면 양쪽 모두: (0.5, -1.5, 약 0.537430)
# print(reinforce_walkthrough(action=1, use_baseline=False))

def reinforce_demo():
    # Two answers: incorrect 7 / correct 12. This is a one-step policy.
    logits = torch.nn.Parameter(torch.zeros(2))
    opt = torch.optim.SGD([logits], lr=0.2)
    before = logits.softmax(0).detach().tolist()
    for _ in range(120):
        dist = torch.distributions.Categorical(logits=logits)
        action = dist.sample((64,))
        reward = (action == 1).float()
        # Baseline known exactly for this toy; detach before policy loss.
        baseline = logits.softmax(0)[1].detach()
        loss = -(dist.log_prob(action) * (reward - baseline)).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return before, logits.softmax(0).detach().tolist()

def actor_critic_one_step():
    # S0: choose bridge (1). State values initially match the example.
    z = torch.nn.Parameter(torch.tensor(math.log(0.4/0.6)))
    values = torch.nn.Parameter(torch.tensor([3.14, 6.5]))
    actor_opt = torch.optim.SGD([z], lr=0.1)
    critic_opt = torch.optim.SGD([values], lr=0.1)
    # Observe r=-1, next state=S1; freeze the TD target.
    target = (-1 + 0.9*values[1]).detach()
    advantage = (target-values[0]).detach()  # +1.71
    logp = F.logsigmoid(z)                 # log P(bridge)
    actor_loss = -logp*advantage
    critic_loss = 0.5*(values[0]-target).square()
    actor_opt.zero_grad(); actor_loss.backward(); actor_opt.step()
    critic_opt.zero_grad(); critic_loss.backward(); critic_opt.step()
    return float(z.sigmoid().detach()), values.detach().tolist()

def gae(rewards, values, done, gamma=1.0, lam=0.95):
    # One trajectory: rewards [T], values [T+1], done [T].
    advantage = torch.zeros_like(rewards)
    carry = torch.zeros((), dtype=rewards.dtype)
    for t in reversed(range(len(rewards))):
        alive = 1.0 - done[t]
        delta = rewards[t] + gamma * values[t+1] * alive - values[t]
        carry = delta + gamma * lam * alive * carry
        advantage[t] = carry
    returns = advantage + values[:-1]
    return advantage.detach(), returns.detach()

def trpo_scalar(delta=0.01):
    # ONE trainable Bernoulli logit at S0, hold S1 policy fixed.
    # This is an analytic miniature, not a scalable full TRPO trainer.
    p_old = 0.4
    z_old = math.log(p_old/(1-p_old))
    q_difference = 4.85-2.0
    g = q_difference*p_old*(1-p_old)
    fisher = p_old*(1-p_old)  # scalar Hessian of KL at old policy
    direction = g/fisher
    full_step = math.sqrt(2*delta/(g*direction))*direction
    rows = []
    for fraction in [1.0, 0.5, 0.25, 0.125, 0.0625]:
        p = 1/(1+math.exp(-(z_old+fraction*full_step)))
        kl = p_old*math.log(p_old/p)+(1-p_old)*math.log((1-p_old)/(1-p))
        improvement = (p-p_old)*q_difference
        expected = g*fraction*full_step
        accept = kl <= delta and improvement >= 0.1*expected
        rows.append(dict(fraction=fraction,p=p,kl=kl,gain=improvement,accept=accept))
        if accept:
            return rows
    return rows  # if no candidate passes, retain old parameters

def ppo_loss(new_lp, old_lp, advantage, mask, eps=0.2):
    # old_lp is the policy that GENERATED this batch, not the SFT reference.
    ratio = (new_lp - old_lp.detach()).exp()
    a = advantage.detach()
    unclipped = ratio * a
    clipped = ratio.clamp(1-eps, 1+eps) * a
    return -(torch.minimum(unclipped, clipped) * mask).sum() / mask.sum()

def ppo_total(new_lp, old_lp, a, mask, value, returns,
              entropy, value_coef=0.5, entropy_coef=0.01):
    actor = ppo_loss(new_lp, old_lp, a, mask)
    critic = (((value - returns.detach()) ** 2) * mask).sum() / mask.sum()
    ent = (entropy * mask).sum() / mask.sum()
    # If KL was included in rewards before GAE, don't add it again here.
    return actor + value_coef * critic - entropy_coef * ent

def robot_ppo(rounds=50):
    # Two states, actions 0/1. S0: safe/bridge; S1: rush/careful.
    torch.manual_seed(23)
    logits = torch.nn.Parameter(torch.tensor([math.log(.4/.6), math.log(.75/.25)]))
    values = torch.nn.Parameter(torch.zeros(2))
    opt = torch.optim.Adam([logits, values], lr=0.04)
    gamma, lam = 0.9, 0.95
    history = []
    for iteration in range(rounds):
        states, actions, old_lps, advantages, targets = [], [], [], [], []
        # Collection and target construction have NO autograd graph.
        with torch.no_grad():
            for _ in range(128):
                trajectory = []
                state = 0
                while True:
                    dist = torch.distributions.Bernoulli(logits=logits[state])
                    action = dist.sample()
                    if state == 0:
                        reward, terminal = ((2.,True) if action == 0 else (-1.,False))
                    else:
                        reward = 8. if action == 1 else (10. if torch.rand(()) < .5 else -6.)
                        terminal = True
                    trajectory.append((state, action, dist.log_prob(action),reward))
                    if terminal: break
                    state = 1
                rs = torch.tensor([row[3] for row in trajectory])
                vs = torch.stack([values[row[0]] for row in trajectory]+[torch.tensor(0.)])
                done = torch.zeros(len(rs)); done[-1] = 1
                adv, ret = gae(rs, vs, done, gamma, lam)
                for j,(s,a,lp,r) in enumerate(trajectory):
                    states.append(s); actions.append(a); old_lps.append(lp)
                    advantages.append(adv[j]); targets.append(ret[j])
        s = torch.tensor(states)
        a, old = torch.stack(actions), torch.stack(old_lps)
        adv, ret = torch.stack(advantages), torch.stack(targets)
        # Full-buffer updates for clarity; a real trainer shuffles minibatches.
        for _ in range(4):
            dist = torch.distributions.Bernoulli(logits=logits[s])
            actor_loss = ppo_loss(dist.log_prob(a), old, adv, torch.ones_like(adv))
            value_loss = 0.5*(values[s]-ret).square().mean()
            loss = actor_loss + value_loss - 0.01*dist.entropy().mean()
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_([logits,values],0.5)
            opt.step()
        p,q = logits.sigmoid().detach().tolist()
        history.append((iteration+1,p,q,robot_values(p,q)['V0']))
    return history

def response_logps(logits, input_ids, response_mask):
    # logits[b,t] predicts input_ids[b,t+1]. Mask labels, not inputs.
    logp = logits[:, :-1].log_softmax(-1)
    target = input_ids[:, 1:]
    chosen = logp.gather(-1, target.unsqueeze(-1)).squeeze(-1)
    mask = response_mask[:, 1:].to(chosen.dtype)
    return chosen, mask, (chosen * mask).sum(-1)

def sft_loss(logits, ids, response_mask):
    token_lp, mask, _ = response_logps(logits, ids, response_mask)
    return -(token_lp * mask).sum() / mask.sum().clamp_min(1)

def reward_model_demo():
    # Learn scalar scores for 3 fixed answers from human comparisons.
    # A: polite but 7; B: terse 12; C: explains 3*4=12.
    score = torch.nn.Parameter(torch.zeros(3))
    opt = torch.optim.SGD([score], lr=0.2)
    chosen = torch.tensor([2, 1, 2])
    rejected = torch.tensor([0, 0, 1])
    for _ in range(80):
        loss = -F.logsigmoid(score[chosen] - score[rejected]).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return score.detach()

def tiny_rlhf_ppo():
    # Complete RM -> rollout -> value baseline -> PPO loop on 3 answers.
    rm = reward_model_demo()                 # freeze learned human reward
    logits = torch.nn.Parameter(torch.zeros(3))
    critic = torch.nn.Parameter(torch.tensor(0.0))
    ref_lp = torch.full((3,), -math.log(3))  # fixed SFT reference
    opt = torch.optim.Adam([logits, critic], lr=0.03)
    for _ in range(60):
        with torch.no_grad():
            old_all = logits.log_softmax(0)
            actions = torch.distributions.Categorical(logits=logits).sample((96,))
            old_lp = old_all[actions]
            reward = rm[actions] - 0.1 * (old_lp - ref_lp[actions])
            advantage = reward - critic.detach()
        for _ in range(3):                  # reuse this frozen rollout batch
            new_lp = logits.log_softmax(0)[actions]
            actor = ppo_loss(new_lp, old_lp, advantage, torch.ones_like(reward))
            loss = actor + 0.5 * (critic - reward).square().mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
    return logits.softmax(0).detach().tolist()

def dpo_loss(chosen_lp, rejected_lp, ref_chosen_lp, ref_rejected_lp, beta=0.1):
    # All inputs are RESPONSE sequence log-probability SUMS, shape [B].
    chosen_ratio = chosen_lp - ref_chosen_lp.detach()
    rejected_ratio = rejected_lp - ref_rejected_lp.detach()
    return -F.logsigmoid(beta * (chosen_ratio - rejected_ratio)).mean()

def dpo_numeric():
    return dpo_loss(torch.tensor([math.log(0.5)]), torch.tensor([math.log(0.1)]),
                    torch.tensor([math.log(0.25)]), torch.tensor([math.log(0.25)]),
                    beta=1.0).item()  # 0.182322 = -log(5/6)

def kto_loss(policy_lp, ref_lp, desirable, mismatch_lp, mismatch_ref_lp,
             beta=1.0, lambda_good=1.0, lambda_bad=1.0):
    # mismatched (x_i, y_j), j != i; needs separate model forward passes.
    # This is the paper's convenient BIASED KL-reference estimator.
    z = (mismatch_lp - mismatch_ref_lp).mean().clamp_min(0).detach()
    logratio = policy_lp - ref_lp.detach()
    good = lambda_good * (1 - torch.sigmoid(beta * (logratio - z)))
    bad = lambda_bad * (1 - torch.sigmoid(beta * (z - logratio)))
    return torch.where(desirable.bool(), good, bad).mean()

def verify_integer(text, gold=12):
    # Toy contract: exactly one final tag, containing a signed integer.
    tags = re.findall(r"<final>(.*?)</final>", text, flags=re.S)
    if len(tags) != 1 or not re.fullmatch(r"\s*[+-]?\d+\s*", tags[0]):
        return 0.0
    return float(int(tags[0]) == gold)

samples = ["3+4=7 <final>7</final>", "3*4=12 <final>12</final>",
           "<final>12</final>", "3*4=14 <final>14</final>"]
# [0, 1, 1, 0]; correct final answer does NOT certify the reasoning.

def grpo_loss(new_lp, old_lp, a, mask, ref_lp=None, beta=0.0, eps=0.2):
    ratio = (new_lp - old_lp.detach()).exp()
    a = a.detach().unsqueeze(-1)
    obj = torch.minimum(ratio*a, ratio.clamp(1-eps, 1+eps)*a)
    if beta:
        assert ref_lp is not None
        # Nonnegative sampled KL surrogate used in GRPO.
        # Off-policy batches make this an approximation, not exact KL.
        d = ref_lp.detach() - new_lp
        kl = d.exp() - d - 1
        obj = obj - beta * kl
    return -seq_mean(obj, mask)

def policy_weight_steps():
    # RLHF: one chosen token vs all other tokens, one term of a 3-token mean.
    z = torch.nn.Parameter(torch.tensor(math.log(.55/.45)))
    opt = torch.optim.SGD([z],lr=.1)
    old_logp = torch.tensor(math.log(.5))
    a = torch.tensor(.728685677)
    rho = (F.logsigmoid(z)-old_logp).exp()
    loss = -torch.minimum(rho*a,rho.clamp(.8,1.2)*a)/3
    opt.zero_grad(); loss.backward()
    rlhf_gradient = z.grad.item()
    opt.step()
    rlhf_probability = z.sigmoid().detach().item()
    # RLVR: four fixed actions, one sample of each, [bad,good,good,bad].
    logits = torch.nn.Parameter(torch.zeros(4))
    opt = torch.optim.SGD([logits],lr=.2)
    old_lp = torch.full((4,1),-math.log(4))
    a = torch.tensor([-1.,1.,1.,-1.])
    loss = grpo_loss(logits.log_softmax(0)[:,None],old_lp,a,torch.ones(4,1))
    opt.zero_grad(); loss.backward()
    gradient = logits.grad.detach().tolist()
    opt.step()
    return dict(rlhf_logit_gradient=rlhf_gradient,rlhf_new_probability=rlhf_probability,
                rlvr_logit_gradient=gradient,rlvr_new_probability=logits.softmax(0).detach().tolist())

def group_advantage(rewards, normalize=True):
    # [number_of_prompts, G]; population std, correction=0.
    centered = rewards - rewards.mean(-1, keepdim=True)
    if normalize:
        centered = centered / rewards.std(-1, keepdim=True, unbiased=False).clamp_min(1e-8)
    return centered.detach()

def rloo_advantage(rewards):
    g = rewards.shape[-1]
    assert g >= 2
    baseline = (rewards.sum(-1, keepdim=True) - rewards) / (g-1)
    return (rewards - baseline).detach()

def rloo_loss(sequence_lp, rewards):
    # [prompts, G], fresh on-policy samples; sequence_lp is a token SUM.
    return -(rloo_advantage(rewards) * sequence_lp).mean()

def dr_grpo_loss(new_lp, old_lp, raw_a, mask, max_length, eps=0.2):
    # raw_a = R - group_mean(R), WITHOUT dividing by group std.
    ratio = (new_lp - old_lp.detach()).exp()
    a = raw_a.detach().unsqueeze(-1)
    obj = torch.minimum(ratio*a, ratio.clamp(1-eps, 1+eps)*a)
    return -(obj * mask).sum() / (new_lp.shape[0] * max_length)

def dapo_loss(new_lp, old_lp, a, mask, eps_low=0.2, eps_high=0.28):
    # Call after retaining mixed-correctness groups; no explicit KL here.
    ratio = (new_lp - old_lp.detach()).exp()
    a = a.detach().unsqueeze(-1)
    obj = torch.minimum(ratio*a, ratio.clamp(1-eps_low, 1+eps_high)*a)
    return -(obj * mask).sum() / mask.sum()   # global TOKEN mean

def useful_groups(binary_rewards):
    correct = binary_rewards.sum(-1)
    return (correct > 0) & (correct < binary_rewards.shape[-1])

def overlong_penalty(length, max_length, soft_window):
    assert 0 < soft_window <= max_length
    return -min(1.0, max(0.0, (length-(max_length-soft_window))/soft_window))

def gspo_loss(new_lp, old_lp, a, mask, eps=3e-4):
    # Full sequence geometric-mean ratio, not arithmetic token mean.
    mean_logratio = ((new_lp-old_lp.detach())*mask).sum(-1) / mask.sum(-1)
    ratio = mean_logratio.exp()
    a = a.detach()
    obj = torch.minimum(ratio*a, ratio.clamp(1-eps, 1+eps)*a)
    return -obj.mean()   # differentiation through ratio gives 1/length

def sapo_loss(new_lp, old_lp, a, mask, tau_pos=1.0, tau_neg=1.05):
    ratio = (new_lp-old_lp.detach()).exp()
    a = a.detach().unsqueeze(-1)
    tau = torch.where(a > 0, tau_pos, tau_neg)
    surrogate = (4/tau) * torch.sigmoid(tau*(ratio-1))
    return -seq_mean(surrogate*a, mask)

def pnpo_loss(new_lp, behavior_lp, a, mask, eps_low=7e-4, eps_high=9.5e-4):
    # Assumes valid response tokens followed by RIGHT padding.
    t = torch.arange(1, new_lp.shape[-1]+1, device=new_lp.device)
    with torch.no_grad():
        d = (new_lp-behavior_lp) * mask
        weight = (d.cumsum(-1) / t).exp()
        h = (mask.sum(-1, keepdim=True) / t).sqrt()
        gate = (weight >= 1-eps_low*h) & (weight <= 1+eps_high*h)
        coefficient = gate * weight * a.unsqueeze(-1)
    # Stop gradient through weights; do NOT divide by accepted-token count.
    return -seq_mean(coefficient * new_lp, mask)

def gapo_upper(correct, group_size, low=0.003, high_max=0.005):
    assert group_size > 1 and 1 <= correct <= group_size
    return low + (high_max-low)*(group_size-correct)/(group_size-1)

def gapo_group_loss(new_lp, old_lp, rewards, mask, low=0.003, high_max=0.005):
    # One prompt's G responses; main sequence-ratio surrogate (Eq.12).
    # Token-aggregation variants in experiments are a separate choice.
    c, g = int(rewards.sum().item()), len(rewards)
    a = (rewards-rewards.mean()).detach()   # no std normalization
    if c == 0 or c == g:
        return new_lp.sum()*0              # no group reward signal
    high = gapo_upper(c, g, low, high_max)
    s = (((new_lp-old_lp.detach())*mask).sum(-1)/mask.sum(-1)).exp()
    return -torch.minimum(s*a, s.clamp(1-low, 1+high)*a).mean()

def cliff_advantage(rewards, pitfall_token, mask, teacher_ok=True, lam=0.0):
    # One prompt. pitfall_token is first wrong token, ZERO-based.
    # A teacher identifies steps; preprocessing maps step boundaries to tokens.
    base = group_advantage(rewards[None])[0]
    if not teacher_ok:
        return base[:, None].expand_as(mask).detach()
    mu = rewards.mean()
    if mu <= 0 or mu >= 1:
        return torch.zeros_like(mask)  # explicit toy handling: skip zero variance
    sd = (mu*(1-mu)).sqrt()
    a_good, a_bad = (1-mu)/sd, -mu/sd
    t = torch.arange(mask.shape[-1], device=mask.device)[None, :]
    raw = torch.where(rewards[:,None] == 1, a_good,
                      torch.where(t < pitfall_token[:,None], lam*a_good, a_bad))
    offset = (raw*mask).sum()/mask.sum()
    return ((raw-offset)*mask).detach()

python3 -m venv .venv
# macOS / Linux
source .venv/bin/activate
# Windows: .venv\Scriptsctivate
python -m pip install torch==2.8.0 numpy==2.0.2
python rl_examples.py

def tiny_rlvr_grpo():
    # A complete RLVR loop over four fixed strings: teaches probabilities,
    # NOT arithmetic generalization or language generation.
    logits = torch.nn.Parameter(torch.zeros(4))
    opt = torch.optim.Adam([logits], lr=0.04)
    for _ in range(100):
        with torch.no_grad():
            ids = torch.distributions.Categorical(logits=logits).sample((16,))
            rewards = torch.tensor([verify_integer(samples[i]) for i in ids.tolist()])
            a = group_advantage(rewards[None])[0]
            old_lp = logits.log_softmax(0)[ids, None].clone()
        for _ in range(2):
            new_lp = logits.log_softmax(0)[ids, None]
            loss = grpo_loss(new_lp, old_lp, a, torch.ones_like(new_lp))
            opt.zero_grad()
            loss.backward()
            opt.step()
    probs = logits.softmax(0).detach()
    return probs.tolist(), float(probs[1]+probs[2])

Q update: 0.464
REINFORCE P(correct): 0.5000 → 0.9775
Reward model scores A/B/C: [-1.8876, 0.0000, 1.8876]
RLHF + PPO probabilities A/B/C: [0.0073, 0.0049, 0.9878]
DPO loss: 0.182322
RLVR total correct probability: 0.9909