add DeepMimic helper utils

examples/pybullet/gym/pybullet_utils/mpi_util.py

@@ -0,0 +1,52 @@
+import numpy as np
+from mpi4py import MPI
+
+ROOT_PROC_RANK = 0
+
+def get_num_procs():
+    return MPI.COMM_WORLD.Get_size()
+
+def get_proc_rank():
+    return MPI.COMM_WORLD.Get_rank()
+
+def is_root_proc():
+    rank = get_proc_rank()
+    return rank == ROOT_PROC_RANK
+
+def bcast(x):
+    MPI.COMM_WORLD.Bcast(x, root=ROOT_PROC_RANK)
+    return
+
+def reduce_sum(x):
+    return reduce_all(x, MPI.SUM)
+
+def reduce_prod(x):
+    return reduce_all(x, MPI.PROD)
+
+def reduce_avg(x):
+    buffer = reduce_sum(x)
+    buffer /= get_num_procs()
+    return buffer
+
+def reduce_min(x):
+    return reduce_all(x, MPI.MIN)
+
+def reduce_max(x):
+    return reduce_all(x, MPI.MAX)
+
+def reduce_all(x, op):
+    is_array = isinstance(x, np.ndarray)
+    x_buf = x if is_array else np.array([x])
+    buffer = np.zeros_like(x_buf)
+    MPI.COMM_WORLD.Allreduce(x_buf, buffer, op=op)
+    buffer = buffer if is_array else buffer[0]
+    return buffer
+
+def gather_all(x):
+    is_array = isinstance(x, np.ndarray)
+    x_buf = np.array([x])
+    buffer = np.zeros_like(x_buf)
+    buffer = np.repeat(buffer, get_num_procs(), axis=0)
+    MPI.COMM_WORLD.Allgather(x_buf, buffer)
+    buffer = list(buffer)
+    return buffer