refractoring of DMP environmets to fit gym interface better.

alr_envs/classic_control/__init__.py

@@ -1 +1,3 @@
 from alr_envs.classic_control.simple_reacher import SimpleReacherEnv
+from alr_envs.classic_control.viapoint_reacher import ViaPointReacher
+from alr_envs.classic_control.hole_reacher import HoleReacher

alr_envs/classic_control/hole_reacher.py

@@ -3,9 +3,12 @@ import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib import patches
 
+from alr_envs import DmpWrapper
+from alr_envs.utils.wrapper.detpmp_wrapper import DetPMPWrapper
+
 
 def ccw(A, B, C):
-    return (C[1]-A[1]) * (B[0]-A[0]) - (B[1]-A[1]) * (C[0]-A[0]) > 1e-12
+    return (C[1] - A[1]) * (B[0] - A[0]) - (B[1] - A[1]) * (C[0] - A[0]) > 1e-12
 
 
 # Return true if line segments AB and CD intersect
@@ -13,37 +16,66 @@ def intersect(A, B, C, D):
     return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
 
 
+def holereacher_dmp(**kwargs):
+    _env = gym.make("alr_envs:HoleReacher-v0")
+    # _env = HoleReacher(**kwargs)
+    return DmpWrapper(_env, num_dof=5, num_basis=5, duration=2, dt=_env.dt, learn_goal=True, alpha_phase=3.5,
+                      start_pos=_env.start_pos, policy_type="velocity", weights_scale=100, goal_scale=0.1)
+
+
+def holereacher_fix_goal_dmp(**kwargs):
+    _env = gym.make("alr_envs:HoleReacher-v0")
+    # _env = HoleReacher(**kwargs)
+    return DmpWrapper(_env, num_dof=5, num_basis=5, duration=2, dt=_env.dt, learn_goal=False, alpha_phase=3.5,
+                      start_pos=_env.start_pos, policy_type="velocity", weights_scale=50, goal_scale=1,
+                      final_pos=np.array([2.02669572, -1.25966385, -1.51618198, -0.80946476, 0.02012344]))
+
+
+def holereacher_detpmp(**kwargs):
+    _env = gym.make("alr_envs:HoleReacher-v0")
+    # _env = HoleReacher(**kwargs)
+    return DetPMPWrapper(_env, num_dof=5, num_basis=5, width=0.005, policy_type="velocity", start_pos=_env.start_pos,
+                         duration=2, post_traj_time=0, dt=_env.dt, weights_scale=0.25, zero_start=True, zero_goal=False)
+
+
 class HoleReacher(gym.Env):
 
-    def __init__(self, num_links, hole_x, hole_width, hole_depth, allow_self_collision=False,
+    def __init__(self, n_links, hole_x, hole_width, hole_depth, allow_self_collision=False,
                  allow_wall_collision=False, collision_penalty=1000):
+
+        self.n_links = n_links
+        self.link_lengths = np.ones((n_links, 1))
+
+        # task
         self.hole_x = hole_x  # x-position of center of hole
         self.hole_width = hole_width  # width of hole
         self.hole_depth = hole_depth  # depth of hole
-        self.num_links = num_links
-        self.link_lengths = np.ones((num_links, 1))
+
         self.bottom_center_of_hole = np.hstack([hole_x, -hole_depth])
         self.top_center_of_hole = np.hstack([hole_x, 0])
-        self.left_wall_edge = np.hstack([hole_x - self.hole_width/2, 0])
+        self.left_wall_edge = np.hstack([hole_x - self.hole_width / 2, 0])
         self.right_wall_edge = np.hstack([hole_x + self.hole_width / 2, 0])
+
+        # collision
         self.allow_self_collision = allow_self_collision
         self.allow_wall_collision = allow_wall_collision
         self.collision_penalty = collision_penalty
 
+        # state
         self._joints = None
         self._joint_angles = None
         self._angle_velocity = None
-        self.start_pos = np.hstack([[np.pi/2], np.zeros(self.num_links - 1)])
-        self.start_vel = np.zeros(self.num_links)
+        self.start_pos = np.hstack([[np.pi / 2], np.zeros(self.n_links - 1)])
+        self.start_vel = np.zeros(self.n_links)
 
         self.dt = 0.01
-        self.time_limit = 2
+        # self.time_limit = 2
 
-        action_bound = np.pi * np.ones((self.num_links,))
+        action_bound = np.pi * np.ones((self.n_links,))
         state_bound = np.hstack([
-            [np.pi] * self.num_links,  # cos
-            [np.pi] * self.num_links,  # sin
-            [np.inf] * self.num_links,  # velocity
+            [np.pi] * self.n_links,  # cos
+            [np.pi] * self.n_links,  # sin
+            [np.inf] * self.n_links,  # velocity
             [np.inf] * 2,  # x-y coordinates of target distance
             [np.inf]  # env steps, because reward start after n steps TODO: Maybe
         ])
@@ -51,11 +83,11 @@ class HoleReacher(gym.Env):
         self.observation_space = gym.spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape)
 
         self.fig = None
