act.py
离散策略动作层
原文件 · 194 行 · 9725 字节
algorithms/utils/act.py
SHA-256: 82a45f6ed7338619ab9695d30934e9c92f7aba3b432de921db0184af0f8d2a4e
1from .distributions import Bernoulli, Categorical, DiagGaussian2import torch3import torch.nn as nn45class ACTLayer(nn.Module):6 """7 MLP Module to compute actions.8 :param action_space: (gym.Space) action space.9 :param inputs_dim: (int) dimension of network input.10 :param use_orthogonal: (bool) whether to use orthogonal initialization.11 :param gain: (float) gain of the output layer of the network.12 """13 def __init__(self, action_space, inputs_dim, use_orthogonal, gain):14 super(ACTLayer, self).__init__()15 self.mixed_action = False16 self.multi_discrete = False17 self.continuous_action = False1819 if action_space.__class__.__name__ == "Discrete":20 action_dim = action_space.n21 self.action_out = Categorical(inputs_dim, action_dim, use_orthogonal, gain)22 elif action_space.__class__.__name__ == "Box":23 self.continuous_action = True24 action_dim = action_space.shape[0]25 self.action_out = DiagGaussian(inputs_dim, action_dim, use_orthogonal, gain)26 elif action_space.__class__.__name__ == "MultiBinary":27 action_dim = action_space.shape[0]28 self.action_out = Bernoulli(inputs_dim, action_dim, use_orthogonal, gain)29 elif action_space.__class__.__name__ == "MultiDiscrete":30 self.multi_discrete = True31 action_dims = action_space.high - action_space.low + 132 self.action_outs = []33 for action_dim in action_dims:34 self.action_outs.append(Categorical(inputs_dim, action_dim, use_orthogonal, gain))35 self.action_outs = nn.ModuleList(self.action_outs)36 else: # discrete + continous37 self.mixed_action = True38 continous_dim = action_space[0].shape[0]39 discrete_dim = action_space[1].n40 self.action_outs = nn.ModuleList([DiagGaussian(inputs_dim, continous_dim, use_orthogonal, gain), Categorical(41 inputs_dim, discrete_dim, use_orthogonal, gain)])42 43 def forward(self, x, available_actions=None, deterministic=False):44 """45 Compute actions and action logprobs from given input.46 :param x: (torch.Tensor) input to network.47 :param available_actions: (torch.Tensor) denotes which actions are available to agent48 (if None, all actions available)49 :param deterministic: (bool) whether to sample from action distribution or return the mode.5051 :return actions: (torch.Tensor) actions to take.52 :return action_log_probs: (torch.Tensor) log probabilities of taken actions.53 """54 if self.mixed_action:55 actions = []56 action_log_probs = []57 for action_out in self.action_outs:58 action_logit = action_out(x)59 action = action_logit.mode() if deterministic else action_logit.sample()60 action_log_prob = action_logit.log_probs(action)61 actions.append(action.float())62 action_log_probs.append(action_log_prob)6364 actions = torch.cat(actions, -1)65 action_log_probs = torch.sum(torch.cat(action_log_probs, -1), -1, keepdim=True)6667 elif self.multi_discrete:68 actions = []69 action_log_probs = []70 for action_out in self.action_outs:71 action_logit = action_out(x)72 action = action_logit.mode() if deterministic else action_logit.sample()73 action_log_prob = action_logit.log_probs(action)74 actions.append(action)75 action_log_probs.append(action_log_prob)7677 actions = torch.cat(actions, -1)78 action_log_probs = torch.cat(action_log_probs, -1)79 elif self.continuous_action:80 # actions = []81 # action_log_probs = []82 action_logit = self.action_out(x)83 actions = action_logit.mode() if deterministic else action_logit.sample()84 action_log_probs = action_logit.log_probs(actions)85 # actions.append(action.float())86 # action_log_probs.append(action_log_prob)87 # actions = torch.cat(actions, -1)88 # action_log_probs = torch.sum(torch.cat(action_log_probs, -1), -1, keepdim=True)89 else:90 action_logits = self.action_out(x, available_actions)91 actions = action_logits.mode() if deterministic else action_logits.sample() 92 action_log_probs = action_logits.log_probs(actions)93 94 return actions, action_log_probs9596 def get_probs(self, x, available_actions=None):97 """98 Compute action probabilities from inputs.99 :param x: (torch.Tensor) input to network.100 :param available_actions: (torch.Tensor) denotes which actions are available to agent101 (if None, all actions available)102103 :return action_probs: (torch.Tensor)104 """105 if self.mixed_action or self.multi_discrete:106 action_probs = []107 for action_out in self.action_outs:108 action_logit = action_out(x)109 action_prob = action_logit.probs110 action_probs.append(action_prob)111 action_probs = torch.cat(action_probs, -1)112 else:113 action_logits = self.action_out(x, available_actions)114 action_probs = action_logits.probs115 116 return action_probs117118 def evaluate_actions(self, x, action, available_actions=None, active_masks=None):119 """120 Compute log probability and entropy of given actions.121 :param x: (torch.Tensor) input to network.122 :param action: (torch.Tensor) actions whose entropy and log probability to evaluate.123 :param available_actions: (torch.Tensor) denotes which actions are available to agent124 (if None, all actions available)125 :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.126127 :return action_log_probs: (torch.Tensor) log probabilities of the input actions.128 :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.129 """130 if self.mixed_action:131 a, b = action.split((2, 1), -1)132 b = b.long()133 action = [a, b] 134 action_log_probs = [] 135 dist_entropy = []136 for action_out, act in zip(self.action_outs, action):137 action_logit = action_out(x)138 action_log_probs.append(action_logit.log_probs(act))139 if active_masks is not None:140 if len(action_logit.entropy().shape) == len(active_masks.shape):141 dist_entropy.append((action_logit.entropy() * active_masks).sum()/active_masks.sum()) 142 else:143 dist_entropy.append((action_logit.entropy() * active_masks.squeeze(-1)).sum()/active_masks.sum())144 else:145 dist_entropy.append(action_logit.entropy().mean())146 147 action_log_probs = torch.sum(torch.cat(action_log_probs, -1), -1, keepdim=True)148 dist_entropy = dist_entropy[0] / 2.0 + dist_entropy[1] / 0.98 #! dosen't make sense149150 elif self.multi_discrete:151 action = torch.transpose(action, 0, 1)152 action_log_probs = []153 dist_entropy = []154 for action_out, act in zip(self.action_outs, action):155 action_logit = action_out(x)156 action_log_probs.append(action_logit.log_probs(act))157 if active_masks is not None:158 dist_entropy.append((action_logit.entropy()*active_masks.squeeze(-1)).sum()/active_masks.sum())159 else:160 dist_entropy.append(action_logit.entropy().mean())161162 action_log_probs = torch.cat(action_log_probs, -1) # ! could be wrong163 dist_entropy = torch.tensor(dist_entropy).mean()164165 elif self.continuous_action:166 # a, b = action.split((2, 1), -1)167 # b = b.long()168 # action = [a, b]169 action_log_probs = []170 dist_entropy = []171 # for action_out, act in zip(self.action_outs, action):172 action_logit = self.action_out(x)173 action_log_probs.append(action_logit.log_probs(action))174 if active_masks is not None:175 if len(action_logit.entropy().shape) == len(active_masks.shape):176 dist_entropy.append((action_logit.entropy() * active_masks).sum() / active_masks.sum())177 else:178 dist_entropy.append(179 (action_logit.entropy() * active_masks.squeeze(-1)).sum() / active_masks.sum())180 else:181 dist_entropy.append(action_logit.entropy().mean())182183 action_log_probs = torch.sum(torch.cat(action_log_probs, -1), -1, keepdim=True)184 dist_entropy = dist_entropy[0] # / 2.0 + dist_entropy[1] / 0.98 # ! dosen't make sense185186 else:187 action_logits = self.action_out(x, available_actions)188 action_log_probs = action_logits.log_probs(action)189 if active_masks is not None:190 dist_entropy = (action_logits.entropy()*active_masks.squeeze(-1)).sum()/active_masks.sum()191 else:192 dist_entropy = action_logits.entropy().mean()193 194 return action_log_probs, dist_entropy