start refactor and biac dev merge
This commit is contained in:
Maximilian Huettenrauch
@@ -1,40 +0,0 @@
|
||||
from abc import abstractmethod, ABC
|
||||
from typing import Union
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
||||
|
||||
class AlrEnv(gym.Env, ABC):
|
||||
|
||||
@property
|
||||
def active_obs(self):
|
||||
"""Returns boolean mask for each observation entry
|
||||
whether the observation is returned for the contextual case or not.
|
||||
This effectively allows to filter unwanted or unnecessary observations from the full step-based case.
|
||||
"""
|
||||
return np.ones(self.observation_space.shape, dtype=bool)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def start_pos(self) -> Union[float, int, np.ndarray]:
|
||||
"""
|
||||
Returns the starting position of the joints
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
def goal_pos(self) -> Union[float, int, np.ndarray]:
|
||||
"""
|
||||
Returns the current final position of the joints for the MP.
|
||||
By default this returns the starting position.
|
||||
"""
|
||||
return self.start_pos
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def dt(self) -> Union[float, int]:
|
||||
"""
|
||||
Returns the time between two simulated steps of the environment
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -1,54 +0,0 @@
|
||||
import gym
|
||||
import numpy as np
|
||||
from mp_lib import det_promp
|
||||
|
||||
from alr_envs.utils.mps.alr_env import AlrEnv
|
||||
from alr_envs.utils.mps.mp_wrapper import MPWrapper
|
||||
|
||||
|
||||
class DetPMPWrapper(MPWrapper):
|
||||
def __init__(self, env: AlrEnv, num_dof, num_basis, width, start_pos=None, duration=1, post_traj_time=0.,
|
||||
policy_type=None, weights_scale=1, zero_start=False, zero_goal=False, learn_mp_length: bool =True,
|
||||
**mp_kwargs):
|
||||
self.duration = duration # seconds
|
||||
|
||||
super().__init__(env=env, num_dof=num_dof, duration=duration, post_traj_time=post_traj_time,
|
||||
policy_type=policy_type, weights_scale=weights_scale, num_basis=num_basis,
|
||||
width=width, zero_start=zero_start, zero_goal=zero_goal,
|
||||
**mp_kwargs)
|
||||
|
||||
self.learn_mp_length = learn_mp_length
|
||||
if self.learn_mp_length:
|
||||
parameter_space_shape = (1+num_basis*num_dof,)
|
||||
else:
|
||||
parameter_space_shape = (num_basis * num_dof,)
|
||||
self.min_param = -np.inf
|
||||
self.max_param = np.inf
|
||||
self.parameterization_space = gym.spaces.Box(low=self.min_param, high=self.max_param,
|
||||
shape=parameter_space_shape, dtype=np.float32)
|
||||
|
||||
self.start_pos = start_pos
|
||||
|
||||
def initialize_mp(self, num_dof: int, duration: int, num_basis: int = 5, width: float = None,
|
||||
zero_start: bool = False, zero_goal: bool = False, **kwargs):
|
||||
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):
|
||||
if self.learn_mp_length:
|
||||
duration = max(1, self.duration*abs(action[0]))
|
||||
params = np.reshape(action[1:], (self.mp.n_basis, -1)) * self.weights_scale # TODO: Fix Bug when zero_start is true
|
||||
else:
|
||||
duration = self.duration
|
||||
params = np.reshape(action, (self.mp.n_basis, -1)) * self.weights_scale # TODO: Fix Bug when zero_start is true
|
||||
self.mp.set_weights(1., params)
|
||||
_, des_pos, des_vel, _ = self.mp.compute_trajectory(frequency=max(1, duration))
|
||||
if self.mp.zero_start:
|
||||
des_pos += self.start_pos
|
||||
|
||||
return des_pos, des_vel
|
||||
@@ -1,76 +0,0 @@
|
||||
import gym
|
||||
import numpy as np
|
||||
from mp_lib import dmps
|
||||
from mp_lib.basis import DMPBasisGenerator
|
||||
from mp_lib.phase import ExpDecayPhaseGenerator
|
||||
|
||||
from alr_envs.utils.mps.alr_env import AlrEnv
|
||||
from alr_envs.utils.mps.mp_wrapper import MPWrapper
|
||||
|
||||
|
||||
class DmpWrapper(MPWrapper):
|
||||
|
||||
def __init__(self, env: AlrEnv, num_dof: int, num_basis: int,
|
||||
duration: int = 1, alpha_phase: float = 2., dt: float = None,
|
||||
learn_goal: bool = False, post_traj_time: float = 0.,
|
||||
weights_scale: float = 1., goal_scale: float = 1., bandwidth_factor: float = 3.,
|
||||
policy_type: str = None, render_mode: str = None):
|
||||
|
||||
"""
|
||||
This Wrapper generates a trajectory based on a DMP and will only return episodic performances.
|
||||
Args:
|
||||
env:
|
||||
num_dof:
|
||||
num_basis:
|
||||
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=env, num_dof=num_dof, duration=duration, post_traj_time=post_traj_time,
|
||||
policy_type=policy_type, weights_scale=weights_scale, render_mode=render_mode,
|
||||
num_basis=num_basis, alpha_phase=alpha_phase, bandwidth_factor=bandwidth_factor)
|
||||
|
||||
action_bounds = np.inf * np.ones((np.prod(self.mp.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, num_basis: int, alpha_phase: float = 2.,
|
||||
bandwidth_factor: int = 3, **kwargs):
|
||||
|
||||
phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
|
||||
basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=num_basis,
|
||||
basis_bandwidth_factor=bandwidth_factor)
|
||||
|
||||
dmp = dmps.DMP(num_dof=num_dof, basis_generator=basis_generator, phase_generator=phase_generator,
|
||||
dt=self.dt)
|
||||
|
||||
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]
|
||||
else:
|
||||
goal_pos = self.env.goal_pos
|
||||
assert goal_pos is not None
|
||||
|
||||
weight_matrix = np.reshape(params, self.mp.weights.shape) # [num_basis, num_dof]
|
||||
return goal_pos * self.goal_scale, weight_matrix * self.weights_scale
|
||||
|
||||
def mp_rollout(self, action):
|
||||
self.mp.dmp_start_pos = self.env.start_pos
|
||||
goal_pos, weight_matrix = self.goal_and_weights(action)
|
||||
self.mp.set_weights(weight_matrix, goal_pos)
|
||||
return self.mp.reference_trajectory(self.t)
|
||||
@@ -1,134 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Union
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
||||
from alr_envs.utils.mps.alr_env import AlrEnv
|
||||
from alr_envs.utils.policies import get_policy_class, BaseController
|
||||
|
||||
|
||||
class MPWrapper(gym.Wrapper, ABC):
|
||||
"""
|
||||
Base class for movement primitive based gym.Wrapper implementations.
|
||||
|
||||
:param env: The (wrapped) environment this wrapper is applied on
|
||||
:param num_dof: Dimension of the action space of the wrapped env
|
||||
:param duration: Number of timesteps in the trajectory of the movement primitive
|
||||
:param post_traj_time: Time for which the last position of the trajectory is fed to the environment to continue
|
||||
simulation
|
||||
:param policy_type: Type or object defining the policy that is used to generate action based on the trajectory
|
||||
:param weight_scale: Scaling parameter for the actions given to this wrapper
|
||||
:param render_mode: Equivalent to gym render mode
|
||||
"""
|
||||
def __init__(self,
|
||||
env: AlrEnv,
|
||||
num_dof: int,
|
||||
duration: int = 1,
|
||||
post_traj_time: float = 0.,
|
||||
policy_type: Union[str, BaseController] = None,
|
||||
weights_scale: float = 1.,
|
||||
render_mode: str = None,
|
||||
**mp_kwargs
|
||||
):
|
||||
super().__init__(env)
|
||||
|
||||
# adjust observation space to reduce version
|
||||
obs_sp = self.env.observation_space
|
||||
self.observation_space = gym.spaces.Box(low=obs_sp.low[self.env.active_obs],
|
||||
high=obs_sp.high[self.env.active_obs],
|
||||
dtype=obs_sp.dtype)
|
||||
|
||||
self.post_traj_steps = int(post_traj_time / env.dt)
|
||||
|
||||
self.mp = self.initialize_mp(num_dof=num_dof, duration=duration, **mp_kwargs)
|
||||
self.weights_scale = weights_scale
|
||||
|
||||
if type(policy_type) is str:
|
||||
self.policy = get_policy_class(policy_type, env, mp_kwargs)
|
||||
else:
|
||||
self.policy = policy_type
|
||||
|
||||
# rendering
|
||||
self.render_mode = render_mode
|
||||
self.render_kwargs = {}
|
||||
|
||||
# TODO: @Max I think this should not be in this class, this functionality should be part of your sampler.
|
||||
def __call__(self, params, contexts=None):
|
||||
"""
|
||||
Can be used to provide a batch of parameter sets
|
||||
"""
|
||||
params = np.atleast_2d(params)
|
||||
obs = []
|
||||
rewards = []
|
||||
dones = []
|
||||
infos = []
|
||||
# for p, c in zip(params, contexts):
|
||||
for p in params:
|
||||
# self.configure(c)
|
||||
ob, reward, done, info = self.step(p)
|
||||
obs.append(ob)
|
||||
rewards.append(reward)
|
||||
dones.append(done)
|
||||
infos.append(info)
|
||||
|
||||
return obs, np.array(rewards), dones, infos
|
||||
|
||||
def reset(self):
|
||||
return self.env.reset()[self.env.active_obs]
|
||||
|
||||
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.mp.num_dimensions))])
|
||||
|
||||
trajectory_length = len(trajectory)
|
||||
actions = np.zeros(shape=(trajectory_length, self.mp.num_dimensions))
|
||||
observations= np.zeros(shape=(trajectory_length,) + self.env.observation_space.shape)
|
||||
rewards = np.zeros(shape=(trajectory_length,))
|
||||
trajectory_return = 0
|
||||
infos = dict(step_infos =[])
|
||||
for t, pos_vel in enumerate(zip(trajectory, velocity)):
|
||||
actions[t,:] = self.policy.get_action(pos_vel[0], pos_vel[1])
|
||||
observations[t,:], rewards[t], done, info = self.env.step(actions[t,:])
|
||||
trajectory_return += rewards[t]
|
||||
infos['step_infos'].append(info)
|
||||
if self.render_mode:
|
||||
self.env.render(mode=self.render_mode, **self.render_kwargs)
|
||||
if done:
|
||||
break
|
||||
|
||||
infos['step_actions'] = actions[:t+1]
|
||||
infos['step_observations'] = observations[:t+1]
|
||||
infos['step_rewards'] = rewards[:t+1]
|
||||
infos['trajectory_length'] = t+1
|
||||
done = True
|
||||
return observations[t][self.env.active_obs], trajectory_return, 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
|
||||
|
||||
@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: float, **kwargs):
|
||||
"""
|
||||
Create respective instance of MP
|
||||
Returns:
|
||||
MP instance
|
||||
"""
|
||||
|
||||
raise NotImplementedError
|
||||
Reference in New Issue
Block a user