-        rect_1 = patches.Rectangle((-self.num_links, -1),
-                                   self.num_links + self.hole_x - self.hole_width / 2, 1,
+        rect_1 = patches.Rectangle((-self.n_links, -1),
+                                   self.n_links + self.hole_x - self.hole_width / 2, 1,
                                    fill=True, edgecolor='k', facecolor='k')
         rect_2 = patches.Rectangle((self.hole_x + self.hole_width / 2, -1),
-                                   self.num_links - self.hole_x + self.hole_width / 2, 1,
+                                   self.n_links - self.hole_x + self.hole_width / 2, 1,
                                    fill=True, edgecolor='k', facecolor='k')
         rect_3 = patches.Rectangle((self.hole_x - self.hole_width / 2, -1), self.hole_width,
                                    1 - self.hole_depth,
@@ -65,7 +97,7 @@ class HoleReacher(gym.Env):
 
     @property
     def end_effector(self):
-        return self._joints[self.num_links].T
+        return self._joints[self.n_links].T
 
     def configure(self, context):
         pass
@@ -73,13 +105,13 @@ class HoleReacher(gym.Env):
     def reset(self):
         self._joint_angles = self.start_pos
         self._angle_velocity = self.start_vel
-        self._joints = np.zeros((self.num_links + 1, 2))
+        self._joints = np.zeros((self.n_links + 1, 2))
         self._update_joints()
         self._steps = 0
 
         return self._get_obs().copy()
 
-    def step(self, action):
+    def step(self, action: np.ndarray):
         """
         a single step with an action in joint velocity space
         """
@@ -90,16 +122,12 @@ class HoleReacher(gym.Env):
 
         self._update_joints()
 
-        # rew = self._reward()
-
-        # compute reward directly in step function
-
         reward = 0
-        if not self._is_collided:
+        if self._is_collided:
+            reward = -self.collision_penalty
+        else:
             if self._steps == 199:
                 reward = - np.linalg.norm(self.end_effector - self.bottom_center_of_hole) ** 2
-        else:
-            reward = -self.collision_penalty
 
         reward -= 5e-8 * np.sum(acc ** 2)
 
@@ -107,7 +135,8 @@ class HoleReacher(gym.Env):
 
         self._steps += 1
 
-        done = self._steps * self.dt > self.time_limit or self._is_collided
+        # done = self._steps * self.dt > self.time_limit or self._is_collided
+        done = self._is_collided
 
         return self._get_obs().copy(), reward, done, info
 
@@ -145,18 +174,6 @@ class HoleReacher(gym.Env):
             self._steps
         ])
 
-    # def _reward(self):
-    #     dist_reward = 0
-    #     if not self._is_collided:
-    #         if self._steps == 180:
-    #             dist_reward = np.linalg.norm(self.end_effector - self.bottom_center_of_hole)
-    #     else:
-    #         dist_reward = np.linalg.norm(self.end_effector - self.bottom_center_of_hole)
-    #
-    #     out = - dist_reward ** 2
-    #
-    #     return out
-
     def get_forward_kinematics(self, num_points_per_link=1):
         theta = self._joint_angles[:, None]
 
@@ -167,7 +184,7 @@ class HoleReacher(gym.Env):
 
         accumulated_theta = np.cumsum(theta, axis=0)
 
-        endeffector = np.zeros(shape=(self.num_links, num_points_per_link, 2))
+        endeffector = np.zeros(shape=(self.n_links, num_points_per_link, 2))
 
         x = np.cos(accumulated_theta) * self.link_lengths * intermediate_points
         y = np.sin(accumulated_theta) * self.link_lengths * intermediate_points
@@ -175,7 +192,7 @@ class HoleReacher(gym.Env):
         endeffector[0, :, 0] = x[0, :]
         endeffector[0, :, 1] = y[0, :]
 
