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 "

astrai/model/__init__.py

@@ -13,6 +13,7 @@ from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.encoder import EmbeddingEncoder
 from astrai.model.transformer import AutoRegressiveLM
+from astrai.model.value import ValueModel
 
 __all__ = [
     # Modules
@@ -26,6 +27,7 @@ __all__ = [
     "AutoRegressiveLM",
     "EmbeddingEncoder",
     "AutoModel",
+    "ValueModel",
     # LoRA
     "LoRAConfig",
     "inject_lora",

astrai/model/value.py

@@ -0,0 +1,55 @@
+"""Value (critic) model for actor-critic RL training."""
+
+from typing import Dict, Optional
+
+import torch.nn as nn
+from torch import Tensor
+
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.automodel import ModelFactory
+from astrai.model.components.linear import Linear
+from astrai.model.transformer import AutoRegressiveLM, process_attention_mask
+
+
+@ModelFactory.register("value_model")
+class ValueModel(AutoRegressiveLM):
+    """Critic scoring each state with a scalar instead of vocab logits.
+
+    Inherits the ``AutoRegressiveLM`` components so a policy checkpoint can
+    warm-start the critic backbone (``load_state_dict(..., strict=False)``);
+    only ``value_head`` keeps its fresh initialization.  The inherited
+    ``lm_head`` parameters stay dormant — the forward below never projects
+    through them — so checkpoints round-trip with stable keys.  The trunk
+    pass mirrors ``AutoRegressiveLM.forward`` for training-style input;
+    ``tests/trainer/test_ppo_strategy.py`` pins the two to identical
+    hidden states.
+    """
+
+    def __init__(self, config: AutoRegressiveLMConfig):
+        super().__init__(config)
+        self.value_head = Linear(config.hidden_size, 1, bias=True)
+        # Zero head so training starts from V(s) == 0 and the first GAE
+        # advantages are driven purely by rewards.
+        nn.init.zeros_(self.value_head.weight)
+        nn.init.zeros_(self.value_head.bias)
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        input_mask: Optional[Tensor] = None,
+        position_ids: Optional[Tensor] = None,
+    ) -> Dict[str, Tensor]:
+        if input_ids.ndim != 2:
+            raise ValueError("critic input_ids must be [batch, seq_len]")
+        x = self.embed_tokens(input_ids)
+        rotary_emb = self.rotary_embedding(x, position_ids)
+        attn_mask = process_attention_mask(input_mask)
+        use_sdpa_causal_mask = attn_mask is None
+
+        for layer in self.layers:
+            x = layer(x, rotary_emb, attn_mask, None, use_sdpa_causal_mask, None)[
+                "hidden_states"
+            ]
+        hidden_states = self.norm(x)
+        values = self.value_head(hidden_states).squeeze(-1)
+        return {"values": values}

astrai/trainer/rollout.py

@@ -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):

astrai/trainer/strategy.py

@@ -1,7 +1,7 @@
 """Training strategy implementations with factory pattern."""
 
 from abc import ABC
-from typing import Any, Callable, Dict, List, Optional, TypedDict, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union
 
 import torch
 import torch.nn as nn
@@ -85,6 +85,166 @@ def get_logprobs(
     }
 
 
+def rollout_sequences(
+    prompts: Tensor,
+    prompt_mask: Tensor,
+    responses: Tensor,
+    response_masks: Tensor,
+) -> Tuple[Tensor, Tensor]:
+    """Concatenate grouped prompts with responses for sequence scoring.
+
+    Expands ``prompts`` [B, P] across the group dimension of ``responses``
+    [B, G, R] and builds the combined key-padding + causal attention mask.
+
+    Returns:
+        ``(full_sequences, attn_mask)`` each shaped [B*G, P + R]; the
+        attention mask is 4-D boolean.
+    """
+    group_size = responses.size(1)
+    responses_flat = responses.view(-1, responses.size(-1))
+    masks_flat = response_masks.view(-1, responses.size(-1)).bool()
+    prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
+    prompt_mask_expanded = (
+        prompt_mask.unsqueeze(1).expand(-1, group_size, -1).flatten(0, 1).bool()
+    )
+
+    full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
+    # Build full attention mask: key-padding + causal
+    key_pad = torch.cat([prompt_mask_expanded, masks_flat], dim=-1)[:, None, None, :]
+    S = key_pad.shape[-1]
+    causal = torch.tril(
+        torch.ones(S, S, dtype=torch.bool, device=full_sequences.device)
+    )[None, None, :, :]
+    attn_mask = key_pad & causal
+    return full_sequences, attn_mask
+
+
+def rollout_token_logprobs(
+    model: nn.Module,
+    prompts: Tensor,
+    prompt_mask: Tensor,
+    responses: Tensor,
+    response_masks: Tensor,
+) -> LogprobsOutput:
+    """Per-response-token log probabilities for a grouped rollout batch.
+
+    Prompt tokens are masked out (0) so logprobs are computed only for
+    response tokens.  ``get_logprobs`` shifts the mask by one position, so
+    the first response token's logprob (predicted from the last prompt
+    token) is correctly included.
+
+    Returns:
+        ``logprobs`` reshaped to [B, G, R]: position j is the log-probability
+        of response token j under ``model``.
+    """
+    batch_size, group_size, response_len = responses.shape
+    prompt_len = prompts.size(1)
+    full_sequences, attn_mask = rollout_sequences(
+        prompts, prompt_mask, responses, response_masks
+    )
+    masks_flat = response_masks.view(-1, response_len)
+    full_masks = torch.cat(
+        [
+            torch.zeros(
+                batch_size * group_size,
+                prompt_len,
+                dtype=torch.bool,
+                device=full_sequences.device,
+            ),
+            masks_flat,
+        ],
+        dim=-1,
+    )
+
+    # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
+    # Response token logprobs occupy the last ``response_len`` positions.
+    output = get_logprobs(model, full_sequences, attn_mask, full_masks, "none")
+    output["logprobs"] = output["logprobs"][:, prompt_len - 1 :].view(
+        batch_size, group_size, response_len
+    )
+    return output
+
+
+def rollout_token_values(
+    model: nn.Module,
+    prompts: Tensor,
+    prompt_mask: Tensor,
+    responses: Tensor,
+    response_masks: Tensor,
+) -> Tensor:
+    """Critic values [B, G, R] aligned with response token positions.
+
+    Position j holds V(s_j) — the value of the state right before response
+    token j is emitted — matching the logprob alignment of
+    :func:`rollout_token_logprobs`.
+    """
+    prompt_len = prompts.size(1)
+    full_sequences, attn_mask = rollout_sequences(
+        prompts, prompt_mask, responses, response_masks
+    )
+    output = model(full_sequences, input_mask=attn_mask)
+    values = output["values"].float()[
+        :, prompt_len - 1 : prompt_len - 1 + responses.size(-1)
+    ]
+    return values.view(responses.shape)
+
+
+def compute_gae(
+    rewards: Tensor,
+    values: Tensor,
+    mask: Tensor,
+    gamma: float,
+    gae_lambda: float,
+) -> Tuple[Tensor, Tensor]:
+    """Generalized advantage estimation over padded response tokens.
+
+    Args:
+        rewards: [B, G, R] per-token rewards; the terminal reward must sit
+            at each response's last valid position, padded positions 0.
+        values: [B, G, R] rollout-time critic values V(s_t) (see
+            :func:`rollout_token_values`).
+        mask: [B, G, R] valid-token mask; padded positions are excluded
+            and cannot leak into valid advantages.
+        gamma: Discount factor.
+        gae_lambda: GAE bias/variance trade-off.
+
+    Returns:
+        ``(advantages, returns)`` shaped [B, G, R].  The episode ends at the
+        last valid token (no bootstrap value beyond truncation).
+    """
+    response_len = rewards.size(-1)
+    flat_rewards = rewards.reshape(-1, response_len)
+    flat_mask = mask.reshape(-1, response_len).to(values.dtype)
+    # Padded values must be zero or the backward scan would leak them.
+    flat_values = values.reshape(-1, response_len) * flat_mask
+
+    advantages = torch.zeros_like(flat_rewards)
+    gae = torch.zeros_like(flat_values[:, 0])
+    for t in range(response_len - 1, -1, -1):
+        next_values = (
+            flat_values[:, t + 1]
+            if t + 1 < response_len
+            else torch.zeros_like(flat_values[:, t])
+        )
+        delta = flat_rewards[:, t] + gamma * next_values - flat_values[:, t]
+        gae = flat_mask[:, t] * (delta + gamma * gae_lambda * gae)
+        advantages[:, t] = gae
+    returns = advantages + flat_values
+    return advantages.view_as(rewards), returns.view_as(rewards)
+
+
+def _validate_behavior_logprobs(behavior_logprobs: Tensor, responses: Tensor) -> None:
+    """Reject behaviour-policy logprobs that do not match the responses."""
+    if behavior_logprobs.shape != responses.shape:
+        raise ValueError(
+            "logprobs_old shape must match responses: "
+            f"got {tuple(behavior_logprobs.shape)}, "
+            f"expected {tuple(responses.shape)}"
+        )
+    if not torch.isfinite(behavior_logprobs).all():
+        raise ValueError("logprobs_old must contain only finite values")
+
+
 def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
     S = position_ids.size(1)
     device = position_ids.device
