Regularized Policy Optimization

Dive into Deep Learning · §15.3

Regularized policy optimization
preference-based rewards · proxy-reward optimization · fixed-reference penalties · a closed-form optimum

Learning a Reward Model

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 rarely entered it, so the data contain little evidence there. Its near-zero estimate follows from initialization and regularization.

Optimizing an Imperfect Reward

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 fitted-reward optimum enters the hazard lane. It scores higher under the fitted reward and substantially lower under the true reward, exposing unsupported reward estimates.

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 throughout; true return reaches +0.53 at \approx 4 nats and then decreases sharply
  • the same curve at language-model scale: :cite:Gao.Schulman.Hilton.2023
  • this motivates direct control of divergence from the reference

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}.

The Closed-Form Optimum

The optimum is a product rather than a reward ranking: the reference’s dip at a_2 remains in \pi^\star. The divergences are 0.15 at \beta = 2 and 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), the optimizer of the entropy-regularized one-step objective
  • 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).

Connections to SAC

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 combines pathwise gradients, twin critics, and this section’s entropy-regularized objective
  • Ahead: the same objective, with a language model as \pi_{\textrm{ref}}