From 7395a8394e594a82accc53c712231ede6bc4fda4 Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sat, 25 Jul 2026 17:55:51 +0100
Subject: [PATCH 01/11] feat(models): add AR value head and GAE track wiring

Wire per-token critic values through Qwen3 replay and compute GAE
advantages on RolloutTrack for the PPO critic path (issue #86, part 2/3).
---
 unirl/models/qwen3/ar.py         | 162 ++++++++++++++++++++++++++-----
 unirl/models/qwen3/bundle.py     |  10 +-
 unirl/models/qwen3/config.py     |   3 +
 unirl/models/types/value_head.py |  24 +++++
 unirl/types/advantages.py        |  42 ++++++++
 unirl/types/segments/text.py     |   4 +
 6 files changed, 216 insertions(+), 29 deletions(-)
 create mode 100644 unirl/models/types/value_head.py

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index 1e072ac76..6bdccf2d4 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -23,7 +23,7 @@
 from dataclasses import dataclass
 from dataclasses import field as dc_field
 from types import MethodType
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
@@ -31,6 +31,7 @@
 from torch.utils.checkpoint import checkpoint
 
 from unirl.models.types.ar import ARSamplingParams, ARStage, ARStep, left_pad_prompt
+from unirl.models.types.ar_replay import ARReplayOutput
 from unirl.types.segments import TextSegment
 from unirl.utils.dtypes import parse_torch_dtype
 
@@ -41,6 +42,11 @@
 
 _SPARSE_PACKED_ATTN = ("flex_attention", "flash_attention_2", "flash_attention_3", "flash_attention_4")
 
+try:
+    from transformers.masking_utils import find_packed_sequence_indices
+except Exception:
+    find_packed_sequence_indices = None
+
 
 @functools.lru_cache(maxsize=None)
 def _warn_packed_disabled(attn_impl: str) -> None:
@@ -72,9 +78,7 @@ def _packed_replay_supported(attn_impl: Optional[str]) -> bool:
     if attn_impl not in _SPARSE_PACKED_ATTN:
         _warn_packed_disabled(str(attn_impl))
         return False
-    try:
-        from transformers.masking_utils import find_packed_sequence_indices  # noqa: F401
-    except Exception:
+    if find_packed_sequence_indices is None:
         return False
     return True
 
@@ -87,6 +91,7 @@ def _replay_aware_forward(
     temperature: float = 1.0,
     autocast_dtype: Optional[torch.dtype] = None,
     packed_predict_index: Optional[torch.Tensor] = None,
+    return_values: bool = False,
     **kw: Any,
 ) -> Any:
     """Dual-mode ``forward`` installed on the Qwen3 CausalLM instance.
@@ -129,6 +134,7 @@ def _replay_aware_forward(
     # [B, chunk, vocab] FP32 transient stays ~1.2 GiB, and each chunk is
     # gradient-checkpointed (recomputed in backward rather than held).
     T = float(temperature) if float(temperature) > 0.0 else 1.0
+    value_head = getattr(self, "value_head", None) if return_values else None
 
     if packed_predict_index is not None:
         # Packed varlen replay: ``hidden`` is one packed row [1, L_total, H]
@@ -155,8 +161,18 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
             else:
                 flat_parts.append(_flat_logp_chunk(h, tok))
         if not flat_parts:
-            return hidden.new_zeros((0,), dtype=torch.float32)
-        return torch.cat(flat_parts, dim=0)
+            empty = hidden.new_zeros((0,), dtype=torch.float32)
+            if value_head is None:
+                return empty
+            return ARReplayOutput(log_probs=empty, values=empty)
+        log_probs = torch.cat(flat_parts, dim=0)
+        if value_head is None:
+            return log_probs
+        value_parts: List[torch.Tensor] = []
+        for s in range(0, int(h_pred.size(0)), flat_chunk):
+            value_parts.append(value_head(h_pred[s : s + flat_chunk]))
+        values = torch.cat(value_parts, dim=0) if value_parts else log_probs.new_zeros(0)
+        return ARReplayOutput(log_probs=log_probs, values=values)
     T_max = int(response_tokens.size(1))
     resp_hidden = hidden[:, prompt_len - 1 : prompt_len - 1 + T_max, :]
 
@@ -176,8 +192,63 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
         else:
             parts.append(_logp_chunk(h, tok))
     if not parts:
-        return resp_hidden.new_zeros((bsz, 0), dtype=torch.float32)  # T_max == 0
-    return torch.cat(parts, dim=1)
+        empty = resp_hidden.new_zeros((bsz, 0), dtype=torch.float32)
+        if value_head is None:
+            return empty
+        return ARReplayOutput(log_probs=empty, values=empty)
+    log_probs = torch.cat(parts, dim=1)
+    if value_head is None:
+        return log_probs
+    value_parts = []
+    for s in range(0, T_max, chunk):
+        value_parts.append(value_head(resp_hidden[:, s : s + chunk, :]))
+    values = torch.cat(value_parts, dim=1) if value_parts else log_probs.new_zeros((bsz, 0))
+    return ARReplayOutput(log_probs=log_probs, values=values)
+
+
+def _require_value_head_for_replay(model: Any, return_values: bool) -> None:
+    if return_values and getattr(model, "value_head", None) is None:
+        raise ValueError(
+            "Qwen3ARStage.replay: return_values=True requires a value head "
+            "(set use_value_head=True in the pipeline config)"
+        )
+
+
+def _finalize_replay_output(
+    out: Union[torch.Tensor, ARReplayOutput],
+    *,
+    segment: TextSegment,
+    return_values: bool,
+    logprob_dtype: torch.dtype,
+    device: torch.device,
+) -> Union[torch.Tensor, ARReplayOutput]:
+    """Cast log-probs and flatten packed values to match ``segment`` layout."""
+    if isinstance(out, ARReplayOutput):
+        log_probs = out.log_probs.to(dtype=logprob_dtype)
+        if not return_values:
+            return log_probs
+        values = out.values
+        if values is None:
+            raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
+        if log_probs.ndim == 1:
+            return ARReplayOutput(log_probs=log_probs, values=values.to(device=device))
+        if segment.cu_seqlens is None or segment.lengths is None:
+            raise ValueError("Qwen3ARStage.replay: segment requires cu_seqlens to flatten values")
+        lengths = [int(n) for n in segment.lengths.tolist()]
+        cu = [int(c) for c in segment.cu_seqlens.tolist()]
+        flat: List[torch.Tensor] = []
+        for b, n in enumerate(lengths):
+            if n <= 0:
+                continue
+            flat.append(values[b, :n])
+        packed_values = torch.cat(flat, dim=0) if flat else values.new_zeros(0, device=device)
+        return ARReplayOutput(log_probs=log_probs, values=packed_values.to(device=device))
+    if return_values:
+        raise ValueError(
+            "Qwen3ARStage.replay: return_values=True but critic returned no values "
+            "(set use_value_head=True in the pipeline config)"
+        )
+    return out.to(dtype=logprob_dtype)
 
 
 # Attention backends with a sparse packed kernel (skip cross-sequence blocks):
@@ -421,22 +492,28 @@ def replay(
         *,
         segment: TextSegment,
         temperature: float = 1.0,
-    ) -> torch.Tensor:
+        return_values: bool = False,
+    ) -> Union[torch.Tensor, ARReplayOutput]:
         """Per-token log-prob replay over a stored rollout segment.
 
         Branch: prefer :meth:`packed_replay` (packed-varlen, zero padding, B > 1)
         and fall back to :meth:`padding_replay` (the dense ``[B, P_max + T_max]``
         padded path) when packing does not apply. Returns packed varlen
-        ``[total_tokens]`` aligned with ``segment.log_probs``; caller controls
-        grad / ``.train()`` scope. ``temperature`` divides logits before
-        ``log_softmax`` to match SGLang's sampler (``1.0`` is a no-op).
+        ``[total_tokens]`` aligned with ``segment.log_probs`` unless
+        ``return_values=True``, in which case an :class:`ARReplayOutput` with
+        packed ``values`` is returned alongside log-probs.
         """
+        _require_value_head_for_replay(self.model.transformer, return_values)
         attn_impl = getattr(getattr(self.model.transformer, "config", None), "_attn_implementation", None)
         if _packed_replay_supported(attn_impl):
-            packed = self.packed_replay(conditions, segment=segment, temperature=temperature)
+            packed = self.packed_replay(
+                conditions, segment=segment, temperature=temperature, return_values=return_values
+            )
             if packed is not None:
                 return packed
-        return self.padding_replay(conditions, segment=segment, temperature=temperature)
+        return self.padding_replay(
+            conditions, segment=segment, temperature=temperature, return_values=return_values
+        )
 
     def packed_replay(
         self,
@@ -444,7 +521,8 @@ def packed_replay(
         *,
         segment: TextSegment,
         temperature: float = 1.0,
-    ) -> Optional[torch.Tensor]:
+        return_values: bool = False,
+    ) -> Optional[Union[torch.Tensor, ARReplayOutput]]:
         """Packed-varlen replay (B > 1): zero padding anywhere.
 
         Concatenate every sample's REAL prompt tokens + its flat response tokens
@@ -520,9 +598,16 @@ def packed_replay(
             packed_predict_index=predict_index,
             prompt_len=0,
             temperature=temperature,
+            return_values=return_values,
             autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None),
         )
