| """Behaviour-cloning warm-start from the classical controller. |
| |
| Collects (observation, classical action) pairs by running the classical |
| PD + potential-field policy on GENERAL layouts only (trap=False, where |
| the teacher succeeds 10/10 in sanity checks). Trap layouts are excluded |
| deliberately: imitating the teacher there would teach the stall. |
| Trains only the actor (mean head) of ActorCritic with MSE on the |
| tanh-bounded action. The critic is left for PPO fine-tuning. |
| Outputs: bc_policy.pt, bc.log via stdout. |
| """ |
| import argparse |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| from sim import AMAX, BatchedSim, generate_layout |
| from policies import ActorCritic, classical_action |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--episodes", type=int, default=400) |
| ap.add_argument("--n-agents", type=int, default=4) |
| ap.add_argument("--epochs", type=int, default=30) |
| ap.add_argument("--seed", type=int, default=123) |
| ap.add_argument("--out", default="bc_policy.pt") |
| args = ap.parse_args() |
| torch.manual_seed(args.seed) |
| torch.set_num_threads(2) |
| rng = np.random.default_rng(args.seed) |
| sim = BatchedSim(16, args.n_agents, seed=args.seed) |
| obs_list, act_list = [], [] |
| layouts = [generate_layout(rng, args.n_agents, trap=False) for _ in range(16)] |
| obs = sim.reset(layouts) |
| done_eps = 0 |
| while done_eps < args.episodes: |
| act = classical_action(obs) |
| obs_list.append(obs.reshape(-1, obs.shape[-1]).copy()) |
| act_list.append(act.reshape(-1, 2).copy()) |
| obs, _, ep_done, _ = sim.step(act) |
| fin = np.where(ep_done)[0] |
| if len(fin): |
| done_eps += len(fin) |
| obs = sim.reset_envs(fin.tolist(), |
| [generate_layout(rng, args.n_agents, trap=False) for _ in fin]) |
| O = torch.as_tensor(np.concatenate(obs_list, axis=0)) |
| A = torch.as_tensor(np.concatenate(act_list, axis=0)) |
| print(f"BC dataset: {O.shape[0]} samples, {done_eps} episodes", flush=True) |
| net = ActorCritic() |
| opt = torch.optim.Adam(list(net.actor.parameters()) + list(net.mean.parameters()), lr=1e-3) |
| n = O.shape[0] |
| for ep in range(args.epochs): |
| perm = torch.randperm(n) |
| tot = 0.0 |
| for s in range(0, n, 1024): |
| b = perm[s:s + 1024] |
| mean, _ = net.forward(O[b]) |
| pred = torch.tanh(mean) * AMAX |
| loss = nn.functional.mse_loss(pred, A[b]) |
| opt.zero_grad(); loss.backward(); opt.step() |
| tot += float(loss) * len(b) |
| print(f"bc epoch {ep + 1}/{args.epochs} mse={tot / n:.4f}", flush=True) |
| |
| with torch.no_grad(): |
| mean, _ = net.forward(O[:4096]) |
| pred = torch.tanh(mean) * AMAX |
| print(f"BC final: mean |action error| = {float((pred - A[:4096]).abs().mean()):.3f} " |
| f"(accel units)", flush=True) |
| torch.save(net.state_dict(), args.out) |
| print(f"saved {args.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
