arena/environments/drop_table_inference/drop_table_inference/table.py

"""The monster, its drop table, and the kills that came out of it.

A task is one monster. Its table is a set of invented items, each with a rate the agent
never sees, and the log is a sample from it: a visible slice the agent reads and a held-out
continuation it is graded on. Both come out of one kill stream, so the graded kills are
literally the *next* ones the monster would have dropped — not a differently-seeded
population that could differ in ways the agent could not have anticipated.

Four properties of the generator are load-bearing, and none is decoration:

A LADDER WITH NO REPEATS. Rates are not merely "structured", they are drawn WITHOUT
REPLACEMENT from two published sets: commons are distinct multiples of 4/128, and the rest
are distinct unit fractions from 1/128 down to 1/24576. Distinctness is the first half of
why the tail is learnable. Nothing in a kill log distinguishes one item that dropped zero
times from another — any estimator that is a function of (count, kills) alone must give them
the same number, and if that number is also the same across TASKS then a memorised constant
is a perfect tail estimate and the term is inert. That is exactly what the first version of
this environment did: its reference gave every zero-count item (0+0.5)/(kills+0.5(n+1)), a
value with four distinct settings across ninety-six tasks, so a reply that never opened the
kill log scored 1.000 on the tail. With no repeats, the rungs the OBSERVED items occupy are
the rungs the unobserved ones cannot: a table whose counts pin 1/128 and 1/384 leaves its
silent item somewhere below, and one that pins nothing above 1/3072 leaves it somewhere much
higher.

Distinctness alone was not enough, and the number that says so is 8%: with the ladder and
the elimination but no price, the reference beat the best memorised constant on the tail
loss by eight percent, and a frequency copy with a constant floor held 0.86 of the term.
Rungs eliminate, but two SILENT items are still interchangeable, and the irreducible spread
inside that pair swamped everything the elimination revealed. That is what VALUE_SCALE
below is for.

A PRICE THAT RANKS WHAT THE LOG CANNOT. Every item carries a shop value, lognormal around
SCALE * (1/rate) ** EXPONENT — rare drops are worth more, which is a fact about the genre
rather than anybody's copy. It is the only channel that separates two silent items, and
separating them is what turns the tail from a lottery into an inference: with it the
reference beats the best memorised constant on the tail by 4.75x rather than by 8%, and the
frequency-copy-with-a-floor attack falls from 0.86 of the term to 0.28. The scatter is
deliberate and it is large — sigma 0.55 in log value against 3x rung spacing — so the price
ranks the candidates and never decides them, and an agent that reads only the price scores
0.46 where one that reads price, counts and leftover rungs together scores 1.000.

A ZERO-COUNT ITEM, ALWAYS. The kill stream is rejected and redrawn until at least one item
on the list dropped zero times in the visible slice. That case is the entire environment:
the item list says the item exists, the log says it never appeared, and an estimate of 0 is
an infinitely strong claim about something the graded kills will demonstrate is false. A
task where every item happened to show up does not pose the question, so it is not used.
The rejection is part of the model — `estimate._posterior` divides it back out,
because conditioning on "some item was silent" is evidence about the table and an estimator
that ignores it is leaving information in the prompt.

HELD-OUT KILLS UNTIL EVERY ITEM APPEARS. The item list is a promise that all of these drop;
the graded slice has to keep it, or an estimate of zero for the rarest item costs nothing.
It is also what keeps the counterweight finite: `estimate.restraint` divides by what the
continuation produced, and an item the continuation never produced would make it infinite.

No wiki content: every item, monster and word here is invented. The MECHANICS are the
borrowed part, and mechanics are facts.
"""

from __future__ import annotations

import bisect
import math
import random
from dataclasses import dataclass

