Train from feedback

"""Fit a human-preference reward head and perform a real PPO policy update.

This entry point requires human-exported vote files. It never creates labels.
"""
import argparse
import copy
import hashlib
import json
import random
import sys
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F

parser = argparse.ArgumentParser()
parser.add_argument('--source', type=Path, required=True)
parser.add_argument('--votes', type=Path, nargs='+', required=True)
parser.add_argument('--feedback-kind', choices=['human', 'ai'], default='human')
parser.add_argument('--device', default='cpu')
args = parser.parse_args()
sys.path.insert(0, str(args.source))
from tiny_gpt.checkpoint import load_checkpoint, save_checkpoint
from tiny_gpt.tokenizer import CharacterTokenizer

torch.set_num_threads(4)
tokenizer = CharacterTokenizer()
public = Path('public/tiny-llm')
checkpoint = Path('scripts/tiny-llm/runs/finetune/best.pt')
run = Path('scripts/tiny-llm/runs/rlhf')
if run.exists(): raise RuntimeError('RL run already exists. Preserve it and declare a new experiment before rerunning.')
def sha(path):
    if path.suffix == '.pt':
        metadata = torch.load(path, map_location='cpu', weights_only=True)['metadata']
        if metadata.get('restored_from_arrays'): return metadata['origin_checkpoint_sha256']
    return hashlib.sha256(path.read_bytes()).hexdigest()
write = lambda p, o: p.write_text(json.dumps(o, indent=2, allow_nan=False) + '\n')
collection = json.loads((public / 'preferences.json').read_text())
if collection['checkpoint_sha256'] != sha(checkpoint): raise ValueError('Candidates belong to a different supervised policy')
pairs = {p['id']: p for p in collection['pairs']}
labels, seen, provenance = [], set(), []
for path in args.votes:
    file = json.loads(path.read_text())
    if file.get('format') != 1: raise ValueError('Unknown preference format')
    if args.feedback_kind == 'ai' and file.get('candidate_file_sha256') != sha(public / 'preferences.json'):
        raise ValueError('AI labels refer to a different candidate corpus')
    if args.feedback_kind == 'human' and (file.get('consent') is not True or file.get('human_attestation') is not True):
        raise ValueError('Human votes require explicit consent and human attestation')
    if args.feedback_kind == 'ai' and (file.get('feedback_kind') != 'ai' or file.get('human_attestation') is not False or not file.get('labeler')):
        raise ValueError('AI preferences must identify their labeler and must not claim human attestation')
    if file.get('checkpoint_sha256') != collection['checkpoint_sha256'] or not (file.get('participant') or file.get('labeler')):
        raise ValueError('Vote file provenance is missing or mismatched')
    provenance.append({'vote_file_sha256': sha(path), 'votes': len(file['votes'])})
    for vote in file['votes']:
        participant = file.get('participant', file.get('labeler'))
        key = (participant, vote['pair_id'])
        if key in seen: raise ValueError('Duplicate participant vote')
        seen.add(key)
        if vote['pair_id'] not in pairs or vote['choice'] not in ['a', 'b', 'tie', 'reject'] or not vote.get('timestamp'):
            raise ValueError('Unknown or malformed human vote')
        labels.append({**vote, 'participant': participant, 'pair': pairs[vote['pair_id']]})
binary = [l for l in labels if l['choice'] in ['a', 'b']]
splits = {split: [l for l in binary if l['pair']['split'] == split] for split in ['train', 'validation', 'test']}
counts = {split: len(values) for split, values in splits.items()}
if counts['train'] < 30 or counts['validation'] < 4 or counts['test'] < 8:
    raise ValueError(f'Insufficient genuine binary labels: {counts}. Need at least 30 train, 4 validation and 8 test.')
run.mkdir(parents=True)
write(run / 'labels.json', {'feedback_kind': args.feedback_kind, 'provenance': provenance, 'labels': labels})
torch.manual_seed(91)
encoder, _ = load_checkpoint(checkpoint, device=args.device)
encoder.requires_grad_(False)
reward = nn.Linear(encoder.config.width, 1).to(args.device)

