Tpo-Torch – Target Policy Optimization for Stable RLHF Alignment in PyTorch

Target Policy Optimization — a simpler alternative to PPO for RLHF.

Based on arXiv:2604.06159 (Kaddour, 2026).

TPO is an RLHF algorithm. It is not a replacement for cross-entropy (CE) training. CE and TPO do different things:

CE is supervised learning. TPO is reinforcement learning. They solve different problems. You don't compare them.

TPO competes with PPO , not CE. The advantage: TPO gives you the same RLHF capability with much less complexity — no value function, no clipping, no importance sampling ratios.

PPO: loss = -min(ratio * A, clip(ratio, 1-eps, 1+eps) * A) # needs V(s), clipping TPO: loss = -target_prob * log P_policy(token) # needs nothing extra Install pip install -e . Quick Start from tpo_torch import TPOTrainer from transformers import AutoModelForCausalLM , AutoTokenizer model = AutoModelForCausalLM . from_pretrained ( "Qwen/Qwen2.5-0.5B-Instruct" ) ref_model = AutoModelForCausalLM . from_pretrained ( "Qwen/Qwen2.5-0.5B-Instruct" ) tokenizer = AutoTokenizer . from_pretrained ( "Qwen/Qwen2.5-0.5B-Instruct" ) tokenizer . pad_token = tokenizer . eos_token # Dataset needs: prompt, labels, and 'advantages' (higher = better response) trainer = TPOTrainer ( model = model , ref_model = ref_model , beta = 0.1 , train_dataset = dataset , processing_class = tokenizer , ) trainer . train () Architecture ┌──────────────────────┐ │ Reference Model │ │ (frozen) │ └──────────┬───────────┘ │ ref_logits ▼ Prompt + Labels ──▶ ┌──────────────────────┐ │ TPO Loss Function │ │ │ │ 1. log-odds(P_ref) │ │ 2. + advantage/beta │ │ 3. sigmoid -> target │ │ 4. CE loss vs target │ └──────────┬───────────┘ │ loss ▼ ┌──────────────────────┐ │ Policy Model │ │ (trained) │ └──────────────────────┘ API Function Description tpo_loss_from_logits(logits, ref_logits, labels, advantages, beta) Loss from raw model logits tpo_loss(policy_logprobs, ref_logprobs, advantages, beta) Loss from pre-computed log-probs TPOTrainer(model, ref_model, beta, ...) HuggingFace Trainer with TPO TPODataCollator(tokenizer) Preserves advantages in batches Benchmarks All benchmarks run on NVIDIA RTX 3050 Laptop (4GB VRAM) .

Both are RLHF methods evaluated on the same metric : held-out perplexity.

Note: Our PPO-clip is simplified (no value function, no GAE). Full PPO with a critic would perform better but requires significantly more code and compute. TPO achieves competitive results with none of that infrastructure.

python benchmarks/run_benchmarks.py CLI tpo train --max-steps 10 tpo bench tpo info Development git clone https://github.com/Griffith-7/tpo-torch.git cd tpo-torch pip install -e " .[dev] " Run tests:

pytest tests/ -v ruff check tpo_torch/ Requirements Python >= 3.9 PyTorch >= 2.0 transformers >= 4.40 Citation @misc { kaddour2026targetpolicyoptimization , title = { Target Policy Optimization } , author = { Jean Kaddour } , year = { 2026 } , eprint = { 2604.06159 } , archivePrefix = { arXiv } , primaryClass = { cs.LG } } License MIT