updates

alr_envs/classic_control/hole_reacher.py

@@ -1,7 +1,5 @@
 import gym
 import numpy as np
-import matplotlib
-matplotlib.use('TkAgg')
 import matplotlib.pyplot as plt
 from matplotlib import patches
 
@@ -112,7 +110,7 @@ class HoleReacher(gym.Env):
         if self._is_collided:
             reward -= self.collision_penalty
 
-        info = {}
+        info = {"is_collided": self._is_collided}
 
         self._steps += 1
 
@@ -286,6 +284,10 @@ class HoleReacher(gym.Env):
 
                 plt.pause(0.01)
 
+    def close(self):
+        if self.fig is not None:
+            plt.close(self.fig)
+
 
 if __name__ == '__main__':
     nl = 5
@@ -306,3 +308,5 @@ if __name__ == '__main__':
 
         if done:
             break
+
+    env.close()

alr_envs/classic_control/utils.py

@@ -0,0 +1,34 @@
+from alr_envs.classic_control.hole_reacher import HoleReacher
+from alr_envs.utils.dmp_env_wrapper import DmpEnvWrapperVel
+
+
+def make_env(rank, seed=0):
+    """
+    Utility function for multiprocessed env.
+
+    :param env_id: (str) the environment ID
+    :param num_env: (int) the number of environments you wish to have in subprocesses
+    :param seed: (int) the initial seed for RNG
+    :param rank: (int) index of the subprocess
+    :returns a function that generates an environment
+    """
+
+    def _init():
+        env = HoleReacher(num_links=5,
+                          allow_self_collision=False,
+                          allow_wall_collision=False,
+                          hole_width=0.15,
+                          hole_depth=1,
+                          hole_x=1,
+                          collision_penalty=100000)
+
+        env = DmpEnvWrapperVel(env,
+                               num_dof=5,
+                               num_basis=5,
+                               duration=2,
+                               dt=env._dt,
+                               learn_goal=True)
+        env.seed(seed + rank)
+        return env
+
+    return _init