@torch.no_grad()
def features(prompt, output):
    ids = tokenizer.encode(prompt, add_bos=True) + output
    if len(ids) > 32 or not output: raise ValueError('Candidate does not fit the original context')
    return encoder.hidden_states(torch.tensor([ids], device=args.device))[0, -1].detach().clone()

feature_cache = {}
for pair in pairs.values():
    for side in ['a', 'b']:
        feature_cache[(pair['id'], side)] = features(pair['prompt'], pair[side]['token_ids'])
def tensors(values):
    preferred = torch.stack([feature_cache[(l['pair_id'], l['choice'])] for l in values])
    rejected = torch.stack([feature_cache[(l['pair_id'], 'b' if l['choice'] == 'a' else 'a')] for l in values])
    return preferred, rejected
train_a, train_b = tensors(splits['train'])
validation_a, validation_b = tensors(splits['validation'])
optimizer = torch.optim.AdamW(reward.parameters(), lr=0.001)
generator = torch.Generator().manual_seed(91)
best, best_state, history = float('inf'), None, []
for epoch in range(1, 101):
    order = torch.randperm(len(train_a), generator=generator)
    losses = []
    for start in range(0, len(order), 16):
        chosen = order[start:start + 16].to(args.device)
        loss = F.softplus(-(reward(train_a[chosen]) - reward(train_b[chosen]))).mean()
        optimizer.zero_grad(); loss.backward(); optimizer.step(); losses.append(float(loss.detach()))
    with torch.no_grad(): val_loss = float(F.softplus(-(reward(validation_a) - reward(validation_b))).mean())
    if val_loss < best: best, best_state = val_loss, copy.deepcopy(reward.state_dict())
    history.append({'epoch': epoch, 'training_pair_loss': sum(losses) / len(losses), 'validation_pair_loss': val_loss})
reward.load_state_dict(best_state)
reward.eval().requires_grad_(False)
torch.save(reward.state_dict(), run / 'reward-head.pt')
test_a, test_b = tensors(splits['test'])
with torch.no_grad():
    differences = (reward(test_a) - reward(test_b)).flatten()
    heldout = {'examples': len(differences), 'correct': int((differences > 0).sum()), 'score_ties': int((differences == 0).sum()), 'pair_loss': float(F.softplus(-differences).mean())}
reference, _ = load_checkpoint(checkpoint, device=args.device)
reference.requires_grad_(False)
policy, _ = load_checkpoint(checkpoint, device=args.device)
value = nn.Linear(policy.config.width, 1).to(args.device)
torch.manual_seed(101)
# Initialize the value head after declaring the PPO seed.
value.reset_parameters()
optimizer = torch.optim.AdamW(list(policy.parameters()) + list(value.parameters()), lr=0.00001)
rng = random.Random(101)
prompts = sorted(set(l['pair']['prompt'] for l in splits['train']))

@torch.no_grad()
def rollout(prompt):
    ids = tokenizer.encode(prompt, add_bos=True)
    context, response, logp, values, kls, entropies = [], [], [], [], [], []
    for _ in range(min(16, 32 - len(ids))):
        x = torch.tensor([ids], device=args.device)
        hidden = policy.hidden_states(x)
        logits = policy.lm_head(hidden)[0, -1]
        p = logits.softmax(-1)
        selected = int(torch.multinomial(p, 1))
        current_logp = logits.log_softmax(-1)
        ref_logp = reference(x)[0, -1].log_softmax(-1)
        context.append(list(ids)); response.append(selected)
        logp.append(current_logp[selected].detach())
        values.append(value(hidden[0, -1]).squeeze().detach())
        kls.append((p * (current_logp - ref_logp)).sum().detach())
        entropies.append(-(p * current_logp).sum().detach())
        ids.append(selected)
        if selected == 2: break
    score = reward(features(prompt, response)).squeeze().detach()
    # Undiscounted terminal reward with per-step reference KL penalties.
    rewards = [-0.05 * k for k in kls]
    rewards[-1] = rewards[-1] + score
    returns, cumulative = [], torch.zeros((), device=args.device)
    for r in reversed(rewards): cumulative = cumulative + r; returns.insert(0, cumulative.clone())
    return {'contexts': context, 'response': response, 'old_logp': logp, 'old_values': values, 'returns': returns, 'reward': float(score), 'kl': float(torch.stack(kls).mean()), 'entropy': float(torch.stack(entropies).mean())}

