PIG-Demo/envs/wordle_five/wordle_five/reward.py

"""The reward. This file is quoted verbatim on the demo page, so it is written
to be read by someone who does not write Python.

Three components, weights summing to 1.0. Two of them pull toward winning. The
third pulls against them on purpose.

# region: pig-demo/reward
"""

from __future__ import annotations

from dataclasses import dataclass

from .engine import GREEN, WORD_LENGTH, score_guess

SOLVED_WEIGHT = 0.50
ECONOMY_WEIGHT = 0.30
CONSISTENCY_WEIGHT = 0.20


@dataclass
class Episode:
    """Everything the reward needs from one played game."""

    answer: str
    guesses: list[str]
    patterns: list[str]
    rejected: int
    reference_depth: int
    max_guesses: int = 6


def solved(ep: Episode) -> float:
    """Did it win? The objective, and half the reward."""
    return 1.0 if ep.patterns and ep.patterns[-1] == GREEN * WORD_LENGTH else 0.0


def economy(ep: Episode) -> float:
    """How few turns it took, against how few the reference player needed.

    Always a ratio against a reference on the SAME hidden word, never an
    absolute count — otherwise the reward would mostly measure whether the word
    happened to be easy, and a model would be punished for a hard draw.
    """
    if not solved(ep):
        return 0.0
    return min(1.0, ep.reference_depth / max(1, len(ep.guesses)))


def consistency(ep: Episode) -> float:
    """The share of guesses that were still possible answers when they were made.

    This is the counterweight, and it is genuinely in tension with the other
    two. A player maximising information will deliberately guess words that
    CANNOT win, because a word that splits the remaining candidates evenly
    teaches it more than a word that might happen to be right. That is good
    play, and it costs consistency. The reference solver scores below 1.0 here.

    So the three components describe a real trade: win, win fast, and do not
    spend turns on moves that could not have won. You cannot max all three, and
    which one you weight is a decision about what you actually want — which is
    the entire argument this demo exists to make.

    Scored over every turn SPENT, not every guess accepted — see below. Without
    that, a policy which never produces a parseable word scores a vacuous 1.0
    here (no guesses, therefore no contradictions) and collects 0.2 for doing
    nothing at all.
    """
    spent = len(ep.guesses) + ep.rejected
    if spent == 0:
        return 0.0

    viable = 0
    for i, guess in enumerate(ep.guesses):
        # A guess is viable iff it is consistent with every earlier piece of
        # feedback — that is, it was still a possible answer at the moment it
        # was played.
        if all(score_guess(ep.guesses[k], guess) == ep.patterns[k] for k in range(i)):
            viable += 1

    # The denominator is every turn SPENT, including replies the game refused.
    # Counting only legal guesses would hand a free 1.0 to a policy that plays
    # one word and then jams the parser five times: one guess, no
    # contradictions, perfect score. Measuring against turns spent asks the
    # question that actually matters — of the attempts you made, how many were
    # a move that could have won?
    return viable / spent


def score(ep: Episode) -> dict[str, float]:
    """The weighted components, unweighted. The caller applies the weights."""
    return {
        "solved": solved(ep),
        "economy": economy(ep),
        "consistency": consistency(ep),
    }


def total(ep: Episode) -> float:
    parts = score(ep)
    return (
        parts["solved"] * SOLVED_WEIGHT
        + parts["economy"] * ECONOMY_WEIGHT
        + parts["consistency"] * CONSISTENCY_WEIGHT
    )


# endregion: pig-demo/reward


def metrics(ep: Episode) -> dict[str, float]:
    """Diagnostics. Weight zero — reported, never summed into the reward."""
    return {
        "guesses_used": float(len(ep.guesses)),
        "rejected_replies": float(ep.rejected),
        "reference_depth": float(ep.reference_depth),
        "turns_granted": float(ep.max_guesses),
        "inconsistent_guesses": float(round((1.0 - consistency(ep)) * len(ep.guesses))),
    }


WEIGHTS = {
    "solved": SOLVED_WEIGHT,
    "economy": ECONOMY_WEIGHT,
    "consistency": CONSISTENCY_WEIGHT,
}

ROLES = {
    "solved": "objective",
    "economy": "objective",
    "consistency": "counterweight",
}