from typing import Optional, Tuple, Type import torch import torch.nn as nn from torch.distributions import Categorical from rl_algo_impls.shared.actor import Actor, PiForward from rl_algo_impls.shared.module.module import mlp class MaskedCategorical(Categorical): def __init__( self, probs=None, logits=None, validate_args=None, mask: Optional[torch.Tensor] = None, ): if mask is not None: assert logits is not None, "mask requires logits and not probs" logits = torch.where(mask, logits, -1e8) self.mask = mask super().__init__(probs, logits, validate_args) def entropy(self) -> torch.Tensor: if self.mask is None: return super().entropy() # If mask set, then use approximation for entropy p_log_p = self.logits * self.probs # type: ignore masked = torch.where(self.mask, p_log_p, 0) return -masked.sum(-1) class CategoricalActorHead(Actor): def __init__( self, act_dim: int, in_dim: int, hidden_sizes: Tuple[int, ...] = (32,), activation: Type[nn.Module] = nn.Tanh, init_layers_orthogonal: bool = True, ) -> None: super().__init__() layer_sizes = (in_dim,) + hidden_sizes + (act_dim,) self._fc = mlp( layer_sizes, activation, init_layers_orthogonal=init_layers_orthogonal, final_layer_gain=0.01, ) def forward( self, obs: torch.Tensor, actions: Optional[torch.Tensor] = None, action_masks: Optional[torch.Tensor] = None, ) -> PiForward: logits = self._fc(obs) pi = MaskedCategorical(logits=logits, mask=action_masks) return self.pi_forward(pi, actions) @property def action_shape(self) -> Tuple[int, ...]: return ()