Blame for train.py

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