code: def test_dist():
      """
      test the running mean std
      """
      np.random.seed(0)
      p_1, p_2, p_3 = (
          np.random.randn(3, 1),
          np.random.randn(4, 1),
          np.random.randn(5, 1),
      )
      q_1, q_2, q_3 = (
          np.random.randn(6, 1),
          np.random.randn(7, 1),
          np.random.randn(8, 1),
      )
  
      comm = mpi4py.MPI.COMM_WORLD
      assert comm.Get_size() == 2
      if comm.Get_rank() == 0:
          x_1, x_2, x_3 = p_1, p_2, p_3
      elif comm.Get_rank() == 1:
          x_1, x_2, x_3 = q_1, q_2, q_3
      else:
          assert False
  
      rms = RunningMeanStd(epsilon=0.0, shape=(1,))
      tf_util.initialize()
  
      rms.update(x_1)
      rms.update(x_2)
      rms.update(x_3)
  
      bigvec = np.concatenate([p_1, p_2, p_3, q_1, q_2, q_3])
  
      def checkallclose(var_1, var_2):
          print(var_1, var_2)
          return np.allclose(var_1, var_2)
  
      assert checkallclose(
          bigvec.mean(axis=0),
          rms.mean.eval(),
      )
      assert checkallclose(
          bigvec.std(axis=0),
          rms.std.eval(),
      )
  
file_path: stable-baselines/stable_baselines/common/mpi_running_mean_std.py
imports: import mpi4py
  import tensorflow
  import numpy
  import stable_baselines.common.tf_util
label_desc: Use image node_id 1 for calling a global function with example usage: test_dist() and returns: np