from __future__ import annotations import argparse from dataclasses import dataclass, replace import sys import time import types from typing import Any import torch import config as _checkpoint_config from config import V2Config from env import FastFloatingBeaconEnv from mapgen import mapgen_summary from replay import V2RunWriter from runtime import resolve_device, sync_device from runner import BeaconController, LocalBeaconPlannerController, TokenAttentionController # Old checkpoints were pickled with e2e_runner_v2.config.V2Config. # Keep this alias only for torch.load compatibility. _compat_e2e = sys.modules.setdefault("e2e_runner_v2", types.ModuleType("e2e_runner_v2")) _compat_e2e.config = _checkpoint_config _compat_e2e.train_fast_beacon_planner = sys.modules[__name__] sys.modules.setdefault("e2e_runner_v2.config", _checkpoint_config) sys.modules.setdefault("e2e_runner_v2.train_fast_beacon_planner", sys.modules[__name__]) @dataclass(slots=True) class StageSpec: route_jumps: int distractors: int max_updates: int threshold: float def main() -> None: parser = argparse.ArgumentParser(description="Fast pure-Torch egocentric beacon planner PPO.") parser.add_argument("--run-name", default="fast_beacon_planner_v1") parser.add_argument("--device", default="cuda") parser.add_argument("--envs", type=int, default=32768) parser.add_argument("--eval-envs", type=int, default=2048) parser.add_argument("--steps", type=int, default=64) parser.add_argument("--epochs", type=int, default=1) parser.add_argument("--minibatches", type=int, default=16) parser.add_argument("--hidden", type=int, default=128) parser.add_argument("--depth", type=int, default=2) parser.add_argument("--policy-arch", choices=("mlp", "attention"), default="mlp") parser.add_argument("--planner-mode", choices=("none", "local-beacon"), default="local-beacon") parser.add_argument("--planner-hidden", type=int, default=128) parser.add_argument("--planner-depth", type=int, default=2) parser.add_argument("--planner-max-xy", type=float, default=8.0) parser.add_argument("--planner-max-z", type=float, default=1.0) parser.add_argument("--planner-blend-bias", type=float, default=-2.0) parser.add_argument("--reset-planner-head", action="store_true") parser.add_argument("--capture", type=int, default=8) parser.add_argument("--log-every", type=int, default=10) parser.add_argument("--min-stage-updates", type=int, default=10) parser.add_argument("--entropy-coef", type=float, default=0.012) parser.add_argument("--sil-coef", type=float, default=0.0) parser.add_argument("--resume-checkpoint", default=None) parser.add_argument("--anchor-checkpoint", default=None) parser.add_argument("--anchor-bc-coef", type=float, default=0.0) parser.add_argument("--map-height-offset", type=float, default=2.75) parser.add_argument("--map-vertical-scale", type=float, default=1.45) parser.add_argument("--map-max-height", type=float, default=5.75) parser.add_argument("--map-loopback-depth", type=float, default=0.58) parser.add_argument("--map-loopback-lateral", type=float, default=0.52) parser.add_argument("--map-loopback-period", type=int, default=6) parser.add_argument("--map-curve-scale", type=float, default=1.30) parser.add_argument("--map-braid-lanes", type=int, default=3) parser.add_argument("--map-lane-width", type=float, default=1.55) parser.add_argument("--map-style-jitter", type=float, default=0.35) parser.add_argument("--map-switchback-period", type=int, default=4) parser.add_argument("--map-cleanup-iters", type=int, default=24) parser.add_argument("--map-illegal-decoy-fraction", type=float, default=0.05) parser.add_argument("--floating-platforms", action="store_true") parser.add_argument("--goal-progress-reward-scale", type=float, default=2.0) parser.add_argument("--first-visit-reward", type=float, default=0.04) parser.add_argument("--curriculum", default="12:30:240:0.35,16:42:320:0.28,20:54:420:0.22,24:66:540:0.16") parser.add_argument("--quick", action="store_true") args = parser.parse_args() cfg = V2Config( seed=17, device=str(args.device), envs=int(args.envs), eval_envs=int(args.eval_envs), ppo_steps=int(args.steps), ppo_epochs=int(args.epochs), minibatches=int(args.minibatches), hidden=int(args.hidden), lr=3.0e-4, gamma=0.992, gae_lambda=0.94, entropy_coef=float(args.entropy_coef), value_coef=0.55, max_steps_factor=24.0, map_height_offset=float(args.map_height_offset), map_vertical_scale=float(args.map_vertical_scale), map_max_height=float(args.map_max_height), map_braid_loopback_depth=float(args.map_loopback_depth), map_braid_loopback_lateral=float(args.map_loopback_lateral), map_braid_loopback_period=int(args.map_loopback_period), map_braid_curve_scale=float(args.map_curve_scale), map_braid_lanes=int(args.map_braid_lanes), map_braid_lane_width=float(args.map_lane_width), map_braid_style_jitter=float(args.map_style_jitter), map_switchback_period=int(args.map_switchback_period), map_overlap_cleanup_iters=int(args.map_cleanup_iters), map_illegal_decoy_fraction=float(args.map_illegal_decoy_fraction), goal_progress_reward_scale=float(args.goal_progress_reward_scale), first_visit_reward=float(args.first_visit_reward), pillar_platforms=not bool(args.floating_platforms), ) if args.quick: cfg = replace(cfg, envs=min(int(args.envs), 2048), eval_envs=min(int(args.eval_envs), 256), ppo_steps=24, ppo_epochs=1, minibatches=4) args.curriculum = "4:8:2:0.05" args.log_every = 1 args.min_stage_updates = 1 args.capture = min(int(args.capture), 2) result = train( cfg, str(args.run_name), depth=int(args.depth), policy_arch=str(args.policy_arch), planner_mode=str(args.planner_mode), planner_hidden=int(args.planner_hidden), planner_depth=int(args.planner_depth), planner_max_xy=float(args.planner_max_xy), planner_max_z=float(args.planner_max_z), planner_blend_bias=float(args.planner_blend_bias), capture=int(args.capture), log_every=int(args.log_every), min_stage_updates=int(args.min_stage_updates), resume_checkpoint=args.resume_checkpoint, reset_planner_head=bool(args.reset_planner_head), curriculum=parse_curriculum(str(args.curriculum)), sil_coef=float(args.sil_coef), anchor_checkpoint=args.anchor_checkpoint, anchor_bc_coef=float(args.anchor_bc_coef), ) print(result, flush=True) def parse_curriculum(value: str) -> list[StageSpec]: stages: list[StageSpec] = [] for chunk in str(value or "").split(","): parts = [x.strip() for x in chunk.split(":")] if len(parts) not in (3, 4): continue try: jumps = int(parts[0]) distractors = int(parts[1]) updates = int(parts[2]) threshold = float(parts[3]) if len(parts) == 4 else 0.50 except ValueError: continue stages.append(StageSpec(jumps, distractors, updates, min(max(threshold, 0.0), 1.0))) return stages or [StageSpec(12, 30, 240, 0.35)] def make_model( obs_size: int, cfg: V2Config, *, depth: int, planner_mode: str, planner_hidden: int, planner_depth: int, planner_max_xy: float, planner_max_z: float, planner_blend_bias: float, policy_arch: str = "mlp", ) -> torch.nn.Module: if policy_arch == "attention" and planner_mode != "local-beacon": return TokenAttentionController(obs_size, hidden=int(cfg.hidden), depth=int(depth)) if planner_mode == "local-beacon": return LocalBeaconPlannerController( obs_size, hidden=int(cfg.hidden), depth=int(depth), planner_hidden=int(planner_hidden), planner_depth=int(planner_depth), planner_max_xy=float(planner_max_xy), planner_max_z=float(planner_max_z), planner_blend_bias=float(planner_blend_bias), ) return BeaconController(obs_size, hidden=int(cfg.hidden), depth=int(depth)) def load_compatible_state(module: torch.nn.Module, state: dict[str, torch.Tensor]) -> tuple[int, int]: """Load matching tensors, widening top-K token input weights when needed.""" target = module.state_dict() adapted: dict[str, torch.Tensor] = {} exact = 0 widened = 0 for key, target_value in target.items(): source_value = state.get(key) if source_value is None: continue if tuple(source_value.shape) == tuple(target_value.shape): adapted[key] = source_value exact += 1 continue if ( key.endswith("encoder.0.weight") and source_value.ndim == 2 and target_value.ndim == 2 and source_value.shape[0] == target_value.shape[0] ): widened_value = target_value.detach().clone() old_cols = int(source_value.shape[1]) new_cols = int(target_value.shape[1]) base = int(FastFloatingBeaconEnv.base_features) old_token = 8 new_token = int(FastFloatingBeaconEnv.token_features) if old_cols >= base and new_cols >= base and (old_cols - base) % old_token == 0 and (new_cols - base) % new_token == 0: old_k = (old_cols - base) // old_token new_k = (new_cols - base) // new_token if old_k == new_k: widened_value[:, :base] = source_value[:, :base] for token in range(old_k): old_start = base + token * old_token new_start = base + token * new_token widened_value[:, new_start : new_start + old_token] = source_value[:, old_start : old_start + old_token] adapted[key] = widened_value widened += 1 module.load_state_dict(adapted, strict=False) return exact, widened def train( cfg: V2Config, run_name: str, *, depth: int, planner_mode: str, planner_hidden: int, planner_depth: int, planner_max_xy: float, planner_max_z: float, planner_blend_bias: float, policy_arch: str = "mlp", capture: int, log_every: int, min_stage_updates: int, resume_checkpoint: str | None, reset_planner_head: bool, curriculum: list[StageSpec], sil_coef: float = 0.0, anchor_checkpoint: str | None = None, anchor_bc_coef: float = 0.0, ) -> dict[str, Any]: torch.set_float32_matmul_precision("high") device = resolve_device(cfg.device) probe = FastFloatingBeaconEnv(cfg, envs=1, seed=int(cfg.seed)) model = make_model( probe.obs_size, cfg, depth=depth, policy_arch=policy_arch, planner_mode=planner_mode, planner_hidden=planner_hidden, planner_depth=planner_depth, planner_max_xy=planner_max_xy, planner_max_z=planner_max_z, planner_blend_bias=planner_blend_bias, ).to(device) if resume_checkpoint: payload = torch.load(str(resume_checkpoint), map_location=device, weights_only=False) state = payload.get("model", payload) if isinstance(payload, dict) else payload if isinstance(state, dict): exact, widened = load_compatible_state(model, state) if exact + widened > 0: print(f"loaded_resume_checkpoint path={resume_checkpoint} exact={exact} widened={widened}", flush=True) else: if any(key.startswith("runner.") for key in state): runner_state = { key[len("runner.") :]: value for key, value in state.items() if key.startswith("runner.") } else: runner_state = state if isinstance(model, LocalBeaconPlannerController): exact, widened = load_compatible_state(model.runner, runner_state) else: exact, widened = load_compatible_state(model, runner_state) print(f"loaded_resume_runner_checkpoint path={resume_checkpoint} exact={exact} widened={widened}", flush=True) if reset_planner_head and isinstance(model, LocalBeaconPlannerController): model.reset_planner_head(float(planner_blend_bias)) print(f"reset_planner_head blend_bias={float(planner_blend_bias):.3f}", flush=True) anchor_model: torch.nn.Module | None = None if float(anchor_bc_coef) > 0.0: anchor_path = str(anchor_checkpoint or resume_checkpoint or "") if not anchor_path: print("anchor_checkpoint_missing path=none; disabling anchor_bc", flush=True) anchor_bc_coef = 0.0 else: anchor_model = make_model( probe.obs_size, cfg, depth=depth, policy_arch=policy_arch, planner_mode=planner_mode, planner_hidden=planner_hidden, planner_depth=planner_depth, planner_max_xy=planner_max_xy, planner_max_z=planner_max_z, planner_blend_bias=planner_blend_bias, ).to(device) payload = torch.load(anchor_path, map_location=device, weights_only=False) state = payload.get("model", payload) if isinstance(payload, dict) else payload if isinstance(state, dict): exact, widened = load_compatible_state(anchor_model, state) if exact + widened == 0 and any(key.startswith("runner.") for key in state): runner_state = { key[len("runner.") :]: value for key, value in state.items() if key.startswith("runner.") } if isinstance(anchor_model, LocalBeaconPlannerController): exact, widened = load_compatible_state(anchor_model.runner, runner_state) else: exact, widened = load_compatible_state(anchor_model, runner_state) print( f"loaded_anchor_checkpoint path={anchor_path} exact={exact} widened={widened} coef={float(anchor_bc_coef):.4f}", flush=True, ) anchor_model.eval() for parameter in anchor_model.parameters(): parameter.requires_grad_(False) optimizer = torch.optim.AdamW(model.parameters(), lr=float(cfg.lr), weight_decay=1e-4) writer = V2RunWriter(cfg.run_dir, run_name, cfg, algorithm="fast_beacon_planner") started = time.perf_counter() write_index = 0 global_update = 0 last_stats: dict[str, float] = {} def checkpoint(stage_cfg: V2Config, stage: StageSpec, name: str, local: int, total: int, stats: dict[str, float]) -> bool: nonlocal write_index eval_stats, frames, successes = evaluate(model, stage_cfg, capture=capture) mastery_success = float(eval_stats["eval_success"]) mastered = local >= int(min_stage_updates) and mastery_success >= float(stage.threshold) merged = { **stats, **eval_stats, "update": float(total), "stage_update": float(local), "stage_threshold": float(stage.threshold), "stage_mastered": 1.0 if mastered else 0.0, "phase_seconds": time.perf_counter() - started, } maps, rollouts = replay_payload(frames, successes) write_index += 1 writer.write_replay( write_index, f"fast_beacon_{name}_{int(local):04d}", stage_cfg, maps, rollouts, merged, meta_extra={ "controller": "egocentric goal-beacon runner PPO", "planner": "learned egocentric local-beacon delta PPO" if planner_mode == "local-beacon" else "none", "teacher": "none", "route_observation": "none", "physics_backend": "torch_fast_kinematic", "goal_beacon": "true egocentric final-goal beacon; planner may rewrite local beacon fields from the same observation", "map_distribution": "GPU-generated floating-platform courses with controllable lanes, switchbacks, vertical variation, and legal/illegal decoys", "planner_mode": str(planner_mode), "resume_checkpoint": str(resume_checkpoint or ""), "self_imitation": f"successful on-policy trajectories only; coef={float(sil_coef):.4f}", "anchor_bc": f"stable policy action imitation on current egocentric obs only; coef={float(anchor_bc_coef):.4f}", }, ) torch.save( { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "config": cfg, "stage_config": stage_cfg, "stage": stage, "local_update": int(local), "total_update": int(total), "curriculum": curriculum, "sil_coef": float(sil_coef), "stats": merged, }, writer.path / "fast_beacon_latest.pt", ) print(format_stats(f"fast_beacon_{name}_{int(local):04d}", merged), flush=True) return mastered for stage_i, stage in enumerate(curriculum, start=1): stage_cfg = replace(cfg, route_jumps=int(stage.route_jumps), distractors=int(stage.distractors), ppo_updates=int(stage.max_updates)) stage_name = f"s{stage_i}_g{stage.route_jumps}_d{stage.distractors}" offset = int(global_update) last_stats, used, mastered = train_stage( model, optimizer, stage_cfg, log_every=log_every, stage=stage, callback=lambda local, stats, scfg=stage_cfg, spec=stage, name=stage_name, off=offset: checkpoint( scfg, spec, name, local, off + local, stats ), sil_coef=float(sil_coef), anchor_model=anchor_model, anchor_bc_coef=float(anchor_bc_coef), ) global_update += int(used) if not mastered: print(f"fast_beacon_{stage_name}_not_mastered max_updates={used} threshold={stage.threshold:.3f}", flush=True) return {"run": writer.id, "stats": last_stats} def train_stage( model: torch.nn.Module, optimizer: torch.optim.Optimizer, cfg: V2Config, *, log_every: int, stage: StageSpec, callback, sil_coef: float = 0.0, anchor_model: torch.nn.Module | None = None, anchor_bc_coef: float = 0.0, ) -> tuple[dict[str, float], int, bool]: device = next(model.parameters()).device env = FastFloatingBeaconEnv(cfg, envs=int(cfg.envs), seed=int(cfg.seed) + 70_000 + int(cfg.route_jumps) * 997 + int(cfg.distractors)) started = time.perf_counter() env_steps = 0 last: dict[str, float] = {} mastered = False for update in range(int(cfg.ppo_updates)): collect_start = time.perf_counter() batch, last_value = collect_rollout(model, env, int(cfg.ppo_steps), device) sync_device(device) collect_elapsed = max(1e-6, time.perf_counter() - collect_start) rollout_steps = int(cfg.ppo_steps) * int(env.n) env_steps += rollout_steps adv, returns = advantages(batch, last_value, cfg) valid_adv = adv[batch["masks"] > 0.5] adv = ((adv - valid_adv.mean()) / (valid_adv.std() + 1e-8)).clamp(-10.0, 10.0) flat_obs = batch["obs"].reshape(-1, batch["obs"].shape[-1]) flat_actions = batch["actions"].reshape(-1, batch["actions"].shape[-1]) flat_logprobs = batch["logprobs"].reshape(-1) flat_adv = adv.reshape(-1) flat_returns = returns.reshape(-1) flat_masks = batch["masks"].reshape(-1) flat_success_env = batch["success"].reshape(1, -1).expand(batch["masks"].shape[0], -1).reshape(-1) total = int(flat_obs.shape[0]) mb = max(1, total // int(cfg.minibatches)) last_loss = torch.zeros((), device=device) last_sil_loss = torch.zeros((), device=device) last_anchor_loss = torch.zeros((), device=device) update_start = time.perf_counter() for _epoch in range(int(cfg.ppo_epochs)): order = torch.randperm(total, device=device) for start_i in range(0, total, mb): ids = order[start_i : start_i + mb] mask = flat_masks[ids] > 0.5 if int(mask.sum()) == 0: continue logs, entropy, value = model.evaluate_actions(flat_obs[ids], flat_actions[ids]) ratio = (logs - flat_logprobs[ids]).clamp(-10.0, 10.0).exp() pg_all = -torch.minimum( flat_adv[ids] * ratio, flat_adv[ids] * ratio.clamp(1.0 - float(cfg.clip), 1.0 + float(cfg.clip)), ) pg = pg_all[mask].mean() vf = 0.5 * (value - flat_returns[ids]).square()[mask].mean() ent = entropy[mask].mean() sil_loss = torch.zeros((), device=device) if float(sil_coef) > 0.0: sil_mask = mask & flat_success_env[ids] if bool(sil_mask.any().detach().cpu()): sil_weight = flat_adv[ids].detach().clamp_min(0.0)[sil_mask] if bool((sil_weight > 0.0).any().detach().cpu()): sil_weight = sil_weight / (sil_weight.mean().detach() + 1e-6) sil_loss = -(logs[sil_mask] * sil_weight).mean() else: sil_loss = -logs[sil_mask].mean() anchor_loss = torch.zeros((), device=device) if anchor_model is not None and float(anchor_bc_coef) > 0.0: with torch.no_grad(): anchor_action, _anchor_out = anchor_model.choose_action( flat_obs[ids], deterministic_controls=True, deterministic_buttons=True, ) anchor_logs, _anchor_entropy, _anchor_value = model.evaluate_actions(flat_obs[ids], anchor_action) anchor_loss = -anchor_logs[mask].mean() loss = ( pg + float(cfg.value_coef) * vf - float(cfg.entropy_coef) * ent + float(sil_coef) * sil_loss + float(anchor_bc_coef) * anchor_loss ) optimizer.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() last_loss = loss.detach() last_sil_loss = sil_loss.detach() last_anchor_loss = anchor_loss.detach() sync_device(device) update_elapsed = max(1e-6, time.perf_counter() - update_start) elapsed = max(1e-6, time.perf_counter() - started) local = update + 1 last = { "ppo_loss": float(last_loss.detach().cpu()), "ppo_return": float(batch["rewards"].sum(0).mean().detach().cpu()), "ppo_success": float(batch["success"].float().mean().detach().cpu()), "episode_success_rate": float(batch["success_events"].sum().detach().cpu() / (batch["dones"].sum().detach().cpu() + 1e-6)), "ppo_progress": float(batch["progress"].mean().detach().cpu()), "sil_loss": float(last_sil_loss.detach().cpu()), "anchor_bc_loss": float(last_anchor_loss.detach().cpu()), "sil_fraction": float(((flat_masks > 0.5) & flat_success_env).float().sum().detach().cpu() / ((flat_masks > 0.5).float().sum().detach().cpu() + 1e-6)), "sil_positive_fraction": float((((flat_masks > 0.5) & flat_success_env) & (flat_adv > 0.0)).float().sum().detach().cpu() / ((flat_masks > 0.5).float().sum().detach().cpu() + 1e-6)), "jump_fraction": float(batch["actions"][..., 3].float().mean().detach().cpu()), "grounded_jump_fraction": float((batch["actions"][..., 3].float() * (batch["obs"][..., 3] > 0.5).float()).sum().detach().cpu() / ((batch["obs"][..., 3] > 0.5).float().sum().detach().cpu() + 1e-6)), "sprint_fraction": float(batch["actions"][..., 4].float().mean().detach().cpu()), "forward_mean": float(batch["actions"][..., 0].float().mean().detach().cpu()), "side_abs": float(batch["actions"][..., 1].float().abs().mean().detach().cpu()), "turn_abs": float(batch["actions"][..., 2].float().abs().mean().detach().cpu()), "control_std": float(model.policy_log_std().exp().mean().detach().cpu()), "env_sps": float(env_steps / elapsed), "collect_sps": float(rollout_steps / collect_elapsed), "update_sps": float(rollout_steps / update_elapsed), "env_steps": float(env_steps), "ppo_envs": float(env.n), "stage_route_jumps": float(cfg.route_jumps), "stage_distractors": float(cfg.distractors), } if hasattr(model, "planner_summary"): last.update(model.planner_summary(flat_obs[flat_masks > 0.5])) if callback is not None and int(log_every) > 0 and local % int(log_every) == 0: mastered = bool(callback(local, last)) if mastered: return last, local, True return last, int(cfg.ppo_updates), mastered @torch.no_grad() def collect_rollout(model: torch.nn.Module, env: FastFloatingBeaconEnv, steps: int, device: torch.device) -> tuple[dict[str, torch.Tensor], torch.Tensor]: obs = env.reset_done().to(device=device, dtype=torch.float32) obs_dim = int(obs.shape[-1]) count = int(env.n) obs_buf = torch.empty((steps, count, obs_dim), device=device) action_buf = torch.empty((steps, count, model.action_size), device=device) logprob_buf = torch.empty((steps, count), device=device) reward_buf = torch.empty((steps, count), device=device) done_buf = torch.empty((steps, count), device=device) mask_buf = torch.empty((steps, count), device=device) value_buf = torch.empty((steps, count), device=device) progress_buf = torch.empty((steps, count), device=device) success_event_buf = torch.empty((steps, count), device=device) success_seen = torch.zeros(count, dtype=torch.bool, device=device) for t in range(int(steps)): active = ~env.done obs_buf[t].copy_(obs) action, logprob, _entropy, value = model.act(obs, deterministic=False) step = env.step(action) success = step.info["success"].to(device=device) success_seen |= success action_buf[t].copy_(action) logprob_buf[t].copy_(logprob) reward_buf[t].copy_(step.reward.to(device=device)) done_buf[t].copy_(step.done.to(device=device, dtype=torch.float32)) mask_buf[t].copy_(active.to(device=device, dtype=torch.float32)) value_buf[t].copy_(value) progress_buf[t].copy_(step.info["progress"].to(device=device)) success_event_buf[t].copy_(success.to(dtype=torch.float32)) obs = step.observation.to(device=device, dtype=torch.float32) last_value = model(obs).value return { "obs": obs_buf, "actions": action_buf, "logprobs": logprob_buf, "rewards": reward_buf, "dones": done_buf, "masks": mask_buf, "values": value_buf, "progress": progress_buf, "success_events": success_event_buf, "success": success_seen, }, last_value @torch.no_grad() def advantages(batch: dict[str, torch.Tensor], last_value: torch.Tensor, cfg: V2Config) -> tuple[torch.Tensor, torch.Tensor]: rewards, dones, values, masks = batch["rewards"], batch["dones"], batch["values"], batch["masks"] adv = torch.zeros_like(rewards) gae = torch.zeros(rewards.shape[1], device=rewards.device) for t in reversed(range(int(rewards.shape[0]))): next_value = last_value if t == int(rewards.shape[0]) - 1 else values[t + 1] nonterminal = 1.0 - dones[t] delta = rewards[t] + float(cfg.gamma) * next_value * nonterminal - values[t] gae = (delta + float(cfg.gamma) * float(cfg.gae_lambda) * nonterminal * gae) * masks[t] adv[t] = gae return adv, adv + values @torch.no_grad() def evaluate(model: torch.nn.Module, cfg: V2Config, *, capture: int) -> tuple[dict[str, float], list[list[dict[str, Any]]], list[bool]]: device = next(model.parameters()).device env = FastFloatingBeaconEnv(cfg, envs=int(cfg.eval_envs), seed=int(cfg.seed) + 300_000) obs = env.reset().to(device=device, dtype=torch.float32) geometry = mapgen_summary(env) capture_ids = list(range(min(int(capture), env.n))) frames = [[] for _ in capture_ids] recording = [True for _ in capture_ids] episode_done = torch.zeros(env.n, dtype=torch.bool, device=device) success_seen = torch.zeros(env.n, dtype=torch.bool, device=device) for row, snap in enumerate(env.snapshots(capture_ids)): frames[row].append(snap) for _ in range(env.max_steps): action, _out = model.choose_action(obs, deterministic_controls=True, deterministic_buttons=True) step = env.step(action) done = step.done.to(device=device) success_seen |= (~episode_done) & step.info["success"].to(device=device) active_done = (~episode_done) & done episode_done |= done obs = step.observation.to(device=device, dtype=torch.float32) for row, snap in enumerate(env.snapshots(capture_ids)): if not recording[row]: continue frames[row].append(snap) if bool(active_done[capture_ids[row]].detach().cpu()): recording[row] = False if bool(episode_done.all().detach().cpu()): break successes = [bool(success_seen[i].detach().cpu()) for i in capture_ids] return { "eval_success": float(success_seen.float().mean().detach().cpu()), "eval_progress": float(env.progress_tensor().mean().detach().cpu()), "eval_steps": float(env.steps.float().mean().detach().cpu()), "map_goal_xy_min": float(geometry.get("goal_xy_min", 0.0)), "map_goal_xy_mean": float(geometry.get("goal_xy_mean", 0.0)), "map_route_path_xy_mean": float(geometry.get("route_path_xy_mean", 0.0)), "map_goal_xy_to_route_path_ratio_mean": float(geometry.get("goal_xy_to_route_path_ratio_mean", 0.0)), "map_extra_after_goal_fraction": float(geometry.get("extra_after_goal_fraction", 0.0)), "map_shortest_hops_mean": float(geometry.get("all_path_shortest_hops_mean", 0.0)), "map_shortest_hops_p50": float(geometry.get("all_path_shortest_hops_p50", 0.0)), "map_shortcut_ratio_mean": float(geometry.get("all_path_shortcut_ratio_mean", 0.0)), "map_shortcut_le_half_fraction": float(geometry.get("all_path_shortcut_le_half_fraction", 0.0)), "map_shortcut_le_five_fraction": float(geometry.get("all_path_shortcut_le_five_fraction", 0.0)), "map_nearest_edge_gap_mean": float(geometry.get("all_platform_nearest_edge_gap_mean", 0.0)), }, frames, successes def replay_payload(frames_by_env: list[list[dict[str, Any]]], successes: list[bool]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: rollouts: list[dict[str, Any]] = [] maps: list[dict[str, Any]] = [] for i, frames in enumerate(frames_by_env): terminal = frames[-1] if frames else {} success = bool(successes[i]) or bool(terminal.get("success")) fallen = bool(terminal.get("fallen")) timeout = bool(terminal.get("timeout")) course, compact = compact_ego_frames(frames) metrics = { "success": success, "clear_rate": 1.0 if success else 0.0, "done": bool(terminal.get("done")) or success or fallen or timeout, "fallen": fallen, "burned": False, "timeout": timeout, "outcome": terminal.get("outcome", "success" if success else "running"), "maxStage": max((float(f.get("stage", 0.0)) for f in frames), default=0.0), "max_stage": max((float(f.get("stage", 0.0)) for f in frames), default=0.0), "frames": len(frames), "source": "fast_beacon_planner", } rollouts.append({"course": course, "frames": compact, "metrics": metrics, "policy": "fast_beacon_planner"}) maps.append({"index": i, "metrics": metrics, "rollouts": [i]}) return maps, rollouts STATIC_KEYS = { "activeNodes", "egoReplay", "fireMask", "goal", "launchPadMask", "movingMask", "routeOrder", "sizes", } def compact_ego_frames(frames: list[dict[str, Any]]) -> tuple[dict[str, Any], list[dict[str, Any]]]: if not frames: return {}, [] course = {key: frames[0][key] for key in STATIC_KEYS if key in frames[0]} compact = [{key: value for key, value in frame.items() if key not in STATIC_KEYS} for frame in frames] return course, compact def format_stats(name: str, stats: dict[str, float]) -> str: keys = ( "ppo_loss", "ppo_return", "ppo_success", "episode_success_rate", "ppo_progress", "sil_loss", "anchor_bc_loss", "sil_fraction", "sil_positive_fraction", "jump_fraction", "grounded_jump_fraction", "sprint_fraction", "forward_mean", "side_abs", "turn_abs", "control_std", "planner_blend", "planner_delta_x_m", "planner_delta_y_m", "planner_delta_z_m", "planner_delta_xy_abs_m", "env_sps", "collect_sps", "update_sps", "eval_success", "eval_progress", "map_goal_xy_mean", "map_route_path_xy_mean", "map_goal_xy_to_route_path_ratio_mean", "map_extra_after_goal_fraction", "map_shortest_hops_mean", "map_shortcut_ratio_mean", "map_shortcut_le_half_fraction", "stage_mastered", "phase_seconds", ) return name + " " + " ".join(f"{key}={stats.get(key, 0.0):.4f}" for key in keys) if __name__ == "__main__": main()