code/prepare_structured_hints.py

"""Build leakage-controlled privileged SFT data from ALFWorld expert plans.

The privileged hint is an episode-level structured task sketch.  It exposes the
expert-known *type* of the target's source support and an abstract milestone
plan, but never exposes the current gold command or TextWorld instance ids.

TextWorld environments are process-local, so extraction uses multiple worker
processes rather than threads.  Each worker replays complete expert episodes and
writes a private JSONL shard; the parent merges shards deterministically.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import multiprocessing as mp
import os
import re
import shutil
import signal
import sys
import uuid
from collections import Counter, defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any

import pandas as pd

PROJECT = Path("/share/shenmingxiang/closing-square")
WS = PROJECT / "verl_workspace"
AGENT_R1 = WS / "third_party" / "Agent-R1"
sys.path.insert(0, str(AGENT_R1))

from recipes.alfworld.env.alfworld_wrapper import AlfworldTextworldEnv


ENTITY_NAMES = {
    "alarmclock": "alarm clock",
    "armchair": "arm chair",
    "bathtubbasin": "bathtub basin",
    "butterknife": "butter knife",
    "cart": "utility cart",
    "cellphone": "cell phone",
    "coffeemachine": "coffee machine",
    "coffeetable": "coffee table",
    "countertop": "counter top",
    "creditcard": "credit card",
    "desklamp": "desk lamp",
    "diningtable": "dining table",
    "dishsponge": "dish sponge",
    "floorlamp": "floor lamp",
    "garbagecan": "garbage can",
    "glassbottle": "glass bottle",
    "handtowel": "hand towel",
    "laundryhamper": "laundry hamper",
    "lightswitch": "light switch",
    "remotecontrol": "remote control",
    "safebox": "safe",
    "saltshaker": "salt shaker",
    "shelvingunit": "shelf",
    "sidetable": "side table",
    "sinkbasin": "sink basin",
    "soapbar": "soap bar",
    "soapbottle": "soap bottle",
    "spraybottle": "spray bottle",
    "stoveburner": "stove burner",
    "tissuebox": "tissue box",
    "toiletpaper": "toilet paper",
    "toiletpaperhanger": "toilet paper hanger",
    "towelholder": "towel holder",
    "tvstand": "tv stand",
    "wateringcan": "watering can",
    "winebottle": "wine bottle",
}

SKILL_NAMES = {
    "pick_and_place": "find and place one target object",
    "pick_clean_then_place": "find, clean, and place one target object",
    "pick_cool_then_place": "find, cool, and place one target object",
    "pick_heat_then_place": "find, heat, and place one target object",
    "pick_two_obj_and_place": "find and place two target objects",
    "look_at_obj_in_light": "find a target object and inspect it under an activated light",
}

MILESTONES = {
    "pick_and_place": ["locate the target", "acquire the target", "deliver the target"],
    "pick_clean_then_place": [
        "locate the target", "acquire the target", "clean the target", "deliver the processed target"
    ],
    "pick_cool_then_place": [
        "locate the target", "acquire the target", "cool the target", "deliver the processed target"
    ],
    "pick_heat_then_place": [
        "locate the target", "acquire the target", "heat the target", "deliver the processed target"
    ],
    "pick_two_obj_and_place": [
        "locate a remaining target", "acquire that target", "deliver that target", "repeat for the second target"
    ],
    "look_at_obj_in_light": [
        "locate the target", "acquire the target", "reach the task light", "activate the light while holding the target"
    ],
}

TEMPLATE_TEXT = {
    "NAVIGATE": "navigate to the next relevant support, station, or destination",
    "ACQUIRE": "acquire the relevant target object",
    "PLACE": "place the held target at the intended destination",
    "OPEN": "open the relevant container",
    "CLOSE": "close the relevant container",
    "CLEAN": "clean the held target at the processing station",
    "COOL": "cool the held target at the processing station",
    "HEAT": "heat the held target at the processing station",
    "ACTIVATE": "activate the relevant task device",
    "INSPECT": "inspect the relevant target or support",
    "OBSERVE": "gather another observation",
    "OTHER": "perform the next task-relevant interaction",
}


def canonical_entity(value: Any) -> str:
    text = str(value or "").strip()
    if not text:
        return ""
    # TextWorld commands suffix concrete instances with an integer.  H is
    # deliberately type-level, so discard that id before any normalization.
    text = re.sub(r"\s+\d+$", "", text)
    text = re.sub(r"([a-z])([A-Z])", r"\1 \2", text).lower()
    key = re.sub(r"[^a-z]", "", text)
    return ENTITY_NAMES.get(key, re.sub(r"\s+", " ", text).strip())


def _load_traj(game_relative_path: str) -> dict[str, Any]:
    trial_dir = WS / "data" / "alfworld" / "games" / Path(game_relative_path).parent
    with (trial_dir / "traj_data.json").open("r", encoding="utf-8") as f:
        return json.load(f)


def _source_supports(expert_policy: list[str]) -> list[str]:
    """Return TextWorld support *types* used by target-acquisition actions.

    ALFRED ``high_pddl`` can describe a different initial placement from the
    compiled TextWorld game.  The TextWorld expert policy is the authoritative
    trajectory for this dataset.  We therefore abstract the support argument
    from commands such as ``take alarmclock 1 from sidetable 2`` and immediately
    canonicalize it to ``side table``; instance ids and the full command never
    enter H.
    """
    supports: list[str] = []
    for command in expert_policy:
        match = re.match(r"^take\s+.+?\s+from\s+(.+)$", command.strip(), re.IGNORECASE)
        if not match:
            continue
        support = canonical_entity(match.group(1))
        if support and support not in supports:
            supports.append(support)
    return supports


def action_template(command: str) -> str:
    text = command.strip().lower()
    if text.startswith("go to "):
        return "NAVIGATE"
    if text.startswith("take "):
        return "ACQUIRE"
    if text.startswith("move ") or text.startswith("put "):
        return "PLACE"
    if text.startswith("open "):
        return "OPEN"
    if text.startswith("close "):
        return "CLOSE"
    if text.startswith("clean ") or text.startswith("wash "):
        return "CLEAN"
    if text.startswith("cool "):
        return "COOL"
    if text.startswith("heat "):
        return "HEAT"
    if text.startswith("use ") or text.startswith("toggle "):
        return "ACTIVATE"
    if text.startswith("examine "):
        return "INSPECT"
    if text in {"look", "inventory"}:
        return "OBSERVE"
    return "OTHER"


def build_hint_fields(task_family: str, traj: dict[str, Any], expert_policy: list[str]) -> dict[str, Any]:
    params = traj.get("pddl_params") or {}
    target = canonical_entity(params.get("object_target"))
    destination = canonical_entity(params.get("parent_target"))
    sources = _source_supports(expert_policy)

    if task_family == "pick_clean_then_place":
        station = "sink basin"
    elif task_family == "pick_heat_then_place":
        station = "microwave"
    elif task_family == "pick_cool_then_place":
        station = "fridge"
    elif task_family == "look_at_obj_in_light":
        station = canonical_entity(params.get("toggle_target")) or "task light"
    else:
        station = ""

    return {
        "skill": task_family,
        "skill_description": SKILL_NAMES[task_family],
        "target_type": target,
        "target_count": "two" if task_family == "pick_two_obj_and_place" else "one",
        "source_support_types": sources,
        "processing_or_device_type": station,
        "goal_receptacle_type": destination,
        # TextWorld's compiled expert can order interactions differently from
        # ALFRED high_pddl.  Preserve only its action-template sequence, never
        # its arguments or instance ids.
        "action_templates": [action_template(command) for command in expert_policy],
    }


def render_hint(fields: dict[str, Any]) -> str:
    sources = fields["source_support_types"]
    if not sources:
        source_text = "unknown support type"
    elif len(sources) == 1:
        source_text = sources[0]
    else:
        source_text = ", then ".join(sources)

    facts = [
        f"- Target object type: {fields['target_type']}.",
        f"- Number of required targets: {fields['target_count']}.",
        f"- Expert-known source support type: {source_text}.",
    ]
    if fields["processing_or_device_type"]:
        facts.append(f"- Required processing station or task device type: {fields['processing_or_device_type']}.")
    if fields["goal_receptacle_type"]:
        facts.append(f"- Goal receptacle type: {fields['goal_receptacle_type']}.")

    milestones = "\n".join(f"- {TEMPLATE_TEXT[x]}." for x in fields["action_templates"])
    return (
        "Privileged task sketch for training:\n"
        f"Skill: {fields['skill_description']}.\n"
        "Hidden world facts:\n"
        + "\n".join(facts)
        + "\nAbstract milestones:\n"
        + milestones
        + "\nGround this sketch in the current observation, interaction history, and admissible command list."
    )


def hint_signature(fields: dict[str, Any]) -> str:
    core = {
        "skill": fields["skill"],
        "target": fields["target_type"],
        "count": fields["target_count"],
        "sources": fields["source_support_types"],
        "station": fields["processing_or_device_type"],
        "destination": fields["goal_receptacle_type"],
        "templates": fields["action_templates"],
    }
    return json.dumps(core, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def _process_game(job: dict[str, Any], env: AlfworldTextworldEnv) -> tuple[list[dict[str, Any]], str | None]:
    extra = job["extra_info"]
    game = str(extra["game_relative_path"])
    try:
        traj = _load_traj(game)
        obs, info = env.reset_with_info(game)
        initial_policy = [str(x) for x in info.get("policy_commands", [])]
        fields = build_hint_fields(str(extra["task_family"]), traj, initial_policy)
        hint = render_hint(fields)
        signature = hint_signature(fields)
        history: list[str] = []
        records: list[dict[str, Any]] = []

        for step_idx in range(50):
            admissible = [str(x) for x in info.get("admissible_commands", [])]
            policy = [str(x) for x in info.get("policy_commands", [])]
            if not policy or bool(info.get("success", False)):
                break
            gold_action = policy[0]
            records.append(
                {
                    "game_id": game,
                    "task_id": str(extra.get("task_id", "")),
                    "task_family": str(extra["task_family"]),
                    "goal": str(extra.get("goal_text", "")),
                    "step_idx": step_idx,
                    "history": list(history),
                    "obs": str(obs),
                    "admissible": admissible,
                    "H": hint,
                    "H_id": signature,
                    "H_fields": fields,
                    "Y": gold_action,
                    "gold_is_admissible": gold_action in admissible,
                    "sft_split": job["sft_split"],
                }
            )
            next_obs, _, done, info = env.step(gold_action)
            history.append(f"Action taken: {gold_action}\nEnvironment response: {next_obs}")
            obs = next_obs
            if done:
                break
        return records, None
    except Exception as exc:
        return [], f"{type(exc).__name__}: {exc}"
    finally:
        env._close_env()


def _game_timeout(_signum, _frame) -> None:
    raise TimeoutError("ALFWorld game replay exceeded the per-game timeout")


def _worker(
    chunk_id: int,
    jobs: list[dict[str, Any]],
    part_dir: str,
    game_timeout: int,
) -> dict[str, Any]:
    env = AlfworldTextworldEnv(data_root=str(WS / "data" / "alfworld"), max_episode_steps=50)
    signal.signal(signal.SIGALRM, _game_timeout)
    part_path = Path(part_dir) / f"part_{chunk_id:04d}.jsonl"
    error_path = Path(part_dir) / f"errors_{chunk_id:04d}.jsonl"
    n_records = 0
    n_failed = 0
    with part_path.open("w", encoding="utf-8") as out, error_path.open("w", encoding="utf-8") as err:
        for job in jobs:
            signal.alarm(game_timeout)
            try:
                records, failure = _process_game(job, env)
            except TimeoutError as exc:
                records, failure = [], f"TimeoutError: {exc}"
            finally:
                signal.alarm(0)
            if failure:
                n_failed += 1
                err.write(json.dumps({"game_id": job["extra_info"]["game_relative_path"], "error": failure}) + "\n")
                continue
            for record in records:
                out.write(json.dumps(record, ensure_ascii=False) + "\n")
            n_records += len(records)
    return {"chunk": chunk_id, "games": len(jobs), "records": n_records, "failed": n_failed}


def _assign_game_splits(df: pd.DataFrame, val_fraction: float, seed: int) -> list[dict[str, Any]]:
    del seed  # stable hashes make the split independent of input ordering
    jobs: list[dict[str, Any]] = []
    by_family: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for _, row in df.iterrows():
        extra = dict(row["extra_info"])
        by_family[str(extra["task_family"])].append(extra)

    cutoff = int(round(val_fraction * 10_000))
    for family in sorted(by_family):
        family_jobs = sorted(by_family[family], key=lambda x: str(x["game_relative_path"]))
        for extra in family_jobs:
            digest = hashlib.sha1(str(extra["game_relative_path"]).encode("utf-8")).hexdigest()
            bucket = int(digest[:8], 16) % 10_000
            jobs.append({"extra_info": extra, "sft_split": "val" if bucket < cutoff else "train"})
    return jobs


def _normal(text: str) -> str:
    return re.sub(r"\s+", " ", text.strip().lower())


def build_manifest(output_path: Path, game_count: int, failures: int) -> dict[str, Any]:
    task_counts: Counter[str] = Counter()
    split_counts: Counter[str] = Counter()
    split_games: dict[str, set[str]] = defaultdict(set)
    hint_ids: set[str] = set()
    records = 0
    exact_hint_action = 0
    action_inside_hint = 0
    hint_with_digit = 0
    invalid_gold = 0
    total_hint_chars = 0
    with output_path.open("r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            item = json.loads(line)
            records += 1
            task_counts[item["task_family"]] += 1
            split_counts[item["sft_split"]] += 1
            split_games[item["sft_split"]].add(item["game_id"])
            hint_ids.add(item["H_id"])
            h, y = _normal(item["H"]), _normal(item["Y"])
            exact_hint_action += int(h == y)
            action_inside_hint += int(y in h)
            hint_with_digit += int(bool(re.search(r"\d", item["H"])))
            invalid_gold += int(not item["gold_is_admissible"])
            total_hint_chars += len(item["H"])

    manifest = {
        "schema_version": 1,
        "source": str(WS / "data" / "alfworld" / "train.parquet"),
        "output": str(output_path),
        "games_requested": game_count,
        "games_failed": failures,
        "records": records,
        "games_by_sft_split": {k: len(v) for k, v in sorted(split_games.items())},
        "records_by_sft_split": dict(sorted(split_counts.items())),
        "records_by_task_family": dict(sorted(task_counts.items())),
        "unique_hint_signatures": len(hint_ids),
        "mean_hint_chars": total_hint_chars / max(records, 1),
        "leakage_checks": {
            "exact_hint_equals_action": exact_hint_action,
            "full_gold_action_inside_hint": action_inside_hint,
            "hint_contains_digit": hint_with_digit,
            "gold_action_not_admissible": invalid_gold,
        },
    }
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default=str(WS / "data" / "alfworld" / "train.parquet"))
    parser.add_argument("--output", default=str(WS / "data" / "privileged_plan_sft.jsonl"))
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--chunks", type=int, default=64)
    parser.add_argument("--limit", type=int, default=0, help="Optional game limit for smoke tests")
    parser.add_argument("--val_fraction", type=float, default=0.05)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--game_timeout", type=int, default=120)
    parser.add_argument(
        "--resume_part_dir",
        type=Path,
        default=None,
        help="Read-only shard directory from a stopped earlier run",
    )
    parser.add_argument(
        "--completed_chunks",
        default="",
        help="Comma-separated chunk ids known to be complete in --resume_part_dir",
    )
    args = parser.parse_args()

    output_path = Path(args.output).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    # Every extraction owns a private shard directory.  This avoids one run
    # deleting shards that another live run is still writing.
    run_token = f"{os.getpid()}.{uuid.uuid4().hex}"
    part_dir = output_path.parent / f".{output_path.name}.parts.{run_token}"
    part_dir.mkdir(parents=True, exist_ok=False)

    df = pd.read_parquet(args.input)
    if args.limit > 0:
        df = df.head(args.limit)
    jobs = _assign_game_splits(df, args.val_fraction, args.seed)
    chunks = [jobs[i::args.chunks] for i in range(args.chunks)]
    chunks = [chunk for chunk in chunks if chunk]
    print(f"extracting {len(jobs)} games with {args.workers} processes in {len(chunks)} chunks", flush=True)

    summaries: list[dict[str, Any]] = []
    completed = {
        int(value) for value in args.completed_chunks.split(",") if value.strip()
    }
    if completed and args.resume_part_dir is None:
        raise ValueError("--completed_chunks requires --resume_part_dir")
    if args.resume_part_dir is not None:
        source_dir = args.resume_part_dir.resolve()
        for chunk_id in sorted(completed):
            source_part = source_dir / f"part_{chunk_id:04d}.jsonl"
            source_errors = source_dir / f"errors_{chunk_id:04d}.jsonl"
            if not source_part.is_file() or not source_errors.is_file():
                raise FileNotFoundError(f"missing completed resume shard for chunk {chunk_id}")
            shutil.copy2(source_part, part_dir / source_part.name)
            shutil.copy2(source_errors, part_dir / source_errors.name)
            with source_part.open("r", encoding="utf-8") as handle:
                record_count = sum(1 for line in handle if line.strip())
            with source_errors.open("r", encoding="utf-8") as handle:
                failure_count = sum(1 for line in handle if line.strip())
            summaries.append(
                {
                    "chunk": chunk_id,
                    "games": len(chunks[chunk_id]),
                    "records": record_count,
                    "failed": failure_count,
                }
            )
        print(f"resumed {len(completed)} completed chunks from {source_dir}", flush=True)

    context = mp.get_context("spawn")
    with ProcessPoolExecutor(max_workers=args.workers, mp_context=context) as pool:
        futures = [
            pool.submit(_worker, i, chunk, str(part_dir), args.game_timeout)
            for i, chunk in enumerate(chunks)
            if i not in completed
        ]
        for future in as_completed(futures):
            summary = future.result()
            summaries.append(summary)
            print(
                f"chunk {summary['chunk']:02d}: games={summary['games']} "
                f"records={summary['records']} failed={summary['failed']}",
                flush=True,
            )

    temp_output = output_path.parent / f".{output_path.name}.tmp.{run_token}"
    with temp_output.open("w", encoding="utf-8") as out:
        for part in sorted(part_dir.glob("part_*.jsonl")):
            with part.open("r", encoding="utf-8") as src:
                shutil.copyfileobj(src, out)
        out.flush()
        os.fsync(out.fileno())
    os.replace(temp_output, output_path)

    failures = sum(x["failed"] for x in summaries)
    manifest = build_manifest(output_path, len(jobs), failures)
    manifest_path = output_path.with_suffix(output_path.suffix + ".manifest.json")
    temp_manifest = manifest_path.parent / f".{manifest_path.name}.tmp.{run_token}"
    with temp_manifest.open("w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)
        f.flush()
        os.fsync(f.fileno())
    os.replace(temp_manifest, manifest_path)
    error_output = output_path.with_suffix(output_path.suffix + ".errors.jsonl")
    temp_errors = error_output.parent / f".{error_output.name}.tmp.{run_token}"
    with temp_errors.open("w", encoding="utf-8") as destination:
        for error_file in sorted(part_dir.glob("errors_*.jsonl")):
            with error_file.open("r", encoding="utf-8") as source:
                shutil.copyfileobj(source, destination)
        destination.flush()
        os.fsync(destination.fileno())
    os.replace(temp_errors, error_output)
    # All futures have joined and the merged output is durable.  Successful
    # private shards are now safe to remove; failed runs are retained so a
    # repair process never races against deletion of its diagnosis inputs.
    if failures == 0:
        shutil.rmtree(part_dir)
    else:
        print(f"retaining failed run shards for repair: {part_dir}", flush=True)
    print(json.dumps(manifest, ensure_ascii=False, indent=2), flush=True)
    if any(manifest["leakage_checks"].values()):
        raise SystemExit("leakage or validity checks failed; inspect manifest before SFT")


if __name__ == "__main__":
    main()