train.py
训练入口、环境与 runner 选择
原文件 · 551 行 · 21141 字节
train/train.py
SHA-256: 8e4e909d00f0866db14e6760596f56c1ff4b0876d10b8ce416b1e5953f79995c
1"""2# @Time : 2021/6/30 10:07 下午3# @Author : hezhiqiang4# @Email : tinyzqh@163.com5# @File : train.py6"""78# !/usr/bin/env python9import sys10import os11import re12import socket13import setproctitle14import numpy as np15from pathlib import Path16import torch17import yaml18from datetime import datetime1920# Get the parent directory of the current file21parent_dir = os.path.abspath(os.path.join(os.getcwd(), "."))2223# Append the parent directory to sys.path, otherwise the following import will fail24sys.path.append(parent_dir)2526from config import get_args, get_config27from envs.macf.env_wrappers import DummyVecEnv2829"""Train script for MPEs."""303132REPO_ROOT = Path(__file__).resolve().parents[1]333435def _safe_name(value):36 safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value)).strip("._")37 if not safe:38 raise ValueError(f"Invalid path name: {value!r}")39 return safe404142def _resolve_repo_path(value):43 path = Path(value).expanduser()44 return path if path.is_absolute() else REPO_ROOT / path454647def _unique_path(path):48 if not path.exists():49 return path50 for index in range(2, 1000):51 candidate = path.with_name(f"{path.name}-{index}")52 if not candidate.exists():53 return candidate54 raise RuntimeError(f"Could not allocate a unique run directory for {path}")555657def _link_tensorboard_dir(tb_dir, link_dir):58 link_dir.parent.mkdir(parents=True, exist_ok=True)59 if link_dir.is_symlink():60 if link_dir.resolve() == tb_dir.resolve():61 return62 raise FileExistsError(f"TensorBoard link already exists: {link_dir}")63 if link_dir.exists():64 raise FileExistsError(f"TensorBoard link path already exists: {link_dir}")6566 relative_target = os.path.relpath(tb_dir, start=link_dir.parent)67 link_dir.symlink_to(relative_target, target_is_directory=True)686970def _prepare_run_dirs(config):71 configured_run_dir = getattr(config, "run_dir", None)72 if configured_run_dir:73 run_dir = _resolve_repo_path(configured_run_dir)74 else:75 results_root = _resolve_repo_path(getattr(config, "results_root", "runs"))76 group = _safe_name(77 getattr(config, "group_name", None)78 or getattr(config, "experiment_group", "single")79 )80 env_name = _safe_name(config.env_name)81 timestamp = datetime.now().strftime("%Y%m%d-%H%M")82 group_dir = _unique_path(results_root / f"{timestamp}_{env_name}_{group}")83 run_name = f"{_safe_name(config.algorithm_name)}_{_safe_name(config.experiment_name)}"84 run_dir = group_dir / run_name8586 tb_dir = run_dir / "tb"87 tb_dir.mkdir(parents=True, exist_ok=False)8889 log_root = _resolve_repo_path(getattr(config, "log_root", "logs"))90 tb_link_dir = log_root / run_dir.parent.name / run_dir.name91 _link_tensorboard_dir(tb_dir, tb_link_dir)9293 return run_dir, tb_dir, tb_link_dir949596def _write_config_snapshot(config, run_dir):97 snapshot = {}98 for key, value in vars(config).items():99 if key in {"envs", "eval_envs", "device"}:100 continue101 snapshot[key] = str(value) if isinstance(value, Path) else value102 with (run_dir / "config.yaml").open("w", encoding="utf-8") as file:103 yaml.safe_dump(snapshot, file, sort_keys=False, allow_unicode=True)104105106def _discrete_env_class(config):107 """按 env_name 选择具体环境包。108109 taskcf 是原任务能力联盟组建环境；coadel 是 coalition delivery 环境。110 macf/MyEnv/coalition_delivery 作为历史配置别名保留。111 """112113 env_name = str(getattr(config, "env_name", "taskcf"))114 if env_name in ("capcoal", "macf", "MyEnv", "taskcf"):115 from envs.capcoal.env_discrete import DiscreteActionEnv116 return DiscreteActionEnv117 if env_name in ("delivcoal", "coalition_delivery", "coadel"):118 from envs.delivcoal.env_discrete import DiscreteActionEnv119 return DiscreteActionEnv120 if env_name in ("gymrt2a", "mrt2a"):121 from envs.gymrt2a.env_discrete import DiscreteActionEnv122 return DiscreteActionEnv123 raise ValueError(124 f"Unknown env_name={env_name!r}; expected capcoal, delivcoal, or gymrt2a."125 )126127128def make_train_env(config):129 def get_env_fn(rank):130 def init_env():131 DiscreteActionEnv = _discrete_env_class(config)132 env = DiscreteActionEnv(config)133134 env.seed(config.seed + rank * 1000)135 return env136137 return init_env138139 return DummyVecEnv([get_env_fn(i) for i in range(config.n_rollout_threads)])140141142def make_eval_env(config):143 def get_env_fn(rank):144 def init_env():145 DiscreteActionEnv = _discrete_env_class(config)146 env = DiscreteActionEnv(config)147 eval_reset_mode = getattr(config, "eval_reset_mode", None)148 if eval_reset_mode is not None:149 env.set_reset_mode(eval_reset_mode)150 env.seed(config.eval_seed + rank * 1000)151 return env152153 return init_env154155 return DummyVecEnv([get_env_fn(i) for i in range(config.n_eval_rollout_threads)])156157158def parse_args(args, parser):159 parser.add_argument("--scenario_name", type=str, default="MyEnv", help="Which scenario to run on")160 parser.add_argument("--num_landmarks", type=int, default=3)161162 add_my_args(parser)163164 config = parser.parse_known_args(args)[0]165166 return config167168169def add_my_args(parser):170171 parser.add_argument("--config_file", type=str, default=None, help="path of config file (*.yaml)")172173 # GAT174 parser.add_argument("--attr_dim", type=int, default=4, help="")175 parser.add_argument("--gat_hidden_dim", type=int, default=64, help="")176 parser.add_argument("--gat_out_dim", type=int, default=32, help="")177 parser.add_argument("--gat_mode", type=str, default="gated", choices=["gat", "cat", "gated"], help="GAT fusion mode")178 parser.add_argument(179 "--graph_model",180 type=str,181 default="legacy_gat",182 choices=[183 "legacy_gat",184 "hatt",185 "hatt_v2_0",186 "hatt_v2_0_raw",187 "hatt_v2_0_physics",188 "hatt_v2_0_taskattn",189 "hatt_v2_0_noidle",190 "hatt_v2_0_resource",191 "hatt_v2_0_setctx",192 "hatt_v2_0_plain_noidle",193 "hatt_v2_0_detached",194 "hatt_v2_0_split",195 "hatt_v2_0_type_embed",196 "hatt_v2_0_type_gate",197 "hatt_v2_1",198 "hatt_v2",199 "hatt_v2_3a",200 "hatt_v2_3b",201 "hatt_v2_3c",202 "hatt_v2_4",203 "hatt_v2_5a",204 "hatt_v2_5b",205 "hatt_v2_5c",206 "hatt_v2_5",207 "hatt_v2_6a",208 "hatt_v2_6b",209 "hatt_v2_6c",210 "hatt_v2_6",211 "hatt_v2_7a",212 "hatt_v2_7b",213 "hatt_v2_7",214 "hatt_v3_0a",215 "hatt_v3_0b",216 "hatt_v3_0c",217 "hatt_v3_1a",218 "hatt_v3_1b",219 "hatt_v3_2a",220 "hatt_v3_2b",221 "hatt_v3_2c",222 "hatt_v3_2d",223 "hatt_v3_2r",224 "hatt_v3_2r_no_transition",225 "hatt_v3_2r_no_member",226 "hatt_v3_3aa",227 "hatt_v3_4",228 "hatt_gymrt2a",229 ],230 )231 parser.add_argument("--hatt_hidden_dim", type=int, default=128)232 parser.add_argument("--hatt_num_heads", type=int, default=4)233 parser.add_argument("--hatt_dropout", type=float, default=0.0)234 parser.add_argument(235 "--hatt_capability_scale",236 type=float,237 default=5.0,238 help="scale raw agent capabilities into the normalized task requirement space for HATT residuals",239 )240 parser.add_argument(241 "--hatt_v20_task_input_mode",242 choices=("native", "restore_req_scale", "compact_projection"),243 default="native",244 help="controlled HATT v2.0 task-schema intervention used for B/C attribution",245 )246 parser.add_argument(247 "--hatt_paired_init_seed",248 type=int,249 default=-1,250 help="local actor/critic initialization seed; -1 keeps historical RNG behavior",251 )252 parser.add_argument(253 "--hatt_v20_init_scheme",254 choices=("default", "compact_function_preserving"),255 default="default",256 help="optional exact compact6-to-full12 actor initialization for schema attribution",257 )258 parser.add_argument(259 "--hatt_shared_task_obs",260 type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),261 default=True,262 help="whether obs['tasks'] is shared across all agents and can be encoded once",263 )264 parser.add_argument(265 "--hatt_use_aa_attention",266 type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),267 default=True,268 help="enable HATT v2.0 agent-agent relation attention",269 )270 parser.add_argument(271 "--hatt_use_coalition_context",272 type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),273 default=True,274 help="enable HATT v2.0 task-indexed coalition context",275 )276 parser.add_argument(277 "--hatt_candidate_aa_mode",278 type=str,279 default="learned",280 choices=["learned", "uniform", "zero"],281 help=(282 "HATT v3.3 A-A gate mode: learned for the method, uniform for the "283 "matched training control, zero for checkpoint intervention"284 ),285 )286 parser.add_argument(287 "--hatt_aa_task_context",288 type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),289 default=True,290 help=(291 "HATT v3.4: feed the serving task's residual demand/priority into "292 "the A-A edge gate; false reproduces the v3.3AA 3-dim edge"293 ),294 )295 parser.add_argument(296 "--reconfiguration_mode",297 choices=["legacy", "always", "periodic", "event", "urgency", "learned"],298 default="legacy",299 help="coalition reconsideration schedule; learned enables the HATT v2.0 PPO gate",300 )301 parser.add_argument(302 "--allow_reconfiguration_mode_mismatch",303 type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),304 default=False,305 help=(306 "allow an explicit evaluation-only intervention between non-learned "307 "reconfiguration schedules; learned-gate mismatches remain forbidden"308 ),309 )310 parser.add_argument("--reconfiguration_interval", type=int, default=10)311 parser.add_argument("--reconfiguration_urgency_threshold", type=float, default=0.7)312 parser.add_argument(313 "--reconfiguration_urgency_aggregation",314 choices=["mean", "sum"],315 default="mean",316 )317 parser.add_argument("--reconfiguration_switch_cost", type=float, default=0.0)318 parser.add_argument("--reconfiguration_gate_init", type=float, default=0.8)319 parser.add_argument("--reconfiguration_gate_floor", type=float, default=0.02)320 parser.add_argument(321 "--reconfiguration_gate_parameterization",322 choices=["mixture", "switch", "hierarchical"],323 default="mixture",324 help=(325 "mixture preserves the historical latent reconsideration gate; "326 "switch makes the gate the identifiable probability of leaving "327 "the current task"328 ),329 )330 parser.add_argument(331 "--reconfiguration_gate_override",332 choices=["none", "open", "closed", "event", "event_or_learned"],333 default="none",334 help=(335 "gate intervention; event_or_learned is a safe hybrid that forces "336 "event decisions open and learns only the residual periodic switches"337 ),338 )339 parser.add_argument(340 "--reconfiguration_decision_coef",341 type=float,342 default=0.0,343 help="regularization cost for learned (non-mandatory) hierarchical redecisions",344 )345 parser.add_argument(346 "--reconfiguration_context_features",347 type=lambda value: str(value).lower() in ("1", "true", "yes", "y"),348 default=False,349 help="expose explicit event/time context to the learned reconfiguration gate",350 )351 parser.add_argument("--reconfiguration_context_dim", type=int, default=9)352353 # Act-Crs354 parser.add_argument("--obs_shape", type=int, default=32, help="")355 parser.add_argument("--share_obs_shape", type=int, default=32, help="")356357 # Enviroment358 parser.add_argument("--map_size", type=float, default=10000, help="")359 parser.add_argument("--num_uav", type=int, default=8, help="")360 parser.add_argument("--num_ugv", type=int, default=4, help="")361 parser.add_argument("--num_agents", type=int, default=12, help="")362363 parser.add_argument("--uav_perception_rg", type=list, default=[50, 200], help="")364 parser.add_argument("--ugv_perception_rg", type=list, default=[10, 30], help="")365366 parser.add_argument("--target_num", type=int, default=20, help="任务数量")367 parser.add_argument("--gat_dim", type=int, default=4, help="输出维度")368 parser.add_argument("--env_reset_mode", type=str, default="jitter",369 choices=["jitter", "full_random"],370 help="'jitter': base env + task jitter; 'full_random': full reinit each episode")371 parser.add_argument("--eval_reset_mode", type=str, default=None,372 choices=["jitter", "full_random"],373 help="optional reset mode used only by evaluation envs")374 parser.add_argument("--curriculum_switch_ep", type=int, default=-1,375 help="switch env_reset_mode to full_random after this env_episode; -1 = no switch")376 parser.add_argument("--max_steps", type=int, default=1000, help="最大步数")377 parser.add_argument(378 "--episode_time_limit",379 type=int,380 default=None,381 help="optional termination-only horizon; max_steps keeps time/deadline semantics",382 )383 parser.add_argument("--cap_dim", type=int, default=4, help="能力维度")384385 # Runner386 parser.add_argument("--use_gat", action="store_true", default=False, help="")387 parser.add_argument("--rule_based", action="store_true", default=False, help="")388 parser.add_argument("--log_info_interval", type=int, default=1, help="")389 parser.add_argument("--log_root", type=str, default='logs', help="")390 parser.add_argument("--results_root", type=str, default="runs", help="")391 parser.add_argument("--group_name", type=str, default=None, help="")392 parser.add_argument("--experiment_group", type=str, default="single", help="")393 parser.add_argument("--rewards_config_file", type=str, default='', help="")394 parser.add_argument("--gymrt2a_map_size", type=int, default=20)395 parser.add_argument("--gymrt2a_num_agents", type=int, default=10)396 parser.add_argument("--gymrt2a_objects", type=int, nargs="+", default=[4, 3, 3])397 parser.add_argument("--gymrt2a_view_range", type=int, default=5)398 parser.add_argument("--gymrt2a_comm_range", type=int, default=8)399 parser.add_argument("--gymrt2a_comm_dim", type=int, default=0)400 parser.add_argument("--gymrt2a_share_intention", type=str, default="path")401 parser.add_argument("--gymrt2a_reward_exponent", type=float, default=2.0)402 parser.add_argument("--gymrt2a_sectors", type=int, default=2)403 parser.add_argument("--gymrt2a_respawn", type=lambda v: str(v).lower() in ("1", "true", "yes"), default=True)404 parser.add_argument("--gymrt2a_ego_dim", type=int, default=3)405 parser.add_argument("--gymrt2a_other_dim", type=int, default=4)406407408def main(args):409 parser = get_args()410 config = parse_args(args, parser)411412 get_config(config)413414 if (415 (416 getattr(config, "explore_coverage_features", False)417 or getattr(config, "explore_observer_relative_geometry", False)418 )419 and getattr(config, "hatt_shared_task_obs", False)420 and getattr(config, "use_gat", False)421 and str(getattr(config, "graph_model", "")).startswith("hatt")422 ):423 raise ValueError(424 "Observer-relative exploration features require "425 "hatt_shared_task_obs=false; broadcasting agent 0 task features "426 "would discard the other agents' candidate geometry and coverage."427 )428429 if getattr(config, "reconfiguration_mode", "legacy") == "learned":430 if not getattr(config, "use_gat", False) or config.graph_model != "hatt_v2_0":431 raise ValueError(432 "reconfiguration_mode='learned' is frozen to use_gat=true and "433 "graph_model='hatt_v2_0'; use periodic/event/urgency for other policies."434 )435436 trainer_algorithm_name = (437 getattr(config, "trainer_algorithm_name", None) or config.algorithm_name438 )439 value_decomposition_algorithms = {"qmix", "vdn", "hatt_qmix"}440 if str(config.algorithm_name).lower() in value_decomposition_algorithms:441 trainer_algorithm_name = str(config.algorithm_name).lower()442 if not getattr(config, "use_pcfa", False):443 if trainer_algorithm_name not in ("rmappo", "mappo", *value_decomposition_algorithms):444 raise NotImplementedError(445 f"unsupported trainer_algorithm_name {trainer_algorithm_name!r}"446 )447 config.trainer_algorithm_name = trainer_algorithm_name448449 if trainer_algorithm_name == "rmappo":450 assert config.use_recurrent_policy or config.use_naive_recurrent_policy, "check recurrent policy!"451 elif trainer_algorithm_name == "mappo":452 assert (453 config.use_recurrent_policy == False and config.use_naive_recurrent_policy == False454 ), "check recurrent policy!"455 else:456 # PCFA does not use a neural trainer; keep the field for logging only.457 config.trainer_algorithm_name = trainer_algorithm_name or "pcfa"458459 assert (460 config.share_policy == True and config.scenario_name == "simple_speaker_listener"461 ) == False, "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."462463 # cuda464 if config.cuda >= 0 and torch.cuda.is_available():465 print("choose to use gpu...")466 device = torch.device(f"cuda:{config.cuda}")467 torch.set_num_threads(config.n_training_threads)468 if config.cuda_deterministic:469 torch.backends.cudnn.benchmark = False470 torch.backends.cudnn.deterministic = True471 else:472 print("choose to use cpu...")473 device = torch.device("cpu")474 torch.set_num_threads(config.n_training_threads)475476 run_dir, tb_dir, tb_link_dir = _prepare_run_dirs(config)477478 setproctitle.setproctitle(479 str(config.algorithm_name)480 + "-"481 + str(config.env_name)482 + "-"483 + str(config.experiment_name)484 + "@"485 + str(config.user_name)486 )487488 # seed489 torch.manual_seed(config.seed)490 torch.cuda.manual_seed_all(config.seed)491 np.random.seed(config.seed)492493 if str(config.env_name) in ("gymrt2a", "mrt2a"):494 config.num_agents = int(config.gymrt2a_num_agents)495 config.eval_reset_mode = "full_random"496 else:497 config.num_agents = config.num_uav + config.num_ugv498499 # env init500 envs = make_train_env(config)501 eval_envs = make_eval_env(config) if config.use_eval else None502503 setattr(config, 'envs', envs)504 setattr(config, 'eval_envs', eval_envs)505 setattr(config, 'device', device)506 setattr(config, 'run_dir', run_dir)507 setattr(config, 'tb_dir', tb_dir)508 setattr(config, 'tb_link_dir', tb_link_dir)509 _write_config_snapshot(config, run_dir)510511 # run experiments512 # if config.share_policy:513 # from runner.shared.env_runner import EnvRunner as Runner514515 # if config["use_gat"]:516 if config.rule_based:517 from runner.rule_based.macf_runner import MACFEnvRunner as Runner518 elif getattr(config, "use_pcfa", False):519 from runner.pcfa.macf_runner import MACFEnvRunner as Runner520 elif str(config.trainer_algorithm_name).lower() in value_decomposition_algorithms:521 if str(config.env_name) in ("gymrt2a", "mrt2a"):522 raise NotImplementedError("QMIX-family training currently targets capcoal/delivcoal Dict interfaces.")523 from runner.qmix.macf_runner import MACFRunner as Runner524 else:525 if str(config.env_name) in ("gymrt2a", "mrt2a"):526 from runner.shared.gymrt2a_runner import GymRT2ARunner as Runner527 else:528 from runner.shared.macf_runner import MACFEnvRunner as Runner529 # else:530 # from runner.separated_nogat.env_runner import EnvRunner as Runner531 # from runner.pcfa.macf_runner import MACFEnvRunner as Runner532533 # from runner.qmix.macf_runner import MACFRunner as Runner534535 runner = Runner(config)536 runner.run()537538 # post process539 envs.close()540 if config.use_eval and eval_envs is not envs:541 eval_envs.close()542543 runner.writter.export_scalars_to_json(str(run_dir / "summary.json"))544 runner.writter.close()545 if hasattr(runner, "writer") and runner.writer is not runner.writter:546 runner.writer.flush()547 runner.writer.close()548549550if __name__ == "__main__":551 main(sys.argv[1:])