r_actor_critic.py
MLP/recurrent actor 与集中式 critic
原文件 · 173 行 · 8523 字节
algorithms/algorithm/r_actor_critic.py
SHA-256: cf854e669d44f8d341171a2136079656288f517ddd980fa56814cbd379ac45f2
1"""2# @Time : 2021/7/1 6:53 下午3# @Author : hezhiqiang014# @Email : hezhiqiang01@baidu.com5# @File : r_actor_critic.py6"""78import torch9import torch.nn as nn10from algorithms.utils.util import init, check11from algorithms.utils.cnn import CNNBase12from algorithms.utils.mlp import MLPBase13from algorithms.utils.rnn import RNNLayer14from algorithms.utils.act import ACTLayer15from algorithms.utils.popart import PopArt16from utils.util import get_shape_from_obs_space17from utils.obs_utils import flatten_obs_dict181920class R_Actor(nn.Module):21 """22 Actor network class for MAPPO. Outputs actions given observations.23 :param args: (argparse.Namespace) arguments containing relevant model information.24 :param obs_space: (gym.Space) observation space.25 :param action_space: (gym.Space) action space.26 :param device: (torch.device) specifies the device to run on (cpu/gpu).27 """28 def __init__(self, args, obs_space, action_space, device=torch.device("cpu")):29 super(R_Actor, self).__init__()30 self.hidden_size = args.hidden_size3132 self._gain = args.gain33 self._use_orthogonal = args.use_orthogonal34 self._use_policy_active_masks = args.use_policy_active_masks35 self._use_naive_recurrent_policy = args.use_naive_recurrent_policy36 self._use_recurrent_policy = args.use_recurrent_policy37 self._recurrent_N = args.recurrent_N38 self.tpdv = dict(dtype=torch.float32, device=device)3940 obs_shape = get_shape_from_obs_space(obs_space)41 base = CNNBase if len(obs_shape) == 3 else MLPBase42 self.base = base(args, obs_shape)4344 if self._use_naive_recurrent_policy or self._use_recurrent_policy:45 self.rnn = RNNLayer(self.hidden_size, self.hidden_size, self._recurrent_N, self._use_orthogonal)4647 self.act = ACTLayer(action_space, self.hidden_size, self._use_orthogonal, self._gain)4849 self.to(device)5051 def forward(self, obs, rnn_states, masks, available_actions=None, deterministic=False):52 """53 Compute actions from the given inputs.54 :param obs: (np.ndarray / torch.Tensor) observation inputs into network.55 :param rnn_states: (np.ndarray / torch.Tensor) if RNN network, hidden states for RNN.56 :param masks: (np.ndarray / torch.Tensor) mask tensor denoting if hidden states should be reinitialized to zeros.57 :param available_actions: (np.ndarray / torch.Tensor) denotes which actions are available to agent58 (if None, all actions available)59 :param deterministic: (bool) whether to sample from action distribution or return the mode.6061 :return actions: (torch.Tensor) actions to take.62 :return action_log_probs: (torch.Tensor) log probabilities of taken actions.63 :return rnn_states: (torch.Tensor) updated RNN hidden states.64 """65 obs = check(flatten_obs_dict(obs)).to(**self.tpdv)66 rnn_states = check(rnn_states).to(**self.tpdv)67 masks = check(masks).to(**self.tpdv)68 if available_actions is not None:69 available_actions = check(available_actions).to(**self.tpdv)7071 actor_features = self.base(obs)7273 if self._use_naive_recurrent_policy or self._use_recurrent_policy:74 actor_features, rnn_states = self.rnn(actor_features, rnn_states, masks)7576 actions, action_log_probs = self.act(actor_features, available_actions, deterministic)7778 return actions, action_log_probs, rnn_states7980 def evaluate_actions(self, obs, rnn_states, action, masks, available_actions=None, active_masks=None):81 """82 Compute log probability and entropy of given actions.83 :param obs: (torch.Tensor) observation inputs into network.84 :param action: (torch.Tensor) actions whose entropy and log probability to evaluate.85 :param rnn_states: (torch.Tensor) if RNN network, hidden states for RNN.86 :param masks: (torch.Tensor) mask tensor denoting if hidden states should be reinitialized to zeros.87 :param available_actions: (torch.Tensor) denotes which actions are available to agent88 (if None, all actions available)89 :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.9091 :return action_log_probs: (torch.Tensor) log probabilities of the input actions.92 :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.93 """94 obs = check(flatten_obs_dict(obs)).to(**self.tpdv)95 rnn_states = check(rnn_states).to(**self.tpdv)96 action = check(action).to(**self.tpdv)97 masks = check(masks).to(**self.tpdv)98 if available_actions is not None:99 available_actions = check(available_actions).to(**self.tpdv)100101 if active_masks is not None:102 active_masks = check(active_masks).to(**self.tpdv)103104 actor_features = self.base(obs)105106 if self._use_naive_recurrent_policy or self._use_recurrent_policy:107 actor_features, rnn_states = self.rnn(actor_features, rnn_states, masks)108109 action_log_probs, dist_entropy = self.act.evaluate_actions(actor_features,110 action, available_actions,111 active_masks=112 active_masks if self._use_policy_active_masks113 else None)114115 return action_log_probs, dist_entropy116117118class R_Critic(nn.Module):119 """120 Critic network class for MAPPO. Outputs value function predictions given centralized input (MAPPO) or121 local observations (IPPO).122 :param args: (argparse.Namespace) arguments containing relevant model information.123 :param cent_obs_space: (gym.Space) (centralized) observation space.124 :param device: (torch.device) specifies the device to run on (cpu/gpu).125 """126 def __init__(self, args, cent_obs_space, device=torch.device("cpu")):127 super(R_Critic, self).__init__()128 self.hidden_size = args.hidden_size129 self._use_orthogonal = args.use_orthogonal130 self._use_naive_recurrent_policy = args.use_naive_recurrent_policy131 self._use_recurrent_policy = args.use_recurrent_policy132 self._recurrent_N = args.recurrent_N133 self._use_popart = args.use_popart134 self.tpdv = dict(dtype=torch.float32, device=device)135 init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][self._use_orthogonal]136137 cent_obs_shape = get_shape_from_obs_space(cent_obs_space)138 base = CNNBase if len(cent_obs_shape) == 3 else MLPBase139 self.base = base(args, cent_obs_shape)140141 if self._use_naive_recurrent_policy or self._use_recurrent_policy:142 self.rnn = RNNLayer(self.hidden_size, self.hidden_size, self._recurrent_N, self._use_orthogonal)143144 def init_(m):145 return init(m, init_method, lambda x: nn.init.constant_(x, 0))146147 if self._use_popart:148 self.v_out = init_(PopArt(self.hidden_size, 1, device=device))149 else:150 self.v_out = init_(nn.Linear(self.hidden_size, 1))151152 self.to(device)153154 def forward(self, cent_obs, rnn_states, masks):155 """156 Compute actions from the given inputs.157 :param cent_obs: (np.ndarray / torch.Tensor) observation inputs into network.158 :param rnn_states: (np.ndarray / torch.Tensor) if RNN network, hidden states for RNN.159 :param masks: (np.ndarray / torch.Tensor) mask tensor denoting if RNN states should be reinitialized to zeros.160161 :return values: (torch.Tensor) value function predictions.162 :return rnn_states: (torch.Tensor) updated RNN hidden states.163 """164 cent_obs = check(cent_obs).to(**self.tpdv)165 rnn_states = check(rnn_states).to(**self.tpdv)166 masks = check(masks).to(**self.tpdv)167168 critic_features = self.base(cent_obs)169 if self._use_naive_recurrent_policy or self._use_recurrent_policy:170 critic_features, rnn_states = self.rnn(critic_features, rnn_states, masks)171 values = self.v_out(critic_features)172173 return values, rnn_states