""" # @Time : 2021/6/30 10:07 下午 # @Author : hezhiqiang # @Email : tinyzqh@163.com # @File : train.py """ # !/usr/bin/env python import sys import os import re import socket import setproctitle import numpy as np from pathlib import Path import torch import yaml from datetime import datetime # Get the parent directory of the current file parent_dir = os.path.abspath(os.path.join(os.getcwd(), ".")) # Append the parent directory to sys.path, otherwise the following import will fail sys.path.append(parent_dir) from config import get_args, get_config from envs.macf.env_wrappers import DummyVecEnv """Train script for MPEs.""" REPO_ROOT = Path(__file__).resolve().parents[1] def _safe_name(value): safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value)).strip("._") if not safe: raise ValueError(f"Invalid path name: {value!r}") return safe def _resolve_repo_path(value): path = Path(value).expanduser() return path if path.is_absolute() else REPO_ROOT / path def _unique_path(path): if not path.exists(): return path for index in range(2, 1000): candidate = path.with_name(f"{path.name}-{index}") if not candidate.exists(): return candidate raise RuntimeError(f"Could not allocate a unique run directory for {path}") def _link_tensorboard_dir(tb_dir, link_dir): link_dir.parent.mkdir(parents=True, exist_ok=True) if link_dir.is_symlink(): if link_dir.resolve() == tb_dir.resolve(): return raise FileExistsError(f"TensorBoard link already exists: {link_dir}") if link_dir.exists(): raise FileExistsError(f"TensorBoard link path already exists: {link_dir}") relative_target = os.path.relpath(tb_dir, start=link_dir.parent) link_dir.symlink_to(relative_target, target_is_directory=True) def _prepare_run_dirs(config): configured_run_dir = getattr(config, "run_dir", None) if configured_run_dir: run_dir = _resolve_repo_path(configured_run_dir) else: results_root = _resolve_repo_path(getattr(config, "results_root", "runs")) group = _safe_name( getattr(config, "group_name", None) or getattr(config, "experiment_group", "single") ) env_name = _safe_name(config.env_name) timestamp = datetime.now().strftime("%Y%m%d-%H%M") group_dir = _unique_path(results_root / f"{timestamp}_{env_name}_{group}") run_name = f"{_safe_name(config.algorithm_name)}_{_safe_name(config.experiment_name)}" run_dir = group_dir / run_name tb_dir = run_dir / "tb" tb_dir.mkdir(parents=True, exist_ok=False) log_root = _resolve_repo_path(getattr(config, "log_root", "logs")) tb_link_dir = log_root / run_dir.parent.name / run_dir.name _link_tensorboard_dir(tb_dir, tb_link_dir) return run_dir, tb_dir, tb_link_dir def _write_config_snapshot(config, run_dir): snapshot = {} for key, value in vars(config).items(): if key in {"envs", "eval_envs", "device"}: continue snapshot[key] = str(value) if isinstance(value, Path) else value with (run_dir / "config.yaml").open("w", encoding="utf-8") as file: yaml.safe_dump(snapshot, file, sort_keys=False, allow_unicode=True) def _discrete_env_class(config): """按 env_name 选择具体环境包。 taskcf 是原任务能力联盟组建环境；coadel 是 coalition delivery 环境。 macf/MyEnv/coalition_delivery 作为历史配置别名保留。 """ env_name = str(getattr(config, "env_name", "taskcf")) if env_name in ("capcoal", "macf", "MyEnv", "taskcf"): from envs.capcoal.env_discrete import DiscreteActionEnv return DiscreteActionEnv if env_name in ("delivcoal", "coalition_delivery", "coadel"): from envs.delivcoal.env_discrete import DiscreteActionEnv return DiscreteActionEnv if env_name in ("gymrt2a", "mrt2a"): from envs.gymrt2a.env_discrete import DiscreteActionEnv return DiscreteActionEnv raise ValueError( f"Unknown env_name={env_name!r}; expected capcoal, delivcoal, or gymrt2a." ) def make_train_env(config): def get_env_fn(rank): def init_env(): DiscreteActionEnv = _discrete_env_class(config) env = DiscreteActionEnv(config) env.seed(config.seed + rank * 1000) return env return init_env return DummyVecEnv([get_env_fn(i) for i in range(config.n_rollout_threads)]) def make_eval_env(config): def get_env_fn(rank): def init_env(): DiscreteActionEnv = _discrete_env_class(config) env = DiscreteActionEnv(config) eval_reset_mode = getattr(config, "eval_reset_mode", None) if eval_reset_mode is not None: env.set_reset_mode(eval_reset_mode) env.seed(config.eval_seed + rank * 1000) return env return init_env return DummyVecEnv([get_env_fn(i) for i in range(config.n_eval_rollout_threads)]) def parse_args(args, parser): parser.add_argument("--scenario_name", type=str, default="MyEnv", help="Which scenario to run on") parser.add_argument("--num_landmarks", type=int, default=3) add_my_args(parser) config = parser.parse_known_args(args)[0] return config def add_my_args(parser): parser.add_argument("--config_file", type=str, default=None, help="path of config file (*.yaml)") # GAT parser.add_argument("--attr_dim", type=int, default=4, help="") parser.add_argument("--gat_hidden_dim", type=int, default=64, help="") parser.add_argument("--gat_out_dim", type=int, default=32, help="") parser.add_argument("--gat_mode", type=str, default="gated", choices=["gat", "cat", "gated"], help="GAT fusion mode") parser.add_argument( "--graph_model", type=str, default="legacy_gat", choices=[ "legacy_gat", "hatt", "hatt_v2_0", "hatt_v2_0_raw", "hatt_v2_0_physics", "hatt_v2_0_taskattn", "hatt_v2_0_noidle", "hatt_v2_0_resource", "hatt_v2_0_setctx", "hatt_v2_0_plain_noidle", "hatt_v2_0_detached", "hatt_v2_0_split", "hatt_v2_0_type_embed", "hatt_v2_0_type_gate", "hatt_v2_1", "hatt_v2", "hatt_v2_3a", "hatt_v2_3b", "hatt_v2_3c", "hatt_v2_4", "hatt_v2_5a", "hatt_v2_5b", "hatt_v2_5c", "hatt_v2_5", "hatt_v2_6a", "hatt_v2_6b", "hatt_v2_6c", "hatt_v2_6", "hatt_v2_7a", "hatt_v2_7b", "hatt_v2_7", "hatt_v3_0a", "hatt_v3_0b", "hatt_v3_0c", "hatt_v3_1a", "hatt_v3_1b", "hatt_v3_2a", "hatt_v3_2b", "hatt_v3_2c", "hatt_v3_2d", "hatt_v3_2r", "hatt_v3_2r_no_transition", "hatt_v3_2r_no_member", "hatt_v3_3aa", "hatt_v3_4", "hatt_gymrt2a", ], ) parser.add_argument("--hatt_hidden_dim", type=int, default=128) parser.add_argument("--hatt_num_heads", type=int, default=4) parser.add_argument("--hatt_dropout", type=float, default=0.0) parser.add_argument( "--hatt_capability_scale", type=float, default=5.0, help="scale raw agent capabilities into the normalized task requirement space for HATT residuals", ) parser.add_argument( "--hatt_v20_task_input_mode", choices=("native", "restore_req_scale", "compact_projection"), default="native", help="controlled HATT v2.0 task-schema intervention used for B/C attribution", ) parser.add_argument( "--hatt_paired_init_seed", type=int, default=-1, help="local actor/critic initialization seed; -1 keeps historical RNG behavior", ) parser.add_argument( "--hatt_v20_init_scheme", choices=("default", "compact_function_preserving"), default="default", help="optional exact compact6-to-full12 actor initialization for schema attribution", ) parser.add_argument( "--hatt_shared_task_obs", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"), default=True, help="whether obs['tasks'] is shared across all agents and can be encoded once", ) parser.add_argument( "--hatt_use_aa_attention", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"), default=True, help="enable HATT v2.0 agent-agent relation attention", ) parser.add_argument( "--hatt_use_coalition_context", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"), default=True, help="enable HATT v2.0 task-indexed coalition context", ) parser.add_argument( "--hatt_candidate_aa_mode", type=str, default="learned", choices=["learned", "uniform", "zero"], help=( "HATT v3.3 A-A gate mode: learned for the method, uniform for the " "matched training control, zero for checkpoint intervention" ), ) parser.add_argument( "--hatt_aa_task_context", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"), default=True, help=( "HATT v3.4: feed the serving task's residual demand/priority into " "the A-A edge gate; false reproduces the v3.3AA 3-dim edge" ), ) parser.add_argument( "--reconfiguration_mode", choices=["legacy", "always", "periodic", "event", "urgency", "learned"], default="legacy", help="coalition reconsideration schedule; learned enables the HATT v2.0 PPO gate", ) parser.add_argument( "--allow_reconfiguration_mode_mismatch", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"), default=False, help=( "allow an explicit evaluation-only intervention between non-learned " "reconfiguration schedules; learned-gate mismatches remain forbidden" ), ) parser.add_argument("--reconfiguration_interval", type=int, default=10) parser.add_argument("--reconfiguration_urgency_threshold", type=float, default=0.7) parser.add_argument( "--reconfiguration_urgency_aggregation", choices=["mean", "sum"], default="mean", ) parser.add_argument("--reconfiguration_switch_cost", type=float, default=0.0) parser.add_argument("--reconfiguration_gate_init", type=float, default=0.8) parser.add_argument("--reconfiguration_gate_floor", type=float, default=0.02) parser.add_argument( "--reconfiguration_gate_parameterization", choices=["mixture", "switch", "hierarchical"], default="mixture", help=( "mixture preserves the historical latent reconsideration gate; " "switch makes the gate the identifiable probability of leaving " "the current task" ), ) parser.add_argument( "--reconfiguration_gate_override", choices=["none", "open", "closed", "event", "event_or_learned"], default="none", help=( "gate intervention; event_or_learned is a safe hybrid that forces " "event decisions open and learns only the residual periodic switches" ), ) parser.add_argument( "--reconfiguration_decision_coef", type=float, default=0.0, help="regularization cost for learned (non-mandatory) hierarchical redecisions", ) parser.add_argument( "--reconfiguration_context_features", type=lambda value: str(value).lower() in ("1", "true", "yes", "y"), default=False, help="expose explicit event/time context to the learned reconfiguration gate", ) parser.add_argument("--reconfiguration_context_dim", type=int, default=9) # Act-Crs parser.add_argument("--obs_shape", type=int, default=32, help="") parser.add_argument("--share_obs_shape", type=int, default=32, help="") # Enviroment parser.add_argument("--map_size", type=float, default=10000, help="") parser.add_argument("--num_uav", type=int, default=8, help="") parser.add_argument("--num_ugv", type=int, default=4, help="") parser.add_argument("--num_agents", type=int, default=12, help="") parser.add_argument("--uav_perception_rg", type=list, default=[50, 200], help="") parser.add_argument("--ugv_perception_rg", type=list, default=[10, 30], help="") parser.add_argument("--target_num", type=int, default=20, help="任务数量") parser.add_argument("--gat_dim", type=int, default=4, help="输出维度") parser.add_argument("--env_reset_mode", type=str, default="jitter", choices=["jitter", "full_random"], help="'jitter': base env + task jitter; 'full_random': full reinit each episode") parser.add_argument("--eval_reset_mode", type=str, default=None, choices=["jitter", "full_random"], help="optional reset mode used only by evaluation envs") parser.add_argument("--curriculum_switch_ep", type=int, default=-1, help="switch env_reset_mode to full_random after this env_episode; -1 = no switch") parser.add_argument("--max_steps", type=int, default=1000, help="最大步数") parser.add_argument( "--episode_time_limit", type=int, default=None, help="optional termination-only horizon; max_steps keeps time/deadline semantics", ) parser.add_argument("--cap_dim", type=int, default=4, help="能力维度") # Runner parser.add_argument("--use_gat", action="store_true", default=False, help="") parser.add_argument("--rule_based", action="store_true", default=False, help="") parser.add_argument("--log_info_interval", type=int, default=1, help="") parser.add_argument("--log_root", type=str, default='logs', help="") parser.add_argument("--results_root", type=str, default="runs", help="") parser.add_argument("--group_name", type=str, default=None, help="") parser.add_argument("--experiment_group", type=str, default="single", help="") parser.add_argument("--rewards_config_file", type=str, default='', help="") parser.add_argument("--gymrt2a_map_size", type=int, default=20) parser.add_argument("--gymrt2a_num_agents", type=int, default=10) parser.add_argument("--gymrt2a_objects", type=int, nargs="+", default=[4, 3, 3]) parser.add_argument("--gymrt2a_view_range", type=int, default=5) parser.add_argument("--gymrt2a_comm_range", type=int, default=8) parser.add_argument("--gymrt2a_comm_dim", type=int, default=0) parser.add_argument("--gymrt2a_share_intention", type=str, default="path") parser.add_argument("--gymrt2a_reward_exponent", type=float, default=2.0) parser.add_argument("--gymrt2a_sectors", type=int, default=2) parser.add_argument("--gymrt2a_respawn", type=lambda v: str(v).lower() in ("1", "true", "yes"), default=True) parser.add_argument("--gymrt2a_ego_dim", type=int, default=3) parser.add_argument("--gymrt2a_other_dim", type=int, default=4) def main(args): parser = get_args() config = parse_args(args, parser) get_config(config) if ( ( getattr(config, "explore_coverage_features", False) or getattr(config, "explore_observer_relative_geometry", False) ) and getattr(config, "hatt_shared_task_obs", False) and getattr(config, "use_gat", False) and str(getattr(config, "graph_model", "")).startswith("hatt") ): raise ValueError( "Observer-relative exploration features require " "hatt_shared_task_obs=false; broadcasting agent 0 task features " "would discard the other agents' candidate geometry and coverage." ) if getattr(config, "reconfiguration_mode", "legacy") == "learned": if not getattr(config, "use_gat", False) or config.graph_model != "hatt_v2_0": raise ValueError( "reconfiguration_mode='learned' is frozen to use_gat=true and " "graph_model='hatt_v2_0'; use periodic/event/urgency for other policies." ) trainer_algorithm_name = ( getattr(config, "trainer_algorithm_name", None) or config.algorithm_name ) value_decomposition_algorithms = {"qmix", "vdn", "hatt_qmix"} if str(config.algorithm_name).lower() in value_decomposition_algorithms: trainer_algorithm_name = str(config.algorithm_name).lower() if not getattr(config, "use_pcfa", False): if trainer_algorithm_name not in ("rmappo", "mappo", *value_decomposition_algorithms): raise NotImplementedError( f"unsupported trainer_algorithm_name {trainer_algorithm_name!r}" ) config.trainer_algorithm_name = trainer_algorithm_name if trainer_algorithm_name == "rmappo": assert config.use_recurrent_policy or config.use_naive_recurrent_policy, "check recurrent policy!" elif trainer_algorithm_name == "mappo": assert ( config.use_recurrent_policy == False and config.use_naive_recurrent_policy == False ), "check recurrent policy!" else: # PCFA does not use a neural trainer; keep the field for logging only. config.trainer_algorithm_name = trainer_algorithm_name or "pcfa" assert ( config.share_policy == True and config.scenario_name == "simple_speaker_listener" ) == False, "The simple_speaker_listener scenario can not use shared policy. Please check the config.py." # cuda if config.cuda >= 0 and torch.cuda.is_available(): print("choose to use gpu...") device = torch.device(f"cuda:{config.cuda}") torch.set_num_threads(config.n_training_threads) if config.cuda_deterministic: torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True else: print("choose to use cpu...") device = torch.device("cpu") torch.set_num_threads(config.n_training_threads) run_dir, tb_dir, tb_link_dir = _prepare_run_dirs(config) setproctitle.setproctitle( str(config.algorithm_name) + "-" + str(config.env_name) + "-" + str(config.experiment_name) + "@" + str(config.user_name) ) # seed torch.manual_seed(config.seed) torch.cuda.manual_seed_all(config.seed) np.random.seed(config.seed) if str(config.env_name) in ("gymrt2a", "mrt2a"): config.num_agents = int(config.gymrt2a_num_agents) config.eval_reset_mode = "full_random" else: config.num_agents = config.num_uav + config.num_ugv # env init envs = make_train_env(config) eval_envs = make_eval_env(config) if config.use_eval else None setattr(config, 'envs', envs) setattr(config, 'eval_envs', eval_envs) setattr(config, 'device', device) setattr(config, 'run_dir', run_dir) setattr(config, 'tb_dir', tb_dir) setattr(config, 'tb_link_dir', tb_link_dir) _write_config_snapshot(config, run_dir) # run experiments # if config.share_policy: # from runner.shared.env_runner import EnvRunner as Runner # if config["use_gat"]: if config.rule_based: from runner.rule_based.macf_runner import MACFEnvRunner as Runner elif getattr(config, "use_pcfa", False): from runner.pcfa.macf_runner import MACFEnvRunner as Runner elif str(config.trainer_algorithm_name).lower() in value_decomposition_algorithms: if str(config.env_name) in ("gymrt2a", "mrt2a"): raise NotImplementedError("QMIX-family training currently targets capcoal/delivcoal Dict interfaces.") from runner.qmix.macf_runner import MACFRunner as Runner else: if str(config.env_name) in ("gymrt2a", "mrt2a"): from runner.shared.gymrt2a_runner import GymRT2ARunner as Runner else: from runner.shared.macf_runner import MACFEnvRunner as Runner # else: # from runner.separated_nogat.env_runner import EnvRunner as Runner # from runner.pcfa.macf_runner import MACFEnvRunner as Runner # from runner.qmix.macf_runner import MACFRunner as Runner runner = Runner(config) runner.run() # post process envs.close() if config.use_eval and eval_envs is not envs: eval_envs.close() runner.writter.export_scalars_to_json(str(run_dir / "summary.json")) runner.writter.close() if hasattr(runner, "writer") and runner.writer is not runner.writter: runner.writer.flush() runner.writer.close() if __name__ == "__main__": main(sys.argv[1:])