ADJECTIVES = [
    "Ashen", "Brackish", "Cindered", "Duskbound", "Ember", "Frostbit", "Gilded", "Hollow",
    "Ivory", "Jagged", "Kelpish", "Lodestone", "Murkwrought", "Nettled", "Pitchblack",
    "Quarried", "Saltworn", "Tarnished", "Umbral", "Verdigris",
]
NOUNS = [
    "talon", "sigil", "vertebra", "censer", "ingot", "lantern", "chitin", "warhorn",
    "tessera", "phial", "reliquary", "scale", "cog", "wick", "grimoire", "spur",
    "anvil-shard", "orb", "tabard", "quill",
]
MONSTER_PREFIX = ["Marrow", "Gloom", "Rime", "Slagborn", "Fenwater", "Barrow", "Cinder", "Hush"]
MONSTER_KIND = ["Warden", "Basilisk", "Revenant", "Colossus", "Hierophant", "Drake", "Herald", "Ossifier"]

# The common grid, and it is coarse on purpose. Adjacent numerators sit 33-50% apart, so a
# few hundred kills can actually tell them apart — a continuous numerator would put the
# notches inside the sampling error, the structure would be real but unresolvable, and
# "spot the grid" would be a slogan rather than a strategy.
COMMON_DENOMINATOR = 128
COMMON_NUMERATORS = (8, 12, 16, 24, 32, 40, 48)

# The rare ladder: unit fractions spaced by roughly a factor of two. No sample of 600 kills
# resolves 1/1024 from 1/2048 on its own evidence, and that is the point — down here the
# information is in which rungs the OTHER items have already taken.
LADDER_DENOMINATORS = (128, 384, 1024, 3072, 8192, 24576)

# How many of each. Both are published to the agent, because the reference estimator uses
# them and a reference the agent cannot compute is an unreachable ceiling wearing a hat.
N_COMMON = (3, 4)
N_LADDER = (4, 5, 6)

# The price law, and it is the reason the tail is a skill rather than a lottery.
#
# Two items that both dropped zero times are EXCHANGEABLE on the kill log: no statistic of
# the counts distinguishes them, so every estimator must give them the same number, and
# the reference's advantage over a memorised constant collapses to whatever the free-rung
# set alone reveals. Measured, that was four to eight percent — the reference beat the best
# constant on the tail by less than a tenth, which is not an environment, it is a coin.
#
# A price breaks the tie, and it is the OSRS-shaped way to break it: rare drops are worth
# more. `value` is lognormal around SCALE * (1/rate) ** EXPONENT, so it ranks the silent
# items without deciding them — at 3x rung spacing and this sigma, the price alone gets an
# adjacent pair the right way round about six times in seven and gets nothing for free two
# rungs down. The agent has to combine it with the counts and the leftover rungs, which is
# the entire task.
VALUE_SCALE = 15.0
VALUE_EXPONENT = 0.8
VALUE_LOG_SD = 0.55

# The share of kills that drop something. Bounded away from 1 so "no drop" is always a
# substantial category: it is where the mass an agent does not assign to items has to go,
# and it is what makes the sum-to-at-most-one constraint bite instead of being cosmetic.
MIN_ITEM_MASS = 0.55
MAX_ITEM_MASS = 0.90

# Held-out kills stop when every item has been seen; this caps the walk in case a table
# with a 1/24576 tertiary draws a very long silence. Hitting it costs an item's evidence
# rather than anything worse, and at twelve times the requested slice it is a 1e-8 event.
HELD_OUT_CEILING = 12


@dataclass(frozen=True)
class DropTable:
    """One monster's table. `rates` is in DISPLAY order, deliberately shuffled: listing
    items by rarity would hand over the ordering the agent is being asked to recover."""

    monster: str
    rates: dict[str, float]
    values: dict[str, int]

    @property
    def items(self) -> list[str]:
        return list(self.rates)

    @property
    def nothing(self) -> float:
        return 1.0 - sum(self.rates.values())


@dataclass(frozen=True)
class Stream:
    """One task: the table, and both slices as tallies.

    Tallies rather than kill sequences because the tally is the sufficient statistic — the
    order kills came in carries no information about the table, so serialising 5,000 lines
    would cost tokens and teach nothing.
    """

    table: DropTable
    visible: dict[str, int]
    visible_kills: int
    held: dict[str, int]
    held_kills: int

    @property
    def visible_nothing(self) -> int:
        return self.visible_kills - sum(self.visible.values())

    @property
    def held_nothing(self) -> int:
        return self.held_kills - sum(self.held.values())


