text: import os
  import torch
  import signal
  import torch.nn.functional as F
  import torch.distributed as dist
  import numpy as np
  
  # https://github.com/ikostrikov/pytorch-ddpg-naf/blob/master/ddpg.py#L11
  def soft_update(target, source, tau):
      """
      Perform DDPG soft update (move target params toward source based on weight
      factor tau)
      Inputs:
          target (torch.nn.Module): Net to copy parameters to
          source (torch.nn.Module): Net whose parameters to copy
          tau (float, 0 < x < 1): Weight factor for update
      """
      for target_param, param in zip(target.parameters(), source.parameters()):
          target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)
  
  # https://github.com/ikostrikov/pytorch-ddpg-naf/blob/master/ddpg.py#L15
  def hard_update(target, source):
      """
      Copy network parameters from source to target
      Inputs:
          target (torch.nn.Module): Net to copy parameters to
          source (torch.nn.Module): Net whose parameters to copy
      """
      for target_param, param in zip(target.parameters(), source.parameters()):
          target_param.data.copy_(param.data)
  
  # https://github.com/seba-1511/dist_tuto.pth/blob/gh-pages/train_dist.py
  def average_gradients(model):
      """ Gradient averaging. """
      size = float(dist.get_world_size())
      for param in model.parameters():
          dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM, group=0)
          param.grad.data /= size
  
  # https://github.com/seba-1511/dist_tuto.pth/blob/gh-pages/train_dist.py
  def init_processes(rank, size, fn, backend='gloo'):
      """ Initialize the distributed environment. """
      os.environ['MASTER_ADDR'] = '127.0.0.1'
      os.environ['MASTER_PORT'] = '29500'
      dist.init_process_group(backend, rank=rank, world_size=size)
      fn(rank, size)
  
  def onehot_from_logits(logits, dim=1):
      """
      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
      return (logits == logits.max(dim, keepdim=True)[0]).float()
  
  # modified for PyTorch from https://github.com/ericjang/gumbel-softmax/blob/master/Categorical%20VAE.ipynb
  def sample_gumbel(shape, eps=1e-20, device='cpu'):
      """Sample from Gumbel(0, 1)"""
      U = torch.rand(shape, device=device)
      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, dim=1):
      """ Draw a sample from the Gumbel-Softmax distribution"""
      y = logits + sample_gumbel(logits.shape, device=logits.device)
      return F.softmax(y / temperature, dim=dim)
  
  # modified for PyTorch from https://github.com/ericjang/gumbel-softmax/blob/master/Categorical%20VAE.ipynb
  def gumbel_softmax(logits, temperature=1.0, hard=False, dim=1):
      """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, dim=dim)
      if hard:
          y_hard = onehot_from_logits(y, dim=dim)
          y = (y_hard - y).detach() + y
      return y
  
  def categorical_sample(probs):
      int_acs = torch.multinomial(probs, 1)
      acs = torch.zeros(probs.shape, device=probs.device).scatter_(1, int_acs, 1)
      return int_acs, acs
  
  def disable_gradients(module):
      for p in module.parameters():
          p.requires_grad = False
  
  def enable_gradients(module):
      for p in module.parameters():
          p.requires_grad = True
  
  def sep_clip_grad_norm(parameters, max_norm, norm_type=2):
      """
      Clips gradient norms calculated on a per-parameter basis, rather than over
      the whole list of parameters as in torch.nn.utils.clip_grad_norm.
      Code based on torch.nn.utils.clip_grad_norm
      """
      parameters = list(filter(lambda p: p.grad is not None, parameters))
      max_norm = float(max_norm)
      norm_type = float(norm_type)
      for p in parameters:
          if norm_type == float('inf'):
              p_norm = p.grad.data.abs().max()
          else:
              p_norm = p.grad.data.norm(norm_type)
          clip_coef = max_norm / (p_norm + 1e-6)
          if clip_coef < 1:
              p.grad.data.mul_(clip_coef)
  
  def pol_kl(p_probs, p_log_probs, q_log_probs):
      """
      KL divergence between 2 policies, assuming input probabilities are
      calculated from the same states
      """
      return -(p_probs * (q_log_probs - p_log_probs)).sum(dim=1).mean()
  
  def apply_to_all_elements(x, fn):
      """
      Recursively apply function to all elements of nested lists/tuples
      """
      if type(x) not in (list, tuple):
          return fn(x)
      return [apply_to_all_elements(y, fn) for y in x]
  
  class timeout:
      def __init__(self, seconds=1, error_message='Timeout'):
          self.seconds = seconds
          self.error_message = error_message
      def handle_timeout(self, signum, frame):
          raise TimeoutError(self.error_message)
      def __enter__(self):
          signal.signal(signal.SIGALRM, self.handle_timeout)
          signal.alarm(self.seconds)
      def __exit__(self, type, value, traceback):
          signal.alarm(0)
  
  
  class RunningMeanStd(object):
      # https://github.com/openai/random-network-distillation/blob/master/mpi_util.py
      def __init__(self, epsilon=1e-4, shape=(), gamma=None):
          self.mean = np.zeros(shape, 'float64')
          self.var = np.ones(shape, 'float64')
          self.count = epsilon
          self.gamma = gamma  # keep track of discounted sum (per env) if provided
          self.disc_sum = None
  
      def update(self, x, active_envs=None):
          if self.gamma is not None:
              assert active_envs is not None
              if self.disc_sum is None:
                  self.disc_sum = x
              else:
                  self.disc_sum[active_envs.astype(bool)] = self.gamma * self.disc_sum[active_envs.astype(bool)] + x
                  x = self.disc_sum[active_envs.astype(bool)]
  
          batch_mean, batch_std, batch_count = np.mean(x, axis=0), np.std(x, axis=0), x.shape[0]
          batch_var = np.square(batch_std)
          self.update_from_moments(batch_mean, batch_var, batch_count)
  
      def update_from_moments(self, batch_mean, batch_var, batch_count):
          delta = batch_mean - self.mean
          tot_count = self.count + batch_count
  
          new_mean = self.mean + delta * batch_count / tot_count
          m_a = self.var * (self.count)
          m_b = batch_var * (batch_count)
          M2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count)
          new_var = M2 / (self.count + batch_count)
  
          new_count = batch_count + self.count
  
          self.mean = new_mean
          self.var = new_var
          self.count = new_count
  