跳转至

PPO

PPO (Schulman et al., 2017) bounds how far a single update can move the policy. It clips the importance ratio (\frac{p_\theta(x)}{q(x)}) so the learner cannot stray too far from the sampler (q) — useful when you take several gradient steps on the same rollout batch.

The clipping and the loss are computed token by token. The clipped objective is:

[ \mathcal{L}{\text{CLIP}}(\theta) = -\mathbb{E}\right) A(x)\right] ]}\left[\text{clip}\left(\frac{p_\theta(x)}{q(x)},\, 1-\epsilon_{\text{low}},\, 1+\epsilon_{\text{high}

and the final loss takes the more conservative of the clipped and unclipped objectives:

[ \mathcal{L}{\text{PPO}}(\theta) = -\mathbb{E}\right) A(x)\right)\right] ]}\left[\min\left(\frac{p_\theta(x)}{q(x)} A(x),\ \text{clip}\left(\frac{p_\theta(x)}{q(x)},\, 1-\epsilon_{\text{low}},\, 1+\epsilon_{\text{high}

In MinT the clip bounds are fixed to (\epsilon_{\text{low}}=\epsilon_{\text{high}}=0.2).

Equivalent code:

prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
clipped_ratio = torch.clamp(prob_ratio, clip_low_threshold, clip_high_threshold)
unclipped = prob_ratio * advantages
clipped = clipped_ratio * advantages
ppo_objective = torch.min(unclipped, clipped)   # most conservative
loss = -ppo_objective.sum()

Inputs

  • target_tokens: array[(N,), int] — IDs sampled by (q).
  • logprobs: array[(N,), float]sampling_logprobs.
  • advantages: array[(N,), float] — per-token advantage.

Outputs

  • logprobstarget_logprobs.

Diagnostics

  • loss:sum — total clipped loss.

Custom clip bounds

Pass loss_fn_config to override the defaults:

fb = training_client.forward_backward(
    data=data,
    loss_fn="ppo",
    loss_fn_config={"clip_low_threshold": 0.9, "clip_high_threshold": 1.1},
)
result = fb.result()