"""Grounded task generator for Brain 1: one base generator, two output formats. gen: propose (home, utterance, gold call) -> execute in SandboxHome -> keep only verified records (actions must produce the expected diff; rejections must fail with a NAME miss). Records land in data/processed/tasks.jsonl, self-contained. sft: tasks.jsonl -> intent_sft-shaped rows (Brain 1 stage-2 SFT) rl: tasks.jsonl -> prompt + expected-diff rows (RLVR; sandbox replays = reward) Run: uv run --group sandbox --group train python train/taskgen.py gen|sft|rl """ import json import random import sys from collections import Counter from dataclasses import asdict, dataclass from datetime import datetime, timedelta from pathlib import Path import yaml ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT)) from sandbox.harness import SandboxHome # noqa: E402 _CFG = yaml.safe_load((Path(__file__).parent / "config.yaml").read_text()) PROCESSED = ROOT / "data" / "processed" TASKS_FILE = PROCESSED / "tasks.jsonl" HOME_DIRS = [ ROOT / "data" / "homes", ROOT / "plans" / "home-assistant-datasets-src" / "datasets" / "assist-mini", ] OFF_LIKE = {"off", "closed", "unlocked", "docked", "idle"} TEMPLATES = { ("light", "on"): ["turn on the {name}", "switch the {name} on", "can you turn the {name} on", "{name} on please"], ("light", "off"): ["turn off the {name}", "switch the {name} off", "kill the {name}", "{name} off please"], ("cover", "on"): ["open the {name}", "can you open the {name}", "pull up the {name}"], ("cover", "off"): ["close the {name}", "shut the {name}", "can you close the {name}"], ("lock", "on"): ["lock the {name}", "secure the {name}", "make sure the {name} is locked"], ("lock", "off"): ["unlock the {name}", "open up the {name}", "can you unlock the {name}"], ("valve", "on"): ["open the {name}", "turn on the {name}"], ("valve", "off"): ["close the {name}", "shut off the {name}"], (None, "on"): ["turn on the {name}", "switch on the {name}", "start the {name}"], (None, "off"): ["turn off the {name}", "switch off the {name}", "shut down the {name}"], } ATTR_TEMPLATES = { "brightness": ["set the {name} to {pct} percent", "dim the {name} to {pct}%", "make the {name} {pct} percent bright"], "position": ["set the {name} to {pct} percent open", "open the {name} to {pct}%"], } MEDIA = { "playing": ("HassMediaPause", ["pause the {name}", "pause whatever is on the {name}"]), "paused": ("HassMediaUnpause", ["resume the {name}", "unpause the {name}"]), } # automation requests are NOT templated here: subagents write them per home from # train/prompts/routegen.md into ROUTES_DIR; gen() ingests + validates them. ROUTES_DIR = PROCESSED / "route_requests" BATCH_TEMPLATES = { "on": ["turn on all the lights", "switch every light on", "all lights on please"], "off": ["turn off all the lights", "switch every light off", "kill all the lights"], } EXCLUSION_TEMPLATES = [ "turn off all the lights except the {kept}", "switch off every light but the {kept}", "turn everything off except the {kept}", ] DENIALS = [ "Sorry, I can't find a {name} to control.", "I don't see a {name} in this home.", "There's no {name} listed here.", ] FOREIGN = [ ("hot tub heater", "switch"), ("pool light", "light"), ("wine fridge", "switch"), ("doggy door", "cover"), ("side gate lock", "lock"), ("attic fan", "fan"), ("sprinkler valve", "valve"), ("sauna", "climate"), ("garage heater", "switch"), ("greenhouse light", "light"), ] CONTROLLABLE = {"light", "switch", "fan", "cover", "lock", "valve", "media_player"} @dataclass class TaskRecord: home: str category: str # control | attr | disambig | reject_missing utterance: str devices_block: str time: str # list of {"name": intent, "slots": {...}} (batch => several); None => no call # (rejections AND route_automation — those route verbatim, see to_sft) gold: list | None # intent-relevant outcomes only: {entity: {"state"|attr: after_value}}. # Reward = containment (replay diff includes these after-values), NOT exact # diff equality — incidental keys (color_mode, hidden brightness restores) # aren't reproducible from yaml and aren't what the record asserts. expect: dict response: str | None # denial sentence for rejections inventory: str # raw home yaml so RL reward can rebuild the sandbox meta: dict | None = None # route records: {"trigger": eid, "actuator": eid} for the Brain-2 reward def _rand_time(rng: random.Random) -> str: dt = datetime(2024, 1, 1) + timedelta(minutes=rng.randrange(2 * 365 * 24 * 60)) hour = dt.strftime("%I").lstrip("0") return f"{hour}:{dt.minute} {dt.strftime('%p')} on {dt.strftime('%A, %B')} {dt.day} {dt.year}" # dup-named entities in verb-ambiguous domains get the domain noun appended in the # UTTERANCE (the slot keeps the real name) so the user text justifies the domain slot _DOMAIN_NOUN = {"light": "light", "fan": "fan", "switch": "switch"} def _propose(entity: dict, state: str, attrs: dict, dup: bool, rng: random.Random): """Pick an intent that should change this entity, given its current state.""" domain, name = entity["id"].split(".")[0], entity.get("name", entity["id"]) slots = {"name": name} | ({"domain": domain} if dup else {}) spoken = name if dup: if domain in _DOMAIN_NOUN: spoken = f"{name} {_DOMAIN_NOUN[domain]}" elif domain not in ("cover", "lock", "valve"): return None # no way to phrase the disambiguation unambiguously if domain == "media_player" and state in MEDIA: intent, templates = MEDIA[state] return "control", intent, slots, rng.choice(templates).format(name=spoken) if domain == "light" and state == "on" and rng.random() < 0.5: current = round((attrs.get("brightness") or 0) / 255 * 100) if choices := [p for p in (10, 25, 40, 60, 75, 90) if abs(p - current) >= 15]: pct = rng.choice(choices) slots["brightness"] = pct return "attr", "HassLightSet", slots, rng.choice(ATTR_TEMPLATES["brightness"]).format(name=spoken, pct=pct) if domain == "cover" and rng.random() < 0.4: current = attrs.get("current_position") or 0 if choices := [p for p in (20, 35, 50, 65, 80) if abs(p - current) >= 15]: pct = rng.choice(choices) slots["position"] = pct return "attr", "HassSetPosition", slots, rng.choice(ATTR_TEMPLATES["position"]).format(name=spoken, pct=pct) direction = "on" if state in OFF_LIKE else "off" intent = {"on": "HassTurnOn", "off": "HassTurnOff"}[direction] templates = TEMPLATES.get((domain, direction)) or TEMPLATES[(None, direction)] return "control", intent, slots, rng.choice(templates).format(name=spoken) # every attr devices_block renders must be patched back, or replayed homes # won't reproduce the prompt's device lines _PATCH_ATTRS = ("brightness", "current_position", "percentage", "volume_level", "media_title") def _live_inventory(home: SandboxHome) -> str: """Inventory yaml with current state written back, so a fresh boot reproduces exactly the world the record's devices_block describes (drift-consistent).""" snap = home.snapshot() inv = yaml.safe_load(yaml.safe_dump(home.inventory)) for e in inv.get("entities", []): state, attrs = snap[e["id"]] e["state"] = state for k in _PATCH_ATTRS: if attrs.get(k) is not None: e.setdefault("attributes", {})[k] = attrs[k] else: e.get("attributes", {}).pop(k, None) # stale value would resurrect on boot return yaml.safe_dump(inv, sort_keys=False) def gen_home(path: Path, n: int, rng: random.Random) -> list[TaskRecord]: home = SandboxHome.from_file(path) home_id = path.parent.name if path.name == "_fixtures.yaml" else path.stem entities = [e for e in home.inventory.get("entities", []) if e["id"].split(".")[0] in CONTROLLABLE] name_counts = Counter(e.get("name", e["id"]) for e in home.inventory.get("entities", [])) records: list[TaskRecord] = [] attempts = 0 while len(records) < n and attempts < n * 4 and entities: attempts += 1 block, time = home.devices_block(), _rand_time(rng) inv_yaml = _live_inventory(home) if rng.random() < 0.2: # rejection: a plausible device this home doesn't have name, domain = rng.choice([f for f in FOREIGN if f[0] not in name_counts]) templates = TEMPLATES.get((domain, "on")) or TEMPLATES[(None, "on")] probe = home.execute("HassTurnOn", {"name": name}) if probe.ok or "NAME" not in (probe.error or ""): continue # grounding: only keep if HA truly can't match it records.append(TaskRecord( home_id, "reject_missing", rng.choice(templates).format(name=name), block, time, None, {}, rng.choice(DENIALS).format(name=name), inv_yaml, )) continue entity = rng.choice(entities) state, attrs = home.snapshot()[entity["id"]] dup = name_counts[entity.get("name", entity["id"])] > 1 proposal = _propose(entity, state, attrs, dup, rng) if proposal is None: continue category, intent, slots, utterance = proposal result = home.execute(intent, slots) if not result.ok or entity["id"] not in result.changed: continue # grounding: drop anything that didn't verifiably act delta = result.changed[entity["id"]] want = {} if "state" in delta: want["state"] = delta["state"][1] for slot_key, diff_key in (("brightness", "brightness"), ("position", "current_position")): if slot_key in slots and diff_key in delta: want[diff_key] = delta[diff_key][1] if not want: continue records.append(TaskRecord( home_id, "disambig" if dup else category, utterance, block, time, [{"name": intent, "slots": slots}], {entity["id"]: want}, None, inv_yaml, )) records += _gen_routes(home, home_id, rng) records += _gen_batches(home, home_id, name_counts, rng) home.close() return records def _gen_routes(home: SandboxHome, home_id: str, rng: random.Random) -> list[TaskRecord]: """Ingest subagent-written automation requests (ROUTES_DIR/.json), validating each against the live inventory. Doubles as the Brain-2 GRPO prompt set; meta carries trigger/actuator for its reward.""" path = ROUTES_DIR / f"{home_id}.json" if not path.exists(): return [] 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} records, seen = [], set() for req in json.loads(path.read_text()): utt = (req.get("utterance") or "").strip() trig, act = req.get("trigger"), req.get("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 from utterance" if name_of[act].lower() not in utt.lower() else None ) if problem: print(f"[gen] drop route ({home_id}, {problem}): {utt[:60]!r}") continue seen.add(utt.lower()) records.append(TaskRecord( home_id, "route_automation", utt, home.devices_block(), _rand_time(rng), None, {}, None, _live_inventory(home), meta={"trigger": trig, "actuator": act}, )) return records def _gen_batches(home: SandboxHome, home_id: str, name_counts: Counter, rng: random.Random) -> list[TaskRecord]: """One batch + one exclusion record per home with enough lights. Gold lists every targeted light; grounded by executing all calls and requiring the changers to change.""" records = [] lights = [e for e in home.inventory.get("entities", []) if e["id"].startswith("light.")] def call(e, intent): slots = {"name": e.get("name", e["id"])} if name_counts[e.get("name", e["id"])] > 1: slots["domain"] = "light" return {"name": intent, "slots": slots} def run_all(targets, intent, movers): block, time, inv = home.devices_block(), _rand_time(rng), _live_inventory(home) gold = [call(e, intent) for e in targets] expect = {} for e in movers: c = call(e, intent) r = home.execute(c["name"], c["slots"]) if not r.ok or e["id"] not in r.changed or "state" not in r.changed[e["id"]]: return None expect[e["id"]] = {"state": r.changed[e["id"]]["state"][1]} return block, time, inv, gold, expect live = home.states() on = [e for e in lights if live[e["id"]] == "on"] off = [e for e in lights if live[e["id"]] == "off"] direction, movers = ("off", on) if len(on) >= len(off) else ("on", off) if len(movers) >= 2: done = run_all(lights, "HassTurnOff" if direction == "off" else "HassTurnOn", movers) if done: block, time, inv, gold, expect = done records.append(TaskRecord( home_id, "batch", rng.choice(BATCH_TEMPLATES[direction]), block, time, gold, expect, None, inv, )) live = home.states() on = [e for e in lights if live[e["id"]] == "on"] if len(on) >= 3: kept = rng.choice(on) targets = [e for e in on if e["id"] != kept["id"]] done = run_all(targets, "HassTurnOff", targets) if done: block, time, inv, gold, expect = done records.append(TaskRecord( home_id, "exclusion", rng.choice(EXCLUSION_TEMPLATES).format(kept=kept.get("name", kept["id"])), block, time, gold, expect, None, inv, )) return records def _find_homes(roots: list[Path] = HOME_DIRS) -> list[Path]: files: list[Path] = [] for root in roots: if root.exists(): files += sorted(root.glob("**/_fixtures.yaml")) + sorted(root.glob("*.yaml")) return files def gen(n_per_home: int = 8, seed: int = 0) -> None: PROCESSED.mkdir(parents=True, exist_ok=True) rng = random.Random(seed) records: list[TaskRecord] = [] # publishable output: ONLY our generated pool, never third-party fixtures for path in _find_homes([HOME_DIRS[0]]): recs = gen_home(path, n_per_home, rng) records += recs print(f"[gen] {path.parent.name}: {len(recs)} records") with TASKS_FILE.open("w") as f: for r in records: f.write(json.dumps(asdict(r), ensure_ascii=False) + "\n") print(f"[gen] {len(records)} total {dict(Counter(r.category for r in records))} -> {TASKS_FILE}") def _system(record: dict) -> str: return ( f"{_CFG['system_prompt']}\n\nActions:\n{_CFG['intents']}" f"\n\nDevices:\n{record['devices_block']}\nCurrent time: {record['time']}" ) def to_sft(record: dict) -> dict: if record["category"] == "route_automation": payload = {"automation": record["utterance"]} # verbatim handoff to Brain 2 elif record["gold"]: payload = {"intents": record["gold"]} else: payload = {"response": record["response"]} return { "messages": [ {"role": "system", "content": _system(record)}, {"role": "user", "content": record["utterance"]}, {"role": "assistant", "content": json.dumps(payload, ensure_ascii=False, separators=(",", ":"))}, ], "category": record["category"], "home": record["home"], } def to_rl(record: dict) -> dict: meta = record.get("meta") or {} return { "messages": [ {"role": "system", "content": _system(record)}, {"role": "user", "content": record["utterance"]}, ], "category": record["category"], "home": record["home"], "inventory": record["inventory"], "expect": json.dumps(record["expect"], ensure_ascii=False), "gold": json.dumps(record["gold"], ensure_ascii=False), "utterance": record["utterance"], "devices_block": record["devices_block"], "trigger": meta.get("trigger", ""), "actuator": meta.get("actuator", ""), } def _format(kind: str) -> None: from datasets import Dataset fmt = {"sft": to_sft, "rl": to_rl}[kind] rows = [fmt(json.loads(line)) for line in TASKS_FILE.read_text().splitlines()] out = PROCESSED / "taskgen" / f"{kind}.parquet" out.parent.mkdir(exist_ok=True) Dataset.from_list(rows).to_parquet(str(out)) print(f"[{kind}] {len(rows)} rows -> {out}") def expect_met(changed: dict, expect: dict) -> bool: """The RLVR reward core: every expected after-value appears in the diff.""" return all( eid in changed and key in changed[eid] and changed[eid][key][1] == after for eid, wants in expect.items() for key, after in wants.items() ) def audit(n: int = 12, seed: int = 11) -> None: """Replay n random action records on fresh boots; block + expect must reproduce.""" lines = [json.loads(line) for line in TASKS_FILE.read_text().splitlines()] sample = random.Random(seed).sample([r for r in lines if r["gold"]], n) ok = 0 for rec in sample: home = SandboxHome(rec["inventory"]) block_ok = rec["devices_block"] == home.devices_block() merged: dict = {} calls_ok = True for call in rec["gold"]: result = home.execute(call["name"], call["slots"]) calls_ok = calls_ok and result.ok for eid, delta in result.changed.items(): merged.setdefault(eid, {}).update(delta) met = calls_ok and expect_met(merged, rec["expect"]) home.close() ok += block_ok and met if not (block_ok and met): print(f"[audit] MISMATCH {rec['home']} {rec['utterance']!r} block={block_ok} expect={met}") print(f"[audit] {ok}/{n} records replay perfectly") def diversity() -> None: """Score route_automation utterances on diversity; flag template-y sameness. Checks: near-duplicate pairs (Jaccard>0.7), opener concentration, type-token ratio, trigger-kind spread. Run after ingesting a new generation batch.""" lines = [json.loads(line) for line in TASKS_FILE.read_text().splitlines()] routes = [r for r in lines if r["category"] == "route_automation"] if not routes: print("[diversity] no route_automation records") return utts = [r["utterance"].lower().replace(",", "") for r in routes] word_sets = [set(u.split()) for u in utts] dups = [ (utts[i], utts[j]) for i in range(len(utts)) for j in range(i + 1, len(utts)) if len(word_sets[i] & word_sets[j]) / len(word_sets[i] | word_sets[j]) > 0.7 ] trigrams = [" ".join(u.split()[i:i + 3]) for u in utts for i in range(len(u.split()) - 2)] distinct_tri = len(set(trigrams)) / len(trigrams) openers = Counter(" ".join(u.split()[:2]) for u in utts) kinds = Counter( "time/sun" if r["meta"]["trigger"] in ("time", "sun") else r["meta"]["trigger"].split(".")[0] for r in routes ) top_opener, top_n = openers.most_common(1)[0] top_kind, kind_n = kinds.most_common(1)[0] print(f"[diversity] {len(routes)} routes over {len({r['home'] for r in routes})} homes") print(f"[diversity] trigger kinds: {dict(kinds)} — top {top_kind} {kind_n / len(routes):.0%} (flag > 50%)") print(f"[diversity] distinct trigrams: {distinct_tri:.2f} (flag < 0.55)") print(f"[diversity] top opener: {top_opener!r} x{top_n} ({top_n / len(utts):.0%}, flag > 25%)") print(f"[diversity] near-duplicate pairs (jaccard>0.7): {len(dups)} (flag > {len(utts) // 20})") for a, b in dups[:5]: print(f" ~ {a!r}\n {b!r}") flags = ( (distinct_tri < 0.55) + (top_n / len(utts) > 0.25) + (len(dups) > len(utts) // 20) + (kind_n / len(routes) > 0.5) ) print(f"[diversity] {'FAIL — regenerate the flagged areas' if flags else 'PASS'}") def show(per_category: int = 2, seed: int = 3) -> None: """Pretty-print sample records per category, plus one full SFT and RL row.""" lines = [json.loads(line) for line in TASKS_FILE.read_text().splitlines()] rng = random.Random(seed) by_cat: dict[str, list] = {} for r in lines: by_cat.setdefault(r["category"], []).append(r) for cat, recs in sorted(by_cat.items()): print(f"\n=== {cat} ({len(recs)} records) ===") for rec in rng.sample(recs, min(per_category, len(recs))): print(f" [{rec['home']}] user: {rec['utterance']!r}") if rec["gold"]: print(f" gold: {json.dumps(rec['gold'], separators=(',', ':'))}") print(f" expect: {json.dumps(rec['expect'], separators=(',', ':'))}") elif rec["category"] == "route_automation": print(f" gold: route -> Brain 2, verbatim (meta {rec['meta']})") else: print(f" gold: (no tool call) -> {rec['response']!r}") rec = rng.choice([r for r in lines if r["gold"]]) sft_row = to_sft(rec) print("\n=== one full SFT row (what Brain 1 trains on) ===") for m in sft_row["messages"]: print(f"--- {m['role']} ---\n{m['content']}") rl_row = to_rl(rec) print("\n=== same record as RL row (no assistant turn; sandbox judges) ===") print("messages: system + user as above, NO assistant") print(f"expect: {rl_row['expect']}") print(f"inventory: {len(rl_row['inventory'])} chars of patched home yaml") def check() -> None: """Boot-validate every home in the pool; report stats or the failure.""" ok = bad = 0 for path in _find_homes(): label = path.parent.name if path.name == "_fixtures.yaml" else path.stem try: home = SandboxHome.from_file(path) except Exception as err: # noqa: BLE001 bad += 1 print(f"[check] FAIL {label}: {str(err)[:120]}") continue ents = home.inventory.get("entities", []) domains = Counter(e["id"].split(".")[0] for e in ents) dups = [n for n, c in Counter(e.get("name") for e in ents).items() if c > 1] home.close() ok += 1 print(f"[check] ok {label}: {len(ents)} entities {dict(domains)}" + (f" dups={dups}" if dups else "")) print(f"[check] {ok} ok, {bad} failed") if __name__ == "__main__": cmd = sys.argv[1] if len(sys.argv) > 1 else "gen" {"gen": gen, "sft": lambda: _format("sft"), "rl": lambda: _format("rl"), "check": check, "audit": audit, "show": show, "diversity": diversity}[cmd]()