Phase 5: Composable rubric grader + dense step shaping

graders.py

@@ -1,4 +1,267 @@
 """Composable rubric grader + dense step-level reward shaping.
 
-Will be filled in Phase 5.
+Reward flow per step:
+    compute_step_shaping(state, action, response)  -> float  (dense signal)
+    build_grader()(state)                          -> RewardBreakdown (terminal-ish)
+Environment combines them into reward.shaped_total and reward.binary.
 """
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from models import Action, EpisodeState, RewardBreakdown, ToolResponse
+
+
+# ─────────────────────────────────────────────────────────────────
+# Rubric base
+# ─────────────────────────────────────────────────────────────────
+
+class Rubric:
+    name: str = ""
+
+    def score(self, state: EpisodeState) -> tuple[str, float, dict]:
+        raise NotImplementedError
+
+
+# ─────────────────────────────────────────────────────────────────
+# Concrete rubrics
+# ─────────────────────────────────────────────────────────────────
+
+class CompletionRubric(Rubric):
+    name = "task_completion"
+
+    def score(self, state: EpisodeState) -> tuple[str, float, dict]:
+        gt = state.ground_truth_final_state
+        if not gt:
+            return (self.name, 0.0, {"reason": "no ground truth"})
+        satisfied = 0
+        for key, expected in gt.items():
+            if _matches(state.agent_state.get(key), expected):
+                satisfied += 1
+        total = len(gt)
+        return (
+            self.name,
+            satisfied / total,
+            {"satisfied": satisfied, "total": total},
+        )
+
+
+class DriftDetectionRubric(Rubric):
+    name = "drift_detection"
+
+    def score(self, state: EpisodeState) -> tuple[str, float, dict]:
+        fired = [d for d in state.drift_plan if d.fires_at_step <= state.step]
+        if not fired:
+            return (self.name, 0.0, {"reason": "no drifts fired yet"})
+        detected = sum(1 for d in fired if d.detected_by_agent)
+        return (
+            self.name,
+            detected / len(fired),
+            {"detected": detected, "fired_total": len(fired)},
+        )
+
+
+class AdaptationRubric(Rubric):
+    name = "adaptation_quality"
+
+    def score(self, state: EpisodeState) -> tuple[str, float, dict]:
+        fired = [d for d in state.drift_plan if d.fires_at_step <= state.step]
+        if not fired:
+            return (self.name, 0.0, {"reason": "no drifts fired yet"})
+
+        adapted = 0
+        opportunities = 0
+        for drift in fired:
+            first_post = None
+            for h in state.history:
+                if h.step <= drift.fires_at_step:
+                    continue
+                act = h.action
+                if act.type == "call_tool" and act.tool_call is not None and act.tool_call.tool == drift.tool:
+                    first_post = h
+                    break
+                if act.type == "retry_with_variant" and act.retry is not None and act.retry.tool == drift.tool:
+                    first_post = h
+                    break
+            if first_post is not None:
+                opportunities += 1
+                if first_post.response is not None and first_post.response.ok:
+                    adapted += 1
+
+        if opportunities == 0:
+            return (self.name, 0.0, {"reason": "no post-drift calls yet"})
+        return (
+            self.name,
+            adapted / opportunities,
+            {"adapted": adapted, "opportunities": opportunities},
+        )
+
+
+class EfficiencyRubric(Rubric):
+    name = "efficiency"
+
+    def score(self, state: EpisodeState) -> tuple[str, float, dict]:
+        step_eff = max(0.0, 1.0 - state.step / state.max_steps)
+        tok_eff = state.token_budget_remaining / max(1, state.token_budget)
+        return (
+            self.name,
+            0.5 * step_eff + 0.5 * tok_eff,
+            {"step_eff": step_eff, "tok_eff": tok_eff},
+        )
+
+
+# ─────────────────────────────────────────────────────────────────
+# Composition primitives
+# ─────────────────────────────────────────────────────────────────
+
+class WeightedSum:
+    def __init__(self, rubrics_with_weights: list[tuple[Rubric, float]]) -> None:
+        total_w = sum(w for _, w in rubrics_with_weights)
+        assert abs(total_w - 1.0) < 1e-6, f"Weights must sum to 1.0, got {total_w}"
+        self.items = rubrics_with_weights
+
+    def score(self, state: EpisodeState) -> tuple[float, dict]:
+        total = 0.0
+        breakdown: dict[str, dict] = {}
+        for rubric, w in self.items:
+            name, val, details = rubric.score(state)
+            total += w * val
+            breakdown[name] = {
+                "value": val,
+                "weight": w,
+                "weighted": w * val,
+                "details": details,
+            }
+        return total, breakdown
+
+
+class Gate:
+    def __init__(self, name: str, predicate: Callable[[EpisodeState], bool]) -> None:
+        self.name = name
+        self.predicate = predicate
+
+    def score(self, state: EpisodeState) -> tuple[str, float]:
+        return (self.name, 1.0 if self.predicate(state) else 0.0)
+
+
+# ─────────────────────────────────────────────────────────────────
+# Dense step-level shaping (GRPO convergence fix)
+# ─────────────────────────────────────────────────────────────────
+
+def compute_step_shaping(
+    state: EpisodeState,
+    action: Action,
+    response: ToolResponse | None,
+) -> float:
+    """Dense step-level reward — fires during episode to densify GRPO signal."""
+    shaped = 0.0
+    history = state.history
+    prev_response = history[-1].response if history else None
+
+    # +0.10 — inspecting schema right after a failure
+    if action.type == "inspect_schema":
+        if prev_response is not None and not prev_response.ok:
+            shaped += 0.10
+
+    # +0.15 — correct drift report (matching tool + kind, already fired, not yet detected)
+    if action.type == "report_drift" and action.report is not None:
+        for d in state.drift_plan:
+            matches = (
+                d.tool == action.report.tool
+                and d.kind == action.report.drift_kind
+                and d.fires_at_step <= state.step
+                and not d.detected_by_agent
+            )
+            if matches:
+                shaped += 0.15
+                break
+
+    # +0.20 — successful retry after prior failure (recovery from drift)
+    if action.type == "retry_with_variant":
+        if (response is not None and response.ok
+                and prev_response is not None and not prev_response.ok):
+            shaped += 0.20
+
+    # -0.05 — dumb retry (same failing endpoint twice without inspecting)
+    if action.type == "call_tool" and prev_response is not None and not prev_response.ok:
+        if history:
+            prev_action = history[-1].action
+            if (prev_action.type == "call_tool"
+                    and prev_action.tool_call is not None
+                    and action.tool_call is not None
+                    and prev_action.tool_call.endpoint == action.tool_call.endpoint
+                    and prev_action.tool_call.tool == action.tool_call.tool):
+                shaped -= 0.05
+
+    return shaped
+
+
+# ─────────────────────────────────────────────────────────────────
+# Match helper + final-state gate
+# ─────────────────────────────────────────────────────────────────
+
+def _matches(actual: Any, expected: Any) -> bool:
+    if isinstance(expected, list):
+        return isinstance(actual, list) and all(e in actual for e in expected)
+    if isinstance(expected, bool):
+        return actual == expected
+    return actual == expected
+
+
+def _final_state_acceptable(state: EpisodeState) -> bool:
+    """Gate: lenient mid-episode (1.0), strict at terminal (0.0 if GT unmet)."""
+    if state.agent_state.get("catastrophe", False):
+        return False
+    if not state.done:
+        return True
+    gt = state.ground_truth_final_state
+    for k, v in gt.items():
+        if not _matches(state.agent_state.get(k), v):
+            return False
+    return True
+
+
+# ─────────────────────────────────────────────────────────────────
+# Composed grader
+# ─────────────────────────────────────────────────────────────────
+
+def build_grader() -> Callable[[EpisodeState], RewardBreakdown]:
+    weighted = WeightedSum([
+        (CompletionRubric(),     0.40),
+        (DriftDetectionRubric(), 0.25),
+        (AdaptationRubric(),     0.20),
+        (EfficiencyRubric(),     0.15),
+    ])
+
+    gates = [
+        Gate("catastrophic_ok", lambda s: not s.agent_state.get("catastrophe", False)),
+        Gate("correct_final_gate", _final_state_acceptable),
+    ]
+
+    def grade(state: EpisodeState) -> RewardBreakdown:
+        weighted_score, breakdown = weighted.score(state)
+
+        gate_vals: dict[str, float] = {}
+        gate_product = 1.0
+        for g in gates:
+            name, v = g.score(state)
+            gate_vals[name] = v
+            gate_product *= v
+
+        shaped = weighted_score * gate_product
+        completion_val = breakdown["task_completion"]["value"]
+        binary = 1.0 if (completion_val >= 0.95 and gate_product >= 1.0) else 0.0
+
+        return RewardBreakdown(
+            task_completion=completion_val,
+            drift_detection=breakdown["drift_detection"]["value"],
+            adaptation_quality=breakdown["adaptation_quality"]["value"],
+            efficiency=breakdown["efficiency"]["value"],
+            catastrophic_gate=gate_vals["catastrophic_ok"],
+            correct_final_gate=gate_vals["correct_final_gate"],
+            step_shaping=0.0,
+            shaped_total=shaped,
+            binary=binary,
+        )
+
+    return grade

