feat(graders): implement task graders with metrics and registry integration

Introduce a modular grader system for evaluating episode trajectories across tasks 1-4, including a registry for callable graders, metric computation functions, and integration into the inference pipeline for deterministic scoring. Add comprehensive unit tests covering basic functionality, invalid action penalties, and range validation. Update constants with normalization scales and penalties. Include new documentation directory for related specs.

docs/graders.md

@@ -0,0 +1,29 @@
+# Grader Implementation Details
+
+The thermal plant environment evaluates external LLM agents using deterministic graders implemented in `graders/task*_grader.py`.
+
+## Core Principles
+- **Deterministic**: Graders rely only on mathematical computations over `trajectory.steps[].raw_state` (full precision) and not on rounded observations.
+- **Normalization**: Computed metrics are mapped to a normalized range `[0,1]` using linear clamping. The default metric scale is `0.5`, defined in `utils/constants.py` under the dictionary `METRIC_NORM_SCALES`.
+- **Score constraints**: All final scores are bounded algebraically and strictly clamped to `[0,1]` via `max(0.0, min(1.0, score))`.
+
+## Shared Metrics (`graders/_metrics.py`)
+All graders share deterministic extraction functions:
+- **TE** (Tracking Error): `(1/N) * sum(|P - L|)`
+- **OS** (Overshoot): `max(0, P - L)` over all steps
+- **SV** (Safety Violations): `(1/N) * sum(max(0, T - 1.0) + max(0, Pr - 1.0))`
+- **OC** (Oscillation): `(1/(N-1)) * sum(|U_t - U_{t-1}|)`
+- **SL** (Stress Level): `(1/N) * sum(S)`
+- **Failure Flag (FF)**: `1` if trajectory ended due to a catastrophic state violation of temperature, pressure, or stress, else `0`
+- **Invalid Output Penalty**: Graders subtract `INVALID_ACTION_PENALTY` (default `0.2`) from the computed task score for *each step* where the LLM's payload was unparseable. 
+
+## Special Operations
+- Task 4 (Fault Recovery) computes **RT**, the Recovery Time. This defines the step index when the control tracks within bounds again (`|P - L| < 0.1` and `T < 1.0`). If recovery fails completely, it returns `N`. 
+RT is normalized directly as `Clamp(RT / N)` instead of using the generic `0.5` scaling value.
+
+## Debugging and Diagnostics
+The `inference.py` runner includes an optional debug mode. Call inference with `DEBUG=1` to print un-rounded raw values, actual metric outputs, and penalty flags to `stderr`. 
+```bash
+DEBUG=1 python inference.py
+```
+This data is printed only to `stderr` and never `stdout`, protecting the `[START]/[STEP]/[END]` requirement established by the framework specification.

graders/_metrics.py

