train/heldout_prep.py

"""Brain-2 heldout eval set: agent-authored automation requests for the 10
reserved homes (data/homes/heldout/), validated against the live inventory with
taskgen's exact rules. These homes never appear in warmup or GRPO data.

Output: data/processed/grpo/heldout.parquet (same columns the probe/eval reads).

Run: uv run --group sandbox --group train python train/heldout_prep.py
"""

import json
import pathlib
import sys

_REPO = pathlib.Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_REPO))

HOMES = _REPO / "data" / "homes" / "heldout"
REQS = _REPO / "data" / "processed" / "route_requests" / "heldout"
OUT = _REPO / "data" / "processed" / "grpo"


def main() -> None:
    from datasets import Dataset

    from sandbox.harness import SandboxHome

    rows, dropped = [], 0
    for path in sorted(HOMES.glob("*.yaml")):
        home_id = path.stem
        reqs = json.loads((REQS / f"{home_id}.json").read_text())
        inventory = path.read_text()
        home = SandboxHome(inventory)
        ents = home.inventory.get("entities", [])
        ids = {e["id"] for e in ents}
        name_of = {e["id"]: e.get("name", e["id"]) for e in ents}
        seen = set()
        for req in reqs:
            utt, trig, act = req["utterance"].strip(), req["trigger"], req["actuator"]
            problem = (
                "empty" if not utt
                else "dup" if utt.lower() in seen
                else "bad actuator" if act not in ids
                else "bad trigger" if trig not in ids and trig not in ("time", "sun")
                else "actuator name missing" if name_of[act].lower() not in utt.lower()
                else None
            )
            if problem:
                dropped += 1
                print(f"[heldout] drop ({home_id}, {problem}): {utt[:60]!r}")
                continue
            seen.add(utt.lower())
            rows.append(
                {
                    "utterance": utt,
                    "devices_block": home.devices_block(),
                    "trigger": trig,
                    "actuator": act,
                    "inventory": inventory,
                    "home": home_id,
                    "category": "route_automation",
                }
            )
        home.close()

    OUT.mkdir(parents=True, exist_ok=True)
    Dataset.from_list(rows).to_parquet(str(OUT / "heldout.parquet"))
    homes = len({r["home"] for r in rows})
    print(f"[heldout] {len(rows)} prompts across {homes} homes ({dropped} dropped) -> {OUT / 'heldout.parquet'}")


if __name__ == "__main__":
    main()