diff --git a/tests/test_rl_helpers.py b/tests/test_rl_helpers.py new file mode 100644 index 0000000..6e8381e --- /dev/null +++ b/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)) diff --git a/xorl_client/__init__.py b/xorl_client/__init__.py index 0b067fd..b1e5c57 100644 --- a/xorl_client/__init__.py +++ b/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", diff --git a/xorl_client/client/service_client.py b/xorl_client/client/service_client.py index 60ad39c..77b7747 100644 --- a/xorl_client/client/service_client.py +++ b/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) diff --git a/xorl_client/rl.py b/xorl_client/rl.py new file mode 100644 index 0000000..8a1bb52 --- /dev/null +++ b/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 diff --git a/xorl_client/types/datum.py b/xorl_client/types/datum.py index 6b484cb..a229dc6 100644 --- a/xorl_client/types/datum.py +++ b/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", } diff --git a/xorl_client/types/forward_backward_output.py b/xorl_client/types/forward_backward_output.py index ada5a26..2bdf515 100644 --- a/xorl_client/types/forward_backward_output.py +++ b/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 diff --git a/xorl_client/types/lora_config.py b/xorl_client/types/lora_config.py index 02979c8..87a11a1 100644 --- a/xorl_client/types/lora_config.py +++ b/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