from typing import Union import gym import matplotlib.pyplot as plt import numpy as np from gym.utils import seeding from matplotlib import patches from alr_envs.classic_control.utils import check_self_collision class HoleReacherEnv(gym.Env): def __init__(self, n_links: int, hole_x: Union[None, float] = None, hole_depth: Union[None, float] = None, hole_width: float = 1., random_start: bool = False, allow_self_collision: bool = False, allow_wall_collision: bool = False, collision_penalty: float = 1000): self.n_links = n_links self.link_lengths = np.ones((n_links, 1)) self.random_start = random_start # provided initial parameters self.initial_x = hole_x # x-position of center of hole self.initial_width = hole_width # width of hole self.initial_depth = hole_depth # depth of hole # temp container for current env state self._tmp_x = None self._tmp_width = None self._tmp_depth = None self._goal = None # x-y coordinates for reaching the center at the bottom of the hole # 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 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], # hole width # [np.inf], # hole depth [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) # containers for plotting self.metadata = {'render.modes': ["human", "partial"]} self.fig = None self._steps = 0 self.seed() @property def dt(self) -> Union[float, int]: return self._dt # @property # def start_pos(self): # return self._start_pos @property def current_pos(self): return self._joint_angles.copy() @property def current_vel(self): return self._angle_velocity.copy() def step(self, action: np.ndarray): """ A single step with an action in joint velocity space """ acc = (action - self._angle_velocity) / self.dt self._angle_velocity = action self._joint_angles = self._joint_angles + self.dt * self._angle_velocity # + 0.001 * np.random.randn(5) self._update_joints() reward, info = self._get_reward(acc) info.update({"is_collided": self._is_collided}) self.end_effector_traj.append(np.copy(self.end_effector)) self._steps += 1 done = self._is_collided return self._get_obs().copy(), reward, done, info def reset(self): if self.random_start: # Maybe change more than first seed first_joint = self.np_random.uniform(np.pi / 4, 3 * np.pi / 4) self._joint_angles = np.hstack([[first_joint], np.zeros(self.n_links - 1)]) self._start_pos = self._joint_angles.copy() else: self._joint_angles = self._start_pos self._generate_hole() self._set_patches() self._angle_velocity = self._start_vel self._joints = np.zeros((self.n_links + 1, 2)) self._update_joints() self._steps = 0 self.end_effector_traj = [] return self._get_obs().copy() def _generate_hole(self): if self.initial_width is None: width = self.np_random.uniform(0.15, 0.5, 1) else: width = np.copy(self.initial_width) if self.initial_x is None: # sample whole on left or right side direction = np.random.choice([-1, 1]) # Hole center needs to be half the width away from the arm to give a valid setting. x = direction * self.np_random.uniform(width / 2, 3.5, 1) else: x = np.copy(self.initial_x) if self.initial_depth is None: # TODO we do not want this right now. depth = self.np_random.uniform(1, 1, 1) else: depth = np.copy(self.initial_depth) self._tmp_hole_width = width self._tmp_hole_x = x self._tmp_hole_depth = depth self._goal = np.hstack([self._tmp_hole_x, -self._tmp_hole_depth]) 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_reward(self, acc: np.ndarray): reward = 0 # success = False if self._steps == 199 or self._is_collided: # return reward only in last time step # Episode also terminates when colliding, hence return reward dist = np.linalg.norm(self.end_effector - self._goal) # success = dist < 0.005 and not self._is_collided reward = - dist ** 2 - self.collision_penalty * self._is_collided reward -= 5e-8 * np.sum(acc ** 2) # info = {"is_success": success} return reward, {} # info def _get_obs(self): theta = self._joint_angles return np.hstack([ np.cos(theta), np.sin(theta), self._angle_velocity, self._tmp_width, # self._tmp_hole_depth, self.end_effector - self._goal, self._steps ]) def _get_forward_kinematics(self, num_points_per_link=1): theta = self._joint_angles[:, None] 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) end_effector = 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 end_effector[0, :, 0] = x[0, :] end_effector[0, :, 1] = y[0, :] for i in range(1, self.n_links): end_effector[i, :, 0] = x[i, :] + end_effector[i - 1, -1, 0] end_effector[i, :, 1] = y[i, :] + end_effector[i - 1, -1, 1] return np.squeeze(end_effector + 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] < (self._tmp_x - self._tmp_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._tmp_x + self._tmp_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._tmp_x - self._tmp_width / 2)) & ( line_points[:, :, 0] < (self._tmp_x + self._tmp_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._tmp_depth) if nr_line_points_below_surface_in_hole > 0: return True return False def render(self, mode='human'): if self.fig is None: # Create base figure once on the beginning. Afterwards only update plt.ion() self.fig = plt.figure() ax = self.fig.add_subplot(1, 1, 1) # limits lim = np.sum(self.link_lengths) + 0.5 ax.set_xlim([-lim, lim]) ax.set_ylim([-1.1, lim]) self.line, = ax.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k') self._set_patches() self.fig.show() self.fig.gca().set_title( f"Iteration: {self._steps}, distance: {self.end_effector - self._goal}") if mode == "human": # arm self.line.set_data(self._joints[:, 0], self._joints[:, 1]) self.fig.canvas.draw() self.fig.canvas.flush_events() elif mode == "partial": 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) def _set_patches(self): if self.fig is not None: self.fig.gca().patches = [] left_block = patches.Rectangle((-self.n_links, -self._tmp_depth), self.n_links + self._tmp_x - self._tmp_width / 2, self._tmp_depth, fill=True, edgecolor='k', facecolor='k') right_block = patches.Rectangle((self._tmp_x + self._tmp_width / 2, -self._tmp_depth), self.n_links - self._tmp_x + self._tmp_width / 2, self._tmp_depth, fill=True, edgecolor='k', facecolor='k') hole_floor = patches.Rectangle((self._tmp_x - self._tmp_width / 2, -self._tmp_depth), self._tmp_width, 1 - self._tmp_depth, fill=True, edgecolor='k', facecolor='k') # Add the patch to the Axes self.fig.gca().add_patch(left_block) self.fig.gca().add_patch(right_block) self.fig.gca().add_patch(hole_floor) def seed(self, seed=None): self.np_random, seed = seeding.np_random(seed) return [seed] @property def end_effector(self): return self._joints[self.n_links].T def close(self): super().close() if self.fig is not None: plt.close(self.fig)