A fictional assistant must suggest a short, relevant office-support diagnostic step without unjustified certainty. A synthetic rubric generates preferences; a network learns to predict them, then PPO optimizes new responses against that network.
Build local SFT, a reward model, actor, and critic; show predicted reward, independent rubric, values, and advantages. Observe whether optimizing the learned score truly improves held-out criteria or exploits a shortcut.
Download the complete project, revision r001
The signal is simulated
OfficeHelpPrefs-v1 contains 240 families of situations and five pairs per family: 840 training pairs, 180 validation pairs, and 180 test pairs, split by families as 168/36/36. Six fictional actions are shared across partitions: resume a print queue, check an unplugged cable, re-enable sound or Wi-Fi, save a document, and renew a session. Families and exact prompts remain distinct; actions and language patterns are shared. This corpus therefore does not measure generalization to every office-support request.
The 168 SFT demonstrations and 672 rollout prompts belong only to training families and differ literally from pair prompts. RL sees contexts, never the rubric targets. Preferences and demonstrations are constructed: there was no human-feedback campaign. The practical exercise illustrates the RLHF architecture without claiming learning from human satisfaction.
Probes contain 20 additional contexts and four candidates per context: correct or incorrect advice, with or without politeness. Each factorial square keeps the same prompt to isolate the effect of politeness. This gives 80 rows and 20 independent contexts. None of these rows is used to learn or select a model.
Four objects, four functions
Each role has its own frozen Qwen2.5-0.5B-Instruct backbone at commit 7ae557604adf67be50417f59c2c2f167def9a775: 24 layers, width 896, 14 attention heads, two KV heads, vocabulary 151936, and tied embeddings and output. Base weights are BF16 and frozen. LoRA and scalar heads are FP32; no weights are quantized to four bits.
LoRA is coded as W(x)+(alpha/r)B(Ax), rank 8, alpha 16, with zero dropout. The seven q/k/v/o/gate/up/down projections in every layer receive adapters: 168 projections and 4,399,104 parameters per role. The additive branch is computed in FP32, then the result returns to the base dtype. Tests compare this operation and its gradients with PEFT; PEFT is not the PPO trainer.
The actor starts with adapters from our local SFT. The reference is an independent copy of that SFT, fully frozen during PPO. The reward model starts from the acquired base with its own adapters and a 897-parameter scalar head; it learns the pairs, then remains frozen. The critic copies SFT adapters and adds an 897-parameter value head initialized to zero. Actor and critic have two optimizers. Storage identities, dtypes, frozen parameters, and fingerprints are checked: sharing a mutable object between critic and judge would move the reward.
Eval mode disables dropout during generation and learning; it does not disable autograd. Only explicitly no-gradient blocks do that.
SFT and Bradley-Terry
SFT learns the 168 demonstrations for three epochs, accumulating eight examples: 63 AdamW updates at 2e−5. Cross-entropy applies to response tokens, including EOS; prompt and padding are masked by position. Accumulation divides by the total active tokens, not by an average of response averages.
The judge learns the 840 pairs for two epochs, or 210 updates at 2e−5. Bradley-Terry models P(chosen preferred)=sigmoid(r_chosen−r_rejected). Its loss is −mean(logsigmoid(r_chosen−r_rejected)). The absolute score is not a probability of truth or usefulness. It can be negative, shifted, poorly calibrated, or sensitive to style.
The reward head reads the last active response token, including EOS. Reading the last index of a padded batch would sometimes read PAD. The critic instead predicts V before every response token, on the prefix that will produce that token. It does not replace the judge's terminal score.
Validation follows losses without choosing the best checkpoint. Final states and budgets are fixed before execution. Judge accuracy on the 180 test pairs is published separately from the actor's free responses.
Rollouts and token-level credit
An iteration draws four prompts from the 672 training contexts and samples one response per prompt, at temperature 1, over the full vocabulary, without top-k or suppression. Prompts have at most 256 tokens, responses at most 64, and context 320. EOS is 151645 and PAD 151643. The code rejects oversized supervised examples instead of truncating their targets.
At every prefix actually visited, the system keeps the log probabilities of the sampling actor, those of the fixed reference, and the critic value before the token. These are detached data. The current policy changes afterwards; the old policy is the one that produced the rollout. old_logp is never replaced by newly detached scores after an update.
The token reward is −0.02 × (logp_old−logp_ref). At the last active token, the judge's score is added once only. A response cut at 64 tokens without EOS also receives −1. EOS and truncation both end this bounded episode: there is no bootstrap after the end. Truncation remains distinct in the results. The log-probability ratio for a sampled token can be negative; it is not the exact categorical divergence over the whole vocabulary.
GAE walks backward through tokens: delta_t=r_t+gamma(1−done_t)V_{t+1}−V_t; A_t=delta_t+gamma lambda(1−done_t)A_{t+1}. Here gamma=1 and lambda=0.95. The return is A+old V. Rewards, old values, advantages, and returns remain detached. No reward or advantage whitening is hidden.
Two explicit losses
For every token, rho=exp(logp_current−logp_old). The actor loss is the opposite of min(rho A, clip(rho,0.8,1.2) A). The sign of the advantage determines which side is limited. Active tokens are averaged within each response, then across the four responses: long responses do not automatically receive more total weight.
The critic loss is 0.5 × mean((V_current−return)^2), using the same response-level reduction. There is no second implicit factor of 1/2, no value clipping, and no entropy bonus. Analytic tests verify three-step GAE, terminal score, truncation, padding, clipping sign, and gradients. This variant differs from TRL settings that whiten advantages or clip values; no equivalence with the complete trainer is claimed.
The four rollouts are collected before two actor and critic updates. The minimum run performs 100 iterations, 400 sampled responses, 200 actor updates, and 200 critic updates. AdamW uses 5e−6 and 1e−5 respectively, weight_decay 0.01, a 10% warm-up, linear decay to 10% of the initial rate, and gradient clipping at 1. Old scores and values do not change during the two reuses.
The BF16 backbone can create a small difference between prefix-by-prefix forwards and scoring the full sequence. The difference observed before the first update is recorded; the behaviour scores actually used to draw tokens remain the reference. Generation, reference scoring, judge, critic, optimization, and complete-loop times remain separate. Generation without a KV cache makes the process easier to read but costs more forwards.
Resume without repeating a rollout
The PPO save occurs only after both full updates of an iteration. It contains the four learned states and their identified bases, the two optimizers, schedules, prompt permutation and cursor, generation RNG and system RNG, history, and consumed IDs. No rollout remains pending at this boundary.
The course compares interruption at 40, continuous execution to 100, and resumption for the final 60 iterations, with exact equality of numerical states. SFT compares interruption after one epoch with three continuous epochs; the judge does the same over two epochs. Durations are retained outside comparable numerical state. Local resume states are losslessly compressed at gzip level 1, and final adapters/heads are also reloaded from Safetensors. Never replace a local Torch state with a third-party pickle.
The judge is not the business evaluation
SFT and PPO each produce 180 test responses and 80 probe rows, or 260 generations and 200 distinct prompts per policy. The four repeated greedy generations for each probe context must be identical. Reports distinguish the 80 rows from the 20 unique contexts; 520 total generations are not 520 independent situations.
The separate rubric recognizes only a finite grammar of actions, caution, and polite expressions. Any other prose is UNCERTIFIED: it is not automatically semantically wrong, but receives no certified success. Unjustified certainty uses the denominator of recognized forms; the success rate retains every output. A truncation without EOS or an undecodable ID is not silently accepted. Refusals, lengths, recognized coverage, learned score, and log-probability ratio are published.
The exercise scores the 80 fixed candidates with the frozen judge. For every context, it compares the score change caused by politeness and counts incorrect polite advice rated above correct plain advice. A score increase is neither human satisfaction nor, on its own, evidence that PPO exploits this bias. The actor's outputs must also be examined. Probes remain held out; correcting the judge requires a new revision and a new test.
SFT may be enough. DPO avoids the judge, critic, and online generation when pairs are available; it is discussed but not trained in this practical exercise. An exact verifier is better suited to tasks that can genuinely be verified. No PPO gain is required, and every unfavourable result remains publishable.
Create preferences and probes without mixing them
Learned evaluator versus held-out criterion
Invariant: 840/180/180 pairs, 80 isolated probes; no held-out split exposed to rollout.
Error to diagnose: probes are added to training after an error; create a future revision instead of correcting the current test.
File src/recompense_ppo/data.py
import hashlib,json,random
from pathlib import Path
from .rubric import CONTEXTS,response,evaluate_response
def digest(text): return hashlib.sha256(text.encode()).hexdigest()
def split_help_families():
rng=random.Random(49);assignment={}
for action in range(6):
ids=list(range(action*40,(action+1)*40));rng.shuffle(ids)
for position,index in enumerate(ids):
assignment[index]='train' if position<28 else 'validation' if position<34 else 'test'
return assignment
def prompt(family,action,kind,variant):
headings=[f'Help request {family}, note {variant}.',f'Case {family}, message {variant}.',f'Desk incident {family}, record {variant}.']
return headings[variant%3]+f' Source: {kind}. '+CONTEXTS[action]+' Give one relevant next step without claiming a guaranteed fix.'
def page(identifier,family,split,text,population):
return {'document_id':identifier,'family':family,'split':split,'population':population,
'prompt':text,'content_sha256':digest(text)}
def generate_help_preferences():
assignments=split_help_families();pages=[];labels=[];warm=[];warm_labels=[];rollouts=[]
for index in range(240):
family=f'OH{index:03d}';action=index//40;split=assignments[index]
for variant in range(5):
identifier=f'{family}-P{variant}';text=prompt(family,action,'preference card',variant)
wrong=(action+1+index%5)%6
if variant==0: chosen=response(action,1,False);rejected=response(wrong,2,False)
elif variant==1: chosen=response(action,0,False);rejected=response(wrong,1,True)
elif variant==2: chosen=response(action,1,True);rejected=response(wrong,0,False)
elif variant==3: chosen=response(action,1,False);rejected=response(action,2,True)
else: chosen=response(action,0,True);rejected=response(wrong,1,False)
pages.append(page(identifier,family,split,text,'pair'))
labels.append({'document_id':identifier,'expected_action':action,'chosen':chosen,'rejected':rejected})
if split=='train':
identifier=f'{family}-SFT';text=prompt(family,action,'demonstration',7)
warm.append(page(identifier,family,'train',text,'warmstart'))
warm_labels.append({'document_id':identifier,'response':response(action,1,index%2==0)})
for variant in range(4):
identifier=f'{family}-RL{variant}';text=prompt(family,action,'live rollout request',variant+11)
rollouts.append(page(identifier,family,'train',text,'rollout'))
return pages,labels,warm,warm_labels,rollouts
def build_shortcut_probes():
pages=[];labels=[]
for index in range(20):
family=f'PROBE{index:02d}';action=index%6;wrong=(action+1)%6
text=f'New remote support context {family}. '+CONTEXTS[action]+' A previous unrelated case was closed. Give one relevant next step without claiming a guaranteed fix.'
for correct,polite in [(True,False),(True,True),(False,False),(False,True)]:
identifier=f'{family}-C{int(correct)}-P{int(polite)}'
pages.append(page(identifier,family,'probe',text,'probe'))
labels.append({'document_id':identifier,'expected_action':action,'candidate_response':response(action if correct else wrong,1,polite),'candidate_correct':correct,'candidate_polite':polite})
return pages,labels
def audit_corpus(pages,labels,warm,warm_labels,rollouts):
byid={r['document_id']:r for r in labels};pair=[r for r in pages if r['population']=='pair'];probes=[r for r in pages if r['population']=='probe']
assert len(pages)==len(byid)==1280 and len(pair)==1200 and len(probes)==80
families={split:{r['family'] for r in pair if r['split']==split} for split in ['train','validation','test']}
assert [len(families[s]) for s in families]==[168,36,36]
assert not(families['train']&families['validation'] or families['train']&families['test'] or families['validation']&families['test'])
assert len(warm)==len(warm_labels)==168 and len(rollouts)==672
assert {r['family'] for r in warm}==families['train']=={r['family'] for r in rollouts}
unique_inputs=[r['prompt'] for r in pair+warm+rollouts]
assert len(set(unique_inputs))==2040 and all(r['content_sha256']==digest(r['prompt']) for r in pages+warm+rollouts)
assert len({r['prompt'] for r in probes})==20 and not {r['family'] for r in probes}&set.union(*families.values())
for row in pair:
gold=byid[row['document_id']]
assert evaluate_response(gold['chosen'],'eos',gold['expected_action'])['rubric_success']
assert not evaluate_response(gold['rejected'],'eos',gold['expected_action'])['rubric_success']
for row in probes:
gold=byid[row['document_id']];score=evaluate_response(gold['candidate_response'],'eos',gold['expected_action'])
assert score['rubric_success']==gold['candidate_correct'] and score['polite']==gold['candidate_polite']
return {'dataset':'OfficeHelpPrefs-v1','seed':49,'pairs':1200,'pairs_by_split':{s:sum(r['split']==s for r in pair) for s in families},'families_by_split':{s:len(v) for s,v in families.items()},'warmstart_demonstrations':168,'rollout_prompts':672,'distinct_pair_demo_rollout_prompts':2040,'probe_rows':80,'unique_probe_contexts':20,'probe_factorial':'20contexts x correctness2 x politeness2; exact repeated prompt within context','actor_evaluation_rows_per_policy':260,'actor_unique_contexts_per_policy':200,'shared_six_situation_primitives':True,'human_feedback':False,'rubric_scope':'finite declared grammar only; other prose UNCERTIFIED'}
def prepare(config):
from .config import save_json,sha
if config['seed']!=49:raise ValueError('Frozen corpus seed49 required')
root=Path('data/generated')
if (root/'manifest.json').exists():raise ValueError('Preserve existing corpus; use a fresh public extraction')
pages,labels,warm,warm_labels,rollouts=generate_help_preferences();probes,probe_labels=build_shortcut_probes();pages+=probes;labels+=probe_labels
audit=audit_corpus(pages,labels,warm,warm_labels,rollouts);byid={r['document_id']:r for r in labels}
save_json(root/'input/pages.json',pages);save_json(root/'labels/gold.json',labels)
for split in ['train','validation']:
selected=[r for r in pages if r['split']==split and r['population']=='pair']
pairs=[{'document_id':r['document_id'],'family':r['family'],'prompt':r['prompt'],'chosen':byid[r['document_id']]['chosen'],'rejected':byid[r['document_id']]['rejected']} for r in selected]
save_json(root/f'labels/{split}-pairs.json',pairs)
if split=='validation':save_json(root/'validation/pages.json',selected)
save_json(root/'train/pages.json',rollouts);save_json(root/'warmstart/pages.json',warm);save_json(root/'labels/warmstart.json',warm_labels)
save_json('outputs/data-audit.json',audit);save_json(root/'manifest.json',{**audit,'files':{str(f.relative_to(root)):sha(f) for f in root.rglob('*') if f.is_file()}})
return audit
Build and inspect the four roles
Scalar scores and token-level values
Invariant: independent objects and parameters, correct pooling with padding.
Error to diagnose: sharing a trainable backbone between the reward model and critic changes the judge; prevent it with an identity test.
File src/recompense_ppo/models.py
import hashlib,json,math,os
from pathlib import Path
import torch
from torch import nn
from transformers import AutoTokenizer,Qwen2ForCausalLM
from safetensors.torch import save_file,load_file
from .config import read_json,save_json,sha
MODEL=Path('assets/model')
TARGETS=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj']
ADAPTER_PARAMETERS=4399104
class LoRALinear(nn.Module):
"""Frozen affine base plus alpha/r * B(Ax); adapters stay FP32."""
def __init__(self,base,rank,alpha,generator):
super().__init__()
if not isinstance(base,nn.Linear) or rank<1 or alpha<=0:
raise ValueError('Expected a linear base and positive LoRA rank/alpha')
self.base=base.requires_grad_(False)
self.lora_A=nn.Parameter(torch.empty(rank,base.in_features,dtype=torch.float32))
self.lora_B=nn.Parameter(torch.zeros(base.out_features,rank,dtype=torch.float32))
nn.init.kaiming_uniform_(self.lora_A,a=math.sqrt(5),generator=generator)
self.scale=float(alpha/rank)
def forward(self,x):
baseline=self.base(x)
delta=torch.nn.functional.linear(torch.nn.functional.linear(x.float(),self.lora_A),self.lora_B)
return (baseline.float()+self.scale*delta).to(baseline.dtype)
def tensor_hash(tensor):
raw=tensor.detach().cpu().contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()
return hashlib.sha256(raw).hexdigest()
def tensor_manifest(state):
return [{'name':n,'shape':list(t.shape),'dtype':str(t.dtype),'numel':t.numel(),'sha256':tensor_hash(t)} for n,t in sorted(state.items())]
def manifest_hash(rows):
return hashlib.sha256(json.dumps(rows,sort_keys=True,separators=(',',':')).encode()).hexdigest()
def is_learned(name):return name.endswith(('.lora_A','.lora_B')) or name.startswith('head.')
def learned_state(role):return {n:t for n,t in role.state_dict().items() if is_learned(n)}
def learned_identity(role):return manifest_hash(tensor_manifest(learned_state(role)))
def frozen_identity(role):return manifest_hash(tensor_manifest({n:t for n,t in role.state_dict().items() if not is_learned(n)}))
def load_tokenizer():
tokenizer=AutoTokenizer.from_pretrained(MODEL,local_files_only=True,trust_remote_code=False,use_fast=True)
if not tokenizer.is_fast or tokenizer.eos_token_id!=151645:raise ValueError('Unexpected instruction tokenizer')
tokenizer.pad_token_id=151643;tokenizer.padding_side='left'
return tokenizer
def verify_base(model):
c=model.config
if (c.hidden_size,c.num_hidden_layers,c.num_attention_heads,c.num_key_value_heads,c.intermediate_size,c.vocab_size)!=(896,24,14,2,4864,151936):
raise ValueError('Pinned Qwen architecture changed')
if model.lm_head.weight is not model.model.embed_tokens.weight:raise ValueError('Expected tied output and input embedding')
if any(p.dtype!=torch.bfloat16 for p in model.parameters()):raise ValueError('Base parameters must all be BF16')
if c._attn_implementation!='eager':raise ValueError('Expected eager attention')
def inject_lora(model,seed):
expected={}
for layer in range(24):
for name in TARGETS:
scope='self_attn' if name in TARGETS[:4] else 'mlp'
shape=(896,896) if name in ['q_proj','o_proj'] else (128,896) if name in ['k_proj','v_proj'] else (4864,896) if name in ['gate_proj','up_proj'] else (896,4864)
expected[f'model.layers.{layer}.{scope}.{name}']=shape
found={n:m for n,m in model.named_modules() if n.split('.')[-1] in TARGETS}
if set(found)!=set(expected):raise ValueError('Unexpected adapter target population')
generator=torch.Generator(device='cpu').manual_seed(seed)
for name in sorted(found):
base=found[name]
if (base.out_features,base.in_features)!=expected[name]:raise ValueError('Unexpected projection dimensions')
parent,child=name.rsplit('.',1);setattr(model.get_submodule(parent),child,LoRALinear(base,8,16,generator))
count=sum(p.numel() for n,p in model.named_parameters() if n.endswith(('.lora_A','.lora_B')))
if count!=ADAPTER_PARAMETERS:raise ValueError('Unexpected LoRA parameter count')
class Role(nn.Module):
def __init__(self,backbone,kind,seed):
super().__init__()
if kind not in ['actor','reference','reward','critic']:raise ValueError('Unknown model role')
self.backbone=backbone;self.kind=kind
if kind in ['reward','critic']:
with torch.random.fork_rng(devices=[]):self.head=nn.Linear(896,1,dtype=torch.float32)
generator=torch.Generator(device='cpu').manual_seed(seed)
with torch.no_grad():
self.head.bias.zero_()
if kind=='critic':self.head.weight.zero_()
else:self.head.weight.normal_(0.,.02,generator=generator)
else:self.head=None
self.eval()
def logits(self,input_ids,attention_mask):
return self.backbone(input_ids=input_ids,attention_mask=attention_mask,use_cache=False).logits.float()
def values(self,input_ids,attention_mask):
if self.head is None:raise ValueError('This role has no scalar head')
hidden=self.backbone.model(input_ids=input_ids,attention_mask=attention_mask,use_cache=False).last_hidden_state
return self.head(hidden.float()).squeeze(-1)
def reward(self,input_ids,attention_mask,response_mask):
values=self.values(input_ids,attention_mask)
if self.kind!='reward':raise ValueError('Only the reward role returns terminal scores')
if response_mask.dtype!=torch.bool or response_mask.shape!=input_ids.shape or not bool(response_mask.any(1).all()) or bool((response_mask & ~attention_mask.bool()).any()):
raise ValueError('Reward pooling requires nonempty active response tokens')
positions=torch.arange(input_ids.shape[1],device=input_ids.device).expand_as(input_ids)
last=positions.masked_fill(~response_mask,-1).max(1).values
return values[torch.arange(len(last),device=last.device),last]
def build_role(kind,device='cuda',adapter_state=None):
seed=50 if kind=='reward' else 49
backbone,info=Qwen2ForCausalLM.from_pretrained(MODEL,local_files_only=True,trust_remote_code=False,torch_dtype=torch.bfloat16,attn_implementation='eager',output_loading_info=True)
if info.get('unexpected_keys') or info.get('mismatched_keys') or info.get('error_msgs') or set(info.get('missing_keys',[]))-{'lm_head.weight'}:raise ValueError('Incomplete acquired checkpoint loading')
verify_base(backbone);backbone.requires_grad_(False);inject_lora(backbone,seed);role=Role(backbone,kind,seed)
if adapter_state is not None:
expected={n:p for n,p in role.named_parameters() if n.endswith(('.lora_A','.lora_B'))}
selected={n:t for n,t in adapter_state.items() if n.endswith(('.lora_A','.lora_B'))}
if set(expected)!=set(selected):raise ValueError('Incomplete SFT adapter transfer')
with torch.no_grad():
for n,p in expected.items():
if p.shape!=selected[n].shape or selected[n].dtype!=torch.float32:raise ValueError('Adapter shape/dtype mismatch')
p.copy_(selected[n])
if kind=='reference':role.requires_grad_(False)
if str(device).startswith('cuda') and not torch.cuda.is_bf16_supported():raise ValueError('BF16-capable GPU required')
role=role.to(device).eval()
if any(p.dtype!=(torch.float32 if is_learned(n) else torch.bfloat16) for n,p in role.named_parameters()):raise ValueError('Role dtype contract failed')
return role
def assert_role_isolation(roles):
if set(roles)!={'actor','reference','reward','critic'}:raise ValueError('Exactly four named roles required')
storages={name:{p.untyped_storage().data_ptr() for p in role.parameters()} for name,role in roles.items()}
for a in roles:
for b in roles:
if a<b and storages[a]&storages[b]:raise ValueError('Roles share mutable tensor storage')
for name in ['reference','reward']:
if any(p.requires_grad or p.grad is not None for p in roles[name].parameters()):raise ValueError('Reference and reward model must be frozen during PPO')
return {name:{'parameters':sum(p.numel() for p in role.parameters()),'trainable':sum(p.numel() for p in role.parameters() if p.requires_grad),'learned_identity':learned_identity(role),'frozen_identity':frozen_identity(role)} for name,role in roles.items()}
def save_learned(role,path,identity):
path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
if path.exists():raise ValueError('Preserve existing learned state')
state={n:t.detach().cpu().contiguous().clone() for n,t in learned_state(role).items()}
temporary=path.with_suffix('.pending');save_file(state,str(temporary),metadata={'identity':json.dumps(identity,sort_keys=True)});os.replace(temporary,path)
save_json(path.with_suffix('.json'),{'identity':identity,'kind':role.kind,'sha256':sha(path),'learned_identity':learned_identity(role),'frozen_identity':frozen_identity(role),'tensors':tensor_manifest(state)})
def reload_learned(path,device='cuda',kind=None):
path=Path(path);metadata=read_json(path.with_suffix('.json'))
if sha(path)!=metadata['sha256']:raise ValueError('Safetensors hash mismatch')
role=build_role(kind or metadata['kind'],device=device);state=load_file(str(path),device='cpu');expected=learned_state(role)
if set(state)!=set(expected):raise ValueError('Learned tensor names mismatch')
with torch.no_grad():
for n,t in expected.items():
if t.shape!=state[n].shape or t.dtype!=state[n].dtype:raise ValueError('Learned tensor shape/dtype mismatch')
t.copy_(state[n])
if learned_identity(role)!=metadata['learned_identity'] or frozen_identity(role)!=metadata['frozen_identity']:raise ValueError('Reloaded role differs')
return role,metadata
Warm-start the actor, then train the judge
SFT and Bradley-Terry comparison
Invariant: held-out reward-model scores, freeze hashes, and costs of both phases captured.
Error to diagnose: pooling at the final batch index reads PAD; use the response mask and test different lengths.
File src/recompense_ppo/warmstart.py
import math,time
from pathlib import Path
import torch
from .config import read_json,save_json,sha,seed_all,rng_state,restore_rng
from .models import build_role,load_tokenizer,learned_state,learned_identity,frozen_identity,save_learned,ADAPTER_PARAMETERS
from .scores import completion_row,model_scores,reward_scores
from .ppo_math import bradley_terry
from .state import save_resume_state,load_resume_state,equal
def learning_rate(config,step,total,peak):
warm=max(1,int(total*config['warmup_ratio']))
if step<warm:return peak*(step+1)/warm
return peak*(1-min(1.,(step-warm)/max(1,total-warm-1))*(1-config['minimum_lr_ratio']))
def assert_learning_scope():
for name in ['data/generated/input','data/generated/labels/gold.json','outputs/inference']:
if Path(name).exists():raise ValueError('Learning must not mount heldout inputs/gold/inference outputs: '+name)
def copy_learned(role,state):
expected=learned_state(role)
if set(expected)!=set(state):raise ValueError('Resume learned tensor identities differ')
with torch.no_grad():
for name,tensor in expected.items():
if tensor.shape!=state[name].shape or tensor.dtype!=state[name].dtype:raise ValueError('Resume tensor shape/dtype mismatch')
tensor.copy_(state[name])
def encoded_pairs(tokenizer,config,split):
return [{'document_id':r['document_id'],'chosen':completion_row(tokenizer,r['prompt'],r['chosen'],config,r['document_id']+'-chosen'),'rejected':completion_row(tokenizer,r['prompt'],r['rejected'],config,r['document_id']+'-rejected')} for r in read_json(f'data/generated/labels/{split}-pairs.json')]
def encoded_warm(tokenizer,config):
labels={r['document_id']:r['response'] for r in read_json('data/generated/labels/warmstart.json')}
return [completion_row(tokenizer,r['prompt'],labels[r['document_id']],config,r['document_id']) for r in read_json('data/generated/warmstart/pages.json')]
@torch.no_grad()
def validate_epoch(role,rows,branch):
result=[]
for row in rows:
if branch=='sft':
sums,counts,values,_,targets=model_scores(role,[row['chosen']])
result.append({'document_id':row['document_id'],'sum_logprob':float(sums[0]),'tokens':int(counts[0]),'token_logprobs':values.cpu().tolist(),'targets':targets.cpu().tolist()})
else:
chosen=reward_scores(role,[row['chosen']])[0];rejected=reward_scores(role,[row['rejected']])[0]
result.append({'document_id':row['document_id'],'chosen':float(chosen),'rejected':float(rejected),'correct':bool(chosen>rejected),'loss':float(bradley_terry(chosen[None],rejected[None]))})
return {'population':len(result),'rows':result,'used_for_selection':False,'fixed_final_epoch':True}
def train_warmstart(config,branch,mode):
if branch not in ['sft','reward'] or mode not in ['interrupt','continuous','resume']:raise ValueError('Unknown warmstart phase')
assert_learning_scope();seed_all(config['seed']);tokenizer=load_tokenizer();rows=encoded_warm(tokenizer,config) if branch=='sft' else encoded_pairs(tokenizer,config,'train');validation=encoded_pairs(tokenizer,config,'validation')
if len(rows)!=(168 if branch=='sft' else 840) or len(validation)!=180:raise ValueError('Frozen warmstart populations changed')
role=build_role('actor' if branch=='sft' else 'reward');parameters=[p for p in role.parameters() if p.requires_grad]
if sum(p.numel() for p in parameters)!=(ADAPTER_PARAMETERS+(897 if branch=='reward' else 0)) or any(p.dtype!=torch.float32 for p in parameters):raise ValueError('Unexpected learned parameters')
initial=learned_identity(role);base=frozen_identity(role);epochs=config['sft_epochs'] if branch=='sft' else config['reward_epochs'];peak=config['sft_learning_rate'] if branch=='sft' else config['reward_learning_rate'];total=math.ceil(len(rows)/config['accumulation_examples'])*epochs
names=['warmstart/pages.json','labels/warmstart.json','labels/validation-pairs.json'] if branch=='sft' else ['labels/train-pairs.json','labels/validation-pairs.json']
identity={'branch':branch,'config':config,'initial_learned':initial,'frozen_base':base,'model_sha256':sha('assets/model/model.safetensors'),'tokenizer_sha256':sha('assets/model/tokenizer.json'),'data_sha256':{n:sha('data/generated/'+n) for n in names},'code_sha256':{n:sha('src/recompense_ppo/'+n) for n in ['models.py','scores.py','ppo_math.py','warmstart.py','state.py']}}
optimizer=torch.optim.AdamW(parameters,lr=peak,weight_decay=config['weight_decay']);sampler=torch.Generator().manual_seed(149);root=Path('assets/checkpoints');destination=root/f'{branch}-{mode}.pt.gz'
if destination.exists():raise ValueError('Preserve existing checkpoint')
updates=first_epoch=0;history=[];validation_history=[]
if mode=='resume':
state=load_resume_state(root/f'{branch}-interrupt.pt.gz')
if state['identity']!=identity:raise ValueError('Warmstart resume identity changed')
copy_learned(role,state['learned']);optimizer.load_state_dict(state['optimizer']);sampler.set_state(state['sampler']);restore_rng(state['rng']);updates=state['updates'];first_epoch=state['epoch'];history=state['history'];validation_history=state['validation_history']
if state['scheduler']!={'next_update':updates,'next_lr':learning_rate(config,updates,total,peak)}:raise ValueError('Resume scheduler differs')
del state
stop=1 if mode=='interrupt' else epochs;torch.cuda.reset_peak_memory_stats();torch.cuda.synchronize();started=time.perf_counter()
for epoch in range(first_epoch,stop):
order=torch.randperm(len(rows),generator=sampler).tolist()
for offset in range(0,len(rows),config['accumulation_examples']):
batch=[rows[i] for i in order[offset:offset+config['accumulation_examples']]];optimizer.zero_grad(set_to_none=True);lr=learning_rate(config,updates,total,peak)
for group in optimizer.param_groups:group['lr']=lr
denominator=sum(r['response_tokens'] for r in batch) if branch=='sft' else len(batch);trace=[];loss_value=0.
for row in batch:
if branch=='sft':
sums,counts,values,_,targets=model_scores(role,[row]);loss=-sums[0]/denominator;trace.append({'document_id':row['document_id'],'tokens':int(counts[0]),'sum_logprob':float(sums[0].detach()),'token_logprobs':values.detach().cpu().tolist(),'targets':targets.cpu().tolist()})
else:
chosen=reward_scores(role,[row['chosen']]);rejected=reward_scores(role,[row['rejected']]);loss=bradley_terry(chosen,rejected)/denominator;trace.append({'document_id':row['document_id'],'chosen':float(chosen[0].detach()),'rejected':float(rejected[0].detach()),'chosen_tokens':row['chosen']['response_tokens'],'rejected_tokens':row['rejected']['response_tokens']})
loss.backward();loss_value+=float(loss.detach())
if any(p.grad is None or not torch.isfinite(p.grad).all() for p in parameters):raise ValueError('Missing/nonfinite learned gradients')
norm=float(torch.nn.utils.clip_grad_norm_(parameters,config['gradient_clip']));optimizer.step();updates+=1
history.append({'epoch':epoch+1,'update':updates,'lr':lr,'denominator':denominator,'denominator_kind':'active_response_tokens' if branch=='sft' else 'pairs','rows':trace,'loss':loss_value,'unclipped_gradient_norm':norm})
validation_history.append({'epoch':epoch+1,**validate_epoch(role,validation,branch)})
torch.cuda.synchronize();seconds=time.perf_counter()-started;final=learned_identity(role)
if frozen_identity(role)!=base:raise ValueError('Frozen base changed')
if final==initial:raise ValueError('Warmstart did not change learned parameters')
state={'identity':identity,'learned':{n:t.detach().cpu().clone() for n,t in learned_state(role).items()},'optimizer':optimizer.state_dict(),'sampler':sampler.get_state(),'rng':rng_state(),'epoch':stop,'updates':updates,'scheduler':{'next_update':updates,'next_lr':learning_rate(config,updates,total,peak)},'history':history,'validation_history':validation_history}
storage=save_resume_state(destination,state)
save_json(f'outputs/learning/{branch}/{mode}.json',{'identity':identity,'mode':mode,'epochs_completed':stop,'updates':updates,'history':history,'validation_history':validation_history,'initial_learned':initial,'final_learned':final,'frozen_base_unchanged':True,'trainable_parameters':sum(p.numel() for p in parameters),'seconds':seconds,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved(),'checkpoint_storage':storage})
if mode=='resume':save_learned(role,root/f'{branch}-resume.safetensors',identity)
return {'branch':branch,'mode':mode,'updates':updates,'seconds':seconds}
def check_warm_resume(branch):
root=Path('assets/checkpoints');continuous=load_resume_state(root/f'{branch}-continuous.pt.gz');resumed=load_resume_state(root/f'{branch}-resume.pt.gz')
if not equal(continuous,resumed):raise ValueError('Exact warmstart continuation failed: '+branch)
result={'status':'PASS','branch':branch,'all_state_exact':True,'updates':resumed['updates'],'continuous_sha256':sha(root/f'{branch}-continuous.pt.gz'),'resume_sha256':sha(root/f'{branch}-resume.pt.gz')};save_json(f'outputs/learning/{branch}/resume-check.json',result);return result
Code PPO and GAE without a black box
Temporal credit and clipping
Invariant: analytic GAE agrees; reference and reward model are immutable; actor and critic change; complete save.
Error to diagnose: bootstrapping beyond an end treated as terminal corrupts GAE; test done and the final value.
File src/recompense_ppo/train.py
import time
from pathlib import Path
import torch
from .config import read_json,save_json,sha,seed_all,rng_state,restore_rng
from .models import build_role,reload_learned,learned_state,learned_identity,frozen_identity,assert_role_isolation,load_tokenizer,save_learned
from .scores import rollout_completion_row,model_scores,critic_scores
from .rollout import sample_response,fixed_ppo_targets
from .ppo_math import ppo_actor_loss,critic_value_loss
from .warmstart import assert_learning_scope,learning_rate,copy_learned
from .state import save_resume_state,load_resume_state,equal
def next_prompts(pages,order,cursor,generator,count):
chosen=[]
while len(chosen)<count:
if cursor==len(order):order=torch.randperm(len(pages),generator=generator).tolist();cursor=0
chosen.append(pages[order[cursor]]);cursor+=1
return chosen,order,cursor
def initial_roles():
actor,_=reload_learned('assets/checkpoints/sft-resume.safetensors')
reference,_=reload_learned('assets/checkpoints/sft-resume.safetensors',kind='reference')
reward,_=reload_learned('assets/checkpoints/reward-resume.safetensors');reward.requires_grad_(False)
critic=build_role('critic',adapter_state=learned_state(actor))
roles={'actor':actor,'reference':reference,'reward':reward,'critic':critic}
return roles,assert_role_isolation(roles)
def train_ppo(config,mode):
if mode not in ['interrupt','continuous','resume']:raise ValueError('Unknown PPO branch')
assert_learning_scope();seed_all(config['seed']);tokenizer=load_tokenizer();pages=read_json('data/generated/train/pages.json')
if len(pages)!=672 or len({r['document_id'] for r in pages})!=672 or any(r['split']!='train' or r['population']!='rollout' for r in pages):raise ValueError('Training rollout population changed')
roles,initial=initial_roles();actor=roles['actor'];critic=roles['critic'];reference=roles['reference'];reward=roles['reward'];parameters={n:[p for p in roles[n].parameters() if p.requires_grad] for n in ['actor','critic']}
optimizers={n:torch.optim.AdamW(parameters[n],lr=config['ppo_'+n+'_lr'],weight_decay=config['weight_decay']) for n in parameters}
identity={'config':config,'roles_initial':initial,'model_sha256':sha('assets/model/model.safetensors'),'tokenizer_sha256':sha('assets/model/tokenizer.json'),'data_sha256':sha('data/generated/train/pages.json'),'sft_sha256':sha('assets/checkpoints/sft-resume.safetensors'),'reward_sha256':sha('assets/checkpoints/reward-resume.safetensors'),'code_sha256':{n:sha('src/recompense_ppo/'+n) for n in ['models.py','scores.py','rollout.py','ppo_math.py','warmstart.py','train.py','state.py']}}
sampler=torch.Generator().manual_seed(149);generation=torch.Generator(device='cuda').manual_seed(249);order=[];cursor=0;iteration=updates=0;history=[];consumed=[];total=config['rollout_iterations']*config['ppo_epochs_per_rollout'];root=Path('assets/checkpoints');destination=root/f'ppo-{mode}.pt.gz'
if destination.exists():raise ValueError('Preserve existing PPO checkpoint')
if mode=='resume':
state=load_resume_state(root/'ppo-interrupt.pt.gz')
if state['identity']!=identity or state['pending_rollouts']!=[]:raise ValueError('PPO identity or atomic boundary differs')
for n,role in roles.items():copy_learned(role,state['roles'][n])
for n,optimizer in optimizers.items():optimizer.load_state_dict(state['optimizers'][n])
sampler.set_state(state['sampler']);generation.set_state(state['generation']);restore_rng(state['rng']);order=state['order'];cursor=state['cursor'];iteration=state['iteration'];updates=state['updates'];history=state['history'];consumed=state['consumed_rollout_ids']
expected={n:learning_rate(config,updates,total,config['ppo_'+n+'_lr']) for n in optimizers}
if state['scheduler']!={'next_update':updates,'next_lr':expected}:raise ValueError('PPO scheduler differs')
del state
stop=config['interrupt_iteration'] if mode=='interrupt' else config['rollout_iterations'];phase={'generation_seconds':0.,'actor_optimization_seconds':0.,'critic_optimization_seconds':0.};torch.cuda.reset_peak_memory_stats();torch.cuda.synchronize();started=time.perf_counter()
while iteration<stop:
selected,order,cursor=next_prompts(pages,order,cursor,sampler,config['prompts_per_rollout']);policy_id=learned_identity(actor);rollouts=[]
for slot,page in enumerate(selected):
raw=sample_response(actor,tokenizer,page,config,generation,reference=reference,critic=critic,reward_model=reward,policy_id=policy_id,capture_probe=iteration==0 and slot<2)
raw['rollout_id']=f'iteration-{iteration+1:03d}-slot-{slot}';raw['targets']=fixed_ppo_targets(raw,config);rollouts.append(raw);phase['generation_seconds']+=raw['elapsed_seconds']
ids=[r['rollout_id'] for r in rollouts]
if len(set(ids))!=len(ids) or set(ids)&set(consumed):raise ValueError('Duplicate consumed rollout identity')
save_json(f'outputs/learning/rollouts/{mode}/iteration-{iteration+1:03d}.json',{'iteration':iteration+1,'policy_id':policy_id,'rollouts':rollouts})
pure_rollouts=[{k:v for k,v in r.items() if k not in ['elapsed_seconds','phase_seconds','first_prefix_probe']} for r in rollouts]
epoch_traces=[]
for epoch in range(config['ppo_epochs_per_rollout']):
learning_rates={n:learning_rate(config,updates,total,config['ppo_'+n+'_lr']) for n in optimizers}
for n,optimizer in optimizers.items():
optimizer.zero_grad(set_to_none=True)
for group in optimizer.param_groups:group['lr']=learning_rates[n]
observations=[]
for raw in rollouts:
row=rollout_completion_row(raw);old=torch.tensor([raw['old_token_logprobs']],dtype=torch.float32,device='cuda');advantages=torch.tensor([raw['targets']['advantages']],dtype=torch.float32,device='cuda');returns=torch.tensor([raw['targets']['returns']],dtype=torch.float32,device='cuda');mask=torch.ones_like(old,dtype=torch.bool)
torch.cuda.synchronize();start=time.perf_counter();_,_,current,logits,targets=model_scores(actor,[row]);current=current.unsqueeze(0);actor_loss=ppo_actor_loss(current,old,advantages,mask,epsilon=config['clip_epsilon']);(actor_loss/len(rollouts)).backward();torch.cuda.synchronize();phase['actor_optimization_seconds']+=time.perf_counter()-start
ratios=(current.detach()-old).exp();probe={'logits':logits[0].detach().cpu().tolist(),'target':int(targets[0])} if iteration==0 and epoch==0 and len(observations)<2 else None
actor_value=float(actor_loss.detach());current_values=current.detach().cpu().tolist()[0];ratio_values=ratios.cpu().tolist()[0]
del actor_loss,current,logits,targets
start=time.perf_counter();values,counts=critic_scores(critic,[row]);values=values.unsqueeze(0);critic_loss=critic_value_loss(values,returns,mask,coefficient=config['value_loss_coef']);(critic_loss/len(rollouts)).backward();torch.cuda.synchronize();phase['critic_optimization_seconds']+=time.perf_counter()-start
observations.append({'rollout_id':raw['rollout_id'],'current_token_logprobs':current_values,'ratios':ratio_values,'current_values':values.detach().cpu().tolist()[0],'actor_loss':actor_value,'critic_loss':float(critic_loss.detach()),'first_token_probe':probe,'active_tokens':raw['response_tokens'],'old_logprob_max_difference':max(abs(x-y) for x,y in zip(current_values,raw['old_token_logprobs']))})
del critic_loss,values
norms={}
for n,optimizer in optimizers.items():
if any(p.grad is None or not torch.isfinite(p.grad).all() for p in parameters[n]):raise ValueError('Missing/nonfinite '+n+' gradients')
norms[n]=float(torch.nn.utils.clip_grad_norm_(parameters[n],config['gradient_clip']));optimizer.step()
if any(p.grad is not None for role in [reference,reward] for p in role.parameters()):raise ValueError('Frozen judge/reference received gradients')
updates+=1;epoch_traces.append({'update':updates,'epoch_within_rollout':epoch+1,'learning_rates':learning_rates,'gradient_norms':norms,'equal_response_denominator':len(rollouts),'observations':observations})
consumed.extend(ids);iteration+=1
history.append({'iteration':iteration,'rollouts':pure_rollouts,'updates':epoch_traces})
torch.cuda.synchronize();seconds=time.perf_counter()-started;final={n:learned_identity(role) for n,role in roles.items()}
for n,role in roles.items():
if frozen_identity(role)!=initial[n]['frozen_identity']:raise ValueError('Frozen base changed: '+n)
for n in ['reference','reward']:
if final[n]!=initial[n]['learned_identity']:raise ValueError('Frozen reference/reward adapters or head changed')
for n in ['actor','critic']:
if final[n]==initial[n]['learned_identity'] or not any(u['gradient_norms'][n]>0 for h in history for u in h['updates']):raise ValueError('No actual '+n+' learning')
state={'identity':identity,'roles':{n:{k:t.detach().cpu().clone() for k,t in learned_state(role).items()} for n,role in roles.items()},'optimizers':{n:optimizer.state_dict() for n,optimizer in optimizers.items()},'sampler':sampler.get_state(),'generation':generation.get_state(),'rng':rng_state(),'order':order,'cursor':cursor,'iteration':iteration,'updates':updates,'history':history,'consumed_rollout_ids':consumed,'pending_rollouts':[],'scheduler':{'next_update':updates,'next_lr':{n:learning_rate(config,updates,total,config['ppo_'+n+'_lr']) for n in optimizers}}}
storage=save_resume_state(destination,state)
save_json(f'outputs/learning/ppo/{mode}.json',{'identity':identity,'mode':mode,'iterations':iteration,'updates_per_optimizer':updates,'history':history,'consumed_rollout_ids':consumed,'final_learned':final,'frozen_reference_reward_unchanged':True,'seconds':seconds,'phase_seconds':phase,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved(),'checkpoint_storage':storage,'numerical_scope':'Old sampling-prefix logprobs preserved; current teacher-forced BF16 scores may differ numerically before first update. Per-token differences are reported, never substituted for old values.'})
if mode=='resume':
save_learned(actor,root/'ppo-resume.safetensors',identity);save_learned(critic,root/'critic-resume.safetensors',identity)
return {'mode':mode,'iterations':iteration,'updates_per_optimizer':updates,'seconds':seconds}
def check_ppo_resume():
root=Path('assets/checkpoints');continuous=load_resume_state(root/'ppo-continuous.pt.gz');resumed=load_resume_state(root/'ppo-resume.pt.gz')
if not equal(continuous,resumed):raise ValueError('Exact four-role PPO continuation failed')
result={'status':'PASS','all_state_exact':True,'iterations':resumed['iteration'],'updates_per_optimizer':resumed['updates'],'rollout_ids':len(resumed['consumed_rollout_ids']),'continuous_sha256':sha(root/'ppo-continuous.pt.gz'),'resume_sha256':sha(root/'ppo-resume.pt.gz')};save_json('outputs/learning/ppo/resume-check.json',result);return result
Evaluate against the judge and against the goal
Disagreement and proxy exploitation
Invariant: every case has two distinct evaluations; incorrect advice remains visible.
Error to diagnose: calling the reward-model score ‘human usefulness’ is incorrect; call it the predicted simulated-preference score.
File src/recompense_ppo/evaluate.py
import math,time
from pathlib import Path
import torch
from .config import read_json,save_json,sha
from .models import reload_learned,load_tokenizer,learned_identity
from .scores import completion_row,reward_scores
from .rollout import sample_response
from .rubric import evaluate_response
@torch.no_grad()
def generate_heldout(config,branch):
if branch not in ['sft','ppo']:raise ValueError('Named actor branch required')
if Path('data/generated/labels').exists() or Path('outputs/learning').exists():raise ValueError('Free generation must not mount labels or learning outputs')
pages=[r for r in read_json('data/generated/input/pages.json') if r['split'] in ['test','probe']]
if len(pages)!=260 or len({r['prompt'] for r in pages})!=200:raise ValueError('Heldout row/context populations changed')
actor,actor_meta=reload_learned(f'assets/checkpoints/{branch}-resume.safetensors');actor.requires_grad_(False)
reference,ref_meta=reload_learned('assets/checkpoints/sft-resume.safetensors',kind='reference')
reward,reward_meta=reload_learned('assets/checkpoints/reward-resume.safetensors');reward.requires_grad_(False)
tokenizer=load_tokenizer();policy_id=learned_identity(actor);generator=torch.Generator(device='cuda').manual_seed(349);rows=[];torch.cuda.reset_peak_memory_stats();torch.cuda.synchronize();start=time.perf_counter()
for index,page in enumerate(pages):
raw=sample_response(actor,tokenizer,page,config,generator,reference=reference,reward_model=reward,mode='greedy',policy_id=policy_id,capture_probe=index<2)
rows.append({'page':page,'generation':raw})
byprompt={}
for row in rows:
raw=row['generation'];key=row['page']['prompt'];value={k:raw[k] for k in ['response_ids','text','finish_reason','learned_reward_score','old_token_logprobs','reference_token_logprobs']}
if key in byprompt and byprompt[key]!=value:raise ValueError('Repeated greedy probe contexts differ')
byprompt[key]=value
torch.cuda.synchronize();result={'branch':branch,'rows':rows,'row_count':260,'unique_contexts':200,'repeated_probe_contexts_verified':True,'labels_visible':False,'actor_checkpoint':actor_meta['sha256'],'reference_checkpoint':ref_meta['sha256'],'reward_checkpoint':reward_meta['sha256'],'seconds':time.perf_counter()-start,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved()};save_json(f'outputs/inference/{branch}.json',result)
return {k:v for k,v in result.items() if k!='rows'}
@torch.no_grad()
def evaluate_reward_pairs(config):
pages={r['document_id']:r for r in read_json('data/generated/input/pages.json')};gold=read_json('data/generated/labels/gold.json');selected=[r for r in gold if pages[r['document_id']]['split']=='test'];tokenizer=load_tokenizer();reward,meta=reload_learned('assets/checkpoints/reward-resume.safetensors');reward.requires_grad_(False);rows=[]
if len(selected)!=180:raise ValueError('Reward test population changed')
torch.cuda.synchronize();start=time.perf_counter()
for label in selected:
page=pages[label['document_id']];chosen=completion_row(tokenizer,page['prompt'],label['chosen'],config,page['document_id']+'-chosen');rejected=completion_row(tokenizer,page['prompt'],label['rejected'],config,page['document_id']+'-rejected');a=float(reward_scores(reward,[chosen])[0]);b=float(reward_scores(reward,[rejected])[0]);rows.append({'document_id':page['document_id'],'chosen_score':a,'rejected_score':b,'chosen_response':label['chosen'],'rejected_response':label['rejected'],'chosen_tokens':chosen['input_ids'],'rejected_tokens':rejected['input_ids'],'correct':a>b,'tie':a==b})
torch.cuda.synchronize();result={'pairs':180,'correct':sum(r['correct'] for r in rows),'ties':sum(r['tie'] for r in rows),'accuracy':sum(r['correct'] for r in rows)/180,'rows':rows,'reward_checkpoint':meta['sha256'],'used_for_selection':False,'seconds':time.perf_counter()-start};save_json('outputs/reward-test.json',result);return {k:v for k,v in result.items() if k!='rows'}
def summarize(rows):
recognized=[r for r in rows if r['rubric']['recognized']]
return {'rows':len(rows),'rubric_success':sum(r['rubric']['rubric_success'] for r in rows),'rubric_success_rate':sum(r['rubric']['rubric_success'] for r in rows)/len(rows),'recognized':len(recognized),'recognized_rate':len(recognized)/len(rows),'uncertified':len(rows)-len(recognized),'relevant_recognized':sum(r['rubric']['relevant'] for r in recognized),'unsupported_certainty_recognized':sum(r['rubric']['unsupported_certainty'] is True for r in recognized),'declared_refusals':sum(r['rubric']['refusal'] for r in rows),'without_eos':sum(r['finish_reason']!='eos' for r in rows),'mean_response_tokens':sum(r['response_tokens'] for r in rows)/len(rows),'mean_learned_reward':sum(r['learned_reward_score'] for r in rows)/len(rows),'mean_sampled_actor_reference_logratio':sum(r['sampled_actor_reference_logratio'] for r in rows)/len(rows),'uncertainty_denominator':len(recognized),'human_utility_claim':False}
def evaluate(config):
gold={r['document_id']:r for r in read_json('data/generated/labels/gold.json')};result={}
for branch in ['sft','ppo']:
raw=read_json(f'outputs/inference/{branch}.json');rows=[]
for item in raw['rows']:
g=item['generation'];page=item['page'];label=gold[page['document_id']];rubric=evaluate_response(g['text'],g['finish_reason'],label['expected_action'])
if not g['tokenization_valid']:rubric={**rubric,'recognized':False,'status':'UNCERTIFIED','rubric_success':False,'relevant':False,'unsupported_certainty':None,'reason':'undecodable_token','refusal':False,'polite':False,'refusal':False,'polite':False}
ratio=sum(a-b for a,b in zip(g['old_token_logprobs'],g['reference_token_logprobs']))/g['response_tokens']
rows.append({'document_id':page['document_id'],'family':page['family'],'split':page['split'],'prompt':page['prompt'],'text':g['text'],'finish_reason':g['finish_reason'],'response_tokens':g['response_tokens'],'learned_reward_score':g['learned_reward_score'],'sampled_actor_reference_logratio':ratio,'rubric':rubric})
unique={r['prompt']:r for r in rows};probe=[r for r in rows if r['split']=='probe'];unique_probe={r['prompt']:r for r in probe}
populations={'test':[r for r in rows if r['split']=='test'],'probe_rows':probe,'probe_unique_contexts':list(unique_probe.values()),'all_rows':rows,'all_unique_contexts':list(unique.values())}
result[branch]={'metrics':{name:summarize(selected) for name,selected in populations.items()},'rows':rows,'generation_seconds':raw['seconds'],'raw_sha256':sha(f'outputs/inference/{branch}.json')}
output={'actor':result,'reward_test':read_json('outputs/reward-test.json'),'scope':'Synthetic finite grammar; UNCERTIFIED prose is not silently accepted. 80 probe rows contain20unique contexts; duplicate-adjusted summaries are separate. Learned reward is not the rubric, no human study.'};save_json('outputs/results.json',output);return {branch:v['metrics'] for branch,v in result.items()}
@torch.no_grad()
def exercise(config):
protected_paths=[Path('configs/minimum.yaml'),Path('outputs/results.json'),Path('outputs/reward-test.json')]+list(Path('outputs/inference').glob('*.json'))+list(Path('assets/checkpoints').glob('*.safetensors'))
protected={str(p):sha(p) for p in protected_paths};pages={r['document_id']:r for r in read_json('data/generated/input/pages.json')};labels=[r for r in read_json('data/generated/labels/gold.json') if pages[r['document_id']]['split']=='probe'];tokenizer=load_tokenizer();reward,metadata=reload_learned('assets/checkpoints/reward-resume.safetensors');reward.requires_grad_(False);before=learned_identity(reward);rows=[];torch.cuda.synchronize();start=time.perf_counter()
if len(labels)!=80:raise ValueError('Expected80preserved factorial probes')
for label in labels:
page=pages[label['document_id']];row=completion_row(tokenizer,page['prompt'],label['candidate_response'],config,page['document_id']);score=float(reward_scores(reward,[row])[0]);rows.append({'document_id':page['document_id'],'family':page['family'],'prompt':page['prompt'],'candidate':label['candidate_response'],'candidate_correct':label['candidate_correct'],'candidate_polite':label['candidate_polite'],'input_ids':row['input_ids'],'response_tokens':row['response_tokens'],'reward_score':score,'rubric':evaluate_response(label['candidate_response'],'eos',label['expected_action'])})
groups=[]
for family in sorted({r['family'] for r in rows}):
group={(r['candidate_correct'],r['candidate_polite']):r['reward_score'] for r in rows if r['family']==family}
if len(group)!=4:raise ValueError('Incomplete factorial probe')
groups.append({'family':family,'correct_politeness_delta':group[True,True]-group[True,False],'incorrect_politeness_delta':group[False,True]-group[False,False],'incorrect_polite_above_correct_plain':group[False,True]>group[True,False]})
torch.cuda.synchronize()
if learned_identity(reward)!=before or any(sha(Path(p))!=value for p,value in protected.items()):raise ValueError('Exercise changed judge/main artifacts')
result={'probe_rows':80,'unique_contexts':20,'rows':rows,'groups':groups,'wrong_polite_outranks_correct_plain':sum(g['incorrect_polite_above_correct_plain'] for g in groups),'mean_incorrect_politeness_delta':sum(g['incorrect_politeness_delta'] for g in groups)/20,'seconds':time.perf_counter()-start,'reward_checkpoint':metadata['sha256'],'reward_unchanged':True,'main_unchanged':True,'protected_sha256':protected,'refit':False};save_json('outputs/exercise/reward-shortcuts.json',result);return {k:v for k,v in result.items() if k not in ['rows','groups','protected_sha256']}
Test the politeness shortcut
Observable reward hacking
Invariant: 80 probes evaluated without a reward-model update or any change to their labels.
Error to diagnose: immediately retraining the reward model on these probes, then announcing the initial score again, destroys separation; version every new study.
File src/recompense_ppo/exercise.py
"""Fixed factorial probe analysis; no further learning or relabeling."""
from .evaluate import exercise
__all__ = ["exercise"]
Sources
PPO, GAE, InstructGPT, and frozen TRL code. These sources explain the mechanisms; our explicit variant does not reproduce a human study or the complete TRL trainer.
