feat: add online ppo with value-model critic and gae advantages

- register online_ppo train type backed by PPOStrategy: token-level clipped surrogate over GAE advantages plus masked value regression against rollout-pinned returns, with explained-variance metrics - fold the reference-KL penalty (k3 estimator) into per-token rewards before GAE and pin advantages/returns on RolloutResult so replayed gradient steps optimize fixed targets - add self-contained ValueModel critic with a zero-initialized value head and backbone warm-started from policy weights; AutoRegressiveLM stays untouched and trunk parity is pinned by tests - step the critic's own optimizer outside the policy-version lock with the same max_grad_norm clipping as the policy - persist critic state as value_model.pt/value_optimizer.pt checkpoint extras; resume restores it, fails loudly when missing, and the train.sh completeness check requires the extras for online_ppo configs - extract shared rollout sequence/logprob helpers from GRPO (behavior unchanged) and add ppo_gamma/ppo_gae_lambda/ppo_vf_coef CLI options
2026-09-05 01:59:50 +08:00
parent 816b96a58a
commit 350e4a1849
17 changed files with 1390 additions and 72 deletions
@@ -68,9 +68,16 @@ class RolloutResult(RawRollout):

    Fields:
        rewards: Reward per response, shape ``[B, G]``.
+        advantages: Optional GAE advantages ``[B, G, R_max]`` pinned at
+            rollout time by actor-critic strategies (PPO).  ``None`` until
+            a strategy computes them.
+        returns: Optional GAE value targets ``[B, G, R_max]`` matching
+            ``advantages``.
    """

    rewards: Tensor
+    advantages: Optional[Tensor] = None
+    returns: Optional[Tensor] = None


 class BaseRewardModel(ABC):