@@ -630,82 +790,36 @@ class GRPOStrategy(BaseStrategy):
         masks = batch["masks"]
         rewards = batch["rewards"]
 
-        batch_size, group_size, response_len = responses.shape
         behavior_logprobs = batch.get("logprobs_old")
         if behavior_logprobs is not None:
-            if behavior_logprobs.shape != responses.shape:
+            _validate_behavior_logprobs(behavior_logprobs, responses)
-                raise ValueError(
-                    "logprobs_old shape must match responses: "
-                    f"got {tuple(behavior_logprobs.shape)}, "
-                    f"expected {tuple(responses.shape)}"
-                )
-            if not torch.isfinite(behavior_logprobs).all():
-                raise ValueError("logprobs_old must contain only finite values")
             behavior_logprobs = behavior_logprobs.detach().float()
         elif self.old_model is None:
             raise ValueError(
                 "GRPO batches must provide logprobs_old when no old_model is configured"
             )
 
-        responses_flat = responses.view(-1, response_len)
-        masks_flat = masks.view(-1, response_len)
-        prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
         prompt_mask = batch.get("prompt_mask")
         if prompt_mask is None:
             prompt_mask = prompts.ne(0)
-        prompt_mask_expanded = (
-            prompt_mask.unsqueeze(1).expand(-1, group_size, -1).flatten(0, 1)
-        )
-        prompt_len = prompt_expanded.size(1)
 
-        full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
+        policy_output = rollout_token_logprobs(
-        # Prompt tokens are masked out (0) so logprobs are computed only for
+            self.model, prompts, prompt_mask, responses, masks
-        # response tokens.  get_logprobs shifts the mask by one position, so
-        # the first response token's logprob (predicted from the last prompt
-        # token) is correctly included.
-        full_masks = torch.cat(
-            [torch.zeros_like(prompt_expanded, dtype=torch.bool), masks_flat], dim=-1
-        )
-
-        # Build full attention mask: key-padding + causal
-        key_pad = torch.cat([prompt_mask_expanded, masks_flat.bool()], dim=-1)[
-            :, None, None, :
-        ]
-        S = key_pad.shape[-1]
-        causal = torch.tril(
-            torch.ones(S, S, dtype=torch.bool, device=full_sequences.device)
-        )[None, None, :, :]
-        attn_mask = key_pad & causal
-
-        # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
-        # Response token logprobs occupy the last ``response_len`` positions
-        # (the first response token is predicted from the last prompt token).
-        policy_output = get_logprobs(
-            self.model, full_sequences, attn_mask, full_masks, "none"
         )
         token_log_probs_policy = policy_output["logprobs"]
         aux_loss = policy_output["aux_loss"]
-        token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :]
         with torch.no_grad():
             if behavior_logprobs is None:
-                old_output = get_logprobs(
+                token_log_probs_old = rollout_token_logprobs(
-                    self.old_model, full_sequences, attn_mask, full_masks, "none"
+                    self.old_model, prompts, prompt_mask, responses, masks
-                )
+                )["logprobs"]
-                token_log_probs_old = old_output["logprobs"]
-                token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
             else:
                 token_log_probs_old = behavior_logprobs
-            ref_output = get_logprobs(
+            token_log_probs_ref = rollout_token_logprobs(
-                self.ref_model, full_sequences, attn_mask, full_masks, "none"
+                self.ref_model, prompts, prompt_mask, responses, masks
-            )
+            )["logprobs"]
-            token_log_probs_ref = ref_output["logprobs"]
-            token_log_probs_ref = token_log_probs_ref[:, prompt_len - 1 :]
 
-        # Reshape to [B, G, response_len]
+        token_masks = masks.float()
-        token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
-        token_log_probs_old = token_log_probs_old.view(batch_size, group_size, -1)
-        token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1)
-        token_masks = masks_flat.view(batch_size, group_size, -1).float()
 
         # Group-normalized advantages from scalar per-response rewards.
         eps = 1e-8
@@ -754,8 +868,228 @@ class GRPOStrategy(BaseStrategy):
         }
 
 
