"""Current-environment runner for recurrent QMIX, VDN and HATT-QMIX.""" import os import time from pathlib import Path import numpy as np import torch from tensorboardX import SummaryWriter from algorithms.qmix.algorithm.QMixPolicy import QMixPolicy from algorithms.qmix.qmix import QMix from runner.shared.macf_runner import MACFEnvRunner from utils.obs_utils import flatten_obs_dict, flatten_state_dict from utils.qmix_buffer import EpisodeReplayBuffer def _env_item(tree, index): if isinstance(tree, dict): return {key: np.asarray(value[index], dtype=np.float32).copy() for key, value in tree.items()} return np.asarray(tree[index], dtype=np.float32).copy() def _append_obs(storage, obs): if isinstance(obs, dict): for key, value in obs.items(): storage.setdefault(key, []).append(value) else: storage.append(obs) def _stack_obs(storage): if isinstance(storage, dict): return {key: np.stack(values).astype(np.float32) for key, values in storage.items()} return np.stack(storage).astype(np.float32) class MACFRunner(MACFEnvRunner): """Off-policy runner aligned with ``DummyVecEnv`` and sweep checkpoints.""" def __init__(self, config): self.config = config self.envs = config.envs self.eval_envs = config.eval_envs self.device = config.device self.num_agents = int(config.num_agents) self.env_name = config.env_name self.algorithm_name = config.algorithm_name self.experiment_name = config.experiment_name self.num_env_steps = int(config.num_env_steps) self.episode_length = int(config.episode_length) self.n_rollout_threads = int(config.n_rollout_threads) self.n_eval_rollout_threads = int(config.n_eval_rollout_threads) self.recurrent_N = 1 self.hidden_size = int(config.qmix_agent_hidden_dim) self.use_eval = bool(config.use_eval) self.eval_interval = config.eval_interval self.eval_episodes = int(config.eval_episodes) self.eval_seed = int(config.eval_seed) self.save_interval = config.save_interval self.log_interval = int(config.log_interval) self.log_info_interval = int(config.log_info_interval) self.model_dir = config.model_dir self.run_dir = Path(config.run_dir) self.log_dir = str(getattr(config, "eval_output_dir", self.run_dir)) self.tb_dir = str(getattr(config, "tb_dir", self.run_dir / "tb")) os.makedirs(self.log_dir, exist_ok=True) os.makedirs(self.tb_dir, exist_ok=True) self.writter = SummaryWriter(self.tb_dir) self.writer = self.writter self.save_dir = str(self.run_dir / "models") os.makedirs(self.save_dir, exist_ok=True) obs_space = self.envs.observation_space[0] action_space = self.envs.action_space[0] # The PPO shared-state space also declares a per-focal one-hot field. # QMIX consumes the environment's single joint state directly. state_spaces = self.envs.share_observation_space[0].spaces state_dim = sum( int(np.prod(space.shape)) for key, space in state_spaces.items() if key != "focal" ) self.policy = QMixPolicy( config, obs_space, action_space, self.num_agents, self.device ) self.trainer = QMix( config, self.policy, state_dim, self.num_agents, self.device ) self.buffer = EpisodeReplayBuffer(config.qmix_buffer_size, config.seed) self.env_episode = 0 self.episode = 0 self.best_eval_score = None self.best_eval_summary = None self.last_save_env_episode = 0 self.last_eval_env_episode = 0 self.last_train_env_episode = 0 self.last_target_env_episode = 0 if self.model_dir: self.restore() def _policy_obs(self, obs): if self.policy.use_hatt: return obs return flatten_obs_dict(obs) def _new_episode(self, obs, state, available, active, env_index): initial_obs = _env_item(self._policy_obs(obs), env_index) obs_storage = {} if isinstance(initial_obs, dict) else [] _append_obs(obs_storage, initial_obs) return { "obs": obs_storage, "states": [np.asarray(state[env_index], dtype=np.float32).copy()], "available_actions": [np.asarray(available[env_index], dtype=np.float32).copy()], "active_masks": [np.asarray(active[env_index], dtype=np.float32).copy()], "alive_masks": [np.ones((self.num_agents, 1), dtype=np.float32)], "actions": [], "rewards": [], "dones": [], "held_actions": [], "return": 0.0, "sent": 0.0, "sent_goal": 0.0, } def _finish_episode(self, episode): packed = { "obs": _stack_obs(episode["obs"]), "states": np.stack(episode["states"]).astype(np.float32), "available_actions": np.stack(episode["available_actions"]).astype(np.float32), "actions": np.stack(episode["actions"]).astype(np.int64), "rewards": np.asarray(episode["rewards"], dtype=np.float32).reshape(-1, 1), "dones": np.asarray(episode["dones"], dtype=np.float32).reshape(-1, 1), "active_masks": np.stack(episode["active_masks"]).astype(np.float32), "alive_masks": np.stack(episode["alive_masks"]).astype(np.float32), "held_actions": np.stack(episode["held_actions"]).astype(np.int64), } packed = self._pack_transition_mode(packed) self.buffer.insert(packed) def _pack_transition_mode(self, packed): """Attach physical discounts or compress an episode to decision events.""" gamma = float(self.config.gamma) transitions = len(packed["actions"]) if self.config.qmix_transition_mode != "event_smdp": packed["discounts"] = np.full( (transitions, 1), gamma, dtype=np.float32 ) return packed starts = [ step for step in range(transitions) if np.any(packed["active_masks"][step] > 0.5) ] if not starts: starts = [0] endpoints = starts[1:] + [transitions] obs = packed["obs"] event = { "obs": ({key: [value[starts[0]]] for key, value in obs.items()} if isinstance(obs, dict) else [obs[starts[0]]]), "states": [packed["states"][starts[0]]], "available_actions": [packed["available_actions"][starts[0]]], "active_masks": [packed["active_masks"][starts[0]]], "alive_masks": [packed["alive_masks"][starts[0]]], "actions": [], "rewards": [], "dones": [], "held_actions": [], "discounts": [], } for start, end in zip(starts, endpoints): duration = end - start weights = gamma ** np.arange(duration, dtype=np.float32) reward = float(np.sum(packed["rewards"][start:end, 0] * weights)) event["actions"].append(packed["actions"][start]) event["rewards"].append([reward]) event["dones"].append(packed["dones"][end - 1]) event["held_actions"].append(packed["held_actions"][end - 1]) event["discounts"].append([gamma ** duration]) if isinstance(obs, dict): for key, value in obs.items(): event["obs"][key].append(value[end]) else: event["obs"].append(obs[end]) for key in ( "states", "available_actions", "active_masks", "alive_masks" ): event[key].append(packed[key][end]) for key, value in event.items(): if key == "obs" and isinstance(value, dict): event[key] = {name: np.stack(items) for name, items in value.items()} elif key in ("actions", "held_actions"): event[key] = np.asarray(value, dtype=np.int64) else: event[key] = np.asarray(value, dtype=np.float32) return event def _epsilon(self, total_steps): fraction = min(1.0, total_steps / max(1, int(self.config.qmix_epsilon_anneal_steps))) return float(self.config.qmix_epsilon_start + fraction * ( self.config.qmix_epsilon_finish - self.config.qmix_epsilon_start )) @staticmethod def _advance_rollout_hidden(next_hidden, agent_dones, active_masks=None, previous_hidden=None, event_mode=False): """Advance recurrent state on every physical transition. Event-mode active masks control whether a newly proposed action is executed. They do not pause observation history: replay training and deterministic evaluation both feed every physical transition through the recurrent agent network. Keep collection aligned with those paths and reset only environments whose episodes ended. """ hidden = np.asarray(next_hidden, dtype=np.float32).copy() if event_mode: update = np.any( np.asarray(active_masks) > 0.5, axis=tuple(range(1, np.asarray(active_masks).ndim)) )[:, None, None] hidden = np.where(update, hidden, np.asarray(previous_hidden)) hidden[np.asarray(agent_dones, dtype=bool)] = 0.0 return hidden @staticmethod def _completed_interval_due(current_episode, last_episode, interval): intervals = interval if isinstance(interval, (list, tuple)) else [interval] return current_episode > last_episode and any( int(value) > 0 and current_episode % int(value) == 0 for value in intervals ) def run(self): obs, active_masks, available = self.envs.reset(0) state = flatten_state_dict(self.envs.get_global_state()) episodes = [ self._new_episode(obs, state, available, active_masks, i) for i in range(self.n_rollout_threads) ] hidden = np.zeros( (self.n_rollout_threads, self.num_agents, self.hidden_size), dtype=np.float32 ) last_actions = np.zeros((self.n_rollout_threads, self.num_agents, 1), dtype=np.int64) max_episodes = getattr(self.config, "max_train_episodes", None) total_steps = 0 updates = 0 last_info = {} last_episode_return = 0.0 last_sent_percentage = 0.0 last_log_episode = -1 start = time.time() while total_steps < self.num_env_steps and ( max_episodes is None or self.env_episode < int(max_episodes) ): epsilon = self._epsilon(total_steps) self.trainer.prep_rollout() with torch.no_grad(): actions_t, next_hidden_t = self.policy.select_actions( obs, hidden, available, epsilon=epsilon ) actions = actions_t.cpu().numpy().astype(np.int64) next_hidden = next_hidden_t.cpu().numpy() actions = active_masks.astype(np.int64) * actions + ( 1 - active_masks.astype(np.int64) ) * last_actions next_obs, rewards, dones, infos, next_active, _, next_available = self.envs.step( actions, self.env_episode ) next_state = flatten_state_dict(self.envs.get_global_state()) env_dones = np.all(dones, axis=1) completed_this_step = int(np.count_nonzero(env_dones)) for index in range(self.n_rollout_threads): episode = episodes[index] episode["actions"].append(actions[index, :, 0].copy()) episode["held_actions"].append(actions[index, :, 0].copy()) team_reward = float(np.mean(rewards[index])) * float(self.config.qmix_reward_scale) episode["rewards"].append(team_reward) episode["dones"].append(float(env_dones[index])) episode["return"] += team_reward episode["sent"] += float(infos[index].get("sent", 0.0)) episode["sent_goal"] = float(infos[index].get("sent_goal", 0.0)) _append_obs(episode["obs"], _env_item(self._policy_obs(next_obs), index)) episode["states"].append(np.asarray(next_state[index], dtype=np.float32).copy()) episode["available_actions"].append(np.asarray(next_available[index], dtype=np.float32).copy()) episode["active_masks"].append( np.asarray(next_active[index], dtype=np.float32).copy() ) episode["alive_masks"].append( (~np.asarray(dones[index], dtype=bool)).astype(np.float32)[:, None] ) if env_dones[index]: self._finish_episode(episode) self.env_episode += 1 last_info = infos[index] last_episode_return = episode["return"] last_sent_percentage = 100.0 * episode["sent"] / max( 1.0, episode["sent_goal"] ) self.writer.add_scalar("reward/episode", episode["return"], self.env_episode) self.writer.add_scalar("train/epsilon", epsilon, self.env_episode) episodes[index] = self._new_episode( next_obs, next_state, next_available, next_active, index ) hidden = self._advance_rollout_hidden( next_hidden, dones, active_masks=active_masks, previous_hidden=hidden, event_mode=self.config.qmix_transition_mode == "event_smdp", ) last_actions = actions.copy() for index, done in enumerate(env_dones): if done: last_actions[index] = 0 obs, state, active_masks, available = next_obs, next_state, next_active, next_available total_steps += self.n_rollout_threads train_interval = int(self.config.qmix_train_interval) learning_ready = len(self.buffer) >= int(self.config.qmix_learning_starts) if completed_this_step > 0 and not learning_ready: # Warm-up episodes fill replay without accumulating a gradient # update debt that would be paid all at once at learning start. self.last_train_env_episode = self.env_episode elif (completed_this_step > 0 and learning_ready and self.env_episode - self.last_train_env_episode >= train_interval): due_groups = ( self.env_episode - self.last_train_env_episode ) // train_interval for _ in range(due_groups * int(self.config.qmix_updates_per_train)): batch = self.buffer.sample( self.config.qmix_batch_size, self.config.qmix_sequence_length, self.config.qmix_burn_in, ) train_info = self.trainer.train_batch(batch) updates += 1 for key, value in train_info.items(): self.writer.add_scalar(f"train/{key}", value, total_steps) self.last_train_env_episode = self.env_episode target_interval = int(self.config.qmix_target_update_interval) if (completed_this_step > 0 and target_interval > 0 and self.env_episode - self.last_target_env_episode >= target_interval): self.trainer.hard_target_updates() self.last_target_env_episode = self.env_episode self.writer.add_scalar("train/target_updates", 1, self.env_episode) save_due = completed_this_step > 0 and self._completed_interval_due( self.env_episode, self.last_save_env_episode, self.save_interval ) if save_due: self.save() if getattr(self.config, "save_episode_checkpoints", False): self.save(tag=f"ep{self.env_episode:06d}") self.last_save_env_episode = self.env_episode eval_due = completed_this_step > 0 and self._completed_interval_due( self.env_episode, self.last_eval_env_episode, self.eval_interval ) if eval_due and self.use_eval: tag = f"ep{self.env_episode:06d}" self.save(tag=tag) summary = self.periodic_eval( total_steps, Path(self.save_dir) / f"actor_{tag}.pt", ) self._save_eval_checkpoints(summary, total_steps) self.last_eval_env_episode = self.env_episode if (self.env_episode != last_log_episode and self.env_episode and self.env_episode % max(1, self.log_interval) == 0): elapsed = max(time.time() - start, 1e-6) expected = elapsed * self.num_env_steps / max(1, total_steps) completed = 0 if last_info: completed = sum( int(task.get("completed", False)) for task in last_info.get("tasks", []) if int(task.get("type", -1)) == 1 ) target_episodes = int(max_episodes) if max_episodes else max( 1, self.num_env_steps // max(1, self.episode_length) ) print( f"\n rollout {self.env_episode}/{target_episodes}, " f"env_episode {self.env_episode}, {last_sent_percentage:.1f}% sent, " f"{completed} completed, ep_reward: {last_episode_return:.2f}, " f"timesteps {total_steps}/{self.num_env_steps}, " f"FPS {int(total_steps / elapsed)}, time cost {elapsed:.1f}, " f"expected time {expected:.1f}.\n" ) last_log_episode = self.env_episode self.save() def save(self, tag=None): suffix = f"_{tag}" if tag else "" torch.save( self.policy.get_actor_checkpoint(), os.path.join(self.save_dir, f"actor{suffix}.pt"), ) torch.save( self.trainer.state_dict(), os.path.join(self.save_dir, f"trainer{suffix}.pt"), ) def restore(self): checkpoint = torch.load( os.path.join(str(self.model_dir), "actor.pt"), map_location=self.device ) self.policy.load_actor_checkpoint(checkpoint)