Regularized Policy Optimization

Dive into Deep Learning · §15.3

Regularized policy optimization
rewards learned from comparisons · an optimizer finds the estimate’s errors · penalize distance to a frozen reference · the optimum in closed form

Rewards You Learn

No formula scores a helpful answer; people can compare.

P(\tau \succ \tau') = \sigma\big( r(\tau) - r(\tau') \big)

Bradley-Terry: fitting r is logistic regression on feature differences.

Fit on 1000 trajectories from a competent reference: rms error 0.19 where the data lives, and the hazard lane, truly -1.5 to enter, is priced at \approx 0: the reference knew better, so the data is silent. Silence reads as zero, a property of the zero-initialized linear fit.

Reward Hacking, Produced

Plan optimally against the fitted reward; grade under the truth.

plan on true reward  : fitted return -0.24, true return +0.59
plan on fitted reward: fitted return -0.03, true return -2.11

The hacked plan wins by its own yardstick and drives straight down the hazard lane. Goodhart’s law: an optimizer pointed at an estimated objective finds the estimate’s errors.

True Return Against KL Budget

Per \beta: solve the regularized Bellman equations (soft, reference-weighted backup), residual printed; every point is an exact optimum.

largest soft Bellman residual across the sweep: 0.0e+00
true return -0.10 at KL 0, peak +0.53 at KL 3.7, then -1.83 at KL 14.2
  • fitted return rises the whole way; true return rises to +0.53 at \approx 4 nats, then falls off a cliff
  • the same curve at language-model scale: :cite:Gao.Schulman.Hilton.2023
  • so control the x axis directly

The Regularized Objective

\max_{\pi}\ E_{a \sim \pi}\big[ r(a) \big] - \beta\, D_{\textrm{KL}}\big( \pi \Vert \pi_{\textrm{ref}} \big)

Proposition. \ \pi^\star(a) = \pi_{\textrm{ref}}(a)\, e^{r(a)/\beta} / Z, with value \beta \log Z.

Proof. The objective equals \beta \log Z - \beta\, D_{\textrm{KL}}(\pi \Vert \pi^\star); apply Gibbs’ inequality. \blacksquare

Verified numerically: largest gap over a ladder of \beta is \approx 10^{-16}.

One Picture

A product, not a ranking: the reference’s dip at a_2 survives into \pi^\star. Spent divergence: 0.15 at \beta = 2, 1.84 at \beta = 0.2.

Four Consequences

  • no penalty \Rightarrow a point-mass optimum (ties: the reference over the tied set): PPO’s entropy decay had a destination built into the objective
  • \beta interpolates: reference \leftarrow \beta \to \infty, greedy \leftarrow \beta \to 0
  • uniform reference = entropy bonus = max-ent RL; \pi^\star = \mathrm{softmax}(r/\beta): Boltzmann exploration was optimal all along
  • Bayes’ rule: prior \pi_{\textrm{ref}}, likelihood e^{r/\beta}, posterior \pi^\star :cite:Levine.2018,Korbak.Perez.Buckley.2022

The Two KLs

Trust region (the clip): against the previous iterate; constrains the path; gone at convergence.

Penalty (here): against a frozen reference; in the objective; changes the optimum.

PPO-RLHF runs both at once :cite:ouyang2022training.

Direction matters: D_{\textrm{KL}}(\pi \Vert \pi_{\textrm{ref}}) is mode-seeking, so post-training sharpens (:numref:sec_mdl-fwd-vs-rev-kl).

Where This Goes

V(s) = \beta \log \sum_a \pi_{\textrm{ref}}(a \mid s)\, e^{Q(s, a)/\beta} \ \xrightarrow{\ \beta \to 0\ } \ \max_a Q(s, a)

  • soft backup: DQN’s \max is the sharp corner of a family
  • DDPG \to TD3 \to SAC: pathwise gradient + twin critics
    • this section’s objective; three parts you already own
  • Ahead: the same objective, with a language model as \pi_{\textrm{ref}}