diff --git a/alr_envs/classic_control/hole_reacher_v2.py b/alr_envs/classic_control/hole_reacher_v2.py deleted file mode 100644 index 028ab26..0000000 --- a/alr_envs/classic_control/hole_reacher_v2.py +++ /dev/null @@ -1,307 +0,0 @@ -import gym -import numpy as np -import matplotlib.pyplot as plt -from matplotlib import patches -from alr_envs.classic_control.utils import check_self_collision - - -class HoleReacher(gym.Env): - - 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.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.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.n_links - 1)]) - self.start_vel = np.zeros(self.n_links) - - self.dt = 0.01 - # self.time_limit = 2 - - action_bound = np.pi * np.ones((self.n_links,)) - state_bound = np.hstack([ - [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 - ]) - self.action_space = gym.spaces.Box(low=-action_bound, high=action_bound, shape=action_bound.shape) - self.observation_space = gym.spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape) - - self.fig = None - 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.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, - fill=True, edgecolor='k', facecolor='k') - - rect_4 = patches.Rectangle((-1, 0), # south west corner - 0.5, # width - self.n_links, # height - fill=True, edgecolor='k', facecolor='k') - - self.patches = [rect_1, rect_2, rect_3, rect_4] - - @property - def end_effector(self): - return self._joints[self.n_links].T - - def configure(self, context): - pass - - def reset(self): - self._joint_angles = self.start_pos - self._angle_velocity = self.start_vel - self._joints = np.zeros((self.n_links + 1, 2)) - self._update_joints() - self._steps = 0 - - return self._get_obs().copy() - - def step(self, action: np.ndarray): - """ - a single step with an action in joint velocity space - """ - vel = action # + 0.01 * np.random.randn(self.num_links) - acc = (vel - self._angle_velocity) / self.dt - self._angle_velocity = vel - self._joint_angles = self._joint_angles + self.dt * self._angle_velocity - - self._update_joints() - - # rew = self._reward() - - # compute reward directly in step function - - success = False - reward = 0 - if not self._is_collided: - if self._steps == 199: - dist = np.linalg.norm(self.end_effector - self.bottom_center_of_hole) - reward = - dist ** 2 - success = dist < 0.005 - else: - dist = np.linalg.norm(self.end_effector - self.bottom_center_of_hole) - # if self.collision_penalty != 0: - # reward = -self.collision_penalty - # else: - reward = - dist ** 2 - self.collision_penalty - - reward -= 5e-8 * np.sum(acc ** 2) - - info = {"is_collided": self._is_collided, "is_success": success} - - self._steps += 1 - - # done = self._steps * self.dt > self.time_limit or self._is_collided - done = self._is_collided - - return self._get_obs().copy(), reward, done, info - - def _update_joints(self): - """ - update _joints to get new end effector position. The other links are only required for rendering. - Returns: - - """ - line_points_in_taskspace = self.get_forward_kinematics(num_points_per_link=20) - - self._joints[1:, 0] = self._joints[0, 0] + line_points_in_taskspace[:, -1, 0] - self._joints[1:, 1] = self._joints[0, 1] + line_points_in_taskspace[:, -1, 1] - - self_collision = False - wall_collision = False - - if not self.allow_self_collision: - self_collision = 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 = True - - if not self.allow_wall_collision: - wall_collision = self.check_wall_collision(line_points_in_taskspace) - - self._is_collided = self_collision or wall_collision - - def _get_obs(self): - theta = self._joint_angles - return np.hstack([ - np.cos(theta), - np.sin(theta), - self._angle_velocity, - self.end_effector - self.bottom_center_of_hole, - self._steps - ]) - - 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 - - accumulated_theta = np.cumsum(theta, axis=0) - - 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 - - endeffector[0, :, 0] = x[0, :] - endeffector[0, :, 1] = y[0, :] - - 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] - - return np.squeeze(endeffector + self._joints[0, :]) - - def check_wall_collision(self, line_points): - - # all points that are before the hole in x - r, c = np.where((line_points[:, :, 0] > -1) & (line_points[:, :, 0] < -0.5) & - (line_points[:, :, 1] > 0) & (line_points[:, :, 1] < self.n_links)) - - if len(r) > 0: - return True - - # 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() - # plt.ion() - # plt.pause(0.01) - else: - plt.figure(self.fig.number) - - if mode == "human": - plt.cla() - plt.title(f"Iteration: {self._steps}, distance: {self.end_effector - self.bottom_center_of_hole}") - - # Arm - plt.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k') - - # Add the patch to the Axes - [plt.gca().add_patch(rect) for rect in self.patches] - - lim = np.sum(self.link_lengths) + 0.5 - plt.xlim([-lim, lim]) - plt.ylim([-1.1, lim]) - # plt.draw() - plt.pause(1e-4) # pushes window to foreground, which is annoying. - # self.fig.canvas.flush_events() - - elif mode == "partial": - if self._steps == 1: - # fig, ax = plt.subplots() - # Add the patch to the Axes - [plt.gca().add_patch(rect) for rect in self.patches] - # plt.pause(0.01) - - if self._steps % 20 == 0 or self._steps in [1, 199] or self._is_collided: - # Arm - plt.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k', alpha=self._steps / 200) - # ax.plot(line_points_in_taskspace[:, 0, 0], - # line_points_in_taskspace[:, 0, 1], - # line_points_in_taskspace[:, -1, 0], - # line_points_in_taskspace[:, -1, 1], marker='o', color='k', alpha=t / 200) - - lim = np.sum(self.link_lengths) + 0.5 - plt.xlim([-lim, lim]) - plt.ylim([-1.1, lim]) - plt.pause(0.01) - - elif mode == "final": - if self._steps == 199 or self._is_collided: - # fig, ax = plt.subplots() - - # Add the patch to the Axes - [plt.gca().add_patch(rect) for rect in self.patches] - - 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') - - plt.pause(0.01) - - def close(self): - if self.fig is not None: - plt.close(self.fig) - - -if __name__ == '__main__': - nl = 5 - render_mode = "human" # "human" or "partial" or "final" - 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) - - for i in range(200): - # objective.load_result("/tmp/cma") - # test with random actions - ac = 2 * env.action_space.sample() - # ac[0] += np.pi/2 - obs, rew, d, info = env.step(ac) - env.render(mode=render_mode) - - print(rew) - - if d: - break - - env.close()