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

Researchers have released Tpo-Torch, a PyTorch implementation of Target Policy Optimization (TPO), a reinforcement learning from human feedback (RLHF) algorithm that simplifies the standard PPO approach by eliminating the need for a value function, clipping, and importance sampling ratios. Based on a 2026 arXiv paper by Kaddour, TPO achieves a held-out perplexity of 74.51 compared to 825.66 for a simplified PPO-clip, with zero NaN gradients and a maximum gradient norm of 0.09 across advantage values from 0.01 to 1000. The library, which includes a TPOTrainer compatible with HuggingFace models, is designed to provide stable RLHF alignment with reduced complexity.

Target Policy Optimization — a simpler alternative to PPO for RLHF. Based on arXiv:2604.06159 https://arxiv.org/abs/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: | Cross-Entropy SFT | PPO | TPO | | |---|---|---|---| | Purpose | Predict next token | Optimize reward via RL | Optimize reward via RL | | Needs labeled data | Yes | No | No | | Needs reward model | No | Yes | Yes | | Needs value/critic head | No | Yes | No | | Needs clipping | No | Yes | No | | Needs importance ratios | No | Yes | No | | Training stability | Stable | Often unstable | Stable | 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 pip install -e . python 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 ┌──────────────────────┐ │ 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 │ └──────────────────────┘ | 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 | All benchmarks run on NVIDIA RTX 3050 Laptop 4GB VRAM . Both are RLHF methods evaluated on the same metric : held-out perplexity. | Method | Final Perplexity | Implementation Complexity | |---|---|---| | TPO beta=0.1 | 74.51 | Loss function only | | PPO-clip simplified | 825.66 | Loss + ratio + clipping | 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. Zero NaNs at advantage values from 0.01 to 1000- Max gradient norm: 0.09 — stable across all regimes | Seq Len | Latency | Throughput | |---|---|---| | 32 | 4.17ms | 61,345 tok/s | | 64 | 8.26ms | 62,021 tok/s | | 128 | 16.50ms | 62,046 tok/s | | 256 | 32.81ms | 62,418 tok/s | | 512 | 65.42ms | 62,614 tok/s | | 1024 | 929.89ms | 8,810 tok/s | python benchmarks/run benchmarks.py tpo train --max-steps 10 tpo bench tpo info 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/ - Python = 3.9 - PyTorch = 2.0 - transformers = 4.40 @misc{kaddour2026targetpolicyoptimization, title={Target Policy Optimization}, author={Jean Kaddour}, year={2026}, eprint={2604.06159}, archivePrefix={arXiv}, primaryClass={cs.LG} }