"""PPO training (torch only; BC warm-start via --init, see BUILD_REPORT.md). Training: 16 parallel envs, random layouts regenerated on every episode reset (trap probability ramped by --trap-prob/--trap-ramp-iters), fixed seed 0. Rollout = 128 simulator steps per iteration. Default 120 iterations = 245,760 environment steps (983,040 agent-steps). PPO: gamma 0.99, GAE lambda 0.95, clip 0.2, lr 3e-4, 4 epochs, minibatch 2048, entropy coef 0.005 (default; overridable), value coef 0.5, max-grad-norm 0.5. Reward is the shaping reward returned by BatchedSim.step (progress + success/collision bonuses + step/proximity penalties). Outputs: ppo_policy.pt, training_curve.csv (per-iteration mean return, success rate over finished agent-episodes, collision rate, entropy, value loss, policy loss), train.log via stdout redirection by caller. """ import argparse import csv import json import time import numpy as np import torch import torch.nn as nn from sim import AMAX, OBS_DIM, BatchedSim, generate_layout from policies import ActorCritic def main(): ap = argparse.ArgumentParser() ap.add_argument("--iterations", type=int, default=120) ap.add_argument("--rollout", type=int, default=128) ap.add_argument("--n-envs", type=int, default=16) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--out", default="ppo_policy.pt") ap.add_argument("--curve", default="training_curve.csv") ap.add_argument("--n-agents", type=int, default=4) ap.add_argument("--init", default="") ap.add_argument("--ent-coef", type=float, default=0.005) ap.add_argument("--trap-prob", type=float, default=0.5) ap.add_argument("--trap-ramp-iters", type=int, default=150) args = ap.parse_args() torch.manual_seed(args.seed) np.random.seed(args.seed) torch.set_num_threads(2) device = "cpu" net = ActorCritic().to(device) if args.init: net.load_state_dict(torch.load(args.init, map_location=device, weights_only=True)) print(f"initialised from {args.init}", flush=True) opt = torch.optim.Adam(net.parameters(), lr=3e-4) layout_rng = np.random.default_rng(args.seed) def sample_layout(): # Curriculum: trap probability ramps 0 -> trap_prob over the # first trap_ramp_iters iterations, so the policy first learns # steering around convex obstacles (short detours) before the # wall-trap detour (long, progress-negative detour) appears. p = args.trap_prob * min(1.0, (sample_layout.it - 1) / max(args.trap_ramp_iters, 1)) return generate_layout(layout_rng, args.n_agents, trap=bool(layout_rng.random() < p)) sample_layout.it = 1 sim = BatchedSim(args.n_envs, args.n_agents, seed=args.seed) layouts = [sample_layout() for _ in range(args.n_envs)] obs = sim.reset(layouts) gamma, lam, clip, ent_coef, vf_coef = 0.99, 0.95, 0.2, args.ent_coef, 0.5 rows = [] total_env_steps = 0 t0 = time.time() for it in range(1, args.iterations + 1): sample_layout.it = it obs_buf, act_buf, logp_buf, rew_buf, val_buf, done_buf = [], [], [], [], [], [] ep_ret, ep_succ, ep_coll, ep_n = [], [], [], 0 cur_ret = np.zeros((args.n_envs, args.n_agents)) cur_succ = np.zeros((args.n_envs, args.n_agents), dtype=bool) cur_coll = np.zeros((args.n_envs, args.n_agents), dtype=bool) prev_done = np.zeros((args.n_envs, args.n_agents), dtype=bool) for t in range(args.rollout): o = torch.as_tensor(obs) with torch.no_grad(): # pre-tanh Gaussian sample (PPO logprob is on this sample) mean, value = net.forward(o) std = net.log_std.exp().expand_as(mean) dist = torch.distributions.Normal(mean, std) raw_s = dist.sample() logp = dist.log_prob(raw_s).sum(-1) action = torch.tanh(raw_s) * AMAX act_np = action.numpy() obs_buf.append(obs.copy()); act_buf.append(raw_s.numpy()) logp_buf.append(logp.numpy()); val_buf.append(value.numpy()) obs2, rew, ep_done, _ = sim.step(act_np) rew_buf.append(rew) # per-agent episode end: done_agent flips False->True once just_done = sim.done_agent.copy() done_buf.append(np.zeros_like(rew, dtype=bool)) cur_ret += rew cur_succ |= sim.success cur_coll |= sim.collision # episode-level done for GAE: per-agent terminal when its # own episode ends (done_agent) or the env timed out obs = obs2 total_env_steps += args.n_envs # flush finished agents: done_agent transitions False->True newly = just_done & ~prev_done if newly.any(): ep_ret.extend(cur_ret[newly].tolist()) ep_succ.extend(cur_succ[newly].tolist()) ep_coll.extend(cur_coll[newly].tolist()) ep_n += int(newly.sum()) cur_ret[newly] = 0.0 cur_succ[newly] = False cur_coll[newly] = False prev_done = just_done # terminal mask for this step: agent done or env done done_buf[-1] = just_done | np.broadcast_to(ep_done[:, None], just_done.shape) # partial resets for finished envs fin = np.where(ep_done)[0] if len(fin): new_layouts = [sample_layout() for _ in fin] obs = sim.reset_envs(fin.tolist(), new_layouts) prev_done[fin] = False cur_ret[fin] = 0.0 cur_succ[fin] = False cur_coll[fin] = False obs_arr = np.stack(obs_buf) # (T,E,A,15) act_arr = np.stack(act_buf) # (T,E,A,2) pre-tanh samples logp_arr = np.stack(logp_buf) # (T,E,A) rew_arr = np.stack(rew_buf).astype(np.float32) val_arr = np.stack(val_buf) # (T,E,A) done_arr = np.stack(done_buf) # (T,E,A) # GAE with torch.no_grad(): _, last_val = net.forward(torch.as_tensor(obs)) last_val = last_val.numpy() adv = np.zeros_like(rew_arr) gae = np.zeros(rew_arr.shape[1:]) for t in reversed(range(args.rollout)): nxt_val = last_val if t == args.rollout - 1 else val_arr[t + 1] nxt_nonterm = 1.0 - done_arr[t].astype(np.float32) delta = rew_arr[t] + gamma * nxt_val * nxt_nonterm - val_arr[t] gae = delta + gamma * lam * nxt_nonterm * gae adv[t] = gae ret_arr = adv + val_arr # flatten (T*E*A, ...) O = torch.as_tensor(obs_arr.reshape(-1, OBS_DIM)) A_ = torch.as_tensor(act_arr.reshape(-1, 2)) LP = torch.as_tensor(logp_arr.reshape(-1)) ADV = torch.as_tensor(adv.reshape(-1)) RET = torch.as_tensor(ret_arr.reshape(-1)) ADV = (ADV - ADV.mean()) / (ADV.std() + 1e-8) n = O.shape[0] idx = np.arange(n) ploss_sum = vloss_sum = ent_sum = 0.0 nb = 0 for _epoch in range(4): np.random.shuffle(idx) for s in range(0, n, 2048): b = idx[s:s + 2048] logp_new, ent, val_new = net.evaluate(O[b], A_[b]) ratio = torch.exp(logp_new - LP[b]) s1 = ratio * ADV[b] s2 = torch.clamp(ratio, 1 - clip, 1 + clip) * ADV[b] ploss = -torch.min(s1, s2).mean() vloss = ((val_new - RET[b]) ** 2).mean() loss = ploss + vf_coef * vloss - ent_coef * ent.mean() opt.zero_grad() loss.backward() nn.utils.clip_grad_norm_(net.parameters(), 0.5) opt.step() ploss_sum += float(ploss.detach()); vloss_sum += float(vloss.detach()); ent_sum += float(ent.detach().mean()) nb += 1 row = { "iteration": it, "env_steps_total": total_env_steps, "agent_steps_total": total_env_steps * args.n_agents, "mean_episode_return": float(np.mean(ep_ret)) if ep_ret else float("nan"), "episodes_finished": ep_n, "success_rate_finished": float(np.mean(ep_succ)) if ep_succ else float("nan"), "collision_rate_finished": float(np.mean(ep_coll)) if ep_coll else float("nan"), "policy_loss": ploss_sum / max(nb, 1), "value_loss": vloss_sum / max(nb, 1), "entropy": ent_sum / max(nb, 1), "elapsed_s": round(time.time() - t0, 1), } rows.append(row) print(f"iter {it}/{args.iterations} env_steps={total_env_steps} " f"episodes={ep_n} mean_return={row['mean_episode_return']:.2f} " f"success={row['success_rate_finished']:.3f} coll={row['collision_rate_finished']:.3f} " f"ploss={row['policy_loss']:.3f} ent={row['entropy']:.3f} " f"elapsed={row['elapsed_s']:.0f}s", flush=True) if it % 20 == 0 or it == args.iterations: torch.save(net.state_dict(), args.out) with open(args.curve, "w", newline="") as f: w = csv.DictWriter(f, fieldnames=list(rows[0].keys())) w.writeheader(); w.writerows(rows) torch.save(net.state_dict(), args.out) with open(args.curve, "w", newline="") as f: w = csv.DictWriter(f, fieldnames=list(rows[0].keys())) w.writeheader(); w.writerows(rows) print(f"DONE total_env_steps={total_env_steps} agent_steps={total_env_steps * args.n_agents} " f"elapsed={time.time() - t0:.1f}s weights={args.out}") if __name__ == "__main__": main()