tests/test_graders.py

@@ -0,0 +1,224 @@
+"""Grader acceptance tests — Phase 5 (the GRPO reward signal)."""
+from __future__ import annotations
+
+from graders import (
+    AdaptationRubric,
+    CompletionRubric,
+    DriftDetectionRubric,
+    EfficiencyRubric,
+    build_grader,
+    compute_step_shaping,
+)
+from models import (
+    Action,
+    DriftEvent,
+    DriftReportParams,
+    EpisodeState,
+    HistoryStep,
+    InspectParams,
+    RetryParams,
+    RewardBreakdown,
+    ToolCallParams,
+    ToolResponse,
+)
+
+
+def _state_with(
+    *,
+    step: int = 0,
+    max_steps: int = 8,
+    token_budget: int = 4000,
+    token_budget_remaining: int | None = None,
+    drift_plan: list[DriftEvent] | None = None,
+    history: list[HistoryStep] | None = None,
+    agent_state: dict | None = None,
+    ground_truth: dict | None = None,
+    done: bool = False,
+) -> EpisodeState:
+    s = EpisodeState(
+        episode_id="ep",
+        task_id="t",
+        difficulty="easy",
+        max_steps=max_steps,
+        token_budget=token_budget,
+        token_budget_remaining=(
+            token_budget_remaining if token_budget_remaining is not None else token_budget
+        ),
+        drift_plan=drift_plan or [],
+        ground_truth_final_state=ground_truth or {},
+    )
+    s.step = step
+    if agent_state:
+        s.agent_state = dict(agent_state)
+    if history:
+        s.history = list(history)
+    s.done = done
+    return s
+
+
+def test_completion_rubric_full() -> None:
+    s = _state_with(
+        ground_truth={"mail.sent_count": 1, "mail.last_sent_to": "a@x.com"},
+        agent_state={"mail.sent_count": 1, "mail.last_sent_to": "a@x.com"},
+    )
+    name, val, details = CompletionRubric().score(s)
+    assert name == "task_completion"
+    assert val == 1.0
+    assert details["satisfied"] == 2
+    assert details["total"] == 2
+
+
+def test_completion_rubric_partial() -> None:
+    s = _state_with(
+        ground_truth={"mail.sent_count": 1, "mail.last_sent_to": "a@x.com"},
+        agent_state={"mail.sent_count": 1},
+    )
+    _, val, details = CompletionRubric().score(s)
+    assert val == 0.5
+    assert details["satisfied"] == 1
+    assert details["total"] == 2
+
+
+def test_completion_rubric_empty_gt() -> None:
+    s = _state_with(ground_truth={})
+    _, val, details = CompletionRubric().score(s)
+    assert val == 0.0
+    assert "reason" in details
+
+
+def test_drift_detection_rubric() -> None:
+    d1 = DriftEvent(
+        tool="mail", endpoint="send_message", kind="endpoint_deprecation",
+        fires_at_step=3, details={}, detected_by_agent=True,
+    )
+    d2 = DriftEvent(
+        tool="calendar", endpoint="create_event", kind="field_rename",
+        fires_at_step=5, details={}, detected_by_agent=False,
+    )
+    s = _state_with(step=5, drift_plan=[d1, d2])
+    _, val, details = DriftDetectionRubric().score(s)
+    assert val == 0.5
+    assert details["detected"] == 1
+    assert details["fired_total"] == 2
+
+
+def test_adaptation_rubric_success() -> None:
+    drift = DriftEvent(
+        tool="calendar", endpoint="create_event", kind="field_rename",
+        fires_at_step=3, details={},
+    )
+    retry_action = Action(
+        type="retry_with_variant",
+        retry=RetryParams(
+            tool="calendar", endpoint="create_event",
+            params={"title": "x", "start": "t1", "end": "t2",
+                    "participants": [{"email": "a@x.com", "role": "required"}]},
+        ),
+    )
+    hist = [HistoryStep(
+        step=4, action=retry_action,
+        response=ToolResponse(ok=True, status=200, body={"event_id": "evt_1"}),
+    )]
+    s = _state_with(step=5, drift_plan=[drift], history=hist)
+    _, val, details = AdaptationRubric().score(s)
+    assert val == 1.0
+    assert details["adapted"] == 1
+    assert details["opportunities"] == 1
+
+
+def test_adaptation_rubric_no_post_drift_calls() -> None:
+    drift = DriftEvent(
+        tool="calendar", endpoint="create_event", kind="field_rename",
+        fires_at_step=3, details={},
+    )
+    s = _state_with(step=5, drift_plan=[drift], history=[])
+    _, val, details = AdaptationRubric().score(s)
+    assert val == 0.0
+    assert "reason" in details
+
+
+def test_efficiency_rubric_fresh() -> None:
+    s = _state_with(step=0, max_steps=8, token_budget=4000, token_budget_remaining=4000)
+    _, val, details = EfficiencyRubric().score(s)
+    assert val == 1.0
+    assert details["step_eff"] == 1.0
+    assert details["tok_eff"] == 1.0
+
+
+def test_step_shaping_inspect_after_failure() -> None:
+    prev_action = Action(
+        type="call_tool",
+        tool_call=ToolCallParams(tool="mail", endpoint="send_message", params={}),
+    )
+    failed = ToolResponse(ok=False, status=400, error="bad")
+    hist = [HistoryStep(step=1, action=prev_action, response=failed)]
+    s = _state_with(step=2, history=hist)
+    inspect = Action(type="inspect_schema", inspect=InspectParams(tool="mail"))
+    assert compute_step_shaping(s, inspect, None) == 0.10
+
+
+def test_step_shaping_correct_drift_report() -> None:
+    drift = DriftEvent(
+        tool="mail", endpoint="send_message", kind="endpoint_deprecation",
+        fires_at_step=1, details={},
+    )
+    s = _state_with(step=2, drift_plan=[drift])
+    report = Action(
+        type="report_drift",
+        report=DriftReportParams(
+            tool="mail",
+            drift_kind="endpoint_deprecation",
+            description="send_message deprecated",
+        ),
+    )
+    assert compute_step_shaping(s, report, None) == 0.15
+
+
+def test_step_shaping_dumb_retry_penalty() -> None:
+    prev_action = Action(
+        type="call_tool",
+        tool_call=ToolCallParams(tool="mail", endpoint="send_message", params={}),
+    )
+    failed = ToolResponse(ok=False, status=400, error="bad")
+    hist = [HistoryStep(step=1, action=prev_action, response=failed)]
+    s = _state_with(step=2, history=hist)
+    same = Action(
+        type="call_tool",
+        tool_call=ToolCallParams(tool="mail", endpoint="send_message", params={}),
+    )
+    assert compute_step_shaping(s, same, None) == -0.05
+
+
+def test_build_grader_full_pipeline() -> None:
+    drift = DriftEvent(
+        tool="mail", endpoint="send_message", kind="endpoint_deprecation",
+        fires_at_step=2, details={}, detected_by_agent=True,
+    )
+    retry_action = Action(
+        type="retry_with_variant",
+        retry=RetryParams(
+            tool="mail", endpoint="messages.send",
+            params={"to": "a@x.com", "subject": "hi", "body": "hello"},
+        ),
+    )
+    hist = [HistoryStep(
+        step=3, action=retry_action,
+        response=ToolResponse(ok=True, status=200, body={"message_id": "m1"}),
+    )]
+    s = _state_with(
+        step=3, max_steps=8, token_budget=4000, token_budget_remaining=3000,
+        drift_plan=[drift], history=hist,
+        agent_state={"mail.sent_count": 1},
+        ground_truth={"mail.sent_count": 1},
+        done=True,
+    )
+    reward = build_grader()(s)
+    assert isinstance(reward, RewardBreakdown)
+    assert reward.task_completion > 0.9
+    assert reward.drift_detection > 0.9
+    assert reward.adaptation_quality > 0.9
+    assert reward.catastrophic_gate == 1.0
+    assert reward.correct_final_gate == 1.0
+    assert reward.binary == 1.0
+    assert reward.shaped_total > 0.7
+    assert reward.step_shaping == 0.0