Skip to content

fix(reward): trainer adapters live in openadapt_evals.reward; flow keeps the worker and the client #453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
abrichr merged 2 commits into from
Sep 2, 2026
Merged

fix(reward): trainer adapters live in openadapt_evals.reward; flow keeps the worker and the client #453

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7292837
fix(reward): trainer adapters live in openadapt_evals.reward; flow ke…
abrichr Sep 2, 2026
3d12e7b
docs(reward): the OpenAI graders reference lists six types, not five
abrichr Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -468,8 +468,10 @@ openadapt-flow serve-reward --seed-mockmed --port 8788

A reward receipt isn't an Execute Seal. A model rollout isn't a qualified
program, so it never gets one, and the receipt never says Flow governed the
policy. Adapters for TRL's `GRPOTrainer` and verl's reward manager are in
`openadapt_flow.reward.callables`. See
policy. The adapters for TRL's `GRPOTrainer` and verl's reward manager live in
`openadapt_evals.reward` (`pip install 'openadapt-evals>=0.96.0'`), because
TRL trains a `None` reward as 0.0 and the evals adapters drop an unscored
episode instead. This package keeps the worker and the HTTP client. See
[docs/REWARD_WORKER.md](docs/REWARD_WORKER.md).

## Development
Expand Down
55 changes: 35 additions & 20 deletions docs/REWARD_WORKER.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,10 @@ shape `openadapt_evals.reward.receipts.EpisodeDescriptor` sends:
`reward_contract_digest`, and optional `task_id`, `environment_id`, and
`metadata`. The digest must be the contract this worker serves, or the
episode is refused. The trainer never gets a credential for the system of
record. `openadapt_flow.reward.callables` carries the adapters for TRL's
`GRPOTrainer` and verl's reward manager, and `HttpRewardClient` for the trip
between the two machines.
record. The trainer-side adapters live in `openadapt_evals.reward`
(`pip install 'openadapt-evals>=0.96.0'`), not in this package.
`openadapt_flow.reward.callables` keeps only `HttpRewardClient`, the payload
builder, and the receipt's scalar, for the trip between the two machines.

The oracle still has to know which record to read. That identity comes from
one of three places: `metadata.oracle_identity` on the descriptor, an
Expand Down Expand Up @@ -183,23 +184,37 @@ sees it.

## Trainer adapters

`trl_reward_function(scorer, policy_checkpoint_id=..., reward_contract_digest=...)`
returns a function
with the signature TRL's `GRPOTrainer` expects for `reward_funcs`:
`(prompts, completions, completion_ids, trainer_state, **kwargs) ->
list[float | None]`. The dataset carries `episode_id` and `oracle_identity`
columns; the policy update is `trainer_state.global_step`. An unscored
episode returns `None`, which TRL documents as "this reward function does
not apply to this sample" and excludes.

`verl_compute_score(scorer, policy_checkpoint_id=..., reward_contract_digest=...)`
returns a
`compute_score(data_source, solution_str, ground_truth, extra_info)` for
`custom_reward_function.path`. verl has no None sentinel, so an unscored
episode returns `{"score": nan, "openadapt_unscored": true, ...}`. NaN on
purpose: a group that keeps it produces a NaN loss instead of a quiet 0.
`drop_unscored(rewards, *aligned)` and `scored_groups(groups)` remove those
samples before the advantage is computed.
The adapters for TRL's `GRPOTrainer` and verl's reward manager live in
`openadapt_evals.reward`: `CertifiedRewardFunction` (TRL) and
`CertifiedRewardManager` (verl). Install them on the trainer node:

```bash
pip install 'openadapt-evals>=0.96.0'
```

