File size: 3,053 Bytes
6953bd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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)
# action agreement on a held-out-style check: mean |pred - teacher|
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()
|