-        for i in range(1, self.num_links):
+        for i in range(1, self.n_links):
             endeffector[i, :, 0] = x[i, :] + endeffector[i - 1, -1, 0]
             endeffector[i, :, 1] = y[i, :] + endeffector[i - 1, -1, 1]
 
@@ -183,7 +200,7 @@ class HoleReacher(gym.Env):
 
     def check_self_collision(self, line_points):
         for i, line1 in enumerate(line_points):
-            for line2 in line_points[i+2:, :, :]:
+            for line2 in line_points[i + 2:, :, :]:
                 # if line1 != line2:
                 if intersect(line1[0], line1[-1], line2[0], line2[-1]):
                     return True
@@ -211,7 +228,7 @@ class HoleReacher(gym.Env):
 
         # all points that are above the hole
         r, c = np.where((line_points[:, :, 0] > (self.hole_x - self.hole_width / 2)) & (
-                           line_points[:, :, 0] < (self.hole_x + self.hole_width / 2)))
+                line_points[:, :, 0] < (self.hole_x + self.hole_width / 2)))
 
         # check if any of those points are below surface
         nr_line_points_below_surface_in_hole = np.sum(line_points[r, c, 1] < -self.hole_depth)
@@ -243,7 +260,7 @@ class HoleReacher(gym.Env):
             plt.xlim([-lim, lim])
             plt.ylim([-1.1, lim])
             # plt.draw()
-            plt.pause(1e-4) #  pushes window to foreground, which is annoying.
+            plt.pause(1e-4)  # pushes window to foreground, which is annoying.
             # self.fig.canvas.flush_events()
 
         elif mode == "partial":
@@ -273,7 +290,7 @@ class HoleReacher(gym.Env):
                 # Add the patch to the Axes
                 [plt.gca().add_patch(rect) for rect in self.patches]
 
-                plt.xlim(-self.num_links, self.num_links), plt.ylim(-1, self.num_links)
+                plt.xlim(-self.n_links, self.n_links), plt.ylim(-1, self.n_links)
                 # Arm
                 plt.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k')
 
@@ -287,7 +304,8 @@ class HoleReacher(gym.Env):
 if __name__ == '__main__':
     nl = 5
     render_mode = "human"  # "human" or "partial" or "final"
-    env = HoleReacher(num_links=nl, allow_self_collision=False, allow_wall_collision=False, hole_width=0.15, hole_depth=1, hole_x=1)
+    env = HoleReacher(n_links=nl, allow_self_collision=False, allow_wall_collision=False, hole_width=0.15,
+                      hole_depth=1, hole_x=1)
     env.reset()
     # env.render(mode=render_mode)
 

alr_envs/classic_control/simple_reacher.py

@@ -1,11 +1,9 @@
-import os
-
 import gym
-import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 from gym import spaces
 from gym.utils import seeding
+
 from alr_envs.utils.utils import angle_normalize
 
 
@@ -33,7 +31,7 @@ class SimpleReacherEnv(gym.Env):
         self._angle_velocity = None
 
         self.max_torque = 1  # 10
