train_cfm_delivcoal.py
CFM pointer/POMO 配送训练
原文件 · 378 行 · 15241 字节
experiments/reproduction/train_cfm_delivcoal.py
SHA-256: aec960321ec1956619cb83465977e2cd7ee81cd8ccf2a76c2edcc7b489800b51
1"""在 DelivCoal 环境中训练 HeteroMRTA 项目提供的 CFM 分配网络。23这里直接使用原项目的 AttentionNet，并按智能体顺序选择任务。场景、可见任务、4可选动作、环境推进、奖励和统计指标仍由 DelivCoal 提供。5"""67from __future__ import annotations89import argparse10import json11import os12import sys13import time14from datetime import datetime15from pathlib import Path1617import numpy as np18import torch19from torch.distributions import Categorical2021REPO_ROOT = Path(__file__).resolve().parents[2]22if str(REPO_ROOT) not in sys.path:23 sys.path.insert(0, str(REPO_ROOT))2425from experiments.reproduction.delivcoal_adapter_common import (26 DEFAULT_CONFIG,27 aggregate_rows,28 atomic_json,29 config_dict,30 load_env_config,31 seed_everything,32 summarize_episode,33)34from envs.delivcoal.env_discrete import DiscreteActionEnv353637CFM_SOURCE = REPO_ROOT / "third_party/heteromrta"38if str(CFM_SOURCE) not in sys.path:39 sys.path.insert(0, str(CFM_SOURCE))4041from attention import AttentionNet # noqa: E402424344def now():45 return datetime.now().isoformat(timespec="seconds")464748def model_input_dims(config):49 """根据 DelivCoal 当前的观测格式计算 CFM 的输入维度。"""50 agent_dim = 2 + int(config.cap_dim) + 551 if (52 bool(config.use_coalition_delivery)53 and getattr(config, "coalition_task_feature_mode", "compact") == "compact"54 ):55 task_dim = 656 else:57 task_dim = 2 + int(config.cap_dim) + 1 + 1 + 3 + 158 return agent_dim, task_dim596061def atomic_torch_save(path, payload):62 """先写临时文件，再替换正式 checkpoint，避免读到未写完的文件。"""63 path = Path(path)64 temporary = path.with_suffix(path.suffix + ".tmp")65 torch.save(payload, temporary)66 os.replace(temporary, path)676869def team_agent_features(env):70 """读取所有智能体的完整特征，供 CFM 依次分配任务。"""71 return env.get_global_state()["agents"].astype(np.float32, copy=True)727374def choose_actions(network, env, obs, active_mask, available_actions, device,75 greedy=False, record=False):76 """按 active_ids 顺序为智能体选择任务，并返回训练所需的输入快照。"""77 n_agents, n_actions = available_actions.shape78 actions = np.asarray(79 [agent.coalition for agent in env.env.agents], dtype=np.int6480 )81 agent_features = team_agent_features(env)82 transitions = []83 active_ids = np.flatnonzero(np.asarray(active_mask).reshape(-1) > 0.5)8485 # CFM 在同一个决策时刻逐个处理智能体。86 for agent_id in active_ids:87 available = np.asarray(available_actions[agent_id], dtype=bool)88 if not available.any():89 continue90 # AttentionNet 会使用任务序列的第 0 行汇总信息。先把可选任务移到序列前面，91 # 网络选出序列位置后，再通过 task_order 转回 DelivCoal 的任务编号。92 task_order = np.concatenate(93 [np.flatnonzero(available), np.flatnonzero(~available)]94 ).astype(np.int64)95 ordered_available = available[task_order]96 tasks = torch.as_tensor(97 obs["tasks"][agent_id, task_order], dtype=torch.float32, device=device98 ).unsqueeze(0)99 agents = torch.as_tensor(100 agent_features, dtype=torch.float32, device=device101 ).unsqueeze(0)102 mask = torch.as_tensor(103 ~ordered_available, dtype=torch.bool, device=device104 ).reshape(1, 1, n_actions)105 index = torch.tensor([[[agent_id]]], dtype=torch.long, device=device)106 with torch.no_grad():107 # greedy 用于确定性评估；训练时从概率分布中采样。108 probabilities, _ = network(tasks, agents, mask, index)109 distribution = Categorical(probs=probabilities)110 action = (111 torch.argmax(probabilities, dim=-1)112 if greedy113 else distribution.sample()114 )115 sequence_action = int(action.item())116 action_id = int(task_order[sequence_action])117 actions[agent_id] = action_id118 if record:119 # 保存动作发生时的输入。后续智能体会看到更新后的 agent_features。120 transitions.append(121 (122 obs["tasks"][agent_id, task_order].copy(),123 agent_features.copy(),124 (~ordered_available).copy(),125 int(agent_id),126 sequence_action,127 )128 )129 # 第 8 项保存联盟编号。立刻写回本次选择，让后续智能体看到前面的分配结果。130 agent_features[agent_id, 8] = action_id / max(1, n_actions - 1)131 return actions, transitions132133134def rollout(network, config, seed, device, greedy, record=False,135 policy_seed=None):136 """运行一个完整 episode，并按需保存其中每次任务选择。"""137 seed_everything(seed)138 env = DiscreteActionEnv(config)139 env.seed(seed)140 obs, active, available = env.reset(0)141 if policy_seed is not None:142 # 环境已经用 seed 固定场景。这里单独设置采样种子，让同一场景的多次143 # rollout 可以抽到不同动作。144 seed_everything(policy_seed)145 total_reward = 0.0146 sent_total = 0.0147 transitions = []148 info = {}149 length = 0150 for step in range(int(config.max_steps)):151 actions, current = choose_actions(152 network, env, obs, active, available, device, greedy, record153 )154 obs, rewards, dones, info, active, _, available = env.step(actions)155 total_reward += float(np.mean(rewards))156 sent_total += float(info.get("sent", 0.0))157 transitions.extend(current)158 length = step + 1159 if bool(np.all(dones)):160 break161 row = summarize_episode(info, total_reward, length, sent_total)162 row["seed"] = int(seed)163 return row, transitions164165166def update_policy(network, optimizer, samples, device, batch_size):167 """用每次 rollout 相对组内平均回报的差值更新 CFM。"""168 optimizer.zero_grad(set_to_none=True)169 rollout_terms = []170 entropy_terms = []171 decision_count = 0172 for transitions, advantage in samples:173 # advantage 是该 rollout 回报减去同组 rollout 的平均回报。174 if not transitions:175 continue176 log_prob_parts = []177 entropy_parts = []178 for start in range(0, len(transitions), batch_size):179 # transition 中保存的是当时实际输入网络的数据，重新前向计算所选动作概率。180 batch = transitions[start : start + batch_size]181 tasks = torch.as_tensor(182 np.stack([item[0] for item in batch]), device=device183 )184 agents = torch.as_tensor(185 np.stack([item[1] for item in batch]), device=device186 )187 masks = torch.as_tensor(188 np.stack([item[2] for item in batch]), device=device189 ).unsqueeze(1)190 indices = torch.as_tensor(191 [item[3] for item in batch], dtype=torch.long, device=device192 ).reshape(-1, 1, 1)193 actions = torch.as_tensor(194 [item[4] for item in batch], dtype=torch.long, device=device195 )196 probabilities, _ = network(tasks, agents, masks, indices)197 distribution = Categorical(probs=probabilities)198 log_prob_parts.append(distribution.log_prob(actions))199 entropy_parts.append(distribution.entropy())200 mean_log_prob = torch.cat(log_prob_parts).mean()201 # 提高高于平均回报动作的概率，降低低于平均回报动作的概率。202 rollout_terms.append(-float(advantage) * mean_log_prob)203 entropy_terms.append(torch.cat(entropy_parts).mean())204 decision_count += len(transitions)205 if not rollout_terms:206 return {"loss": 0.0, "entropy": 0.0, "grad_norm": 0.0,207 "decisions": 0}208 policy_loss = torch.stack(rollout_terms).mean()209 entropy = torch.stack(entropy_terms).mean()210 loss = policy_loss211 loss.backward()212 # 限制梯度大小，避免单组回报差值过大导致更新突变。213 grad_norm = torch.nn.utils.clip_grad_norm_(network.parameters(), 10.0)214 optimizer.step()215 return {216 "loss": float(loss.detach().cpu()),217 "entropy": float(entropy.detach().cpu()),218 "grad_norm": float(torch.as_tensor(grad_norm).detach().cpu()),219 "decisions": int(decision_count),220 }221222223@torch.no_grad()224def evaluate(network, config, device, seeds):225 training = network.training226 network.eval()227 rows = [rollout(network, config, seed, device, True)[0] for seed in seeds]228 if training:229 network.train()230 return {"episodes": rows, "summary": aggregate_rows(rows)}231232233def checkpoint_payload(network, optimizer, scheduler, episode, args, config):234 agent_input_dim, task_input_dim = model_input_dims(config)235 return {236 "method": "cfm-delivcoal",237 "episode": int(episode),238 "model": network.state_dict(),239 "optimizer": optimizer.state_dict(),240 "scheduler": scheduler.state_dict(),241 "agent_input_dim": agent_input_dim,242 "task_input_dim": task_input_dim,243 "embedding_dim": int(args.embedding_dim),244 "args": vars(args),245 "environment": config_dict(config),246 }247248249def train(args):250 output = Path(args.output).resolve()251 models = output / "models"252 trainers = output / "trainers"253 output.mkdir(parents=True, exist_ok=True)254 models.mkdir(exist_ok=True)255 trainers.mkdir(exist_ok=True)256 config = load_env_config(args.env_config)257 seed_everything(args.seed)258 device = torch.device(259 f"cuda:{args.cuda}" if args.cuda >= 0 and torch.cuda.is_available() else "cpu"260 )261 agent_input_dim, task_input_dim = model_input_dims(config)262 network = AttentionNet(263 agent_input_dim, task_input_dim, args.embedding_dim264 ).to(device)265 optimizer = torch.optim.Adam(network.parameters(), lr=args.lr)266 scheduler = torch.optim.lr_scheduler.StepLR(267 optimizer, step_size=args.lr_decay_interval, gamma=args.lr_decay268 )269 status = {270 "name": f"cfm-delivcoal-s{args.seed}",271 "state": "training",272 "created_at": now(),273 "method": "CFM",274 "environment": "DelivCoal",275 "environment_config": config_dict(config),276 "seed": args.seed,277 "completed_episodes": 0,278 "total_episodes": args.episodes,279 }280 atomic_json(output / "status.json", status)281 atomic_json(output / "reproduction_config.json", {282 "method": "official HeteroMRTA AttentionNet / CFM",283 "adapter": "DelivCoal sequential allocation adapter",284 "args": vars(args),285 "environment": config_dict(config),286 })287 metrics_path = output / "metrics.jsonl"288 started = time.monotonic()289 for episode in range(1, args.episodes + 1):290 network.train()291 rows = []292 samples = []293 scenario_seed = args.seed * 1_000_000 + episode294 for pomo_id in range(args.pomo_size):295 row, transitions = rollout(296 network, config, scenario_seed, device, False, record=True,297 policy_seed=scenario_seed + 10_000 * (pomo_id + 1),298 )299 rows.append(row)300 samples.append((transitions, 0.0))301 rewards = np.asarray([row["episode_reward"] for row in rows])302 baseline = float(rewards.mean())303 samples = [304 (transitions, float(reward - baseline))305 for (transitions, _), reward in zip(samples, rewards)306 ]307 update = update_policy(network, optimizer, samples, device, args.batch_size)308 scheduler.step()309 payload = checkpoint_payload(310 network, optimizer, scheduler, episode, args, config311 )312 if episode % args.actor_save_interval == 0 or episode == args.episodes:313 atomic_torch_save(models / f"actor_ep{episode:06d}.pt", {314 key: payload[key] for key in (315 "method", "episode", "model", "agent_input_dim",316 "task_input_dim", "embedding_dim", "environment"317 )318 })319 if episode % args.trainer_save_interval == 0 or episode == args.episodes:320 atomic_torch_save(trainers / f"trainer_ep{episode:06d}.pt", payload)321 evaluation = None322 if episode % args.eval_interval == 0 or episode == args.episodes:323 seeds = range(args.validation_seed,324 args.validation_seed + args.validation_episodes)325 evaluation = evaluate(network, config, device, seeds)326 record = {327 "episode": episode,328 "environment_steps": int(episode * args.pomo_size * config.max_steps),329 "train": aggregate_rows(rows),330 "update": update,331 "validation": evaluation,332 "recorded_at": now(),333 }334 with metrics_path.open("a", encoding="utf-8") as file:335 file.write(json.dumps(record, ensure_ascii=False) + "\n")336 status["completed_episodes"] = episode337 status["updated_at"] = now()338 atomic_json(output / "status.json", status)339 elapsed = time.monotonic() - started340 fps = episode * args.pomo_size * config.max_steps / max(elapsed, 1e-6)341 eta = elapsed / episode * (args.episodes - episode)342 print(343 f"rollout {episode}/{args.episodes}, timesteps "344 f"{episode * args.pomo_size * config.max_steps}, FPS {fps:.1f}, "345 f"time cost {elapsed:.1f}s, expected time {eta:.1f}s, "346 f"reward {rewards.mean():.4f}, success "347 f"{np.mean([row['success'] for row in rows]):.3f}",348 flush=True,349 )350 status.update(state="completed", finished_at=now(), actor_checkpoint=str(351 models / f"actor_ep{args.episodes:06d}.pt"352 ))353 atomic_json(output / "status.json", status)354355356def parse_args():357 parser = argparse.ArgumentParser()358 parser.add_argument("--output", required=True)359 parser.add_argument("--env-config", default=str(DEFAULT_CONFIG))360 parser.add_argument("--seed", type=int, default=10)361 parser.add_argument("--episodes", type=int, default=1500)362 parser.add_argument("--pomo-size", type=int, default=4)363 parser.add_argument("--embedding-dim", type=int, default=128)364 parser.add_argument("--lr", type=float, default=1e-5)365 parser.add_argument("--lr-decay-interval", type=int, default=32)366 parser.add_argument("--lr-decay", type=float, default=0.98)367 parser.add_argument("--batch-size", type=int, default=128)368 parser.add_argument("--actor-save-interval", type=int, default=1)369 parser.add_argument("--trainer-save-interval", type=int, default=100)370 parser.add_argument("--eval-interval", type=int, default=50)371 parser.add_argument("--validation-episodes", type=int, default=32)372 parser.add_argument("--validation-seed", type=int, default=10_000)373 parser.add_argument("--cuda", type=int, default=0)374 return parser.parse_args()375376377if __name__ == "__main__":378 train(parse_args())