build_quick_data_set.py

import argparse
import csv
import json
import os

from build_imitation_dataset import build_dataset_from_episode_rows
from pretrain_imitation_model import pretrain_imitation_model
from profiles import CURRENT_MODEL_GENERATION, get_profile_paths, list_profile_names


ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
ANALYSIS_DIR = os.path.join(ROOT_DIR, "runs", "analysis")
DEFAULT_REVIEWS_CSV = os.path.join(ANALYSIS_DIR, "episode_review_labels.csv")
DEFAULT_RECORDED_LIVE_POSITION_CSV = os.path.join(
    ANALYSIS_DIR,
    "two_gem_archive",
    f"recorded_live_state_{CURRENT_MODEL_GENERATION}",
    "index.csv",
)
DEFAULT_RECORDED_LIVE_VISION_CSV = os.path.join(
    ANALYSIS_DIR,
    "two_gem_archive",
    f"recorded_live_vision_{CURRENT_MODEL_GENERATION}",
    "index.csv",
)
DEFAULT_RECORDED_LIVE_LEGACY_CSV = os.path.join(
    ANALYSIS_DIR,
    "two_gem_archive",
    "backfilled_legacy",
    "index.csv",
)
DEFAULT_LEGACY_TWO_GEM_CSV = os.path.join(ANALYSIS_DIR, "two_gem_episodes.csv")
LEDGE_REGION_BOUNDS = {
    "bottom_right": {
        "label": "Small right ledge",
        "x_min": 0.90,
        "x_max": 0.98,
        "y_min": 0.74,
        "y_max": 0.84,
    },
    "left_large": {
        "label": "Big left ledge",
        "x_min": 0.75,
        "x_max": 0.89,
        "y_min": 0.80,
        "y_max": 0.88,
    },
}


def _parse_float(value):
    if value in (None, ""):
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _parse_int(value):
    if value in (None, ""):
        return None
    try:
        return int(float(value))
    except (TypeError, ValueError):
        return None


