Closed form of the KL-regularised RL optimum

ProofGraduate⏱ 1 hourLLMsrlhfkl regularisationgibbsvariational

Problem

Prove that the maximiser of $\E_{y\sim\pi}[r(x,y)] - \beta\,\mathrm{KL}(\pi\|\pi_{\mathrm{ref}})$ is $\pi^\star(y\mid x)\propto\pi_{\mathrm{ref}}(y\mid x)e^{r(x,y)/\beta}$ , and explain why RLHF uses PPO rather than this formula.

Write a complete argument. Every inequality you use should be one you could state precisely.

Solution

Opening this before you have a real attempt on paper is the fastest way to feel like you understand something you cannot reproduce.

Related problems

All problems

Problem	Field	Type	Difficulty	Time
Rewards are identifiable only up to a prompt-dependent shift	LLMs	Proof	Intermediate	30 min
The DPO gradient is a weighted contrastive step	LLMs	Proof	Advanced	30 min
The RAG evidence lower bound	RAG	Proof	Advanced	30 min
Derive the DPO objective abyss	LLMs	Proof	Graduate	1 hour