rMAPPOPolicy.py
标准 MAPPO 策略对照
原文件 · 142 行 · 7843 字节
algorithms/algorithm/rMAPPOPolicy.py
SHA-256: b53b5014190cdcdbe9fa32d20c2281589fbc1d7e070e7c9574e4b454f41e31c6
1"""2# @Time : 2021/7/1 6:53 下午3# @Author : hezhiqiang014# @Email : hezhiqiang01@baidu.com5# @File : rMAPPOPolicy.py6"""78import torch9from algorithms.algorithm.r_actor_critic import R_Actor, R_Critic10from utils.util import update_linear_schedule11from utils.obs_utils import flatten_state_dict121314class RMAPPOPolicy:15 """16 MAPPO Policy class. Wraps actor and critic networks to compute actions and value function predictions.1718 :param args: (argparse.Namespace) arguments containing relevant model and policy information.19 :param obs_space: (gym.Space) observation space.20 :param cent_obs_space: (gym.Space) value function input space (centralized input for MAPPO, decentralized for IPPO).21 :param action_space: (gym.Space) action space.22 :param device: (torch.device) specifies the device to run on (cpu/gpu).23 """2425 def __init__(self, args, obs_space, cent_obs_space, act_space, device=torch.device("cpu")):26 self.device = device27 self.lr = args.lr28 self.critic_lr = args.critic_lr29 self.opti_eps = args.opti_eps30 self.weight_decay = args.weight_decay3132 self.obs_space = obs_space33 self.share_obs_space = cent_obs_space34 self.act_space = act_space3536 self.actor = R_Actor(args, self.obs_space, self.act_space, self.device)37 self.critic = R_Critic(args, self.share_obs_space, self.device)3839 self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),40 lr=self.lr, eps=self.opti_eps,41 weight_decay=self.weight_decay)42 self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),43 lr=self.critic_lr,44 eps=self.opti_eps,45 weight_decay=self.weight_decay)4647 def lr_decay(self, episode, episodes):48 """49 Decay the actor and critic learning rates.50 :param episode: (int) current training episode.51 :param episodes: (int) total number of training episodes.52 """53 update_linear_schedule(self.actor_optimizer, episode, episodes, self.lr)54 update_linear_schedule(self.critic_optimizer, episode, episodes, self.critic_lr)5556 @staticmethod57 def _critic_input(cent_obs):58 cent_obs = flatten_state_dict(cent_obs)59 if getattr(cent_obs, "ndim", 0) > 2:60 cent_obs = cent_obs.reshape(-1, cent_obs.shape[-1])61 return cent_obs6263 def get_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, masks, available_actions=None,64 deterministic=False):65 """66 Compute actions and value function predictions for the given inputs.67 :param cent_obs (np.ndarray): centralized input to the critic.68 :param obs (np.ndarray): local agent inputs to the actor.69 :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.70 :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.71 :param masks: (np.ndarray) denotes points at which RNN states should be reset.72 :param available_actions: (np.ndarray) denotes which actions are available to agent73 (if None, all actions available)74 :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.7576 :return values: (torch.Tensor) value function predictions.77 :return actions: (torch.Tensor) actions to take.78 :return action_log_probs: (torch.Tensor) log probabilities of chosen actions.79 :return rnn_states_actor: (torch.Tensor) updated actor network RNN states.80 :return rnn_states_critic: (torch.Tensor) updated critic network RNN states.81 """82 actions, action_log_probs, rnn_states_actor = self.actor(obs,83 rnn_states_actor,84 masks,85 available_actions,86 deterministic)8788 values, rnn_states_critic = self.critic(self._critic_input(cent_obs), rnn_states_critic, masks)89 return values, actions, action_log_probs, rnn_states_actor, rnn_states_critic9091 def get_values(self, cent_obs, rnn_states_critic, masks):92 """93 Get value function predictions.94 :param cent_obs (np.ndarray): centralized input to the critic.95 :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.96 :param masks: (np.ndarray) denotes points at which RNN states should be reset.9798 :return values: (torch.Tensor) value function predictions.99 """100 values, _ = self.critic(self._critic_input(cent_obs), rnn_states_critic, masks)101 return values102103 def evaluate_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,104 available_actions=None, active_masks=None):105 """106 Get action logprobs / entropy and value function predictions for actor update.107 :param cent_obs (np.ndarray): centralized input to the critic.108 :param obs (np.ndarray): local agent inputs to the actor.109 :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.110 :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.111 :param action: (np.ndarray) actions whose log probabilites and entropy to compute.112 :param masks: (np.ndarray) denotes points at which RNN states should be reset.113 :param available_actions: (np.ndarray) denotes which actions are available to agent114 (if None, all actions available)115 :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.116117 :return values: (torch.Tensor) value function predictions.118 :return action_log_probs: (torch.Tensor) log probabilities of the input actions.119 :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.120 """121 action_log_probs, dist_entropy = self.actor.evaluate_actions(obs,122 rnn_states_actor,123 action,124 masks,125 available_actions,126 active_masks)127128 values, _ = self.critic(self._critic_input(cent_obs), rnn_states_critic, masks)129 return values, action_log_probs, dist_entropy130131 def act(self, obs, rnn_states_actor, masks, available_actions=None, deterministic=False):132 """133 Compute actions using the given inputs.134 :param obs (np.ndarray): local agent inputs to the actor.135 :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.136 :param masks: (np.ndarray) denotes points at which RNN states should be reset.137 :param available_actions: (np.ndarray) denotes which actions are available to agent138 (if None, all actions available)139 :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.140 """141 actions, _, rnn_states_actor = self.actor(obs, rnn_states_actor, masks, available_actions, deterministic)142 return actions, rnn_states_actor