fix(rollout): preserve sampled EOS probabilities - #2148
Draft
bvolpato wants to merge 1 commit into
Draft
Conversation
Signed-off-by: bvolpato <brunocvcunha@gmail.com>
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.
Single-assistant-message rollouts remove even a sampled final EOS, then append it with logprob
0.0and an active loss mask. An unchanged policy with EOS probability 0.01 consequently gets a sequence importance ratio of 0.01 instead of 1; outlier filtering can reject the whole on-policy trajectory.Keep the sampled EOS and its probability on the final environment turn. Intermediate EOS tokens are still removed while the assistant message continues. If finalization adds a genuinely synthetic EOS, exclude it from policy/importance-correction loss and leave the reward on the last generated token. Empty and EOS-only responses are covered as well.
Reproduction and validation
-0.5and synthetic EOS gradient0, with total reward preserved.ANYSCALE_CLI_TOKEN; it did not execute tests. Upstream code-quality and gym checks passed; remaining CPU jobs were still running when submitted.The existing executor timing test stalled in the sandbox but passed outside it, along with the complete 40-test suite.
The sample-support work in #2082 preserves the EOS support row but retains the zero rollout logprob; this fix addresses that separate bookkeeping defect.
Downsides
Synthetic EOS tokens no longer receive a policy gradient or terminal reward. This corrects the objective but can change loss and learning curves for existing single-assistant-message runs that previously trained on those fabricated actions.
Risk and rollback
The affected path is non-batched generation with
use_conversation_multi_turn=false. Reward placement and token alignment are the main regression risks; tests cover final observations, intermediate turns, truncation, and routed-expert metadata. Revert this commit to roll back the behavior.