fix: require explicit off-policy correction for async PPO training - #1829
Open
Shi-Dong wants to merge 1 commit into
Open
fix: require explicit off-policy correction for async PPO training#1829Shi-Dong wants to merge 1 commit into
Shi-Dong wants to merge 1 commit into
Conversation
Shi-Dong
requested review from
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23 and
yueming-yuan
as code owners
July 27, 2026 20:26
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Zhichenzzz
approved these changes
Jul 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
In the async train loop (
train_async.py), the next rollout is launched before the actor trains on the current one, and new weights are only published at the end of eachupdate_weights_intervalwindow. Samples therefore come from a policy that is at least one update behind the actor doing the training.With the default flags this staleness is silently ignored for PPO:
use_rollout_logprobs=Falsemakes the trainer recomputelog_probswith the current actor and use them as the PPO ratio denominator. Clipping is then anchored to a policy that never generated the trajectory, and no behavior-to-proximal importance correction is applied (with KL reward shaping, the same mis-attributed log probs also leak into the advantages). Theweight_versionsrecorded on samples are only a metric — nothing enforces consistency.The framework already ships three explicit ways to handle this, but nothing requires picking one:
--use-rollout-logprobs— use the rollout engine's log probs directly as the ratio denominator (true behavior policy).--use-tis— recompute the denominator but apply truncated importance sampling correction against the rollout log probs.--keep-old-actor— recompute the denominator with a retained copy of the weights the rollout engines actually used (the rollout_actor → old_actor queue keeps it version-aligned with the in-flight generation).This PR adds
validate_async_off_policy_correctiontomiles/utils/arguments.pyand calls it at the top oftrain_async.py::train: async PPO training (--advantage-estimator ppo) now fails fast with an actionable error unless one of the three flags is set. Other estimators (e.g. GRPO) are unaffected, and the synchronous entrypoint (train.py) is untouched — there the rollout and the recomputed log probs use the same weights, so no contract is needed.Test plan
tests/fast/utils/test_arguments.py:--use-rollout-logprobs/--use-tis/--keep-old-actorindividually satisfies the contractuse_critic=False) pass without any flagpython -m py_compileon all three touched filesblack --check,isort --check-only, andruff checkclean on all touched filesNote:
tests/fast/utils/test_arguments.pyimportsmiles.utils.arguments, which requires sglang, so the new tests were verified on this branch by exercising the extracted function body directly (same cases as committed); CI runs the real test file.