| 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 |
|
|
|
|
| |
| |
| _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() |
|
|
