distributions.py

动作分布与 log probability

原文件 · 118 行 · 3592 字节

algorithms/utils/distributions.py
SHA-256: 811944a1317bdc40030c927b7c821a0468db4d49ff18cb7c0c5011752e3d986a

1import torch2import torch.nn as nn3from .util import init45"""6Modify standard PyTorch distributions so they to make compatible with this codebase. 7"""89#10# Standardize distribution interfaces11#1213# Categorical14class FixedCategorical(torch.distributions.Categorical):15    def sample(self):16        return super().sample().unsqueeze(-1)1718    def log_probs(self, actions):19        return (20            super()21            .log_prob(actions.squeeze(-1))22            .view(actions.size(0), -1)23            .sum(-1)24            .unsqueeze(-1)25        )2627    def mode(self):28        return self.probs.argmax(dim=-1, keepdim=True)293031# Normal32class FixedNormal(torch.distributions.Normal):33    def log_probs(self, actions):34        return super().log_prob(actions).sum(-1, keepdim=True)3536    def entrop(self):37        return super.entropy().sum(-1)3839    def mode(self):40        return self.mean414243# Bernoulli44class FixedBernoulli(torch.distributions.Bernoulli):45    def log_probs(self, actions):46        return super.log_prob(actions).view(actions.size(0), -1).sum(-1).unsqueeze(-1)4748    def entropy(self):49        return super().entropy().sum(-1)5051    def mode(self):52        return torch.gt(self.probs, 0.5).float()535455class Categorical(nn.Module):56    def __init__(self, num_inputs, num_outputs, use_orthogonal=True, gain=0.01):57        super(Categorical, self).__init__()58        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]59        def init_(m): 60            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain)6162        self.linear = init_(nn.Linear(num_inputs, num_outputs))6364    def forward(self, x, available_actions=None):65        x = self.linear(x)66        if available_actions is not None:67            x[available_actions == 0] = -1e1068        return FixedCategorical(logits=x)697071class DiagGaussian(nn.Module):72    def __init__(self, num_inputs, num_outputs, use_orthogonal=True, gain=0.01):73        super(DiagGaussian, self).__init__()7475        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]76        def init_(m): 77            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain)7879        self.fc_mean = init_(nn.Linear(num_inputs, num_outputs))80        self.logstd = AddBias(torch.zeros(num_outputs))8182    def forward(self, x):83        action_mean = self.fc_mean(x)8485        #  An ugly hack for my KFAC implementation.86        zeros = torch.zeros(action_mean.size())87        if x.is_cuda:88            zeros = zeros.cuda()8990        action_logstd = self.logstd(zeros)91        return FixedNormal(action_mean, action_logstd.exp())929394class Bernoulli(nn.Module):95    def __init__(self, num_inputs, num_outputs, use_orthogonal=True, gain=0.01):96        super(Bernoulli, self).__init__()97        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]98        def init_(m): 99            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain)100        101        self.linear = init_(nn.Linear(num_inputs, num_outputs))102103    def forward(self, x):104        x = self.linear(x)105        return FixedBernoulli(logits=x)106107class AddBias(nn.Module):108    def __init__(self, bias):109        super(AddBias, self).__init__()110        self._bias = nn.Parameter(bias.unsqueeze(1))111112    def forward(self, x):113        if x.dim() == 2:114            bias = self._bias.t().view(1, -1)115        else:116            bias = self._bias.t().view(1, -1, 1, 1)117118        return x + bias