-        return per_token_flat.to(dtype=self.logprob_dtype)
+        return _finalize_replay_output(
+            per_token_flat,
+            segment=segment,
+            return_values=return_values,
+            logprob_dtype=self.logprob_dtype,
+            device=device,
+        )
 
     def padding_replay(
         self,
@@ -530,7 +615,8 @@ def padding_replay(
         *,
         segment: TextSegment,
         temperature: float = 1.0,
-    ) -> torch.Tensor:
+        return_values: bool = False,
+    ) -> Union[torch.Tensor, ARReplayOutput]:
         """Dense ``[B, P_max + T_max]`` padded replay — the default / fallback path.
 
         One teacher-forced forward over padded ``prompt + response``; gather
@@ -629,28 +715,54 @@ def padding_replay(
         # root-wrapped or plain) and never materializes [B, L, vocab] logits.
         # The cuda-vs-cpu autocast decision lives here; dtype validity and the
         # autocast scope live in the patched forward.
-        per_token = self.model.transformer(
+        out = self.model.transformer(
             input_ids=full_ids,
             attention_mask=full_mask,
             position_ids=position_ids,
             response_tokens=response_tokens,
             prompt_len=prompt_len,
             temperature=temperature,
+            return_values=return_values,
             autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None),
-        )  # [B, T_max] FP32
+        )  # [B, T_max] FP32 or ARReplayOutput
 
         if T_max == 0:
-            return torch.zeros(0, dtype=self.logprob_dtype, device=device)
+            empty = torch.zeros(0, dtype=self.logprob_dtype, device=device)
+            if return_values:
+                return ARReplayOutput(log_probs=empty, values=empty)
+            return empty
+
+        if isinstance(out, ARReplayOutput):
+            per_token = out.log_probs
+            per_value = out.values
+        else:
+            per_token = out
+            per_value = None
+        if return_values and per_value is None:
+            raise ValueError(
+                "Qwen3ARStage.replay: return_values=True but critic returned no values "
+                "(set use_value_head=True in the pipeline config)"
+            )
 
-        flat: List[torch.Tensor] = []
+        flat_logp: List[torch.Tensor] = []
+        flat_val: List[torch.Tensor] = []
         for b in range(batch_size):
             n = lengths[b]
             if n == 0:
                 continue
-            flat.append(per_token[b, :n])
-        if not flat:
-            return torch.zeros(0, dtype=self.logprob_dtype, device=device)
-        return torch.cat(flat, dim=0).to(dtype=self.logprob_dtype)
+            flat_logp.append(per_token[b, :n])
+            if per_value is not None:
+                flat_val.append(per_value[b, :n])
+        if not flat_logp:
+            empty = torch.zeros(0, dtype=self.logprob_dtype, device=device)
+            if return_values:
+                return ARReplayOutput(log_probs=empty, values=empty)
+            return empty
+        log_probs = torch.cat(flat_logp, dim=0).to(dtype=self.logprob_dtype)
+        if not return_values:
+            return log_probs
+        values = torch.cat(flat_val, dim=0)
+        return ARReplayOutput(log_probs=log_probs, values=values)
 
     def _resolve_stop_ids(
         self,
diff --git a/unirl/models/qwen3/bundle.py b/unirl/models/qwen3/bundle.py
index 3c13308a8..e3d3325f5 100644
--- a/unirl/models/qwen3/bundle.py
+++ b/unirl/models/qwen3/bundle.py
@@ -25,9 +25,11 @@
 
 import torch
 import torch.nn as nn
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
 from unirl.models.types.bundle import Bundle
 from unirl.models.types.meta_init import build_meta_init_transformer
+from unirl.models.types.value_head import ValueHead
 from unirl.utils.dtypes import parse_torch_dtype
 
 from .config import Qwen3PipelineConfig
@@ -57,8 +59,6 @@ def __init__(
     @classmethod
     def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
         """Load the Qwen3 transformer + tokenizer from a HuggingFace-layout checkpoint."""
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-
         path = config.pretrained_model_ckpt_path
         tokenizer_path = config.tokenizer_ckpt_path or path
 
@@ -78,8 +78,6 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
             # checkpoint, so to_empty later clobbers them -> garbage RoPE. It
             # captures them; meta_init_state is stashed on the BUNDLE below and
             # restored by load_trainable_weights after the sharded weight load.
-            from transformers import AutoConfig
-
             hf_config = AutoConfig.from_pretrained(path, trust_remote_code=bool(config.trust_remote_code))
             transformer, meta_init_state = build_meta_init_transformer(
                 lambda: AutoModelForCausalLM.from_config(hf_config, trust_remote_code=bool(config.trust_remote_code)),
@@ -115,6 +113,10 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
         if tokenizer.pad_token is None and tokenizer.eos_token is not None:
             tokenizer.pad_token = tokenizer.eos_token
 
+        if config.use_value_head:
+            hidden_size = int(getattr(transformer.config, "hidden_size"))
+            transformer.value_head = ValueHead(hidden_size).to(device)
+
         bundle = cls(
             transformer=transformer,
             tokenizer=tokenizer,
diff --git a/unirl/models/qwen3/config.py b/unirl/models/qwen3/config.py
index 37f63b5ef..42a6c51cb 100644
--- a/unirl/models/qwen3/config.py
+++ b/unirl/models/qwen3/config.py
@@ -69,6 +69,9 @@ class Qwen3PipelineConfig:
     use_lora: bool = False
     lora_target_modules: Optional[List[str]] = None
 
+    # Attach a scalar value head on the transformer for PPO / GAE training.
+    use_value_head: bool = False
+
     system_instruction: Optional[str] = None
     # Chat-template thinking switch; MUST agree with the rollout engine's
     # chat_template_kwargs.enable_thinking or train/rollout prompts diverge.
diff --git a/unirl/models/types/value_head.py b/unirl/models/types/value_head.py
new file mode 100644
index 000000000..5a2d1969e
--- /dev/null
+++ b/unirl/models/types/value_head.py
@@ -0,0 +1,24 @@
+"""Scalar value head for PPO-style critic training on AR hidden states."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ValueHead(nn.Module):
+    """Linear critic ``V(h)`` on last hidden states.
+
+    Kept in FP32 for stable value loss math (mirrors replay log-prob FP32 policy).
+    """
+
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, 1, bias=True, dtype=torch.float32)
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        """Map ``[..., H]`` hidden states to ``[...,]`` scalar values."""
+        return self.proj(hidden.float()).squeeze(-1)
+
+
+__all__ = ["ValueHead"]
diff --git a/unirl/types/advantages.py b/unirl/types/advantages.py
index f49f96b6e..1140d23c9 100644
--- a/unirl/types/advantages.py
+++ b/unirl/types/advantages.py
@@ -71,6 +71,48 @@ def compute_gae_advantages(
     return advantages, returns
 
 
+def scatter_terminal_rewards(
+    rewards_per_sample: torch.Tensor,
+    *,
+    lengths: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+) -> torch.Tensor:
+    """Place each sample's scalar reward on its last response token in packed layout.
+
+    Args:
+        rewards_per_sample: Per-trajectory rewards ``[B]``.
+        lengths: Response token counts per sample ``[B]``.
+        cu_seqlens: Packed cumulative offsets ``[B + 1]`` (``TextSegment.cu_seqlens``).
+
+    Returns:
+        Packed per-token rewards ``[total_tokens]`` (zero except terminal positions).
+    """
+    if rewards_per_sample.ndim != 1:
+        raise ValueError(
+            f"scatter_terminal_rewards: rewards_per_sample must be 1D, got shape {tuple(rewards_per_sample.shape)}"
+        )
+    batch_size = int(lengths.shape[0])
+    if int(rewards_per_sample.shape[0]) != batch_size:
+        raise ValueError(
+            f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.shape[0])}) "
+            f"!= lengths batch ({batch_size})"
+        )
+    if int(cu_seqlens.shape[0]) != batch_size + 1:
+        raise ValueError(
+            f"scatter_terminal_rewards: cu_seqlens length ({int(cu_seqlens.shape[0])}) "
+            f"!= batch_size + 1 ({batch_size + 1})"
+        )
+    total = int(cu_seqlens[-1].item())
+    out = rewards_per_sample.new_zeros(total)
+    cu = [int(c) for c in cu_seqlens.tolist()]
+    for b in range(batch_size):
+        n = int(lengths[b].item())
+        if n <= 0:
+            continue
+        out[cu[b] + n - 1] = rewards_per_sample[b]
+    return out
+
+
 def _gae_1d(
     rewards: torch.Tensor,
     values: torch.Tensor,
diff --git a/unirl/types/segments/text.py b/unirl/types/segments/text.py
index 7a1570f23..fc34aba9c 100644
--- a/unirl/types/segments/text.py
+++ b/unirl/types/segments/text.py
@@ -43,6 +43,10 @@ class TextSegment(Segment):
     # Original engine emission retained when an algorithm replaces ``log_probs``
     # with a train-side replay anchor.
     rollout_log_probs: Optional[torch.Tensor] = packed_field(default=None)
+    # PPO / GAE path (optional): per-token critic and advantage plumbing.
+    values: Optional[torch.Tensor] = packed_field(default=None)
+    returns: Optional[torch.Tensor] = packed_field(default=None)
+    token_advantages: Optional[torch.Tensor] = packed_field(default=None)
 
     def as_condition_with(self, encoder: Callable[..., Any]) -> Condition:
         """Re-embed packed tokens via the supplied encoder into a TextEmbedCondition.

From e858e8ea4b06a29639c51c37981362a934a97236 Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sat, 25 Jul 2026 17:58:52 +0100
Subject: [PATCH 02/11] fix(models): restore lazy third-party imports in Qwen3

Keep transformers imports inside from_config / _packed_replay_supported
as in upstream; only UniRL imports (ValueHead, ARReplayOutput) stay top-level.
---
 unirl/models/qwen3/ar.py     | 9 +++------
 unirl/models/qwen3/bundle.py | 5 ++++-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index 6bdccf2d4..50bdf6d1b 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -42,11 +42,6 @@
 
 _SPARSE_PACKED_ATTN = ("flex_attention", "flash_attention_2", "flash_attention_3", "flash_attention_4")
 
-try:
-    from transformers.masking_utils import find_packed_sequence_indices
-except Exception:
-    find_packed_sequence_indices = None
-
 
 @functools.lru_cache(maxsize=None)
 def _warn_packed_disabled(attn_impl: str) -> None:
@@ -78,7 +73,9 @@ def _packed_replay_supported(attn_impl: Optional[str]) -> bool:
     if attn_impl not in _SPARSE_PACKED_ATTN:
         _warn_packed_disabled(str(attn_impl))
         return False
-    if find_packed_sequence_indices is None:
+    try:
+        from transformers.masking_utils import find_packed_sequence_indices  # noqa: F401
+    except Exception:
         return False
     return True
 
diff --git a/unirl/models/qwen3/bundle.py b/unirl/models/qwen3/bundle.py
index e3d3325f5..4b2182d39 100644
--- a/unirl/models/qwen3/bundle.py
+++ b/unirl/models/qwen3/bundle.py
@@ -25,7 +25,6 @@
 
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
 from unirl.models.types.bundle import Bundle
 from unirl.models.types.meta_init import build_meta_init_transformer
@@ -59,6 +58,8 @@ def __init__(
     @classmethod
     def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
         """Load the Qwen3 transformer + tokenizer from a HuggingFace-layout checkpoint."""
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
         path = config.pretrained_model_ckpt_path
         tokenizer_path = config.tokenizer_ckpt_path or path
 
@@ -78,6 +79,8 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
             # checkpoint, so to_empty later clobbers them -> garbage RoPE. It
             # captures them; meta_init_state is stashed on the BUNDLE below and
             # restored by load_trainable_weights after the sharded weight load.
+            from transformers import AutoConfig
+
             hf_config = AutoConfig.from_pretrained(path, trust_remote_code=bool(config.trust_remote_code))
             transformer, meta_init_state = build_meta_init_transformer(
                 lambda: AutoModelForCausalLM.from_config(hf_config, trust_remote_code=bool(config.trust_remote_code)),

From 06aa5e352aa1cc61a6d3e85e5b6d8c179d18ccbe Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sat, 25 Jul 2026 18:07:34 +0100
Subject: [PATCH 03/11] fix(models): require value head in replay forward path

Call _require_value_head_for_replay at the start of _replay_aware_forward
so return_values=True fails fast before the transformer forward, not only
from Qwen3ARStage.replay().
---
 unirl/models/qwen3/ar.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index 50bdf6d1b..0e2cb17b7 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -110,6 +110,8 @@ def _replay_aware_forward(
                 return f(self, **kw)
         raise RuntimeError("_replay_aware_forward: no class-level forward found in the MRO")
 
+    _require_value_head_for_replay(self, return_values)
+
     # cuDNN's fused SDPA backward (ScaledDotProductCudnnAttentionBackward0) returns
     # NaN grads on some bf16 sequences while the forward stays finite (confirmed via
     # torch.autograd.detect_anomaly): it floods every parameter grad and forces the
@@ -206,7 +208,7 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
 def _require_value_head_for_replay(model: Any, return_values: bool) -> None:
     if return_values and getattr(model, "value_head", None) is None:
         raise ValueError(
-            "Qwen3ARStage.replay: return_values=True requires a value head "
+            "Qwen3 replay: return_values=True requires a value head "
             "(set use_value_head=True in the pipeline config)"
         )
 

From e913991027c6251c37222e8b7747483e0d2dab27 Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sat, 25 Jul 2026 18:14:46 +0100
Subject: [PATCH 04/11] refactor(models): use ReplayResult for AR replay values

Extend ReplayResult with optional per-token values for PPO/GAE and
replace ARReplayOutput in Qwen3 replay with the shared type.
---
 unirl/models/qwen3/ar.py            | 38 ++++++++++++++---------------
 unirl/models/types/replay_result.py | 17 ++++++++-----
 2 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index 0e2cb17b7..4f41bc522 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -31,7 +31,7 @@
 from torch.utils.checkpoint import checkpoint
 
 from unirl.models.types.ar import ARSamplingParams, ARStage, ARStep, left_pad_prompt
-from unirl.models.types.ar_replay import ARReplayOutput
+from unirl.models.types.replay_result import ReplayResult
 from unirl.types.segments import TextSegment
 from unirl.utils.dtypes import parse_torch_dtype
 
@@ -163,7 +163,7 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
             empty = hidden.new_zeros((0,), dtype=torch.float32)
             if value_head is None:
                 return empty
-            return ARReplayOutput(log_probs=empty, values=empty)
+            return ReplayResult(log_probs=empty, values=empty)
         log_probs = torch.cat(flat_parts, dim=0)
         if value_head is None:
             return log_probs
@@ -171,7 +171,7 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
         for s in range(0, int(h_pred.size(0)), flat_chunk):
             value_parts.append(value_head(h_pred[s : s + flat_chunk]))
         values = torch.cat(value_parts, dim=0) if value_parts else log_probs.new_zeros(0)
-        return ARReplayOutput(log_probs=log_probs, values=values)
+        return ReplayResult(log_probs=log_probs, values=values)
     T_max = int(response_tokens.size(1))
     resp_hidden = hidden[:, prompt_len - 1 : prompt_len - 1 + T_max, :]
 
@@ -194,7 +194,7 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
         empty = resp_hidden.new_zeros((bsz, 0), dtype=torch.float32)
         if value_head is None:
             return empty
-        return ARReplayOutput(log_probs=empty, values=empty)
+        return ReplayResult(log_probs=empty, values=empty)
     log_probs = torch.cat(parts, dim=1)
     if value_head is None:
         return log_probs
@@ -202,7 +202,7 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
     for s in range(0, T_max, chunk):
         value_parts.append(value_head(resp_hidden[:, s : s + chunk, :]))
     values = torch.cat(value_parts, dim=1) if value_parts else log_probs.new_zeros((bsz, 0))
-    return ARReplayOutput(log_probs=log_probs, values=values)
+    return ReplayResult(log_probs=log_probs, values=values)
 
 
 def _require_value_head_for_replay(model: Any, return_values: bool) -> None:
@@ -214,15 +214,15 @@ def _require_value_head_for_replay(model: Any, return_values: bool) -> None:
 
 
 def _finalize_replay_output(
-    out: Union[torch.Tensor, ARReplayOutput],
+    out: Union[torch.Tensor, ReplayResult],
     *,
     segment: TextSegment,
     return_values: bool,
     logprob_dtype: torch.dtype,
     device: torch.device,
-) -> Union[torch.Tensor, ARReplayOutput]:
+) -> Union[torch.Tensor, ReplayResult]:
     """Cast log-probs and flatten packed values to match ``segment`` layout."""
-    if isinstance(out, ARReplayOutput):
+    if isinstance(out, ReplayResult):
         log_probs = out.log_probs.to(dtype=logprob_dtype)
         if not return_values:
             return log_probs
@@ -230,7 +230,7 @@ def _finalize_replay_output(
         if values is None:
             raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
         if log_probs.ndim == 1:
-            return ARReplayOutput(log_probs=log_probs, values=values.to(device=device))
+            return ReplayResult(log_probs=log_probs, values=values.to(device=device))
         if segment.cu_seqlens is None or segment.lengths is None:
             raise ValueError("Qwen3ARStage.replay: segment requires cu_seqlens to flatten values")
         lengths = [int(n) for n in segment.lengths.tolist()]
@@ -241,7 +241,7 @@ def _finalize_replay_output(
                 continue
             flat.append(values[b, :n])
         packed_values = torch.cat(flat, dim=0) if flat else values.new_zeros(0, device=device)
-        return ARReplayOutput(log_probs=log_probs, values=packed_values.to(device=device))
+        return ReplayResult(log_probs=log_probs, values=packed_values.to(device=device))
     if return_values:
         raise ValueError(
             "Qwen3ARStage.replay: return_values=True but critic returned no values "
@@ -492,14 +492,14 @@ def replay(
         segment: TextSegment,
         temperature: float = 1.0,
         return_values: bool = False,
-    ) -> Union[torch.Tensor, ARReplayOutput]:
+    ) -> Union[torch.Tensor, ReplayResult]:
         """Per-token log-prob replay over a stored rollout segment.
 
         Branch: prefer :meth:`packed_replay` (packed-varlen, zero padding, B > 1)
         and fall back to :meth:`padding_replay` (the dense ``[B, P_max + T_max]``
         padded path) when packing does not apply. Returns packed varlen
         ``[total_tokens]`` aligned with ``segment.log_probs`` unless
-        ``return_values=True``, in which case an :class:`ARReplayOutput` with
+        ``return_values=True``, in which case an :class:`ReplayResult` with
         packed ``values`` is returned alongside log-probs.
         """
         _require_value_head_for_replay(self.model.transformer, return_values)
@@ -521,7 +521,7 @@ def packed_replay(
         segment: TextSegment,
         temperature: float = 1.0,
         return_values: bool = False,
-    ) -> Optional[Union[torch.Tensor, ARReplayOutput]]:
+    ) -> Optional[Union[torch.Tensor, ReplayResult]]:
         """Packed-varlen replay (B > 1): zero padding anywhere.
 
         Concatenate every sample's REAL prompt tokens + its flat response tokens
@@ -615,7 +615,7 @@ def padding_replay(
         segment: TextSegment,
         temperature: float = 1.0,
         return_values: bool = False,
-    ) -> Union[torch.Tensor, ARReplayOutput]:
+    ) -> Union[torch.Tensor, ReplayResult]:
         """Dense ``[B, P_max + T_max]`` padded replay — the default / fallback path.
 
         One teacher-forced forward over padded ``prompt + response``; gather
@@ -723,15 +723,15 @@ def padding_replay(
             temperature=temperature,
             return_values=return_values,
             autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None),
-        )  # [B, T_max] FP32 or ARReplayOutput
+        )  # [B, T_max] FP32 or ReplayResult
 
         if T_max == 0:
             empty = torch.zeros(0, dtype=self.logprob_dtype, device=device)
             if return_values:
-                return ARReplayOutput(log_probs=empty, values=empty)
+                return ReplayResult(log_probs=empty, values=empty)
             return empty
 
-        if isinstance(out, ARReplayOutput):
+        if isinstance(out, ReplayResult):
             per_token = out.log_probs
             per_value = out.values
         else:
@@ -755,13 +755,13 @@ def padding_replay(
         if not flat_logp:
             empty = torch.zeros(0, dtype=self.logprob_dtype, device=device)
             if return_values:
-                return ARReplayOutput(log_probs=empty, values=empty)
+                return ReplayResult(log_probs=empty, values=empty)
             return empty
         log_probs = torch.cat(flat_logp, dim=0).to(dtype=self.logprob_dtype)
         if not return_values:
             return log_probs
         values = torch.cat(flat_val, dim=0)
-        return ARReplayOutput(log_probs=log_probs, values=values)
+        return ReplayResult(log_probs=log_probs, values=values)
 
     def _resolve_stop_ids(
         self,
diff --git a/unirl/models/types/replay_result.py b/unirl/models/types/replay_result.py
index 3bf7a7222..f7e262105 100644
--- a/unirl/models/types/replay_result.py
+++ b/unirl/models/types/replay_result.py
@@ -8,10 +8,9 @@
 
 Diffusion stages populate ``log_probs`` and ``prev_sample_means`` (the
 mean of the SDE Gaussian — μ_θ — used as the second moment in the KL
-penalty). AR stages currently return a plain ``Tensor`` (signature
-divergence with diffusion is intentional for now); when AR replay grows
-``logits``-based KL support, the ``logits`` field on this result will be
-the canonical home.
+penalty). AR stages return a plain ``Tensor`` for GRPO-style replay, or a
+:class:`ReplayResult` when optional critic ``values`` (or future ``logits``)
+are requested.
 """
 
 from __future__ import annotations
