improved project structure and exposed methods
This commit is contained in:
ottofabian
@@ -0,0 +1,21 @@
|
||||
### Classic Control
|
||||
|
||||
## Step-based Environments
|
||||
|Name| Description|Horizon|Action Dimension|Observation Dimension
|
||||
|---|---|---|---|---|
|
||||
|`SimpleReacher-v0`| Simple reaching task (2 links) without any physics simulation. Provides no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions towards the end of the trajectory.| 200 | 2 | 9
|
||||
|`LongSimpleReacher-v0`| Simple reaching task (5 links) without any physics simulation. Provides no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions towards the end of the trajectory.| 200 | 5 | 18
|
||||
|`ViaPointReacher-v0`| Simple reaching task leveraging a via point, which supports self collision detection. Provides a reward only at 100 and 199 for reaching the viapoint and goal point, respectively.| 200 | 5 | 18
|
||||
|`HoleReacher-v0`| 5 link reaching task where the end-effector needs to reach into a narrow hole without collding with itself or walls | 200 | 5 | 18
|
||||
|
||||
## MP Environments
|
||||
|Name| Description|Horizon|Action Dimension|Context Dimension
|
||||
|---|---|---|---|---|
|
||||
|`ViaPointReacherDMP-v0`| A DMP provides a trajectory for the `ViaPointReacher-v0` task. | 200 | 25
|
||||
|`HoleReacherFixedGoalDMP-v0`| A DMP provides a trajectory for the `HoleReacher-v0` task with a fixed goal attractor. | 200 | 25
|
||||
|`HoleReacherDMP-v0`| A DMP provides a trajectory for the `HoleReacher-v0` task. The goal attractor needs to be learned. | 200 | 30
|
||||
|`ALRBallInACupSimpleDMP-v0`| A DMP provides a trajectory for the `ALRBallInACupSimple-v0` task where only 3 joints are actuated. | 4000 | 15
|
||||
|`ALRBallInACupDMP-v0`| A DMP provides a trajectory for the `ALRBallInACup-v0` task. | 4000 | 35
|
||||
|`ALRBallInACupGoalDMP-v0`| A DMP provides a trajectory for the `ALRBallInACupGoal-v0` task. | 4000 | 35 | 3
|
||||
|
||||
[//]: |`HoleReacherDetPMP-v0`|
|
||||
@@ -0,0 +1,3 @@
|
||||
from .hole_reacher.hole_reacher import HoleReacherEnv
|
||||
from .simple_reacher.simple_reacher import SimpleReacherEnv
|
||||
from .viapoint_reacher.viapoint_reacher import ViaPointReacherEnv
|
||||
@@ -0,0 +1 @@
|
||||
from .mp_wrapper import MPWrapper
|
||||
@@ -0,0 +1,314 @@
|
||||
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.alr.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)
|
||||
else:
|
||||
width = np.copy(self.initial_width)
|
||||
if self.initial_x is None:
|
||||
# sample whole on left or right side
|
||||
direction = self.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)
|
||||
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)
|
||||
else:
|
||||
depth = np.copy(self.initial_depth)
|
||||
|
||||
self._tmp_width = width
|
||||
self._tmp_x = x
|
||||
self._tmp_depth = depth
|
||||
self._goal = np.hstack([self._tmp_x, -self._tmp_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: {np.linalg.norm(self.end_effector - self._goal) ** 2}")
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from mp_env_api import MPEnvWrapper
|
||||
|
||||
|
||||
class MPWrapper(MPEnvWrapper):
|
||||
@property
|
||||
def active_obs(self):
|
||||
return np.hstack([
|
||||
[self.env.random_start] * self.env.n_links, # cos
|
||||
[self.env.random_start] * self.env.n_links, # sin
|
||||
[self.env.random_start] * self.env.n_links, # velocity
|
||||
[self.env.initial_width is None], # hole width
|
||||
# [self.env.hole_depth is None], # hole depth
|
||||
[True] * 2, # x-y coordinates of target distance
|
||||
[False] # env steps
|
||||
])
|
||||
|
||||
@property
|
||||
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.current_pos
|
||||
|
||||
@property
|
||||
def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.current_vel
|
||||
|
||||
@property
|
||||
def goal_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
raise ValueError("Goal position is not available and has to be learnt based on the environment.")
|
||||
|
||||
@property
|
||||
def dt(self) -> Union[float, int]:
|
||||
return self.env.dt
|
||||
@@ -0,0 +1 @@
|
||||
from .mp_wrapper import MPWrapper
|
||||
@@ -0,0 +1,33 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from mp_env_api import MPEnvWrapper
|
||||
|
||||
|
||||
class MPWrapper(MPEnvWrapper):
|
||||
@property
|
||||
def active_obs(self):
|
||||
return np.hstack([
|
||||
[self.env.random_start] * self.env.n_links, # cos
|
||||
[self.env.random_start] * self.env.n_links, # sin
|
||||
[self.env.random_start] * self.env.n_links, # velocity
|
||||
[True] * 2, # x-y coordinates of target distance
|
||||
[False] # env steps
|
||||
])
|
||||
|
||||
@property
|
||||
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.current_pos
|
||||
|
||||
@property
|
||||
def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.current_vel
|
||||
|
||||
@property
|
||||
def goal_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
raise ValueError("Goal position is not available and has to be learnt based on the environment.")
|
||||
|
||||
@property
|
||||
def dt(self) -> Union[float, int]:
|
||||
return self.env.dt
|
||||
@@ -0,0 +1,202 @@
|
||||
from typing import Iterable, Union
|
||||
|
||||
import gym
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from gym import spaces
|
||||
from gym.utils import seeding
|
||||
|
||||
|
||||
class SimpleReacherEnv(gym.Env):
|
||||
"""
|
||||
Simple Reaching Task without any physics simulation.
|
||||
Returns no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions
|
||||
towards the end of the trajectory.
|
||||
"""
|
||||
|
||||
def __init__(self, n_links: int, target: Union[None, Iterable] = None, random_start: bool = True):
|
||||
super().__init__()
|
||||
self.link_lengths = np.ones(n_links)
|
||||
self.n_links = n_links
|
||||
self._dt = 0.1
|
||||
|
||||
self.random_start = random_start
|
||||
|
||||
# provided initial parameters
|
||||
self.inital_target = target
|
||||
|
||||
# temp container for current env state
|
||||
self._goal = None
|
||||
|
||||
self._joints = None
|
||||
self._joint_angles = None
|
||||
self._angle_velocity = None
|
||||
self._start_pos = np.zeros(self.n_links)
|
||||
self._start_vel = np.zeros(self.n_links)
|
||||
|
||||
self.max_torque = 1
|
||||
self.steps_before_reward = 199
|
||||
|
||||
action_bound = np.ones((self.n_links,)) * self.max_torque
|
||||
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 = spaces.Box(low=-action_bound, high=action_bound, shape=action_bound.shape)
|
||||
self.observation_space = spaces.Box(low=-state_bound, high=state_bound, shape=state_bound.shape)
|
||||
|
||||
# containers for plotting
|
||||
self.metadata = {'render.modes': ["human"]}
|
||||
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
|
||||
|
||||
@property
|
||||
def current_vel(self):
|
||||
return self._angle_velocity
|
||||
|
||||
def step(self, action: np.ndarray):
|
||||
"""
|
||||
A single step with action in torque space
|
||||
"""
|
||||
|
||||
# action = self._add_action_noise(action)
|
||||
ac = np.clip(action, -self.max_torque, self.max_torque)
|
||||
|
||||
self._angle_velocity = self._angle_velocity + self.dt * ac
|
||||
self._joint_angles = self._joint_angles + self.dt * self._angle_velocity
|
||||
self._update_joints()
|
||||
|
||||
reward, info = self._get_reward(action)
|
||||
|
||||
self._steps += 1
|
||||
done = False
|
||||
|
||||
return self._get_obs().copy(), reward, done, info
|
||||
|
||||
def reset(self):
|
||||
|
||||
# TODO: maybe do initialisation more random?
|
||||
# Sample only orientation of first link, i.e. the arm is always straight.
|
||||
if self.random_start:
|
||||
self._joint_angles = np.hstack([[self.np_random.uniform(-np.pi, np.pi)], np.zeros(self.n_links - 1)])
|
||||
self._start_pos = self._joint_angles.copy()
|
||||
else:
|
||||
self._joint_angles = self._start_pos
|
||||
|
||||
self._generate_goal()
|
||||
|
||||
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 _update_joints(self):
|
||||
"""
|
||||
update joints to get new end-effector position. The other links are only required for rendering.
|
||||
Returns:
|
||||
|
||||
"""
|
||||
angles = np.cumsum(self._joint_angles)
|
||||
x = self.link_lengths * np.vstack([np.cos(angles), np.sin(angles)])
|
||||
self._joints[1:] = self._joints[0] + np.cumsum(x.T, axis=0)
|
||||
|
||||
def _get_reward(self, action: np.ndarray):
|
||||
diff = self.end_effector - self._goal
|
||||
reward_dist = 0
|
||||
|
||||
if self._steps >= self.steps_before_reward:
|
||||
reward_dist -= np.linalg.norm(diff)
|
||||
# reward_dist = np.exp(-0.1 * diff ** 2).mean()
|
||||
# reward_dist = - (diff ** 2).mean()
|
||||
|
||||
reward_ctrl = (action ** 2).sum()
|
||||
reward = reward_dist - reward_ctrl
|
||||
return reward, dict(reward_dist=reward_dist, reward_ctrl=reward_ctrl)
|
||||
|
||||
def _get_obs(self):
|
||||
theta = self._joint_angles
|
||||
return np.hstack([
|
||||
np.cos(theta),
|
||||
np.sin(theta),
|
||||
self._angle_velocity,
|
||||
self.end_effector - self._goal,
|
||||
self._steps
|
||||
])
|
||||
|
||||
def _generate_goal(self):
|
||||
|
||||
if self.inital_target is None:
|
||||
|
||||
total_length = np.sum(self.link_lengths)
|
||||
goal = np.array([total_length, total_length])
|
||||
while np.linalg.norm(goal) >= total_length:
|
||||
goal = self.np_random.uniform(low=-total_length, high=total_length, size=2)
|
||||
else:
|
||||
goal = np.copy(self.inital_target)
|
||||
|
||||
self._goal = goal
|
||||
|
||||
def render(self, mode='human'): # pragma: no cover
|
||||
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([-lim, lim])
|
||||
|
||||
self.line, = ax.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k')
|
||||
goal_pos = self._goal.T
|
||||
self.goal_point, = ax.plot(goal_pos[0], goal_pos[1], 'gx')
|
||||
self.goal_dist, = ax.plot([self.end_effector[0], goal_pos[0]], [self.end_effector[1], goal_pos[1]], 'g--')
|
||||
|
||||
self.fig.show()
|
||||
|
||||
self.fig.gca().set_title(f"Iteration: {self._steps}, distance: {self.end_effector - self._goal}")
|
||||
|
||||
# goal
|
||||
goal_pos = self._goal.T
|
||||
if self._steps == 1:
|
||||
self.goal_point.set_data(goal_pos[0], goal_pos[1])
|
||||
|
||||
# arm
|
||||
self.line.set_data(self._joints[:, 0], self._joints[:, 1])
|
||||
|
||||
# distance between end effector and goal
|
||||
self.goal_dist.set_data([self.end_effector[0], goal_pos[0]], [self.end_effector[1], goal_pos[1]])
|
||||
|
||||
self.fig.canvas.draw()
|
||||
self.fig.canvas.flush_events()
|
||||
|
||||
def seed(self, seed=None):
|
||||
self.np_random, seed = seeding.np_random(seed)
|
||||
return [seed]
|
||||
|
||||
def close(self):
|
||||
del self.fig
|
||||
|
||||
@property
|
||||
def end_effector(self):
|
||||
return self._joints[self.n_links].T
|
||||
@@ -0,0 +1,18 @@
|
||||
def ccw(A, B, C):
|
||||
return (C[1] - A[1]) * (B[0] - A[0]) - (B[1] - A[1]) * (C[0] - A[0]) > 1e-12
|
||||
|
||||
|
||||
def intersect(A, B, C, D):
|
||||
"""
|
||||
Checks whether line segments AB and CD intersect
|
||||
"""
|
||||
return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
|
||||
|
||||
|
||||
def check_self_collision(line_points):
|
||||
"""Checks whether line segments and intersect"""
|
||||
for i, line1 in enumerate(line_points):
|
||||
for line2 in line_points[i + 2:, :, :]:
|
||||
if intersect(line1[0], line1[-1], line2[0], line2[-1]):
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1 @@
|
||||
from .mp_wrapper import MPWrapper
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from mp_env_api import MPEnvWrapper
|
||||
|
||||
|
||||
class MPWrapper(MPEnvWrapper):
|
||||
@property
|
||||
def active_obs(self):
|
||||
return np.hstack([
|
||||
[self.env.random_start] * self.env.n_links, # cos
|
||||
[self.env.random_start] * self.env.n_links, # sin
|
||||
[self.env.random_start] * self.env.n_links, # velocity
|
||||
[self.env.initial_via_target is None] * 2, # x-y coordinates of via point distance
|
||||
[True] * 2, # x-y coordinates of target distance
|
||||
[False] # env steps
|
||||
])
|
||||
|
||||
@property
|
||||
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.current_pos
|
||||
|
||||
@property
|
||||
def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
return self.env.current_vel
|
||||
|
||||
@property
|
||||
def goal_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
raise ValueError("Goal position is not available and has to be learnt based on the environment.")
|
||||
|
||||
@property
|
||||
def dt(self) -> Union[float, int]:
|
||||
return self.env.dt
|
||||
@@ -0,0 +1,294 @@
|
||||
from typing import Iterable, Union
|
||||
|
||||
import gym
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from gym.utils import seeding
|
||||
|
||||
from alr_envs.alr.classic_control.utils import check_self_collision
|
||||
|
||||
|
||||
class ViaPointReacherEnv(gym.Env):
|
||||
|
||||
def __init__(self, n_links, random_start: bool = False, via_target: Union[None, Iterable] = None,
|
||||
target: Union[None, Iterable] = None, allow_self_collision=False, collision_penalty=1000):
|
||||
|
||||
self.n_links = n_links
|
||||
self.link_lengths = np.ones((n_links, 1))
|
||||
|
||||
self.random_start = random_start
|
||||
|
||||
# provided initial parameters
|
||||
self.intitial_target = target # provided target value
|
||||
self.initial_via_target = via_target # provided via point target value
|
||||
|
||||
# temp container for current env state
|
||||
self._via_point = np.ones(2)
|
||||
self._goal = np.array((n_links, 0))
|
||||
|
||||
# collision
|
||||
self.allow_self_collision = allow_self_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.weight_matrix_scale = 1
|
||||
|
||||
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] * 2, # x-y coordinates of via point distance
|
||||
[np.inf] * 2, # x-y coordinates of target distance
|
||||
[np.inf] # env steps, because reward start after n steps
|
||||
])
|
||||
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):
|
||||
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
|
||||
"""
|
||||
vel = action
|
||||
self._angle_velocity = vel
|
||||
self._joint_angles = self._joint_angles + self.dt * self._angle_velocity
|
||||
self._update_joints()
|
||||
|
||||
acc = (vel - self._angle_velocity) / self.dt
|
||||
reward, info = self._get_reward(acc)
|
||||
|
||||
info.update({"is_collided": self._is_collided})
|
||||
|
||||
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 dirst 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_goal()
|
||||
|
||||
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 _generate_goal(self):
|
||||
# TODO: Maybe improve this later, this can yield quite a lot of invalid settings
|
||||
|
||||
total_length = np.sum(self.link_lengths)
|
||||
|
||||
# rejection sampled point in inner circle with 0.5*Radius
|
||||
if self.initial_via_target is None:
|
||||
via_target = np.array([total_length, total_length])
|
||||
while np.linalg.norm(via_target) >= 0.5 * total_length:
|
||||
via_target = self.np_random.uniform(low=-0.5 * total_length, high=0.5 * total_length, size=2)
|
||||
else:
|
||||
via_target = np.copy(self.initial_via_target)
|
||||
|
||||
# rejection sampled point in outer circle
|
||||
if self.intitial_target is None:
|
||||
goal = np.array([total_length, total_length])
|
||||
while np.linalg.norm(goal) >= total_length or np.linalg.norm(goal) <= 0.5 * total_length:
|
||||
goal = self.np_random.uniform(low=-total_length, high=total_length, size=2)
|
||||
else:
|
||||
goal = np.copy(self.intitial_target)
|
||||
|
||||
self._via_point = via_target
|
||||
self._goal = goal
|
||||
|
||||
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
|
||||
|
||||
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):
|
||||
self_collision = True
|
||||
|
||||
self._is_collided = self_collision
|
||||
|
||||
def _get_reward(self, acc):
|
||||
success = False
|
||||
reward = -np.inf
|
||||
if not self._is_collided:
|
||||
dist = np.inf
|
||||
# return intermediate reward for via point
|
||||
if self._steps == 100:
|
||||
dist = np.linalg.norm(self.end_effector - self._via_point)
|
||||
# return reward in last time step for goal
|
||||
elif self._steps == 199:
|
||||
dist = np.linalg.norm(self.end_effector - self._goal)
|
||||
|
||||
success = dist < 0.005
|
||||
else:
|
||||
# Episode terminates when colliding, hence return reward
|
||||
dist = np.linalg.norm(self.end_effector - self._goal)
|
||||
reward = -self.collision_penalty
|
||||
|
||||
reward -= dist ** 2
|
||||
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.end_effector - self._via_point,
|
||||
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)
|
||||
|
||||
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 render(self, mode='human'):
|
||||
goal_pos = self._goal.T
|
||||
via_pos = self._via_point.T
|
||||
|
||||
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([-lim, lim])
|
||||
|
||||
self.line, = ax.plot(self._joints[:, 0], self._joints[:, 1], 'ro-', markerfacecolor='k')
|
||||
self.goal_point_plot, = ax.plot(goal_pos[0], goal_pos[1], 'go')
|
||||
self.via_point_plot, = ax.plot(via_pos[0], via_pos[1], 'gx')
|
||||
|
||||
self.fig.show()
|
||||
|
||||
self.fig.gca().set_title(f"Iteration: {self._steps}, distance: {self.end_effector - self._goal}")
|
||||
|
||||
if mode == "human":
|
||||
# goal
|
||||
if self._steps == 1:
|
||||
self.goal_point_plot.set_data(goal_pos[0], goal_pos[1])
|
||||
self.via_point_plot.set_data(via_pos[0], goal_pos[1])
|
||||
|
||||
# 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 == 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 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):
|
||||
if self.fig is not None:
|
||||
plt.close(self.fig)
|
||||
Reference in New Issue
Block a user