feat(experimental): Divergence Proximal Policy Optimization#5117
feat(experimental): Divergence Proximal Policy Optimization#5117LeonEricsson wants to merge 21 commits intohuggingface:mainfrom
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
oh, I completely missed that there was already a draft PR for this #5065, sorry @catherinelee274 |
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
casinca
left a comment
casinca
left a comment
There was a problem hiding this comment.
Hey Leon, imho, the implementation looks nice.
The other one sounds more PPO oriented with GAE and value function I guess. I just wanted to clarify with you the sampled token bucket
Co-authored-by: casinca <47400729+casinca@users.noreply.github.com>
The paper defines A'_t = TopK(μ, K) ∪ {a_t}, requiring the sampled
token to always be present in the top-K set. Both vLLM and transformers
generation paths were blindly truncating to K entries, which dropped
the sampled token when it ranked outside top-K. Now the K-th entry is
evicted and replaced with the sampled token when it is absent.
# Conflicts: # trl/generation/vllm_client.py # trl/generation/vllm_generation.py # trl/scripts/vllm_serve.py
qgallouedec
left a comment
qgallouedec
left a comment
There was a problem hiding this comment.
Ok for me
For the record, I didn't check deeply if the implementation is faithful to the paper. It probably needs to be confirmed in the future.
I have a few nits though, I'll also check with codex if it can spot issues
|
@codex review |
| def _make_trainer(self, divergence_type="binary_tv", epsilon=0.2, epsilon_high=0.28): | ||
| """Create a minimal DPPOTrainer-like object with just the attributes needed for _compute_divergence_mask.""" | ||
|
|
||
| class Stub: | ||
| pass | ||
|
|
||
| stub = Stub() | ||
| stub.divergence_type = divergence_type | ||
| stub.epsilon_low = epsilon | ||
| stub.epsilon_high = epsilon_high | ||
| return stub |
There was a problem hiding this comment.
prefer using a public static method instead
There was a problem hiding this comment.
instead of DPPOTrainer._compute_divergence_mask
There was a problem hiding this comment.
could you have a look at the proposed fix, wasn't quite sure I understood what you meant
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5ad2ab6f07
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
|
Hi Leon, huge thanks for this PR. Overall it looks good to me. I leave several small comments above, mainly for typo or default hyperparameters. Thanks for your great efforts again! |
5 participants
note this PR incorporates #5107
Read "Rethinking the Trust Region in LLM Reinforcement Learning" over the weekend. Really like this approach. It continues the recent push toward improved off-policy regulation. The work comes from the same authors as the DAPO paper.
What does this PR do?
Implements the proposed method, DPPO, as an experimental trainer.
DPPO replaces PPO/GRPO clipping with a principled trust region based on direct policy divergence estimates. The paper argues that PPO-style clipping, being based on the probability ratio of the sampled token, acts as a noisy single-sample Monte Carlo estimate of the true policy divergence, over-penalizing low-probability tokens and under-penalizing high-probability ones.
Implementation diffs to
GRPOTrainerDPPO uses a two-policy setup, in contrast to GRPO’s three-policy formulation: a sampling policy$\mu_{\text{old}}$ (used to generate rollouts) and a current policy $\pi$ (updated during training). Unlike GRPO’s old/new-policy structure, DPPO computes its trust-region terms directly between $\pi$ and $\mu_{\text{old}}$ , using rollout-time statistics collected from $\mu_{\text{old}}$ . We remove any recomputation of the logprobs at rollout time.
Depending on the divergence type, DPPO requires additional statistics from the policy distributions that are not computed in GRPO. The$\mu_{\text{old}}$ ’s top-K distribution, as well as $\pi$ evaluated on those same token IDs. This constitutes the main deviation from
binary_*objectives only require sampled-token log-probabilities (already available in GRPO-style rollouts). In contrast, thetopk_*objectives additionally requireGRPOTrainer.To support this, we modify the generation step to compute top-K log-probabilities at rollout time. During each gradient step, we then evaluate$\pi$ on the token IDs corresponding to $\mu_{\text{old}}$ ’s top-K positions.
Before submitting
Pull Request section?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.