eps_greedy_gumbel_softmax.py

eps_greedy_gumbel_softmax.py

import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Gumbel


B = 1000

logits = torch.tensor([np.log(.1), np.log(.2), np.log(.7)]).unsqueeze(0).expand(B, -1)

# Standard sample of gumbel softmax
standard_samples = F.gumbel_softmax(logits, tau=0.01, hard=True)
print(standard_samples.mean(0))

random_noise = Gumbel(torch.zeros((B, 3)), torch.full((B, 3), 10000.)).sample()
eps_samples = F.gumbel_softmax(logits + random_noise, tau=0.01, hard=True)
print(eps_samples.mean(0))