ppo_history = []
for update in range(1, 31):
    trajectories = [rollout(rng.choice(prompts)) for _ in range(16)]
    samples = [(context, selected, old_logp, ret, ret - old_value)
               for t in trajectories for context, selected, old_logp, ret, old_value in
               zip(t['contexts'], t['response'], t['old_logp'], t['returns'], t['old_values'])]
    advantages = torch.stack([s[4] for s in samples])
    mean, std = advantages.mean(), advantages.std(unbiased=False).clamp_min(1e-8)
    clipping, measured_kl, measured_entropy, losses = [], [], [], []
    for _ in range(2):
        optimizer.zero_grad()
        total = torch.zeros((), device=args.device)
        for context, selected, old_logp, ret, advantage in samples:
            x = torch.tensor([context], device=args.device)
            hidden = policy.hidden_states(x)
            logits = policy.lm_head(hidden)[0, -1]
            current = logits.log_softmax(-1)
            p = current.exp()
            with torch.no_grad(): ref_logp = reference(x)[0, -1].log_softmax(-1)
            kl = (p * (current - ref_logp)).sum()
            entropy = -(p * current).sum()
            ratio = (current[selected] - old_logp).exp()
            normalized = (advantage - mean) / std
            clipped = ratio.clamp(0.8, 1.2)
            policy_loss = -torch.minimum(ratio * normalized, clipped * normalized)
            value_loss = (value(hidden[0, -1]).squeeze() - ret).square()
            loss = policy_loss + 0.5 * value_loss + 0.05 * kl - 0.001 * entropy
            total = total + loss / len(samples)
            clipping.append(float((ratio.detach() != clipped.detach()).float()))
            measured_kl.append(float(kl.detach())); measured_entropy.append(float(entropy.detach()))
        if not bool(torch.isfinite(total)): raise RuntimeError('PPO loss became non-finite')
        total.backward()
        torch.nn.utils.clip_grad_norm_(list(policy.parameters()) + list(value.parameters()), 1, error_if_nonfinite=True)
        optimizer.step(); losses.append(float(total.detach()))
    entry = {'update': update, 'reward': sum(t['reward'] for t in trajectories) / 16, 'reference_kl': sum(measured_kl) / len(measured_kl), 'clipping_fraction': sum(clipping) / len(clipping), 'entropy': sum(measured_entropy) / len(measured_entropy), 'mean_output_length': sum(len(t['response']) for t in trajectories) / 16, 'loss': sum(losses) / len(losses)}
    ppo_history.append(entry); write(run / 'ppo-history.json', ppo_history)
    print('PPO', update, 'reward', entry['reward'], 'KL', entry['reference_kl'], flush=True)
policy.eval()
save_checkpoint(policy, run / 'policy.pt', metadata={'stage': args.feedback_kind + '-reward-ppo', 'steps': 60, 'seed': 101, 'supervised_sha256': sha(checkpoint), 'label_file_sha256': [p['vote_file_sha256'] for p in provenance]})
torch.save(value.state_dict(), run / 'value-head.pt')
write(public / 'rlhf.json', {'status': 'trained; frozen evaluation and policy export still required', 'supervised_sha256': sha(checkpoint), 'policy_sha256': sha(run / 'policy.pt'), 'feedback_kind': args.feedback_kind, 'label_files': provenance, 'binary_counts': counts, 'ties': sum(l['choice'] == 'tie' for l in labels), 'reject_both': sum(l['choice'] == 'reject' for l in labels), 'reward_validation_loss': best, 'reward_heldout': heldout, 'reward_history': history, 'ppo_history': ppo_history, 'reward_source': args.feedback_kind + ' preferences; scalar reward head with frozen supervised encoder', 'test_used_for_selection': False})
print('Saved actual PPO policy. Evaluate arithmetic, preference outcome and confidence before publication.')