scripts/audit_map_distribution.py

from __future__ import annotations

import argparse
import json
from dataclasses import replace
from typing import Any

import torch

from config import V2Config
from env import FastFloatingBeaconEnv
from mapgen import configure_motif_mapgen, install_motif_mapgen, mapgen_summary
from settings import ENV_DEFAULTS, MAPGEN_KWARGS, PHYSICS_DEFAULTS


STAGES: dict[str, tuple[int, int]] = {
    "rookie": (10, 6),
    "standard": (14, 10),
    "cursed": (18, 14),
    "nightmare": (22, 18),
}


def cfg_for(stage: str, route_jumps: int, distractors: int, *, envs: int, seed: int, device: str) -> V2Config:
    fields = V2Config.__dataclass_fields__
    cfg = V2Config(
        seed=int(seed),
        device=str(device),
        envs=int(envs),
        eval_envs=int(envs),
        route_jumps=int(route_jumps),
        distractors=int(distractors),
        pillar_platforms=False,
        sensor_mode="topk",
        sensor_topk=16,
        sensor_token_range=9.5,
        sensor_fov_degrees=120.0,
        sensor_token_sort="distance",
        observe_progress=False,
        **{key: value for key, value in PHYSICS_DEFAULTS.items() if key in fields},
        **{key: value for key, value in ENV_DEFAULTS.items() if key in fields},
        **{key: value for key, value in MAPGEN_KWARGS.items() if key in fields},
    )
    return replace(cfg, run_dir=f"data/audit/{stage}")


def audit_stage(stage: str, route_jumps: int, distractors: int, *, envs: int, seed: int, device: str) -> dict[str, Any]:
    cfg = cfg_for(stage, route_jumps, distractors, envs=envs, seed=seed, device=device)
    env = FastFloatingBeaconEnv(cfg, envs=envs, seed=seed)
    env.reset()
    summary = mapgen_summary(env)
    route = max(1, int(route_jumps))
    shortcut_ratio = float(summary["all_path_shortcut_ratio_mean"])
    return {
        "stage": stage,
        "seed": int(seed),
        "envs": int(envs),
        "route_jumps": int(route_jumps),
        "distractors": int(distractors),
        "platforms": int(summary["platforms"]),
        "goal_xy_min": float(summary["goal_xy_min"]),
        "goal_xy_mean": float(summary["goal_xy_mean"]),
        "height_span_mean": float(summary["height_span_mean"]),
        "route_path_xy_mean": float(summary["route_path_xy_mean"]),
        "goal_xy_to_route_path_ratio_mean": float(summary["goal_xy_to_route_path_ratio_mean"]),
        "extra_after_goal_fraction": float(summary["extra_after_goal_fraction"]),
        "extra_before_start_fraction": float(summary["extra_before_start_fraction"]),
        "extra_between_start_goal_fraction": float(summary["extra_between_start_goal_fraction"]),
        "route_edge_gap_mean": float(summary["route_edge_gap_mean"]),
        "route_edge_gap_min": float(summary["route_edge_gap_min"]),
        "shortest_hops_mean": float(summary["all_path_shortest_hops_mean"]),
        "shortest_hops_p50": float(summary["all_path_shortest_hops_p50"]),
        "shortcut_ratio_mean": shortcut_ratio,
        "shortcut_le_half_fraction": float(summary["all_path_shortcut_le_half_fraction"]),
        "shortcut_le_five_fraction": float(summary["all_path_shortcut_le_five_fraction"]),
        "nearest_edge_gap_mean": float(summary["all_platform_nearest_edge_gap_mean"]),
        "nearest_edge_gap_min": float(summary["all_platform_nearest_edge_gap_min"]),
        "tiny_platform_fraction": float(summary["tiny_platform_fraction"]),
        "route_away_goal_fraction": float(summary["route_away_goal_fraction"]),
        "route_side_goal_fraction": float(summary["route_side_goal_fraction"]),
        "route_down_step_fraction": float(summary["route_down_step_fraction"]),
        "pass_basic_hardness": bool(
            float(summary["goal_xy_min"]) >= max(16.0, route * 1.45)
            and float(summary["goal_xy_mean"]) >= max(18.0, route * 1.75)
            and float(summary["extra_after_goal_fraction"]) <= 0.02
            and shortcut_ratio >= 0.58
            and float(summary["all_path_shortcut_le_five_fraction"]) <= 0.05
            and float(summary["all_path_shortcut_le_half_fraction"]) <= 0.25
        ),
    }


def parse_stage(value: str) -> tuple[str, int, int]:
    if ":" in value:
        name, route, distractors = value.split(":", 2)
        return name.strip(), int(route), int(distractors)
    route, distractors = STAGES[value]
    return value, route, distractors


def main() -> None:
    parser = argparse.ArgumentParser(description="Audit procedural parkour maps for shortcut/easiness inflation.")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--envs", type=int, default=512)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument(
        "--stage",
        action="append",
        default=[],
        help="Stage name or name:route_jumps:distractors. Defaults to all app stages.",
    )
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    configure_motif_mapgen(replace=True, **MAPGEN_KWARGS)
    install_motif_mapgen()
    stages = [parse_stage(item) for item in args.stage] if args.stage else [
        (name, route, distractors) for name, (route, distractors) in STAGES.items()
    ]
    rows = [
        audit_stage(name, route, distractors, envs=int(args.envs), seed=int(args.seed) + i * 1009, device=str(args.device))
        for i, (name, route, distractors) in enumerate(stages)
    ]
    if args.json:
        print(json.dumps({"rows": rows}, indent=2, sort_keys=True))
        return
    print(
        "stage route distr goal_min goal_mean path_xy straight extra_after shortest_p50 ratio le_half le_five edge_gap tiny away side pass",
        flush=True,
    )
    for row in rows:
        print(
            f"{row['stage']} {row['route_jumps']} {row['distractors']} "
            f"{row['goal_xy_min']:.2f} {row['goal_xy_mean']:.2f} "
            f"{row['route_path_xy_mean']:.2f} {row['goal_xy_to_route_path_ratio_mean']:.2f} "
            f"{row['extra_after_goal_fraction']:.3f} "
            f"{row['shortest_hops_p50']:.2f} {row['shortcut_ratio_mean']:.2f} "
            f"{row['shortcut_le_half_fraction']:.3f} {row['shortcut_le_five_fraction']:.3f} "
            f"{row['nearest_edge_gap_mean']:.2f} {row['tiny_platform_fraction']:.2f} "
            f"{row['route_away_goal_fraction']:.2f} {row['route_side_goal_fraction']:.2f} "
            f"{'yes' if row['pass_basic_hardness'] else 'NO'}",
            flush=True,
        )


if __name__ == "__main__":
    torch.set_float32_matmul_precision("high")
    main()