From ccdf0a7acec47e85bb58f7067db7c5ba9819a8e5 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Sat, 29 Aug 2026 06:08:50 -0700 Subject: [PATCH] Add explicit rewarded OPD profile --- CHANGELOG.md | 15 ++ PYPI.md | 2 + README.md | 2 + README.zh-CN.md | 2 + docs/compatibility.md | 19 +- docs/for-verl-users.md | 17 +- docs/generated/upstream-support-matrix.json | 8 + docs/profiles/index.md | 20 +- docs/verl-opd-runtime.md | 13 +- scripts/check_docs_visual.py | 1 + src/miniverl/bridge/export.py | 50 +++- src/miniverl/bridge/opd_pg_v08.py | 122 ++++++++- src/miniverl/bridge/opd_runtime.py | 38 ++- src/miniverl/bridge/profiles.py | 76 +++++- src/miniverl/cache/store.py | 16 ++ src/miniverl/config/__init__.py | 4 + src/miniverl/config/models.py | 94 ++++++- src/miniverl/data/verl_parquet.py | 38 ++- src/miniverl/losses/chunked.py | 59 +++-- src/miniverl/losses/verl_pg.py | 97 +++++++ src/miniverl/rewards/__init__.py | 35 +++ src/miniverl/rewards/composer.py | 87 +++++++ src/miniverl/rewards/models.py | 136 ++++++++++ src/miniverl/rewards/providers.py | 184 ++++++++++++++ src/miniverl/schemas/cache.py | 6 +- src/miniverl/teachers/local.py | 20 +- src/miniverl/training/checkpoint.py | 4 + src/miniverl/training/trainer.py | 173 ++++++++++++- tests/unit/test_cache.py | 37 +++ tests/unit/test_packaging.py | 1 + tests/unit/test_profile_registry.py | 1 + tests/unit/test_rewarded_pg.py | 268 ++++++++++++++++++++ tests/unit/test_rewards.py | 148 +++++++++++ tests/unit/test_verl_opd_export.py | 33 +++ tests/unit/test_verl_parquet_source.py | 37 +++ tests/unit/test_verl_pg_profile.py | 50 ++++ 36 files changed, 1820 insertions(+), 93 deletions(-) create mode 100644 src/miniverl/rewards/__init__.py create mode 100644 src/miniverl/rewards/composer.py create mode 100644 src/miniverl/rewards/models.py create mode 100644 src/miniverl/rewards/providers.py create mode 100644 tests/unit/test_rewarded_pg.py create mode 100644 tests/unit/test_rewards.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 02d61e4..0dab219 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,21 @@ All notable changes to miniVERL are recorded here. The format follows prompt/group cursors, trajectory counts, backend synchronization, cache identity and export provenance to exact interruption/resume state. +### Rewarded OPD profile + +- Added the separate conformance-only + `verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1` profile. It combines an + explicit deterministic task-reward advantage with the existing sampled-k1 + distillation advantage; the direct-GKD and reward-free PG profiles retain + their original contracts and identities. +- Added torch-free reward request/result/provider types, a closed exact-answer + provider, versioned group transformations, fail-closed Parquet validation, + and provider/composer identity binding across reward logs, teacher caches, + checkpoints and exports. +- Kept the objective local and critic-free: there is no value model, global + normalization, arbitrary reward code, GRPO claim or distributed execution + claim. No task-quality benchmark was run for this profile. + ### Documentation narrative - Rebuilt the English, Chinese and PyPI landing narratives around the local diff --git a/PYPI.md b/PYPI.md index 6a65f28..9266821 100644 --- a/PYPI.md +++ b/PYPI.md @@ -135,6 +135,8 @@ the complete mapping. [Compatibility profiles](https://github.com/DaoyuanLi2816/ how profile identity follows plans, caches, checkpoints and exports. Separate conformance-only grouped profiles add transactional Parquet `n>1` samples without changing either measured `n=1` profile or introducing GRPO semantics. +A separate conformance-only rewarded profile adds deterministic exact-answer +rewards and explicit group advantage composition; it has no task-quality claim. ## Measured systems evidence diff --git a/README.md b/README.md index 2f716ce..a0543a7 100644 --- a/README.md +++ b/README.md @@ -135,6 +135,8 @@ the complete mapping. [Compatibility profiles](docs/profiles/index.md) explains how profile identity follows plans, caches, checkpoints and exports. Separate conformance-only grouped profiles add transactional Parquet `n>1` samples without changing either measured `n=1` profile or introducing GRPO semantics. +A separate conformance-only rewarded profile adds deterministic exact-answer +rewards and explicit group advantage composition; it has no task-quality claim. ## Measured systems evidence diff --git a/README.zh-CN.md b/README.zh-CN.md index e674318..81a702d 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -124,6 +124,8 @@ distillation: [兼容 profile](docs/profiles/index.md)说明 profile identity 如何跟随 plan、cache、checkpoint 与 export。另有两个仅完成一致性验证的 grouped profile，为 Parquet prompt 提供事务化 `n>1` 独立样本；它们不改变两个实测 `n=1` profile，也不引入 GRPO 语义。 +另有一个仅完成一致性验证的 rewarded profile，加入确定性的 exact-answer reward 与显式 +group advantage 组合；目前不对任务质量作结论。 ## 实测系统证据 diff --git a/docs/compatibility.md b/docs/compatibility.md index 37d0acc..c5dcb0e 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -22,13 +22,16 @@ migrations, and readers retain the schema-specific behavior needed to load them. ## Pinned OPD profiles -v0.8 adds the separately versioned `verl-opd-v0.8-single-gpu-v1` executable -profile: one local actor, one teacher, one generation, pure GKD -`forward_kl_topk`, token-mean aggregation, LoRA/QLoRA and Parquet prompts. +v0.8 introduced `verl-opd-v0.8-single-gpu-v1`: one local actor, one teacher, +pure GKD `forward_kl_topk`, token-mean aggregation, LoRA/QLoRA and Parquet +prompts. The registry now also contains sampled-k1, grouped-independent and a +separate rewarded sampled-k1 contract. `import-verl --config` publishes a canonical profile plus field report; -`export-verl` recognizes a completed compatible run and emits reward-free OPD -overrides. PG OPD, task-reward mixtures, multiple teachers, multimodal and all -distributed semantics fail closed. Local physical phase scheduling is always +`export-verl` recognizes a completed compatible run and emits its profile-bound +OPD overrides. The rewarded profile accepts only a closed deterministic +exact-answer provider and remains critic/value/GRPO-free. Multiple teachers, +arbitrary reward code, multimodal and all distributed semantics fail closed. +Local physical phase scheduling is always labelled as a reinterpretation of upstream resource fields. The [machine-readable field matrix](generated/verl-opd-v0.8-compatibility.json) @@ -54,8 +57,8 @@ Optimizer state, distributed RNG, FSDP/Megatron native checkpoints, Ray runtime state and teacher-cache conversion sit beyond that exchange contract. The release smoke records artifact/config validation and distributed execution as separate evidence fields. Legacy reward-scaffold bundles use -`launch.template.sh`; current pure-OPD bundles become launchable through the -documented materialization checks. +`launch.template.sh`; current profile exports remain explicit about their +individual launch blockers and semantic-parity status. See the [current local runtime](verl-opd-runtime.md), [current scale-out contract](verl-opd-scaleout.md), and [legacy bridge](legacy-verl-bridge.md). diff --git a/docs/for-verl-users.md b/docs/for-verl-users.md index ce22acb..f1040b2 100644 --- a/docs/for-verl-users.md +++ b/docs/for-verl-users.md @@ -1,7 +1,8 @@ # For verl users -miniVERL brings two documented verl v0.8 OPD profiles into a local, inspectable -workflow. Familiar field names and Parquet data flow through a typed compiler, +miniVERL brings five closed verl v0.8 OPD profiles into a local, inspectable +workflow; two have measured systems records and three are conformance-only. +Familiar field names and Parquet data flow through a typed compiler, which turns resource intent into sequential phases on one CUDA GPU and records every local mapping in the resulting plan. @@ -18,7 +19,8 @@ by plans, caches, checkpoints and exports. ## Supported input shape - A resolved YAML using the `verl-opd-v0.8-single-gpu-v1` field subset. -- Reward-free verl-style Parquet prompts with structured chat messages. +- Verl-style Parquet prompts with structured chat messages; the rewarded + profile additionally requires bound deterministic exact-answer metadata. - One actor, one teacher, token-mean aggregation, and either direct forward-top-k GKD or sampled-k1 vanilla policy loss. The measured profiles remain `n=1`; their grouped counterparts accept independent `n>1` samples. @@ -100,8 +102,9 @@ produces top-k token IDs and log-probabilities on actor-generated tokens; top-k mass and overlap are diagnostics, not an explicit tail bucket. In the PG k1 profile the target is only the sampled token's teacher log-probability, bound to the rollout actor log-probability and policy version; current actor -log-probability is recomputed at update time. Both profiles use reward-free -distillation with one actor and one teacher. Local runtime strategies +log-probability is recomputed at update time. The rewarded PG profile adds a +separately logged task advantage from the closed exact-answer provider; it has +no critic, value model or GRPO semantics. Local runtime strategies may keep quantized roles resident, swap movable unquantized roles, or share a compatible backbone while preserving distinct role identities in provenance. @@ -127,8 +130,8 @@ has progressed. unsupported values. - **Algorithm field rejected:** select the explicit PG-k1 profile for its narrow sampled policy-loss path, or a grouped profile for Parquet `n>1`. - Rewards, critics, external advantages, KL penalties, group-relative - baselines, multiple teachers and distributed counts remain unsupported. + Arbitrary reward code, critics, value models, KL penalties, multiple teachers + and distributed counts remain unsupported. - **Interpolation rejected:** resolve Hydra/OmegaConf in your trusted verl environment first. miniVERL will not execute `${...}`. - **High-risk reinterpretation not accepted:** inspect `miniverl plan`, then diff --git a/docs/generated/upstream-support-matrix.json b/docs/generated/upstream-support-matrix.json index bad5235..f5e7aa2 100644 --- a/docs/generated/upstream-support-matrix.json +++ b/docs/generated/upstream-support-matrix.json @@ -34,6 +34,14 @@ "upstream_commit": "7aed6b230776f963fa09509c10d9c3a767d1102c", "objective": "sampled_k1_grouped_independent_samples", "distributed_execution_tested": false + }, + { + "name": "verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1", + "status": "active", + "upstream_tag": "v0.8.0", + "upstream_commit": "7aed6b230776f963fa09509c10d9c3a767d1102c", + "objective": "task_reward_plus_sampled_k1_vanilla_policy_loss", + "distributed_execution_tested": false } ], "unprofiled_versions": "unsupported" diff --git a/docs/profiles/index.md b/docs/profiles/index.md index 6e602ec..06c4273 100644 --- a/docs/profiles/index.md +++ b/docs/profiles/index.md @@ -45,13 +45,19 @@ name. | `verl-opd-v0.8-single-gpu-pg-k1-v1` | sampled `k1` + vanilla policy loss | sampled-token teacher log-probability | sampled-token signal; smaller target artifact | measured | | `verl-opd-v0.8-single-gpu-grouped-v1` | direct GKD over independent grouped samples | top-k token IDs and log-probabilities | configurable `n>1`; no group baseline | conformance only | | `verl-opd-v0.8-single-gpu-pg-k1-grouped-v1` | sampled `k1` over independent grouped samples | sampled-token teacher log-probability | configurable `n>1`; no GRPO estimator | conformance only | - -All four are reward-free, strict current-policy, one-actor/one-teacher profiles -on one CUDA GPU. The PG profiles use a detached distillation-derived advantage -with the pinned vanilla policy-loss form. The measured records compare runtime -and semantic conformance; task-quality questions belong to a declared benchmark. -The two grouped profiles add transactional Parquet prompt groups while leaving -the published `n=1` identities unchanged; each sample is optimized independently. +| `verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1` | task reward + sampled `k1` through vanilla policy loss | sampled-token teacher log-probability + deterministic exact-answer reward | explicit group transform; no critic or value model | conformance only | + +The first four profiles are reward-free. The rewarded profile is a separate, +strict current-policy contract with a closed deterministic provider and +versioned task/distillation advantage composition. All five use one actor and +one teacher on one CUDA GPU. The PG profiles retain the pinned vanilla +policy-loss form. The measured records compare runtime and semantic +conformance; the rewarded profile has no task-quality result yet. +Library callers may inject an object implementing `RewardProvider` with +`reward.provider: python_api`; YAML and run artifacts never name or load Python +modules. The built-in verl-shaped profile remains fixed to `exact_answer`. +The group-capable profiles use transactional Parquet prompt groups while +leaving the published `n=1` identities unchanged. See [For verl users](../for-verl-users.md), the [PG-k1 ADR](../adr/0010-verl-v0.8-pg-k1-contract.md), and the runtime evidence for the measured boundary. diff --git a/docs/verl-opd-runtime.md b/docs/verl-opd-runtime.md index 1ac0040..cea04c9 100644 --- a/docs/verl-opd-runtime.md +++ b/docs/verl-opd-runtime.md @@ -1,8 +1,9 @@ # Current verl-style OPD runtime -This is miniVERL's current executable path: two documented, resolved profiles -from official verl `v0.8.0`, pinned at `7aed6b23`, compiled into local phases -on one NVIDIA CUDA GPU. +This is miniVERL's current executable path: five closed profiles from official +verl `v0.8.0`, pinned at `7aed6b23`, compiled into local phases on one NVIDIA +CUDA GPU. Two profiles have measured systems records; the grouped and rewarded +extensions are conformance-only. ## Start from the pinned profile @@ -44,6 +45,12 @@ policy-loss semantics. Its export carries sampled-token signals in place of a top-k target requirement. See [Which profile should I use?](profiles/index.md) for the choice guide. +`verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1` keeps that sampled-k1 path and +adds an explicit deterministic task-reward advantage. The provider, group +transform and both coefficients are versioned and carried by cache, checkpoint +and export identity. This is a local task-plus-distillation objective, not a +GRPO or distributed verl result. + ## Placement strategy - NF4/int8 roles use resident local phases. diff --git a/scripts/check_docs_visual.py b/scripts/check_docs_visual.py index 96a47b7..9d74435 100644 --- a/scripts/check_docs_visual.py +++ b/scripts/check_docs_visual.py @@ -34,6 +34,7 @@ "/for-verl-users/", "/config-overrides/", "/immutable-plans/", + "/profiles/", "/hardware-planning/", "/verl-opd-runtime/", "/verl-opd-reference-workload/", diff --git a/src/miniverl/bridge/export.py b/src/miniverl/bridge/export.py index 63f7a22..c380dec 100644 --- a/src/miniverl/bridge/export.py +++ b/src/miniverl/bridge/export.py @@ -25,7 +25,10 @@ required_verl_text, validate_target_verl, ) -from miniverl.bridge.opd_pg_v08 import VERL_OPD_PG_K1_V08_PROFILE +from miniverl.bridge.opd_pg_v08 import ( + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, + VERL_OPD_PG_K1_V08_PROFILE, +) from miniverl.bridge.opd_v08 import VERL_OPD_V08_PROFILE from miniverl.bridge.profiles import ( VERL_OPD_GROUPED_V08_PROFILE, @@ -432,13 +435,20 @@ def _write_hashes(root: Path) -> None: write_text(checksum, "\n".join(lines) + "\n") -_PG_PROFILES = frozenset({VERL_OPD_PG_K1_V08_PROFILE, VERL_OPD_PG_K1_GROUPED_V08_PROFILE}) +_PG_PROFILES = frozenset( + { + VERL_OPD_PG_K1_V08_PROFILE, + VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, + } +) _OPD_PROFILES = frozenset( { VERL_OPD_V08_PROFILE, VERL_OPD_PG_K1_V08_PROFILE, VERL_OPD_GROUPED_V08_PROFILE, VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, } ) @@ -460,7 +470,9 @@ def _opd_run_contract(run: Path) -> tuple[dict[str, Any], dict[str, Any]] | None if report.get("executable") is not True or not isinstance(source, dict): raise ConfigError("source run does not carry an executable verl OPD compatibility plan") required: dict[str, Any] = { - "distillation.distillation_loss.use_task_rewards": False, + "distillation.distillation_loss.use_task_rewards": ( + profile == VERL_OPD_PG_K1_REWARDED_V08_PROFILE + ), "actor_rollout_ref.actor.use_kl_loss": False, "algorithm.use_kl_in_reward": False, } @@ -591,8 +603,11 @@ def _opd_launch_script(adapter: dict[str, Any], teacher: dict[str, Any]) -> str: def _opd_bundle_readme(profile: str) -> str: + rewarded = profile == VERL_OPD_PG_K1_REWARDED_V08_PROFILE semantics = ( - "sampled k1 teacher log-probabilities and pinned vanilla policy-gradient loss" + "explicit task and sampled-k1 distillation advantages through the pinned vanilla policy loss" + if rewarded + else "sampled k1 teacher log-probabilities and pinned vanilla policy-gradient loss" if profile in _PG_PROFILES else "pure GKD `forward_kl_topk` targets" ) @@ -601,8 +616,7 @@ def _opd_bundle_readme(profile: str) -> str: This checksummed bundle preserves a local `{profile}` run as standard PEFT, tokenizer and Parquet artifacts plus {semantics} overrides for [`verl {VERL_TAG}`]({VERL_REPOSITORY}/tree/{VERL_TAG}) at -`{VERL_COMMIT}`. It contains no reward scaffold because task rewards are -disabled by contract. +`{VERL_COMMIT}`. {"The rewarded profile uses a closed deterministic provider and no critic, value model or GRPO normalization." if rewarded else "Task rewards are disabled by contract."} `recipe/launch.template.sh` remains fail-closed until the exact student and teacher snapshots are materialized. The bundle has not run distributed verl; @@ -657,6 +671,11 @@ def _export_opd_bundle( "teacher base snapshot is not bundled", "distributed verl execution was not tested", ] + rewarded = profile == VERL_OPD_PG_K1_REWARDED_V08_PROFILE + if rewarded: + blockers.append( + "rewarded task-plus-distillation composition has no claimed distributed verl parity" + ) if teacher_adapter_path is not None: blockers.append( "teacher adapter requires an explicit merge/materialization step for upstream verl" @@ -711,7 +730,9 @@ def _export_opd_bundle( }, "miniverl_version": __version__, "target_semantics": ( - "pure OPD sampled k1 policy gradient" + "rewarded OPD task plus sampled-k1 distillation policy gradient" + if rewarded + else "pure OPD sampled k1 policy gradient" if profile in _PG_PROFILES else "pure GKD forward_kl_topk OPD" ), @@ -722,8 +743,13 @@ def _export_opd_bundle( "samples_per_prompt": samples_per_prompt, "trajectory_schema_version": 3, "sample_semantics": "independent_current_policy_trajectory", - "group_baseline": False, + "group_baseline": ( + _get(source, "distillation.distillation_loss.task_advantage_mode") + if rewarded + else False + ), "grpo": False, + **({"critic": False, "value_model": False} if rewarded else {}), }, "artifact_complete": False, "artifact_bundle_complete": True, @@ -738,7 +764,13 @@ def _export_opd_bundle( "launchable": False, "distributed_execution_tested": False, "algorithm_semantic_parity": False, - "reward_required": False, + "reward_required": rewarded, + "reward_provider": ( + _get(source, "distillation.distillation_loss.reward_provider") if rewarded else None + ), + "advantage_composer_version": ( + "miniverl-task-distill-advantage-v1" if rewarded else None + ), "launch_blockers": blockers, "data_round_trip": data_evidence, "unsupported_semantics": [ diff --git a/src/miniverl/bridge/opd_pg_v08.py b/src/miniverl/bridge/opd_pg_v08.py index 6d7ca38..f156a59 100644 --- a/src/miniverl/bridge/opd_pg_v08.py +++ b/src/miniverl/bridge/opd_pg_v08.py @@ -5,7 +5,7 @@ import copy from collections.abc import Mapping, Sequence from pathlib import Path -from typing import Any, Literal +from typing import Any, Literal, cast import yaml from pydantic import BaseModel, Field, ValidationError, model_validator @@ -37,9 +37,15 @@ "load_verl_pg_k1_v08_source", "pg_compatibility_rule", "pg_field_rules_digest", + "VERL_OPD_PG_K1_REWARDED_V08_PROFILE", + "rewarded_pg_compatibility_rule", + "rewarded_pg_field_rules_digest", + "VerlRewardedPGK1V08Profile", + "VerlRewardedPGK1LossConfig", ] VERL_OPD_PG_K1_V08_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-v1" +VERL_OPD_PG_K1_REWARDED_V08_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1" class VerlPGK1LossConfig(VerlOPDLossConfig): @@ -69,6 +75,22 @@ class VerlPGK1V08Profile(VerlOPDV08Profile): distillation: VerlPGK1DistillationConfig +class VerlRewardedPGK1LossConfig(VerlPGK1LossConfig): + """Reward fields owned only by the rewarded profile.""" + + task_reward_coef: float = Field(gt=0) + task_advantage_mode: Literal["raw", "group_center", "group_standardize", "leave_one_out"] + reward_provider: Literal["exact_answer"] + + +class VerlRewardedPGK1DistillationConfig(VerlOPDDistillationConfig): + distillation_loss: VerlRewardedPGK1LossConfig + + +class VerlRewardedPGK1V08Profile(VerlOPDV08Profile): + distillation: VerlRewardedPGK1DistillationConfig + + _PG_FIELD_RULES: dict[str, _Rule] = dict(_FIELD_RULES) _PG_FIELD_RULES.update( { @@ -103,6 +125,24 @@ class VerlPGK1V08Profile(VerlOPDV08Profile): } ) +_REWARDED_PG_FIELD_RULES = dict(_PG_FIELD_RULES) +_REWARDED_PG_FIELD_RULES.update( + { + "distillation.distillation_loss.use_task_rewards": _rule( + "source.use_task_rewards", "exact", "enabled only by the rewarded profile" + ), + "distillation.distillation_loss.task_reward_coef": _rule( + "loss.task_reward_coef", "exact", "explicit task-advantage coefficient" + ), + "distillation.distillation_loss.task_advantage_mode": _rule( + "loss.advantage_mode", "exact", "versioned group reward transformation" + ), + "distillation.distillation_loss.reward_provider": _rule( + "reward.provider", "semantically_conformant", "closed built-in provider registry" + ), + } +) + def pg_compatibility_rule(field: str) -> dict[str, Any]: """Return one immutable PG-profile field rule.""" @@ -127,6 +167,29 @@ def pg_field_rules_digest() -> str: return _canonical_digest([pg_compatibility_rule(field) for field in sorted(_PG_FIELD_RULES)]) +def rewarded_pg_compatibility_rule(field: str) -> dict[str, Any]: + try: + rule = _REWARDED_PG_FIELD_RULES[field] + except KeyError as exc: + raise ConfigError( + f"field {field!r} is not part of profile {VERL_OPD_PG_K1_REWARDED_V08_PROFILE!r}" + ) from exc + return { + "upstream_field": field, + "local_target": rule.target, + "classification": rule.classification, + "reason": rule.reason, + "semantic_risk": rule.risk, + "user_confirmation_required": rule.confirmation, + } + + +def rewarded_pg_field_rules_digest() -> str: + return _canonical_digest( + [rewarded_pg_compatibility_rule(field) for field in sorted(_REWARDED_PG_FIELD_RULES)] + ) + + def _pg_semantic_blocker( path: str, value: Any, *, allow_grouped_samples: bool = False ) -> str | None: @@ -234,6 +297,7 @@ def compile_verl_pg_k1_v08( require_executable: bool = True, allow_grouped_samples: bool = False, profile_name: str = VERL_OPD_PG_K1_V08_PROFILE, + rewarded: bool = False, ) -> CompiledLocalExecutionPlan: """Compile the closed sampled-k1 profile into a deterministic local plan.""" merged = copy.deepcopy(dict(payload)) @@ -259,14 +323,19 @@ def compile_verl_pg_k1_v08( flat = _flatten(merged) _reject_non_finite_source_numbers(flat) - unknown = sorted(path for path in flat if path not in _PG_FIELD_RULES) + rules = _REWARDED_PG_FIELD_RULES if rewarded else _PG_FIELD_RULES + unknown = sorted(path for path in flat if path not in rules) if unknown: raise ConfigError( f"verl PG-k1 OPD config is not executable: unsupported field {unknown[0]!r}", hint="the profile accepts only its documented resolved verl v0.8 subset", ) try: - source = VerlPGK1V08Profile.model_validate(merged) + source = ( + VerlRewardedPGK1V08Profile.model_validate(merged) + if rewarded + else VerlPGK1V08Profile.model_validate(merged) + ) except ValidationError as exc: finite = any("finite" in error.get("msg", "").lower() for error in exc.errors()) message = "verl PG-k1 numeric fields must be finite" if finite else "invalid PG-k1 profile" @@ -275,12 +344,27 @@ def compile_verl_pg_k1_v08( compatibility: list[CompatibilityEntry] = [] blockers: list[str] = [] for path, value in sorted(flat.items()): - rule = _PG_FIELD_RULES[path] + rule = rules[path] blocker = _pg_semantic_blocker( path, value, allow_grouped_samples=allow_grouped_samples, ) + if rewarded: + if path == "distillation.distillation_loss.use_task_rewards": + blocker = None if value is True else "the rewarded profile requires task rewards" + elif ( + path == "distillation.distillation_loss.reward_provider" and value != "exact_answer" + ): + blocker = "only the built-in exact_answer provider is supported" + elif ( + path == "distillation.distillation_loss.task_advantage_mode" + and value in {"group_center", "group_standardize", "leave_one_out"} + and source.actor_rollout_ref.rollout.n == 1 + ): + blocker = "group reward transformations require actor_rollout_ref.rollout.n > 1" + elif path == "distillation.distillation_loss.use_task_rewards" and value is not False: + blocker = "task-reward mixtures are unsupported" if blocker: blockers.append(path) classification: FieldClassification = "unsupported" @@ -323,6 +407,7 @@ def compile_verl_pg_k1_v08( ) executable = not blockers + rewarded_source = cast(VerlRewardedPGK1V08Profile, source) local_execution = { "device_count": 1, "distributed_execution": False, @@ -334,16 +419,37 @@ def compile_verl_pg_k1_v08( "teacher_model": source.distillation.teacher_models.teacher_model.model_path, "loss_mode": "k1", "loss_reduction": "token-mean", - "task_rewards": False, + "task_rewards": rewarded, "policy_gradient": True, "policy_loss_mode": "vanilla", "teacher_target": "sampled_token_log_probability", "samples_per_prompt": source.actor_rollout_ref.rollout.n, "group_semantics": "independent_current_policy_trajectories", "trajectory_schema_version": 3 if allow_grouped_samples else 2, + **( + { + "reward_provider": rewarded_source.distillation.distillation_loss.reward_provider, + "task_advantage_mode": ( + rewarded_source.distillation.distillation_loss.task_advantage_mode + ), + "distillation_coef": ( + rewarded_source.distillation.distillation_loss.distillation_loss_coef + ), + "task_reward_coef": ( + rewarded_source.distillation.distillation_loss.task_reward_coef + ), + "advantage_composer_version": "miniverl-task-distill-advantage-v1", + "critic": False, + "value_model": False, + "global_advantage_normalization": False, + "grpo": False, + } + if rewarded + else {} + ), } high_risk = [item.upstream_field for item in compatibility if item.user_confirmation_required] - common = { + common: dict[str, Any] = { "profile": profile_name, "upstream": {"repository": VERL_REPOSITORY, "tag": VERL_TAG, "commit": VERL_COMMIT}, "source_digest": _canonical_digest(merged), @@ -387,6 +493,7 @@ def load_verl_pg_k1_v08( require_executable: bool = True, allow_grouped_samples: bool = False, profile_name: str = VERL_OPD_PG_K1_V08_PROFILE, + rewarded: bool = False, ) -> CompiledLocalExecutionPlan: """Load one resolved YAML document without executing interpolation.""" if path.suffix.lower() in {".sh", ".bash", ".ps1", ".cmd", ".bat"}: @@ -407,6 +514,7 @@ def load_verl_pg_k1_v08( require_executable=require_executable, allow_grouped_samples=allow_grouped_samples, profile_name=profile_name, + rewarded=rewarded, ) @@ -420,6 +528,7 @@ def load_verl_pg_k1_v08_source( require_executable: bool = True, allow_grouped_samples: bool = False, profile_name: str = VERL_OPD_PG_K1_V08_PROFILE, + rewarded: bool = False, ) -> CompiledLocalExecutionPlan: """Load the PG profile from an explicit resolved YAML path.""" return load_verl_pg_k1_v08( @@ -431,4 +540,5 @@ def load_verl_pg_k1_v08_source( require_executable=require_executable, allow_grouped_samples=allow_grouped_samples, profile_name=profile_name, + rewarded=rewarded, ) diff --git a/src/miniverl/bridge/opd_runtime.py b/src/miniverl/bridge/opd_runtime.py index e5ba228..ee99bb3 100644 --- a/src/miniverl/bridge/opd_runtime.py +++ b/src/miniverl/bridge/opd_runtime.py @@ -4,12 +4,16 @@ import math import re -from typing import Any, Literal +from typing import Any, Literal, cast from pydantic import BaseModel, ConfigDict from miniverl.bridge.opd_capabilities import decide_placement -from miniverl.bridge.opd_pg_v08 import VERL_OPD_PG_K1_V08_PROFILE +from miniverl.bridge.opd_pg_v08 import ( + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, + VERL_OPD_PG_K1_V08_PROFILE, + VerlRewardedPGK1LossConfig, +) from miniverl.bridge.opd_v08 import CompiledLocalExecutionPlan from miniverl.bridge.profiles import ( VERL_OPD_GROUPED_V08_PROFILE, @@ -106,6 +110,7 @@ def build_system_plan(compiled: CompiledLocalExecutionPlan) -> OPDSystemPlan: pg_k1 = compiled.profile in { VERL_OPD_PG_K1_V08_PROFILE, VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, } if pg_k1: # token id + position + actor/teacher log-probs + response-mask weight @@ -202,7 +207,7 @@ def build_system_plan(compiled: CompiledLocalExecutionPlan) -> OPDSystemPlan: "teacher_target": "sampled_token_log_probability", "policy_gradient": True, "policy_loss_mode": "vanilla", - "task_rewards": False, + "task_rewards": compiled.profile == VERL_OPD_PG_K1_REWARDED_V08_PROFILE, } if pg_k1 else { @@ -331,10 +336,13 @@ def compile_native_run_config( pg_k1 = compiled.profile in { VERL_OPD_PG_K1_V08_PROFILE, VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, } + rewarded = compiled.profile == VERL_OPD_PG_K1_REWARDED_V08_PROFILE + rewarded_loss = cast(VerlRewardedPGK1LossConfig, source.distillation.distillation_loss) loss_payload = ( { - "mode": "verl_pg_k1", + "mode": "verl_pg_k1_rewarded" if rewarded else "verl_pg_k1", "aggregation": "token-mean", "divergence": "reverse_kl", "temperature": 1.0, @@ -348,6 +356,18 @@ def compile_native_run_config( "clip_ratio_high": 0.2, "clip_ratio_c": 3.0, "estimator_implementation_version": "verl-v0.8-pg-k1-v1", + **( + { + "advantage_composer_version": "miniverl-task-distill-advantage-v1", + "advantage_mode": rewarded_loss.task_advantage_mode, + "distillation_coef": ( + source.distillation.distillation_loss.distillation_loss_coef + ), + "task_reward_coef": rewarded_loss.task_reward_coef, + } + if rewarded + else {} + ), } if pg_k1 else { @@ -373,7 +393,7 @@ def compile_native_run_config( "tags": [ compiled.profile, "verl-v0.8.0", - "pure-opd", + "rewarded-opd" if rewarded else "pure-opd", "pg-k1" if pg_k1 else "direct-gkd", ], "profile_identity": profile_identity, @@ -418,7 +438,7 @@ def compile_native_run_config( "val_files": source.data.val_files, "prompt_key": source.data.prompt_key, "allow_plain_string_prompts": False, - "use_task_rewards": False, + "use_task_rewards": rewarded, "max_prompt_length": source.data.max_prompt_length, "max_response_length": source.data.max_response_length, "truncation": source.data.truncation, @@ -430,6 +450,7 @@ def compile_native_run_config( "hf_cached" if compiled.profile in {VERL_OPD_GROUPED_V08_PROFILE, VERL_OPD_PG_K1_GROUPED_V08_PROFILE} + or rewarded else "hf_reference" ), "samples_per_prompt": source.actor_rollout_ref.rollout.n, @@ -449,6 +470,11 @@ def compile_native_run_config( }, "selection": {"selector": "all_model_tokens"}, "loss": loss_payload, + "reward": { + "enabled": rewarded, + "provider": "exact_answer" if rewarded else None, + "error_policy": "fail", + }, "train": { "cycles": total_steps, "rollouts_per_cycle": logical_batch, diff --git a/src/miniverl/bridge/profiles.py b/src/miniverl/bridge/profiles.py index c47da9e..22ca2e3 100644 --- a/src/miniverl/bridge/profiles.py +++ b/src/miniverl/bridge/profiles.py @@ -17,11 +17,15 @@ from miniverl.bridge.contract import VERL_COMMIT, VERL_REPOSITORY, VERL_TAG from miniverl.bridge.opd_pg_v08 import ( + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, VERL_OPD_PG_K1_V08_PROFILE, VerlPGK1V08Profile, + VerlRewardedPGK1V08Profile, load_verl_pg_k1_v08_source, pg_compatibility_rule, pg_field_rules_digest, + rewarded_pg_compatibility_rule, + rewarded_pg_field_rules_digest, ) from miniverl.bridge.opd_v08 import ( VERL_OPD_V08_PROFILE, @@ -45,6 +49,7 @@ "load_profile_source", "VERL_OPD_GROUPED_V08_PROFILE", "VERL_OPD_PG_K1_GROUPED_V08_PROFILE", + "VERL_OPD_PG_K1_REWARDED_V08_PROFILE", ] VERL_OPD_GROUPED_V08_PROFILE = "verl-opd-v0.8-single-gpu-grouped-v1" @@ -284,6 +289,7 @@ def show(self) -> dict[str, Any]: class _PGK1Profile(CompatibilityProfile): name: str = VERL_OPD_PG_K1_V08_PROFILE grouped: bool = False + rewarded: bool = False @property def identity(self) -> ProfileIdentity: @@ -293,29 +299,49 @@ def identity(self) -> ProfileIdentity: upstream_repository=VERL_REPOSITORY, upstream_tag=VERL_TAG, upstream_commit=VERL_COMMIT, - field_rule_digest=pg_field_rules_digest(), + field_rule_digest=( + rewarded_pg_field_rules_digest() if self.rewarded else pg_field_rules_digest() + ), native_compiler_version=( - "pg-k1-native-grouped-v1" if self.grouped else "pg-k1-native-v1" + "pg-k1-rewarded-native-v1" + if self.rewarded + else "pg-k1-native-grouped-v1" + if self.grouped + else "pg-k1-native-v1" + ), + loss_conformance_version=( + "pg-k1-task-distill-composer-v1" if self.rewarded else "pg-k1-verl-v0.8-v1" + ), + export_version=( + "verl-opd-pg-k1-rewarded-export-v1" if self.rewarded else "verl-opd-pg-k1-export-v1" ), - loss_conformance_version="pg-k1-verl-v0.8-v1", - export_version="verl-opd-pg-k1-export-v1", ) @property def summary(self) -> ProfileSummary: return ProfileSummary( name=self.name, - objective="k1 + vanilla policy loss", + objective=( + "task reward + k1 distillation through vanilla policy loss" + if self.rewarded + else "k1 + vanilla policy loss" + ), teacher_target="sampled-token teacher log-probability", - status="conformance_only" if self.grouped else "measured", + status="conformance_only" if self.grouped or self.rewarded else "measured", identity=self.identity, ) def config_schema(self) -> dict[str, Any]: - return VerlPGK1V08Profile.model_json_schema() + return ( + VerlRewardedPGK1V08Profile.model_json_schema() + if self.rewarded + else VerlPGK1V08Profile.model_json_schema() + ) def explain(self, field: str) -> CompatibilityExplanation: - rule = pg_compatibility_rule(field) + rule = ( + rewarded_pg_compatibility_rule(field) if self.rewarded else pg_compatibility_rule(field) + ) classification = str(rule["classification"]) unsupported = classification == "unsupported" informational = classification == "informational_only" @@ -347,6 +373,7 @@ def check( require_executable=False, allow_grouped_samples=self.grouped, profile_name=self.name, + rewarded=self.rewarded, ) fields = [item.model_dump(mode="json") for item in compiled.compatibility] unsupported = sum(not item.executable for item in compiled.compatibility) @@ -414,15 +441,29 @@ def show(self) -> dict[str, Any]: payload["algorithm"]["adv_estimator"] = None if self.grouped: payload["actor_rollout_ref"]["rollout"]["n"] = 4 + if self.rewarded: + payload["actor_rollout_ref"]["rollout"]["n"] = 4 + payload["distillation"]["distillation_loss"].update( + { + "use_task_rewards": True, + "task_reward_coef": 1.0, + "task_advantage_mode": "group_center", + "reward_provider": "exact_answer", + } + ) return { **self.summary.model_dump(mode="json"), "identity": self.identity.model_dump(mode="json"), "algorithm_contract": { "actor_count": 1, "teacher_count": 1, - "generations_per_prompt": 4 if self.grouped else 1, - "objective": "sampled k1 through vanilla policy loss", - "task_rewards": False, + "generations_per_prompt": 4 if self.grouped or self.rewarded else 1, + "objective": ( + "explicit task plus distillation advantages through vanilla policy loss" + if self.rewarded + else "sampled k1 through vanilla policy loss" + ), + "task_rewards": self.rewarded, "distributed_execution": False, "device": "one local CUDA GPU", }, @@ -445,6 +486,11 @@ def show(self) -> dict[str, Any]: name=VERL_OPD_PG_K1_GROUPED_V08_PROFILE, grouped=True, ), + VERL_OPD_PG_K1_REWARDED_V08_PROFILE: _PGK1Profile( + name=VERL_OPD_PG_K1_REWARDED_V08_PROFILE, + grouped=True, + rewarded=True, + ), } @@ -500,6 +546,7 @@ def load_profile_source( if name in { VERL_OPD_PG_K1_V08_PROFILE, VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, }: return load_verl_pg_k1_v08_source( source, @@ -508,7 +555,12 @@ def load_profile_source( trailing_overrides=trailing_overrides, accept_local_reinterpretations=accept_local_reinterpretations, require_executable=require_executable, - allow_grouped_samples=name == VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + allow_grouped_samples=name + in { + VERL_OPD_PG_K1_GROUPED_V08_PROFILE, + VERL_OPD_PG_K1_REWARDED_V08_PROFILE, + }, profile_name=name, + rewarded=name == VERL_OPD_PG_K1_REWARDED_V08_PROFILE, ) raise AssertionError(f"registered profile {name!r} has no compiler dispatch") diff --git a/src/miniverl/cache/store.py b/src/miniverl/cache/store.py index cfabea1..0ad8ad5 100644 --- a/src/miniverl/cache/store.py +++ b/src/miniverl/cache/store.py @@ -188,6 +188,8 @@ def create( target_representation: str = "topk_distribution", score_implementation_version: str | None = None, estimator_implementation_version: str | None = None, + reward_provider_identity_digest: str | None = None, + advantage_composer_version: str | None = None, execution_plan_digest: str | None = None, profile_identity: dict[str, Any] | None = None, samples_per_prompt: int = 1, @@ -235,6 +237,8 @@ def create( target_representation=target_representation, score_implementation_version=score_implementation_version, estimator_implementation_version=estimator_implementation_version, + reward_provider_identity_digest=reward_provider_identity_digest, + advantage_composer_version=advantage_composer_version, execution_plan_digest=execution_plan_digest, profile_identity=dict(profile_identity or {}), samples_per_prompt=samples_per_prompt, @@ -294,6 +298,8 @@ def assert_compatible( target_representation: str = "topk_distribution", score_implementation_version: str | None = None, estimator_implementation_version: str | None = None, + reward_provider_identity_digest: str | None = None, + advantage_composer_version: str | None = None, execution_plan_digest: str | None = None, profile_identity: dict[str, Any] | None = None, samples_per_prompt: int = 1, @@ -325,6 +331,13 @@ def assert_compatible( "legacy teacher cache cannot verify grouped rollout identity", hint="re-score the cache with trajectory schema v3 and grouped provenance", ) + if self.index.schema_version < 5 and ( + reward_provider_identity_digest is not None or advantage_composer_version is not None + ): + raise StaleCacheError( + "legacy teacher cache cannot verify reward-provider/composer identity", + hint="re-score the cache with the rewarded profile identity recorded", + ) expected: dict[str, Any] = { "teacher_model_id": teacher_model_id, "teacher_model_revision": teacher_model_revision, @@ -340,6 +353,9 @@ def assert_compatible( expected["score_implementation_version"] = score_implementation_version if estimator_implementation_version is not None: expected["estimator_implementation_version"] = estimator_implementation_version + if self.index.schema_version >= 5: + expected["reward_provider_identity_digest"] = reward_provider_identity_digest + expected["advantage_composer_version"] = advantage_composer_version if self.index.schema_version >= 2: expected["tokenizer_identity"] = dict(tokenizer_identity or {}) expected["teacher_adapter_provenance"] = ( diff --git a/src/miniverl/config/__init__.py b/src/miniverl/config/__init__.py index 45998de..867db97 100644 --- a/src/miniverl/config/__init__.py +++ b/src/miniverl/config/__init__.py @@ -31,6 +31,8 @@ Quantization, ReferenceModelConfig, ReportConfig, + RewardConfig, + RewardProviderKind, RolloutConfig, RunConfig, RunMeta, @@ -77,6 +79,8 @@ "PromptTruncation", "Quantization", "ReferenceModelConfig", + "RewardConfig", + "RewardProviderKind", "ReportConfig", "RolloutConfig", "RunConfig", diff --git a/src/miniverl/config/models.py b/src/miniverl/config/models.py index 474a5f3..43e05f8 100644 --- a/src/miniverl/config/models.py +++ b/src/miniverl/config/models.py @@ -27,6 +27,7 @@ TeacherMode, ) from miniverl.errors import ConfigError +from miniverl.rewards import ADVANTAGE_COMPOSER_VERSION, AdvantageMode from miniverl.runtime.generation import RolloutBackendKind from miniverl.utils.runs import write_text @@ -46,6 +47,8 @@ "Quantization", "TeacherContextMode", "LossMode", + "RewardConfig", + "RewardProviderKind", "LossAggregation", "Divergence", "SelectorName", @@ -158,6 +161,7 @@ class LossMode(str, Enum): BUCKETED_TOPK_TAIL = "bucketed_topk_tail" VERL_FORWARD_KL_TOPK = "forward_kl_topk" VERL_PG_K1 = "verl_pg_k1" + VERL_PG_K1_REWARDED = "verl_pg_k1_rewarded" class LossAggregation(str, Enum): @@ -486,6 +490,10 @@ class LossConfig(_Base): clip_ratio_high: float = Field(default=0.2, ge=0.0, lt=1.0) clip_ratio_c: float = Field(default=3.0, gt=1.0) estimator_implementation_version: str | None = None + advantage_composer_version: str | None = None + advantage_mode: AdvantageMode = AdvantageMode.NONE + distillation_coef: float = Field(default=1.0, ge=0.0) + task_reward_coef: float = Field(default=0.0, ge=0.0) @model_validator(mode="before") @classmethod @@ -771,6 +779,22 @@ class ReportConfig(_Base): max_tokens_per_trajectory: int = Field(default=400, ge=0, le=8192) +class RewardProviderKind(str, Enum): + """Built-in reward-provider registry; no artifact-driven imports.""" + + ENVIRONMENT_VERIFIER = "environment_verifier" + EXACT_ANSWER = "exact_answer" + PYTHON_API = "python_api" + + +class RewardConfig(_Base): + """Fail-closed deterministic task-reward configuration.""" + + enabled: bool = False + provider: RewardProviderKind | None = None + error_policy: Literal["fail"] = "fail" + + class RunMeta(_Base): """Identity and reproducibility settings for the run.""" @@ -809,6 +833,7 @@ class RunConfig(_Base): offline_kd: OfflineKDConfig = Field(default_factory=OfflineKDConfig) eval: EvalConfig = Field(default_factory=EvalConfig) report: ReportConfig = Field(default_factory=ReportConfig) + reward: RewardConfig = Field(default_factory=RewardConfig) alignment: AlignmentConfig | None = None # -- cross-field validation ---------------------------------------- @@ -875,10 +900,10 @@ def _validate_combination(self) -> RunConfig: "source.kind=verl_parquet requires the actor and teacher to score the " "same rendered prompt; privileged_context is outside the v0.8 profile" ) - if self.source.use_task_rewards: + if self.source.use_task_rewards and self.loss.mode is not LossMode.VERL_PG_K1_REWARDED: raise ValueError( - "source.use_task_rewards=true is not implemented for the pure OPD prompt " - "runtime; refusing to ignore row reward_model metadata" + "existing PG-k1 profile is reward-free; source.use_task_rewards=true " + "requires the explicit rewarded PG profile" ) if mode is TrainingMode.OPD and self.cache.reuse_across_policy_versions: raise ValueError( @@ -949,7 +974,8 @@ def _validate_combination(self) -> RunConfig: "supported verl v0.8 profile" ) - if self.loss.mode is LossMode.VERL_PG_K1: + pg_modes = {LossMode.VERL_PG_K1, LossMode.VERL_PG_K1_REWARDED} + if self.loss.mode in pg_modes: from miniverl.bridge.opd_pg_contract import VERL_PG_K1_IMPLEMENTATION_VERSION if mode is not TrainingMode.OPD: @@ -1000,6 +1026,66 @@ def _validate_combination(self) -> RunConfig: raise ValueError( "loss.mode=verl_pg_k1 requires every response token with unit response-mask weight" ) + if self.loss.mode is LossMode.VERL_PG_K1: + if self.source.use_task_rewards or self.reward.enabled: + raise ValueError( + "existing PG-k1 profile is reward-free; select the explicit rewarded mode" + ) + if ( + self.loss.advantage_composer_version is not None + or self.loss.advantage_mode is not AdvantageMode.NONE + or self.loss.distillation_coef != 1.0 + or self.loss.task_reward_coef != 0.0 + ): + raise ValueError( + "existing PG-k1 profile cannot enable reward advantage composition" + ) + else: + if not self.source.use_task_rewards: + raise ValueError( + "loss.mode=verl_pg_k1_rewarded requires source.use_task_rewards=true" + ) + if not self.reward.enabled: + raise ValueError("loss.mode=verl_pg_k1_rewarded requires reward.enabled=true") + if self.reward.provider not in { + RewardProviderKind.EXACT_ANSWER, + RewardProviderKind.PYTHON_API, + }: + raise ValueError( + "rewarded PG requires exact_answer or an injected python_api provider" + ) + if self.loss.advantage_composer_version != ADVANTAGE_COMPOSER_VERSION: + raise ValueError( + "rewarded PG requires advantage_composer_version=" + f"{ADVANTAGE_COMPOSER_VERSION}" + ) + if self.loss.task_reward_coef <= 0.0: + raise ValueError("rewarded PG requires task_reward_coef > 0") + if self.loss.advantage_mode is AdvantageMode.NONE: + raise ValueError("rewarded PG requires an active task advantage_mode") + if ( + self.loss.advantage_mode + in { + AdvantageMode.GROUP_CENTER, + AdvantageMode.GROUP_STANDARDIZE, + AdvantageMode.LEAVE_ONE_OUT, + } + and self.rollout.samples_per_prompt == 1 + ): + raise ValueError( + f"advantage_mode={self.loss.advantage_mode.value} requires " + "samples_per_prompt > 1" + ) + elif ( + self.reward != RewardConfig() + or self.loss.advantage_composer_version is not None + or self.loss.advantage_mode is not AdvantageMode.NONE + or self.loss.distillation_coef != 1.0 + or self.loss.task_reward_coef != 0.0 + ): + raise ValueError( + "reward and advantage settings apply only to the explicit rewarded PG profile" + ) if mode is TrainingMode.SFT and self.loss.sampled_token_nll_weight not in (0.0, 1.0): raise ValueError( "run.mode=sft trains with oracle cross-entropy only; " diff --git a/src/miniverl/data/verl_parquet.py b/src/miniverl/data/verl_parquet.py index 47f2992..7b06571 100644 --- a/src/miniverl/data/verl_parquet.py +++ b/src/miniverl/data/verl_parquet.py @@ -15,6 +15,8 @@ SplitName = Literal["train", "val"] _PRESERVED_FIELDS = ("data_source", "ability", "reward_model", "extra_info") +_REWARD_MODEL_FIELDS = frozenset({"style", "ground_truth"}) +_REWARD_SCALARS = (str, int, float, bool) def _canonical(value: Any) -> str: @@ -39,6 +41,35 @@ def _parquet() -> Any: return pq +def _validate_task_reward_metadata(preserved: dict[str, Any], *, location: str) -> None: + reward_model = preserved["reward_model"] + if not isinstance(reward_model, dict): + raise ConfigError(f"{location} reward_model must be an object") + data_source = preserved["data_source"] + if not isinstance(data_source, str) or not data_source: + raise ConfigError(f"{location} requires a non-empty string data_source for task rewards") + unknown = sorted(set(reward_model).difference(_REWARD_MODEL_FIELDS)) + if unknown: + raise ConfigError( + f"{location} reward_model contains unsupported fields: {', '.join(unknown)}", + hint="only deterministic exact/rule metadata is accepted; Python modules and code are not", + ) + if reward_model.get("style") not in {"exact", "rule"}: + raise ConfigError(f"{location} reward_model.style must be exact/rule") + extra = preserved["extra_info"] + if extra is not None and not isinstance(extra, dict): + raise ConfigError(f"{location} extra_info must be an object or null") + primary = reward_model.get("ground_truth") + secondary = extra.get("ground_truth") if isinstance(extra, dict) else None + if primary is None and secondary is None: + raise ConfigError(f"{location} task reward metadata requires a scalar ground_truth") + if primary is not None and secondary is not None and primary != secondary: + raise ConfigError(f"{location} reward_model and extra_info ground_truth values disagree") + ground_truth = primary if primary is not None else secondary + if not isinstance(ground_truth, _REWARD_SCALARS): + raise ConfigError(f"{location} ground_truth must be a JSON scalar") + + @dataclass(frozen=True) class PromptRecord: """One validated prompt row with its complete supported provenance.""" @@ -160,11 +191,8 @@ def _validate_row(self, row: dict[str, Any], *, path: Path, source_index: int) - " or an explicitly enabled plain string" ) preserved = {name: row.get(name) for name in _PRESERVED_FIELDS} - if self.config.use_task_rewards and preserved["reward_model"] is None: - raise ConfigError( - f"{location} has no reward_model but source.use_task_rewards=true", - hint="provide reward_model per row or disable task rewards for pure OPD", - ) + if self.config.use_task_rewards: + _validate_task_reward_metadata(preserved, location=location) payload = {"prompt": validated, **preserved} canonical = _canonical(payload) return PromptRecord( diff --git a/src/miniverl/losses/chunked.py b/src/miniverl/losses/chunked.py index 8a394e7..aca2662 100644 --- a/src/miniverl/losses/chunked.py +++ b/src/miniverl/losses/chunked.py @@ -30,7 +30,7 @@ from collections.abc import Callable from dataclasses import dataclass, field -from typing import Protocol +from typing import Any, Protocol import torch @@ -187,11 +187,15 @@ class VerlPGK1TargetProvider: clip_ratio_high: float = 0.2 clip_ratio_c: float = 3.0 loss_max_clamp: float | None = None + rewarded: bool = False + task_advantage: float = 0.0 + distillation_coef: float = 1.0 + task_reward_coef: float = 0.0 kind: str = "verl_pg_k1" diagnostics: list[dict[str, torch.Tensor | float]] = field(default_factory=list) def divergence(self, start: int, end: int, student_logits: torch.Tensor) -> torch.Tensor: - from miniverl.losses.verl_pg import verl_pg_k1_loss + from miniverl.losses.verl_pg import rewarded_pg_k1_loss, verl_pg_k1_loss targets = self.target_token_ids[start:end].to(student_logits.device) current = ( @@ -199,24 +203,49 @@ def divergence(self, start: int, end: int, student_logits: torch.Tensor) -> torc .gather(-1, targets.unsqueeze(-1)) .squeeze(-1) ) - output = verl_pg_k1_loss( - current_log_probs=current, - old_log_probs=self.old_actor_log_probs[start:end].to(student_logits.device), - teacher_log_probs=self.teacher_sampled_token_log_probs[start:end].to( - student_logits.device - ), - weights=torch.ones(end - start, device=student_logits.device), - clip_ratio=self.clip_ratio, - clip_ratio_low=self.clip_ratio_low, - clip_ratio_high=self.clip_ratio_high, - clip_ratio_c=self.clip_ratio_c, - loss_max_clamp=self.loss_max_clamp, - ) + old = self.old_actor_log_probs[start:end].to(student_logits.device) + teacher = self.teacher_sampled_token_log_probs[start:end].to(student_logits.device) + weights = torch.ones(end - start, device=student_logits.device) + output: Any + if self.rewarded: + output = rewarded_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=teacher, + weights=weights, + task_advantage=self.task_advantage, + distillation_coef=self.distillation_coef, + task_reward_coef=self.task_reward_coef, + clip_ratio=self.clip_ratio, + clip_ratio_low=self.clip_ratio_low, + clip_ratio_high=self.clip_ratio_high, + clip_ratio_c=self.clip_ratio_c, + loss_max_clamp=self.loss_max_clamp, + ) + rewarded_output = output + extra = { + "distillation_advantages": rewarded_output.distillation_advantages.detach().cpu(), + "task_advantages": rewarded_output.task_advantages.detach().cpu(), + } + else: + output = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=teacher, + weights=weights, + clip_ratio=self.clip_ratio, + clip_ratio_low=self.clip_ratio_low, + clip_ratio_high=self.clip_ratio_high, + clip_ratio_c=self.clip_ratio_c, + loss_max_clamp=self.loss_max_clamp, + ) + extra = {} self.diagnostics.append( { "estimator": output.estimator.detach().cpu(), "advantages": output.advantages.detach().cpu(), "ratio": output.ratio.detach().cpu(), + **extra, **output.metrics, } ) diff --git a/src/miniverl/losses/verl_pg.py b/src/miniverl/losses/verl_pg.py index 9d74c8b..1d63d8f 100644 --- a/src/miniverl/losses/verl_pg.py +++ b/src/miniverl/losses/verl_pg.py @@ -18,8 +18,10 @@ __all__ = [ "VERL_PG_K1_IMPLEMENTATION_VERSION", "VerlPGK1Output", + "RewardedPGK1Output", "verl_k1_estimator", "verl_pg_k1_loss", + "rewarded_pg_k1_loss", ] @@ -35,6 +37,20 @@ class VerlPGK1Output: metrics: dict[str, float] +@dataclass(frozen=True) +class RewardedPGK1Output: + """Explicit distillation and task components for the rewarded profile.""" + + loss: torch.Tensor + estimator: torch.Tensor + distillation_advantages: torch.Tensor + task_advantages: torch.Tensor + advantages: torch.Tensor + per_token_loss: torch.Tensor + ratio: torch.Tensor + metrics: dict[str, float] + + def _require_vector(name: str, value: torch.Tensor, length: int | None = None) -> None: if value.ndim != 1: raise ValueError(f"{name} must be a one-dimensional token vector") @@ -136,3 +152,84 @@ def verl_pg_k1_loss( ratio=ratio, metrics=metrics, ) + + +def rewarded_pg_k1_loss( + *, + current_log_probs: torch.Tensor, + old_log_probs: torch.Tensor, + teacher_log_probs: torch.Tensor, + weights: torch.Tensor, + task_advantage: float, + distillation_coef: float, + task_reward_coef: float, + clip_ratio: float = 0.2, + clip_ratio_low: float | None = None, + clip_ratio_high: float | None = None, + clip_ratio_c: float = 3.0, + loss_max_clamp: float | None = None, +) -> RewardedPGK1Output: + """Compose task and distillation advantages for the rewarded PG profile.""" + + n = current_log_probs.numel() + _require_vector("current_log_probs", current_log_probs) + _require_vector("old_log_probs", old_log_probs, n) + _require_vector("teacher_log_probs", teacher_log_probs, n) + _require_vector("weights", weights, n) + if bool((weights < 0).any()): + raise ValueError("weights must be non-negative") + if not torch.isfinite(torch.tensor(task_advantage)): + raise ValueError("task_advantage must be finite") + if not torch.isfinite(torch.tensor([distillation_coef, task_reward_coef])).all() or ( + distillation_coef < 0 or task_reward_coef < 0 + ): + raise ValueError("advantage coefficients must be finite and non-negative") + if not 0.0 <= clip_ratio < 1.0: + raise ValueError("clip_ratio must be in [0, 1)") + low = clip_ratio if clip_ratio_low is None else clip_ratio_low + high = clip_ratio if clip_ratio_high is None else clip_ratio_high + if not 0.0 <= low < 1.0 or not 0.0 <= high < 1.0: + raise ValueError("clip ratios must be in [0, 1)") + if clip_ratio_c <= 1.0: + raise ValueError("clip_ratio_c must be greater than 1") + + work_weights = weights.to(device=current_log_probs.device, dtype=torch.float32) + estimator = verl_k1_estimator(old_log_probs, teacher_log_probs).to(current_log_probs.device) + if loss_max_clamp is not None: + estimator = estimator.clamp(min=-loss_max_clamp, max=loss_max_clamp) + distillation_advantages = -estimator.detach() + task_advantages = torch.full_like(distillation_advantages, float(task_advantage)) + advantages = ( + distillation_coef * distillation_advantages + task_reward_coef * task_advantages + ).detach() + negative_approx_kl = ( + current_log_probs.to(torch.float32) - old_log_probs.to(torch.float32) + ).clamp(min=-20.0, max=20.0) + ratio = negative_approx_kl.exp() + losses1 = -advantages * ratio + losses2 = -advantages * ratio.clamp(1.0 - low, 1.0 + high) + clipped = torch.maximum(losses1, losses2) + dual_clipped = torch.minimum(-advantages * clip_ratio_c, clipped) + per_token_loss = torch.where(advantages < 0, dual_clipped, clipped) + loss = _masked_mean(per_token_loss, work_weights) + metrics = { + "distillation_advantage_mean": float( + _masked_mean(distillation_advantages, work_weights).detach() + ), + "task_advantage_mean": float(_masked_mean(task_advantages, work_weights).detach()), + "total_advantage_mean": float(_masked_mean(advantages, work_weights).detach()), + "pg_clipfrac": float( + _masked_mean((losses2 > losses1).to(torch.float32), work_weights).detach() + ), + "ppo_kl": float(_masked_mean(-negative_approx_kl, work_weights).detach()), + } + return RewardedPGK1Output( + loss=loss, + estimator=estimator, + distillation_advantages=distillation_advantages, + task_advantages=task_advantages, + advantages=advantages, + per_token_loss=per_token_loss, + ratio=ratio, + metrics=metrics, + ) diff --git a/src/miniverl/rewards/__init__.py b/src/miniverl/rewards/__init__.py new file mode 100644 index 0000000..18f14b0 --- /dev/null +++ b/src/miniverl/rewards/__init__.py @@ -0,0 +1,35 @@ +"""Typed deterministic task rewards and advantage composition.""" + +from miniverl.rewards.composer import ( + ADVANTAGE_COMPOSER_VERSION, + AdvantageComposer, + AdvantageComposition, + AdvantageMode, +) +from miniverl.rewards.models import ( + RewardComponent, + RewardProviderIdentity, + RewardRequest, + RewardResult, + RewardStatus, +) +from miniverl.rewards.providers import ( + EnvironmentVerifierRewardProvider, + ExactAnswerRewardProvider, + RewardProvider, +) + +__all__ = [ + "ADVANTAGE_COMPOSER_VERSION", + "AdvantageComposition", + "AdvantageComposer", + "AdvantageMode", + "EnvironmentVerifierRewardProvider", + "ExactAnswerRewardProvider", + "RewardComponent", + "RewardProvider", + "RewardProviderIdentity", + "RewardRequest", + "RewardResult", + "RewardStatus", +] diff --git a/src/miniverl/rewards/composer.py b/src/miniverl/rewards/composer.py new file mode 100644 index 0000000..42756fb --- /dev/null +++ b/src/miniverl/rewards/composer.py @@ -0,0 +1,87 @@ +"""Versioned task-reward advantage composition.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import Enum + +__all__ = [ + "ADVANTAGE_COMPOSER_VERSION", + "AdvantageComposition", + "AdvantageComposer", + "AdvantageMode", +] + +ADVANTAGE_COMPOSER_VERSION = "miniverl-task-distill-advantage-v1" + + +class AdvantageMode(str, Enum): + NONE = "none" + RAW = "raw" + GROUP_CENTER = "group_center" + GROUP_STANDARDIZE = "group_standardize" + LEAVE_ONE_OUT = "leave_one_out" + + +@dataclass(frozen=True) +class AdvantageComposition: + task_advantages: list[float] + mode: AdvantageMode + distillation_coef: float + task_reward_coef: float + zero_variance: bool + implementation_version: str = ADVANTAGE_COMPOSER_VERSION + + +@dataclass(frozen=True) +class AdvantageComposer: + mode: AdvantageMode + distillation_coef: float = 1.0 + task_reward_coef: float = 1.0 + + def __post_init__(self) -> None: + for name, value in ( + ("distillation_coef", self.distillation_coef), + ("task_reward_coef", self.task_reward_coef), + ): + if not math.isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and non-negative") + + def compose_group(self, rewards: list[float]) -> AdvantageComposition: + if not rewards: + raise ValueError("reward group cannot be empty") + if any(not math.isfinite(value) for value in rewards): + raise ValueError("reward group must contain only finite values") + if self.mode is AdvantageMode.LEAVE_ONE_OUT and len(rewards) < 2: + raise ValueError("leave_one_out requires at least two samples per prompt") + mean = sum(rewards) / len(rewards) + variance = sum((value - mean) ** 2 for value in rewards) / len(rewards) + zero_variance = variance == 0.0 + if self.mode is AdvantageMode.NONE: + advantages = [0.0] * len(rewards) + elif self.mode is AdvantageMode.RAW: + advantages = list(rewards) + elif self.mode is AdvantageMode.GROUP_CENTER: + advantages = [value - mean for value in rewards] + elif self.mode is AdvantageMode.GROUP_STANDARDIZE: + scale = math.sqrt(variance) + advantages = ( + [0.0] * len(rewards) + if zero_variance + else [(value - mean) / scale for value in rewards] + ) + else: + advantages = [value - (sum(rewards) - value) / (len(rewards) - 1) for value in rewards] + return AdvantageComposition( + task_advantages=advantages, + mode=self.mode, + distillation_coef=self.distillation_coef, + task_reward_coef=self.task_reward_coef, + zero_variance=zero_variance, + ) + + def expand_task_advantage(self, value: float, assistant_mask: list[bool]) -> list[float]: + if not math.isfinite(value): + raise ValueError("task advantage must be finite") + return [value if selected else 0.0 for selected in assistant_mask] diff --git a/src/miniverl/rewards/models.py b/src/miniverl/rewards/models.py new file mode 100644 index 0000000..43cff21 --- /dev/null +++ b/src/miniverl/rewards/models.py @@ -0,0 +1,136 @@ +"""Torch-free reward requests, results, and provider identities.""" + +from __future__ import annotations + +import hashlib +import json +from enum import Enum +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +__all__ = [ + "RewardComponent", + "RewardProviderIdentity", + "RewardRequest", + "RewardResult", + "RewardStatus", +] + + +def _digest(payload: Any) -> str: + return hashlib.sha256( + json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + ).hexdigest() + + +class _FrozenModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class RewardStatus(str, Enum): + OK = "ok" + ERROR = "error" + NOT_APPLICABLE = "not_applicable" + + +class RewardComponent(_FrozenModel): + name: str = Field(min_length=1) + value: float + detail: str | None = None + + +class RewardProviderIdentity(_FrozenModel): + name: str = Field(min_length=1) + version: str = Field(min_length=1) + config_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + package_name: str | None = None + package_version: str | None = None + deterministic: bool = True + + @property + def digest(self) -> str: + return _digest(self.model_dump(mode="json")) + + +class RewardRequest(_FrozenModel): + schema_version: int = 1 + trajectory_id: str = Field(min_length=1) + prompt_group_id: str = Field(min_length=1) + sample_index: int = Field(ge=0) + samples_per_prompt: int = Field(ge=1) + row_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + response_text: str + reward_model: Any + ground_truth: Any + data_source: str = Field(min_length=1) + input_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + + @classmethod + def create( + cls, + *, + trajectory_id: str, + prompt_group_id: str, + sample_index: int, + samples_per_prompt: int, + row_digest: str, + response_text: str, + reward_model: Any, + ground_truth: Any, + data_source: str, + ) -> RewardRequest: + inputs = { + "trajectory_id": trajectory_id, + "prompt_group_id": prompt_group_id, + "sample_index": sample_index, + "samples_per_prompt": samples_per_prompt, + "row_digest": row_digest, + "response_text": response_text, + "reward_model": reward_model, + "ground_truth": ground_truth, + "data_source": data_source, + } + return cls(**inputs, input_digest=_digest(inputs)) + + @model_validator(mode="after") + def _identity_is_consistent(self) -> RewardRequest: + if self.sample_index >= self.samples_per_prompt: + raise ValueError("sample_index must be smaller than samples_per_prompt") + expected = _digest(self.model_dump(mode="json", exclude={"schema_version", "input_digest"})) + if self.input_digest != expected: + raise ValueError("reward input_digest does not bind the request inputs") + return self + + +class RewardResult(_FrozenModel): + schema_version: int = 1 + trajectory_id: str = Field(min_length=1) + prompt_group_id: str = Field(min_length=1) + sample_index: int = Field(ge=0) + samples_per_prompt: int = Field(ge=1) + provider: RewardProviderIdentity + input_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + raw_reward: float | None + components: tuple[RewardComponent, ...] = () + status: RewardStatus + failure_category: str | None = None + detail: str | None = None + duration_ms: float = Field(ge=0.0) + deterministic: bool + + @model_validator(mode="after") + def _status_matches_value(self) -> RewardResult: + if self.status is RewardStatus.OK and self.raw_reward is None: + raise ValueError("successful reward results require raw_reward") + if self.status is not RewardStatus.OK and self.raw_reward is not None: + raise ValueError("non-success reward results cannot substitute a numeric reward") + if self.raw_reward is not None and not (-float("inf") < self.raw_reward < float("inf")): + raise ValueError("raw_reward must be finite") + return self diff --git a/src/miniverl/rewards/providers.py b/src/miniverl/rewards/providers.py new file mode 100644 index 0000000..5a7c621 --- /dev/null +++ b/src/miniverl/rewards/providers.py @@ -0,0 +1,184 @@ +"""Built-in deterministic reward providers; no artifact-driven code loading.""" + +from __future__ import annotations + +import hashlib +import json +import time +from typing import Any, Protocol, runtime_checkable + +from miniverl.rewards.models import ( + RewardComponent, + RewardProviderIdentity, + RewardRequest, + RewardResult, + RewardStatus, +) + +__all__ = [ + "EnvironmentVerifierRewardProvider", + "ExactAnswerRewardProvider", + "RewardProvider", +] + + +def _config_digest(payload: Any) -> str: + return hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +@runtime_checkable +class RewardProvider(Protocol): + @property + def identity(self) -> RewardProviderIdentity: ... + + def score(self, request: RewardRequest) -> RewardResult: ... + + +class ExactAnswerRewardProvider: + """Strict string equality over declared Parquet ground truth.""" + + def __init__(self, *, strip_whitespace: bool = True) -> None: + self.strip_whitespace = strip_whitespace + self._identity = RewardProviderIdentity( + name="builtin_exact_answer", + version="miniverl-exact-answer-v1", + config_digest=_config_digest({"strip_whitespace": strip_whitespace}), + package_name="miniverl", + deterministic=True, + ) + + @property + def identity(self) -> RewardProviderIdentity: + return self._identity + + def score(self, request: RewardRequest) -> RewardResult: + started = time.perf_counter() + metadata = request.reward_model + valid = isinstance(metadata, dict) and metadata.get("style") in {"exact", "rule"} + declared = metadata.get("ground_truth") if valid else None + if declared is None and valid: + declared = request.ground_truth + if not valid or not isinstance(declared, (str, int, float, bool)): + return self._result( + request, + started=started, + status=RewardStatus.ERROR, + raw_reward=None, + failure_category="invalid_reward_metadata", + detail=( + "reward_model must declare style exact/rule and a scalar ground_truth; " + "arbitrary Python reward metadata is not executable" + ), + ) + expected = str(declared) + predicted = request.response_text + if self.strip_whitespace: + expected = expected.strip() + predicted = predicted.strip() + solved = predicted == expected + return self._result( + request, + started=started, + status=RewardStatus.OK, + raw_reward=1.0 if solved else 0.0, + components=( + RewardComponent( + name="exact_match", + value=1.0 if solved else 0.0, + detail="deterministic normalized string equality", + ), + ), + failure_category=None if solved else "answer_mismatch", + detail=None if solved else f"expected {expected!r}, received {predicted!r}", + ) + + def _result( + self, + request: RewardRequest, + *, + started: float, + status: RewardStatus, + raw_reward: float | None, + components: tuple[RewardComponent, ...] = (), + failure_category: str | None, + detail: str | None, + ) -> RewardResult: + return RewardResult( + trajectory_id=request.trajectory_id, + prompt_group_id=request.prompt_group_id, + sample_index=request.sample_index, + samples_per_prompt=request.samples_per_prompt, + provider=self.identity, + input_digest=request.input_digest, + raw_reward=raw_reward, + components=components, + status=status, + failure_category=failure_category, + detail=detail, + duration_ms=max((time.perf_counter() - started) * 1000.0, 0.0), + deterministic=True, + ) + + +class EnvironmentVerifierRewardProvider: + """Adapter for the active deterministic ToolEnvironment verifier.""" + + def __init__(self, environment: Any) -> None: + self.environment = environment + name = str(getattr(environment, "name", type(environment).__name__)) + version = str(getattr(environment, "verifier_version", "unspecified")) + self._identity = RewardProviderIdentity( + name=f"environment:{name}", + version=version, + config_digest=_config_digest(dict(getattr(environment, "params", {}))), + package_name="miniverl", + deterministic=True, + ) + + @property + def identity(self) -> RewardProviderIdentity: + return self._identity + + def score(self, request: RewardRequest) -> RewardResult: + started = time.perf_counter() + try: + verification = self.environment.verify(request.response_text) + except Exception as exc: + return RewardResult( + trajectory_id=request.trajectory_id, + prompt_group_id=request.prompt_group_id, + sample_index=request.sample_index, + samples_per_prompt=request.samples_per_prompt, + provider=self.identity, + input_digest=request.input_digest, + raw_reward=None, + status=RewardStatus.ERROR, + failure_category="environment_verifier_error", + detail=str(exc), + duration_ms=max((time.perf_counter() - started) * 1000.0, 0.0), + deterministic=True, + ) + category = getattr(verification.failure_category, "value", verification.failure_category) + return RewardResult( + trajectory_id=request.trajectory_id, + prompt_group_id=request.prompt_group_id, + sample_index=request.sample_index, + samples_per_prompt=request.samples_per_prompt, + provider=self.identity, + input_digest=request.input_digest, + raw_reward=float(verification.reward), + components=( + RewardComponent( + name="environment_verifier", + value=float(verification.reward), + detail=verification.detail, + ), + ), + status=RewardStatus.OK, + failure_category=str(category) if category is not None else None, + detail=verification.detail, + duration_ms=max((time.perf_counter() - started) * 1000.0, 0.0), + deterministic=True, + ) diff --git a/src/miniverl/schemas/cache.py b/src/miniverl/schemas/cache.py index c6ef44c..c63a264 100644 --- a/src/miniverl/schemas/cache.py +++ b/src/miniverl/schemas/cache.py @@ -22,8 +22,8 @@ "CacheCompressionStats", ] -CACHE_SCHEMA_VERSION = 4 -READABLE_CACHE_SCHEMA_VERSIONS = frozenset({1, 2, 3, CACHE_SCHEMA_VERSION}) +CACHE_SCHEMA_VERSION = 5 +READABLE_CACHE_SCHEMA_VERSIONS = frozenset({1, 2, 3, 4, CACHE_SCHEMA_VERSION}) class CacheEntryMeta(BaseModel): @@ -82,6 +82,8 @@ class CacheIndex(BaseModel): ) score_implementation_version: str | None = None estimator_implementation_version: str | None = None + reward_provider_identity_digest: str | None = None + advantage_composer_version: str | None = None execution_plan_digest: str | None = None profile_identity: dict[str, Any] = Field(default_factory=dict) samples_per_prompt: int = Field(default=1, ge=1) diff --git a/src/miniverl/teachers/local.py b/src/miniverl/teachers/local.py index 178a0d2..8c10e15 100644 --- a/src/miniverl/teachers/local.py +++ b/src/miniverl/teachers/local.py @@ -107,8 +107,12 @@ def score( old_actor_log_probs=empty, teacher_sampled_token_log_probs=empty, loss_max_clamp=self.loss.loss_max_clamp, + rewarded=self.loss.mode is LossMode.VERL_PG_K1_REWARDED, + task_advantage=float(student.metadata.get("task_advantage", 0.0)), + distillation_coef=self.loss.distillation_coef, + task_reward_coef=self.loss.task_reward_coef, ) - if self.loss.mode is LossMode.VERL_PG_K1 + if self.loss.mode in {LossMode.VERL_PG_K1, LossMode.VERL_PG_K1_REWARDED} else VerlTopKTargetProvider( topk_indices=torch.zeros(0, 1, dtype=torch.long), topk_log_probs=torch.zeros(0, 1), @@ -143,7 +147,7 @@ def score( with role_context, torch.no_grad(): hidden = self.backend.hidden_states_at(source.token_ids, positions, with_grad=False) - if self.loss.mode is LossMode.VERL_PG_K1: + if self.loss.mode in {LossMode.VERL_PG_K1, LossMode.VERL_PG_K1_REWARDED}: if teacher_view is not None: raise AlignmentError("PG-k1 supports only a standard same-prompt teacher") response_ids = [ @@ -196,6 +200,10 @@ def score( clip_ratio_high=self.loss.clip_ratio_high, clip_ratio_c=self.loss.clip_ratio_c, loss_max_clamp=self.loss.loss_max_clamp, + rewarded=self.loss.mode is LossMode.VERL_PG_K1_REWARDED, + task_advantage=float(student.metadata.get("task_advantage", 0.0)), + distillation_coef=self.loss.distillation_coef, + task_reward_coef=self.loss.task_reward_coef, ) from miniverl.bridge.opd_pg_contract import VERL_PG_K1_IMPLEMENTATION_VERSION @@ -369,10 +377,14 @@ def describe(self) -> dict[str, Any]: if self.loss.mode is LossMode.VERL_FORWARD_KL_TOPK else ( "verl-v0.8-pg-k1-v1" - if self.loss.mode is LossMode.VERL_PG_K1 + if self.loss.mode in {LossMode.VERL_PG_K1, LossMode.VERL_PG_K1_REWARDED} else "miniverl-native-v1" ) ), - "top_k": None if self.loss.mode is LossMode.VERL_PG_K1 else self._effective_top_k(), + "top_k": ( + None + if self.loss.mode in {LossMode.VERL_PG_K1, LossMode.VERL_PG_K1_REWARDED} + else self._effective_top_k() + ), "temperature": self.loss.temperature, } diff --git a/src/miniverl/training/checkpoint.py b/src/miniverl/training/checkpoint.py index 642f3be..1e2b5bb 100644 --- a/src/miniverl/training/checkpoint.py +++ b/src/miniverl/training/checkpoint.py @@ -63,6 +63,8 @@ class CheckpointState: pending_group_identity: list[str] = field(default_factory=list) committed_group_identity: list[str] = field(default_factory=list) backend_sync_identity: str = "" + reward_provider_identity_digest: str | None = None + advantage_composer_version: str | None = None scheduler: dict[str, Any] = field(default_factory=dict) scaler: dict[str, Any] | None = None optimizer_param_groups: list[dict[str, Any]] = field(default_factory=list) @@ -106,6 +108,8 @@ def to_dict(self) -> dict[str, Any]: "pending_group_identity": self.pending_group_identity, "committed_group_identity": self.committed_group_identity, "backend_sync_identity": self.backend_sync_identity, + "reward_provider_identity_digest": self.reward_provider_identity_digest, + "advantage_composer_version": self.advantage_composer_version, "scheduler": self.scheduler, "scaler": self.scaler, "optimizer_param_groups": self.optimizer_param_groups, diff --git a/src/miniverl/training/trainer.py b/src/miniverl/training/trainer.py index ac4ffd8..7021f2b 100644 --- a/src/miniverl/training/trainer.py +++ b/src/miniverl/training/trainer.py @@ -30,7 +30,7 @@ from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, cast from miniverl import __version__ from miniverl.agent.loop import RolloutRunner, RolloutStats @@ -45,6 +45,7 @@ OfflineKDTrajectorySource, OPDFreshness, Quantization, + RewardProviderKind, RunConfig, SourceKind, TeacherContextMode, @@ -69,7 +70,12 @@ resolve_device, ) from miniverl.schemas.alignment import AlignmentMap -from miniverl.schemas.trajectory import TRAJECTORY_SCHEMA_VERSION, Trajectory +from miniverl.schemas.trajectory import ( + TRAJECTORY_SCHEMA_VERSION, + SpanType, + Trajectory, + VerificationRecord, +) from miniverl.selection.selectors import ( SelectionResult, SelectionStats, @@ -197,6 +203,7 @@ def __init__( plan: MemoryPlan, run_lock: RunLock, evaluation_only: bool = False, + reward_provider: Any | None = None, ) -> None: self._closed = False self._state = TrainerState.READY @@ -253,13 +260,33 @@ def __init__( from miniverl.runtime.roles import LocalArtifactBridge, LocalRoleGraph self.reference = reference + self.reward_provider: Any | None = reward_provider + if config.loss.mode is LossMode.VERL_PG_K1_REWARDED: + from miniverl.rewards.providers import ( + ExactAnswerRewardProvider, + RewardProvider, + ) + + if config.reward.provider is RewardProviderKind.EXACT_ANSWER: + if reward_provider is not None: + raise ConfigError("an injected provider requires reward.provider=python_api") + self.reward_provider = ExactAnswerRewardProvider() + elif reward_provider is None or not isinstance(reward_provider, RewardProvider): + raise ConfigError( + "reward.provider=python_api requires a RewardProvider object passed to " + "OPDTrainer.from_config(..., reward_provider=...)" + ) + elif reward_provider is not None: + raise ConfigError("an injected reward provider requires the rewarded PG profile") + self.reward_log = JsonlWriter(paths.root / "rewards.jsonl") + self.advantage_log = JsonlWriter(paths.root / "advantages.jsonl") self.artifact_bridge = LocalArtifactBridge(paths.root) self.role_graph = LocalRoleGraph( actor_policy=self.student, rollout_runtime=self.rollout_runtime, teacher_policy=self.teacher, reference_policy=self.reference, - reward_or_verifier=self.environment, + reward_or_verifier=cast(Any, self.reward_provider or self.environment), target_builder=self.scorer, update_runtime=self, evaluation_runtime=self, @@ -389,6 +416,7 @@ def from_config( for_evaluation: bool = False, lock_timeout: float = 0.0, already_acquired_lock: RunLock | None = None, + reward_provider: Any | None = None, ) -> OPDTrainer: """Validate, seed, load models and create the run directory. @@ -672,6 +700,7 @@ def teacher_fits() -> bool: plan=plan, run_lock=run_lock, evaluation_only=for_evaluation, + reward_provider=reward_provider, ) if write_artifacts and resume_options == 0: trainer._write_startup_artifacts() @@ -1286,7 +1315,10 @@ def _open_cache(self) -> TeacherCache: from miniverl.losses.verl_topk import VERL_TOPK_SCORE_IMPLEMENTATION assert self.teacher is not None - pg_k1 = self.config.loss.mode is LossMode.VERL_PG_K1 + pg_k1 = self.config.loss.mode in { + LossMode.VERL_PG_K1, + LossMode.VERL_PG_K1_REWARDED, + } top_k = ( 1 if pg_k1 @@ -1322,6 +1354,16 @@ def _open_cache(self) -> TeacherCache: "estimator_implementation_version": ( VERL_PG_K1_IMPLEMENTATION_VERSION if pg_k1 else None ), + "reward_provider_identity_digest": ( + self.reward_provider.identity.digest + if self.reward_provider is not None + else None + ), + "advantage_composer_version": ( + self.config.loss.advantage_composer_version + if self.config.loss.mode is LossMode.VERL_PG_K1_REWARDED + else None + ), "execution_plan_digest": self.config.run.execution_plan_digest, "profile_identity": self.config.run.profile_identity, "samples_per_prompt": self.config.rollout.samples_per_prompt, @@ -1698,6 +1740,100 @@ def _build_samples(self, trajectories: list[Trajectory]) -> list[TrainSample]: cache.flush() return samples + def _score_task_rewards(self, trajectories: list[Trajectory]) -> None: + """Score one complete rollout group and attach explicit task advantages.""" + if self.config.loss.mode is not LossMode.VERL_PG_K1_REWARDED: + return + if self.reward_provider is None: + raise ConfigError("rewarded PG requires a configured reward provider") + from miniverl.rewards import AdvantageComposer, AdvantageMode, RewardRequest, RewardStatus + + groups: dict[str, list[tuple[Trajectory, Any]]] = {} + for trajectory in trajectories: + if ( + trajectory.prompt_group_id is None + or trajectory.sample_index is None + or trajectory.samples_per_prompt is None + ): + raise ConfigError("rewarded PG requires schema-v3 grouped trajectories") + metadata = trajectory.metadata + row_digest = metadata.get("row_digest") + data_source = metadata.get("data_source") + if ( + not isinstance(row_digest, str) + or not isinstance(data_source, str) + or not data_source + ): + raise ConfigError("rewarded PG trajectory is missing bound row/data provenance") + extra = metadata.get("extra_info") + ground_truth = extra.get("ground_truth") if isinstance(extra, dict) else None + response = "".join( + span.text + for span in trajectory.spans + if span.span_type in {SpanType.ASSISTANT_TEXT, SpanType.ASSISTANT_FINAL} + ) + request = RewardRequest.create( + trajectory_id=trajectory.trajectory_id, + prompt_group_id=trajectory.prompt_group_id, + sample_index=trajectory.sample_index, + samples_per_prompt=trajectory.samples_per_prompt, + row_digest=row_digest, + response_text=response, + reward_model=metadata.get("reward_model"), + ground_truth=ground_truth, + data_source=data_source, + ) + result = self.reward_provider.score(request) + self.reward_log.write(result.model_dump(mode="json")) + if result.status is not RewardStatus.OK or result.raw_reward is None: + raise ConfigError( + f"reward provider failed closed for {trajectory.trajectory_id}: " + f"{result.failure_category or result.status.value}: {result.detail or ''}" + ) + groups.setdefault(trajectory.prompt_group_id, []).append((trajectory, result)) + + composer = AdvantageComposer( + mode=AdvantageMode(self.config.loss.advantage_mode), + distillation_coef=self.config.loss.distillation_coef, + task_reward_coef=self.config.loss.task_reward_coef, + ) + for group_id, members in groups.items(): + members.sort(key=lambda item: int(item[0].sample_index or 0)) + expected = members[0][0].samples_per_prompt + if len(members) != expected: + raise ConfigError(f"reward group {group_id!r} is incomplete") + composition = composer.compose_group([float(item[1].raw_reward) for item in members]) + for (trajectory, result), advantage in zip( + members, composition.task_advantages, strict=True + ): + trajectory.metadata["task_reward"] = result.model_dump(mode="json") + trajectory.metadata["task_advantage"] = advantage + trajectory.metadata["advantage_composition"] = { + "implementation_version": composition.implementation_version, + "mode": composition.mode.value, + "distillation_coef": composition.distillation_coef, + "task_reward_coef": composition.task_reward_coef, + "zero_variance": composition.zero_variance, + } + self.advantage_log.write( + { + "trajectory_id": trajectory.trajectory_id, + "prompt_group_id": group_id, + "sample_index": trajectory.sample_index, + "reward_input_digest": result.input_digest, + "reward_provider_identity_digest": result.provider.digest, + "raw_task_reward": result.raw_reward, + "task_advantage": advantage, + **trajectory.metadata["advantage_composition"], + } + ) + trajectory.verification = VerificationRecord( + solved=bool(result.raw_reward > 0.0), + reward=float(result.raw_reward), + failure_category=result.failure_category, + detail=result.detail, + ) + # -- optimization -------------------------------------------------------------- def _loss_by_span_type( @@ -2427,6 +2563,7 @@ def _run_cycle(self) -> list[dict[str, Any]]: else None ), ) + self._score_task_rewards(trajectories) if trajectories: rollout_policy_version = trajectories[0].policy_version rollout_seconds = max(time.perf_counter() - rollout_started, 1e-9) @@ -2726,7 +2863,7 @@ def _reload_targets_from_cache(self, samples: list[TrainSample]) -> list[TrainSa ], device=self.plan.device, ) - if config.loss.mode is LossMode.VERL_PG_K1: + if config.loss.mode in {LossMode.VERL_PG_K1, LossMode.VERL_PG_K1_REWARDED}: if ( batch.old_actor_log_probs is None or batch.teacher_sampled_token_log_probs is None @@ -2741,6 +2878,10 @@ def _reload_targets_from_cache(self, samples: list[TrainSample]) -> list[TrainSa clip_ratio_high=config.loss.clip_ratio_high, clip_ratio_c=config.loss.clip_ratio_c, loss_max_clamp=config.loss.loss_max_clamp, + rewarded=config.loss.mode is LossMode.VERL_PG_K1_REWARDED, + task_advantage=float(sample.trajectory.metadata.get("task_advantage", 0.0)), + distillation_coef=config.loss.distillation_coef, + task_reward_coef=config.loss.task_reward_coef, ) elif config.loss.mode is LossMode.VERL_FORWARD_KL_TOPK: sample.teacher.provider = VerlTopKTargetProvider( @@ -3105,6 +3246,14 @@ def _save_checkpoint_impl(self, *, name: str | None = None) -> Path: pending_group_identity=list(self._pending_group_identity), committed_group_identity=list(self._committed_group_identity), backend_sync_identity=self._backend_sync_identity, + reward_provider_identity_digest=( + self.reward_provider.identity.digest if self.reward_provider is not None else None + ), + advantage_composer_version=( + self.config.loss.advantage_composer_version + if self.config.loss.mode is LossMode.VERL_PG_K1_REWARDED + else None + ), scheduler=self.schedule.state_dict(), config_digest=self._config_digest(), resolved_config_digest=self._resolved_config_digest(), @@ -3163,6 +3312,20 @@ def _load_from_checkpoint_impl(self, directory: str | Path) -> CheckpointState: ) if validated.state.group_generation_seed_version != "miniverl-rollout-seed-v1": raise ConfigError("checkpoint uses an unsupported group generation seed version") + expected_reward_digest = ( + self.reward_provider.identity.digest if self.reward_provider is not None else None + ) + if validated.state.reward_provider_identity_digest != expected_reward_digest: + raise ConfigError("checkpoint reward provider identity does not match the current run") + expected_composer = ( + self.config.loss.advantage_composer_version + if self.config.loss.mode is LossMode.VERL_PG_K1_REWARDED + else None + ) + if validated.state.advantage_composer_version != expected_composer: + raise ConfigError( + "checkpoint advantage composer version does not match the current run" + ) if validated.state.pending_group_identity: raise ConfigError( "checkpoint contains a pending partial rollout group", diff --git a/tests/unit/test_cache.py b/tests/unit/test_cache.py index e6bfa66..617be14 100644 --- a/tests/unit/test_cache.py +++ b/tests/unit/test_cache.py @@ -105,6 +105,43 @@ def test_execution_plan_digest_is_persisted_and_required_for_reuse(tmp_path: Pat ) +def test_rewarded_cache_binds_provider_and_advantage_composer(tmp_path: Path) -> None: + cache = TeacherCache.create( + tmp_path / "rewarded", + miniverl_version="0.11.0.dev0", + teacher_model_id="toy-teacher", + teacher_model_revision="rev", + tokenizer_fingerprint="fp", + vocab_size=VOCAB, + top_k=1, + temperature=1.0, + loss_mode="verl_pg_k1_rewarded", + target_representation="sampled_token_log_probs", + estimator_implementation_version="verl-v0.8-pg-k1-v1", + reward_provider_identity_digest="a" * 64, + advantage_composer_version="miniverl-task-distill-advantage-v1", + ) + common = { + "teacher_model_id": "toy-teacher", + "teacher_model_revision": "rev", + "tokenizer_fingerprint": "fp", + "tokenizer_identity": {}, + "teacher_adapter_provenance": None, + "vocab_size": VOCAB, + "top_k": 1, + "temperature": 1.0, + "loss_mode": "verl_pg_k1_rewarded", + "target_representation": "sampled_token_log_probs", + "estimator_implementation_version": "verl-v0.8-pg-k1-v1", + "dtype": "float32", + "reward_provider_identity_digest": "a" * 64, + "advantage_composer_version": "miniverl-task-distill-advantage-v1", + } + cache.assert_compatible(**common) + with pytest.raises(StaleCacheError, match="reward_provider_identity_digest"): + cache.assert_compatible(**{**common, "reward_provider_identity_digest": "b" * 64}) + + # ----------------------------------------------------------- round trip diff --git a/tests/unit/test_packaging.py b/tests/unit/test_packaging.py index 45f745c..999c9db 100644 --- a/tests/unit/test_packaging.py +++ b/tests/unit/test_packaging.py @@ -47,6 +47,7 @@ "losses", "models", "reporting", + "rewards", "runtime", "schemas", "selection", diff --git a/tests/unit/test_profile_registry.py b/tests/unit/test_profile_registry.py index a268ff4..dfe7057 100644 --- a/tests/unit/test_profile_registry.py +++ b/tests/unit/test_profile_registry.py @@ -21,6 +21,7 @@ def test_registry_exposes_immutable_direct_gkd_identity() -> None: PG_PROFILE, GROUPED_PROFILE, GROUPED_PG_PROFILE, + "verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1", ] profile = get_profile(DIRECT_PROFILE) diff --git a/tests/unit/test_rewarded_pg.py b/tests/unit/test_rewarded_pg.py new file mode 100644 index 0000000..551bc4d --- /dev/null +++ b/tests/unit/test_rewarded_pg.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from miniverl.rewards import ADVANTAGE_COMPOSER_VERSION + + +def _config() -> dict: # type: ignore[type-arg] + return { + "run": { + "name": "rewarded-pg", + "mode": "opd", + "seed": 9, + "execution_plan_digest": "d" * 64, + }, + "models": { + "backend": "toy", + "device": "cpu", + "student": {"model_id": "student", "lora": {"enabled": False}}, + "teacher": {"model_id": "teacher", "toy_pretrain_steps": 0}, + }, + "source": { + "kind": "verl_parquet", + "train_files": ["train.parquet"], + "use_task_rewards": True, + }, + "rollout": { + "backend": "hf_cached", + "samples_per_prompt": 2, + "record_logprobs": True, + }, + "selection": {"selector": "all_model_tokens"}, + "loss": { + "mode": "verl_pg_k1_rewarded", + "aggregation": "token-mean", + "divergence": "reverse_kl", + "temperature": 1.0, + "scale_by_temperature_squared": False, + "top_k": 1, + "estimator_implementation_version": "verl-v0.8-pg-k1-v1", + "advantage_composer_version": ADVANTAGE_COMPOSER_VERSION, + "advantage_mode": "group_center", + "distillation_coef": 1.0, + "task_reward_coef": 0.5, + }, + "reward": { + "enabled": True, + "provider": "exact_answer", + "error_policy": "fail", + }, + "train": {"cycles": 1, "rollouts_per_cycle": 1}, + "eval": {"enabled": False}, + "report": {"enabled": False}, + } + + +def test_rewarded_profile_config_requires_group_for_group_normalization() -> None: + from miniverl.config import RunConfig + + payload = _config() + payload["rollout"]["samples_per_prompt"] = 1 + + with pytest.raises(ValidationError, match="requires samples_per_prompt > 1"): + RunConfig.model_validate(payload) + + +def test_rewarded_profile_config_requires_provider_and_fail_closed_policy() -> None: + from miniverl.config import RunConfig + + payload = _config() + payload["reward"]["enabled"] = False + + with pytest.raises(ValidationError, match=r"requires reward\.enabled=true"): + RunConfig.model_validate(payload) + + +def test_existing_pg_profile_still_rejects_task_rewards() -> None: + from miniverl.config import RunConfig + + payload = _config() + payload["loss"].update( + { + "mode": "verl_pg_k1", + "advantage_composer_version": None, + "advantage_mode": "none", + "distillation_coef": 1.0, + "task_reward_coef": 0.0, + } + ) + payload["reward"] = {"enabled": False} + + with pytest.raises(ValidationError, match="existing PG-k1 profile is reward-free"): + RunConfig.model_validate(payload) + + +def test_other_objectives_cannot_silently_ignore_reward_settings() -> None: + from miniverl.config import RunConfig + + payload = _config() + payload["source"]["use_task_rewards"] = False + payload["loss"] = {"mode": "bucketed_topk_tail"} + + with pytest.raises(ValidationError, match="reward and advantage settings apply only"): + RunConfig.model_validate(payload) + + +def test_valid_rewarded_profile_round_trips() -> None: + from miniverl.config import RunConfig + + config = RunConfig.model_validate(deepcopy(_config())) + + assert config.source.use_task_rewards is True + assert config.reward.provider.value == "exact_answer" + assert config.loss.task_reward_coef == 0.5 + + +def test_python_api_provider_is_explicit_config_not_an_artifact_import() -> None: + from miniverl.config import RunConfig + + payload = _config() + payload["reward"]["provider"] = "python_api" + config = RunConfig.model_validate(payload) + + assert config.reward.provider.value == "python_api" + assert "module" not in config.reward.model_dump() + + +@pytest.mark.torch +def test_rewarded_pg_combines_distillation_and_task_components_without_changing_mask() -> None: + import torch + + from miniverl.losses.verl_pg import rewarded_pg_k1_loss + + current = torch.log(torch.tensor([0.4, 0.5, 0.6])) + old = torch.log(torch.tensor([0.3, 0.5, 0.7])) + teacher = torch.log(torch.tensor([0.6, 0.4, 0.8])) + mask = torch.tensor([1.0, 1.0, 0.0]) + result = rewarded_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=teacher, + weights=mask, + task_advantage=0.25, + distillation_coef=0.5, + task_reward_coef=2.0, + ) + + expected_distill = teacher - old + expected_total = 0.5 * expected_distill + 2.0 * 0.25 + torch.testing.assert_close(result.distillation_advantages, expected_distill) + torch.testing.assert_close(result.task_advantages, torch.full_like(old, 0.25)) + torch.testing.assert_close(result.advantages, expected_total) + assert result.metrics["distillation_advantage_mean"] == pytest.approx( + float((expected_distill[:2]).mean()) + ) + assert result.metrics["task_advantage_mean"] == pytest.approx(0.25) + + +@pytest.mark.torch +def test_changing_only_task_reward_reverses_update_direction() -> None: + import torch + + from miniverl.losses.verl_pg import rewarded_pg_k1_loss + + gradients = [] + for task_advantage in (1.0, -1.0): + current = torch.tensor([-1.0], requires_grad=True) + result = rewarded_pg_k1_loss( + current_log_probs=current, + old_log_probs=torch.tensor([-1.0]), + teacher_log_probs=torch.tensor([-1.0]), + weights=torch.ones(1), + task_advantage=task_advantage, + distillation_coef=1.0, + task_reward_coef=1.0, + ) + result.loss.backward() + gradients.append(float(current.grad.item())) + + assert gradients[0] < 0.0 + assert gradients[1] > 0.0 + + +def test_trainer_scores_complete_reward_group_and_persists_provenance(tmp_path: Path) -> None: + from miniverl.config import RunConfig + from miniverl.rewards import ExactAnswerRewardProvider + from miniverl.schemas.trajectory import ( + TRAJECTORY_SCHEMA_VERSION, + Span, + SpanType, + TerminationReason, + Trajectory, + derive_grouped_trajectory_id, + ) + from miniverl.training.trainer import OPDTrainer + from miniverl.utils.runs import JsonlWriter, read_jsonl + + config = RunConfig.model_validate(_config()) + trainer = object.__new__(OPDTrainer) + trainer.config = config + trainer.reward_provider = ExactAnswerRewardProvider() + trainer.reward_log = JsonlWriter(tmp_path / "rewards.jsonl") + trainer.advantage_log = JsonlWriter(tmp_path / "advantages.jsonl") + policy_digest = "a" * 64 + prompt_digest = "b" * 64 + row_digest = "c" * 64 + trajectories = [] + for sample_index, response in enumerate(("answer", "wrong")): + seed = 10 + sample_index + trajectory_id = derive_grouped_trajectory_id( + prompt_group_id="g0", + sample_index=sample_index, + rollout_policy_identity_digest=policy_digest, + generation_seed=seed, + ) + trajectories.append( + Trajectory( + schema_version=TRAJECTORY_SCHEMA_VERSION, + trajectory_id=trajectory_id, + task_id=row_digest, + environment="verl_parquet", + token_ids=[1], + attention_mask=[1], + model_generated_mask=[True], + critical_mask=[False], + spans=[ + Span( + span_type=SpanType.ASSISTANT_TEXT, + start=0, + end=1, + turn_id=0, + text=response, + ) + ], + tokenizer_fingerprint="fp", + model_id="student", + prompt_group_id="g0", + prompt_digest=prompt_digest, + sample_index=sample_index, + samples_per_prompt=2, + generation_seed=seed, + rollout_backend="hf_cached", + rollout_policy_identity_digest=policy_digest, + termination_reason=TerminationReason.EOS_WITHOUT_FINAL, + generated_token_count=1, + metadata={ + "row_digest": row_digest, + "data_source": "unit", + "reward_model": {"style": "exact", "ground_truth": "answer"}, + "extra_info": {}, + }, + ) + ) + + trainer._score_task_rewards(trajectories) + + assert [item.metadata["task_advantage"] for item in trajectories] == [0.5, -0.5] + assert [item.verification.reward for item in trajectories if item.verification] == [1.0, 0.0] + rows = read_jsonl(tmp_path / "rewards.jsonl") + assert [row["raw_reward"] for row in rows] == [1.0, 0.0] + assert all(row["provider"]["version"] == "miniverl-exact-answer-v1" for row in rows) + advantages = read_jsonl(tmp_path / "advantages.jsonl") + assert [row["task_advantage"] for row in advantages] == [0.5, -0.5] + assert all(row["implementation_version"] == ADVANTAGE_COMPOSER_VERSION for row in advantages) diff --git a/tests/unit/test_rewards.py b/tests/unit/test_rewards.py new file mode 100644 index 0000000..bc7fbad --- /dev/null +++ b/tests/unit/test_rewards.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import math + +import pytest + +from miniverl.rewards import ( + ADVANTAGE_COMPOSER_VERSION, + AdvantageComposer, + AdvantageMode, + EnvironmentVerifierRewardProvider, + ExactAnswerRewardProvider, + RewardRequest, + RewardStatus, +) + + +def _request(*, response: str = "42", reward_model=None) -> RewardRequest: # type: ignore[no-untyped-def] + return RewardRequest.create( + trajectory_id="t0", + prompt_group_id="g0", + sample_index=0, + samples_per_prompt=2, + row_digest="a" * 64, + response_text=response, + reward_model=reward_model or {"style": "exact", "ground_truth": "42"}, + ground_truth="42", + data_source="unit", + ) + + +def test_reward_request_digest_binds_all_deterministic_inputs() -> None: + first = _request() + same = _request() + changed = _request(response="41") + + assert first.input_digest == same.input_digest + assert first.input_digest != changed.input_digest + assert len(first.input_digest) == 64 + + +def test_exact_answer_provider_records_success_and_failure_without_code_loading() -> None: + provider = ExactAnswerRewardProvider() + + solved = provider.score(_request()) + failed = provider.score(_request(response="41")) + + assert solved.status is RewardStatus.OK + assert solved.raw_reward == 1.0 + assert solved.components[0].name == "exact_match" + assert failed.status is RewardStatus.OK + assert failed.raw_reward == 0.0 + assert failed.failure_category == "answer_mismatch" + assert solved.provider.name == "builtin_exact_answer" + assert solved.deterministic is True + + +def test_exact_answer_provider_returns_error_not_zero_for_invalid_metadata() -> None: + provider = ExactAnswerRewardProvider() + request = _request(reward_model={"style": "python", "module": "untrusted.py"}) + + result = provider.score(request) + + assert result.status is RewardStatus.ERROR + assert result.raw_reward is None + assert result.failure_category == "invalid_reward_metadata" + + +def test_environment_verifier_adapter_preserves_components() -> None: + from miniverl.environments.base import FailureCategory, VerificationResult + + class Verifier: + name = "unit_env" + verifier_version = "v2" + + def __init__(self) -> None: + self.params = {"mode": "strict"} + + def verify(self, answer: str) -> VerificationResult: + return VerificationResult( + solved=answer == "yes", + reward=0.75, + expected="yes", + predicted=answer, + failure_category=FailureCategory.SOLVED, + detail="unit verifier", + ) + + result = EnvironmentVerifierRewardProvider(Verifier()).score(_request(response="yes")) + + assert result.status is RewardStatus.OK + assert result.raw_reward == 0.75 + assert result.components[0].name == "environment_verifier" + assert result.provider.name == "environment:unit_env" + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (AdvantageMode.NONE, [0.0, 0.0, 0.0]), + (AdvantageMode.RAW, [1.0, 2.0, 4.0]), + (AdvantageMode.GROUP_CENTER, [-4 / 3, -1 / 3, 5 / 3]), + ( + AdvantageMode.GROUP_STANDARDIZE, + [ + (-4 / 3) / math.sqrt(14 / 9), + (-1 / 3) / math.sqrt(14 / 9), + (5 / 3) / math.sqrt(14 / 9), + ], + ), + (AdvantageMode.LEAVE_ONE_OUT, [-2.0, -0.5, 2.5]), + ], +) +def test_advantage_composer_modes_are_explicit( + mode: AdvantageMode, + expected: list[float], +) -> None: + composer = AdvantageComposer( + mode=mode, + distillation_coef=0.25, + task_reward_coef=2.0, + ) + + result = composer.compose_group([1.0, 2.0, 4.0]) + + assert result.task_advantages == pytest.approx(expected) + assert result.distillation_coef == 0.25 + assert result.task_reward_coef == 2.0 + assert result.implementation_version == ADVANTAGE_COMPOSER_VERSION + + +@pytest.mark.parametrize( + "mode", + [AdvantageMode.GROUP_CENTER, AdvantageMode.GROUP_STANDARDIZE, AdvantageMode.LEAVE_ONE_OUT], +) +def test_zero_variance_group_is_defined_as_zero(mode: AdvantageMode) -> None: + result = AdvantageComposer(mode=mode).compose_group([0.5, 0.5, 0.5]) + + assert result.task_advantages == [0.0, 0.0, 0.0] + assert result.zero_variance is True + + +def test_composer_expands_task_scalar_only_to_selected_assistant_tokens() -> None: + composer = AdvantageComposer(mode=AdvantageMode.RAW) + + expanded = composer.expand_task_advantage(2.0, [False, True, True, False]) + + assert expanded == [0.0, 2.0, 2.0, 0.0] diff --git a/tests/unit/test_verl_opd_export.py b/tests/unit/test_verl_opd_export.py index 47ef78f..32a7f7f 100644 --- a/tests/unit/test_verl_opd_export.py +++ b/tests/unit/test_verl_opd_export.py @@ -181,6 +181,16 @@ def _opd_run( } if "grouped" in profile: source["actor_rollout_ref"]["rollout"]["n"] = 4 + if "rewarded" in profile: + source["actor_rollout_ref"]["rollout"]["n"] = 4 + source["distillation"]["distillation_loss"].update( + { + "use_task_rewards": True, + "task_reward_coef": 1.0, + "task_advantage_mode": "group_center", + "reward_provider": "exact_answer", + } + ) (run / "verl-source-config.json").write_text(json.dumps(source), encoding="utf-8") (run / "verl-compatibility-report.json").write_text( json.dumps( @@ -256,6 +266,29 @@ def test_grouped_export_binds_independent_sample_semantics( } +def test_rewarded_export_records_fail_closed_semantic_boundary(tmp_path: Path) -> None: + from miniverl.bridge.contract import VERL_TAG + from miniverl.bridge.export import export_verl_bundle + + profile = "verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1" + run, _, _ = _opd_run(tmp_path, profile=profile) + report = export_verl_bundle(run, target_verl=VERL_TAG, out=tmp_path / "rewarded-bundle") + + assert report["reward_required"] is True + assert report["reward_provider"] == "exact_answer" + assert report["launchable"] is False + assert report["grouped_rollout_semantics"] == { + "samples_per_prompt": 4, + "trajectory_schema_version": 3, + "sample_semantics": "independent_current_policy_trajectory", + "group_baseline": "group_center", + "grpo": False, + "critic": False, + "value_model": False, + } + assert any("no claimed distributed verl parity" in item for item in report["launch_blockers"]) + + def test_pure_opd_export_is_reward_free_and_preserves_data_bytes(tmp_path: Path) -> None: from miniverl.bridge.contract import VERL_TAG from miniverl.bridge.doctor import inspect_bridge_bundle diff --git a/tests/unit/test_verl_parquet_source.py b/tests/unit/test_verl_parquet_source.py index 4c78bda..860195f 100644 --- a/tests/unit/test_verl_parquet_source.py +++ b/tests/unit/test_verl_parquet_source.py @@ -94,6 +94,43 @@ def test_task_rewards_require_reward_model_but_pure_opd_does_not(tmp_path) -> No list(rewarded.iter_split("train")) +@pytest.mark.parametrize( + ("patch", "message"), + [ + ({"data_source": ""}, "data_source"), + ({"reward_model": {"style": "python", "module": "evil"}}, "unsupported"), + ({"reward_model": {"style": "exact"}}, "ground_truth"), + ( + { + "reward_model": {"style": "exact", "ground_truth": "1"}, + "extra_info": {"ground_truth": "2"}, + }, + "disagree", + ), + ], +) +def test_task_reward_rows_fail_closed_on_ambiguous_metadata( + tmp_path, patch: dict[str, object], message: str +) -> None: + row = _row(0) + row.update(patch) + row["reward_model"] = row.get("reward_model", {"style": "exact", "ground_truth": "answer"}) + path = tmp_path / "rewarded.parquet" + _write(path, [row]) + config = VerlParquetSourceConfig(train_files=[str(path)], use_task_rewards=True) + with pytest.raises(ConfigError, match=message): + list(VerlParquetDataset(config).iter_split("train")) + + +def test_task_reward_row_accepts_bound_exact_ground_truth(tmp_path) -> None: + row = _row(0) + row["reward_model"] = {"style": "exact", "ground_truth": "answer"} + path = tmp_path / "rewarded.parquet" + _write(path, [row]) + config = VerlParquetSourceConfig(train_files=[str(path)], use_task_rewards=True) + assert next(VerlParquetDataset(config).iter_split("train")).reward_model == row["reward_model"] + + def test_chat_template_is_applied_once_and_records_provenance(tmp_path) -> None: train = _write(tmp_path / "train.parquet", [_row(1)]) source = VerlParquetSourceConfig(train_files=[str(train)], max_prompt_length=200) diff --git a/tests/unit/test_verl_pg_profile.py b/tests/unit/test_verl_pg_profile.py index 2d47b3e..31ee0b1 100644 --- a/tests/unit/test_verl_pg_profile.py +++ b/tests/unit/test_verl_pg_profile.py @@ -8,6 +8,7 @@ from miniverl.errors import ConfigError PG_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-v1" +REWARDED_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-rewarded-v1" def _payload() -> dict[str, object]: @@ -39,6 +40,21 @@ def _set(payload: dict[str, object], path: str, value: object) -> None: current[parts[-1]] = value +def _rewarded_payload() -> dict[str, object]: + payload = _payload() + loss = payload["distillation"]["distillation_loss"] # type: ignore[index] + loss.update( # type: ignore[union-attr] + { + "use_task_rewards": True, + "task_reward_coef": 1.0, + "task_advantage_mode": "group_center", + "reward_provider": "exact_answer", + } + ) + payload["actor_rollout_ref"]["rollout"]["n"] = 4 # type: ignore[index] + return payload + + def test_pg_profile_compiles_as_a_separate_executable_contract() -> None: from miniverl.bridge.opd_pg_v08 import compile_verl_pg_k1_v08 from miniverl.bridge.opd_runtime import build_system_plan, compile_native_run_config @@ -95,6 +111,7 @@ def test_registry_lists_pg_profile_with_independent_identity() -> None: PG_PROFILE, "verl-opd-v0.8-single-gpu-grouped-v1", "verl-opd-v0.8-single-gpu-pg-k1-grouped-v1", + REWARDED_PROFILE, ] direct = get_profile(names[0]) pg = get_profile(PG_PROFILE) @@ -102,3 +119,36 @@ def test_registry_lists_pg_profile_with_independent_identity() -> None: assert pg.identity.loss_conformance_version == "pg-k1-verl-v0.8-v1" assert pg.summary.objective == "k1 + vanilla policy loss" assert pg.summary.teacher_target == "sampled-token teacher log-probability" + + +def test_rewarded_pg_profile_compiles_to_explicit_reward_contract() -> None: + from miniverl.bridge.opd_pg_v08 import compile_verl_pg_k1_v08 + from miniverl.bridge.opd_runtime import build_system_plan, compile_native_run_config + + plan = compile_verl_pg_k1_v08( + _rewarded_payload(), + profile_name=REWARDED_PROFILE, + allow_grouped_samples=True, + rewarded=True, + ) + native = compile_native_run_config(plan, system_plan=build_system_plan(plan)) + assert native.loss.mode.value == "verl_pg_k1_rewarded" + assert native.loss.advantage_mode.value == "group_center" + assert native.loss.task_reward_coef == pytest.approx(1.0) + assert native.source.use_task_rewards is True + assert native.reward.enabled is True + assert native.reward.provider.value == "exact_answer" + + +def test_rewarded_profile_rejects_n1_group_normalization() -> None: + from miniverl.bridge.opd_pg_v08 import compile_verl_pg_k1_v08 + + payload = _rewarded_payload() + payload["actor_rollout_ref"]["rollout"]["n"] = 1 # type: ignore[index] + with pytest.raises(ConfigError, match="task_advantage_mode"): + compile_verl_pg_k1_v08( + payload, + profile_name=REWARDED_PROFILE, + allow_grouped_samples=True, + rewarded=True, + )