@@ -28,8 +27,9 @@ class ReplayResult:
     others are stage-specific and may be ``None``."""
 
     log_probs: torch.Tensor
-    """Aligned with ``segment.sde_logp`` (or its slice when ``step_indices``
-    subsets). Shape ``[B, S']`` for diffusion replay."""
+    """Aligned with ``segment.sde_logp`` / ``segment.log_probs`` (or a slice
+    when ``step_indices`` subsets). Shape ``[B, S']`` for diffusion replay;
+    packed ``[total_tokens]`` for AR varlen replay."""
 
     prev_sample_means: Optional[torch.Tensor] = None
     """The SDE transition's mean μ_θ at each replayed step. Shape
@@ -42,5 +42,10 @@ class ReplayResult:
     or entropy penalty support; not needed for Binary KL (which uses
     only per-token log-probs). Currently not populated."""
 
+    values: Optional[torch.Tensor] = None
+    """Per-token critic predictions ``V_t`` from replay. Shape ``[B, T]`` or
+    packed ``[total_tokens]`` for AR. Used by PPO / GAE training paths.
+    ``None`` when the stage does not attach a value head."""
+
 
 __all__ = ["ReplayResult"]

From e3aa5964e7c3f3e5608517f3dcc3d68850ec3d34 Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sun, 26 Jul 2026 11:40:55 +0100
Subject: [PATCH 05/11] style: fix ruff format and unused variable for CI

Apply ruff-format changes and remove unused cu in _finalize_replay_output.
---
 unirl/models/qwen3/ar.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index 4f41bc522..cb4028f1e 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -208,8 +208,7 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
 def _require_value_head_for_replay(model: Any, return_values: bool) -> None:
     if return_values and getattr(model, "value_head", None) is None:
         raise ValueError(
-            "Qwen3 replay: return_values=True requires a value head "
-            "(set use_value_head=True in the pipeline config)"
+            "Qwen3 replay: return_values=True requires a value head (set use_value_head=True in the pipeline config)"
         )
 
 
@@ -234,7 +233,6 @@ def _finalize_replay_output(
         if segment.cu_seqlens is None or segment.lengths is None:
             raise ValueError("Qwen3ARStage.replay: segment requires cu_seqlens to flatten values")
         lengths = [int(n) for n in segment.lengths.tolist()]
-        cu = [int(c) for c in segment.cu_seqlens.tolist()]
         flat: List[torch.Tensor] = []
         for b, n in enumerate(lengths):
             if n <= 0:
@@ -510,9 +508,7 @@ def replay(
             )
             if packed is not None:
                 return packed
-        return self.padding_replay(
-            conditions, segment=segment, temperature=temperature, return_values=return_values
-        )
+        return self.padding_replay(conditions, segment=segment, temperature=temperature, return_values=return_values)
 
     def packed_replay(
         self,

From bd91950c1d34bbef3871352a27d43e90e4759429 Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sun, 26 Jul 2026 11:47:43 +0100
Subject: [PATCH 06/11] refactor(types): index terminal rewards via cu_seqlens
 only

Derive last-token positions from packed offsets (end - 1) instead of
passing redundant lengths, with a guard for empty samples.
---
 unirl/types/advantages.py | 18 +++++-------------
 1 file changed, 5 insertions(+), 13 deletions(-)

diff --git a/unirl/types/advantages.py b/unirl/types/advantages.py
index 1140d23c9..c6c367e7e 100644
--- a/unirl/types/advantages.py
+++ b/unirl/types/advantages.py
@@ -74,14 +74,12 @@ def compute_gae_advantages(
 def scatter_terminal_rewards(
     rewards_per_sample: torch.Tensor,
     *,
-    lengths: torch.Tensor,
     cu_seqlens: torch.Tensor,
 ) -> torch.Tensor:
     """Place each sample's scalar reward on its last response token in packed layout.
 
     Args:
         rewards_per_sample: Per-trajectory rewards ``[B]``.
-        lengths: Response token counts per sample ``[B]``.
         cu_seqlens: Packed cumulative offsets ``[B + 1]`` (``TextSegment.cu_seqlens``).
 
     Returns:
@@ -91,25 +89,19 @@ def scatter_terminal_rewards(
         raise ValueError(
             f"scatter_terminal_rewards: rewards_per_sample must be 1D, got shape {tuple(rewards_per_sample.shape)}"
         )
-    batch_size = int(lengths.shape[0])
+    batch_size = int(cu_seqlens.shape[0]) - 1
     if int(rewards_per_sample.shape[0]) != batch_size:
         raise ValueError(
-            f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.shape[0])}) "
-            f"!= lengths batch ({batch_size})"
-        )
-    if int(cu_seqlens.shape[0]) != batch_size + 1:
-        raise ValueError(
-            f"scatter_terminal_rewards: cu_seqlens length ({int(cu_seqlens.shape[0])}) "
-            f"!= batch_size + 1 ({batch_size + 1})"
+            f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.shape[0])}) != batch_size ({batch_size})"
         )
     total = int(cu_seqlens[-1].item())
     out = rewards_per_sample.new_zeros(total)
     cu = [int(c) for c in cu_seqlens.tolist()]
     for b in range(batch_size):
-        n = int(lengths[b].item())
-        if n <= 0:
+        start, end = cu[b], cu[b + 1]
+        if end <= start:
             continue
-        out[cu[b] + n - 1] = rewards_per_sample[b]
+        out[end - 1] = rewards_per_sample[b]
     return out
 
 

From 74bef1c3bc68931623469cb1e49c70e212b8b4a3 Mon Sep 17 00:00:00 2001
From: leviking98z-rgb <leviking98z@gmail.com>
Date: Sun, 26 Jul 2026 18:54:18 +0800
Subject: [PATCH 07/11] fix(models): make value head FSDP dtype-safe

---
 unirl/models/types/value_head.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/unirl/models/types/value_head.py b/unirl/models/types/value_head.py
index 5a2d1969e..b1b11daad 100644
--- a/unirl/models/types/value_head.py
+++ b/unirl/models/types/value_head.py
@@ -4,6 +4,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 
 class ValueHead(nn.Module):
@@ -18,7 +19,13 @@ def __init__(self, hidden_size: int) -> None:
 
     def forward(self, hidden: torch.Tensor) -> torch.Tensor:
         """Map ``[..., H]`` hidden states to ``[...,]`` scalar values."""
-        return self.proj(hidden.float()).squeeze(-1)
+        # FSDP mixed precision may expose gathered parameters in its compute
+        # dtype even though their sharded masters were initialized in FP32.
+        # Cast both operands: casting only ``hidden`` fails when the gathered
+        # projection is BF16.
+        weight = self.proj.weight.float()
+        bias = self.proj.bias.float() if self.proj.bias is not None else None
+        return F.linear(hidden.float(), weight, bias).squeeze(-1)
 
 
 __all__ = ["ValueHead"]

From 09685565bdb3936ef8efd8a7540dd97189fbbce7 Mon Sep 17 00:00:00 2001
From: leviking98z-rgb <leviking98z@gmail.com>
Date: Sun, 2 Aug 2026 02:31:15 +0800
Subject: [PATCH 08/11] fix(models): preserve value-head and GAE correctness on
 current main

---
 unirl/models/qwen3/ar.py            | 127 ++++++++++++----------------
 unirl/models/qwen3/bundle.py        |  26 +++++-
 unirl/models/qwen3/config.py        |   2 +-
 unirl/models/types/replay_result.py |  15 ++--
 unirl/models/types/value_head.py    |  24 +++---
 unirl/types/advantages.py           |  42 ++++-----
 unirl/types/sample.py               |  87 ++++++++++++++++++-
 unirl/types/segments/text.py        |   2 +-
 8 files changed, 206 insertions(+), 119 deletions(-)

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index cb4028f1e..793e06827 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -167,9 +167,7 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
         log_probs = torch.cat(flat_parts, dim=0)
         if value_head is None:
             return log_probs
-        value_parts: List[torch.Tensor] = []
-        for s in range(0, int(h_pred.size(0)), flat_chunk):
-            value_parts.append(value_head(h_pred[s : s + flat_chunk]))
+        value_parts = [value_head(h_pred[s : s + flat_chunk]) for s in range(0, int(h_pred.size(0)), flat_chunk)]
         values = torch.cat(value_parts, dim=0) if value_parts else log_probs.new_zeros(0)
         return ReplayResult(log_probs=log_probs, values=values)
     T_max = int(response_tokens.size(1))
@@ -198,9 +196,7 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
     log_probs = torch.cat(parts, dim=1)
     if value_head is None:
         return log_probs
-    value_parts = []
-    for s in range(0, T_max, chunk):
-        value_parts.append(value_head(resp_hidden[:, s : s + chunk, :]))
+    value_parts = [value_head(resp_hidden[:, s : s + chunk, :]) for s in range(0, T_max, chunk)]
     values = torch.cat(value_parts, dim=1) if value_parts else log_probs.new_zeros((bsz, 0))
     return ReplayResult(log_probs=log_probs, values=values)
 
@@ -218,34 +214,26 @@ def _finalize_replay_output(
     segment: TextSegment,
     return_values: bool,
     logprob_dtype: torch.dtype,
-    device: torch.device,
 ) -> Union[torch.Tensor, ReplayResult]:
-    """Cast log-probs and flatten packed values to match ``segment`` layout."""
-    if isinstance(out, ReplayResult):
-        log_probs = out.log_probs.to(dtype=logprob_dtype)
-        if not return_values:
-            return log_probs
-        values = out.values
-        if values is None:
+    """Cast log-probs and flatten padded critic values to segment order."""
+    if not isinstance(out, ReplayResult):
+        if return_values:
             raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
-        if log_probs.ndim == 1:
-            return ReplayResult(log_probs=log_probs, values=values.to(device=device))
-        if segment.cu_seqlens is None or segment.lengths is None:
-            raise ValueError("Qwen3ARStage.replay: segment requires cu_seqlens to flatten values")
-        lengths = [int(n) for n in segment.lengths.tolist()]
-        flat: List[torch.Tensor] = []
-        for b, n in enumerate(lengths):
-            if n <= 0:
-                continue
-            flat.append(values[b, :n])
-        packed_values = torch.cat(flat, dim=0) if flat else values.new_zeros(0, device=device)
-        return ReplayResult(log_probs=log_probs, values=packed_values.to(device=device))
-    if return_values:
-        raise ValueError(
-            "Qwen3ARStage.replay: return_values=True but critic returned no values "
-            "(set use_value_head=True in the pipeline config)"
-        )
-    return out.to(dtype=logprob_dtype)
+        return out.to(dtype=logprob_dtype)
+
+    log_probs = out.log_probs.to(dtype=logprob_dtype)
+    if not return_values:
+        return log_probs
+    if out.values is None:
+        raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
+    if log_probs.ndim == 1:
+        return ReplayResult(log_probs=log_probs, values=out.values.float())
+    if segment.lengths is None:
+        raise ValueError("Qwen3ARStage.replay: segment requires lengths to flatten critic values")
+
+    flat_values = [out.values[b, : int(length)] for b, length in enumerate(segment.lengths.tolist()) if int(length) > 0]
+    values = torch.cat(flat_values, dim=0) if flat_values else out.values.new_zeros(0)
+    return ReplayResult(log_probs=log_probs, values=values.float())
 
 
 # Attention backends with a sparse packed kernel (skip cross-sequence blocks):
@@ -495,20 +483,26 @@ def replay(
 
         Branch: prefer :meth:`packed_replay` (packed-varlen, zero padding, B > 1)
         and fall back to :meth:`padding_replay` (the dense ``[B, P_max + T_max]``
-        padded path) when packing does not apply. Returns packed varlen
-        ``[total_tokens]`` aligned with ``segment.log_probs`` unless
-        ``return_values=True``, in which case an :class:`ReplayResult` with
-        packed ``values`` is returned alongside log-probs.
+        padded path) when packing does not apply. Returns packed varlen log-probs,
+        or a :class:`ReplayResult` with aligned critic values when requested.
         """
         _require_value_head_for_replay(self.model.transformer, return_values)
         attn_impl = getattr(getattr(self.model.transformer, "config", None), "_attn_implementation", None)
         if _packed_replay_supported(attn_impl):
             packed = self.packed_replay(
-                conditions, segment=segment, temperature=temperature, return_values=return_values
+                conditions,
+                segment=segment,
+                temperature=temperature,
+                return_values=return_values,
             )
             if packed is not None:
                 return packed
-        return self.padding_replay(conditions, segment=segment, temperature=temperature, return_values=return_values)
+        return self.padding_replay(
+            conditions,
+            segment=segment,
+            temperature=temperature,
+            return_values=return_values,
+        )
 
     def packed_replay(
         self,
@@ -601,7 +595,6 @@ def packed_replay(
             segment=segment,
             return_values=return_values,
             logprob_dtype=self.logprob_dtype,
-            device=device,
         )
 
     def padding_replay(
@@ -719,44 +712,34 @@ def padding_replay(
             temperature=temperature,
             return_values=return_values,
             autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None),
-        )  # [B, T_max] FP32 or ReplayResult
+        )  # [B, T_max] FP32 tensor or ReplayResult
 
-        if T_max == 0:
-            empty = torch.zeros(0, dtype=self.logprob_dtype, device=device)
-            if return_values:
-                return ReplayResult(log_probs=empty, values=empty)
-            return empty
-
-        if isinstance(out, ReplayResult):
-            per_token = out.log_probs
-            per_value = out.values
+        finalized = _finalize_replay_output(
+            out,
+            segment=segment,
+            return_values=return_values,
+            logprob_dtype=self.logprob_dtype,
+        )
+        if isinstance(finalized, ReplayResult):
+            per_token = finalized.log_probs
+            values = finalized.values
         else:
-            per_token = out
-            per_value = None
-        if return_values and per_value is None:
-            raise ValueError(
-                "Qwen3ARStage.replay: return_values=True but critic returned no values "
-                "(set use_value_head=True in the pipeline config)"
-            )
+            per_token = finalized
+            values = None
 
-        flat_logp: List[torch.Tensor] = []
-        flat_val: List[torch.Tensor] = []
-        for b in range(batch_size):
-            n = lengths[b]
-            if n == 0:
-                continue
-            flat_logp.append(per_token[b, :n])
-            if per_value is not None:
-                flat_val.append(per_value[b, :n])
-        if not flat_logp:
-            empty = torch.zeros(0, dtype=self.logprob_dtype, device=device)
-            if return_values:
-                return ReplayResult(log_probs=empty, values=empty)
-            return empty
-        log_probs = torch.cat(flat_logp, dim=0).to(dtype=self.logprob_dtype)
+        flat_log_probs: List[torch.Tensor] = []
+        for b, n in enumerate(lengths):
+            if n > 0:
+                flat_log_probs.append(per_token[b, :n])
+        log_probs = (
+            torch.cat(flat_log_probs, dim=0)
+            if flat_log_probs
+            else torch.zeros(0, dtype=self.logprob_dtype, device=device)
+        )
         if not return_values:
             return log_probs
-        values = torch.cat(flat_val, dim=0)
+        if values is None:
+            raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
         return ReplayResult(log_probs=log_probs, values=values)
 
     def _resolve_stop_ids(
diff --git a/unirl/models/qwen3/bundle.py b/unirl/models/qwen3/bundle.py
index 4b2182d39..d44201f0d 100644
--- a/unirl/models/qwen3/bundle.py
+++ b/unirl/models/qwen3/bundle.py
@@ -36,6 +36,24 @@
 logger = logging.getLogger(__name__)
 
 
+def _stamp_value_head_reset(transformer: nn.Module) -> None:
+    """Zero checkpoint-absent value-head params after meta materialization."""
+    from unirl.train.deferred import _stamp
+
+    def _reset(model: nn.Module) -> None:
+        reset: list[str] = []
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if name.startswith("value_head."):
+                    param.zero_()
+                    reset.append(name)
+        if not reset:
+            raise RuntimeError("Qwen3 meta-init: value_head parameters disappeared before post-load reset")
+        logger.info("Qwen3 meta-init: zero-initialized checkpoint-absent value head: %s", reset)
+
+    _stamp(transformer, _reset)
+
+
 class Qwen3Bundle(Bundle):
     """Qwen3 bundle: causal-LM transformer + matching tokenizer."""
 
@@ -118,7 +136,13 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
 
         if config.use_value_head:
             hidden_size = int(getattr(transformer.config, "hidden_size"))
-            transformer.value_head = ValueHead(hidden_size).to(device)
+            transformer_device = next(transformer.parameters()).device
+            transformer.value_head = ValueHead(hidden_size, device=transformer_device)
+            if config.meta_init_transformer:
+                # The base checkpoint has no value_head.* tensors. ``to_empty``
+                # materializes the meta head as uninitialized storage, so reset
+                # it only after the sharded checkpoint load has completed.
+                _stamp_value_head_reset(transformer)
 
         bundle = cls(
             transformer=transformer,
diff --git a/unirl/models/qwen3/config.py b/unirl/models/qwen3/config.py
index 42a6c51cb..2fa5d3504 100644
--- a/unirl/models/qwen3/config.py
+++ b/unirl/models/qwen3/config.py
@@ -69,7 +69,7 @@ class Qwen3PipelineConfig:
     use_lora: bool = False
     lora_target_modules: Optional[List[str]] = None
 
-    # Attach a scalar value head on the transformer for PPO / GAE training.
+    # Attach a scalar critic to the causal LM for PPO / GAE training.
     use_value_head: bool = False
 
     system_instruction: Optional[str] = None
diff --git a/unirl/models/types/replay_result.py b/unirl/models/types/replay_result.py
index f7e262105..381d513c5 100644
--- a/unirl/models/types/replay_result.py
+++ b/unirl/models/types/replay_result.py
@@ -8,9 +8,8 @@
 
 Diffusion stages populate ``log_probs`` and ``prev_sample_means`` (the
 mean of the SDE Gaussian — μ_θ — used as the second moment in the KL
-penalty). AR stages return a plain ``Tensor`` for GRPO-style replay, or a
-:class:`ReplayResult` when optional critic ``values`` (or future ``logits``)
-are requested.
+penalty). AR stages return a plain ``Tensor`` for policy-only replay, or
+a :class:`ReplayResult` when optional critic ``values`` are requested.
 """
 
 from __future__ import annotations
@@ -27,9 +26,8 @@ class ReplayResult:
     others are stage-specific and may be ``None``."""
 
     log_probs: torch.Tensor
-    """Aligned with ``segment.sde_logp`` / ``segment.log_probs`` (or a slice
-    when ``step_indices`` subsets). Shape ``[B, S']`` for diffusion replay;
-    packed ``[total_tokens]`` for AR varlen replay."""
+    """Aligned with ``segment.sde_logp`` (or its slice when ``step_indices``
+    subsets). Shape ``[B, S']`` for diffusion replay."""
 
     prev_sample_means: Optional[torch.Tensor] = None
     """The SDE transition's mean μ_θ at each replayed step. Shape
@@ -43,9 +41,8 @@ class ReplayResult:
     only per-token log-probs). Currently not populated."""
 
     values: Optional[torch.Tensor] = None
-    """Per-token critic predictions ``V_t`` from replay. Shape ``[B, T]`` or
-    packed ``[total_tokens]`` for AR. Used by PPO / GAE training paths.
-    ``None`` when the stage does not attach a value head."""
+    """Per-token critic predictions ``V_t``. Packed ``[total_tokens]`` for AR.
+    ``None`` when replay did not request a value head."""
 
 
 __all__ = ["ReplayResult"]
diff --git a/unirl/models/types/value_head.py b/unirl/models/types/value_head.py
index b1b11daad..1193a8017 100644
--- a/unirl/models/types/value_head.py
+++ b/unirl/models/types/value_head.py
@@ -2,27 +2,31 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 
 class ValueHead(nn.Module):
-    """Linear critic ``V(h)`` on last hidden states.
-
-    Kept in FP32 for stable value loss math (mirrors replay log-prob FP32 policy).
-    """
+    """Linear critic ``V(h)`` on the hidden state that predicts each action."""
 
-    def __init__(self, hidden_size: int) -> None:
+    def __init__(self, hidden_size: int, *, device: Any = None) -> None:
         super().__init__()
-        self.proj = nn.Linear(hidden_size, 1, bias=True, dtype=torch.float32)
+        self.proj = nn.Linear(hidden_size, 1, bias=True, dtype=torch.float32, device=device)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """Start from ``V(s)=0``; deterministic and safe after sharded meta-init."""
+        nn.init.zeros_(self.proj.weight)
+        if self.proj.bias is not None:
+            nn.init.zeros_(self.proj.bias)
 
     def forward(self, hidden: torch.Tensor) -> torch.Tensor:
-        """Map ``[..., H]`` hidden states to ``[...,]`` scalar values."""
+        """Map ``[..., H]`` hidden states to FP32 scalar values ``[...]``."""
         # FSDP mixed precision may expose gathered parameters in its compute
-        # dtype even though their sharded masters were initialized in FP32.
-        # Cast both operands: casting only ``hidden`` fails when the gathered
-        # projection is BF16.
+        # dtype even when the sharded masters are FP32.
         weight = self.proj.weight.float()
         bias = self.proj.bias.float() if self.proj.bias is not None else None
         return F.linear(hidden.float(), weight, bias).squeeze(-1)
diff --git a/unirl/types/advantages.py b/unirl/types/advantages.py
index c6c367e7e..d4c3d4b35 100644
--- a/unirl/types/advantages.py
+++ b/unirl/types/advantages.py
@@ -76,33 +76,27 @@ def scatter_terminal_rewards(
     *,
     cu_seqlens: torch.Tensor,
 ) -> torch.Tensor:
-    """Place each sample's scalar reward on its last response token in packed layout.
-
-    Args:
-        rewards_per_sample: Per-trajectory rewards ``[B]``.
-        cu_seqlens: Packed cumulative offsets ``[B + 1]`` (``TextSegment.cu_seqlens``).
-
-    Returns:
-        Packed per-token rewards ``[total_tokens]`` (zero except terminal positions).
-    """
+    """Scatter each trajectory reward onto its final packed response token."""
     if rewards_per_sample.ndim != 1:
+        raise ValueError(f"scatter_terminal_rewards: expected 1D rewards, got shape {tuple(rewards_per_sample.shape)}")
+    if cu_seqlens.ndim != 1 or cu_seqlens.numel() == 0:
+        raise ValueError("scatter_terminal_rewards: cu_seqlens must be a non-empty 1D tensor")
+    batch_size = int(cu_seqlens.numel()) - 1
+    if int(rewards_per_sample.numel()) != batch_size:
         raise ValueError(
-            f"scatter_terminal_rewards: rewards_per_sample must be 1D, got shape {tuple(rewards_per_sample.shape)}"
+            f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.numel())}) "
+            f"!= packed batch ({batch_size})"
         )
-    batch_size = int(cu_seqlens.shape[0]) - 1
-    if int(rewards_per_sample.shape[0]) != batch_size:
-        raise ValueError(
-            f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.shape[0])}) != batch_size ({batch_size})"
-        )
-    total = int(cu_seqlens[-1].item())
-    out = rewards_per_sample.new_zeros(total)
-    cu = [int(c) for c in cu_seqlens.tolist()]
-    for b in range(batch_size):
-        start, end = cu[b], cu[b + 1]
-        if end <= start:
-            continue
-        out[end - 1] = rewards_per_sample[b]
-    return out
+
+    cu = [int(offset) for offset in cu_seqlens.tolist()]
+    if cu[0] != 0 or any(end < start for start, end in zip(cu, cu[1:])):
+        raise ValueError(f"scatter_terminal_rewards: invalid cumulative offsets {cu}")
+
+    token_rewards = rewards_per_sample.new_zeros(cu[-1])
+    for reward, start, end in zip(rewards_per_sample, cu, cu[1:]):
+        if end > start:
+            token_rewards[end - 1] = reward
+    return token_rewards
 
 
 def _gae_1d(
diff --git a/unirl/types/sample.py b/unirl/types/sample.py
index fede101f6..797207d22 100644
--- a/unirl/types/sample.py
+++ b/unirl/types/sample.py
@@ -31,12 +31,14 @@
     shared_field,
 )
 from unirl.distributed.tensor.ref import hydrate
