| """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(): |
| |
| |
| |
| |
| 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(): |
| |
| 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) |
| |
| 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 |
| |
| |
| obs = obs2 |
| total_env_steps += args.n_envs |
| |
| 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 |
| |
| done_buf[-1] = just_done | np.broadcast_to(ep_done[:, None], just_done.shape) |
| |
| 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) |
| act_arr = np.stack(act_buf) |
| logp_arr = np.stack(logp_buf) |
| rew_arr = np.stack(rew_buf).astype(np.float32) |
| val_arr = np.stack(val_buf) |
| done_arr = np.stack(done_buf) |
| |
| 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 |
| |
| 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() |
|
|
