feat(rl): shared choice-aware parser/decode layer for evaluation and reward (phase 7)

docs/RL_FRAMEWORK.md

@@ -62,6 +62,24 @@ raises on model output and separates format validity from constraint validity;
 the evaluator and the reward adapter both consume it so there is exactly one
 contract implementation.
 
+Phase 6 introduced the prompt-facing choice protocol
+(PromptChoiceRegistry, global Stage 1 ids and local Stage 2 ids), keeping
+the canonical category_id out of prompts and model outputs.
+
+Phase 7 adds the shared choice-aware parser/decode layer in the same parser
+module (`check_stage1_choices` / `check_stage2_choices` ->
+`ChoiceParseResult`): model rollouts answer with choice ids (Stage 1 global
+ids, Stage 2 local bundle ids), which this layer validates (JSON/schema,
+count, uniqueness, known ids / 1..5) and decodes to canonical category_ids
+BEFORE the existing evaluation (`evaluate_stage1_choices` /
+`evaluate_stage2_choices`) and reward entry points
+(`reward_stage1_choices` / `reward_stage2_choices`) apply, completing the
+task-level choice-to-canonical reward contract. Choice validation is
+implemented exactly once; `RewardConfig`, reward numbers, canonical
+registry/ground truth/candidates are unchanged. `RewardResult` gained a
+`constraint_valid` flag (additive) so "valid but wrong" (partial credit) is
+distinguishable from "invalid" (zero).
+
 ### Training adapters
 
 `agent.training.sft` owns messages-Parquet export, validation, the temporary
@@ -117,6 +135,12 @@ Still landing with the first real RL vertical slice (M4):
    reward loop;
 2. a small GPU smoke launcher.
 
+The task-level choice-to-canonical contract is now complete on master: Phase
+6 (prompt-facing choice protocol, PromptChoiceRegistry with global Stage 1
+ids and local Stage 2 ids) and Phase 7 (shared choice-aware parser/decode
+layer consumed by evaluation and reward; see the Evaluation module) are both
+merged.
+
 ## When to add training-algorithm RL code
 
 Do not create empty GRPO/PPO/DAPO packages or `cfg/verl/rl` placeholders.

src/agent/evaluation/classification.py

@@ -10,30 +10,26 @@
 evaluate_stage1/2 consume the unified parser (check_stage1_output /
 check_stage2_output) so the reward adapter and the evaluator share one
 contract implementation. evaluate_stage1_choices / evaluate_stage2_choices
-are thin adapters: they parse and decode the choice protocol BEFORE
-delegating to the canonical evaluators, so choice ids never leak into
-correctness logic and no canonical check logic is duplicated.
+consume the SHARED choice-aware layer (check_stage1_choices /
+check_stage2_choices) and apply the same canonical correctness facts on
+the decoded category ids, so choice validation is implemented exactly
+once across evaluation and reward.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
-import json
 from typing import Sequence
 
 from agent.task import LeafRegistry
 from agent.task.parser import (
-    PredictionFormatError,
+    Stage2Output,
+    check_stage1_choices,
     check_stage1_output,
+    check_stage2_choices,
     check_stage2_output,
-    parse_stage1_output,
-    parse_stage2_output,
-)
-from agent.task.prompt_choices import (
-    PromptChoiceError,
-    PromptChoiceRegistry,
-    decode_stage2_answer,
 )
+from agent.task.prompt_choices import PromptChoiceRegistry
 
 
 @dataclass(frozen=True)
@@ -87,28 +83,23 @@ def evaluate_stage1_choices(
 ) -> Stage1Evaluation:
     """Evaluate a choice-id Stage 1 output; decode BEFORE canonical logic.
 
-    The model answers with global choice ids ("1".."N"); the returned
-    prediction is the decoded canonical category_id tuple. Invalid choice
-    ids / wrong counts / duplicates yield an explicit invalid result with
-    no name or fuzzy fallback. Delegates to evaluate_stage1 with the
-    decoded canonical payload, so the canonical contract logic is never
-    duplicated.
+    Consumes the shared choice-aware check (check_stage1_choices) so
+    evaluation and reward share one choice validation implementation; the
+    returned prediction is the decoded canonical category_id tuple. Invalid
+    choice ids / wrong counts / duplicates yield an explicit invalid result
+    with no name or fuzzy fallback.
     """
     if ground_truth not in registry.ids:
         raise ValueError("ground_truth must belong to the leaf registry")
     choices = choices or PromptChoiceRegistry.from_registry(registry)
