Skip to content

fix(rollout): preserve sampled EOS probabilities - #2148

Draft
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fix-sampled-eos-logprobs
Draft

fix(rollout): preserve sampled EOS probabilities#2148
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fix-sampled-eos-logprobs

Conversation

@bvolpato

@bvolpato bvolpato commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Single-assistant-message rollouts remove even a sampled final EOS, then append it with logprob 0.0 and 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

uv run --isolated --extra skyrl-train --extra dev pytest \
  tests/train/generators/test_singleturn_eos_logprobs.py \
  tests/train/generators/test_skyrl_gym_generator.py -q
  • 40 tests passed in the isolated Python 3.12/PyTorch CPU environment.
  • The nine new regression cases produce 6 failures and 3 passing controls against the original generator, then all pass with the fix.
  • The actual agent loop and importance-correction functions verify ratio 1 and acceptance for an unchanged policy, including an unlikely EOS. Cases cover intermediate turns, route metadata, synthetic EOS, missing logprob capture, empty/EOS-only output, and context/length truncation.
  • An independent check through production tensor preprocessing, advantage normalization, and PPO/TIS gives sampled EOS gradient -0.5 and synthetic EOS gradient 0, with total reward preserved.
  • Ruff, Black, and gitleaks hooks passed.
  • GitHub GPU CI stopped during Anyscale setup because the fork job has no 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.

Signed-off-by: bvolpato <brunocvcunha@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant