text: import numpy as np
  import torch
  import torch.distributed as dist
  import torch.nn.functional as F
  from gym.spaces import Box, Discrete
  from torch.autograd import Variable
  
  
  def onehot_from_logits(logits, eps=0.0):
      """
      Given batch of logits, return one-hot sample using epsilon greedy strategy
      (based on given epsilon)
      """
  
      # get best (according to current policy) actions in one-hot form
      argmax_acs = (logits == logits.max(1, keepdim=True)[0]).float()
  
      # done for IC3Net
      argmax_acs = argmax_acs.double()
  
      if eps == 0.0:
          return argmax_acs
  
      # get random actions in one-hot form
      rand_acs = Variable(torch.eye(logits.shape[1])[[np.random.choice(
          range(logits.shape[1]), size=logits.shape[0])]], requires_grad=False)
      # chooses between best and random actions using epsilon greedy
      return torch.stack([argmax_acs[i] if r > eps else rand_acs[i] for i, r in
                          enumerate(torch.rand(logits.shape[0]))])
  
      # change done for IC3Net
  
  
  # modified for PyTorch from https://github.com/ericjang/gumbel-softmax/blob/master/Categorical%20VAE.ipynb
  def sample_gumbel(shape, eps=1e-20, tens_type=torch.FloatTensor):
      """Sample from Gumbel(0, 1)"""
      U = Variable(tens_type(*shape).uniform_(), requires_grad=False)
      return -torch.log(-torch.log(U + eps) + eps)
  
  # modified for PyTorch from https://github.com/ericjang/gumbel-softmax/blob/master/Categorical%20VAE.ipynb
  def gumbel_softmax_sample(logits, temperature):
      """ Draw a sample from the Gumbel-Softmax distribution"""
      y = logits + sample_gumbel(logits.shape, tens_type=type(logits.data))
      return F.softmax(y / temperature, dim=1)
  
  # modified for PyTorch from https://github.com/ericjang/gumbel-softmax/blob/master/Categorical%20VAE.ipynb
  def gumbel_softmax(logits, temperature=1.0, hard=False):
      """Sample from the Gumbel-Softmax distribution and optionally discretize.
      Args:
        logits: [batch_size, n_class] unnormalized log-probs
        temperature: non-negative scalar
        hard: if True, take argmax, but differentiate w.r.t. soft sample y
      Returns:
        [batch_size, n_class] sample from the Gumbel-Softmax distribution.
        If hard=True, then the returned sample will be one-hot, otherwise it will
        be a probabilitiy distribution that sums to 1 across classes
      """
      y = gumbel_softmax_sample(logits, temperature)
      if hard:
          y_hard = onehot_from_logits(y)
  
          # this was only added to make this work for IC3 Net not fully sure if this the right thing to do
          y_hard = y_hard.double()
  
          y = (y_hard - y).detach() + y
  
      return y
  
  