-    try:
-        output = parse_stage1_output(solution)
-    except PredictionFormatError as exc:
-        return Stage1Evaluation(None, False, False, False, (str(exc),))
-    try:
-        decoded = choices.decode_candidates(output.candidates)
-    except PromptChoiceError as exc:
-        return Stage1Evaluation(None, True, False, False, (str(exc),))
-    return evaluate_stage1(
-        json.dumps({"candidates": list(decoded)}, ensure_ascii=False, separators=(",", ":")),
-        ground_truth=ground_truth,
-        registry=registry,
+    result = check_stage1_choices(solution, choices=choices)
+    if not result.format_valid:
+        return Stage1Evaluation(None, False, False, False, result.errors)
+    if not result.constraint_valid:
+        return Stage1Evaluation(None, True, False, False, result.errors)
+    assert result.decoded is not None
+    return Stage1Evaluation(
+        result.decoded, True, True, ground_truth in result.decoded, ()
     )
 
 
@@ -150,10 +141,11 @@ def evaluate_stage2_choices(
     """Evaluate a local-id Stage 2 output; decode BEFORE canonical logic.
 
     The model answers with a LOCAL bundle id ("1".."5" in candidate order);
-    the returned prediction is the decoded canonical category_id. Anything
-    but an exact local id yields an explicit invalid result with no name or
-    fuzzy fallback. Delegates to evaluate_stage2 with the decoded canonical
-    payload, so the canonical contract logic is never duplicated.
+    the returned prediction is the decoded canonical category_id. Consumes
+    the shared choice-aware check (check_stage2_choices) so evaluation and
+    reward share one choice validation implementation. Anything but an
+    exact local id yields an explicit invalid result with no name or fuzzy
+    fallback.
     """
     if ground_truth not in registry.ids:
         raise ValueError("ground_truth must belong to the leaf registry")
@@ -163,17 +155,15 @@ def evaluate_stage2_choices(
         or any(candidate not in registry.ids for candidate in candidates)
     ):
         raise ValueError("candidates must be 5 unique IDs from the leaf registry")
-    try:
-        output = parse_stage2_output(solution)
-    except PredictionFormatError as exc:
-        return Stage2Evaluation(None, False, False, False, (str(exc),))
-    try:
-        decoded = decode_stage2_answer(output.answer, tuple(candidates))
-    except PromptChoiceError as exc:
-        return Stage2Evaluation(output.answer, True, False, False, (str(exc),))
-    return evaluate_stage2(
-        json.dumps({"answer": decoded}, ensure_ascii=False, separators=(",", ":")),
-        ground_truth=ground_truth,
-        candidates=candidates,
-        registry=registry,
+    result = check_stage2_choices(solution, candidates=candidates)
+    if not result.format_valid:
+        return Stage2Evaluation(None, False, False, False, result.errors)
+    if not result.constraint_valid:
+        assert isinstance(result.output, Stage2Output)
+        return Stage2Evaluation(
+            result.output.answer, True, False, False, result.errors
+        )
+    assert result.decoded is not None
+    return Stage2Evaluation(
+        result.decoded, True, True, result.decoded == ground_truth, ()
     )

src/agent/task/__init__.py

@@ -10,11 +10,14 @@
 from .dataset_config import BUILTIN_DATASET_CONFIGS, DatasetConfig
 from .identity import code_leaf_map, compact, leaf_registry_from_corpus, qualified_category_id
 from .parser import (
+    ChoiceParseResult,
     ParseResult,
     PredictionFormatError,
     Stage1Output,
     Stage2Output,
+    check_stage1_choices,
     check_stage1_output,
+    check_stage2_choices,
     check_stage2_output,
     parse_stage1_output,
     parse_stage2_output,
@@ -62,8 +65,11 @@
     "Stage1Output",
     "Stage2Output",
     "ParseResult",
+    "ChoiceParseResult",
     "check_stage1_output",
     "check_stage2_output",
+    "check_stage1_choices",
+    "check_stage2_choices",
     "parse_stage1_output",
     "parse_stage2_output",
     "PromptChoice",

src/agent/task/parser.py

@@ -1,6 +1,6 @@
 """Strict parsers for model outputs in the two-stage classification task.
 
-Two layers, sharing the exact JSON-shape parsers:
+Three layers, sharing the exact JSON-shape parsers:
 
 - ``parse_stage1_output`` / ``parse_stage2_output`` raise
   ``PredictionFormatError`` on malformed JSON/schema (used by SFT
@@ -11,6 +11,14 @@
   uniqueness / registry membership, or answer membership). RL rewards and
   the evaluator consume this layer so malformed model output can never
   crash a training loop.
+- ``check_stage1_choices`` / ``check_stage2_choices`` are the SHARED
+  choice-aware layer: they run the same JSON/schema parsing, validate the
+  choice protocol (stage1: exact count, uniqueness, known choice ids;
+  stage2: local id 1..5) and decode to canonical category ids, returning a
+  ``ChoiceParseResult``. Both the evaluation adapters and the RL reward
+  entry points consume this layer, so choice validation is implemented
+  exactly once and the model's choice ids never leak into canonical
+  correctness logic.
 """
 
 from __future__ import annotations
@@ -20,6 +28,7 @@
 from typing import Any, Sequence
 
 from .contracts import LeafRegistry
+from .prompt_choices import PromptChoiceRegistry
 
 
 class PredictionFormatError(ValueError):
@@ -167,3 +176,150 @@ def check_stage2_output(text: str, *, candidates: Sequence[str]) -> ParseResult:
     valid = output.answer in candidate_ids
     errors = () if valid else ("stage2 answer must be one of the candidates",)
     return ParseResult(True, valid, output, errors)
+
+
+_STAGE2_LOCAL_IDS = tuple(str(index) for index in range(1, 6))
+
+
+@dataclass(frozen=True)
+class ChoiceParseResult:
+    """Structured never-raised outcome of the choice-aware parse + decode.
+
+    The choice protocol (``agent.task.prompt_choices``) maps model-facing
+    choice ids to canonical category ids. This result mirrors
+    ``ParseResult``'s shape contract and additionally carries ``decoded``:
+    the canonical category ids the model's choice ids map to (stage1: the
+    candidate tuple; stage2: the single answer id). ``output`` always keeps
+    the model-level parsed shape (choice ids) when format_valid, so
+    consumers can see exactly what the model said before decoding.
+
+    - format_valid: the output is a JSON object with exactly the required
+      schema (identical semantics to ParseResult.format_valid).
+    - constraint_valid: the parsed choice ids satisfy the choice protocol
+      (stage1: exactly ``expected_count`` unique known choice ids; stage2:
+      answer is one of the local ids 1..5).
+    - decoded: the canonical category ids when constraint_valid (stage1:
+      tuple of candidate ids; stage2: the single answer id), else None.
+    - errors: human-readable failure reasons; empty when both valid.
+    """
+
+    format_valid: bool
+    constraint_valid: bool
+    decoded: tuple[str, ...] | str | None = None
+    output: "Stage1Output | Stage2Output | None" = None
+    errors: tuple[str, ...] = ()
+
+    @property
+    def ok(self) -> bool:
+        return self.format_valid and self.constraint_valid
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.format_valid, bool) or not isinstance(
+            self.constraint_valid, bool
+        ):
+            raise ValueError("ChoiceParseResult validity flags must be bool")
+        if self.errors and self.format_valid and self.constraint_valid:
+            raise ValueError("a fully valid ChoiceParseResult must carry no errors")
+        if not isinstance(self.errors, tuple) or not all(
+            isinstance(error, str) for error in self.errors
+        ):
+            raise ValueError("ChoiceParseResult errors must be a tuple of strings")
+        if self.constraint_valid and self.decoded is None:
+            raise ValueError(
+                "a constraint-valid ChoiceParseResult must carry decoded canonical ids"
+            )
+
+    def canonical_view(self) -> ParseResult:
+        """Equivalent canonical ``ParseResult`` for the shared eval/reward cores.
+
+        When constraints hold the output is the decoded canonical shape
+        (Stage1Output of decoded ids / Stage2Output of the decoded answer);
+        when only the format held, the model-level shape is kept so
+        consumers can still inspect what the model said.
+        """
+        if not self.format_valid:
+            return ParseResult(False, False, None, self.errors)
+        if self.constraint_valid:
+            output: Stage1Output | Stage2Output | None = (
+                Stage1Output(self.decoded)
+                if isinstance(self.decoded, tuple)
+                else Stage2Output(self.decoded)
+            )
+        else:
+            output = self.output
+        return ParseResult(True, self.constraint_valid, output, self.errors)
+
+
+def check_stage1_choices(
+    text: str,
+    *,
+    choices: PromptChoiceRegistry,
+    expected_count: int = 5,
+) -> ChoiceParseResult:
+    """Shared choice-aware Stage 1 parser + decode; never raises on model output.
+
+    Validates, in order: JSON object, exact schema (only ``candidates``),
+    candidate count, uniqueness, and choice-id membership in the prompt
+    catalog; then decodes the choice ids to canonical category ids.
+    Malformed model output maps to a structured ``ChoiceParseResult``; only
+    a missing/invalid ``choices`` mapping is a programming error and
+    raises.
+    """
+    if expected_count < 1:
+        raise ValueError("expected_count must be positive")
+    if not isinstance(choices, PromptChoiceRegistry):
+        raise ValueError("choices must be a PromptChoiceRegistry")
+    try:
+        output = parse_stage1_output(text)
+    except PredictionFormatError as exc:
+        return ChoiceParseResult(False, False, None, None, (str(exc),))
+    errors: list[str] = []
+    if len(output.candidates) != expected_count:
+        errors.append(
+            f"stage1 prediction must contain exactly {expected_count} candidates"
+        )
+    if len(set(output.candidates)) != len(output.candidates):
+        errors.append("stage1 candidates must be unique")
+    for choice_id in output.candidates:
+        if not choices.contains_choice_id(choice_id):
+            errors.append(f"choice id {choice_id!r} is not in the prompt catalog")
+    if errors:
+        return ChoiceParseResult(True, False, None, output, tuple(errors))
+    return ChoiceParseResult(
+        True,
+        True,
+        tuple(choices.category_id_of(choice_id) for choice_id in output.candidates),
+        output,
+        (),
+    )
+
+
+def check_stage2_choices(text: str, *, candidates: Sequence[str]) -> ChoiceParseResult:
+    """Shared choice-aware Stage 2 parser + decode; never raises on model output.
+
+    Validates: JSON object, exact schema (only ``answer``), and answer
+    membership in the local ids 1..5; then decodes the local id to the
+    canonical category id of the candidate at that position. Candidates
+    originate from the dataset, so an invalid candidate list is a
+    programming error and raises (mirrors check_stage2_output).
+    """
+    if isinstance(candidates, (str, bytes)) or len(candidates) != 5:
+        raise ValueError("stage2 requires exactly 5 candidates for local-id decode")
+    candidate_ids = tuple(candidates)
+    if not all(isinstance(candidate, str) and candidate for candidate in candidate_ids):
+        raise ValueError("stage2 candidates must be non-empty strings")
+    try:
+        output = parse_stage2_output(text)
+    except PredictionFormatError as exc:
+        return ChoiceParseResult(False, False, None, None, (str(exc),))
+    if output.answer not in _STAGE2_LOCAL_IDS:
+        return ChoiceParseResult(
+            True,
+            False,
+            None,
+            output,
+            (f"stage2 answer {output.answer!r} must be one of 1..5",),
+        )
+    return ChoiceParseResult(
+        True, True, candidate_ids[int(output.answer) - 1], output, ()
+    )

src/agent/training/rl/__init__.py

@@ -20,9 +20,12 @@
     STAGE2_PARTIAL_DEFAULT,
     RewardConfig,
     RewardResult,
+    reward_for_choice_result,
     reward_for_parse_result,
     reward_stage1,
+    reward_stage1_choices,
     reward_stage2,
+    reward_stage2_choices,
 )
 from .sample import (
     RewardMeta,
@@ -49,6 +52,9 @@
     "reward_stage1",
     "reward_stage2",
     "reward_for_parse_result",
+    "reward_stage1_choices",
+    "reward_stage2_choices",
+    "reward_for_choice_result",
     "RL_SPLITS",
     "VERL_RL_COLUMNS",
     "export_rl_dataset",