Skip to content

fix: require explicit off-policy correction for async PPO training - #1829

Open
Shi-Dong wants to merge 1 commit into
mainfrom
shi/async-ppo-policy-contract
Open

fix: require explicit off-policy correction for async PPO training#1829
Shi-Dong wants to merge 1 commit into
mainfrom
shi/async-ppo-policy-contract

Conversation

@Shi-Dong

Copy link
Copy Markdown
Contributor

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 each update_weights_interval window. 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=False makes the trainer recompute log_probs with 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). The weight_versions recorded 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_correction to miles/utils/arguments.py and calls it at the top of train_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

  • New tests in tests/fast/utils/test_arguments.py:
    • async PPO with none of the three flags is rejected with the actionable message
    • each of --use-rollout-logprobs / --use-tis / --keep-old-actor individually satisfies the contract
    • non-PPO estimators (use_critic=False) pass without any flag
  • python -m py_compile on all three touched files
  • black --check, isort --check-only, and ruff check clean on all touched files

Note: tests/fast/utils/test_arguments.py imports miles.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.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

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.

2 participants