+from unirl.types.advantages import compute_gae_advantages as _compute_gae
+from unirl.types.advantages import scatter_terminal_rewards
 from unirl.types.conditions import Condition
 from unirl.types.media_preview import MediaPreview
 from unirl.types.primitives import Audios, Images, Texts, Videos, primitive_modality_key
 from unirl.types.sample_id import ancestor_id, child_id, parent_id
 from unirl.types.sampling import BaseSamplingParams
-from unirl.types.segments import Segment
+from unirl.types.segments import Segment, TextSegment
 from unirl.utils.shard_balance import lpt_shard_permutation, shard_token_spread
 
 logger = logging.getLogger(__name__)
@@ -418,6 +420,89 @@ def compute_advantages(
             adv = reshaped - mean
         return _part_with_field(self, "advantages", adv.flatten())
 
+    def compute_gae_advantages(
+        self,
+        *,
+        gamma: float = 1.0,
+        gae_lambda: float = 0.95,
+    ) -> "Part":
+        """Attach packed token advantages and returns for an AR PPO update.
+
+        ``loss_mask`` controls which tokens later contribute to the policy and
+        value losses. It deliberately does not control the GAE recursion:
+        masked actions are still states in the same trajectory, so terminal
+        reward must propagate across them.
+        """
+        if self.rewards is None:
+            raise ValueError("Part.compute_gae_advantages: part has no rewards")
+        if not isinstance(self.segment, TextSegment):
+            raise ValueError("Part.compute_gae_advantages: requires a TextSegment")
+        segment = self.segment
+        if segment.values is None:
+            raise ValueError("Part.compute_gae_advantages: segment.values is None")
+        if segment.cu_seqlens is None or segment.lengths is None:
+            raise ValueError(
+                "Part.compute_gae_advantages: segment requires framework-managed "
+                "cu_seqlens (construct via TextSegment.pack)"
+            )
+
+        values = hydrate(segment.values).to(torch.float32)
+        cu_seqlens = segment.cu_seqlens.to(device=values.device)
+        total_tokens = int(cu_seqlens[-1].item())
+        if values.ndim != 1 or int(values.numel()) != total_tokens:
+            raise ValueError(
+                "Part.compute_gae_advantages: values must be a packed 1D tensor "
+                f"with {total_tokens} elements, got shape {tuple(values.shape)}"
+            )
+
+        rewards = hydrate(self.rewards).to(device=values.device, dtype=torch.float32)
+        token_rewards = scatter_terminal_rewards(rewards, cu_seqlens=cu_seqlens)
+        token_advantages = values.new_zeros(values.shape)
+        token_returns = values.new_zeros(values.shape)
+        cu = [int(offset) for offset in cu_seqlens.tolist()]
+        for start, end in zip(cu, cu[1:]):
+            if end <= start:
+                continue
+            advantages, returns = _compute_gae(
+                token_rewards[start:end],
+                values[start:end],
+                gamma=gamma,
+                gae_lambda=gae_lambda,
+            )
+            token_advantages[start:end] = advantages
+            token_returns[start:end] = returns
+
+        segment_fields = {f.name: getattr(segment, f.name) for f in dc_fields(segment)}
+        segment_fields["token_advantages"] = token_advantages
+        segment_fields["returns"] = token_returns
+        updated_segment = segment._rebuild(segment_fields)
+
+        # Keep the existing per-sample advantage metric meaningful. The loss
+        # mask applies here only as a reduction mask, never as a done signal.
+        loss_mask = None
+        if segment.loss_mask is not None:
+            loss_mask = hydrate(segment.loss_mask).to(device=values.device, dtype=torch.bool)
+            if loss_mask.shape != values.shape:
+                raise ValueError(
+                    "Part.compute_gae_advantages: loss_mask shape "
+                    f"{tuple(loss_mask.shape)} != values shape {tuple(values.shape)}"
+                )
+        sample_advantages: List[torch.Tensor] = []
+        for start, end in zip(cu, cu[1:]):
+            if end <= start:
+                sample_advantages.append(values.new_zeros(()))
+                continue
+            selected = token_advantages[start:end]
+            if loss_mask is not None:
+                selected = selected[loss_mask[start:end]]
+            sample_advantages.append(selected.mean() if selected.numel() else values.new_zeros(()))
+        mean_advantages = (
+            torch.stack(sample_advantages) if sample_advantages else values.new_zeros((0,), dtype=torch.float32)
+        )
+
+        updated = _part_with_field(self, "segment", updated_segment)
+        return _part_with_field(updated, "advantages", mean_advantages)
+
 
 def _part_with_field(part: Part, field_name: str, value: Any) -> Part:
     """Copy of ``part`` with one field replaced."""
diff --git a/unirl/types/segments/text.py b/unirl/types/segments/text.py
index fc34aba9c..8850f902d 100644
--- a/unirl/types/segments/text.py
+++ b/unirl/types/segments/text.py
@@ -43,7 +43,7 @@ class TextSegment(Segment):
     # Original engine emission retained when an algorithm replaces ``log_probs``
     # with a train-side replay anchor.
     rollout_log_probs: Optional[torch.Tensor] = packed_field(default=None)
-    # PPO / GAE path (optional): per-token critic and advantage plumbing.
+    # Optional PPO critic state, aligned one-to-one with packed response tokens.
     values: Optional[torch.Tensor] = packed_field(default=None)
     returns: Optional[torch.Tensor] = packed_field(default=None)
     token_advantages: Optional[torch.Tensor] = packed_field(default=None)

From f40dcd0b5566cfd80758a72eed37670ef73ef9cd Mon Sep 17 00:00:00 2001
From: yhl48 <limjohnyh@gmail.com>
Date: Sun, 26 Jul 2026 17:41:59 +0100
Subject: [PATCH 09/11] feat(algorithms): add AR PPO with GAE and value loss

Wire PPO.prepare_rollout_track for worker-side GAE, policy+value losses,
advantage_mode=gae in AR trainers, and a Qwen3 DAPO recipe (issue #86, 3/3).
---
 .../ar/qwen3_ppo_4b_base_dapo_sglang.yaml     | 147 +++++++++++
 unirl/algorithms/__init__.py                  |   3 +
 unirl/algorithms/base.py                      |  21 ++
 unirl/algorithms/ppo.py                       | 238 ++++++++++++++++++
 unirl/train/stack/base.py                     |   3 +
 unirl/train_ar.py                             |   1 +
 unirl/trainer/ar.py                           |   9 +-
 unirl/trainer/async_ar.py                     |   6 +-
 8 files changed, 426 insertions(+), 2 deletions(-)
 create mode 100644 examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml
 create mode 100644 unirl/algorithms/ppo.py

diff --git a/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml b/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml
new file mode 100644
index 000000000..a2990960e
--- /dev/null
+++ b/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml
@@ -0,0 +1,147 @@
+# @package _global_
+# PPO + GAE Qwen3-4B-Base on DAPO-Math — ARTrainer + SGLang rollout.
+#
+# Sibling of qwen3_grpo_4b_base_dapo_sglang.yaml with a value critic and GAE
+# advantages instead of group-relative GRPO normalization. Requires
+# bundle.config.use_value_head=true and advantage_mode=gae.
+#
+# Data prep (one-time):
+#   python -m unirl.utils.prepare_dapo_math --out-dir data/dapo_math
+#
+# Run (32 GPU example):
+#   DATA_PATH=data/dapo_math/train.jsonl EVAL_DATA_PATH=data/dapo_math/aime_eval.jsonl \
+#   python -m unirl.train_ar --config-name=ar/qwen3_ppo_4b_base_dapo_sglang num_devices=32
+
+num_devices: 32
+batch_size: 64
+num_rollouts: 800
+weight_sync_interval: 1
+eval_interval: 10
+
+# GAE path: advantages computed on workers in PPO.prepare_rollout_track.
+advantage_mode: gae
+adv_normalization_scope: group
+normalize_adv_by_std: false
+
+logging:
+  report_to_wandb: true
+  project_name: unirl-ppo
+  run_name: ppo_qwen3-4b-base_dapo_sglang
+  entity: ${oc.env:WANDB_ENTITY,null}
+  tags: [ppo, gae, qwen3, 4b-base, dapo, sglang, v2]
+
+bundle:
+  _target_: unirl.models.qwen3.bundle.Qwen3Bundle.from_config
+  config:
+    _target_: unirl.models.qwen3.config.Qwen3PipelineConfig
+    pretrained_model_ckpt_path: ${oc.env:QWEN3_PATH,Qwen/Qwen3-4B-Base}
+    model_precision: fp32
+    use_gradient_checkpointing: true
+    use_value_head: true
+    attn_implementation: flex_attention
+
+pipeline:
+  _target_: unirl.models.qwen3.pipeline.Qwen3Pipeline.from_bundle
+  enable_thinking: true
+  autocast_precision: bf16
+  logprob_precision: fp32
+
+backend:
+  _target_: unirl.train.backend.fsdp.FSDPBackend
+  block_class_names: ["Qwen3DecoderLayer"]
+  trainable_attr: transformer
+  fsdp_cfg:
+    _target_: unirl.train.configs.FSDPConfig
+    param_dtype: bf16
+    cpu_offload: false
+    mixed_precision: true
+    fsdp_mode: full
+    reshard_after_forward: true
+    activation_checkpointing: true
+    use_torch_compile: false
+    forward_prefetch: false
+    defer_grad_sync: false
+  optimizer_cfg:
+    _target_: unirl.train.backend.base.OptimizerConfig
+    learning_rate: 1.0e-6
+    adam_beta1: 0.9
+    adam_beta2: 0.999
+    adam_epsilon: 1.0e-8
+    weight_decay: 0.0
+
+rollout:
+  _target_: unirl.rollout.engine.sglang.SGLangRolloutEngine
+  model_path: ${bundle.config.pretrained_model_ckpt_path}
+  concurrency: 16
+  samples_pre_expanded: true
+  chat_template_kwargs:
+    enable_thinking: true
+  engine_kwargs:
+    mem_fraction_static: 0.3
+    skip_server_warmup: true
+    attention_backend: triton
+    disable_cuda_graph: false
+    cuda_graph_max_bs: 16
+    enable_lora: false
+
+reward:
+  _target_: unirl.reward.service.RewardService
+  truncated_reward: keep
+  backend:
+    _target_: unirl.reward.local.mathverify.MathVerifyRewardScorer
+    base_device: cpu
+    config:
+      _target_: unirl.reward.local.mathverify.MathVerifySpec
+
+algorithm:
+  _target_: unirl.algorithms.ppo.PPO
+  stage_attr: ar
+  clip_range: 0.2
+  clip_range_high: null
+  clip_schedule: constant
+  cliprange_value: 0.2
+  vf_coef: 0.5
+  gae_gamma: 1.0
+  gae_lambda: 0.95
+  loss_agg_mode: seq-mean-token-mean
+  horizon: 8192
+  sampling_temperature: 1.0
+  conditions_cls:
+    _target_: hydra.utils.get_class
+    path: unirl.models.qwen3.conditions.Qwen3ARConditions
+
+sync:
+  # Drop train-only value_head weights — SGLang has no critic params to receive them.
+  _target_: unirl.distributed.weight_sync.full.tensor.TensorWeightSync
+  lora_merged: false
+  bucket_size_mb: 64
+  flush_cache: true
+  name_remap:
+    value_head.*: null
+
+stack:
+  _target_: unirl.train.stack.TrainStack
+  micro_batch_size: 1
+  max_grad_norm: 1.0
+  num_updates_per_batch: 4
+  micro_planner:
+    _target_: unirl.train.stack.TokenBudgetPlanner
+    token_budget: 10240
+
+data_source:
+  _target_: unirl.data.data_source.MultimodalRLDataSource
+  args:
+    run:
+      data_path: ${oc.env:DATA_PATH}
+      eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH}}
+      seed: 42
+    algorithm:
+      prompts_per_rollout: 64
+
+sampling:
+  _target_: unirl.types.sampling.ARSamplingParams
+  samples_per_prompt: 8
+  temperature: 1.0
+  top_p: 1.0
+  top_k: 0
+  max_new_tokens: 8192
diff --git a/unirl/algorithms/__init__.py b/unirl/algorithms/__init__.py
index 8604d0b50..d1205ea66 100644
--- a/unirl/algorithms/__init__.py
+++ b/unirl/algorithms/__init__.py
@@ -15,6 +15,7 @@
 from .flowgrpo import FlowGRPO, FlowGRPOConfig
 from .grpo import GRPO, GRPOConfig
 from .gspo import GSPO, GSPOConfig