def _names(rng: random.Random, count: int) -> list[str]:
    seen: list[str] = []
    while len(seen) < count:
        name = f"{rng.choice(ADJECTIVES)} {rng.choice(NOUNS)}"
        if name not in seen:
            seen.append(name)
    return seen


def build_table(seed: int) -> DropTable:
    """The table alone, from the seed. Stable across kill-stream redraws so that a task's
    identity is its monster, not whichever sample happened to be accepted.

    `sample` rather than `choices` in both tiers: with replacement, two items could share a
    rung, the leftover-rung deduction would be unsound, and the tail would go back to being
    a constant nobody has to read the log for.
    """
    rng = random.Random(seed)
    monster = f"{rng.choice(MONSTER_PREFIX)} {rng.choice(MONSTER_KIND)}"
    n_common = rng.choice(N_COMMON)
    n_ladder = rng.choice(N_LADDER)
    names = _names(rng, n_common + n_ladder)

    # Rejection on total mass, with the tier COUNTS fixed outside the loop: redrawing the
    # shape as well would make the split's prior depend on how often that shape lands in
    # the mass window, and the reference's prior — which is P(n_common) * P(n_ladder) —
    # would then be subtly wrong on every task.
    while True:
        numerators = rng.sample(COMMON_NUMERATORS, n_common)
        denominators = rng.sample(LADDER_DENOMINATORS, n_ladder)
        rates = [n / COMMON_DENOMINATOR for n in numerators] + [1.0 / d for d in denominators]
        if MIN_ITEM_MASS <= sum(rates) <= MAX_ITEM_MASS:
            break
    order = list(zip(names, rates))
    rng.shuffle(order)
    table = dict(order)
    # Prices drawn AFTER the shuffle so the display order carries nothing, and from the
    # rate rather than the tier, so a common at 8/128 and a ladder item at 1/128 are priced
    # identically — the price is evidence about the RATE, never about which tier an item
    # came from.
    values = {
        name: max(1, round(VALUE_SCALE * (1.0 / rate) ** VALUE_EXPONENT
                           * math.exp(VALUE_LOG_SD * rng.gauss(0.0, 1.0))))
        for name, rate in table.items()
    }
    return DropTable(monster=monster, rates=table, values=values)


def _tally(rng: random.Random, table: DropTable, kills: int, until_covered: bool) -> tuple[dict[str, int], int]:
    """Roll kills, counting drops. With `until_covered`, keep rolling past `kills` until
    every item has appeared — see the module docstring for why the graded slice owes the
    item list that."""
    items = table.items
    cumulative, running = [], 0.0
    for item in items:
        running += table.rates[item]
        cumulative.append(running)
    counts = {item: 0 for item in items}
    rolled, ceiling = 0, kills * HELD_OUT_CEILING
    while rolled < kills or (until_covered and not all(counts.values()) and rolled < ceiling):
        index = bisect.bisect(cumulative, rng.random())
        if index < len(items):
            counts[items[index]] += 1
        rolled += 1
    return counts, rolled


def build_slices(seed: int, visible_kills: int, held_out_kills: int) -> Stream:
    """The two slices of one task, from one seed.

    The kill stream is redrawn — table held fixed — until the visible slice has a
    zero-count item, because a task where every item showed up does not pose the question
    this environment exists to ask. Rejection is on the visible slice only; the held-out
    continuation is whatever the accepted stream produced next.
    """
    table = build_table(seed)
    for attempt in range(256):
        rng = random.Random((seed << 9) + attempt)
        visible, _ = _tally(rng, table, visible_kills, until_covered=False)
        if all(visible.values()):
            continue
        held, held_kills = _tally(rng, table, held_out_kills, until_covered=True)
        return Stream(table, visible, visible_kills, held, held_kills)
    raise RuntimeError(f"no stream with an unobserved item for seed {seed}")