@@ -0,0 +1,163 @@
+"""Metric extraction helpers for graders.
+
+Functions accept an `EpisodeTrajectory` (utils.schemas.EpisodeTrajectory) and compute
+the canonical raw metrics described in context/06_tasks_and_graders.md.
+"""
+
+from typing import Dict, List, Any
+import math
+
+import utils.constants as C
+
+
+def _safe_get(raw: Dict[str, Any], key: str) -> float:
+    try:
+        return float(raw.get(key, 0.0))
+    except Exception:
+        return 0.0
+
+
+def _extract(trajectory) -> Dict[str, List[float]]:
+    P = []
+    L = []
+    T = []
+    Pr = []
+    U = []
+    F = []
+    S = []
+    for step in getattr(trajectory, "steps", []):
+        raw = getattr(step, "raw_state", {}) or {}
+        P.append(_safe_get(raw, "P"))
+        L.append(_safe_get(raw, "L"))
+        T.append(_safe_get(raw, "T"))
+        Pr.append(_safe_get(raw, "Pr"))
+        U.append(_safe_get(raw, "U"))
+        F.append(_safe_get(raw, "F"))
+        S.append(_safe_get(raw, "S"))
+    return {"P": P, "L": L, "T": T, "Pr": Pr, "U": U, "F": F, "S": S, "N": len(P)}
+
+
+def compute_TE(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    return float(sum(abs(p - l) for p, l in zip(s["P"], s["L"])) / N)
+
+
+def compute_OS(trajectory) -> float:
+    s = _extract(trajectory)
+    if s["N"] == 0:
+        return 0.0
+    diffs = [p - l for p, l in zip(s["P"], s["L"])]
+    return float(max(0.0, max(diffs)))
+
+
+def compute_SV(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    total = 0.0
+    for t, pr in zip(s["T"], s["Pr"]):
+        total += max(0.0, t - 1.0) + max(0.0, pr - 1.0)
+    return float(total / N)
+
+
+def compute_OC(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N <= 1:
+        return 0.0
+    diffs = [abs(u1 - u0) for u0, u1 in zip(s["U"][:-1], s["U"][1:])]
+    return float(sum(diffs) / max(1, (N - 1)))
+
+
+def compute_SL(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    return float(sum(s["S"]) / N)
+
+
+def compute_LP(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    return float(sum(max(0.0, abs(p - l) - 0.1) for p, l in zip(s["P"], s["L"])) / N)
+
+
+def compute_LS(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    return float(sum(max(0.0, st - 0.5) for st in s["S"]) / N)
+
+
+def compute_EMB(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    count = sum(1 for st, t in zip(s["S"], s["T"]) if st < 0.5 and t < 1.0)
+    return float(count / N)
+
+
+def compute_RT(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    for idx, (p, l, t) in enumerate(zip(s["P"], s["L"], s["T"])):
+        if abs(p - l) < 0.1 and t < 1.0:
+            return float(idx + 1)
+    return float(N)
+
+
+def compute_RR(trajectory) -> float:
+    s = _extract(trajectory)
+    N = s["N"]
+    if N == 0:
+        return 0.0
+    return float(sum(max(0.0, t - 1.0) for t in s["T"]) / N)
+
+
+def compute_failure_flag(trajectory) -> int:
+    # FF = 1 if any termination due to catastrophic T/Pr/S breach else 0
+    for step in getattr(trajectory, "steps", []):
+        err = getattr(step, "error", None)
+        if err and isinstance(err, str) and "Catastrophic" in err:
+            return 1
+        raw = getattr(step, "raw_state", {}) or {}
+        if getattr(step, "done", False):
+            if _safe_get(raw, "T") > C.FAIL_T or _safe_get(raw, "Pr") > C.FAIL_PR or _safe_get(raw, "S") > C.FAIL_S:
+                return 1
+    return 0
+
+
+def compute_invalid_count(trajectory) -> int:
+    count = 0
+    for step in getattr(trajectory, "steps", []):
+        if getattr(step, "env_invalid_action", False):
+            count += 1
+            continue
+        parsed = getattr(step, "parsed_action", None)
+        if parsed and getattr(parsed, "invalid_output", False):
+            count += 1
+    return count
+
+
+def normalize_metrics(metrics: Dict[str, float], scales: Dict[str, float] = None) -> Dict[str, float]:
+    if scales is None:
+        scales = getattr(C, "METRIC_NORM_SCALES", {}) or {}
+    normalized: Dict[str, float] = {}
+    for k, v in metrics.items():
+        scale = float(scales.get(k, 0.5)) if scales else 0.5
+        if scale <= 0:
+            normalized[k] = 0.0
+        else:
+            normalized[k] = float(max(0.0, min(1.0, v / scale)))
+    return normalized

graders/registry.py

@@ -0,0 +1,20 @@
+"""Grader registry mapping task ids to grader callables.
+
+Expose `grader_registry()` which returns a dict mapping task id strings
+to functions with signature `grade(trajectory) -> float`.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Dict
+
+from graders import task1_grader, task2_grader, task3_grader, task4_grader
+
+
+def grader_registry() -> Dict[str, Callable]:
+	return {
+		"task1": task1_grader.grade,
+		"task2": task2_grader.grade,
+		"task3": task3_grader.grade,
+		"task4": task4_grader.grade,
+	}

graders/task1_grader.py

@@ -0,0 +1,39 @@
+"""Task 1 grader: Stable baseline operation.
+
+Grader implements the formula from context/06_tasks_and_graders.md and
+returns a float in [0,1].
+"""
+
+from __future__ import annotations
+
+from ._metrics import (
+	compute_TE,
+	compute_SV,
+	compute_OC,
+	compute_failure_flag,
+	compute_invalid_count,
+	normalize_metrics,
+)
+import utils.constants as C
+
+
+def grade(trajectory) -> float:
+	"""Compute Task 1 score from trajectory and return a float in [0,1]."""
+	te = compute_TE(trajectory)
+	sv = compute_SV(trajectory)
+	oc = compute_OC(trajectory)
+	ff = compute_failure_flag(trajectory)
+	invalid_count = compute_invalid_count(trajectory)
+
+	norm = normalize_metrics({"TE": te, "SV": sv, "OC": oc})
+
+	score = 1.0
+	score -= 0.5 * norm.get("TE", 0.0)
+	score -= 0.2 * norm.get("SV", 0.0)
+	score -= 0.2 * norm.get("OC", 0.0)
+	score -= 0.3 * (1 if ff else 0)
+
+	# Apply per-invalid-step penalty
+	score -= invalid_count * getattr(C, "INVALID_ACTION_PENALTY", 0.2)
+
+	return float(max(0.0, min(1.0, score)))

graders/task2_grader.py

@@ -0,0 +1,39 @@
+"""Task 2 grader: Load following.
+
+Formula from context/06_tasks_and_graders.md. Returns a float in [0,1].
+"""
+
+from __future__ import annotations
+
+from ._metrics import (
+	compute_TE,
+	compute_OS,
+	compute_LP,
+	compute_SV,
+	compute_failure_flag,
+	compute_invalid_count,
+	normalize_metrics,
+)
+import utils.constants as C
+
+
+def grade(trajectory) -> float:
+	te = compute_TE(trajectory)
+	os_ = compute_OS(trajectory)
+	lp = compute_LP(trajectory)
+	sv = compute_SV(trajectory)
+	ff = compute_failure_flag(trajectory)
+	invalid_count = compute_invalid_count(trajectory)
+
+	norm = normalize_metrics({"TE": te, "OS": os_, "LP": lp, "SV": sv})
+
+	score = 1.0
+	score -= 0.4 * norm.get("TE", 0.0)
+	score -= 0.2 * norm.get("OS", 0.0)
+	score -= 0.2 * norm.get("LP", 0.0)
+	score -= 0.2 * norm.get("SV", 0.0)
+	score -= 0.3 * (1 if ff else 0)
+
+	score -= invalid_count * getattr(C, "INVALID_ACTION_PENALTY", 0.2)
+
+	return float(max(0.0, min(1.0, score)))

graders/task3_grader.py

@@ -0,0 +1,39 @@
+"""Task 3 grader: Preemptive constraint management.
+
+Implements specification from context/06_tasks_and_graders.md.
+"""
+
+from __future__ import annotations
+
+from ._metrics import (
+	compute_TE,
+	compute_LS,
+	compute_SV,
+	compute_EMB,
+	compute_failure_flag,
+	compute_invalid_count,
+	normalize_metrics,
+)
+import utils.constants as C
+
+
+def grade(trajectory) -> float:
+	te = compute_TE(trajectory)
+	ls = compute_LS(trajectory)
+	sv = compute_SV(trajectory)
+	emb = compute_EMB(trajectory)
+	ff = compute_failure_flag(trajectory)
+	invalid_count = compute_invalid_count(trajectory)
+
+	norm = normalize_metrics({"TE": te, "LS": ls, "SV": sv})
+
+	score = 1.0
+	score -= 0.3 * norm.get("TE", 0.0)
+	score -= 0.3 * norm.get("LS", 0.0)
+	score -= 0.2 * norm.get("SV", 0.0)
+	score += 0.2 * float(emb)
+	score -= 0.3 * (1 if ff else 0)
+
+	score -= invalid_count * getattr(C, "INVALID_ACTION_PENALTY", 0.2)
+
+	return float(max(0.0, min(1.0, score)))

graders/task4_grader.py

@@ -0,0 +1,43 @@
+"""Task 4 grader: Fault recovery with degradation.
+
+Implements the recovery-time based grader from context/06_tasks_and_graders.md.
+"""
+
+from __future__ import annotations
+
+from ._metrics import (
+	compute_RT,
+	compute_RR,
+	compute_SV,
+	compute_OC,
+	compute_failure_flag,
+	compute_invalid_count,
+	normalize_metrics,
+)
+import utils.constants as C
+
+
+def grade(trajectory) -> float:
+	rt = compute_RT(trajectory)
+	rr = compute_RR(trajectory)
+	sv = compute_SV(trajectory)
+	oc = compute_OC(trajectory)
+	ff = compute_failure_flag(trajectory)
+	invalid_count = compute_invalid_count(trajectory)
+
+	N = len(getattr(trajectory, "steps", []))
+	norm_rt = max(0.0, min(1.0, rt / N)) if N > 0 else 0.0
+
+	norm = normalize_metrics({"RR": rr, "SV": sv, "OC": oc})
+
+	# Note: spec uses Norm_RT = Clamp(RT / N) where RT in 1..N
+	score = 1.0
+	score -= 0.3 * norm_rt
+	score -= 0.2 * norm.get("RR", 0.0)
+	score -= 0.2 * norm.get("SV", 0.0)
+	score -= 0.2 * norm.get("OC", 0.0)
+	score -= 0.3 * (1 if ff else 0)
+
+	score -= invalid_count * getattr(C, "INVALID_ACTION_PENALTY", 0.2)
+
+	return float(max(0.0, min(1.0, score)))

inference.py

@@ -282,6 +282,32 @@ def main() -> None:
 		loop_error = str(exc)
 	finally:
 		score = compute_normalized_score(rewards, max_steps)
+		# If grader registry is available, prefer grader-computed score (deterministic, uses raw internals)
+		try:
+			from graders.registry import grader_registry
+			registry = grader_registry()
+			grader_fn = registry.get(TASK_NAME)
+			if grader_fn is not None:
+				# grader functions return float score in [0,1]
+				try:
+					gscore = float(grader_fn(trajectory))
+					# clamp defensively
+					score = max(0.0, min(1.0, gscore))
+					if DEBUG:
+						from graders._metrics import compute_TE, compute_OS, compute_SV, compute_OC, compute_SL, compute_LP, compute_LS, compute_EMB, compute_RT, compute_RR, compute_failure_flag, compute_invalid_count
+						print(
+							f"[DEBUG] Grader Metrics: TE={compute_TE(trajectory):.3f} OS={compute_OS(trajectory):.3f} "
+							f"SV={compute_SV(trajectory):.3f} OC={compute_OC(trajectory):.3f} SL={compute_SL(trajectory):.3f} "
+							f"FF={compute_failure_flag(trajectory)} Invalids={compute_invalid_count(trajectory)}",
+							file=sys.stderr, flush=True
+						)
+				except Exception as e:
+					# Ignore grader errors and fall back to reward-based score
+					if DEBUG:
+						print(f"[DEBUG] Grader exception: {e}", file=sys.stderr, flush=True)
+		except Exception:
+			# missing registry or other import error -> keep reward-based score
+			pass
 		success = score >= SUCCESS_SCORE_THRESHOLD
 		termination_reason = determine_termination_reason(loop_error=loop_error, done=done, steps_taken=steps_taken, max_steps=max_steps)
 		trajectory.summary = TrajectorySummary(

tests/unit/test_graders.py

@@ -0,0 +1,80 @@
+import math
+
+from utils.schemas import EpisodeTrajectory, TrajectoryStep, ParsedAction
+
+from graders.task1_grader import grade as grade_t1
+from graders.task2_grader import grade as grade_t2
+from graders.task3_grader import grade as grade_t3
+from graders.task4_grader import grade as grade_t4
+
+
+def _make_parsed_action(u=0.5, f=0.5, invalid=False):
+    return ParsedAction(
+        u_target=float(u),
+        f_target=float(f),
+        source="json",
+        used_fallback=False,
+        invalid_output=invalid,
+        penalty_applied=0.0,
+        raw_text="",
+        parse_error=None,
+    )
+
+
+def _make_step(idx: int, raw_state: dict, invalid=False, done=False, error=None):
+    pa = _make_parsed_action(u=raw_state.get("U", 0.5), f=raw_state.get("F", 0.5), invalid=invalid)
+    return TrajectoryStep(
+        step=idx,
+        raw_llm_text="",
+        parsed_action=pa,
+        canonical_action={"U_target": pa.u_target, "F_target": pa.f_target},
+        observation={k: float(v) for k, v in raw_state.items()},
+        raw_state={k: float(v) for k, v in raw_state.items()},
+        reward=0.0,
+        done=done,
+        error=error,
+        env_invalid_action=invalid,
+        invalid_penalty_applied=0.0,
+    )
+
+
+def test_task1_basic_deterministic_and_range():
+    traj = EpisodeTrajectory(task="task1", benchmark="test", model="tester")
+    # Two-step simple trajectory
+    s1 = {"P": 0.6, "L": 0.5, "T": 0.5, "Pr": 0.5, "U": 0.5, "F": 0.5, "S": 0.1}
+    s2 = {"P": 0.5, "L": 0.7, "T": 0.4, "Pr": 0.3, "U": 0.4, "F": 0.4, "S": 0.2}
+    traj.steps.append(_make_step(1, s1))
+    traj.steps.append(_make_step(2, s2))
+
+    score1 = grade_t1(traj)
+    score2 = grade_t1(traj)
+    assert isinstance(score1, float)
+    assert 0.0 <= score1 <= 1.0
+    assert math.isclose(score1, score2, rel_tol=1e-9)
+
+
+def test_invalid_penalty_reduces_score():
+    traj = EpisodeTrajectory(task="task1", benchmark="test", model="tester")
+    s = {"P": 0.6, "L": 0.5, "T": 0.5, "Pr": 0.5, "U": 0.5, "F": 0.5, "S": 0.1}
+    traj.steps.append(_make_step(1, s, invalid=False))
+    traj.steps.append(_make_step(2, s, invalid=False))
+    base = grade_t1(traj)
+
+    traj2 = EpisodeTrajectory(task="task1", benchmark="test", model="tester")
+    traj2.steps.append(_make_step(1, s, invalid=True))
+    traj2.steps.append(_make_step(2, s, invalid=False))
+    penalized = grade_t1(traj2)
+
+    assert penalized <= base
+
+
+def test_task2_and_task3_and_task4_basic_ranges():
+    # Build a short trajectory that is safe
+    traj = EpisodeTrajectory(task="task2", benchmark="test", model="tester")
+    for i in range(4):
+        s = {"P": 0.6, "L": 0.6 + 0.0 * i, "T": 0.5, "Pr": 0.5, "U": 0.5, "F": 0.5, "S": 0.05}
+        traj.steps.append(_make_step(i + 1, s))
+
+    assert 0.0 <= grade_t2(traj) <= 1.0
+    assert 0.0 <= grade_t3(traj) <= 1.0
+    assert 0.0 <= grade_t4(traj) <= 1.0

tests/unit/test_graders_edgecases.py

@@ -0,0 +1,25 @@
+import math
+from utils.schemas import EpisodeTrajectory
+from graders.task1_grader import grade as grade_t1
+from graders.task4_grader import grade as grade_t4
+from graders._metrics import compute_RT, compute_TE
+
+def test_empty_trajectory():
+    traj = EpisodeTrajectory(task="task1", benchmark="test", model="tester")
+    score = grade_t1(traj)
+    assert isinstance(score, float)
+    assert 0.0 <= score <= 1.0
+    
+def test_task4_never_recovered():
+    traj = EpisodeTrajectory(task="task4", benchmark="test", model="tester")
+    from tests.unit.test_graders import _make_step
+    # P is 0.5, L is 0.7, diff = 0.2 >= 0.1 -> never recovered
+    traj.steps.append(_make_step(1, {"P": 0.5, "L": 0.7, "T": 0.5, "Pr": 0.5, "U": 0.5, "F": 0.5, "S": 0}))
+    traj.steps.append(_make_step(2, {"P": 0.5, "L": 0.7, "T": 0.5, "Pr": 0.5, "U": 0.5, "F": 0.5, "S": 0}))
+    
+    rt = compute_RT(traj)
+    assert math.isclose(rt, 2.0)
+    score = grade_t4(traj)
+    # Norm_RT should be 1.0 (2 / 2), penalty = 0.3 * 1.0 = 0.3
+    # other penalties 0
+    assert score <= 0.7

utils/constants.py

@@ -264,4 +264,18 @@ PARSER_DEFAULT_U = 0.5
 PARSER_DEFAULT_F = 0.5
 INCLUDE_PARSE_ERROR_IN_STEP = True
 
+# Grader normalization scales and penalties
+METRIC_NORM_SCALES = {
+	"TE": 0.5,
+	"OS": 0.5,
+	"SV": 0.5,
+	"OC": 0.5,
+	"SL": 0.5,
+	"LP": 0.5,
+	"LS": 0.5,
+	"RR": 0.5,
+}
+
+INVALID_ACTION_PENALTY = 0.2
+