| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| from typing import Tuple, Dict | |
| import torch | |
| import numpy as np | |
| import torch.nn.functional as F | |
| class Pd(object): | |
| """ | |
| Overview: | |
| Abstract class for parameterizable probability distributions and sampling functions. | |
| Interfaces: | |
| ``neglogp``, ``entropy``, ``noise_mode``, ``mode``, ``sample`` | |
| .. tip:: | |
| In dereived classes, `logits` should be an attribute member stored in class. | |
| """ | |
| def neglogp(self, x: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Overview: | |
| Calculate cross_entropy between input x and logits | |
| Arguments: | |
| - x (:obj:`torch.Tensor`): the input tensor | |
| Return: | |
| - cross_entropy (:obj:`torch.Tensor`): the returned cross_entropy loss | |
| """ | |
| raise NotImplementedError | |
| def entropy(self) -> torch.Tensor: | |
| """ | |
| Overview: | |
| Calculate the softmax entropy of logits | |
| Arguments: | |
| - reduction (:obj:`str`): support [None, 'mean'], default set to 'mean' | |
| Returns: | |
| - entropy (:obj:`torch.Tensor`): the calculated entropy | |
| """ | |
| raise NotImplementedError | |
| def noise_mode(self): | |
| """ | |
| Overview: | |
| Add noise to logits. This method is designed for randomness | |
| """ | |
| raise NotImplementedError | |
| def mode(self): | |
| """ | |
| Overview: | |
| Return logits argmax result. This method is designed for deterministic. | |
| """ | |
| raise NotImplementedError | |
| def sample(self): | |
| """ | |
| Overview: | |
| Sample from logits's distribution by using softmax. This method is designed for multinomial. | |
| """ | |
| raise NotImplementedError | |
| class CategoricalPd(Pd): | |
| """ | |
| Overview: | |
| Catagorical probility distribution sampler | |
| Interfaces: | |
| ``__init__``, ``neglogp``, ``entropy``, ``noise_mode``, ``mode``, ``sample`` | |
| """ | |
| def __init__(self, logits: torch.Tensor = None) -> None: | |
| """ | |
| Overview: | |
| Init the Pd with logits | |
| Arguments: | |
| - logits (:obj:torch.Tensor): logits to sample from | |
| """ | |
| self.update_logits(logits) | |
| def update_logits(self, logits: torch.Tensor) -> None: | |
| """ | |
| Overview: | |
| Updata logits | |
| Arguments: | |
| - logits (:obj:`torch.Tensor`): logits to update | |
| """ | |
| self.logits = logits | |
| def neglogp(self, x, reduction: str = 'mean') -> torch.Tensor: | |
| """ | |
| Overview: | |
| Calculate cross_entropy between input x and logits | |
| Arguments: | |
| - x (:obj:`torch.Tensor`): the input tensor | |
| - reduction (:obj:`str`): support [None, 'mean'], default set to mean | |
| Return: | |
| - cross_entropy (:obj:`torch.Tensor`): the returned cross_entropy loss | |
| """ | |
| return F.cross_entropy(self.logits, x, reduction=reduction) | |
| def entropy(self, reduction: str = 'mean') -> torch.Tensor: | |
| """ | |
| Overview: | |
| Calculate the softmax entropy of logits | |
| Arguments: | |
| - reduction (:obj:`str`): support [None, 'mean'], default set to mean | |
| Returns: | |
| - entropy (:obj:`torch.Tensor`): the calculated entropy | |
| """ | |
| a = self.logits - self.logits.max(dim=-1, keepdim=True)[0] | |
| ea = torch.exp(a) | |
| z = ea.sum(dim=-1, keepdim=True) | |
| p = ea / z | |
| entropy = (p * (torch.log(z) - a)).sum(dim=-1) | |
| assert (reduction in [None, 'mean']) | |
| if reduction is None: | |
| return entropy | |
| elif reduction == 'mean': | |
| return entropy.mean() | |
| def noise_mode(self, viz: bool = False) -> Tuple[torch.Tensor, Dict[str, np.ndarray]]: | |
| """ | |
| Overview: | |
| add noise to logits | |
| Arguments: | |
| - viz (:obj:`bool`): Whether to return numpy from of logits, noise and noise_logits; \ | |
| Short for ``visualize`` . (Because tensor type cannot visualize in tb or text log) | |
| Returns: | |
| - result (:obj:`torch.Tensor`): noised logits | |
| - viz_feature (:obj:`Dict[str, np.ndarray]`): ndarray type data for visualization. | |
| """ | |
| u = torch.rand_like(self.logits) | |
| u = -torch.log(-torch.log(u)) | |
| noise_logits = self.logits + u | |
| result = noise_logits.argmax(dim=-1) | |
| if viz: | |
| viz_feature = {} | |
| viz_feature['logits'] = self.logits.data.cpu().numpy() | |
| viz_feature['noise'] = u.data.cpu().numpy() | |
| viz_feature['noise_logits'] = noise_logits.data.cpu().numpy() | |
| return result, viz_feature | |
| else: | |
| return result | |
| def mode(self, viz: bool = False) -> Tuple[torch.Tensor, Dict[str, np.ndarray]]: | |
| """ | |
| Overview: | |
| return logits argmax result | |
| Arguments: | |
| - viz (:obj:`bool`): Whether to return numpy from of logits, noise and noise_logits; | |
| Short for ``visualize`` . (Because tensor type cannot visualize in tb or text log) | |
| Returns: | |
| - result (:obj:`torch.Tensor`): the logits argmax result | |
| - viz_feature (:obj:`Dict[str, np.ndarray]`): ndarray type data for visualization. | |
| """ | |
| result = self.logits.argmax(dim=-1) | |
| if viz: | |
| viz_feature = {} | |
| viz_feature['logits'] = self.logits.data.cpu().numpy() | |
| return result, viz_feature | |
| else: | |
| return result | |
| def sample(self, viz: bool = False) -> Tuple[torch.Tensor, Dict[str, np.ndarray]]: | |
| """ | |
| Overview: | |
| Sample from logits's distribution by using softmax | |
| Arguments: | |
| - viz (:obj:`bool`): Whether to return numpy from of logits, noise and noise_logits; \ | |
| Short for ``visualize`` . (Because tensor type cannot visualize in tb or text log) | |
| Returns: | |
| - result (:obj:`torch.Tensor`): the logits sampled result | |
| - viz_feature (:obj:`Dict[str, np.ndarray]`): ndarray type data for visualization. | |
| """ | |
| p = torch.softmax(self.logits, dim=1) | |
| result = torch.multinomial(p, 1).squeeze(1) | |
| if viz: | |
| viz_feature = {} | |
| viz_feature['logits'] = self.logits.data.cpu().numpy() | |
| return result, viz_feature | |
| else: | |
| return result | |
| class CategoricalPdPytorch(torch.distributions.Categorical): | |
| """ | |
| Overview: | |
| Wrapped ``torch.distributions.Categorical`` | |
| Interfaces: | |
| ``__init__``, ``update_logits``, ``update_probs``, ``sample``, ``neglogp``, ``mode``, ``entropy`` | |
| """ | |
| def __init__(self, probs: torch.Tensor = None) -> None: | |
| """ | |
| Overview: | |
| Initialize the CategoricalPdPytorch object. | |
| Arguments: | |
| - probs (:obj:`torch.Tensor`): The tensor of probabilities. | |
| """ | |
| if probs is not None: | |
| self.update_probs(probs) | |
| def update_logits(self, logits: torch.Tensor) -> None: | |
| """ | |
| Overview: | |
| Updata logits | |
| Arguments: | |
| - logits (:obj:`torch.Tensor`): logits to update | |
| """ | |
| super().__init__(logits=logits) | |
| def update_probs(self, probs: torch.Tensor) -> None: | |
| """ | |
| Overview: | |
| Updata probs | |
| Arguments: | |
| - probs (:obj:`torch.Tensor`): probs to update | |
| """ | |
| super().__init__(probs=probs) | |
| def sample(self) -> torch.Tensor: | |
| """ | |
| Overview: | |
| Sample from logits's distribution by using softmax | |
| Return: | |
| - result (:obj:`torch.Tensor`): the logits sampled result | |
| """ | |
| return super().sample() | |
| def neglogp(self, actions: torch.Tensor, reduction: str = 'mean') -> torch.Tensor: | |
| """ | |
| Overview: | |
| Calculate cross_entropy between input x and logits | |
| Arguments: | |
| - actions (:obj:`torch.Tensor`): the input action tensor | |
| - reduction (:obj:`str`): support [None, 'mean'], default set to mean | |
| Return: | |
| - cross_entropy (:obj:`torch.Tensor`): the returned cross_entropy loss | |
| """ | |
| neglogp = super().log_prob(actions) | |
| assert (reduction in ['none', 'mean']) | |
| if reduction == 'none': | |
| return neglogp | |
| elif reduction == 'mean': | |
| return neglogp.mean(dim=0) | |
| def mode(self) -> torch.Tensor: | |
| """ | |
| Overview: | |
| Return logits argmax result | |
| Return: | |
| - result(:obj:`torch.Tensor`): the logits argmax result | |
| """ | |
| return self.probs.argmax(dim=-1) | |
| def entropy(self, reduction: str = None) -> torch.Tensor: | |
| """ | |
| Overview: | |
| Calculate the softmax entropy of logits | |
| Arguments: | |
| - reduction (:obj:`str`): support [None, 'mean'], default set to mean | |
| Returns: | |
| - entropy (:obj:`torch.Tensor`): the calculated entropy | |
| """ | |
| entropy = super().entropy() | |
| assert (reduction in [None, 'mean']) | |
| if reduction is None: | |
| return entropy | |
| elif reduction == 'mean': | |
| return entropy.mean() | |