Both call `openadapt_types.score`, read the receipt's own fields, and refuse
the combinations a trainer must never accept: an unscored episode is removed
from its GRPO group, a `development_only` receipt is never labelled certified
(and the only certificate scope in use today is synthetic), and in
`require_certified` mode an expired certificate stops the run. The wiring,
with a code sample for each trainer, is in the evals package's
`docs/reward/README.md` and on
[docs.openadapt.ai](https://docs.openadapt.ai/commercial/seal-reward/).

This package offers no trainer-facing reward function, and that is on
purpose. TRL lets a reward function return `None` for a sample, but
`GRPOTrainer` turns that `None` into NaN, combines the per-function rewards
with `nansum`, and takes the group mean over the result. With one reward
function the `None` row trains as 0.0, which is what the contract forbids for
`reconciliation_required` and `failed_platform`. verl's per-sample
`compute_score` hook must return a number, so it cannot drop a sample either.
The evals adapters drop an unscored episode the one way a per-completion
scalar allows: the episode gets the mean reward of its scored group-mates, so
its advantage is exactly zero and the scored mean is unchanged.

The dependency runs one way. openadapt-evals depends on openadapt-flow, so
flow cannot import the adapters, and a second copy here would drift from the
first.

## Scope

Expand Down
268 changes: 61 additions & 207 deletions openadapt_flow/reward/callables.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,84 +1,39 @@
"""Reward-function adapters for TRL GRPO and verl.

Both adapters turn a trainer's per-sample data into an episode descriptor,
ask a scorer for a signed receipt, and hand back the scalar. An UNSCORED
episode never becomes 0.0:

* TRL: the adapter returns ``None`` for that sample. TRL documents ``None``
as "this reward function does not apply to this sample" and excludes it
from the reward calculation.
* verl: the reward manager has no such sentinel, so the adapter returns
:data:`UNSCORED_REWARD` (``nan``) together with ``"openadapt_unscored":
True`` in the dict. A NaN poisons a group's advantage instead of
silently training on 0.0. :func:`drop_unscored` removes those samples
from a group before the loss is computed; wire it in, or filter on the
``openadapt_unscored`` key.

TRL contract (read 2026-09-01):
https://huggingface.co/docs/trl/main/en/grpo_trainer#using-a-custom-reward-function
reward_func(prompts, completions, completion_ids, trainer_state, **kwargs)
-> list[float | None]
Every dataset column except ``prompt`` arrives in ``**kwargs`` as a list
aligned with ``completions``. ``trainer_state.global_step`` is the policy
update counter the certificate expiry is denominated in.

verl contract (read 2026-09-01):
https://verl.readthedocs.io/en/latest/preparation/reward_function.html
https://github.com/volcengine/verl/blob/main/verl/workers/reward_manager/naive.py
compute_score(data_source, solution_str, ground_truth, extra_info=None)
-> float | dict (a dict must carry "score"; other keys are logged)
Configured through ``custom_reward_function.path`` and
``custom_reward_function.name``.
"""The trainer-side client for a reward worker. Not a trainer adapter.

This module carries the pieces a trainer node needs to talk to a reward
worker over HTTP: the wire payload for one episode, the ``POST /v1/rewards``
client, and the receipt's scalar. It offers no ``reward_funcs`` entry for
TRL and no ``compute_score`` for verl, on purpose.

The trainer adapters live in ``openadapt_evals.reward``:
``CertifiedRewardFunction`` for TRL's ``GRPOTrainer`` and
``CertifiedRewardManager`` for verl's reward manager. They drop an unscored
episode by group-mean fill: the episode receives the mean reward of its
scored group-mates, so its GRPO advantage is exactly zero and the scored
mean is unchanged.

A ``None`` or NaN reward is not a drop in TRL. ``GRPOTrainer`` turns a
``None`` into NaN, combines the per-function rewards with ``nansum``, and
takes the group mean over the result, so with one reward function an
unscored episode trains as 0.0. That is the outcome the reward contract
forbids for ``reconciliation_required`` and ``failed_platform``. verl's
per-sample ``compute_score`` hook has no sentinel at all; whatever it
returns lands in the reward tensor.

openadapt-evals depends on openadapt-flow, so this package cannot import
the adapters. Install them on the trainer node with
``pip install 'openadapt-evals>=0.96.0'``.
"""

from __future__ import annotations

import math
from typing import (
Any,
Callable,
Iterable,
Mapping,
Optional,
Protocol,
Sequence,
TypeVar,
)
from typing import Any, Mapping, Optional, Protocol

from openadapt_types.reward import RewardEvidenceReceiptV1

from openadapt_flow.reward.models import EpisodeDescriptorV1

#: The verl-side sentinel for an episode the oracle could not score.
#: ``nan`` on purpose: a trainer that forgets to drop it sees a NaN loss,
#: not a quiet 0.0 that teaches the policy that uncertainty is failure.
UNSCORED_REWARD: float = float("nan")

T = TypeVar("T")


def is_unscored(value: Any) -> bool:
"""True for ``None`` (TRL) and for the NaN sentinel (verl)."""

if value is None:
return True
return isinstance(value, float) and math.isnan(value)


def drop_unscored(
rewards: Sequence[Any], *aligned: Sequence[T]
) -> tuple[list[float], list[list[T]]]:
"""Filter a group down to the scored samples.

Returns the kept rewards and, for each aligned sequence passed, the
kept items in the same positions. A group that loses every sample
comes back empty; the trainer must skip that group.
"""

keep = [index for index, value in enumerate(rewards) if not is_unscored(value)]
kept_rewards = [float(rewards[index]) for index in keep]
kept_aligned = [[seq[index] for index in keep] for seq in aligned]
return kept_rewards, kept_aligned
REWARDS_ROUTE = "/v1/rewards"


class RewardScorer(Protocol):
Expand All @@ -92,22 +47,38 @@ def score_episode(self, payload: Mapping[str, Any]) -> dict[str, Any]: ...


class HttpRewardClient:
"""Thin client for ``POST /v1/rewards`` on a reward worker."""
"""Thin client for ``POST /v1/rewards`` on a reward worker.

``transport`` is an optional ``httpx`` transport, for tests
(``httpx.MockTransport``).
"""

def __init__(self, base_url: str, token: str, *, timeout_s: float = 30.0) -> None:
def __init__(
self,
base_url: str,
token: str,
*,
timeout_s: float = 30.0,
transport: Any = None,
) -> None:
self.base_url = base_url.rstrip("/")
self.token = token
self.timeout_s = timeout_s
self.transport = transport

@property
def url(self) -> str:
return f"{self.base_url}{REWARDS_ROUTE}"

def score_episode(self, payload: Mapping[str, Any]) -> dict[str, Any]:
import httpx

response = httpx.post(
f"{self.base_url}/v1/rewards",
json=dict(payload),
headers={"Authorization": f"Bearer {self.token}"},
timeout=self.timeout_s,
)
with httpx.Client(timeout=self.timeout_s, transport=self.transport) as client:
response = client.post(
self.url,
json=dict(payload),
headers={"Authorization": f"Bearer {self.token}"},
)
if response.status_code == 409:
raise RuntimeError(
f"episode already scored: {response.json().get('detail')}"
Expand All @@ -117,7 +88,12 @@ def score_episode(self, payload: Mapping[str, Any]) -> dict[str, Any]:


def scalar_of(envelope: Mapping[str, Any]) -> Optional[float]:
"""The receipt's scalar, or ``None`` when the episode is unscored."""
"""The receipt's scalar, or ``None`` when the episode is unscored.

``None`` here is a client-side value, not a trainer reward. A trainer
must never hand it to TRL or verl as the sample's reward; see the module
docstring.
"""

receipt = RewardEvidenceReceiptV1.model_validate(envelope["receipt"])
return receipt.scalar_reward
Expand All @@ -134,8 +110,10 @@ def episode_from_columns(
) -> dict[str, Any]:
"""Build the wire payload for one episode, in the trainer client's shape.

``oracle_identity`` may be ``None`` when the environment registered the
identity with ``RewardWorker.begin_episode`` before the rollout.
The keys are the ones ``openadapt_evals.reward.receipts.EpisodeDescriptor``
sends, plus the descriptor model's ``schema_version``; the worker accepts
both. ``oracle_identity`` may be ``None`` when the environment registered
the identity with ``RewardWorker.begin_episode`` before the rollout.
"""

metadata: dict[str, Any] = {"runtime_signal": runtime_signal}
Expand All @@ -150,127 +128,3 @@ def episode_from_columns(
reward_contract_digest=reward_contract_digest,
metadata=metadata,
).model_dump(mode="json", exclude_none=True)


# -- TRL ----------------------------------------------------------------------


def trl_reward_function(
scorer: RewardScorer,
*,
policy_checkpoint_id: str,
reward_contract_digest: str,
episode_column: str = "episode_id",
identity_column: str = "oracle_identity",
signal_column: str = "runtime_signal",
) -> Callable[..., list[Optional[float]]]:
"""Build a ``reward_funcs`` entry for ``trl.GRPOTrainer``.

The dataset carries one row per episode with ``episode_column`` (the
episode id the environment ran under), optionally ``identity_column``
(a dict of the oracle identity keys; omit it when the environment
registers the identity with ``begin_episode``), and optionally
``signal_column``. The policy update comes from
``trainer_state.global_step``. Each call returns one float per
completion, or ``None`` for an unscored episode.
"""

def openadapt_verified_effect_reward(
prompts: Sequence[Any],
completions: Sequence[Any],
completion_ids: Optional[Sequence[Any]] = None,
trainer_state: Any = None,
**kwargs: Any,
) -> list[Optional[float]]:
del prompts, completion_ids
episodes = kwargs.get(episode_column)
if episodes is None:
raise KeyError(f"dataset must carry {episode_column!r}")
identities = kwargs.get(identity_column) or [None] * len(completions)
signals = kwargs.get(signal_column) or ["completed"] * len(completions)
policy_update = int(getattr(trainer_state, "global_step", 0) or 0)
rewards: list[Optional[float]] = []
for episode_id, identity, signal in zip(episodes, identities, signals):
envelope = scorer.score_episode(
episode_from_columns(
episode_id=episode_id,
policy_checkpoint_id=policy_checkpoint_id,
policy_update=policy_update,
reward_contract_digest=reward_contract_digest,
oracle_identity=identity,
runtime_signal=signal,
)
)
rewards.append(scalar_of(envelope))
return rewards

return openadapt_verified_effect_reward


# -- verl ---------------------------------------------------------------------


def verl_compute_score(
scorer: RewardScorer,
*,
policy_checkpoint_id: str,
reward_contract_digest: str,
episode_key: str = "openadapt_episode",
) -> Callable[..., dict[str, Any]]:
"""Build a verl ``compute_score`` over a reward worker.

``extra_info[episode_key]`` must hold ``episode_id`` and
``policy_update``, plus ``oracle_identity`` unless the environment
registered it with ``begin_episode``, and optionally ``runtime_signal``. The returned dict
carries ``score`` (the scalar, or :data:`UNSCORED_REWARD`), the receipt
id, and the flags a filter needs. verl stores every extra key in the
batch's non-tensor data, so ``openadapt_unscored`` survives to the point
where :func:`drop_unscored` can act on it.
"""

def compute_score(
data_source: Any,
solution_str: Any,
ground_truth: Any,
extra_info: Optional[Mapping[str, Any]] = None,
) -> dict[str, Any]:
del data_source, solution_str, ground_truth
info = dict(extra_info or {})
episode = info.get(episode_key)
if not isinstance(episode, Mapping):
raise KeyError(f"extra_info must carry {episode_key!r}")
envelope = scorer.score_episode(
episode_from_columns(
episode_id=str(episode["episode_id"]),
policy_checkpoint_id=policy_checkpoint_id,
policy_update=int(episode.get("policy_update", 0)),
reward_contract_digest=reward_contract_digest,
oracle_identity=episode.get("oracle_identity"),
runtime_signal=str(episode.get("runtime_signal", "completed")),
)
)
receipt = envelope["receipt"]
scalar = scalar_of(envelope)
return {
"score": UNSCORED_REWARD if scalar is None else float(scalar),
"openadapt_unscored": scalar is None,
"openadapt_reward_outcome": receipt["reward_outcome"],
"openadapt_receipt_id": receipt["receipt_id"],
"openadapt_certified": bool(receipt["certified"]),
"openadapt_development_only": bool(receipt["development_only"]),
}

return compute_score


def scored_groups(
groups: Iterable[Sequence[Mapping[str, Any]]],
) -> list[list[Mapping[str, Any]]]:
"""Drop unscored samples from each verl group of ``compute_score`` dicts."""

kept: list[list[Mapping[str, Any]]] = []
for group in groups:
rewards = [item["score"] for item in group]
_kept_rewards, (kept_items,) = drop_unscored(rewards, list(group))
kept.append(kept_items)
return kept
Loading