PIG-Demo/envs/probe.py

#!/usr/bin/env python3
"""The probe ladder: does this reward actually measure anything?

A reward function is a claim, and an unprobed one is an unchecked claim. The
ladder plays a set of deliberately-crafted policies — from doing nothing at all
to the best play we can write — and asserts that the reward orders them the way
a person would, per weighted component rather than as one blended float.

A blended float hides the failure this is built to catch: a component that is
flat across every policy is measuring nothing, and it will still move the total
because the other components move.

Run: uv run python envs/probe.py [--taskset wordle-five|alert-triage]
Runs every registered taskset by default. Exits non-zero on any violated
assertion.
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent / "wordle_five"))
sys.path.insert(0, str(Path(__file__).parent / "alert_triage"))

from wordle_five import solver as S  # noqa: E402
from wordle_five.engine import (  # noqa: E402
    MAX_GUESSES,
    Game,
    answer_for_seed,
    answers,
    score_guess,
)
from wordle_five.protocol import parse_guess  # noqa: E402
from wordle_five.reward import ROLES, WEIGHTS, Episode, metrics, score, total  # noqa: E402

from alert_triage import ladder as alert_triage_ladder  # noqa: E402

SEEDS = list(range(16))
COMPONENTS = ("solved", "economy", "consistency")


# ---------------------------------------------------------------- policies --
# Each policy is a function (game) -> reply text. They return TEXT, not moves,
# so the parser and the rejection path are exercised exactly as a model would
# exercise them. A policy that returns unparseable prose is a valid policy.


def policy_inaction(game: Game) -> str:
    """Reasons at length and never commits. The floor: must measure 0.000."""
    return "Let me think carefully about which word would be most informative here."


def policy_spammer(game: Game) -> str:
    """Multiple bracketed tokens. Only the first is read; the rest are noise."""
    return "Maybe [crane] or [slate] or [adieu]?"


def policy_crude(game: Game) -> str:
    """One valid word, forever. Legal, useless, and it repeats itself."""
    return "[tares]"


def policy_plausible(game: Game) -> str:
    """Keeps greens, drops known-absent letters. What a person does unaided."""
    pool = answers()
    if not game.history:
        return "[tares]"
    greens: dict[int, str] = {}
    absent: set[str] = set()
    for guess, pattern in game.history:
        for i, tile in enumerate(pattern):
            if tile == "G":
                greens[i] = guess[i]
            elif tile == "X":
                absent.add(guess[i])
    played = {g for g, _ in game.history}
    for word in pool:
        if word in played:
            continue
        if any(word[i] != c for i, c in greens.items()):
            continue
        if any(c in absent for c in word):
            continue
        return f"[{word}]"
    return "[tares]"


def policy_candidate_only(game: Game) -> str:
    """Only ever guesses words that could still be the answer.

    This is the policy the counterweight rewards. It never wastes a guess, and
    it is measurably slower than the oracle because it cannot buy information
    with a word that has no chance of winning.
    """
    alive = S.consistent_candidates(game.history)
    played = {g for g, _ in game.history}
    for word in alive:
        if word not in played:
            return f"[{word}]"
    return "[tares]"


def policy_oracle(game: Game) -> str:
    """Greedy maximum entropy — the reference player, and the economy denominator."""
    import numpy as np

    pool = answers()
    if not game.history:
        return f"[{pool[S._best_opener()]}]"
    alive = S.consistent_candidates(game.history)
    if not alive:
        return "[tares]"
    index = np.array([pool.index(w) for w in alive])
    return f"[{pool[S.best_guess(index)]}]"


def policy_exhaustive(game: Game) -> str:
    """Knows the answer but burns three turns on known-wrong probes first.

    Proves the turn budget actually binds: this must score strictly below the
    oracle, or `economy` is not doing its job.
    """
    if len(game.history) < 3:
        probes = ["fuzzy", "jumbo", "whelp"]
        pick = probes[len(game.history)]
        return f"[{pick if pick != game.answer else 'tares'}]"
    return f"[{game.answer}]"


POLICIES = {
    "inaction": policy_inaction,
    "spammer": policy_spammer,
    "crude": policy_crude,
    "plausible": policy_plausible,
    "candidate_only": policy_candidate_only,
    "exhaustive": policy_exhaustive,
    "oracle": policy_oracle,
}


# ------------------------------------------------------------------ runner --


def play(policy, seed: int) -> Episode:
    game = Game(seed=seed)
    for _ in range(MAX_GUESSES * 3):   # room for rejected replies
        if game.over:
            break
        guess = parse_guess(policy(game))
        if guess is None:
            game.rejected += 1
            # An unparseable reply still costs a turn of patience, or a policy
            # that never guesses would loop until the budget above runs out and
            # look identical to one that guessed badly.
            if game.rejected >= MAX_GUESSES:
                break
            continue
        game.play(guess)
    return Episode(
        answer=game.answer,
        guesses=[g for g, _ in game.history],
        patterns=[p for _, p in game.history],
        rejected=game.rejected,
        reference_depth=S.reference_depth(game.answer),
    )


def run_wordle() -> int:
    print(f"Probing wordle-five over {len(SEEDS)} seeds\n")
    results: dict[str, dict[str, float]] = {}
    extra: dict[str, dict[str, float]] = {}

    for name, policy in POLICIES.items():
        episodes = [play(policy, s) for s in SEEDS]
        scores = [score(e) for e in episodes]
        results[name] = {c: sum(s[c] for s in scores) / len(scores) for c in COMPONENTS}
        results[name]["total"] = sum(total(e) for e in episodes) / len(episodes)
        mets = [metrics(e) for e in episodes]
        extra[name] = {k: sum(m[k] for m in mets) / len(mets) for k in mets[0]}

    width = max(len(n) for n in POLICIES)
    header = f"{'policy':<{width}}  " + "  ".join(f"{c:>12}" for c in COMPONENTS) + f"  {'TOTAL':>8}"
    print(header)
    print("-" * len(header))
    for name in POLICIES:
        row = results[name]
        cells = "  ".join(f"{row[c]:>12.4f}" for c in COMPONENTS)
        print(f"{name:<{width}}  {cells}  {row['total']:>8.4f}")

    print(f"\nweights: {WEIGHTS}")
    print(f"roles:   {ROLES}")
    print("\nguesses used (mean):", {n: round(extra[n]['guesses_used'], 2) for n in POLICIES})

    failures: list[str] = []

    # 1. Doing nothing must measure exactly nothing. If a policy that never
    #    plays can collect reward, every number above it is inflated.
    if abs(results["inaction"]["total"]) > 1e-9:
        failures.append(f"inaction scored {results['inaction']['total']:.6f}, must be 0.000")

    # 2. The reference player must reach the top of both objective components,
    #    or `economy`'s denominator disagrees with the policy that defines it.
    if abs(results["oracle"]["solved"] - 1.0) > 1e-9:
        failures.append(f"oracle solved {results['oracle']['solved']:.4f}, must be 1.0")
    if abs(results["oracle"]["economy"] - 1.0) > 1e-9:
        failures.append(f"oracle economy {results['oracle']['economy']:.4f}, must be 1.0")

    # 3. The counterweight must actually bite the policy it is aimed at. The
    #    oracle buys information with guesses that cannot win; if that were
    #    free, `consistency` would be a gate wearing a counterweight's name.
    if results["oracle"]["consistency"] >= 1.0 - 1e-9:
        failures.append("oracle scored a perfect consistency — the counterweight is not in tension")

    # 4. And the policy that pays the counterweight's price must collect it.
    if abs(results["candidate_only"]["consistency"] - 1.0) > 1e-9:
        failures.append(
            f"candidate_only consistency {results['candidate_only']['consistency']:.4f}, must be 1.0"
        )

    # 5. Neither of the two good policies may dominate the other on every
    #    component. If one did, the reward would encode a single right answer
    #    and the reward editor would be theatre.
    oracle, cand = results["oracle"], results["candidate_only"]
    if all(oracle[c] >= cand[c] - 1e-9 for c in COMPONENTS):
        failures.append("oracle dominates candidate_only on every component — no real trade-off")
    if all(cand[c] >= oracle[c] - 1e-9 for c in COMPONENTS):
        failures.append("candidate_only dominates oracle on every component — no real trade-off")

    # 6. Burning turns on known-wrong probes must cost.
    if results["exhaustive"]["total"] >= results["oracle"]["total"] - 1e-9:
        failures.append("exhaustive >= oracle — the turn budget does not bind")

    # 7. No weighted component may be flat. A component with the same value for
    #    every policy is measuring nothing and still moves the total.
    for component in COMPONENTS:
        values = [results[n][component] for n in POLICIES]
        if max(values) - min(values) < 1e-9:
            failures.append(f"component '{component}' is flat across the ladder — it measures nothing")

    print()
    if failures:
        for f in failures:
            print(f"FAIL: {f}")
        return 1
    print(f"All {7} ladder assertions hold.")
    return 0


# Each taskset registers its ladder here. The alert-triage rungs live in
# envs/alert_triage/alert_triage/ladder.py beside the policies they probe.
TASKSETS = {
    "wordle-five": run_wordle,
    "alert-triage": alert_triage_ladder.run,
}


def main(argv: list[str]) -> int:
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--taskset", choices=sorted(TASKSETS), action="append")
    args = parser.parse_args(argv)
    chosen = args.taskset or list(TASKSETS)
    status = 0
    for i, name in enumerate(chosen):
        if i:
            print("\n" + "=" * 78 + "\n")
        status |= TASKSETS[name]()
    return status


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))