-        self.steps_before_reward = 180
+        self.steps_before_reward = 199
 
         action_bound = np.ones((self.n_links,))
         state_bound = np.hstack([
@@ -92,7 +90,7 @@ class SimpleReacherEnv(gym.Env):
 
     def _update_joints(self):
         """
-        update _joints to get new end effector position. The other links are only required for rendering.
+        update joints to get new end-effector position. The other links are only required for rendering.
         Returns:
 
         """
@@ -106,7 +104,7 @@ class SimpleReacherEnv(gym.Env):
 
         # TODO: Is this the best option
         if self._steps >= self.steps_before_reward:
-            reward_dist = - np.linalg.norm(diff)
+            reward_dist -= np.linalg.norm(diff)
             # reward_dist = np.exp(-0.1 * diff ** 2).mean()
             # reward_dist = - (diff ** 2).mean()
 

alr_envs/classic_control/utils.py

@@ -1,7 +1,7 @@
 from alr_envs.classic_control.hole_reacher import HoleReacher
 from alr_envs.classic_control.viapoint_reacher import ViaPointReacher
-from alr_envs.utils.dmp_env_wrapper import DmpEnvWrapper
-from alr_envs.utils.detpmp_env_wrapper import DetPMPEnvWrapper
+from alr_envs.utils.wrapper.dmp_wrapper import DmpWrapper
+from alr_envs.utils.wrapper.detpmp_wrapper import DetPMPWrapper
 import numpy as np
 
 
@@ -17,20 +17,20 @@ def make_viapointreacher_env(rank, seed=0):
     """
 
     def _init():
-        _env = ViaPointReacher(num_links=5,
+        _env = ViaPointReacher(n_links=5,
                                allow_self_collision=False,
                                collision_penalty=1000)
 
-        _env = DmpEnvWrapper(_env,
-                             num_dof=5,
-                             num_basis=5,
-                             duration=2,
-                             alpha_phase=2.5,
-                             dt=_env.dt,
-                             start_pos=_env.start_pos,
-                             learn_goal=False,
-                             policy_type="velocity",
-                             weights_scale=50)
+        _env = DmpWrapper(_env,
+                          num_dof=5,
+                          num_basis=5,
+                          duration=2,
+                          alpha_phase=2.5,
+                          dt=_env.dt,
+                          start_pos=_env.start_pos,
+                          learn_goal=False,
+                          policy_type="velocity",
+                          weights_scale=50)
         _env.seed(seed + rank)
         return _env
 
@@ -49,7 +49,7 @@ def make_holereacher_env(rank, seed=0):
     """
 
     def _init():
-        _env = HoleReacher(num_links=5,
+        _env = HoleReacher(n_links=5,
                            allow_self_collision=False,
                            allow_wall_collision=False,
                            hole_width=0.15,
@@ -57,18 +57,18 @@ def make_holereacher_env(rank, seed=0):
                            hole_x=1,
                            collision_penalty=100)
 
-        _env = DmpEnvWrapper(_env,
-                             num_dof=5,
-                             num_basis=5,
-                             duration=2,
-                             dt=_env.dt,
-                             learn_goal=True,
-                             alpha_phase=3.5,
-                             start_pos=_env.start_pos,
-                             policy_type="velocity",
-                             weights_scale=100,
-                             goal_scale=0.1
-                             )
+        _env = DmpWrapper(_env,
+                          num_dof=5,
+                          num_basis=5,
+                          duration=2,
+                          dt=_env.dt,
+                          learn_goal=True,
+                          alpha_phase=3.5,
+                          start_pos=_env.start_pos,
+                          policy_type="velocity",
+                          weights_scale=100,
+                          goal_scale=0.1
+                          )
 
         _env.seed(seed + rank)
         return _env
@@ -88,7 +88,7 @@ def make_holereacher_fix_goal_env(rank, seed=0):
     """
 
     def _init():
-        _env = HoleReacher(num_links=5,
+        _env = HoleReacher(n_links=5,
                            allow_self_collision=False,
                            allow_wall_collision=False,
                            hole_width=0.15,
@@ -96,19 +96,19 @@ def make_holereacher_fix_goal_env(rank, seed=0):
                            hole_x=1,
                            collision_penalty=100)
 
-        _env = DmpEnvWrapper(_env,
-                             num_dof=5,
-                             num_basis=5,
-                             duration=2,
-                             dt=_env.dt,
-                             learn_goal=False,
-                             final_pos=np.array([2.02669572, -1.25966385, -1.51618198, -0.80946476,  0.02012344]),
-                             alpha_phase=3.5,
-                             start_pos=_env.start_pos,
-                             policy_type="velocity",
-                             weights_scale=50,
-                             goal_scale=1
-                             )
+        _env = DmpWrapper(_env,
+                          num_dof=5,
+                          num_basis=5,
+                          duration=2,
+                          dt=_env.dt,
+                          learn_goal=False,
+                          final_pos=np.array([2.02669572, -1.25966385, -1.51618198, -0.80946476,  0.02012344]),
+                          alpha_phase=3.5,
+                          start_pos=_env.start_pos,
+                          policy_type="velocity",
+                          weights_scale=50,
+                          goal_scale=1
+                          )
 
         _env.seed(seed + rank)
         return _env
@@ -128,7 +128,7 @@ def make_holereacher_env_pmp(rank, seed=0):
     """
 
     def _init():
-        _env = HoleReacher(num_links=5,
+        _env = HoleReacher(n_links=5,
                            allow_self_collision=False,
                            allow_wall_collision=False,
                            hole_width=0.15,
@@ -136,19 +136,19 @@ def make_holereacher_env_pmp(rank, seed=0):
                            hole_x=1,
                            collision_penalty=1000)
 
-        _env = DetPMPEnvWrapper(_env,
-                                num_dof=5,
-                                num_basis=5,
-                                width=0.005,
-                                policy_type="velocity",
-                                start_pos=_env.start_pos,
-                                duration=2,
-                                post_traj_time=0,
-                                dt=_env.dt,
-                                weights_scale=0.25,
-                                zero_start=True,
-                                zero_goal=False
-                                )
+        _env = DetPMPWrapper(_env,
+                             num_dof=5,
+                             num_basis=5,
+                             width=0.005,
+                             policy_type="velocity",
+                             start_pos=_env.start_pos,
+                             duration=2,
+                             post_traj_time=0,
+                             dt=_env.dt,
+                             weights_scale=0.25,
+                             zero_start=True,
+                             zero_goal=False
+                             )
         _env.seed(seed + rank)
         return _env
 

alr_envs/classic_control/viapoint_reacher.py

@@ -1,39 +1,43 @@
 import gym
-import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib import patches
+import numpy as np
+
+from alr_envs import DmpWrapper
+from alr_envs.utils.utils import check_self_collision
 
 
-def ccw(A, B, C):
-    return (C[1]-A[1]) * (B[0]-A[0]) - (B[1]-A[1]) * (C[0]-A[0]) > 1e-12
-
-
-# Return true if line segments AB and CD intersect
-def intersect(A, B, C, D):
-    return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
+def viapoint_dmp(**kwargs):
+    _env = gym.make("alr_envs:ViaPointReacher-v0")
+    # _env = ViaPointReacher(**kwargs)
+    return DmpWrapper(_env, num_dof=5, num_basis=5, duration=2, alpha_phase=2.5, dt=_env.dt,
+                      start_pos=_env.start_pos, learn_goal=False, policy_type="velocity", weights_scale=50)
 
 
 class ViaPointReacher(gym.Env):
 
-    def __init__(self, num_links, allow_self_collision=False,
-                 collision_penalty=1000):
-        self.num_links = num_links
-        self.link_lengths = np.ones((num_links, 1))
+    def __init__(self, n_links, allow_self_collision=False, collision_penalty=1000):
+        self.num_links = n_links
+        self.link_lengths = np.ones((n_links, 1))
+
+        # task
+        self.via_point = np.ones(2)
+        self.goal_point = np.array((n_links, 0))
+
+        # collision
         self.allow_self_collision = allow_self_collision
         self.collision_penalty = collision_penalty
 
-        self.via_point = np.ones(2)
-        self.goal_point = np.array((num_links, 0))
-
+        # state
         self._joints = None
         self._joint_angles = None
         self._angle_velocity = None
-        self.start_pos = np.hstack([[np.pi/2], np.zeros(self.num_links - 1)])
+        self.start_pos = np.hstack([[np.pi / 2], np.zeros(self.num_links - 1)])
         self.start_vel = np.zeros(self.num_links)
         self.weight_matrix_scale = 1
 
+        self._steps = 0
         self.dt = 0.01
-        self.time_limit = 2
+        # self.time_limit = 2
 
         action_bound = np.pi * np.ones((self.num_links,))
         state_bound = np.hstack([
@@ -64,7 +68,7 @@ class ViaPointReacher(gym.Env):
 
         return self._get_obs().copy()
 
-    def step(self, action):
+    def step(self, action: np.ndarray):
         """
         a single step with an action in joint velocity space
         """
@@ -75,23 +79,20 @@ class ViaPointReacher(gym.Env):
 
         self._update_joints()
 
-        # rew = self._reward()
-
-        # compute reward directly in step function
-
         dist_reward = 0
         if not self._is_collided:
             if self._steps == 100:
                 dist_reward = np.linalg.norm(self.end_effector - self.via_point)
-            if self._steps == 199:
+            elif self._steps == 199:
                 dist_reward = np.linalg.norm(self.end_effector - self.goal_point)
 
+        # TODO: Do we need that?
         reward = - dist_reward ** 2
 
-        reward -= 1e-6 * np.sum(acc**2)
+        reward -= 1e-6 * np.sum(acc ** 2)
 
         if self._steps == 200:
-            reward -= 0.1 * np.sum(vel**2) ** 2
+            reward -= 0.1 * np.sum(vel ** 2) ** 2
 
         if self._is_collided:
             reward -= self.collision_penalty
@@ -100,7 +101,8 @@ class ViaPointReacher(gym.Env):
 
         self._steps += 1
 
-        done = self._steps * self.dt > self.time_limit or self._is_collided
+        # done = self._steps * self.dt > self.time_limit or self._is_collided
+        done = self._is_collided
 
         return self._get_obs().copy(), reward, done, info
 
@@ -118,8 +120,8 @@ class ViaPointReacher(gym.Env):
         self_collision = False
 
         if not self.allow_self_collision:
-            self_collision = self.check_self_collision(line_points_in_taskspace)
-            if np.any(np.abs(self._joint_angles) > np.pi) and not self.allow_self_collision:
+            self_collision = check_self_collision(line_points_in_taskspace)
+            if np.any(np.abs(self._joint_angles) > np.pi):
                 self_collision = True
 
         self._is_collided = self_collision
@@ -135,25 +137,10 @@ class ViaPointReacher(gym.Env):
             self._steps
         ])
 
-    # def _reward(self):
-    #     dist_reward = 0
-    #     if not self._is_collided:
-    #         if self._steps == 180:
-    #             dist_reward = np.linalg.norm(self.end_effector - self.bottom_center_of_hole)
-    #     else:
-    #         dist_reward = np.linalg.norm(self.end_effector - self.bottom_center_of_hole)
-    #
-    #     out = - dist_reward ** 2
-    #
-    #     return out
-
     def get_forward_kinematics(self, num_points_per_link=1):
         theta = self._joint_angles[:, None]
 
-        if num_points_per_link > 1:
-            intermediate_points = np.linspace(0, 1, num_points_per_link)
-        else:
-            intermediate_points = 1
+        intermediate_points = np.linspace(0, 1, num_points_per_link) if num_points_per_link > 1 else 1
 
         accumulated_theta = np.cumsum(theta, axis=0)
 
@@ -171,46 +158,6 @@ class ViaPointReacher(gym.Env):
 
         return np.squeeze(endeffector + self._joints[0, :])
 
-    def check_self_collision(self, line_points):
-        for i, line1 in enumerate(line_points):
-            for line2 in line_points[i+2:, :, :]:
-                # if line1 != line2:
-                if intersect(line1[0], line1[-1], line2[0], line2[-1]):
-                    return True
-        return False
-
-    def check_wall_collision(self, line_points):
-
-        # all points that are before the hole in x
-        r, c = np.where(line_points[:, :, 0] < (self.hole_x - self.hole_width / 2))
-
-        # check if any of those points are below surface
-        nr_line_points_below_surface_before_hole = np.sum(line_points[r, c, 1] < 0)
-
-        if nr_line_points_below_surface_before_hole > 0:
-            return True
-
-        # all points that are after the hole in x
-        r, c = np.where(line_points[:, :, 0] > (self.hole_x + self.hole_width / 2))
-
-        # check if any of those points are below surface
-        nr_line_points_below_surface_after_hole = np.sum(line_points[r, c, 1] < 0)
-
-        if nr_line_points_below_surface_after_hole > 0:
-            return True
-
-        # all points that are above the hole
-        r, c = np.where((line_points[:, :, 0] > (self.hole_x - self.hole_width / 2)) & (
-                           line_points[:, :, 0] < (self.hole_x + self.hole_width / 2)))
-
-        # check if any of those points are below surface
-        nr_line_points_below_surface_in_hole = np.sum(line_points[r, c, 1] < -self.hole_depth)
-
-        if nr_line_points_below_surface_in_hole > 0:
-            return True
-
-        return False
-
     def render(self, mode='human'):
         if self.fig is None:
             self.fig = plt.figure()
@@ -230,7 +177,7 @@ class ViaPointReacher(gym.Env):
             plt.xlim([-lim, lim])
             plt.ylim([-lim, lim])
             # plt.draw()
-            plt.pause(1e-4) #  pushes window to foreground, which is annoying.
+            plt.pause(1e-4)  # pushes window to foreground, which is annoying.
             # self.fig.canvas.flush_events()
 
         elif mode == "partial":
@@ -274,7 +221,7 @@ class ViaPointReacher(gym.Env):
 if __name__ == '__main__':
     nl = 5
     render_mode = "human"  # "human" or "partial" or "final"
-    env = ViaPointReacher(num_links=nl, allow_self_collision=False)
+    env = ViaPointReacher(n_links=nl, allow_self_collision=False)
     env.reset()
     env.render(mode=render_mode)