refractoring of DMP environmets to fit gym interface better.
This commit is contained in:
ottofabian
@@ -1,87 +0,0 @@
|
||||
from alr_envs.utils.policies import get_policy_class
|
||||
from mp_lib import det_promp
|
||||
import numpy as np
|
||||
import gym
|
||||
|
||||
|
||||
class DetPMPEnvWrapper(gym.Wrapper):
|
||||
def __init__(self,
|
||||
env,
|
||||
num_dof,
|
||||
num_basis,
|
||||
width,
|
||||
start_pos=None,
|
||||
duration=1,
|
||||
dt=0.01,
|
||||
post_traj_time=0.,
|
||||
policy_type=None,
|
||||
weights_scale=1,
|
||||
zero_start=False,
|
||||
zero_goal=False,
|
||||
):
|
||||
super(DetPMPEnvWrapper, self).__init__(env)
|
||||
self.num_dof = num_dof
|
||||
self.num_basis = num_basis
|
||||
self.dim = num_dof * num_basis
|
||||
self.pmp = det_promp.DeterministicProMP(n_basis=num_basis, n_dof=num_dof, width=width, off=0.01,
|
||||
zero_start=zero_start, zero_goal=zero_goal)
|
||||
weights = np.zeros(shape=(num_basis, num_dof))
|
||||
self.pmp.set_weights(duration, weights)
|
||||
self.weights_scale = weights_scale
|
||||
|
||||
self.duration = duration
|
||||
self.dt = dt
|
||||
self.post_traj_steps = int(post_traj_time / dt)
|
||||
|
||||
self.start_pos = start_pos
|
||||
self.zero_start = zero_start
|
||||
|
||||
policy_class = get_policy_class(policy_type)
|
||||
self.policy = policy_class(env)
|
||||
|
||||
def __call__(self, params, contexts=None):
|
||||
params = np.atleast_2d(params)
|
||||
rewards = []
|
||||
infos = []
|
||||
for p, c in zip(params, contexts):
|
||||
reward, info = self.rollout(p, c)
|
||||
rewards.append(reward)
|
||||
infos.append(info)
|
||||
|
||||
return np.array(rewards), infos
|
||||
|
||||
def rollout(self, params, context=None, render=False):
|
||||
""" This function generates a trajectory based on a DMP and then does the usual loop over reset and step"""
|
||||
params = np.reshape(params, newshape=(self.num_basis, self.num_dof)) * self.weights_scale
|
||||
self.pmp.set_weights(self.duration, params)
|
||||
t, des_pos, des_vel, des_acc = self.pmp.compute_trajectory(1 / self.dt, 1.)
|
||||
if self.zero_start:
|
||||
des_pos += self.start_pos[None, :]
|
||||
|
||||
if self.post_traj_steps > 0:
|
||||
des_pos = np.vstack([des_pos, np.tile(des_pos[-1, :], [self.post_traj_steps, 1])])
|
||||
des_vel = np.vstack([des_vel, np.zeros(shape=(self.post_traj_steps, self.num_dof))])
|
||||
|
||||
self._trajectory = des_pos
|
||||
self._velocity = des_vel
|
||||
|
||||
rews = []
|
||||
infos = []
|
||||
|
||||
self.env.configure(context)
|
||||
self.env.reset()
|
||||
|
||||
for t, pos_vel in enumerate(zip(des_pos, des_vel)):
|
||||
ac = self.policy.get_action(pos_vel[0], pos_vel[1])
|
||||
obs, rew, done, info = self.env.step(ac)
|
||||
rews.append(rew)
|
||||
infos.append(info)
|
||||
if render:
|
||||
self.env.render(mode="human")
|
||||
if done:
|
||||
break
|
||||
|
||||
reward = np.sum(rews)
|
||||
|
||||
return reward, info
|
||||
|
||||
@@ -1,121 +0,0 @@
|
||||
from alr_envs.utils.policies import get_policy_class
|
||||
from mp_lib.phase import ExpDecayPhaseGenerator
|
||||
from mp_lib.basis import DMPBasisGenerator
|
||||
from mp_lib import dmps
|
||||
import numpy as np
|
||||
import gym
|
||||
|
||||
|
||||
class DmpEnvWrapper(gym.Wrapper):
|
||||
def __init__(self,
|
||||
env,
|
||||
num_dof,
|
||||
num_basis,
|
||||
start_pos=None,
|
||||
final_pos=None,
|
||||
duration=1,
|
||||
alpha_phase=2,
|
||||
dt=0.01,
|
||||
learn_goal=False,
|
||||
post_traj_time=0.,
|
||||
policy_type=None,
|
||||
weights_scale=1.,
|
||||
goal_scale=1.,
|
||||
):
|
||||
super(DmpEnvWrapper, self).__init__(env)
|
||||
self.num_dof = num_dof
|
||||
self.num_basis = num_basis
|
||||
self.dim = num_dof * num_basis
|
||||
if learn_goal:
|
||||
self.dim += num_dof
|
||||
self.learn_goal = learn_goal
|
||||
self.duration = duration # seconds
|
||||
time_steps = int(duration / dt)
|
||||
self.t = np.linspace(0, duration, time_steps)
|
||||
self.post_traj_steps = int(post_traj_time / dt)
|
||||
|
||||
phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
|
||||
basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=self.num_basis)
|
||||
|
||||
self.dmp = dmps.DMP(num_dof=num_dof,
|
||||
basis_generator=basis_generator,
|
||||
phase_generator=phase_generator,
|
||||
num_time_steps=time_steps,
|
||||
dt=dt
|
||||
)
|
||||
|
||||
self.dmp.dmp_start_pos = start_pos.reshape((1, num_dof))
|
||||
|
||||
dmp_weights = np.zeros((num_basis, num_dof))
|
||||
if learn_goal:
|
||||
dmp_goal_pos = np.zeros(num_dof)
|
||||
else:
|
||||
dmp_goal_pos = final_pos
|
||||
|
||||
self.dmp.set_weights(dmp_weights, dmp_goal_pos)
|
||||
self.weights_scale = weights_scale
|
||||
self.goal_scale = goal_scale
|
||||
|
||||
policy_class = get_policy_class(policy_type)
|
||||
self.policy = policy_class(env)
|
||||
|
||||
def __call__(self, params, contexts=None):
|
||||
params = np.atleast_2d(params)
|
||||
rewards = []
|
||||
infos = []
|
||||
for p, c in zip(params, contexts):
|
||||
reward, info = self.rollout(p, c)
|
||||
rewards.append(reward)
|
||||
infos.append(info)
|
||||
|
||||
return np.array(rewards), infos
|
||||
|
||||
def goal_and_weights(self, params):
|
||||
if len(params.shape) > 1:
|
||||
assert params.shape[1] == self.dim
|
||||
else:
|
||||
assert len(params) == self.dim
|
||||
params = np.reshape(params, [1, self.dim])
|
||||
|
||||
if self.learn_goal:
|
||||
goal_pos = params[0, -self.num_dof:]
|
||||
weight_matrix = np.reshape(params[:, :-self.num_dof], [self.num_basis, self.num_dof])
|
||||
else:
|
||||
goal_pos = self.dmp.dmp_goal_pos.flatten()
|
||||
assert goal_pos is not None
|
||||
weight_matrix = np.reshape(params, [self.num_basis, self.num_dof])
|
||||
|
||||
return goal_pos * self.goal_scale, weight_matrix * self.weights_scale
|
||||
|
||||
def rollout(self, params, context=None, render=False):
|
||||
""" This function generates a trajectory based on a DMP and then does the usual loop over reset and step"""
|
||||
goal_pos, weight_matrix = self.goal_and_weights(params)
|
||||
self.dmp.set_weights(weight_matrix, goal_pos)
|
||||
trajectory, velocity = self.dmp.reference_trajectory(self.t)
|
||||
|
||||
if self.post_traj_steps > 0:
|
||||
trajectory = np.vstack([trajectory, np.tile(trajectory[-1, :], [self.post_traj_steps, 1])])
|
||||
velocity = np.vstack([velocity, np.zeros(shape=(self.post_traj_steps, self.num_dof))])
|
||||
|
||||
self._trajectory = trajectory
|
||||
self._velocity = velocity
|
||||
|
||||
rews = []
|
||||
infos = []
|
||||
|
||||
self.env.configure(context)
|
||||
self.env.reset()
|
||||
|
||||
for t, pos_vel in enumerate(zip(trajectory, velocity)):
|
||||
ac = self.policy.get_action(pos_vel[0], pos_vel[1])
|
||||
obs, rew, done, info = self.env.step(ac)
|
||||
rews.append(rew)
|
||||
infos.append(info)
|
||||
if render:
|
||||
self.env.render(mode="human")
|
||||
if done:
|
||||
break
|
||||
|
||||
reward = np.sum(rews)
|
||||
|
||||
return reward, info
|
||||
@@ -1,8 +1,10 @@
|
||||
from gym import Env
|
||||
|
||||
from alr_envs.mujoco.alr_mujoco_env import AlrMujocoEnv
|
||||
|
||||
|
||||
class BaseController:
|
||||
def __init__(self, env: AlrMujocoEnv):
|
||||
def __init__(self, env: Env):
|
||||
self.env = env
|
||||
|
||||
def get_action(self, des_pos, des_vel):
|
||||
@@ -20,7 +22,7 @@ class VelController(BaseController):
|
||||
|
||||
|
||||
class PDController(BaseController):
|
||||
def __init__(self, env):
|
||||
def __init__(self, env: AlrMujocoEnv):
|
||||
self.p_gains = env.p_gains
|
||||
self.d_gains = env.d_gains
|
||||
super(PDController, self).__init__(env)
|
||||
|
||||
@@ -18,3 +18,31 @@ def angle_normalize(x, type="deg"):
|
||||
return x - two_pi * np.floor((x + np.pi) / two_pi)
|
||||
else:
|
||||
raise ValueError(f"Invalid type {type}. Choose on of 'deg' or 'rad'.")
|
||||
|
||||
|
||||
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):
|
||||
"""
|
||||
Return true if line segments AB and CD intersects
|
||||
Args:
|
||||
A: start point line one
|
||||
B: end point line one
|
||||
C: start point line two
|
||||
D: end point line two
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
|
||||
|
||||
|
||||
def check_self_collision(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
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
import gym
|
||||
import numpy as np
|
||||
from mp_lib import det_promp
|
||||
|
||||
from alr_envs.utils.wrapper.mp_wrapper import MPWrapper
|
||||
|
||||
|
||||
class DetPMPWrapper(MPWrapper):
|
||||
def __init__(self, env, num_dof, num_basis, width, start_pos=None, duration=1, dt=0.01, post_traj_time=0.,
|
||||
policy_type=None, weights_scale=1, zero_start=False, zero_goal=False, **mp_kwargs):
|
||||
# self.duration = duration # seconds
|
||||
|
||||
super().__init__(env, num_dof, duration, dt, post_traj_time, policy_type, weights_scale,
|
||||
num_basis=num_basis, width=width, start_pos=start_pos, zero_start=zero_start,
|
||||
zero_goal=zero_goal)
|
||||
|
||||
action_bounds = np.inf * np.ones((self.mp.n_basis * self.mp.n_dof))
|
||||
self.action_space = gym.spaces.Box(low=-action_bounds, high=action_bounds, dtype=np.float32)
|
||||
|
||||
self.start_pos = start_pos
|
||||
self.dt = dt
|
||||
|
||||
def initialize_mp(self, num_dof: int, duration: int, dt: float, num_basis: int = 5, width: float = None,
|
||||
start_pos: np.ndarray = None, zero_start: bool = False, zero_goal: bool = False):
|
||||
pmp = det_promp.DeterministicProMP(n_basis=num_basis, n_dof=num_dof, width=width, off=0.01,
|
||||
zero_start=zero_start, zero_goal=zero_goal)
|
||||
|
||||
weights = np.zeros(shape=(num_basis, num_dof))
|
||||
pmp.set_weights(duration, weights)
|
||||
|
||||
return pmp
|
||||
|
||||
def mp_rollout(self, action):
|
||||
params = np.reshape(action, (self.mp.n_basis, self.mp.n_dof)) * self.weights_scale
|
||||
self.mp.set_weights(self.duration, params)
|
||||
_, des_pos, des_vel, _ = self.mp.compute_trajectory(1 / self.dt, 1.)
|
||||
if self.mp.zero_start:
|
||||
des_pos += self.start_pos[None, :]
|
||||
|
||||
return des_pos, des_vel
|
||||
@@ -0,0 +1,81 @@
|
||||
from alr_envs.utils.policies import get_policy_class
|
||||
from mp_lib.phase import ExpDecayPhaseGenerator
|
||||
from mp_lib.basis import DMPBasisGenerator
|
||||
from mp_lib import dmps
|
||||
import numpy as np
|
||||
import gym
|
||||
|
||||
from alr_envs.utils.wrapper.mp_wrapper import MPWrapper
|
||||
|
||||
|
||||
class DmpWrapper(MPWrapper):
|
||||
|
||||
def __init__(self, env: gym.Env, num_dof: int, num_basis: int, start_pos: np.ndarray = None,
|
||||
final_pos: np.ndarray = None, duration: int = 1, alpha_phase: float = 2., dt: float = 0.01,
|
||||
learn_goal: bool = False, post_traj_time: float = 0., policy_type: str = None,
|
||||
weights_scale: float = 1., goal_scale: float = 1.):
|
||||
|
||||
"""
|
||||
This Wrapper generates a trajectory based on a DMP and will only return episodic performances.
|
||||
Args:
|
||||
env:
|
||||
num_dof:
|
||||
num_basis:
|
||||
start_pos:
|
||||
final_pos:
|
||||
duration:
|
||||
alpha_phase:
|
||||
dt:
|
||||
learn_goal:
|
||||
post_traj_time:
|
||||
policy_type:
|
||||
weights_scale:
|
||||
goal_scale:
|
||||
"""
|
||||
self.learn_goal = learn_goal
|
||||
self.t = np.linspace(0, duration, int(duration / dt))
|
||||
self.goal_scale = goal_scale
|
||||
|
||||
super().__init__(env, num_dof, duration, dt, post_traj_time, policy_type, weights_scale,
|
||||
num_basis=num_basis, start_pos=start_pos, final_pos=final_pos, alpha_phase=alpha_phase)
|
||||
|
||||
action_bounds = np.inf * np.ones((np.prod(self.mp.dmp_weights.shape) + (num_dof if learn_goal else 0)))
|
||||
self.action_space = gym.spaces.Box(low=-action_bounds, high=action_bounds, dtype=np.float32)
|
||||
|
||||
def initialize_mp(self, num_dof: int, duration: int, dt: float, num_basis: int = 5, start_pos: np.ndarray = None,
|
||||
final_pos: np.ndarray = None, alpha_phase: float = 2.):
|
||||
|
||||
phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
|
||||
basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=num_basis)
|
||||
|
||||
dmp = dmps.DMP(num_dof=num_dof, basis_generator=basis_generator, phase_generator=phase_generator,
|
||||
num_time_steps=int(duration / dt), dt=dt)
|
||||
|
||||
dmp.dmp_start_pos = start_pos.reshape((1, num_dof))
|
||||
|
||||
weights = np.zeros((num_basis, num_dof))
|
||||
goal_pos = np.zeros(num_dof) if self.learn_goal else final_pos
|
||||
|
||||
dmp.set_weights(weights, goal_pos)
|
||||
return dmp
|
||||
|
||||
def goal_and_weights(self, params):
|
||||
assert params.shape[-1] == self.action_space.shape[0]
|
||||
params = np.atleast_2d(params)
|
||||
|
||||
if self.learn_goal:
|
||||
goal_pos = params[0, -self.mp.num_dimensions:] # [num_dof]
|
||||
params = params[:, :-self.mp.num_dimensions] # [1,num_dof]
|
||||
# weight_matrix = np.reshape(params[:, :-self.num_dof], [self.num_basis, self.num_dof])
|
||||
else:
|
||||
goal_pos = self.mp.dmp_goal_pos.flatten()
|
||||
assert goal_pos is not None
|
||||
# weight_matrix = np.reshape(params, [self.num_basis, self.num_dof])
|
||||
|
||||
weight_matrix = np.reshape(params, self.mp.dmp_weights.shape)
|
||||
return goal_pos * self.goal_scale, weight_matrix * self.weights_scale
|
||||
|
||||
def mp_rollout(self, action):
|
||||
goal_pos, weight_matrix = self.goal_and_weights(action)
|
||||
self.mp.set_weights(weight_matrix, goal_pos)
|
||||
return self.mp.reference_trajectory(self.t)
|
||||
@@ -0,0 +1,106 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
||||
from alr_envs.utils.policies import get_policy_class
|
||||
|
||||
|
||||
class MPWrapper(gym.Wrapper, ABC):
|
||||
|
||||
def __init__(self,
|
||||
env: gym.Env,
|
||||
num_dof: int,
|
||||
duration: int = 1,
|
||||
dt: float = 0.01,
|
||||
# learn_goal: bool = False,
|
||||
post_traj_time: float = 0.,
|
||||
policy_type: str = None,
|
||||
weights_scale: float = 1.,
|
||||
**mp_kwargs
|
||||
|
||||
):
|
||||
super().__init__(env)
|
||||
|
||||
# self.num_dof = num_dof
|
||||
# self.num_basis = num_basis
|
||||
# self.duration = duration # seconds
|
||||
self.post_traj_steps = int(post_traj_time / dt)
|
||||
|
||||
self.mp = self.initialize_mp(num_dof, duration, dt, **mp_kwargs)
|
||||
self.weights_scale = weights_scale
|
||||
|
||||
policy_class = get_policy_class(policy_type)
|
||||
self.policy = policy_class(env)
|
||||
|
||||
# rendering
|
||||
self.render_mode = None
|
||||
self.render_kwargs = None
|
||||
|
||||
def step(self, action: np.ndarray):
|
||||
""" This function generates a trajectory based on a DMP and then does the usual loop over reset and step"""
|
||||
trajectory, velocity = self.mp_rollout(action)
|
||||
|
||||
if self.post_traj_steps > 0:
|
||||
trajectory = np.vstack([trajectory, np.tile(trajectory[-1, :], [self.post_traj_steps, 1])])
|
||||
velocity = np.vstack([velocity, np.zeros(shape=(self.post_traj_steps, self.dmp.num_dimensions))])
|
||||
|
||||
# self._trajectory = trajectory
|
||||
# self._velocity = velocity
|
||||
|
||||
rewards = 0
|
||||
infos = []
|
||||
|
||||
# TODO: @Max Why do we need this configure, states should be part of the model
|
||||
# self.env.configure(context)
|
||||
obs = self.env.reset()
|
||||
|
||||
for t, pos_vel in enumerate(zip(trajectory, velocity)):
|
||||
ac = self.policy.get_action(pos_vel[0], pos_vel[1])
|
||||
obs, rew, done, info = self.env.step(ac)
|
||||
rewards += rew
|
||||
infos.append(info)
|
||||
if self.render_mode:
|
||||
self.env.render(mode=self.render_mode, **self.render_kwargs)
|
||||
if done:
|
||||
break
|
||||
|
||||
done = True
|
||||
return obs, rewards, done, infos
|
||||
|
||||
def render(self, mode='human', **kwargs):
|
||||
"""Only set render options here, such that they can be used during the rollout.
|
||||
This only needs to be called once"""
|
||||
self.render_mode = mode
|
||||
self.render_kwargs = kwargs
|
||||
|
||||
def __call__(self, actions):
|
||||
return self.step(actions)
|
||||
# params = np.atleast_2d(params)
|
||||
# rewards = []
|
||||
# infos = []
|
||||
# for p, c in zip(params, contexts):
|
||||
# reward, info = self.rollout(p, c)
|
||||
# rewards.append(reward)
|
||||
# infos.append(info)
|
||||
#
|
||||
# return np.array(rewards), infos
|
||||
|
||||
@abstractmethod
|
||||
def mp_rollout(self, action):
|
||||
"""
|
||||
Generate trajectory and velocity based on the MP
|
||||
Returns:
|
||||
trajectory/positions, velocity
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def initialize_mp(self, num_dof: int, duration: int, dt: float, **kwargs):
|
||||
"""
|
||||
Create respective instance of MP
|
||||
Returns:
|
||||
MP instance
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
Reference in New Issue
Block a user