+@StrategyFactory.register("online_ppo")
+class PPOStrategy(BaseStrategy):
+    """Proximal Policy Optimization with a learned critic (actor-critic).
+
+    Uses the same token-level clipped surrogate as GRPO, but advantages
+    come from GAE(λ) over a :class:`~astrai.model.value.ValueModel` critic
+    instead of group-normalized rewards.
+
+    Roles:
+
+    * **Policy** ``self.model`` — the actor being trained.
+    * **Behaviour policy** — per-token ``logprobs_old`` captured by the
+      rollout sampler; always required (no offline old-model fallback).
+    * **Critic** ``self.critic`` — trained jointly by masked MSE regression
+      against the GAE returns.  It owns a separate optimizer, stepped by
+      :meth:`optimizer_step` *outside* the policy-version lock: the critic
+      never serves generation, so its weights are not part of a policy
+      version publication.
+    * **Reference model** ``self.ref_model`` — optional frozen copy of the
+      initial policy.  When set (and ``kl_coef > 0``), a per-token KL
+      penalty (k3 estimator, ``logπ_old − logπ_ref``) is folded into the
+      rewards before GAE, following InstructGPT-style reward shaping.
+
+    Advantages and returns are computed once per rollout — with rollout-time
+    critic values — and pinned on the :class:`RolloutResult`, so every
+    replayed gradient step optimizes the same fixed targets, mirroring
+    classic PPO's multiple epochs over one batch.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        device: str,
+        critic: nn.Module,
+        critic_optimizer: Optional[Optimizer] = None,
+        ref_model: Optional[nn.Module] = None,
+        clip_eps: float = 0.2,
+        kl_coef: float = 0.01,
+        gamma: float = 1.0,
+        gae_lambda: float = 0.95,
+        vf_coef: float = 0.5,
+        max_grad_norm: Optional[float] = None,
+        **kwargs,
+    ):
+        super().__init__(model, device, **kwargs)
+        self.critic = critic
+        self.critic_optimizer = critic_optimizer
+        self.ref_model = ref_model
+        self.clip_eps = clip_eps
+        self.kl_coef = kl_coef
+        self.gamma = gamma
+        self.gae_lambda = gae_lambda
+        self.vf_coef = vf_coef
+        self.max_grad_norm = max_grad_norm
+
+    def optimizer_step(self, optimizer: Optimizer):
+        """Step the policy under the version lock, then the critic.
+
+        The critic step runs after the policy's atomic version publication:
+        a concurrent rollout may already observe the new policy version,
+        but only the (unused-for-generation) critic lags by one step.
+        """
+        result = super().optimizer_step(optimizer)
+        if self.critic_optimizer is not None:
+            if self.max_grad_norm is not None:
+                torch.nn.utils.clip_grad_norm_(
+                    self.critic.parameters(), self.max_grad_norm
+                )
+            self.critic_optimizer.step()
+            self.critic_optimizer.zero_grad()
+        return result
+
+    @torch.no_grad()
+    def _compute_advantages(
+        self,
+        prompts: Tensor,
+        prompt_mask: Tensor,
+        responses: Tensor,
+        response_masks: Tensor,
+        rewards: Tensor,
+        behavior_logprobs: Tensor,
+    ) -> Tuple[Tensor, Tensor]:
+        """GAE advantages/returns pinned to rollout-time critic values."""
+        device = self.device
+        prompts = prompts.to(device)
+        prompt_mask = prompt_mask.to(device)
+        responses = responses.to(device)
+        response_masks = response_masks.to(device)
+        rewards = rewards.to(device)
+        behavior_logprobs = behavior_logprobs.to(device)
+
+        values = rollout_token_values(
+            self.critic, prompts, prompt_mask, responses, response_masks
+        )
+
+        # Terminal reward lands on each response's last valid token.
+        token_rewards = torch.zeros_like(values)
+        lengths = response_masks.long().sum(dim=-1)
+        terminal = (lengths - 1).clamp(min=0)
+        token_rewards.view(-1, values.size(-1)).scatter_(
+            1, terminal.reshape(-1, 1), rewards.reshape(-1, 1).to(values.dtype)
+        )
+
+        if self.ref_model is not None and self.kl_coef > 0:
+            ref_logprobs = rollout_token_logprobs(
+                self.ref_model, prompts, prompt_mask, responses, response_masks
+            )["logprobs"]
+            # k3 per-token KL estimator folded into the reward.
+            kl_penalty = behavior_logprobs.float() - ref_logprobs
+            token_rewards = (
+                token_rewards
+                - self.kl_coef * kl_penalty * response_masks.to(values.dtype)
+            )
+
+        return compute_gae(
+            token_rewards, values, response_masks, self.gamma, self.gae_lambda
+        )
+
+    def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
+        _validate_behavior_logprobs(result.logprobs_old, result.responses)
+        if result.advantages is None or result.returns is None:
+            result.advantages, result.returns = self._compute_advantages(
+                result.prompts,
+                result.prompt_mask,
+                result.responses,
+                result.response_mask,
+                result.rewards,
+                result.logprobs_old,
+            )
+        return {
+            "prompts": result.prompts,
+            "prompt_mask": result.prompt_mask,
+            "responses": result.responses,
+            "masks": result.response_mask,
+            "rewards": result.rewards,
+            "logprobs_old": result.logprobs_old,
+            "advantages": result.advantages,
+            "returns": result.returns,
+        }
+
+    def supports_online(self) -> bool:
+        return True
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
+        batch = move_to_device(batch, self.device)
+        prompts = batch["prompts"]
+        responses = batch["responses"]
+        masks = batch["masks"]
+
+        behavior_logprobs = batch.get("logprobs_old")
+        if behavior_logprobs is None:
+            raise ValueError(
+                "PPO batches must provide logprobs_old captured at rollout"
+            )
+        _validate_behavior_logprobs(behavior_logprobs, responses)
+        behavior_logprobs = behavior_logprobs.detach().float()
+
+        prompt_mask = batch.get("prompt_mask")
+        if prompt_mask is None:
+            prompt_mask = prompts.ne(0)
+
+        advantages = batch.get("advantages")
+        returns = batch.get("returns")
+        if advantages is None or returns is None:
+            advantages, returns = self._compute_advantages(
+                prompts,
+                prompt_mask,
+                responses,
+                masks,
+                batch["rewards"],
+                behavior_logprobs,
+            )
+
+        policy_output = rollout_token_logprobs(
+            self.model, prompts, prompt_mask, responses, masks
+        )
+        token_log_probs_policy = policy_output["logprobs"]
+
+        # Critic forward with gradients: value regression against the
+        # rollout-pinned GAE returns.
+        values = rollout_token_values(
+            self.critic, prompts, prompt_mask, responses, masks
+        )
+        token_masks = masks.float()
+        token_count = token_masks.sum().clamp(min=1.0)
+
+        # Token-level ratio (π_θ / π_old) and PPO clipping.
+        log_ratio = token_log_probs_policy - behavior_logprobs
+        ratio = torch.exp(log_ratio)
+
+        surr1 = ratio * advantages
+        surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
+        per_token_policy_loss = -torch.min(surr1, surr2)
+        policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
+
+        value_loss = ((values - returns) ** 2 * token_masks).sum() / token_count
+        task_loss = policy_loss + self.vf_coef * value_loss
+
+        def masked_variance(x: Tensor) -> Tensor:
+            mean = (x * token_masks).sum() / token_count
+            return ((x - mean) ** 2 * token_masks).sum() / token_count
+
+        with torch.no_grad():
+            explained_variance = 1.0 - masked_variance(
+                values - returns
+            ) / masked_variance(returns).clamp(min=1e-8)
+
+        return self._loss_output(
+            task_loss,
+            {
+                "policy_loss": policy_loss,
+                "value_loss": value_loss,
+                "explained_variance": explained_variance,
+            },
+            policy_output["aux_loss"],
+            policy_output.get("router_stats"),
+        )
+
+
 # Factory aliases: online variants use the same strategy class; the
 # ``RolloutRunner`` is injected by ``TrainContextBuilder`` to enable
-# online mode, so no separate subclass is needed.
+# online mode, so no separate subclass is needed.  ``PPOStrategy`` is
+# registered under its sole name ``online_ppo`` — it has no offline mode.
 StrategyFactory.register("online_grpo")(GRPOStrategy)
 StrategyFactory.register("online_dpo")(DPOStrategy)

astrai/trainer/train_callback.py

@@ -231,6 +231,12 @@ class CheckpointCallback(TrainCallback):
             obj = getattr(context, name, None)
             if obj:
                 extra[name] = obj.state_dict()
+        critic = getattr(context.strategy, "critic", None)
+        if critic is not None:
+            extra["value_model"] = critic.state_dict()
+            critic_optimizer = getattr(context.strategy, "critic_optimizer", None)
+            if critic_optimizer is not None:
+                extra["value_optimizer"] = critic_optimizer.state_dict()
         return extra
 
 

astrai/trainer/train_context.py

@@ -16,6 +16,7 @@ from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import (
     BaseExecutor,
     ExecutorFactory,
+    broadcast_state_dict,
     create_ref_model,
     strip_compile_prefix,
 )
@@ -297,7 +298,7 @@ class TrainContextBuilder:
         cfg = self.config
         kwargs = dict(cfg.strategy_kwargs)
         kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
-        if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo"):
+        if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo", "online_ppo"):
             kwargs["ref_model"] = create_ref_model(
                 cfg.model_fn,
                 executor=executor,
@@ -313,6 +314,11 @@ class TrainContextBuilder:
             )
         elif cfg.strategy == "online_grpo":
             kwargs["old_model"] = None
+        if cfg.strategy == "online_ppo":
+            critic, critic_optimizer = self._create_critic(context, executor)
+            kwargs["critic"] = critic
+            kwargs["critic_optimizer"] = critic_optimizer
+            kwargs.setdefault("max_grad_norm", cfg.max_grad_norm)
         context.strategy = StrategyFactory.create(
             cfg.strategy,
             model=context.model,
@@ -322,6 +328,63 @@ class TrainContextBuilder:
         )
         return kwargs
 
+    def _create_critic(
+        self, context: TrainContext, executor: BaseExecutor
+    ) -> tuple[nn.Module, OptimizerProtocol]:
+        """Build the PPO critic and its optimizer, restoring persisted state.
+
+        The critic backbone warm-starts from the policy weights (standard
+        actor-critic initialization; the fresh value head is the only
+        randomly-initialized part).  On resume, ``value_model`` /
+        ``value_optimizer`` checkpoint extras override the warm start —
+        and their absence is fatal rather than a silent fresh critic.
+        """
+        cfg = self.config
+        device = get_current_device()
+        checkpoint = context.checkpoint
+        if checkpoint is not None:
+            missing = [
+                name
+                for name in ("value_model", "value_optimizer")
+                if name not in checkpoint.extra
+            ]
+            if missing:
+                raise ValueError(
+                    "online_ppo resume requires critic state in the "
+                    f"checkpoint; missing extras: {', '.join(missing)}"
+                )
+
+        state_dict = executor.unwrap_model(context.model)
+        if executor.use_distributed:
+            state_dict = broadcast_state_dict(state_dict)
+        critic = cfg.critic_model_fn()
+        if state_dict is not None:
+            state_dict = strip_compile_prefix(state_dict)
+            result = critic.load_state_dict(state_dict, strict=False)
+            if result.unexpected_keys:
+                raise ValueError(
+                    "critic model received unexpected keys from the policy "
+                    f"state dict: {result.unexpected_keys[:3]}"
+                )
+            unexpected_missing = [
+                key for key in result.missing_keys if not key.startswith("value_head.")
+            ]
+            if unexpected_missing:
+                raise ValueError(
+                    "critic backbone is missing policy parameters: "
+                    f"{unexpected_missing[:3]}"
+                )
+        if checkpoint is not None:
+            critic.load_state_dict(checkpoint.extra["value_model"])
+        critic = critic.to(device)
+        critic.train()
+
+        optimizer_factory = cfg.critic_optimizer_fn or cfg.optimizer_fn
+        critic_optimizer = optimizer_factory(critic)
+        if checkpoint is not None:
+            critic_optimizer.load_state_dict(checkpoint.extra["value_optimizer"])
+        return critic, critic_optimizer
+
     def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> None:
         cfg = self.config
         if not cfg.strategy.startswith("online_"):

docs/developer/docker-training.md

@@ -186,6 +186,10 @@ scheduler.pt
 manifest.json
 ```
 
+`online_ppo` jobs additionally require `value_model.pt` and
+`value_optimizer.pt` (the critic state); the completeness check derives this
+from the training config's `train_type`.
+
 New checkpoints write `manifest.json` after every payload file, sync the complete
 staging directory, and then atomically rename that directory into place. Legacy
 checkpoints without a manifest remain resumable when the original required files

docs/guides/params.md

@@ -14,7 +14,7 @@
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--config`, `-c` | YAML config file; explicit CLI options override YAML values | None |
-| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`) | required |
+| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`, `online_ppo`) | required |
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--resume` | Resume training from `--param_path` | False |
@@ -139,16 +139,22 @@ with `--optimizer=muon_adamw`.
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo`, `online_dpo` |
 | `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` |
-| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo` |
+| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo`, `online_ppo` |
-| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo`, `online_grpo` |
+| `--grpo_clip_eps` | Clipping epsilon for the PPO-style surrogate loss | 0.2 | `grpo`, `online_grpo`, `online_ppo` |
-| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo`, `online_grpo` |
+| `--grpo_kl_coef` | KL penalty coefficient | 0.01 | `grpo`, `online_grpo`, `online_ppo` |
+| `--ppo_gamma` | PPO reward discount factor | 1.0 | `online_ppo` |
+| `--ppo_gae_lambda` | PPO GAE bias/variance trade-off | 0.95 | `online_ppo` |
+| `--ppo_vf_coef` | PPO value-loss coefficient | 0.5 | `online_ppo` |
 | `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |
 
 ### Online Rollout
 
 `online_grpo` and `online_dpo` are factory aliases for the existing `grpo` and
 `dpo` strategy classes; online behavior is enabled by rollout components rather
-than separate strategy subclasses. These options apply to the online aliases.
+than separate strategy subclasses. `online_ppo` is a dedicated actor-critic
+strategy: a `ValueModel` critic supplies GAE advantages, and its state persists
+as `value_model.pt`/`value_optimizer.pt` checkpoint extras (required for
+resume). These options apply to the online strategies.
 Online strategies require
 a `BaseRewardModel` factory in `TrainConfig`; `train.py` does not currently
 provide a command-line option for configuring one.

docs/guides/training.md

@@ -170,6 +170,18 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every
 behaviour log-probabilities into the loss, so it does not allocate or synchronize
 a separate old-policy model.
 
+`online_ppo` is actor-critic PPO on the same rollout pipeline. A `ValueModel`
+critic (backbone warm-started from the policy, zero-initialized value head)
+scores the rollout states; advantages come from GAE(`--ppo_gamma`,
+`--ppo_gae_lambda`) with the terminal reward on each response's last token and
+the reference-KL penalty (k3 estimator, `--grpo_kl_coef`) folded into per-token
+rewards. Advantages and returns are computed once per rollout and pinned on the
+`RolloutResult`, so replayed steps optimize fixed targets. The critic has its
+own optimizer, stepped outside the policy-version lock, and persists as
+`value_model.pt`/`value_optimizer.pt` checkpoint extras — resume without them
+fails loudly, and `scripts/train.sh` treats a PPO checkpoint as incomplete when
+they are missing.
+
 Every successful optimizer step mutates the shared model and advances its
 monotonic `policy_version` under the same generation lock. The scheduler
 invalidates reusable KV prefixes before accepting the new version, so an async

scripts/docker/lib/train-common.sh

@@ -22,13 +22,26 @@ validate_job_name() {
         die "Invalid TRAIN_JOB_NAME '$1'; use letters, numbers, dot, underscore, or dash"
 }
 
+checkpoint_extra_files() {
+    # Additional files a complete checkpoint must contain for the strategy
+    # configured in the given training YAML. PPO persists critic state as
+    # checkpoint extras (value_model.pt / value_optimizer.pt); a resume
+    # without them must not look complete.
+    local config="$1"
+
+    [[ -n "${config}" && -f "${config}" ]] || return 0
+    if grep -Eq '^[[:space:]]*train_type:[[:space:]]*["'\'']?online_ppo' "${config}"; then
+        printf 'value_model.pt value_optimizer.pt'
+    fi
+}
+
 checkpoint_is_complete() {
     local checkpoint="$1"
     local file
 
     [[ -d "${checkpoint}" ]] || return 1
 
-    for file in meta.json config.json model.safetensors optimizer.pt scheduler.pt; do
+    for file in meta.json config.json model.safetensors optimizer.pt scheduler.pt ${CHECKPOINT_EXTRA_FILES:-}; do
         [[ -s "${checkpoint}/${file}" ]] || return 1
     done
 

scripts/docker/train-entrypoint.sh

@@ -34,6 +34,7 @@ fi
 if [[ -n "${TRAIN_CONFIG}" ]]; then
     [[ -f "${TRAIN_CONFIG}" ]] || die "Training config not found: ${TRAIN_CONFIG}"
 fi
+export CHECKPOINT_EXTRA_FILES="$(checkpoint_extra_files "${TRAIN_CONFIG}")"
 [[ -r /data ]] || die "Training data directory is not readable: /data"
 
 mkdir -p "${CHECKPOINT_DIR}"

scripts/tools/train.py

@@ -21,7 +21,7 @@ from astrai.config.train_config import (
     TRAIN_TYPES,
 )
 from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
-from astrai.model import AutoRegressiveLM
+from astrai.model import AutoRegressiveLM, ValueModel
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.optim import OptimizerFactory
 from astrai.trainer import SchedulerFactory, Trainer
@@ -212,6 +212,27 @@ _SPECS = [
         default=0.01,
         help="GRPO KL penalty coefficient.",
     ),
+    OptSpec(
+        "ppo_gamma",
+        "Algorithm",
+        type=float,
+        default=1.0,
+        help="PPO reward discount factor.",
+    ),
+    OptSpec(
+        "ppo_gae_lambda",
+        "Algorithm",
+        type=float,
+        default=0.95,
+        help="PPO GAE bias/variance trade-off.",
+    ),
+    OptSpec(
+        "ppo_vf_coef",
+        "Algorithm",
+        type=float,
+        default=0.5,
+        help="PPO value-loss coefficient.",
+    ),
     OptSpec(
         "moe_aux_loss_coef",
         "Algorithm",
@@ -408,6 +429,10 @@ def create_model(config):
     return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
 
 
+def create_value_model(config):
+    return ValueModel(config).to(dtype=torch.bfloat16)
+
+
 def create_optimizer(
     model, optimizer_name: str = "muon_adamw", **kwargs
 ) -> optim.Optimizer:
@@ -502,6 +527,9 @@ def train(
         "clip_eps": kwargs.pop("grpo_clip_eps"),
         "kl_coef": kwargs.pop("grpo_kl_coef"),
         "group_size": kwargs.pop("group_size"),
+        "gamma": kwargs.pop("ppo_gamma"),
+        "gae_lambda": kwargs.pop("ppo_gae_lambda"),
+        "vf_coef": kwargs.pop("ppo_vf_coef"),
     }
 
     rollout_interval = kwargs.pop("rollout_interval", 512)
@@ -511,6 +539,11 @@ def train(
     rollout_top_p = kwargs.pop("rollout_top_p", 0.9)
     rollout_max_tokens = kwargs.pop("rollout_max_tokens", 1024)
     reward_model_fn: Callable[[], BaseRewardModel] | None = None
+    critic_model_fn = None
+    if train_type == "online_ppo":
+        # The optimizer defaults to the policy's; critic_optimizer_fn can
+        # override it in the TrainConfig.
+        critic_model_fn = partial(create_value_model, config)
 
     executor_kwargs = {}
     if parallel_mode == "ddp":
@@ -622,7 +655,7 @@ def train(
         collate_fn = dpo_collate_fn
     elif train_type == "grpo":
         collate_fn = grpo_collate_fn
-    elif train_type in ("online_grpo", "online_dpo"):
+    elif train_type in ("online_grpo", "online_dpo", "online_ppo"):
         collate_fn = None
 
     train_config = TrainConfig(
@@ -668,6 +701,7 @@ def train(
         rollout_top_p=rollout_top_p,
         rollout_max_tokens=rollout_max_tokens,
         reward_model_fn=reward_model_fn,
+        critic_model_fn=critic_model_fn,
         moe_aux_loss_coef=kwargs.pop("moe_aux_loss_coef", 0.01),
     )
 

scripts/train.sh

@@ -55,6 +55,7 @@ load_config() {
         die "Failed to load runtime configuration"
     eval "${exports}"
     validate_job_name "${TRAIN_JOB_NAME}"
+    export CHECKPOINT_EXTRA_FILES="$(checkpoint_extra_files "${CONFIG_FILE}")"
 }
 
 compose() {

tests/trainer/test_online_e2e.py

@@ -10,6 +10,7 @@ from torch.utils.data import Dataset
 import astrai.trainer.train_context as train_context
 from astrai.config import TrainConfig
 from astrai.model.transformer import AutoRegressiveLM
+from astrai.model.value import ValueModel
 from astrai.serialization import Checkpoint
 from astrai.trainer.rollout import BaseRewardModel
 from astrai.trainer.schedule import SchedulerFactory
@@ -67,6 +68,10 @@ def _model_fn(model_config):
     return AutoRegressiveLM(model_config).to(dtype=torch.float32)
 
 
+def _value_model_fn(model_config):
+    return ValueModel(model_config).to(dtype=torch.float32)
+
+
 def _optimizer_fn(m):
     return torch.optim.AdamW(m.parameters(), lr=1e-4)
 
@@ -81,16 +86,32 @@ _ONLINE_STRATEGIES = [
     pytest.param(
         "online_grpo",
         {"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2},
+        None,
         id="grpo",
     ),
-    pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, id="dpo"),
+    pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, None, id="dpo"),
+    pytest.param(
+        "online_ppo",
+        {
+            "clip_eps": 0.2,
+            "kl_coef": 0.01,
+            "group_size": 2,
+            "gamma": 1.0,
+            "gae_lambda": 0.95,
+            "vf_coef": 0.5,
+        },
+        True,
+        id="ppo",
+    ),
 ]
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES)
+@pytest.mark.parametrize(
+    ("strategy", "strategy_kwargs", "with_critic"), _ONLINE_STRATEGIES
+)
 def test_online_rollout_end_to_end(
-    base_test_env, strategy, strategy_kwargs, monkeypatch
+    base_test_env, strategy, strategy_kwargs, with_critic, monkeypatch
 ):
     """Run one epoch of online RL rollout with KV-cache-backed generation."""
     created_reference_models = []
@@ -110,7 +131,7 @@ def test_online_rollout_end_to_end(
     tokenizer.set_chat_template(CHAT_TEMPLATE)
     tokenizer.save_pretrained(test_dir)
 
-    train_config = TrainConfig(
+    config_kwargs = dict(
         strategy=strategy,
         model_fn=partial(_model_fn, model_config),
         dataset=InstructionDataset(),
@@ -135,6 +156,10 @@ def test_online_rollout_end_to_end(
         reward_model_fn=LengthRewardModel,
         collate_fn=instruction_collate_fn,
     )
+    if with_critic:
+        config_kwargs["critic_model_fn"] = partial(_value_model_fn, model_config)
+        config_kwargs["critic_optimizer_fn"] = _optimizer_fn
+    train_config = TrainConfig(**config_kwargs)
 
     trainer = Trainer(train_config)
     trainer.train(param_path=test_dir)
@@ -144,6 +169,11 @@ def test_online_rollout_end_to_end(
     checkpoint = Checkpoint.load(checkpoint_dir)
     assert checkpoint.meta["policy_version"] == 2
     assert len(created_reference_models) == 1
+    if with_critic:
+        assert "value_model" in checkpoint.extra
+        assert "value_optimizer" in checkpoint.extra
+    else:
+        assert "value_model" not in checkpoint.extra
 
 
 def _minimal_online_config(**overrides):

tests/trainer/test_ppo_strategy.py

@@ -0,0 +1,469 @@
+"""Unit tests for PPO: GAE numerics, the ValueModel critic, and PPOStrategy."""
+
+import pytest
+import torch
+
+import astrai.trainer.strategy as strategy_module
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.model.value import ValueModel
+from astrai.trainer.rollout import RolloutResult
+from astrai.trainer.strategy import (
+    PPOStrategy,
+    StrategyFactory,
+    compute_gae,
+)
+from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config
+
+
+def _make_batch(
+    batch_size=2, group_size=4, prompt_len=8, response_len=12, device="cpu"
+):
+    """Construct a PPO batch with deterministic shapes.
+
+    Returns dict with prompts [B, P], responses [B, G, R], masks [B, G, R],
+    rewards [B, G], logprobs_old [B, G, R].
+    """
+    return {
+        "prompts": torch.randint(0, 200, (batch_size, prompt_len), device=device),
+        "responses": torch.randint(
+            0, 200, (batch_size, group_size, response_len), device=device
+        ),
+        "masks": torch.ones(batch_size, group_size, response_len, device=device),
+        "rewards": torch.randn(batch_size, group_size, device=device),
+        "logprobs_old": torch.zeros(
+            batch_size, group_size, response_len, device=device
+        ),
+    }
+
+
+def _make_value_model(policy_model, device):
+    """Build a ValueModel whose backbone warm-starts from the policy."""
+    critic = ValueModel(policy_model.config).to(device=device)
+    result = critic.load_state_dict(policy_model.state_dict(), strict=False)
+    assert not result.unexpected_keys
+    assert all(key.startswith("value_head.") for key in result.missing_keys)
+    return critic
+
+
+@pytest.fixture
+def ppo_strategy(device):
+    model, _ = make_model(device)
+    critic = _make_value_model(model, device)
+    strategy = PPOStrategy(
+        model=model,
+        device=device,
+        critic=critic,
+        critic_optimizer=torch.optim.AdamW(critic.parameters(), lr=1e-3),
+        ref_model=make_frozen(model, device),
+        clip_eps=0.2,
+        kl_coef=0.01,
+        gamma=1.0,
+        gae_lambda=0.95,
+        vf_coef=0.5,
+        executor=FakeExecutor(),
+    )
+    return strategy, device
+
+
+# ============== compute_gae ==============
+
+
+def test_gae_monte_carlo_when_values_zero(device):
+    """γ=1, λ=1, V=0: advantage and return equal the terminal reward at
+    every valid position (Monte-Carlo return)."""
+    B, G, R = 2, 3, 4
+    rewards = torch.zeros(B, G, R, device=device)
+    rewards[..., -1] = 1.0
+    values = torch.zeros(B, G, R, device=device)
+    mask = torch.ones(B, G, R, dtype=torch.bool, device=device)
+
+    advantages, returns = compute_gae(rewards, values, mask, gamma=1.0, gae_lambda=1.0)
+
+    assert torch.allclose(advantages, torch.full_like(rewards, 1.0))
+    assert torch.allclose(returns, torch.full_like(rewards, 1.0))
+
+
+def test_gae_lambda_zero_is_one_step_td(device):
+    """λ=0: advantage degenerates to the TD residual δ_t."""
+    torch.manual_seed(0)
+    rewards = torch.zeros(1, 1, 3, device=device)
+    rewards[0, 0, -1] = 2.0
+    values = torch.tensor([[[0.5, 1.0, -0.5]]], device=device)
+    mask = torch.ones(1, 1, 3, dtype=torch.bool, device=device)
+
+    advantages, returns = compute_gae(rewards, values, mask, gamma=0.9, gae_lambda=0.0)
+
+    # δ_2 = r + 0 - V_2 = 2.5;  δ_1 = 0 + 0.9·V_2 - V_1 = -1.45;
+    # δ_0 = 0 + 0.9·V_1 - V_0 = 0.4
+    expected = torch.tensor([[[0.4, -1.45, 2.5]]], device=device)
+    assert torch.allclose(advantages, expected, atol=1e-6)
+    assert torch.allclose(returns, advantages + values, atol=1e-6)
+
+
+def test_gae_hand_computed_discounted_case(device):
+    """γ=0.9, λ=0.8 against a hand-rolled backward accumulation."""
+    rewards = torch.tensor([[[0.0, 0.0, 1.0]]], device=device)
+    values = torch.tensor([[[0.1, 0.2, 0.3]]], device=device)
+    mask = torch.ones(1, 1, 3, dtype=torch.bool, device=device)
+    gamma, lam = 0.9, 0.8
+
+    advantages, returns = compute_gae(rewards, values, mask, gamma, lam)
+
+    delta2 = 1.0 + 0.0 - 0.3
+    gae2 = delta2
+    delta1 = 0.0 + gamma * 0.3 - 0.2
+    gae1 = delta1 + gamma * lam * gae2
+    delta0 = 0.0 + gamma * 0.2 - 0.1
+    gae0 = delta0 + gamma * lam * gae1
+    expected = torch.tensor([[[gae0, gae1, gae2]]], device=device)
+    assert torch.allclose(advantages, expected, atol=1e-6)
+    assert torch.allclose(returns, expected + values, atol=1e-6)
+
+
+def test_gae_padding_does_not_leak(device):
+    """Garbage values at padded positions must not change valid outputs."""
+    torch.manual_seed(1)
+    B, G, R = 2, 2, 5
+    rewards = torch.zeros(B, G, R, device=device)
+    rewards[0, 0, 2] = 1.0  # terminal at position 2 of a length-3 response
+    values = torch.randn(B, G, R, device=device)
+    mask = torch.ones(B, G, R, dtype=torch.bool, device=device)
+    mask[0, 0, 3:] = False
+    mask[1, :, 2:] = False
+    rewards[0, 0, 3:] = 100.0  # reward garbage in padding must be ignored
+
+    advantages, returns = compute_gae(rewards, values, mask, gamma=0.9, gae_lambda=0.9)
+
+    assert torch.allclose(advantages[0, 0, 3:], torch.zeros_like(advantages[0, 0, 3:]))
+    assert torch.allclose(returns[0, 0, 3:], torch.zeros_like(returns[0, 0, 3:]))
+    # The terminal reward at position 2 still drives a finite advantage.
+    assert advantages[0, 0, 2] != 0.0
+
+
+def test_gae_empty_response_is_all_zero(device):
+    """A fully padded response yields zero advantages and returns."""
+    rewards = torch.zeros(1, 1, 3, device=device)
+    values = torch.randn(1, 1, 3, device=device)
+    mask = torch.zeros(1, 1, 3, dtype=torch.bool, device=device)
+
+    advantages, returns = compute_gae(rewards, values, mask, 1.0, 0.95)
+
+    assert torch.count_nonzero(advantages) == 0
+    assert torch.count_nonzero(returns) == 0
+
+
+# ============== ValueModel ==============
+
+
+def test_value_model_trunk_matches_policy_hidden_states(device):
+    """ValueModel's trunk reproduces AutoRegressiveLM's hidden states.
+
+    Pins the duplicated trunk pass in ``ValueModel.forward`` to the policy
+    forward: a ones-initialized value head must return the row-wise sum of
+    the policy's ``hidden_states``.
+    """
+    model, _ = make_model(device)
+    critic = _make_value_model(model, device)
+    with torch.no_grad():
+        critic.value_head.weight.fill_(1.0)
+        critic.value_head.bias.zero_()
+
+    torch.manual_seed(2)
+    input_ids = torch.randint(0, 200, (2, 10), device=device)
+    input_mask = torch.ones(2, 10, dtype=torch.bool, device=device)
+    input_mask[1, :3] = False
+
+    with torch.no_grad():
+        policy_hidden = model(input_ids, input_mask=input_mask)["hidden_states"]
+        values = critic(input_ids, input_mask=input_mask)["values"]
+
+    assert values.shape == (2, 10)
+    assert torch.allclose(values, policy_hidden.sum(dim=-1), atol=1e-5)
+
+
+def test_value_model_zero_head_outputs_zero(device):
+    model, _ = make_model(device)
+    critic = _make_value_model(model, device)
+    input_ids = torch.randint(0, 200, (2, 8), device=device)
+    with torch.no_grad():
+        values = critic(input_ids)["values"]
+    assert torch.count_nonzero(values) == 0
+
+
+def test_value_model_rejects_packed_inference_input(device):
+    critic = ValueModel(make_rollout_config()).to(device=device)
+    with pytest.raises(ValueError, match="critic input_ids"):
+        critic(torch.randint(0, 200, (16,), device=device))
+
+
+# ============== PPOStrategy ==============
+
+
+def test_factory_registers_online_ppo():
+    assert StrategyFactory.is_registered("online_ppo")
+    assert StrategyFactory.get_component_class("online_ppo") is PPOStrategy
+
+
+def test_ppo_supports_online(ppo_strategy):
+    strategy, _ = ppo_strategy
+    assert strategy.supports_online() is True
+
+
+def test_ppo_loss_is_finite_and_differentiable(ppo_strategy):
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    loss = strategy.compute_loss(batch)
+    assert loss.dim() == 0
+    assert torch.isfinite(loss).item()
+    loss.backward()
+    assert any(
+        p.grad is not None and p.grad.abs().sum().item() > 0
+        for p in strategy.model.parameters()
+    )
+    assert any(
+        p.grad is not None and p.grad.abs().sum().item() > 0
+        for p in strategy.critic.parameters()
+    )
+
+
+def test_ppo_requires_behavior_logprobs(ppo_strategy):
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    del batch["logprobs_old"]
+    with pytest.raises(ValueError, match="logprobs_old"):
+        strategy.compute_loss(batch)
+
+
+@pytest.mark.parametrize("invalid", ["shape", "nonfinite"])
+def test_ppo_rejects_invalid_behavior_logprobs(ppo_strategy, invalid):
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    if invalid == "shape":
+        batch["logprobs_old"] = torch.zeros(1, device=device)
+        match = "shape must match responses"
+    else:
+        batch["logprobs_old"] = torch.zeros_like(batch["responses"], dtype=torch.float)
+        batch["logprobs_old"][0, 0, 0] = float("nan")
+        match = "only finite values"
+    with pytest.raises(ValueError, match=match):
+        strategy.compute_loss(batch)
+
+
+def test_ppo_ref_model_not_updated_by_backward(ppo_strategy):
+    strategy, device = ppo_strategy
+    loss = strategy.compute_loss(_make_batch(device=device))
+    loss.backward()
+    for p in strategy.ref_model.parameters():
+        assert p.grad is None
+
+
+def test_ppo_zero_advantage_and_zero_critic_gives_zero_loss(ppo_strategy):
+    """A zero-head critic, zero advantages, and zero returns → zero loss."""
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    batch["advantages"] = torch.zeros_like(batch["responses"], dtype=torch.float)
+    batch["returns"] = torch.zeros_like(batch["responses"], dtype=torch.float)
+    loss = strategy.compute_loss(batch)
+    assert loss.item() == pytest.approx(0.0, abs=1e-6)
+
+
+def test_ppo_all_masked_response_tokens_zero_loss(ppo_strategy):
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    batch["masks"] = torch.zeros_like(batch["masks"])
+    loss = strategy.compute_loss(batch)
+    assert loss.item() == pytest.approx(0.0, abs=1e-6)
+
+
+def test_ppo_uses_supplied_advantages_without_recomputation(ppo_strategy):
+    """Explicit advantages/returns must short-circuit GAE computation."""
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+
+    def _fail(*args, **kwargs):
+        raise AssertionError("advantages were supplied; GAE must not run")
+
+    strategy._compute_advantages = _fail
+    batch["advantages"] = torch.ones_like(batch["responses"], dtype=torch.float)
+    batch["returns"] = torch.zeros_like(batch["responses"], dtype=torch.float)
+    loss = strategy.compute_loss(batch)
+    assert torch.isfinite(loss).item()
+
+
+def test_ppo_optimizer_step_updates_policy_and_critic(ppo_strategy):
+    """optimizer_step steps the policy optimizer and then the critic's."""
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    loss = strategy.compute_loss(batch)
+    loss.backward()
+
+    policy_optimizer = torch.optim.SGD(strategy.model.parameters(), lr=0.1)
+    policy_before = next(strategy.model.parameters()).detach().clone()
+    critic_before = next(strategy.critic.parameters()).detach().clone()
+
+    strategy.optimizer_step(policy_optimizer)
+
+    assert not torch.equal(next(strategy.model.parameters()), policy_before)
+    assert not torch.equal(next(strategy.critic.parameters()), critic_before)
+    # Critic gradients are cleared after its step.
+    assert all(p.grad is None for p in strategy.critic.parameters())
+
+
+def test_ppo_optimizer_step_clips_critic_gradients(ppo_strategy):
+    strategy, device = ppo_strategy
+    batch = _make_batch(device=device)
+    loss = strategy.compute_loss(batch)
+    loss.backward()
+    strategy.critic_optimizer = torch.optim.SGD(strategy.critic.parameters(), lr=1.0)
+    strategy.max_grad_norm = 1e-8
+    before = next(strategy.critic.parameters()).detach().clone()
+
+    strategy.optimizer_step(torch.optim.SGD(strategy.model.parameters(), lr=0.0))
+
+    # Clipped-to-zero critic gradients under SGD (no momentum) leave the
+    # parameters unchanged.
+    assert torch.equal(next(strategy.critic.parameters()), before)
+
+
+# ============== prepare_from_rollout / GAE integration ==============
+
+
+def _make_rollout_result(B=2, G=2, P=6, R=5, device="cpu"):
+    return RolloutResult(
+        prompts=torch.randint(3, 200, (B, P), device=device),
+        prompt_mask=torch.ones(B, P, dtype=torch.bool, device=device),
+        responses=torch.randint(3, 200, (B, G, R), device=device),
+        response_mask=torch.ones(B, G, R, dtype=torch.bool, device=device),
+        rewards=torch.randn(B, G, device=device),
+        logprobs_old=torch.zeros(B, G, R, device=device),
+    )
+
+
+def test_prepare_from_rollout_computes_and_pins_gae(ppo_strategy, monkeypatch):
+    """prepare attaches GAE tensors to the result once, then reuses them."""
+    strategy, device = ppo_strategy
+    result = _make_rollout_result(device=device)
+
+    batch = strategy.prepare_from_rollout(result)
+    assert result.advantages is not None and result.returns is not None
+    assert batch["advantages"] is result.advantages
+    assert batch["returns"] is result.returns
+    assert batch["advantages"].shape == result.responses.shape
+
+    calls = []
+    original = strategy._compute_advantages
+    monkeypatch.setattr(
+        strategy,
+        "_compute_advantages",
+        lambda *a, **k: calls.append(1) or original(*a, **k),
+    )
+    strategy.prepare_from_rollout(result)
+    assert not calls, "pinned advantages must not be recomputed on replay"
+
+
+def test_prepare_from_rollout_respects_response_padding(ppo_strategy):
+    """Padded response positions get zero advantages and returns."""
+    strategy, device = ppo_strategy
+    result = _make_rollout_result(device=device)
+    result.response_mask[0, 0, 3:] = False
+
+    batch = strategy.prepare_from_rollout(result)
+
+    assert torch.count_nonzero(batch["advantages"][0, 0, 3:]) == 0
+    assert torch.count_nonzero(batch["returns"][0, 0, 3:]) == 0
+    assert torch.count_nonzero(batch["advantages"][0, 0, :3]) > 0
+
+
+def test_compute_advantages_matches_hand_computed_gae(ppo_strategy, monkeypatch):
+    """_compute_advantages applies terminal rewards and GAE faithfully."""
+    strategy, device = ppo_strategy
+    result = _make_rollout_result(B=1, G=1, P=4, R=3, device=device)
+    result.rewards = torch.tensor([[2.0]], device=device)
+    result.logprobs_old = torch.zeros(1, 1, 3, device=device)
+    # ref_model == policy at init → zero KL reward shaping only if the
+    # policy and ref agree on logprobs; keep ref out of the picture here.
+    strategy.ref_model = None
+
+    fixed_values = torch.tensor([[[0.1, 0.2, 0.3]]], device=device)
+    monkeypatch.setattr(
+        strategy_module,
+        "rollout_token_values",
+        lambda *args, **kwargs: fixed_values.clone(),
+    )
+
+    advantages, returns = strategy._compute_advantages(
+        result.prompts,
+        result.prompt_mask,
+        result.responses,
+        result.response_mask,
+        result.rewards,
+        result.logprobs_old,
+    )
+
+    rewards = torch.tensor([[[0.0, 0.0, 2.0]]], device=device)
+    expected_adv, expected_ret = compute_gae(
+        rewards, fixed_values, result.response_mask, 1.0, 0.95
+    )
+    assert torch.allclose(advantages, expected_adv, atol=1e-6)
+    assert torch.allclose(returns, expected_ret, atol=1e-6)
+
+
+def test_compute_advantages_folds_kl_penalty_into_rewards(ppo_strategy, monkeypatch):
+    """With a ref model, each valid token's reward loses kl_coef·k3."""
+    strategy, device = ppo_strategy
+    result = _make_rollout_result(B=1, G=1, P=4, R=2, device=device)
+    result.rewards = torch.tensor([[1.0]], device=device)
+    # behaviour policy disagrees with ref by +1 logprob on every token
+    result.logprobs_old = torch.ones(1, 1, 2, device=device)
+
+    fixed_values = torch.zeros(1, 1, 2, device=device)
+    fixed_ref_logprobs = torch.zeros(1, 1, 2, device=device)
+    monkeypatch.setattr(
+        strategy_module,
+        "rollout_token_values",
+        lambda *args, **kwargs: fixed_values.clone(),
+    )
+    monkeypatch.setattr(
+        strategy_module,
+        "rollout_token_logprobs",
+        lambda *args, **kwargs: {"logprobs": fixed_ref_logprobs.clone()},
+    )
+
+    advantages, _ = strategy._compute_advantages(
+        result.prompts,
+        result.prompt_mask,
+        result.responses,
+        result.response_mask,
+        result.rewards,
+        result.logprobs_old,
+    )
+
+    # per-token reward = -kl_coef·(1 - 0) = -0.01; terminal adds 1.0
+    expected_rewards = torch.tensor([[-0.01, 0.99]], device=device)
+    expected_adv, _ = compute_gae(
+        expected_rewards.unsqueeze(0), fixed_values, result.response_mask, 1.0, 0.95
+    )
+    assert torch.allclose(advantages, expected_adv, atol=1e-6)
+
+
+def test_online_call_returns_finite_loss(ppo_strategy):
+    strategy, device = ppo_strategy
+
+    class _RecordingRunner:
+        policy_version = 0
+
+        def __call__(self, batch):
+            return _make_rollout_result(device=device), True
+
+        def step(self):
+            pass
+
+        def apply_weight_update(self, policy_version, update):
+            return update()
+
+    strategy.set_rollout_runner(_RecordingRunner())
+    out = strategy({"instruction": ["x"]})
+    assert torch.isfinite(out["loss"]).item()
+    assert "policy_loss" in out["metrics"]
+    assert "value_loss" in out["metrics"]
+    assert "explained_variance" in out["metrics"]

