text: 
  from mpi4py import MPI
  import multiprocessing
  import gym
  import numpy as np
  import cv2
  from copy import deepcopy
  
  from rlpyt.utils.misc import wrap_print
  
  # class RunningMeanStd(object):
  #     # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
  #     def __init__(self, epsilon=1e-4, shape=()):
  #         self.mean = np.zeros(shape, 'float64')
  #         self.var = np.ones(shape, 'float64')
  #         self.count = epsilon
  
  #     def update(self, x):
  #         batch_mean = np.mean(x, axis=0)
  #         batch_var = np.var(x, axis=0)
  #         batch_count = x.shape[0]
  #         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
  
  class RunningMeanStd(object):
      # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
      def __init__(self, epsilon=1e-4, shape=()):
          self.mean = np.zeros(shape, 'float64')
          self.var = np.ones(shape, 'float64')
          self.count = epsilon
  
      def update(self, x):
          batch_mean = np.mean(x, axis=0)
          batch_var = np.var(x, axis=0)
          batch_count = x.shape[0]
          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
  
  class RewardForwardFilter(object):
      def __init__(self, gamma):
          self.rewems = None
          self.gamma = gamma
  
      def update(self, rews, done=None):
          if done is None:
              if self.rewems is None:
                  self.rewems = rews
              else:
                  self.rewems = self.rewems * self.gamma + rews
              return self.rewems
  
          else:
              if self.rewems is None:
                  self.rewems = rews
              else:
                  mask = np.where(done == 1.0)
                  self.rewems[mask] = self.rewems[mask] * self.gamma + rews[mask]
              return deepcopy(self.rewems)
  
  def generate_observation_stats(env, nsteps=10000):
      '''
      Steps through the environment randomly and produces an observation mean and standard deviation. 
      From https://github.com/openai/large-scale-curiosity/blob/master/utils.py
      '''
      wrap_print('Generating observation mean/std... ({} random steps)'.format(nsteps))
      ob = np.asarray(env.reset())
      obs = [ob]
      for _ in range(nsteps):
          ac = env.action_space.sample()
          ob, _, done, _ = env.step(ac.item(0)) # mainly used for mario with discrete action space (hacky .item() call to make it work)
          if done:
              ob = env.reset()
          obs.append(np.asarray(ob))
      mean = np.mean(obs, 0).astype(np.float32)
      std = np.std(obs, 0).mean().astype(np.float32)
      return mean, std
  
  
  
  