+from .ppo import PPO, PPOConfig
 from .sft import SFT, FlowMatchSFT
 
 __all__ = [
@@ -22,6 +23,8 @@
     "FlowMatchSFT",
     "GRPO",
     "GRPOConfig",
+    "PPO",
+    "PPOConfig",
     "GSPO",
     "GSPOConfig",
     "CPPO",
diff --git a/unirl/algorithms/base.py b/unirl/algorithms/base.py
index e329743d1..cc0a915e4 100644
--- a/unirl/algorithms/base.py
+++ b/unirl/algorithms/base.py
@@ -193,6 +193,27 @@ def _grpo_clip_loss(
     return loss_per_elem, metrics
 
 
+def _ppo_clipped_value_loss(
+    *,
+    values: torch.Tensor,
+    old_values: torch.Tensor,
+    returns: torch.Tensor,
+    clip_range: float,
+) -> torch.Tensor:
+    """PPO clipped value loss. Element-wise; reduction is the caller's job.
+
+    Returns per-token ``0.5 * max((V - R)², (V_clipped - R)²)`` with
+    ``V_clipped = V_old + clip(V - V_old, -clip_range, clip_range)``.
+    """
+    values_f = values.float()
+    old_f = old_values.detach().float()
+    returns_f = returns.detach().float()
+    clipped = old_f + (values_f - old_f).clamp(-clip_range, clip_range)
+    sq1 = (values_f - returns_f).square()
+    sq2 = (clipped - returns_f).square()
+    return 0.5 * torch.maximum(sq1, sq2)
+
+
 # ---------------------------------------------------------------------------
 # Reference-policy KL helpers (FlowGRPO / FlowDPPO ``beta`` term)
 # ---------------------------------------------------------------------------
diff --git a/unirl/algorithms/ppo.py b/unirl/algorithms/ppo.py
new file mode 100644
index 000000000..e4599c69c
--- /dev/null
+++ b/unirl/algorithms/ppo.py
@@ -0,0 +1,238 @@
+"""PPO with GAE for autoregressive ``TextSegment`` training.
+
+Uses per-token GAE advantages from ``segment.token_advantages`` (populated in
+:meth:`prepare_rollout_track`) and clipped value loss against ``segment.returns``.
+Policy clip math is shared with :class:`GRPO` via :func:`_grpo_clip_loss`;
+value clip math via :func:`_ppo_clipped_value_loss` in the same module.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any, Dict, Mapping, Optional, Type
+
+import torch
+
+from unirl.models.types.replay_result import ReplayResult
+from unirl.types.conditions import Condition
+from unirl.types.rollout_resp import RolloutTrack
+from unirl.types.segments.text import TextSegment
+
+from .base import (
+    AlgorithmStepResult,
+    BaseAlgorithmConfig,
+    StageAlgorithm,
+    _grpo_clip_loss,
+    _ppo_clipped_value_loss,
+    _resolve_clip_range_from_schedule,
+    rollout_replay_logp_absdiff,
+    typed_conditions,
+)
+
+
+@dataclass
+class PPOConfig(BaseAlgorithmConfig):
+    stage_attr: str = "ar"
+    conditions_cls: str = ""
+    clip_range: float = 0.2
+    clip_range_high: Optional[float] = None
+    clip_schedule: str = "constant"
+    cliprange_value: float = 0.2
+    vf_coef: float = 0.5
+    gae_gamma: float = 1.0
+    gae_lambda: float = 0.95
+    loss_agg_mode: str = "token-mean"
+    horizon: int = 8192
+
+
+def _aggregate_token_loss(
+    loss_per_elem: torch.Tensor,
+    *,
+    segment: TextSegment,
+    loss_agg_mode: str,
+    horizon: int,
+) -> torch.Tensor:
+    if loss_agg_mode in ("seq-mean-token-sum-norm", "seq-mean-token-mean") and segment.lengths is not None:
+        parts = torch.split(loss_per_elem, segment.lengths.tolist())
+        if loss_agg_mode == "seq-mean-token-sum-norm":
+            return torch.stack([p.sum() for p in parts]).mean() / float(horizon)
+        return torch.stack([p.mean() if p.numel() else p.new_zeros(()) for p in parts]).mean()
+    return loss_per_elem.mean()
+
+
+class PPO(StageAlgorithm):
+    """PPO with GAE over an AR ``TextSegment``.
+
+    :meth:`prepare_rollout_track` runs a no-grad critic replay (``return_values=True``),
+    stores ``segment.values`` as the frozen value anchor, then calls
+    :meth:`RolloutTrack.compute_gae_advantages`. Each train micro-batch replays
+    for fresh log-probs and value predictions for the policy and value losses.
+    """
+
+    supports_multi_update = True
+
+    def __init__(
+        self,
+        *,
+        stage: Any = None,
+        pipeline: Any = None,
+        stage_attr: str = "ar",
+        clip_range: float = 0.2,
+        clip_schedule: str = "constant",
+        clip_range_high: Optional[float] = None,
+        cliprange_value: float = 0.2,
+        vf_coef: float = 0.5,
+        gae_gamma: float = 1.0,
+        gae_lambda: float = 0.95,
+        loss_agg_mode: str = "token-mean",
+        horizon: int = 8192,
+        conditions_cls: Optional[Type[Any]] = None,
+        sampling_temperature: Optional[float] = None,
+    ) -> None:
+        super().__init__()
+        if stage is None and pipeline is None:
+            raise ValueError("PPO: either `stage` or `pipeline` must be provided")
+        if stage is None:
+            stage = getattr(pipeline, stage_attr)
+        self.stage = stage
+        self.clip_range = float(clip_range)
+        self.clip_range_high = None if clip_range_high is None else float(clip_range_high)
+        self.clip_schedule = str(clip_schedule)
+        self.cliprange_value = float(cliprange_value)
+        self.vf_coef = float(vf_coef)
+        self.gae_gamma = float(gae_gamma)
+        self.gae_lambda = float(gae_lambda)
+        self.loss_agg_mode = str(loss_agg_mode)
+        self.horizon = int(horizon)
+        self.conditions_cls = conditions_cls
+        if sampling_temperature is None:
+            from unirl.types.sampling import ARSamplingParams
+
+            sampling_temperature = ARSamplingParams.__dataclass_fields__["temperature"].default
+        self.sampling_temperature = float(sampling_temperature)
+
+    def prepare_rollout_track(self, track: RolloutTrack) -> None:
+        """Replay critic values and compute GAE on the worker shard."""
+        if track.rewards is None:
+            raise ValueError("PPO.prepare_rollout_track: track has no rewards")
+        if track.segment is None or not isinstance(track.segment, TextSegment):
+            raise ValueError("PPO.prepare_rollout_track: requires a TextSegment")
+        segment = track.segment
+        if segment.log_probs is None:
+            raise ValueError("PPO.prepare_rollout_track: segment.log_probs is None")
+
+        typed_conds = typed_conditions(track.conditions, self.conditions_cls)
+        with torch.no_grad():
+            replay_out = self.stage.replay(
+                typed_conds,
+                segment=segment,
+                temperature=self.sampling_temperature,
+                return_values=True,
+            )
+        values = _replay_values(replay_out)
+        track.segment = replace(segment, values=values)
+        updated = track.compute_gae_advantages(gamma=self.gae_gamma, gae_lambda=self.gae_lambda)
+        track.segment = updated.segment
+        track.advantages = updated.advantages
+
+    def compute_loss_and_backward(
+        self,
+        *,
+        conditions: Mapping[str, Condition],
+        segment: TextSegment,
+        advantages: torch.Tensor,
+        training_progress: float,
+        loss_scale: float,
+    ) -> AlgorithmStepResult:
+        del advantages  # GAE uses segment.token_advantages instead of track-level GRPO scalars.
+        if segment.tokens is None or segment.lengths is None or segment.log_probs is None:
+            return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False)
+        if int(segment.tokens.shape[0]) == 0:
+            return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False)
+        if segment.token_advantages is None or segment.returns is None or segment.values is None:
+            raise ValueError(
+                "PPO.compute_loss_and_backward: segment requires token_advantages, returns, and values "
+                "(call prepare_rollout_track first)."
+            )
+
+        typed_conds = typed_conditions(conditions, self.conditions_cls)
+        replay_out = self.stage.replay(
+            typed_conds,
+            segment=segment,
+            temperature=self.sampling_temperature,
+            return_values=True,
+        )
+        new_logp = _replay_log_probs(replay_out)
+        new_values = _replay_values(replay_out)
+
+        old_logp = segment.log_probs.to(dtype=new_logp.dtype, device=new_logp.device)
+        old_values = segment.values.to(dtype=new_values.dtype, device=new_values.device)
+        returns = segment.returns.to(dtype=new_values.dtype, device=new_values.device)
+        adv_per_token = segment.token_advantages.detach().to(dtype=new_logp.dtype, device=new_logp.device)
+
+        clip_range = _resolve_clip_range_from_schedule(self.clip_range, self.clip_schedule, training_progress)
+        clip_high = (
+            None
+            if self.clip_range_high is None
+            else _resolve_clip_range_from_schedule(self.clip_range_high, self.clip_schedule, training_progress)
+        )
+        policy_per_elem, ratio_metrics = _grpo_clip_loss(
+            new_logp=new_logp,
+            old_logp=old_logp,
+            advantages=adv_per_token,
+            clip_range=clip_range,
+            clip_range_high=clip_high,
+        )
+        value_per_elem = _ppo_clipped_value_loss(
+            values=new_values,
+            old_values=old_values,
+            returns=returns,
+            clip_range=self.cliprange_value,
+        )
+        total_per_elem = policy_per_elem + self.vf_coef * value_per_elem
+        loss = _aggregate_token_loss(
+            total_per_elem,
+            segment=segment,
+            loss_agg_mode=self.loss_agg_mode,
+            horizon=self.horizon,
+        )
+        (loss * loss_scale).backward()
+
+        policy_loss = _aggregate_token_loss(
+            policy_per_elem,
+            segment=segment,
+            loss_agg_mode=self.loss_agg_mode,
+            horizon=self.horizon,
+        )
+        value_loss = _aggregate_token_loss(
+            value_per_elem,
+            segment=segment,
+            loss_agg_mode=self.loss_agg_mode,
+            horizon=self.horizon,
+        )
+        metrics: Dict[str, Any] = {
+            "policy_loss": float(policy_loss.detach().item()),
+            "value_loss": float(value_loss.detach().item()),
+            "clip_range": float(clip_range),
+            **rollout_replay_logp_absdiff(new_logp, old_logp),
+            **{k: float(v.item()) for k, v in ratio_metrics.items()},
+        }
+        return AlgorithmStepResult(
+            loss=float(loss.detach().item()),
+            metrics=metrics,
+            num_steps_or_tokens=int(new_logp.shape[0]),
+            has_backward=True,
+        )
+
+
+def _replay_log_probs(replay_out: ReplayResult) -> torch.Tensor:
+    return replay_out.log_probs
+
+
+def _replay_values(replay_out: ReplayResult) -> torch.Tensor:
+    if replay_out.values is None:
+        raise ValueError("PPO: replay with return_values=True must return ReplayResult.values")
+    return replay_out.values
+
+
+__all__ = ["PPO", "PPOConfig"]
diff --git a/unirl/train/stack/base.py b/unirl/train/stack/base.py
index b4160fec7..d035ba5b5 100644
--- a/unirl/train/stack/base.py
+++ b/unirl/train/stack/base.py
@@ -192,6 +192,9 @@ def prepare_segment(self, part: Part, *, plans: Plan) -> None:
         if part.segment is None:
             return
         algorithm = self.algorithm
+        prepare_track = getattr(algorithm, "prepare_rollout_track", None)
+        if prepare_track is not None:
+            prepare_track(resp_track)
         if not algorithm.recomputes_anchor():
             algorithm.prepare_segment(conditions=part.conditions, segment=part.segment)
             return
diff --git a/unirl/train_ar.py b/unirl/train_ar.py
index 063418718..3919b6b08 100755
--- a/unirl/train_ar.py
+++ b/unirl/train_ar.py
@@ -51,6 +51,7 @@ def teardown() -> None:
             logging_cfg=cfg.get("logging"),
             adv_normalization_scope=cfg.get("adv_normalization_scope", "group"),
             normalize_adv_by_std=cfg.get("normalize_adv_by_std", True),
+            advantage_mode=cfg.get("advantage_mode", "grpo"),
             balance_shards=cfg.get("balance_shards", False),
             eval_interval=cfg.get("eval_interval", 0),
             eval_num_prompts=cfg.get("eval_num_prompts", -1),
diff --git a/unirl/trainer/ar.py b/unirl/trainer/ar.py
index 2371dfada..93dafcca6 100644
--- a/unirl/trainer/ar.py
+++ b/unirl/trainer/ar.py
@@ -60,6 +60,7 @@ def __init__(
         logging_cfg: Optional[DictConfig] = None,
         adv_normalization_scope: str = "group",
         normalize_adv_by_std: bool = True,
+        advantage_mode: str = "grpo",
         balance_shards: bool = False,
         eval_interval: int = 0,
         eval_num_prompts: int = -1,
@@ -83,6 +84,8 @@ def __init__(
         # group std. False = mean-center only (reward - group_mean), NO std division —
         # removes the difficulty bias that over-amplifies low-std (hard) prompts.
         self.normalize_adv_by_std = normalize_adv_by_std
+        # "grpo" (default) or "gae" (PPO path — GAE runs in PPO.prepare_rollout_track on workers).
+        self.advantage_mode = str(advantage_mode)
         # verl trainer.balance_batch parity: driver-side reorder of the rollout
         # batch so each DP shard receives a similar total-token workload. FSDP
         # collectives sync all ranks every micro, so a step runs at the SLOWEST
@@ -456,7 +459,11 @@ def train_step(
             if isinstance(part.component_rewards, dict):
                 part.component_rewards = {name: hydrate(value) for name, value in part.component_rewards.items()}
             mean_reward = float(part.rewards.to(torch.float32).mean().item())
-            part = part.compute_advantages(normalize=self.normalize_adv_by_std, scope=self.adv_normalization_scope)
+            if self.advantage_mode == "grpo":
+                part = part.compute_advantages(
+                    normalize=self.normalize_adv_by_std,
+                    scope=self.adv_normalization_scope,
+                )
             sample = sample.with_parts([*sample.parts[:-1], part])
 
         self._dump_rollout_samples(sample, rollout_id)
diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py
index a64e517d2..5199edb8f 100644
--- a/unirl/trainer/async_ar.py
+++ b/unirl/trainer/async_ar.py
@@ -293,7 +293,11 @@ def _advantage_and_train(
         if part.rewards is not None:
             part.rewards = hydrate(part.rewards)
             mean_reward = float(part.rewards.to(torch.float32).mean().item())
-        part = part.compute_advantages(normalize=self.normalize_adv_by_std, scope=self.adv_normalization_scope)
+        if self.advantage_mode == "grpo":
+            part = part.compute_advantages(
+                normalize=self.normalize_adv_by_std,
+                scope=self.adv_normalization_scope,
+            )
         sample = sample.with_parts([*sample.parts[:-1], part])
         train_part = part
         if self.balance_shards:

From 3a3dae09df32d88e8ed23d89e16be358804e15a8 Mon Sep 17 00:00:00 2001
From: Yi Heng Lim <limjohnyh@gmail.com>
Date: Sun, 2 Aug 2026 02:25:52 +0800
Subject: [PATCH 10/11] feat(models): add AR value head and GAE wiring (#256)

Co-authored-by: leviking98z-rgb <leviking98z@gmail.com>
---
 unirl/models/qwen3/ar.py            | 148 ++++++++++++++++++++++------
 unirl/models/qwen3/bundle.py        |  29 ++++++
 unirl/models/qwen3/config.py        |   3 +
 unirl/models/types/replay_result.py |  10 +-
 unirl/models/types/value_head.py    |  35 +++++++
 unirl/types/advantages.py           |  28 ++++++
 unirl/types/sample.py               |  87 +++++++++++++++-
 unirl/types/segments/text.py        |   4 +
 8 files changed, 310 insertions(+), 34 deletions(-)
 create mode 100644 unirl/models/types/value_head.py

diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py
index 1e072ac76..793e06827 100644
--- a/unirl/models/qwen3/ar.py
+++ b/unirl/models/qwen3/ar.py
@@ -23,7 +23,7 @@
 from dataclasses import dataclass
 from dataclasses import field as dc_field
 from types import MethodType
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
@@ -31,6 +31,7 @@
 from torch.utils.checkpoint import checkpoint
 
 from unirl.models.types.ar import ARSamplingParams, ARStage, ARStep, left_pad_prompt
+from unirl.models.types.replay_result import ReplayResult
 from unirl.types.segments import TextSegment
 from unirl.utils.dtypes import parse_torch_dtype
 
@@ -87,6 +88,7 @@ def _replay_aware_forward(
     temperature: float = 1.0,
     autocast_dtype: Optional[torch.dtype] = None,
     packed_predict_index: Optional[torch.Tensor] = None,
+    return_values: bool = False,
     **kw: Any,
 ) -> Any:
     """Dual-mode ``forward`` installed on the Qwen3 CausalLM instance.
@@ -108,6 +110,8 @@ def _replay_aware_forward(
                 return f(self, **kw)
         raise RuntimeError("_replay_aware_forward: no class-level forward found in the MRO")
 
+    _require_value_head_for_replay(self, return_values)
+
     # cuDNN's fused SDPA backward (ScaledDotProductCudnnAttentionBackward0) returns
     # NaN grads on some bf16 sequences while the forward stays finite (confirmed via
     # torch.autograd.detect_anomaly): it floods every parameter grad and forces the
@@ -129,6 +133,7 @@ def _replay_aware_forward(
     # [B, chunk, vocab] FP32 transient stays ~1.2 GiB, and each chunk is
     # gradient-checkpointed (recomputed in backward rather than held).
     T = float(temperature) if float(temperature) > 0.0 else 1.0
+    value_head = getattr(self, "value_head", None) if return_values else None
 
     if packed_predict_index is not None:
         # Packed varlen replay: ``hidden`` is one packed row [1, L_total, H]
@@ -155,8 +160,16 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
             else:
                 flat_parts.append(_flat_logp_chunk(h, tok))
         if not flat_parts:
-            return hidden.new_zeros((0,), dtype=torch.float32)
-        return torch.cat(flat_parts, dim=0)
+            empty = hidden.new_zeros((0,), dtype=torch.float32)
+            if value_head is None:
+                return empty
+            return ReplayResult(log_probs=empty, values=empty)
+        log_probs = torch.cat(flat_parts, dim=0)
+        if value_head is None:
+            return log_probs
+        value_parts = [value_head(h_pred[s : s + flat_chunk]) for s in range(0, int(h_pred.size(0)), flat_chunk)]
+        values = torch.cat(value_parts, dim=0) if value_parts else log_probs.new_zeros(0)
+        return ReplayResult(log_probs=log_probs, values=values)
     T_max = int(response_tokens.size(1))
     resp_hidden = hidden[:, prompt_len - 1 : prompt_len - 1 + T_max, :]
 
@@ -176,8 +189,51 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor:
         else:
             parts.append(_logp_chunk(h, tok))
     if not parts:
-        return resp_hidden.new_zeros((bsz, 0), dtype=torch.float32)  # T_max == 0
-    return torch.cat(parts, dim=1)
+        empty = resp_hidden.new_zeros((bsz, 0), dtype=torch.float32)
+        if value_head is None:
+            return empty
+        return ReplayResult(log_probs=empty, values=empty)
+    log_probs = torch.cat(parts, dim=1)
+    if value_head is None:
+        return log_probs
+    value_parts = [value_head(resp_hidden[:, s : s + chunk, :]) for s in range(0, T_max, chunk)]
+    values = torch.cat(value_parts, dim=1) if value_parts else log_probs.new_zeros((bsz, 0))
+    return ReplayResult(log_probs=log_probs, values=values)
+
+
+def _require_value_head_for_replay(model: Any, return_values: bool) -> None:
+    if return_values and getattr(model, "value_head", None) is None:
+        raise ValueError(
+            "Qwen3 replay: return_values=True requires a value head (set use_value_head=True in the pipeline config)"
+        )
+
+
+def _finalize_replay_output(
+    out: Union[torch.Tensor, ReplayResult],
+    *,
+    segment: TextSegment,
+    return_values: bool,
+    logprob_dtype: torch.dtype,
+) -> Union[torch.Tensor, ReplayResult]:
+    """Cast log-probs and flatten padded critic values to segment order."""
+    if not isinstance(out, ReplayResult):
+        if return_values:
+            raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
+        return out.to(dtype=logprob_dtype)
+
+    log_probs = out.log_probs.to(dtype=logprob_dtype)
+    if not return_values:
+        return log_probs
+    if out.values is None:
+        raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
+    if log_probs.ndim == 1:
+        return ReplayResult(log_probs=log_probs, values=out.values.float())
+    if segment.lengths is None:
+        raise ValueError("Qwen3ARStage.replay: segment requires lengths to flatten critic values")
+
+    flat_values = [out.values[b, : int(length)] for b, length in enumerate(segment.lengths.tolist()) if int(length) > 0]
+    values = torch.cat(flat_values, dim=0) if flat_values else out.values.new_zeros(0)
+    return ReplayResult(log_probs=log_probs, values=values.float())
 
 
 # Attention backends with a sparse packed kernel (skip cross-sequence blocks):
@@ -421,22 +477,32 @@ def replay(
         *,
         segment: TextSegment,
         temperature: float = 1.0,
-    ) -> torch.Tensor:
+        return_values: bool = False,
+    ) -> Union[torch.Tensor, ReplayResult]:
         """Per-token log-prob replay over a stored rollout segment.
 
         Branch: prefer :meth:`packed_replay` (packed-varlen, zero padding, B > 1)
         and fall back to :meth:`padding_replay` (the dense ``[B, P_max + T_max]``
-        padded path) when packing does not apply. Returns packed varlen
-        ``[total_tokens]`` aligned with ``segment.log_probs``; caller controls
-        grad / ``.train()`` scope. ``temperature`` divides logits before
-        ``log_softmax`` to match SGLang's sampler (``1.0`` is a no-op).
+        padded path) when packing does not apply. Returns packed varlen log-probs,
+        or a :class:`ReplayResult` with aligned critic values when requested.
         """
+        _require_value_head_for_replay(self.model.transformer, return_values)
         attn_impl = getattr(getattr(self.model.transformer, "config", None), "_attn_implementation", None)
         if _packed_replay_supported(attn_impl):
-            packed = self.packed_replay(conditions, segment=segment, temperature=temperature)
+            packed = self.packed_replay(
+                conditions,
+                segment=segment,
+                temperature=temperature,
+                return_values=return_values,
+            )
             if packed is not None:
                 return packed
-        return self.padding_replay(conditions, segment=segment, temperature=temperature)
+        return self.padding_replay(
+            conditions,
+            segment=segment,
+            temperature=temperature,
+            return_values=return_values,
+        )
 
     def packed_replay(
         self,
@@ -444,7 +510,8 @@ def packed_replay(
         *,
         segment: TextSegment,
         temperature: float = 1.0,
-    ) -> Optional[torch.Tensor]:
+        return_values: bool = False,
+    ) -> Optional[Union[torch.Tensor, ReplayResult]]:
         """Packed-varlen replay (B > 1): zero padding anywhere.
 
         Concatenate every sample's REAL prompt tokens + its flat response tokens
@@ -520,9 +587,15 @@ def packed_replay(
             packed_predict_index=predict_index,
             prompt_len=0,
             temperature=temperature,
+            return_values=return_values,
             autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None),
         )
-        return per_token_flat.to(dtype=self.logprob_dtype)
+        return _finalize_replay_output(
+            per_token_flat,
+            segment=segment,
+            return_values=return_values,
+            logprob_dtype=self.logprob_dtype,
+        )
 
     def padding_replay(
         self,
@@ -530,7 +603,8 @@ def padding_replay(
         *,
         segment: TextSegment,
         temperature: float = 1.0,
-    ) -> torch.Tensor:
+        return_values: bool = False,
+    ) -> Union[torch.Tensor, ReplayResult]:
         """Dense ``[B, P_max + T_max]`` padded replay — the default / fallback path.
 
         One teacher-forced forward over padded ``prompt + response``; gather
@@ -629,28 +703,44 @@ def padding_replay(
         # root-wrapped or plain) and never materializes [B, L, vocab] logits.
         # The cuda-vs-cpu autocast decision lives here; dtype validity and the
         # autocast scope live in the patched forward.
-        per_token = self.model.transformer(
+        out = self.model.transformer(
             input_ids=full_ids,
             attention_mask=full_mask,
             position_ids=position_ids,
             response_tokens=response_tokens,
             prompt_len=prompt_len,
             temperature=temperature,
+            return_values=return_values,
             autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None),
-        )  # [B, T_max] FP32
-
-        if T_max == 0:
-            return torch.zeros(0, dtype=self.logprob_dtype, device=device)
+        )  # [B, T_max] FP32 tensor or ReplayResult
 
-        flat: List[torch.Tensor] = []
-        for b in range(batch_size):
-            n = lengths[b]
-            if n == 0:
-                continue
-            flat.append(per_token[b, :n])
-        if not flat:
-            return torch.zeros(0, dtype=self.logprob_dtype, device=device)
-        return torch.cat(flat, dim=0).to(dtype=self.logprob_dtype)
+        finalized = _finalize_replay_output(
+            out,
+            segment=segment,
+            return_values=return_values,
+            logprob_dtype=self.logprob_dtype,
+        )
+        if isinstance(finalized, ReplayResult):
+            per_token = finalized.log_probs
+            values = finalized.values
+        else:
+            per_token = finalized
+            values = None
+
+        flat_log_probs: List[torch.Tensor] = []
+        for b, n in enumerate(lengths):
+            if n > 0:
+                flat_log_probs.append(per_token[b, :n])
+        log_probs = (
+            torch.cat(flat_log_probs, dim=0)
+            if flat_log_probs
+            else torch.zeros(0, dtype=self.logprob_dtype, device=device)
+        )
+        if not return_values:
+            return log_probs
+        if values is None:
+            raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values")
+        return ReplayResult(log_probs=log_probs, values=values)
 
     def _resolve_stop_ids(
         self,
diff --git a/unirl/models/qwen3/bundle.py b/unirl/models/qwen3/bundle.py
index 3c13308a8..d44201f0d 100644
--- a/unirl/models/qwen3/bundle.py
+++ b/unirl/models/qwen3/bundle.py
@@ -28,6 +28,7 @@
 
 from unirl.models.types.bundle import Bundle
 from unirl.models.types.meta_init import build_meta_init_transformer
+from unirl.models.types.value_head import ValueHead
 from unirl.utils.dtypes import parse_torch_dtype
 
 from .config import Qwen3PipelineConfig
@@ -35,6 +36,24 @@
 logger = logging.getLogger(__name__)
 
 
+def _stamp_value_head_reset(transformer: nn.Module) -> None:
+    """Zero checkpoint-absent value-head params after meta materialization."""
+    from unirl.train.deferred import _stamp
+
+    def _reset(model: nn.Module) -> None:
+        reset: list[str] = []
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if name.startswith("value_head."):
+                    param.zero_()
+                    reset.append(name)
+        if not reset:
+            raise RuntimeError("Qwen3 meta-init: value_head parameters disappeared before post-load reset")
+        logger.info("Qwen3 meta-init: zero-initialized checkpoint-absent value head: %s", reset)
+
+    _stamp(transformer, _reset)
+
+
 class Qwen3Bundle(Bundle):
     """Qwen3 bundle: causal-LM transformer + matching tokenizer."""
 
@@ -115,6 +134,16 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle":
         if tokenizer.pad_token is None and tokenizer.eos_token is not None:
             tokenizer.pad_token = tokenizer.eos_token
 
+        if config.use_value_head:
+            hidden_size = int(getattr(transformer.config, "hidden_size"))
+            transformer_device = next(transformer.parameters()).device
+            transformer.value_head = ValueHead(hidden_size, device=transformer_device)
+            if config.meta_init_transformer:
+                # The base checkpoint has no value_head.* tensors. ``to_empty``
+                # materializes the meta head as uninitialized storage, so reset
+                # it only after the sharded checkpoint load has completed.
+                _stamp_value_head_reset(transformer)
+
         bundle = cls(
             transformer=transformer,
             tokenizer=tokenizer,
diff --git a/unirl/models/qwen3/config.py b/unirl/models/qwen3/config.py
index 37f63b5ef..2fa5d3504 100644
--- a/unirl/models/qwen3/config.py
+++ b/unirl/models/qwen3/config.py
@@ -69,6 +69,9 @@ class Qwen3PipelineConfig:
     use_lora: bool = False
     lora_target_modules: Optional[List[str]] = None
 
+    # Attach a scalar critic to the causal LM for PPO / GAE training.
+    use_value_head: bool = False
+
     system_instruction: Optional[str] = None
     # Chat-template thinking switch; MUST agree with the rollout engine's
     # chat_template_kwargs.enable_thinking or train/rollout prompts diverge.
diff --git a/unirl/models/types/replay_result.py b/unirl/models/types/replay_result.py
index 3bf7a7222..381d513c5 100644
--- a/unirl/models/types/replay_result.py
+++ b/unirl/models/types/replay_result.py
@@ -8,10 +8,8 @@
 
 Diffusion stages populate ``log_probs`` and ``prev_sample_means`` (the
 mean of the SDE Gaussian — μ_θ — used as the second moment in the KL
-penalty). AR stages currently return a plain ``Tensor`` (signature
-divergence with diffusion is intentional for now); when AR replay grows
-``logits``-based KL support, the ``logits`` field on this result will be
-the canonical home.
+penalty). AR stages return a plain ``Tensor`` for policy-only replay, or
+a :class:`ReplayResult` when optional critic ``values`` are requested.
 """
 
 from __future__ import annotations
@@ -42,5 +40,9 @@ class ReplayResult:
     or entropy penalty support; not needed for Binary KL (which uses
     only per-token log-probs). Currently not populated."""
 
+    values: Optional[torch.Tensor] = None
+    """Per-token critic predictions ``V_t``. Packed ``[total_tokens]`` for AR.
+    ``None`` when replay did not request a value head."""
+
 
 __all__ = ["ReplayResult"]
diff --git a/unirl/models/types/value_head.py b/unirl/models/types/value_head.py
new file mode 100644
index 000000000..1193a8017
--- /dev/null
+++ b/unirl/models/types/value_head.py
@@ -0,0 +1,35 @@
+"""Scalar value head for PPO-style critic training on AR hidden states."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ValueHead(nn.Module):
+    """Linear critic ``V(h)`` on the hidden state that predicts each action."""
+
+    def __init__(self, hidden_size: int, *, device: Any = None) -> None:
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, 1, bias=True, dtype=torch.float32, device=device)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """Start from ``V(s)=0``; deterministic and safe after sharded meta-init."""
+        nn.init.zeros_(self.proj.weight)
+        if self.proj.bias is not None:
+            nn.init.zeros_(self.proj.bias)
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        """Map ``[..., H]`` hidden states to FP32 scalar values ``[...]``."""
+        # FSDP mixed precision may expose gathered parameters in its compute
+        # dtype even when the sharded masters are FP32.
+        weight = self.proj.weight.float()
+        bias = self.proj.bias.float() if self.proj.bias is not None else None
+        return F.linear(hidden.float(), weight, bias).squeeze(-1)
+
+
+__all__ = ["ValueHead"]
diff --git a/unirl/types/advantages.py b/unirl/types/advantages.py
index f49f96b6e..d4c3d4b35 100644
--- a/unirl/types/advantages.py
+++ b/unirl/types/advantages.py
@@ -71,6 +71,34 @@ def compute_gae_advantages(
     return advantages, returns
 
 
+def scatter_terminal_rewards(
+    rewards_per_sample: torch.Tensor,
+    *,
+    cu_seqlens: torch.Tensor,
+) -> torch.Tensor:
+    """Scatter each trajectory reward onto its final packed response token."""
+    if rewards_per_sample.ndim != 1:
+        raise ValueError(f"scatter_terminal_rewards: expected 1D rewards, got shape {tuple(rewards_per_sample.shape)}")
+    if cu_seqlens.ndim != 1 or cu_seqlens.numel() == 0:
+        raise ValueError("scatter_terminal_rewards: cu_seqlens must be a non-empty 1D tensor")
+    batch_size = int(cu_seqlens.numel()) - 1
+    if int(rewards_per_sample.numel()) != batch_size:
+        raise ValueError(
+            f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.numel())}) "
+            f"!= packed batch ({batch_size})"
+        )
+
+    cu = [int(offset) for offset in cu_seqlens.tolist()]
+    if cu[0] != 0 or any(end < start for start, end in zip(cu, cu[1:])):
+        raise ValueError(f"scatter_terminal_rewards: invalid cumulative offsets {cu}")
+
+    token_rewards = rewards_per_sample.new_zeros(cu[-1])
+    for reward, start, end in zip(rewards_per_sample, cu, cu[1:]):
+        if end > start:
+            token_rewards[end - 1] = reward
+    return token_rewards
+
+
 def _gae_1d(
     rewards: torch.Tensor,
     values: torch.Tensor,
diff --git a/unirl/types/sample.py b/unirl/types/sample.py
index fede101f6..797207d22 100644
--- a/unirl/types/sample.py
+++ b/unirl/types/sample.py
@@ -31,12 +31,14 @@
     shared_field,
 )
 from unirl.distributed.tensor.ref import hydrate
+from unirl.types.advantages import compute_gae_advantages as _compute_gae
+from unirl.types.advantages import scatter_terminal_rewards
 from unirl.types.conditions import Condition
 from unirl.types.media_preview import MediaPreview
 from unirl.types.primitives import Audios, Images, Texts, Videos, primitive_modality_key
 from unirl.types.sample_id import ancestor_id, child_id, parent_id
 from unirl.types.sampling import BaseSamplingParams
-from unirl.types.segments import Segment
+from unirl.types.segments import Segment, TextSegment
 from unirl.utils.shard_balance import lpt_shard_permutation, shard_token_spread
 
 logger = logging.getLogger(__name__)
@@ -418,6 +420,89 @@ def compute_advantages(
             adv = reshaped - mean
         return _part_with_field(self, "advantages", adv.flatten())
 
+    def compute_gae_advantages(
+        self,
+        *,
+        gamma: float = 1.0,
+        gae_lambda: float = 0.95,
+    ) -> "Part":
+        """Attach packed token advantages and returns for an AR PPO update.
+
+        ``loss_mask`` controls which tokens later contribute to the policy and
+        value losses. It deliberately does not control the GAE recursion:
+        masked actions are still states in the same trajectory, so terminal
+        reward must propagate across them.
+        """
+        if self.rewards is None:
+            raise ValueError("Part.compute_gae_advantages: part has no rewards")
+        if not isinstance(self.segment, TextSegment):
+            raise ValueError("Part.compute_gae_advantages: requires a TextSegment")
+        segment = self.segment
+        if segment.values is None:
+            raise ValueError("Part.compute_gae_advantages: segment.values is None")
+        if segment.cu_seqlens is None or segment.lengths is None:
+            raise ValueError(
+                "Part.compute_gae_advantages: segment requires framework-managed "
+                "cu_seqlens (construct via TextSegment.pack)"
+            )
+
+        values = hydrate(segment.values).to(torch.float32)
+        cu_seqlens = segment.cu_seqlens.to(device=values.device)
+        total_tokens = int(cu_seqlens[-1].item())
+        if values.ndim != 1 or int(values.numel()) != total_tokens:
+            raise ValueError(
+                "Part.compute_gae_advantages: values must be a packed 1D tensor "
+                f"with {total_tokens} elements, got shape {tuple(values.shape)}"
+            )
+
+        rewards = hydrate(self.rewards).to(device=values.device, dtype=torch.float32)
+        token_rewards = scatter_terminal_rewards(rewards, cu_seqlens=cu_seqlens)
+        token_advantages = values.new_zeros(values.shape)
+        token_returns = values.new_zeros(values.shape)
+        cu = [int(offset) for offset in cu_seqlens.tolist()]
+        for start, end in zip(cu, cu[1:]):
+            if end <= start:
+                continue
+            advantages, returns = _compute_gae(
+                token_rewards[start:end],
+                values[start:end],
+                gamma=gamma,
+                gae_lambda=gae_lambda,
+            )
+            token_advantages[start:end] = advantages
+            token_returns[start:end] = returns
+
+        segment_fields = {f.name: getattr(segment, f.name) for f in dc_fields(segment)}
+        segment_fields["token_advantages"] = token_advantages
+        segment_fields["returns"] = token_returns
+        updated_segment = segment._rebuild(segment_fields)
+
+        # Keep the existing per-sample advantage metric meaningful. The loss
+        # mask applies here only as a reduction mask, never as a done signal.
+        loss_mask = None
+        if segment.loss_mask is not None:
+            loss_mask = hydrate(segment.loss_mask).to(device=values.device, dtype=torch.bool)
+            if loss_mask.shape != values.shape:
+                raise ValueError(
+                    "Part.compute_gae_advantages: loss_mask shape "
+                    f"{tuple(loss_mask.shape)} != values shape {tuple(values.shape)}"
+                )
+        sample_advantages: List[torch.Tensor] = []
+        for start, end in zip(cu, cu[1:]):
+            if end <= start:
+                sample_advantages.append(values.new_zeros(()))
+                continue
+            selected = token_advantages[start:end]
+            if loss_mask is not None:
+                selected = selected[loss_mask[start:end]]
+            sample_advantages.append(selected.mean() if selected.numel() else values.new_zeros(()))
+        mean_advantages = (
+            torch.stack(sample_advantages) if sample_advantages else values.new_zeros((0,), dtype=torch.float32)
+        )
+
+        updated = _part_with_field(self, "segment", updated_segment)
+        return _part_with_field(updated, "advantages", mean_advantages)
+
 
 def _part_with_field(part: Part, field_name: str, value: Any) -> Part:
     """Copy of ``part`` with one field replaced."""
diff --git a/unirl/types/segments/text.py b/unirl/types/segments/text.py
index 7a1570f23..8850f902d 100644
--- a/unirl/types/segments/text.py
+++ b/unirl/types/segments/text.py
@@ -43,6 +43,10 @@ class TextSegment(Segment):
     # Original engine emission retained when an algorithm replaces ``log_probs``
     # with a train-side replay anchor.
     rollout_log_probs: Optional[torch.Tensor] = packed_field(default=None)
+    # Optional PPO critic state, aligned one-to-one with packed response tokens.
+    values: Optional[torch.Tensor] = packed_field(default=None)
+    returns: Optional[torch.Tensor] = packed_field(default=None)
+    token_advantages: Optional[torch.Tensor] = packed_field(default=None)
 
     def as_condition_with(self, encoder: Callable[..., Any]) -> Condition:
         """Re-embed packed tokens via the supplied encoder into a TextEmbedCondition.

From cda258220f2cec891bd851283b715882db09b8c0 Mon Sep 17 00:00:00 2001
From: leviking98z-rgb <leviking98z@gmail.com>
Date: Sun, 2 Aug 2026 02:33:54 +0800
Subject: [PATCH 11/11] fix(algorithms): harden AR PPO integration after review

---
 .../ar/qwen3_ppo_4b_base_dapo_sglang.yaml     |  67 ++---
 unirl/algorithms/__init__.py                  |   4 +-
 unirl/algorithms/base.py                      |  31 +--
 unirl/algorithms/ppo.py                       | 232 ++++++++++++------
 unirl/train/stack/base.py                     |   4 +-
 unirl/train_async_ar.py                       |   1 +
 unirl/trainer/ar.py                           |   5 +-
 unirl/trainer/async_ar.py                     |   4 +
 8 files changed, 211 insertions(+), 137 deletions(-)

diff --git a/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml b/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml
index a2990960e..84d86d663 100644
--- a/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml
+++ b/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml
@@ -1,16 +1,13 @@
 # @package _global_
-# PPO + GAE Qwen3-4B-Base on DAPO-Math — ARTrainer + SGLang rollout.
+# PPO + GAE on Qwen3-4B-Base / DAPO-Math with colocated SGLang rollout.
 #
-# Sibling of qwen3_grpo_4b_base_dapo_sglang.yaml with a value critic and GAE
-# advantages instead of group-relative GRPO normalization. Requires
-# bundle.config.use_value_head=true and advantage_mode=gae.
-#
-# Data prep (one-time):
+# Prepare data:
 #   python -m unirl.utils.prepare_dapo_math --out-dir data/dapo_math
 #
-# Run (32 GPU example):
-#   DATA_PATH=data/dapo_math/train.jsonl EVAL_DATA_PATH=data/dapo_math/aime_eval.jsonl \
-#   python -m unirl.train_ar --config-name=ar/qwen3_ppo_4b_base_dapo_sglang num_devices=32
+# Run:
+#   DATA_PATH=data/dapo_math/train.jsonl \
+#   EVAL_DATA_PATH=data/dapo_math/aime_eval.jsonl \
+#   python -m unirl.train_ar --config-name=ar/qwen3_ppo_4b_base_dapo_sglang
 
 num_devices: 32
 batch_size: 64
@@ -18,10 +15,10 @@ num_rollouts: 800
 weight_sync_interval: 1
 eval_interval: 10
 
-# GAE path: advantages computed on workers in PPO.prepare_rollout_track.
+# PPO prepares per-token GAE on each train worker after the micro-batched
+# pre-update critic replay. The driver must not replace it with GRPO scores.
 advantage_mode: gae
-adv_normalization_scope: group
-normalize_adv_by_std: false
+balance_shards: true
 
 logging:
   report_to_wandb: true
@@ -67,22 +64,35 @@ backend:
     adam_beta1: 0.9
     adam_beta2: 0.999
     adam_epsilon: 1.0e-8
-    weight_decay: 0.0
+    weight_decay: 0.01
+  scheduler_cfg:
+    _target_: unirl.train.backend.base.LrSchedulerConfig
+    type: constant
+    warmup_steps: 0
+    total_steps: 1000
 
 rollout:
-  _target_: unirl.rollout.engine.sglang.SGLangRolloutEngine
-  model_path: ${bundle.config.pretrained_model_ckpt_path}
-  concurrency: 16
-  samples_pre_expanded: true
-  chat_template_kwargs:
-    enable_thinking: true
-  engine_kwargs:
-    mem_fraction_static: 0.3
-    skip_server_warmup: true
-    attention_backend: triton
-    disable_cuda_graph: false
-    cuda_graph_max_bs: 16
-    enable_lora: false
+  _target_: unirl.rollout.engine.sglang.engine.SGLangRolloutEngine
+  config:
+    _target_: unirl.rollout.engine.sglang.config.SGLangEngineConfig
+    backend: native
+    pretrained_model_ckpt_path: ${oc.env:QWEN3_PATH,Qwen/Qwen3-4B-Base}
+    tp_size: 1
+    max_new_tokens: 8192
+    temperature: 1.0
+    top_p: 1.0
+    concurrency: 16
+    samples_pre_expanded: true
+    chat_template_kwargs:
+      enable_thinking: true
+    engine_kwargs:
+      rl_on_policy_target: fsdp
+      mem_fraction_static: 0.3
+      skip_server_warmup: true
+      attention_backend: triton
+      disable_cuda_graph: false
+      cuda_graph_max_bs: 16
+      enable_lora: false
 
 reward:
   _target_: unirl.reward.service.RewardService
@@ -111,13 +121,12 @@ algorithm:
     path: unirl.models.qwen3.conditions.Qwen3ARConditions
 
 sync:
-  # Drop train-only value_head weights — SGLang has no critic params to receive them.
   _target_: unirl.distributed.weight_sync.full.tensor.TensorWeightSync
   lora_merged: false
   bucket_size_mb: 64
   flush_cache: true
-  name_remap:
-    value_head.*: null
+  # The rollout model has no critic; never send train-only value-head tensors.
+  name_remap: {"value_head.*": null}
 
 stack:
   _target_: unirl.train.stack.TrainStack
diff --git a/unirl/algorithms/__init__.py b/unirl/algorithms/__init__.py
index d1205ea66..da5886314 100644
--- a/unirl/algorithms/__init__.py
+++ b/unirl/algorithms/__init__.py
@@ -23,10 +23,10 @@
     "FlowMatchSFT",
     "GRPO",
     "GRPOConfig",
-    "PPO",
-    "PPOConfig",
     "GSPO",
     "GSPOConfig",
+    "PPO",
+    "PPOConfig",
     "CPPO",
     "CPPOConfig",
     "DPPO",
diff --git a/unirl/algorithms/base.py b/unirl/algorithms/base.py
index cc0a915e4..8db2286c2 100644
--- a/unirl/algorithms/base.py
+++ b/unirl/algorithms/base.py
@@ -20,6 +20,7 @@
 
 if TYPE_CHECKING:
     from unirl.types.conditions import Condition
+    from unirl.types.sample import Part
     from unirl.types.segments.base import Segment
 
 
@@ -193,27 +194,6 @@ def _grpo_clip_loss(
     return loss_per_elem, metrics
 
 
-def _ppo_clipped_value_loss(
-    *,
-    values: torch.Tensor,
-    old_values: torch.Tensor,
-    returns: torch.Tensor,
-    clip_range: float,
-) -> torch.Tensor:
-    """PPO clipped value loss. Element-wise; reduction is the caller's job.
-
-    Returns per-token ``0.5 * max((V - R)², (V_clipped - R)²)`` with
-    ``V_clipped = V_old + clip(V - V_old, -clip_range, clip_range)``.
-    """
-    values_f = values.float()
-    old_f = old_values.detach().float()
-    returns_f = returns.detach().float()
-    clipped = old_f + (values_f - old_f).clamp(-clip_range, clip_range)
-    sq1 = (values_f - returns_f).square()
-    sq2 = (clipped - returns_f).square()
-    return 0.5 * torch.maximum(sq1, sq2)
-
-
 # ---------------------------------------------------------------------------
 # Reference-policy KL helpers (FlowGRPO / FlowDPPO ``beta`` term)
 # ---------------------------------------------------------------------------
@@ -461,6 +441,15 @@ def prepare_segment(
         """
         return None
 
+    def prepare_part(self, part: "Part") -> "Part":
+        """Optional post-anchor hook over the complete arranged worker shard.
+
+        Runs after per-micro anchor fields have been reassembled and before any
+        optimizer update. PPO uses it to derive GAE from frozen critic values;
+        other algorithms keep the part unchanged.
+        """
+        return part
+
     @abstractmethod
     def compute_loss_and_backward(
         self,
diff --git a/unirl/algorithms/ppo.py b/unirl/algorithms/ppo.py
index e4599c69c..c393a177e 100644
--- a/unirl/algorithms/ppo.py
+++ b/unirl/algorithms/ppo.py
@@ -1,21 +1,15 @@
-"""PPO with GAE for autoregressive ``TextSegment`` training.
-
-Uses per-token GAE advantages from ``segment.token_advantages`` (populated in
-:meth:`prepare_rollout_track`) and clipped value loss against ``segment.returns``.
-Policy clip math is shared with :class:`GRPO` via :func:`_grpo_clip_loss`;
-value clip math via :func:`_ppo_clipped_value_loss` in the same module.
-"""
+"""PPO with per-token GAE and a clipped value objective for AR training."""
 
 from __future__ import annotations
 
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 from typing import Any, Dict, Mapping, Optional, Type
 
 import torch
 
 from unirl.models.types.replay_result import ReplayResult
 from unirl.types.conditions import Condition
-from unirl.types.rollout_resp import RolloutTrack
+from unirl.types.sample import Part
 from unirl.types.segments.text import TextSegment
 
 from .base import (
@@ -23,12 +17,13 @@
     BaseAlgorithmConfig,
     StageAlgorithm,
     _grpo_clip_loss,
-    _ppo_clipped_value_loss,
     _resolve_clip_range_from_schedule,
     rollout_replay_logp_absdiff,
     typed_conditions,
 )
 
+_LOSS_AGG_MODES = frozenset({"token-mean", "seq-mean-token-sum-norm", "seq-mean-token-mean"})
+
 
 @dataclass
 class PPOConfig(BaseAlgorithmConfig):
@@ -45,31 +40,52 @@ class PPOConfig(BaseAlgorithmConfig):
     horizon: int = 8192
 
 
+def _ppo_clipped_value_loss(
+    *,
+    values: torch.Tensor,
+    old_values: torch.Tensor,
+    returns: torch.Tensor,
+    clip_range: float,
+) -> torch.Tensor:
+    """Elementwise ``0.5 * max((V-R)^2, (V_clipped-R)^2)``."""
+    values_f = values.float()
+    old_f = old_values.detach().float()
+    returns_f = returns.detach().float()
+    clipped = old_f + (values_f - old_f).clamp(-clip_range, clip_range)
+    return 0.5 * torch.maximum((values_f - returns_f).square(), (clipped - returns_f).square())
+
+
 def _aggregate_token_loss(
-    loss_per_elem: torch.Tensor,
+    loss_per_token: torch.Tensor,
     *,
+    active: torch.Tensor,
     segment: TextSegment,
     loss_agg_mode: str,
     horizon: int,
 ) -> torch.Tensor:
-    if loss_agg_mode in ("seq-mean-token-sum-norm", "seq-mean-token-mean") and segment.lengths is not None:
-        parts = torch.split(loss_per_elem, segment.lengths.tolist())
+    """Reduce only trainable tokens while preserving the configured sequence weighting."""
+    if loss_agg_mode == "token-mean":
+        return loss_per_token[active].mean()
+    if segment.lengths is None:
+        raise ValueError(f"PPO: loss_agg_mode={loss_agg_mode!r} requires packed segment lengths")
+
+    loss_chunks = torch.split(loss_per_token, segment.lengths.tolist())
+    mask_chunks = torch.split(active, segment.lengths.tolist())
+    sequence_losses = []
+    for losses, mask in zip(loss_chunks, mask_chunks):
+        selected = losses[mask]
         if loss_agg_mode == "seq-mean-token-sum-norm":
-            return torch.stack([p.sum() for p in parts]).mean() / float(horizon)
-        return torch.stack([p.mean() if p.numel() else p.new_zeros(()) for p in parts]).mean()
-    return loss_per_elem.mean()
+            sequence_losses.append(selected.sum() / float(horizon))
+        else:
+            sequence_losses.append(selected.mean() if selected.numel() else losses.new_zeros(()))
+    return torch.stack(sequence_losses).mean()
 
 
 class PPO(StageAlgorithm):
-    """PPO with GAE over an AR ``TextSegment``.
-
-    :meth:`prepare_rollout_track` runs a no-grad critic replay (``return_values=True``),
-    stores ``segment.values`` as the frozen value anchor, then calls
-    :meth:`RolloutTrack.compute_gae_advantages`. Each train micro-batch replays
-    for fresh log-probs and value predictions for the policy and value losses.
-    """
+    """PPO over an AR ``TextSegment`` with a train-side value head."""
 
     supports_multi_update = True
+    anchor_fields = ("values",)
 
     def __init__(
         self,
@@ -94,6 +110,17 @@ def __init__(
             raise ValueError("PPO: either `stage` or `pipeline` must be provided")
         if stage is None:
             stage = getattr(pipeline, stage_attr)
+        if float(clip_range) < 0.0 or float(cliprange_value) < 0.0:
+            raise ValueError("PPO: policy and value clip ranges must be non-negative")
+        if float(vf_coef) < 0.0:
+            raise ValueError("PPO: vf_coef must be non-negative")
+        if not (0.0 <= float(gae_gamma) <= 1.0 and 0.0 <= float(gae_lambda) <= 1.0):
+            raise ValueError("PPO: gae_gamma and gae_lambda must be in [0, 1]")
+        if str(loss_agg_mode) not in _LOSS_AGG_MODES:
+            raise ValueError(f"PPO: unsupported loss_agg_mode={loss_agg_mode!r}")
+        if int(horizon) <= 0:
+            raise ValueError("PPO: horizon must be positive")
+
         self.stage = stage
         self.clip_range = float(clip_range)
         self.clip_range_high = None if clip_range_high is None else float(clip_range_high)
@@ -103,6 +130,7 @@ def __init__(
         self.gae_gamma = float(gae_gamma)
         self.gae_lambda = float(gae_lambda)
         self.loss_agg_mode = str(loss_agg_mode)
+        self.loss_weighting = "token" if self.loss_agg_mode == "token-mean" else "sample"
         self.horizon = int(horizon)
         self.conditions_cls = conditions_cls
         if sampling_temperature is None:
@@ -111,29 +139,38 @@ def __init__(
             sampling_temperature = ARSamplingParams.__dataclass_fields__["temperature"].default
         self.sampling_temperature = float(sampling_temperature)
 
-    def prepare_rollout_track(self, track: RolloutTrack) -> None:
-        """Replay critic values and compute GAE on the worker shard."""
-        if track.rewards is None:
-            raise ValueError("PPO.prepare_rollout_track: track has no rewards")
-        if track.segment is None or not isinstance(track.segment, TextSegment):
-            raise ValueError("PPO.prepare_rollout_track: requires a TextSegment")
-        segment = track.segment
-        if segment.log_probs is None:
-            raise ValueError("PPO.prepare_rollout_track: segment.log_probs is None")
-
-        typed_conds = typed_conditions(track.conditions, self.conditions_cls)
+    def recomputes_anchor(self) -> bool:
+        """Critic anchors must use the exact micro geometry of the train forward."""
+        return True
+
+    def prepare_segment(
+        self,
+        *,
+        conditions: Mapping[str, Condition],
+        segment: TextSegment,
+    ) -> None:
+        """Freeze pre-update critic values for one planned replay micro."""
+        if segment.tokens is None or segment.log_probs is None:
+            raise ValueError("PPO.prepare_segment: segment requires tokens and rollout log_probs")
+        typed_conds = typed_conditions(conditions, self.conditions_cls)
         with torch.no_grad():
-            replay_out = self.stage.replay(
+            replay = self.stage.replay(
                 typed_conds,
                 segment=segment,
                 temperature=self.sampling_temperature,
                 return_values=True,
             )
-        values = _replay_values(replay_out)
-        track.segment = replace(segment, values=values)
-        updated = track.compute_gae_advantages(gamma=self.gae_gamma, gae_lambda=self.gae_lambda)
-        track.segment = updated.segment
-        track.advantages = updated.advantages
+        values = _replay_values(replay)
+        if values.shape != segment.tokens.shape:
+            raise ValueError(
+                f"PPO.prepare_segment: values shape {tuple(values.shape)} "
+                f"!= packed tokens shape {tuple(segment.tokens.shape)}"
+            )
+        segment.values = values.detach()
+
+    def prepare_part(self, part: Part) -> Part:
+        """Compute GAE after all per-micro critic anchors have been reassembled."""
+        return part.compute_gae_advantages(gamma=self.gae_gamma, gae_lambda=self.gae_lambda)
 
     def compute_loss_and_backward(
         self,
@@ -144,31 +181,57 @@ def compute_loss_and_backward(
         training_progress: float,
         loss_scale: float,
     ) -> AlgorithmStepResult:
-        del advantages  # GAE uses segment.token_advantages instead of track-level GRPO scalars.
+        del advantages  # PPO consumes packed GAE, not the compatibility summary on Part.
         if segment.tokens is None or segment.lengths is None or segment.log_probs is None:
             return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False)
-        if int(segment.tokens.shape[0]) == 0:
+        if int(segment.tokens.numel()) == 0:
             return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False)
         if segment.token_advantages is None or segment.returns is None or segment.values is None:
             raise ValueError(
-                "PPO.compute_loss_and_backward: segment requires token_advantages, returns, and values "
-                "(call prepare_rollout_track first)."
+                "PPO.compute_loss_and_backward: token_advantages, returns, and frozen values "
+                "must be prepared before training"
             )
 
         typed_conds = typed_conditions(conditions, self.conditions_cls)
-        replay_out = self.stage.replay(
+        replay = self.stage.replay(
             typed_conds,
             segment=segment,
             temperature=self.sampling_temperature,
             return_values=True,
         )
-        new_logp = _replay_log_probs(replay_out)
-        new_values = _replay_values(replay_out)
+        if not isinstance(replay, ReplayResult):
+            raise TypeError("PPO: replay with return_values=True must return ReplayResult")
+        new_logp = replay.log_probs
+        new_values = _replay_values(replay)
 
         old_logp = segment.log_probs.to(dtype=new_logp.dtype, device=new_logp.device)
         old_values = segment.values.to(dtype=new_values.dtype, device=new_values.device)
         returns = segment.returns.to(dtype=new_values.dtype, device=new_values.device)
-        adv_per_token = segment.token_advantages.detach().to(dtype=new_logp.dtype, device=new_logp.device)
+        token_advantages = segment.token_advantages.to(dtype=new_logp.dtype, device=new_logp.device).detach()
+        expected = new_logp.shape
+        named_tensors = {
+            "new_values": new_values,
+            "old_logp": old_logp,
+            "old_values": old_values,
+            "returns": returns,
+            "token_advantages": token_advantages,
+        }
+        mismatched = {name: tuple(tensor.shape) for name, tensor in named_tensors.items() if tensor.shape != expected}
+        if mismatched:
+            raise ValueError(
+                f"PPO.compute_loss_and_backward: packed tensor shape mismatch; expected {expected}, got {mismatched}"
+            )
+
+        active = torch.ones(expected, dtype=torch.bool, device=new_logp.device)
+        if segment.loss_mask is not None:
+            active = segment.loss_mask.to(device=new_logp.device, dtype=torch.bool)
+            if active.shape != expected:
+                raise ValueError(
+                    f"PPO.compute_loss_and_backward: loss_mask shape {tuple(active.shape)} != {tuple(expected)}"
+                )
+        active_count = int(active.sum().item())
+        if active_count == 0:
+            return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False)
 
         clip_range = _resolve_clip_range_from_schedule(self.clip_range, self.clip_schedule, training_progress)
         clip_high = (
@@ -176,63 +239,72 @@ def compute_loss_and_backward(
             if self.clip_range_high is None
             else _resolve_clip_range_from_schedule(self.clip_range_high, self.clip_schedule, training_progress)
         )
-        policy_per_elem, ratio_metrics = _grpo_clip_loss(
-            new_logp=new_logp,
-            old_logp=old_logp,
-            advantages=adv_per_token,
+        policy_active, ratio_metrics = _grpo_clip_loss(
+            new_logp=new_logp[active],
+            old_logp=old_logp[active],
+            advantages=token_advantages[active],
             clip_range=clip_range,
             clip_range_high=clip_high,
         )
-        value_per_elem = _ppo_clipped_value_loss(
-            values=new_values,
-            old_values=old_values,
-            returns=returns,
+        value_active = _ppo_clipped_value_loss(
+            values=new_values[active],
+            old_values=old_values[active],
+            returns=returns[active],
             clip_range=self.cliprange_value,
         )
-        total_per_elem = policy_per_elem + self.vf_coef * value_per_elem
-        loss = _aggregate_token_loss(
-            total_per_elem,
-            segment=segment,
-            loss_agg_mode=self.loss_agg_mode,
-            horizon=self.horizon,
-        )
-        (loss * loss_scale).backward()
 
+        policy_per_token = torch.zeros_like(new_logp).masked_scatter(active, policy_active)
+        value_per_token = torch.zeros_like(new_values, dtype=torch.float32).masked_scatter(active, value_active)
         policy_loss = _aggregate_token_loss(
-            policy_per_elem,
+            policy_per_token,
+            active=active,
             segment=segment,
             loss_agg_mode=self.loss_agg_mode,
             horizon=self.horizon,
         )
         value_loss = _aggregate_token_loss(
-            value_per_elem,
+            value_per_token,
+            active=active,
             segment=segment,
             loss_agg_mode=self.loss_agg_mode,
             horizon=self.horizon,
         )
+        loss = policy_loss + self.vf_coef * value_loss
+        (loss * loss_scale).backward()
+
+        active_returns = returns[active].float()
+        active_values = new_values[active].float()
+        return_var = active_returns.var(unbiased=False)
+        explained_variance = (
+            1.0 - (active_returns - active_values).var(unbiased=False) / return_var
+            if float(return_var) > 0.0
+            else return_var.new_zeros(())
+        )
+        value_clip_fraction = ((active_values - old_values[active].float()).abs() > self.cliprange_value).float().mean()
         metrics: Dict[str, Any] = {
-            "policy_loss": float(policy_loss.detach().item()),
-            "value_loss": float(value_loss.detach().item()),
+            "policy_loss": float(policy_loss.detach()),
+            "value_loss": float(value_loss.detach()),
+            "value_mean": float(active_values.detach().mean()),
+            "return_mean": float(active_returns.detach().mean()),
+            "explained_variance": float(explained_variance.detach()),
+            "value_clip_fraction": float(value_clip_fraction.detach()),
             "clip_range": float(clip_range),
-            **rollout_replay_logp_absdiff(new_logp, old_logp),
-            **{k: float(v.item()) for k, v in ratio_metrics.items()},
+            "cliprange_value": self.cliprange_value,
+            **rollout_replay_logp_absdiff(new_logp[active], old_logp[active]),
+            **{name: float(value) for name, value in ratio_metrics.items()},
         }
         return AlgorithmStepResult(
-            loss=float(loss.detach().item()),
+            loss=float(loss.detach()),
             metrics=metrics,
-            num_steps_or_tokens=int(new_logp.shape[0]),
+            num_steps_or_tokens=active_count,
             has_backward=True,
         )
 
 
-def _replay_log_probs(replay_out: ReplayResult) -> torch.Tensor:
-    return replay_out.log_probs
-
-
-def _replay_values(replay_out: ReplayResult) -> torch.Tensor:
-    if replay_out.values is None:
-        raise ValueError("PPO: replay with return_values=True must return ReplayResult.values")
-    return replay_out.values
+def _replay_values(replay: Any) -> torch.Tensor:
+    if not isinstance(replay, ReplayResult) or replay.values is None:
+        raise TypeError("PPO: replay with return_values=True must return ReplayResult.values")
+    return replay.values
 
 
 __all__ = ["PPO", "PPOConfig"]
diff --git a/unirl/train/stack/base.py b/unirl/train/stack/base.py
index d035ba5b5..2f4816a2c 100644
--- a/unirl/train/stack/base.py
+++ b/unirl/train/stack/base.py
@@ -192,9 +192,6 @@ def prepare_segment(self, part: Part, *, plans: Plan) -> None:
         if part.segment is None:
             return
         algorithm = self.algorithm
-        prepare_track = getattr(algorithm, "prepare_rollout_track", None)
-        if prepare_track is not None:
-            prepare_track(resp_track)
         if not algorithm.recomputes_anchor():
             algorithm.prepare_segment(conditions=part.conditions, segment=part.segment)
             return
@@ -508,6 +505,7 @@ def train_track(
         profiler = self._train_step_profiler() if profile_scope() == "train" else None
         with profiler.record("train_track") if profiler is not None else nullcontext():
             self.prepare_segment(part, plans=plans)
+            part = self.algorithm.prepare_part(part)
             self.fsdp_backend.model.train()
             result = self._run_updates(part, plans=plans, training_progress=float(training_progress))
         if profiler is not None:
diff --git a/unirl/train_async_ar.py b/unirl/train_async_ar.py
index 511e5e089..37c359ab9 100755
--- a/unirl/train_async_ar.py
+++ b/unirl/train_async_ar.py
@@ -47,6 +47,7 @@ def main(cfg: DictConfig) -> None:
         logging_cfg=cfg.get("logging"),
         adv_normalization_scope=cfg.get("adv_normalization_scope", "group"),
         normalize_adv_by_std=bool(cfg.get("normalize_adv_by_std", True)),
+        advantage_mode=cfg.get("advantage_mode", "grpo"),
         balance_shards=bool(cfg.get("balance_shards", False)),
         eval_interval=int(cfg.get("eval_interval", 0)),
         eval_num_prompts=int(cfg.get("eval_num_prompts", -1)),
diff --git a/unirl/trainer/ar.py b/unirl/trainer/ar.py
index 93dafcca6..9af9eeed4 100644
--- a/unirl/trainer/ar.py
+++ b/unirl/trainer/ar.py
@@ -84,8 +84,9 @@ def __init__(
         # group std. False = mean-center only (reward - group_mean), NO std division —
         # removes the difficulty bias that over-amplifies low-std (hard) prompts.
         self.normalize_adv_by_std = normalize_adv_by_std
-        # "grpo" (default) or "gae" (PPO path — GAE runs in PPO.prepare_rollout_track on workers).
-        self.advantage_mode = str(advantage_mode)
+        self.advantage_mode = str(advantage_mode).strip().lower()
+        if self.advantage_mode not in ("grpo", "gae"):
+            raise ValueError(f"ARTrainer: advantage_mode must be 'grpo' or 'gae', got {advantage_mode!r}")
         # verl trainer.balance_batch parity: driver-side reorder of the rollout
         # batch so each DP shard receives a similar total-token workload. FSDP
         # collectives sync all ranks every micro, so a step runs at the SLOWEST
diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py
index 5199edb8f..936e43859 100644
--- a/unirl/trainer/async_ar.py
+++ b/unirl/trainer/async_ar.py
@@ -96,6 +96,7 @@ def __init__(
         logging_cfg: Optional[DictConfig] = None,
         adv_normalization_scope: str = "group",
         normalize_adv_by_std: bool = True,
+        advantage_mode: str = "grpo",
         balance_shards: bool = False,
         eval_interval: int = 0,
         eval_num_prompts: int = -1,
@@ -122,6 +123,9 @@ def __init__(
         self.batch_size = batch_size
         self.adv_normalization_scope = adv_normalization_scope
         self.normalize_adv_by_std = normalize_adv_by_std
+        self.advantage_mode = str(advantage_mode).strip().lower()
+        if self.advantage_mode not in ("grpo", "gae"):
+            raise ValueError(f"AsyncARTrainer: advantage_mode must be 'grpo' or 'gae', got {advantage_mode!r}")
         self.balance_shards = bool(balance_shards)
         self.eval_interval = int(eval_interval)
         _num = int(eval_num_prompts)