"""PCFA runner for MACF. This runner executes the non-learning Project Manager-oriented Coalition Formation Algorithm and logs episode-level metrics. """ import os import pickle import json import numpy as np from .base_runner import Runner from algorithms.algorithm.pcfa_policy import PCFAPolicy class MACFEnvRunner(Runner): """Runner for PCFA baseline.""" def __init__(self, config): super(MACFEnvRunner, self).__init__(config) self.episode = 0 self.info_buffer = [] print(self.log_dir) # Retrieve special task IDs from the underlying environment. child_env = self.envs.envs[0] core_env = getattr(child_env, 'env', child_env) idle_task_id = core_env.idle_task_id explore_task_id = core_env.explore_task_id resupply_task_id = core_env.resupply_task_id share_observation_space = ( self.envs.share_observation_space[0] if self.use_centralized_V else self.envs.observation_space[0] ) self.policy = PCFAPolicy( config, self.envs.observation_space[0], share_observation_space, self.envs.action_space[0], device=config.device, ) self.policy.set_task_layout( delivery_task_num=getattr( core_env, 'delivery_task_num', getattr(core_env, 'target_num', explore_task_id) ), explore_task_start=getattr(core_env, 'explore_task_start', explore_task_id), explore_task_end=getattr(core_env, 'explore_task_end', explore_task_id + 1), idle_task_id=idle_task_id, resupply_task_id=resupply_task_id, use_coalition_delivery=getattr(core_env, 'use_coalition_delivery', False), ) if self.n_rollout_threads != 1: raise ValueError( "PCFA evaluation currently requires n_rollout_threads=1 so " "terminal episodes are not mixed after vector-env auto-reset." ) # Dummy RNN states kept only to satisfy the policy interface. self.rnn_states = np.zeros( ( self.episode_length + 1, self.n_rollout_threads, self.num_agents, self.recurrent_N, self.hidden_size, ), dtype=np.float32, ) self.rnn_states_critic = np.zeros_like(self.rnn_states) # Last actions per env thread; used to keep inactive agents on course. self.last_actions = np.full( (self.n_rollout_threads, self.num_agents), self.policy.idle_task_id, dtype=int, ) def run(self): episodes = int(self.num_env_steps) // self.episode_length // self.n_rollout_threads os.makedirs(os.path.join(self.log_dir, 'infos'), exist_ok=True) episode_results = [] next_initial_state = None for self.episode in range(episodes): if next_initial_state is None: obs, active_masks, available_actions = self.envs.reset(self.episode) else: # DummyVecEnv has already reset a terminal environment and # returned the next episode's initial observation. Reuse it # so the deterministic seed stream does not skip scenarios. obs, active_masks, available_actions = next_initial_state next_initial_state = None self.info_buffer = [] self.last_actions[:] = self.policy.idle_task_id episode_reward = 0.0 episode_task_completed = 0 episode_sent = 0 episode_sent_goal = 0 steps_taken = 0 last_info = None for step in range(self.episode_length): new_actions = self._compute_actions(obs, active_masks, available_actions) # Inactive agents keep their previous action to avoid switching to idle. active = np.asarray(active_masks, dtype=bool).reshape( self.n_rollout_threads, self.num_agents ) actions = np.where(active, new_actions, self.last_actions) self.last_actions = actions.copy() obs, rewards, dones, infos, active_masks, _bad_mask, available_actions = self.envs.step(actions) self.info_buffer.append(infos) episode_reward += float(rewards.mean()) episode_sent += infos[0].get('sent', 0) steps_taken = step + 1 last_info = infos[0] if bool(np.all(dones[0])): next_initial_state = (obs, active_masks, available_actions) break delivery_tasks = last_info['tasks'][: self.policy.delivery_task_num] episode_task_completed = sum(1 for t in delivery_tasks if t['completed']) episode_task_discovered = sum(1 for t in delivery_tasks if t['discovered']) episode_sent_goal = last_info.get('sent_goal', 1) success = float(episode_task_completed == self.policy.delivery_task_num) task_completion_ratio = float(episode_task_completed) / max( 1, self.policy.delivery_task_num ) discovered_ratio = float(episode_task_discovered) / max( 1, self.policy.delivery_task_num ) sent_percentage = float(episode_sent) / max(episode_sent_goal, 1) episode_results.append({ 'episode': int(self.episode), 'success': success, 'task_completed': int(episode_task_completed), 'task_completion_ratio': task_completion_ratio, 'task_discovered': int(episode_task_discovered), 'discovered_ratio': discovered_ratio, 'sent': float(episode_sent), 'sent_goal': float(episode_sent_goal), 'sent_percentage': sent_percentage, 'episode_length': int(steps_taken), 'episode_reward': float(episode_reward), }) if self.episode % self.log_interval == 0: print( f"episode {self.episode}/{episodes} | reward {episode_reward:.3f} | " f"completed {episode_task_completed} | sent {episode_sent}/{episode_sent_goal} | " f"success {success:.0f} | steps {steps_taken}" ) self.writter.add_scalar('reward/episode', episode_reward, self.episode) self.writter.add_scalar('reward/task_completed', episode_task_completed, self.episode) self.writter.add_scalar('reward/sent_percentage', sent_percentage, self.episode) self.writter.add_scalar('eval/success', success, self.episode) self.writter.add_scalar( 'eval/task_completion_ratio', task_completion_ratio, self.episode ) self.writter.add_scalar('eval/discovered_ratio', discovered_ratio, self.episode) self.writter.add_scalar('eval/episode_length', steps_taken, self.episode) if self.log_info_interval > 0 and self.episode % self.log_info_interval == 0: with open(os.path.join(self.log_dir, 'infos', f'{self.episode:04d}.pkl'), 'wb') as f: pickle.dump(self.info_buffer, f) print(f'{self.episode:04d}.pkl generated') aggregate = { 'episodes': len(episode_results), 'success_rate': float(np.mean([row['success'] for row in episode_results])), 'task_completion_ratio': float(np.mean([ row['task_completion_ratio'] for row in episode_results ])), 'discovered_ratio': float(np.mean([ row['discovered_ratio'] for row in episode_results ])), 'sent_percentage': float(np.mean([ row['sent_percentage'] for row in episode_results ])), 'episode_length': float(np.mean([ row['episode_length'] for row in episode_results ])), } output = { 'algorithm': 'pcfa-delivcoal' if self.policy.use_coalition_delivery else 'pcfa', 'seed': int(getattr(self.config, 'seed', 0)), 'aggregate': aggregate, 'episodes': episode_results, } output_path = os.path.join(str(self.config.run_dir), 'pcfa_results.json') temp_path = output_path + '.tmp' with open(temp_path, 'w', encoding='utf-8') as file: json.dump(output, file, ensure_ascii=False, indent=2) os.replace(temp_path, output_path) def _compute_actions(self, obs, active_masks, available_actions): env_num = self.n_rollout_threads actions = np.full((env_num, self.num_agents), self.policy.idle_task_id, dtype=int) rnn_states_actor = self.rnn_states[0] rnn_states_critic = self.rnn_states_critic[0] masks = np.ones((env_num, self.num_agents, 1), dtype=np.float32) for env_i in range(env_num): env_obs = {k: v[env_i:env_i + 1] for k, v in obs.items()} env_available_actions = available_actions[env_i:env_i + 1] env_active_masks = active_masks[env_i:env_i + 1] _, acts, _, _, _ = self.policy.get_actions( None, env_obs, rnn_states_actor[env_i:env_i + 1], rnn_states_critic[env_i:env_i + 1], masks[env_i:env_i + 1], env_available_actions, env_active_masks, deterministic=True, ) actions[env_i] = acts[0] return actions