Add value-model (critic) support: state_values, GAE/EV helpers, frozen modules

tests/test_rl_helpers.py

@@ -0,0 +1,69 @@
+"""Value-model client helpers: datum dtypes, state_values output, GAE, EV."""
+
+import math
+
+import pytest
+
+from xorl_client import compute_skip_observation_gae, explained_variance
+from xorl_client.types.datum import Datum
+from xorl_client.types.forward_backward_output import LossFnOutput
+from xorl_client.types.model_input import ModelInput
+from xorl_client.types.tensor_data import TensorData
+
+
+def test_datum_infers_float32_for_value_fields():
+    datum = Datum(
+        model_input=ModelInput.from_ints([1, 2, 3]),
+        loss_fn_inputs={
+            "target_tokens": [2, 3, 4],
+            "weights": [1.0, 1.0, 0.0],
+            "returns": [0.5, 0.25, 0.0],
+            "old_values": [0.4, 0.2, 0.0],
+        },
+    )
+    converted = datum.loss_fn_inputs
+    assert converted["returns"].dtype == "float32"
+    assert converted["old_values"].dtype == "float32"
+    assert converted["target_tokens"].dtype == "int64"
+
+
+def test_loss_fn_output_state_values_roundtrip():
+    payload = {
+        "state_values": {"data": [0.1, 0.2, 0.3], "dtype": "float32", "shape": [3]},
+        "elementwise_loss": {"data": [1.0, 1.0, 1.0], "dtype": "float32", "shape": [3]},
+    }
+    output = LossFnOutput.from_dict(payload)
+    assert isinstance(output.state_values, TensorData)
+    assert output.state_values.data == [0.1, 0.2, 0.3]
+    # Dict-like tinker-compat surface.
+    assert "state_values" in output
+    assert output["state_values"].data == [0.1, 0.2, 0.3]
+    assert "state_values" in output.keys()
+    assert output.logprobs is None
+    assert "logprobs" not in output
+    assert output.to_dict()["state_values"]["data"] == [0.1, 0.2, 0.3]
+
+
+def test_skip_observation_gae_skips_observations():
+    rewards = [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
+    values = [0.5, 9.9, -9.9, 0.4, 9.9, 0.2]
+    action_mask = [1, 0, 0, 1, 0, 1]
+    adv, ret = compute_skip_observation_gae(rewards, values, action_mask, gamma=1.0, lam=1.0)
+    # Action-token chain: bootstraps action-to-action across observation gaps.
+    assert adv[5] == pytest.approx(1.0 - 0.2)
+    assert adv[3] == pytest.approx((0.2 - 0.4) + adv[5])
+    assert adv[0] == pytest.approx((0.4 - 0.5) + adv[3])
+    for t in (1, 2, 4):
+        assert adv[t] == 0.0 and ret[t] == 0.0
+    for t in (0, 3, 5):
+        assert ret[t] == pytest.approx(adv[t] + values[t])
+
+
+def test_explained_variance_from_value_loss_metrics():
+    returns = [0.0, 1.0, 2.0, 3.0]
+    mean_r = 1.5
+    sq_mean = sum(r * r for r in returns) / 4
+    var_r = sq_mean - mean_r**2
+    assert explained_variance(0.0, mean_r, sq_mean) == pytest.approx(1.0)
+    assert explained_variance(var_r, mean_r, sq_mean) == pytest.approx(0.0)
+    assert math.isnan(explained_variance(0.1, 1.0, 1.0))

xorl_client/__init__.py

@@ -9,6 +9,7 @@
 
 # Import types module
 from . import types
+from .rl import compute_skip_observation_gae, explained_variance
 
 # Import clients
 from .client.service_client import ServiceClient
@@ -62,6 +63,9 @@
 )
 
 __all__ = [
+    # RL / value-model helpers
+    "compute_skip_observation_gae",
+    "explained_variance",
     # Core clients
     "ServiceClient",
     "TrainingClient",

xorl_client/client/service_client.py

@@ -102,6 +102,7 @@ def _create_lora_training_client_submit(
         dropout: float = 0.0,
         target_modules: Optional[list[str]] = None,
         model_id: Optional[str] = None,
+        frozen_module_patterns: Optional[list[str]] = None,
     ) -> Future["TrainingClient"]:
         """Helper function that submits the create_lora_training_client request (two-phase pattern)."""
         import time
@@ -126,6 +127,7 @@ def _create_lora_training_client_submit(
             alpha=alpha,
             dropout=dropout,
             target_modules=target_modules,
+            frozen_module_patterns=frozen_module_patterns,
         )
 
         # Send create model request to server
@@ -291,6 +293,7 @@ def create_lora_training_client(
         dropout: float = 0.0,
         target_modules: Optional[list[str]] = None,
         model_id: Optional[str] = None,
+        frozen_module_patterns: Optional[list[str]] = None,
     ) -> "TrainingClient":
         """Create a LoRA training client.
 
@@ -323,7 +326,7 @@ def create_lora_training_client(
             ... )
         """
         return self._create_lora_training_client_submit(
-            base_model, rank, alpha, dropout, target_modules, model_id
+            base_model, rank, alpha, dropout, target_modules, model_id, frozen_module_patterns
         ).result()
 
     async def create_lora_training_client_async(
@@ -334,10 +337,11 @@ async def create_lora_training_client_async(
         dropout: float = 0.0,
         target_modules: Optional[list[str]] = None,
         model_id: Optional[str] = None,
+        frozen_module_patterns: Optional[list[str]] = None,
     ) -> "TrainingClient":
         """Async version of create_lora_training_client."""
         future = self._create_lora_training_client_submit(
-            base_model, rank, alpha, dropout, target_modules, model_id
+            base_model, rank, alpha, dropout, target_modules, model_id, frozen_module_patterns
         )
         return await asyncio.wrap_future(future)
 

xorl_client/rl.py

@@ -0,0 +1,106 @@
+"""Skip-observation token-level GAE (SAO, arXiv:2607.07508, Eq. 4-5).
+
+Agentic trajectories interleave model actions with environment observations
+the model did not generate. Standard GAE bootstraps across every adjacent
+token, which propagates critic noise through observation tokens. The
+skip-observation estimator chains the Bellman recursion across *action tokens
+only*: the TD target of an action token bootstraps from the value of the next
+action token, bridging any observation gap in between.
+
+Pure Python (no torch/numpy), mirroring ``xorl.rl.advantages`` in the server
+repo. Sequences follow the server's target-aligned
+convention: index t describes the t-th target token, matching the per-token
+values returned by the ``value_prediction`` loss and the ``advantages`` /
+``returns`` fields of a training datum.
+
+Note: the server's packer masks tokens where ``advantages == 0.0``; that is
+the desired behavior for non-action tokens (this module emits exactly 0.0
+there), but a true action-token advantage of exactly 0.0 would also be
+masked. Nudge such values by a tiny epsilon if that matters for your reward
+scale.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Sequence, Tuple
+
+
+def compute_skip_observation_gae(
+    rewards: Sequence[float],
+    values: Sequence[float],
+    action_mask: Optional[Sequence[int]] = None,
+    gamma: float = 1.0,
+    lam: float = 1.0,
+    bootstrap_value: float = 0.0,
+) -> Tuple[List[float], List[float]]:
+    """Compute per-token advantages and value targets across action tokens.
+
+    Args:
+        rewards: Per-token rewards, length T (typically zero everywhere except
+            the final action token of the trajectory).
+        values: Per-token value predictions V(s_t), length T (from a
+            ``forward`` call with ``loss_fn="value_prediction"``: the
+            ``state_values`` field of each LossFnOutput).
+        action_mask: Per-token 0/1 mask, length T; 1 marks model-generated
+            (action) tokens. ``None`` treats every token as an action token,
+            which reduces to standard token-level GAE.
+        gamma: Discount factor.
+        lam: GAE lambda. For length-adaptive GAE (VAPO), pass
+            ``1 - 1 / (alpha * num_action_tokens)``.
+        bootstrap_value: Value bootstrapped after the last action token
+            (0.0 for terminated trajectories).
+
+    Returns:
+        ``(advantages, returns)`` lists of length T. Non-action tokens carry
+        0.0 in both; action tokens carry the skip-observation GAE advantage
+        and the corresponding value target ``R_t = A_t + V(s_t)``.
+    """
+    n = len(rewards)
+    if len(values) != n:
+        raise ValueError(f"rewards ({n}) and values ({len(values)}) must have the same length")
+    if action_mask is None:
+        action_indices = list(range(n))
+    else:
+        if len(action_mask) != n:
+            raise ValueError(f"action_mask ({len(action_mask)}) must match rewards ({n})")
+        action_indices = [t for t in range(n) if action_mask[t]]
+
+    advantages = [0.0] * n
+    returns = [0.0] * n
+
+    next_advantage = 0.0
+    next_value = float(bootstrap_value)
+    for t in reversed(action_indices):
+        delta = float(rewards[t]) + gamma * next_value - float(values[t])
+        advantage = delta + gamma * lam * next_advantage
+        advantages[t] = advantage
+        returns[t] = advantage + float(values[t])
+        next_advantage = advantage
+        next_value = float(values[t])
+
+    return advantages, returns
+
+
+def explained_variance(
+    value_error_sq_mean: float,
+    return_mean: float,
+    return_sq_mean: float,
+) -> float:
+    """Critic explained variance from the ``value_loss`` moment metrics.
+
+    ``EV = 1 - E[(R - V)^2] / Var(R)`` with ``Var(R) = E[R^2] - E[R]^2``.
+    The three inputs are exactly the (globally normalized) ``value_error_sq_mean``,
+    ``return_mean``, and ``return_sq_mean`` metrics a ``value_loss``
+    forward_backward reports, so EV composes correctly across micro-batches
+    and ranks. EV is the paper's key critic-health diagnostic (SAO Fig. 4a):
+    it should climb toward 1.0 as the critic converges; near 0 the critic is
+    no better than predicting the mean return.
+
+    Returns NaN when the return distribution is (numerically) constant, where
+    explained variance is undefined.
+    """
+    return_variance = return_sq_mean - return_mean * return_mean
+    if not math.isfinite(return_variance) or return_variance <= 1e-12:
+        return float("nan")
+    return 1.0 - value_error_sq_mean / return_variance

xorl_client/types/datum.py

@@ -19,6 +19,10 @@
     "logprobs": "float32",
     "clip_low_threshold": "float32",
     "clip_high_threshold": "float32",
+    # Value-model (critic) training fields: per-token value targets and
+    # pre-update value predictions (PPO-style value clipping).
+    "returns": "float32",
+    "old_values": "float32",
 }
 
 

xorl_client/types/forward_backward_output.py

@@ -46,6 +46,7 @@ class LossFnOutput:
     loss: Optional[float] = None
     logprobs: Optional[TensorData] = None
     elementwise_loss: Optional[TensorData] = None
+    state_values: Optional[TensorData] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary, excluding None values."""
@@ -56,6 +57,8 @@ def to_dict(self) -> Dict[str, Any]:
             result["logprobs"] = self.logprobs.to_dict() if isinstance(self.logprobs, TensorData) else self.logprobs
         if self.elementwise_loss is not None:
             result["elementwise_loss"] = self.elementwise_loss.to_dict() if isinstance(self.elementwise_loss, TensorData) else self.elementwise_loss
+        if self.state_values is not None:
+            result["state_values"] = self.state_values.to_dict() if isinstance(self.state_values, TensorData) else self.state_values
         return result
 
     @classmethod
@@ -65,6 +68,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "LossFnOutput":
             loss=data.get("loss"),
             logprobs=_to_tensor_data(data.get("logprobs")),
             elementwise_loss=_to_tensor_data(data.get("elementwise_loss")),
+            state_values=_to_tensor_data(data.get("state_values")),
         )
 
     # Dict-like access methods for tinker compatibility
@@ -76,6 +80,8 @@ def __getitem__(self, key: str) -> Any:
             return self.logprobs
         elif key == "elementwise_loss":
             return self.elementwise_loss
+        elif key == "state_values":
+            return self.state_values
         else:
             raise KeyError(key)
 
@@ -87,6 +93,8 @@ def __contains__(self, key: str) -> bool:
             return self.logprobs is not None
         elif key == "elementwise_loss":
             return self.elementwise_loss is not None
+        elif key == "state_values":
+            return self.state_values is not None
         return False
 
     def get(self, key: str, default: Any = None) -> Any:
@@ -106,6 +114,8 @@ def keys(self) -> List[str]:
             result.append("logprobs")
         if self.elementwise_loss is not None:
             result.append("elementwise_loss")
+        if self.state_values is not None:
+            result.append("state_values")
         return result
 
     def values(self) -> List[Any]:
@@ -117,6 +127,8 @@ def values(self) -> List[Any]:
             result.append(self.logprobs)
         if self.elementwise_loss is not None:
             result.append(self.elementwise_loss)
+        if self.state_values is not None:
+            result.append(self.state_values)
         return result
 
     def items(self) -> List[tuple]:
@@ -128,6 +140,8 @@ def items(self) -> List[tuple]:
             result.append(("logprobs", self.logprobs))
         if self.elementwise_loss is not None:
             result.append(("elementwise_loss", self.elementwise_loss))
+        if self.state_values is not None:
+            result.append(("state_values", self.state_values))
         return result
 
 

xorl_client/types/lora_config.py

@@ -16,6 +16,7 @@ class LoraConfig:
     alpha: Optional[int] = None
     dropout: float = 0.0
     target_modules: Optional[List[str]] = None
+    frozen_module_patterns: Optional[List[str]] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary for JSON serialization."""
@@ -24,4 +25,6 @@ def to_dict(self) -> Dict[str, Any]:
             result["alpha"] = self.alpha
         if self.target_modules is not None:
             result["target_modules"] = self.target_modules
+        if self.frozen_module_patterns is not None:
+            result["frozen_module_patterns"] = self.frozen_module_patterns
         return result