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

astrai/config/train_config.py

@@ -11,7 +11,9 @@ from torch.utils.data import Dataset
 from astrai.config.base import BaseConfig
 from astrai.model.components.lora import LoRAConfig
 
-TRAIN_TYPES = frozenset({"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"})
+TRAIN_TYPES = frozenset(
+    {"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo", "online_ppo"}
+)
 PARALLEL_MODES = frozenset({"none", "ddp", "fsdp"})
 BACKENDS = frozenset({"nccl", "gloo"})
 START_METHODS = frozenset({"spawn", "fork", "forkserver"})
@@ -70,6 +72,8 @@ class TrainConfig(BaseConfig):
         rollout_top_p (float): Top-p (nucleus) filtering for online rollout. Defaults to 0.9.
         rollout_max_tokens (int): Maximum generated tokens per response in rollout. Defaults to 1024.
         reward_model_fn (Optional[Callable]): Factory for reward model, required for online RL strategies. Defaults to None.
+        critic_model_fn (Optional[Callable]): Factory for the value (critic) model, required for online_ppo. Defaults to None.
+        critic_optimizer_fn (Optional[Callable]): Factory for the critic optimizer; None reuses optimizer_fn. Defaults to None.
         executor_kwargs (Dict[str, Any]): Extra kwargs passed to ExecutorFactory.create(). Defaults to {}.
         strategy_kwargs (Dict[str, Any]): Extra strategy arguments. Defaults to {}.
     """
@@ -125,6 +129,8 @@ class TrainConfig(BaseConfig):
     rollout_top_p: float = 0.9
     rollout_max_tokens: int = 1024
     reward_model_fn: Optional[Callable] = None
+    critic_model_fn: Optional[Callable] = None
+    critic_optimizer_fn: Optional[Callable] = None
 
     executor_kwargs: Dict[str, Any] = field(default_factory=dict)
     strategy_kwargs: Dict[str, Any] = field(default_factory=dict)
@@ -227,6 +233,10 @@ class TrainConfig(BaseConfig):
                     f"reward_model_fn is required for online RL strategy "
                     f"{self.strategy!r}"
                 )
+            if self.strategy == "online_ppo" and self.critic_model_fn is None:
+                raise ValueError(
+                    "critic_model_fn is required for online RL strategy 'online_ppo'"
+                )
             if self.nprocs > 1:
                 raise ValueError(
                     f"online RL strategy {self.strategy!r} requires single-process "