tests/trainer/test_ppo_train_context.py

@@ -0,0 +1,271 @@
+"""Wiring tests for online PPO: config validation, critic assembly, and
+checkpoint round-trip of critic state."""
+
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import torch
+
+from astrai.config import TrainConfig
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.model.value import ValueModel
+from astrai.serialization import Checkpoint
+from astrai.trainer.rollout import BaseRewardModel
+from astrai.trainer.schedule import SchedulerFactory
+from astrai.trainer.train_callback import CheckpointCallback
+from astrai.trainer.train_context import TrainContext, TrainContextBuilder
+from astrai.trainer.trainer import Trainer
+from tests.helpers import (
+    FakeExecutor,
+    build_test_tokenizer,
+    make_model,
+    make_rollout_config,
+)
+
+
+class _StubRewardModel(BaseRewardModel):
+    def score(self, prompts, responses):
+        return torch.zeros(len(prompts), len(responses[0]) if prompts else 0)
+
+
+class _StubDataset(torch.utils.data.Dataset):
+    def __len__(self):
+        return 2
+
+    def __getitem__(self, idx):
+        return {"instruction": "hello", "input": ""}
+
+
+def _stub_collate(batch):
+    return {
+        "instruction": [b["instruction"] for b in batch],
+        "input": [b.get("input", "") for b in batch],
+    }
+
+
+def _ppo_config(device, **overrides):
+    defaults = dict(
+        strategy="online_ppo",
+        model_fn=lambda: AutoRegressiveLM(make_rollout_config()),
+        dataset=_StubDataset(),
+        optimizer_fn=lambda m: torch.optim.SGD(m.parameters(), lr=0.0),
+        scheduler_fn=lambda o: SchedulerFactory.create(
+            "cosine", o, warmup_steps=1, lr_decay_steps=4, min_rate=0.05
+        ),
+        reward_model_fn=_StubRewardModel,
+        critic_model_fn=lambda: ValueModel(make_rollout_config()),
+        collate_fn=_stub_collate,
+        device_type=device,
+        nprocs=1,
+        parallel_mode="none",
+        strategy_kwargs={"clip_eps": 0.2, "group_size": 2},
+        rollout_interval=1,
+        rollout_max_policy_lag=0,
+        rollout_max_tokens=4,
+        rollout_temperature=1.0,
+    )
+    defaults.update(overrides)
+    return TrainConfig(**defaults)
+
+
+def test_online_ppo_config_requires_critic_model_fn(device):
+    with pytest.raises(ValueError, match="critic_model_fn is required"):
+        _ppo_config(device, critic_model_fn=None)
+
+
+def test_online_ppo_config_accepts_critic(device):
+    config = _ppo_config(device)
+    assert config.strategy == "online_ppo"
+
+
+def test_create_critic_warm_starts_backbone_from_policy(device, monkeypatch):
+    monkeypatch.setenv("LOCAL_DEVICE", device)
+    model, config = make_model(device)
+    cfg = _ppo_config(device)
+    builder = TrainContextBuilder(cfg)
+    context = TrainContext(model=model)
+
+    critic, _ = builder._create_critic(context, FakeExecutor())
+
+    policy_sd = model.state_dict()
+    critic_sd = critic.state_dict()
+    for key in policy_sd:
+        assert torch.equal(critic_sd[key], policy_sd[key])
+    assert torch.count_nonzero(critic_sd["value_head.weight"]) == 0
+    assert torch.count_nonzero(critic_sd["value_head.bias"]) == 0
+
+
+def test_create_critic_restores_checkpoint_extras(device, monkeypatch):
+    monkeypatch.setenv("LOCAL_DEVICE", device)
+    model, config = make_model(device)
+    cfg = _ppo_config(device)
+    builder = TrainContextBuilder(cfg)
+
+    saved_critic = ValueModel(config).to(device)
+    with torch.no_grad():
+        saved_critic.value_head.weight.fill_(1.0)
+    saved_optimizer = torch.optim.SGD(saved_critic.parameters(), lr=0.1)
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(),
+        config=config.to_dict(),
+        extra={
+            "optimizer": {},
+            "scheduler": {},
+            "value_model": saved_critic.state_dict(),
+            "value_optimizer": saved_optimizer.state_dict(),
+        },
+    )
+    context = TrainContext(model=model, checkpoint=checkpoint)
+
+    critic, critic_optimizer = builder._create_critic(context, FakeExecutor())
+
+    assert torch.equal(
+        critic.state_dict()["value_head.weight"],
+        saved_critic.state_dict()["value_head.weight"],
+    )
+    assert (
+        critic_optimizer.state_dict()["param_groups"]
+        == saved_optimizer.state_dict()["param_groups"]
+    )
+
+
+def test_create_critic_resume_without_extras_fails_loudly(device, monkeypatch):
+    monkeypatch.setenv("LOCAL_DEVICE", device)
+    model, _ = make_model(device)
+    cfg = _ppo_config(device)
+    builder = TrainContextBuilder(cfg)
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(),
+        extra={"optimizer": {}, "scheduler": {}},
+    )
+    context = TrainContext(model=model, checkpoint=checkpoint)
+
+    with pytest.raises(
+        ValueError, match="missing extras: value_model, value_optimizer"
+    ):
+        builder._create_critic(context, FakeExecutor())
+
+
+def test_builder_resumes_critic_from_checkpoint(device, temp_dir, monkeypatch):
+    """A full TrainContextBuilder resume restores the persisted critic."""
+    monkeypatch.setenv("LOCAL_DEVICE", device)
+    model, config = make_model(device)
+    saved_critic = ValueModel(config).to(device)
+    with torch.no_grad():
+        saved_critic.value_head.weight.fill_(2.0)
+    saved_optimizer = torch.optim.SGD(saved_critic.parameters(), lr=0.1)
+    policy_optimizer = torch.optim.SGD(model.parameters(), lr=0.0)
+    policy_scheduler = SchedulerFactory.create(
+        "cosine", policy_optimizer, warmup_steps=1, lr_decay_steps=4, min_rate=0.05
+    )
+    checkpoint = Checkpoint(
+        state_dict=model.state_dict(),
+        epoch=0,
+        consumed_samples=2,
+        config=config.to_dict(),
+        extra={
+            "optimizer": policy_optimizer.state_dict(),
+            "scheduler": policy_scheduler.state_dict(),
+            "value_model": saved_critic.state_dict(),
+            "value_optimizer": saved_optimizer.state_dict(),
+        },
+        meta={"policy_version": 3},
+    )
+    checkpoint.save(temp_dir)
+    build_test_tokenizer(vocab_size=200).save_pretrained(temp_dir)
+
+    cfg = _ppo_config(
+        device,
+        model_fn=lambda: AutoRegressiveLM(config),
+        critic_model_fn=lambda: ValueModel(config),
+        ckpt_dir=os.path.join(temp_dir, "ckpt"),
+    )
+    context = TrainContextBuilder(cfg).with_param_path(temp_dir, resume=True).build()
+
+    assert isinstance(context.strategy.critic, ValueModel)
+    assert torch.equal(
+        context.strategy.critic.state_dict()["value_head.weight"],
+        saved_critic.state_dict()["value_head.weight"],
+    )
+    assert context.strategy.policy_version == 3
+
+
+def test_save_extra_persists_critic_state(device):
+    model, _ = make_model(device)
+    critic = ValueModel(make_rollout_config()).to(device)
+    from astrai.trainer.strategy import PPOStrategy
+
+    strategy = PPOStrategy(
+        model=model,
+        device=device,
+        critic=critic,
+        critic_optimizer=torch.optim.SGD(critic.parameters(), lr=0.0),
+        executor=FakeExecutor(),
+    )
+    context = TrainContext(strategy=strategy)
+
+    extra = CheckpointCallback.save_extra(context)
+
+    assert set(extra) == {"value_model", "value_optimizer"}
+    saved = extra["value_model"]
+    live = critic.state_dict()
+    assert set(saved) == set(live)
+    for key in saved:
+        assert torch.equal(saved[key], live[key])
+
+
+def test_save_extra_without_critic_has_no_value_entries(device):
+    model, _ = make_model(device)
+    from astrai.trainer.strategy import GRPOStrategy
+    from tests.helpers import make_frozen
+
+    strategy = GRPOStrategy(
+        model=model,
+        device=device,
+        old_model=None,
+        ref_model=make_frozen(model, device),
+        executor=FakeExecutor(),
+    )
+    context = TrainContext(strategy=strategy)
+
+    extra = CheckpointCallback.save_extra(context)
+
+    assert "value_model" not in extra
+    assert "value_optimizer" not in extra
+
+
+def test_trainer_default_callbacks_do_not_break_ppo(device, temp_dir):
+    """The Trainer's default callback set constructs fine for online_ppo."""
+    cfg = _ppo_config(device, ckpt_dir=os.path.join(temp_dir, "ckpt"))
+    trainer = Trainer(cfg)
+    assert trainer.callbacks
+
+
+def test_sh_checkpoint_extra_files_detects_online_ppo(temp_dir):
+    """The shell completeness helper derives PPO's extra required files."""
+    lib = (
+        Path(__file__).resolve().parents[2]
+        / "scripts"
+        / "docker"
+        / "lib"
+        / "train-common.sh"
+    )
+    ppo_yaml = Path(temp_dir) / "ppo.yaml"
+    ppo_yaml.write_text("train_type: online_ppo\n")
+    grpo_yaml = Path(temp_dir) / "grpo.yaml"
+    grpo_yaml.write_text('train_type: "online_grpo"\n')
+    quoted_yaml = Path(temp_dir) / "quoted.yaml"
+    quoted_yaml.write_text('  train_type:   "online_ppo"\n')
+
+    def extra_files(yaml_path):
+        script = f'source "{lib}"; checkpoint_extra_files "{yaml_path}"'
+        result = subprocess.run(
+            ["bash", "-c", script], capture_output=True, text=True, check=True
+        )
+        return result.stdout.strip()
+
+    assert extra_files(ppo_yaml) == "value_model.pt value_optimizer.pt"
+    assert extra_files(quoted_yaml) == "value_model.pt value_optimizer.pt"
+    assert extra_files(grpo_yaml) == ""