def _load_csv_rows(path):
    if not os.path.exists(path):
        return []
    with open(path, "r", encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


def _load_json(path):
    if not os.path.exists(path):
        return None
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _prompt_text(prompt_text, default=None, allow_blank=False):
    suffix = f" [{default}]" if default else ""
    while True:
        raw = input(f"{prompt_text}{suffix}: ").strip()
        if raw:
            return raw
        if default is not None:
            return default
        if allow_blank:
            return ""
        print("Please enter a value.")


def _prompt_optional_float(prompt_text, default=None, minimum=0.0):
    suffix = f" [{default}]" if default is not None else ""
    while True:
        raw = input(f"{prompt_text}{suffix}: ").strip()
        if not raw:
            return default
        try:
            value = float(raw)
            if value >= minimum:
                return value
        except ValueError:
            pass
        print(f"Please enter a number >= {minimum}, or leave it blank.")


def _prompt_yes_no(prompt_text, default=False):
    suffix = "[Y/n]" if default else "[y/N]"
    while True:
        raw = input(f"{prompt_text} {suffix}: ").strip().lower()
        if not raw:
            return default
        if raw in {"y", "yes"}:
            return True
        if raw in {"n", "no"}:
            return False
        print("Please enter y or n.")


def _prompt_choice(prompt_text, choices, default_index=0):
    print(prompt_text)
    for index, (value, label) in enumerate(choices, start=1):
        default_suffix = " [default]" if index - 1 == default_index else ""
        print(f"  {index}. {label}{default_suffix}")
    while True:
        raw = input("Choose: ").strip()
        if not raw:
            return choices[default_index][0]
        if raw.isdigit():
            idx = int(raw) - 1
            if 0 <= idx < len(choices):
                return choices[idx][0]
        print("Enter one of the listed numbers, or press Enter for the default.")


def _review_status_map(path):
    status = {}
    for row in _load_csv_rows(path):
        key = (row.get("profile", ""), str(row.get("episode", "")))
        if row.get("good_episode") == "yes":
            status[key] = "good"
        elif row.get("viewed") == "yes":
            status[key] = "bad"
    return status


def _candidate_profiles(profile_filter=None):
    if profile_filter:
        return [profile_filter]
    return [name for name in list_profile_names() if name != "default"]


def _default_two_gem_candidates_csv():
    archive_root = os.path.join(ANALYSIS_DIR, "two_gem_archive")
    candidates = [
        DEFAULT_RECORDED_LIVE_POSITION_CSV,
        DEFAULT_RECORDED_LIVE_VISION_CSV,
        DEFAULT_RECORDED_LIVE_LEGACY_CSV,
    ]
    if os.path.isdir(archive_root):
        for name in sorted(os.listdir(archive_root)):
            index_path = os.path.join(archive_root, name, "index.csv")
            if not os.path.isfile(index_path):
                continue
            if index_path not in candidates:
                candidates.append(index_path)
    for candidate in candidates:
        if os.path.exists(candidate):
            return candidate
    return DEFAULT_LEGACY_TWO_GEM_CSV


def _normalized_existing_path(path):
    if not path:
        return path
    normalized = str(path).strip()
    if normalized.startswith("\\\\?\\"):
        normalized = normalized[4:]
    return os.path.normpath(normalized)


def _load_two_gem_candidates(candidates_csv, profile_filter=None):
    candidates = []
    for row in _load_csv_rows(candidates_csv):
        profile = (row.get("profile") or "").strip()
        if profile_filter and profile != profile_filter:
            continue

        episode = str(row.get("episode", "")).strip()
        if not profile or not episode:
            continue

        archive_path = _normalized_existing_path(row.get("archive_path", ""))
        replay_path = archive_path
        if not replay_path or not os.path.exists(replay_path):
            replay_path = os.path.join(
                ROOT_DIR,
                "runs",
                profile,
                "logs",
                "replays",
                f"episode_{episode}.json",
            )

        candidates.append({
            "profile": profile,
            "episode": episode,
            "timestamp": row.get("timestamp", ""),
            "duration": _parse_float(row.get("duration")),
            "reward": _parse_float(row.get("reward")),
            "gems": 2,
            "outcome": (row.get("outcome") or "").strip(),
            "closest_gem_distance": _parse_float(row.get("closest_gem_distance")),
            "observation_mode": (row.get("observation_mode") or "").strip(),
            "archive_mode": row.get("archive_mode", ""),
            "archive_path": archive_path,
            "replay_path": replay_path,
            "_replay_payload": None,
        })
    return candidates


def _load_episode_candidates(profile_filter=None):
    candidates = []
    for profile in _candidate_profiles(profile_filter=profile_filter):
        paths = get_profile_paths(profile)
        episodes_csv = paths["episodes_csv"]
        if not os.path.exists(episodes_csv):
            continue

        for row in _load_csv_rows(episodes_csv):
            episode = str(row.get("Episode", "")).strip()
            if not episode:
                continue
            replay_path = os.path.join(paths["replays_dir"], f"episode_{episode}.json")
            candidates.append({
                "profile": profile,
                "episode": episode,
                "timestamp": row.get("Timestamp", ""),
                "duration": _parse_float(row.get("Duration")),
                "reward": _parse_float(row.get("Reward")),
                "gems": _parse_int(row.get("Gems")),
                "outcome": (row.get("Outcome") or "").strip(),
                "closest_gem_distance": _parse_float(row.get("ClosestGemDistance")),
                "replay_path": replay_path,
                "_replay_payload": None,
            })
    return candidates


def _replay_payload_for_row(row):
    if row.get("_replay_payload") is None:
        row["_replay_payload"] = _load_json(row.get("replay_path", ""))
    return row.get("_replay_payload")


def _character_positions_from_payload(payload, character_name):
    positions = []
    if not payload:
        return positions

    for frame in payload.get("overlay_frames") or []:
        state = frame.get("state") or {}
        char_pos = (state.get(character_name) or [None])[0]
        if char_pos is not None:
            positions.append((float(char_pos[0]), float(char_pos[1])))

    for step in payload.get("demo_steps") or []:
        for key, offset in [("obs_before", 0 if character_name == "fireboy" else 2), ("obs_after", 0 if character_name == "fireboy" else 2)]:
            obs = step.get(key)
            if not obs or len(obs) < offset + 2:
                continue
            x = float(obs[offset])
            y = float(obs[offset + 1])
            if x == -1.0 or y == -1.0:
                continue
            positions.append((x, y))
    return positions


def _resolved_ledge_bounds(args, character_name):
    region_name = getattr(args, "ledge_region", "bottom_right") or "bottom_right"
    bounds = dict(LEDGE_REGION_BOUNDS.get(region_name, LEDGE_REGION_BOUNDS["bottom_right"]))
    bounds.pop("label", None)
    prefix = "fireboy" if character_name == "fireboy" else "watergirl"
    for key in ["x_min", "x_max", "y_min", "y_max"]:
        override = getattr(args, f"{prefix}_ledge_{key}", None)
        if override is not None:
            bounds[key] = override
    return bounds


def _character_hit_ledge(row, character_name, x_min, x_max, y_min, y_max):
    payload = _replay_payload_for_row(row)
    if not payload:
        return False
    for x, y in _character_positions_from_payload(payload, character_name):
        if x_min <= x <= x_max and y_min <= y <= y_max:
            return True
    return False


def _character_step_offset(character_name):
    return 0 if character_name == "fireboy" else 2


def _character_backtracks_for_steps(
    row,
    character_name,
    min_consecutive_steps,
    min_delta,
    start_seconds=None,
    end_seconds=None,
):
    payload = _replay_payload_for_row(row)
    if not payload or min_consecutive_steps is None or min_consecutive_steps <= 0:
        return False

    offset = _character_step_offset(character_name)
    positions = []
    for step in payload.get("demo_steps") or []:
        step_time = _parse_float(step.get("t"))
        if step_time is None:
            continue
        if start_seconds is not None and step_time < start_seconds:
            continue
        if end_seconds is not None and step_time > end_seconds:
            continue

        obs = step.get("obs_after") or step.get("obs_before")
        if not obs or len(obs) < offset + 2:
            continue
        x = float(obs[offset])
        y = float(obs[offset + 1])
        if x == -1.0 or y == -1.0:
            continue
        positions.append((step_time, x, y))

    consecutive = 0
    for index in range(1, len(positions)):
        prev_x = positions[index - 1][1]
        curr_x = positions[index][1]
        if curr_x < (prev_x - min_delta):
            consecutive += 1
            if consecutive >= min_consecutive_steps:
                return True
        else:
            consecutive = 0
    return False


def _matches_rules(row, args, review_status):
    if args.exclude_reviewed_bad:
        if review_status.get((row["profile"], str(row["episode"]))) == "bad":
            return False, "reviewed_bad"

    if args.only_reviewed_good:
        if review_status.get((row["profile"], str(row["episode"]))) != "good":
            return False, "not_reviewed_good"

    reward = row.get("reward")
    if args.min_reward is not None and (reward is None or reward < args.min_reward):
        return False, "reward_too_low"
    if args.max_reward is not None and (reward is None or reward > args.max_reward):
        return False, "reward_too_high"

    duration = row.get("duration")
    if args.max_duration is not None and (duration is None or duration > args.max_duration):
        return False, "duration_too_high"
    if args.min_duration is not None and (duration is None or duration < args.min_duration):
        return False, "duration_too_low"

    gems = row.get("gems")
    if args.min_gems is not None and (gems is None or gems < args.min_gems):
        return False, "gems_too_low"
    if args.max_gems is not None and (gems is None or gems > args.max_gems):
        return False, "gems_too_high"

    if args.outcomes:
        allowed = {item.strip().lower() for item in args.outcomes.split(",") if item.strip()}
        if (row.get("outcome") or "").strip().lower() not in allowed:
            return False, "outcome_filtered"

    closest_gem_distance = row.get("closest_gem_distance")
    if args.max_closest_gem_distance is not None:
        if closest_gem_distance is None or closest_gem_distance > args.max_closest_gem_distance:
            return False, "closest_gem_too_high"

    if not os.path.exists(row.get("replay_path", "")):
        return False, "missing_replay"

    fireboy_on_ledge = _character_hit_ledge(row, "fireboy", **_resolved_ledge_bounds(args, "fireboy"))
    watergirl_on_ledge = _character_hit_ledge(row, "watergirl", **_resolved_ledge_bounds(args, "watergirl"))

    if getattr(args, "ledge_character_mode", None):
        should_reach = bool(args.ledge_should_reach)
        if args.ledge_character_mode == "fireboy" and fireboy_on_ledge != should_reach:
            return False, "fireboy_ledge_mismatch"
        if args.ledge_character_mode == "watergirl" and watergirl_on_ledge != should_reach:
            return False, "watergirl_ledge_mismatch"
        if args.ledge_character_mode == "both":
            result = fireboy_on_ledge and watergirl_on_ledge
            if result != should_reach:
                return False, "both_ledge_mismatch"
        if args.ledge_character_mode == "either":
            result = fireboy_on_ledge or watergirl_on_ledge
            if result != should_reach:
                return False, "either_ledge_mismatch"
    else:
        if args.fireboy_on_ledge is not None and fireboy_on_ledge != args.fireboy_on_ledge:
            return False, "fireboy_ledge_mismatch"
        if args.watergirl_on_ledge is not None and watergirl_on_ledge != args.watergirl_on_ledge:
            return False, "watergirl_ledge_mismatch"

    if args.exclude_backtracking_steps is not None and args.exclude_backtracking_steps > 0:
        fireboy_backtracks = _character_backtracks_for_steps(
            row,
            "fireboy",
            args.exclude_backtracking_steps,
            args.backtrack_min_delta,
            start_seconds=args.step_start_seconds,
            end_seconds=args.step_end_seconds,
        )
        watergirl_backtracks = _character_backtracks_for_steps(
            row,
            "watergirl",
            args.exclude_backtracking_steps,
            args.backtrack_min_delta,
            start_seconds=args.step_start_seconds,
            end_seconds=args.step_end_seconds,
        )

        if args.backtrack_character == "fireboy" and fireboy_backtracks:
            return False, "fireboy_backtracking"
        if args.backtrack_character == "watergirl" and watergirl_backtracks:
            return False, "watergirl_backtracking"
        if args.backtrack_character == "both" and fireboy_backtracks and watergirl_backtracks:
            return False, "both_backtracking"
        if args.backtrack_character == "any" and (fireboy_backtracks or watergirl_backtracks):
            return False, "any_backtracking"

    return True, "matched"


def _sort_candidates(rows, sort_by):
    if sort_by == "duration":
        return sorted(rows, key=lambda row: (row.get("duration") is None, row.get("duration", 0.0), -(row.get("reward") or 0.0)))
    if sort_by == "gems":
        return sorted(rows, key=lambda row: (-(row.get("gems") or 0), -(row.get("reward") or 0.0), row.get("duration") or 0.0))
    if sort_by == "closest_gem":
        return sorted(rows, key=lambda row: (row.get("closest_gem_distance") is None, row.get("closest_gem_distance", 999.0), -(row.get("reward") or 0.0)))
    return sorted(rows, key=lambda row: (-(row.get("reward") or 0.0), row.get("duration") or 0.0, -(row.get("gems") or 0)))


def _resolve_output_name(args):
    if args.output_name:
        return args.output_name
    suggested = f"{args.profile or 'all-profiles'}-quick-imitation"
    return _prompt_text("Name this quick dataset", default=suggested)


def _resolve_time_window(args):
    start = args.step_start_seconds
    end = args.step_end_seconds
    if start is None and end is None:
        start = _prompt_optional_float(
            "Only train on steps starting at what second? Leave blank for episode start",
            default=None,
            minimum=0.0,
        )
        end = _prompt_optional_float(
            "Only train on steps up to what second? Leave blank for episode end",
            default=None,
            minimum=0.0,
        )
    if start is not None and end is not None and end < start:
        raise ValueError("step_end_seconds must be >= step_start_seconds")
    return start, end


def _resolve_optional_rules(args):
    ledge_specified = args.fireboy_on_ledge is not None or args.watergirl_on_ledge is not None
    backtracking_specified = args.exclude_backtracking_steps is not None

    if not ledge_specified:
        if _prompt_yes_no("Apply a ledge rule?", default=False):
            args.ledge_region = _prompt_choice(
                "Which ledge should the rule target?",
                [
                    ("bottom_right", LEDGE_REGION_BOUNDS["bottom_right"]["label"]),
                    ("left_large", LEDGE_REGION_BOUNDS["left_large"]["label"]),
                ],
                default_index=0,
            )
            ledge_character = _prompt_choice(
                "Which character should the ledge rule apply to?",
                [
                    ("watergirl", "Watergirl"),
                    ("fireboy", "Fireboy"),
                    ("both", "Both characters"),
                    ("either", "Either character"),
                ],
                default_index=0,
            )
            should_reach = _prompt_yes_no("Should that character be required to reach the ledge?", default=True)
            args.ledge_character_mode = ledge_character
            args.ledge_should_reach = should_reach
            if ledge_character == "watergirl":
                args.watergirl_on_ledge = should_reach
            elif ledge_character == "fireboy":
                args.fireboy_on_ledge = should_reach
            elif ledge_character == "both":
                args.watergirl_on_ledge = should_reach
                args.fireboy_on_ledge = should_reach

    if not backtracking_specified:
        if _prompt_yes_no("Apply the backtracking rule?", default=False):
            args.exclude_backtracking_steps = int(
                _prompt_optional_float(
                    "Exclude an episode when backtracking lasts at least how many consecutive steps?",
                    default=2,
                    minimum=1.0,
                )
            )
            args.backtrack_character = _prompt_choice(
                "Which character should trigger the backtracking rule?",
                [
                    ("any", "Any character"),
                    ("fireboy", "Fireboy only"),
                    ("watergirl", "Watergirl only"),
                    ("both", "Both characters"),
                ],
                default_index=0,
            )
            args.backtrack_min_delta = _prompt_optional_float(
                "What minimum backward x movement should count as backtracking?",
                default=0.01,
                minimum=0.0,
            )
    return args


def _resolve_training_options(args):
    if args.train_output_profile:
        return args

    if not _prompt_yes_no("Train a model from this dataset after building it?", default=False):
        return args

    suggested = f"{args.output_name}-model" if getattr(args, "output_name", None) else "quick-imitation-model"
    args.train_output_profile = _prompt_text(
        "What should the trained model profile be named",
        default=suggested,
    )

    if not args.source_profile:
        source_profile = _prompt_text(
            "Source profile to initialize from? Leave blank for a fresh PPO policy",
            default="",
            allow_blank=True,
        )
        args.source_profile = source_profile or None

    if args.epochs == 10:
        args.epochs = int(_prompt_optional_float("How many imitation epochs", default=10, minimum=1.0))
    if args.batch_size == 256:
        args.batch_size = int(_prompt_optional_float("Batch size", default=256, minimum=1.0))
    if args.learning_rate == 3e-4:
        args.learning_rate = _prompt_optional_float("Learning rate", default=3e-4, minimum=0.0)
    return args


def _step_in_time_window(args, row, payload, step, step_index):
    step_time = _parse_float(step.get("t"))
    if step_time is None:
        return False
    if args.step_start_seconds is not None and step_time < args.step_start_seconds:
        return False
    if args.step_end_seconds is not None and step_time > args.step_end_seconds:
        return False
    return True


def load_candidates_for_args(args):
    review_status = _review_status_map(args.reviews_csv)
    if args.candidate_source == "episodes":
        candidates = _load_episode_candidates(profile_filter=args.profile)
    else:
        candidates_csv = args.candidates_csv or _default_two_gem_candidates_csv()
        candidates = _load_two_gem_candidates(
            candidates_csv,
            profile_filter=args.profile,
        )
    return candidates, review_status


def evaluate_candidates(args):
    candidates, review_status = load_candidates_for_args(args)
    if not candidates:
        return candidates, [], {}, review_status

    matched = []
    skipped_reason_counts = {}
    for row in candidates:
        include, reason = _matches_rules(row, args, review_status)
        if include:
            matched.append(row)
        else:
            skipped_reason_counts[reason] = skipped_reason_counts.get(reason, 0) + 1

    return candidates, matched, skipped_reason_counts, review_status


def build_quick_data_set(args):
    args.output_name = _resolve_output_name(args)
    args.step_start_seconds, args.step_end_seconds = _resolve_time_window(args)
    args = _resolve_optional_rules(args)
    args = _resolve_training_options(args)

    candidates, matched, skipped_reason_counts, review_status = evaluate_candidates(args)
    del review_status
    if not candidates:
        print("No episode summaries were found for the requested profile(s).")
        return None
    if not matched:
        print("No episodes matched the quick dataset rules.")
        if skipped_reason_counts:
            print(f"Skipped breakdown: {skipped_reason_counts}")
        return None

    matched = _sort_candidates(matched, args.sort_by)
    if args.max_episodes is not None:
        matched = matched[: args.max_episodes]

    dataset_path = build_dataset_from_episode_rows(
        matched,
        output_name=args.output_name,
        step_filter=lambda row, payload, step, step_index: _step_in_time_window(args, row, payload, step, step_index),
        manifest_extra={
            "builder": "quick_rules",
            "profile_filter": args.profile,
            "reviews_csv": args.reviews_csv,
            "candidate_source": args.candidate_source,
            "candidates_csv": None if args.candidate_source == "episodes" else (args.candidates_csv or _default_two_gem_candidates_csv()),
            "filters": {
                "min_reward": args.min_reward,
                "max_reward": args.max_reward,
                "min_gems": args.min_gems,
                "max_gems": args.max_gems,
                "min_duration": args.min_duration,
                "max_duration": args.max_duration,
                "outcomes": args.outcomes,
                "max_closest_gem_distance": args.max_closest_gem_distance,
                "sort_by": args.sort_by,
                "max_episodes": args.max_episodes,
                "exclude_reviewed_bad": args.exclude_reviewed_bad,
                "only_reviewed_good": args.only_reviewed_good,
                "step_start_seconds": args.step_start_seconds,
                "step_end_seconds": args.step_end_seconds,
                "ledge_region": getattr(args, "ledge_region", "bottom_right"),
                "fireboy_on_ledge": args.fireboy_on_ledge,
                "watergirl_on_ledge": args.watergirl_on_ledge,
                "ledge_character_mode": getattr(args, "ledge_character_mode", None),
                "ledge_should_reach": getattr(args, "ledge_should_reach", None),
                "fireboy_ledge_bounds": _resolved_ledge_bounds(args, "fireboy"),
                "watergirl_ledge_bounds": _resolved_ledge_bounds(args, "watergirl"),
                "exclude_backtracking_steps": args.exclude_backtracking_steps,
                "backtrack_character": args.backtrack_character,
                "backtrack_min_delta": args.backtrack_min_delta,
            },
            "matched_episode_count": len(matched),
            "candidate_episode_count": len(candidates),
            "skipped_reason_counts": skipped_reason_counts,
        },
    )
    if not dataset_path:
        return None

    print(f"Matched {len(matched)} episode(s) from {len(candidates)} candidate(s).")
    if skipped_reason_counts:
        print(f"Skipped breakdown: {skipped_reason_counts}")

    if args.train_output_profile:
        pretrain_imitation_model(
            dataset_path=dataset_path,
            output_profile=args.train_output_profile,
            source_profile=args.source_profile,
            epochs=args.epochs,
            batch_size=args.batch_size,
            learning_rate=args.learning_rate,
        )

    return dataset_path


def add_quick_rule_arguments(parser):
    parser.add_argument("--profile", help="Limit to a single profile/run")
    parser.add_argument("--reviews-csv", default=DEFAULT_REVIEWS_CSV)
    parser.add_argument(
        "--candidate-source",
        choices=["two_gem", "episodes"],
        default="two_gem",
        help="Where to pull candidate episodes from. Defaults to the archived two-gem candidate set used by review_good_episodes.",
    )
    parser.add_argument("--candidates-csv", help="Optional override CSV when candidate-source is two_gem")
    parser.add_argument("--min-reward", type=float)
    parser.add_argument("--max-reward", type=float)
    parser.add_argument("--min-gems", type=int)
    parser.add_argument("--max-gems", type=int)
    parser.add_argument("--min-duration", type=float)
    parser.add_argument("--max-duration", type=float)
    parser.add_argument("--outcomes", help="Comma-separated outcomes to keep, e.g. Goal,Death")
    parser.add_argument("--max-closest-gem-distance", type=float)
    parser.add_argument("--max-episodes", type=int, help="Take only the top N matched episodes after sorting")
    parser.add_argument(
        "--sort-by",
        choices=["reward", "duration", "gems", "closest_gem"],
        default="reward",
    )
    parser.add_argument(
        "--exclude-reviewed-bad",
        choices=["true", "false"],
        default="true",
        help="Skip episodes you already reviewed and marked not good",
    )
    parser.add_argument(
        "--only-reviewed-good",
        choices=["true", "false"],
        default="false",
        help="Use the quick rules but still restrict to reviewed-good episodes only",
    )
    parser.add_argument("--step-start-seconds", type=float, help="Only keep demo steps at or after this many seconds into each episode")
    parser.add_argument("--step-end-seconds", type=float, help="Only keep demo steps at or before this many seconds into each episode")
    parser.add_argument("--ledge-region", choices=sorted(LEDGE_REGION_BOUNDS.keys()), default="bottom_right", help="Which named ledge region to use for ledge filtering")
    parser.add_argument("--fireboy-on-ledge", choices=["true", "false"], help="Require whether Fireboy ever reaches the selected ledge region")
    parser.add_argument("--watergirl-on-ledge", choices=["true", "false"], help="Require whether Watergirl ever reaches the selected ledge region")
    parser.add_argument("--fireboy-ledge-x-min", type=float, help="Optional normalized X minimum override for Fireboy ledge filtering")
    parser.add_argument("--watergirl-ledge-x-min", type=float, help="Optional normalized X minimum override for Watergirl ledge filtering")
    parser.add_argument("--fireboy-ledge-x-max", type=float, help="Optional normalized X maximum override for Fireboy ledge filtering")
    parser.add_argument("--watergirl-ledge-x-max", type=float, help="Optional normalized X maximum override for Watergirl ledge filtering")
    parser.add_argument("--fireboy-ledge-y-min", type=float, help="Optional normalized Y minimum override for Fireboy ledge filtering")
    parser.add_argument("--watergirl-ledge-y-min", type=float, help="Optional normalized Y minimum override for Watergirl ledge filtering")
    parser.add_argument("--fireboy-ledge-y-max", type=float, help="Optional normalized Y maximum override for Fireboy ledge filtering")
    parser.add_argument("--watergirl-ledge-y-max", type=float, help="Optional normalized Y maximum override for Watergirl ledge filtering")
    parser.add_argument("--exclude-backtracking-steps", type=int, help="Exclude episodes where a character backtracks for at least this many consecutive demo steps")
    parser.add_argument(
        "--backtrack-character",
        choices=["any", "fireboy", "watergirl", "both"],
        default="any",
        help="Which character backtracking should trigger exclusion",
    )
    parser.add_argument(
        "--backtrack-min-delta",
        type=float,
        default=0.01,
        help="Minimum normalized x decrease counted as a backtracking step",
    )


def normalize_rule_args(args):
    if not hasattr(args, "ledge_character_mode"):
        args.ledge_character_mode = None
    if not hasattr(args, "ledge_should_reach"):
        args.ledge_should_reach = None
    if not hasattr(args, "ledge_region") or not args.ledge_region:
        args.ledge_region = "bottom_right"
    args.exclude_reviewed_bad = args.exclude_reviewed_bad == "true"
    args.only_reviewed_good = args.only_reviewed_good == "true"
    args.fireboy_on_ledge = None if args.fireboy_on_ledge is None else args.fireboy_on_ledge == "true"
    args.watergirl_on_ledge = None if args.watergirl_on_ledge is None else args.watergirl_on_ledge == "true"
    return args


def parse_args():
    parser = argparse.ArgumentParser(
        description="Build a quick imitation dataset from episode rules, with optional model pretraining.",
    )
    parser.add_argument("--output-name", help="Dataset name written under runs/analysis/imitation_datasets")
    add_quick_rule_arguments(parser)
    parser.add_argument("--train-output-profile", help="If set, immediately pretrain a model from the new dataset")
    parser.add_argument("--source-profile", help="Optional PPO profile to initialize from before imitation pretraining")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=3e-4)
    return normalize_rule_args(parser.parse_args())


if __name__ == "__main__":
    build_quick_data_set(parse_args())