added MPEnv
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
import alr_envs.classic_control.hole_reacher as hr
|
||||
import alr_envs.classic_control.viapoint_reacher as vpr
|
||||
from alr_envs.utils.wrapper.dmp_wrapper import DmpWrapper
|
||||
from alr_envs.utils.wrapper.detpmp_wrapper import DetPMPWrapper
|
||||
from alr_envs.utils.mps.dmp_wrapper import DmpWrapper
|
||||
from alr_envs.utils.mps.detpmp_wrapper import DetPMPWrapper
|
||||
import numpy as np
|
||||
|
||||
|
||||
@@ -65,7 +65,7 @@ def make_holereacher_env(rank, seed=0):
|
||||
dt=_env.dt,
|
||||
learn_goal=True,
|
||||
alpha_phase=2,
|
||||
start_pos=_env.start_pos,
|
||||
start_pos=_env._start_pos,
|
||||
policy_type="velocity",
|
||||
weights_scale=50,
|
||||
goal_scale=0.1
|
||||
@@ -105,7 +105,7 @@ def make_holereacher_fix_goal_env(rank, seed=0):
|
||||
learn_goal=False,
|
||||
final_pos=np.array([2.02669572, -1.25966385, -1.51618198, -0.80946476, 0.02012344]),
|
||||
alpha_phase=2,
|
||||
start_pos=_env.start_pos,
|
||||
start_pos=_env._start_pos,
|
||||
policy_type="velocity",
|
||||
weights_scale=50,
|
||||
goal_scale=1
|
||||
@@ -142,7 +142,7 @@ def make_holereacher_env_pmp(rank, seed=0):
|
||||
num_basis=5,
|
||||
width=0.02,
|
||||
policy_type="velocity",
|
||||
start_pos=_env.start_pos,
|
||||
start_pos=_env._start_pos,
|
||||
duration=2,
|
||||
post_traj_time=0,
|
||||
dt=_env.dt,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from alr_envs.utils.wrapper.dmp_wrapper import DmpWrapper
|
||||
from alr_envs.utils.wrapper.detpmp_wrapper import DetPMPWrapper
|
||||
from alr_envs.utils.mps.dmp_wrapper import DmpWrapper
|
||||
from alr_envs.utils.mps.detpmp_wrapper import DetPMPWrapper
|
||||
import gym
|
||||
from gym.vector.utils import write_to_shared_memory
|
||||
import sys
|
||||
|
||||
@@ -2,17 +2,18 @@ import gym
|
||||
import numpy as np
|
||||
from mp_lib import det_promp
|
||||
|
||||
from alr_envs.utils.wrapper.mp_wrapper import MPWrapper
|
||||
from alr_envs.utils.mps.mp_environments import MPEnv
|
||||
from alr_envs.utils.mps.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):
|
||||
def __init__(self, env: MPEnv, num_dof: int, num_basis: int, width: int, start_pos=None, duration: int = 1,
|
||||
dt: float = 0.01, post_traj_time: float = 0., policy_type: str = None, weights_scale: float = 1.,
|
||||
zero_start: bool = False, zero_goal: bool = 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)
|
||||
super().__init__(env, num_dof, dt, duration, 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, **mp_kwargs)
|
||||
|
||||
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)
|
||||
@@ -1,19 +1,18 @@
|
||||
from mp_lib.phase import ExpDecayPhaseGenerator
|
||||
from mp_lib.basis import DMPBasisGenerator
|
||||
from mp_lib import dmps
|
||||
import numpy as np
|
||||
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.wrapper.mp_wrapper import MPWrapper
|
||||
from alr_envs.utils.mps.mp_environments import MPEnv
|
||||
from alr_envs.utils.mps.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,
|
||||
def __init__(self, env: MPEnv, num_dof: int, num_basis: int,
|
||||
duration: int = 1, alpha_phase: float = 2., dt: float = None,
|
||||
learn_goal: bool = False, return_to_start: bool = False, post_traj_time: float = 0.,
|
||||
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):
|
||||
|
||||
@@ -23,8 +22,6 @@ class DmpWrapper(MPWrapper):
|
||||
env:
|
||||
num_dof:
|
||||
num_basis:
|
||||
start_pos:
|
||||
final_pos:
|
||||
duration:
|
||||
alpha_phase:
|
||||
dt:
|
||||
@@ -37,30 +34,17 @@ class DmpWrapper(MPWrapper):
|
||||
self.learn_goal = learn_goal
|
||||
dt = env.dt if hasattr(env, "dt") else dt
|
||||
assert dt is not None
|
||||
# start_pos = start_pos if start_pos is not None else env.start_pos if hasattr(env, "start_pos") else None
|
||||
# TODO: assert start_pos is not None # start_pos will be set in initialize, do we need this here?
|
||||
# if learn_goal:
|
||||
# final_pos = np.zeros_like(start_pos) # arbitrary, will be learned
|
||||
# final_pos = np.zeros((1, num_dof)) # arbitrary, will be learned
|
||||
# else:
|
||||
# final_pos = final_pos if final_pos is not None else start_pos if return_to_start else None
|
||||
# assert final_pos is not None
|
||||
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, render_mode,
|
||||
num_basis=num_basis,
|
||||
# start_pos=start_pos, final_pos=final_pos,
|
||||
alpha_phase=alpha_phase,
|
||||
bandwidth_factor=bandwidth_factor)
|
||||
super().__init__(env, num_dof, dt, duration, post_traj_time, policy_type, weights_scale, render_mode,
|
||||
num_basis=num_basis, alpha_phase=alpha_phase, bandwidth_factor=bandwidth_factor)
|
||||
|
||||
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., bandwidth_factor: float = 3.):
|
||||
def initialize_mp(self, num_dof: int, duration: int, dt: float, num_basis: int = 5, alpha_phase: float = 2.,
|
||||
bandwidth_factor: int = 3):
|
||||
|
||||
phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
|
||||
basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=num_basis,
|
||||
@@ -69,15 +53,6 @@ class DmpWrapper(MPWrapper):
|
||||
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))
|
||||
# in a contextual environment, the start_pos may be not fixed, set in mp_rollout?
|
||||
# TODO: Should we set start_pos in init at all? It's only used after calling rollout anyway...
|
||||
# dmp.dmp_start_pos = start_pos.reshape((1, num_dof)) if start_pos is not None else np.zeros((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):
|
||||
@@ -87,18 +62,15 @@ class DmpWrapper(MPWrapper):
|
||||
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.env.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)
|
||||
weight_matrix = np.reshape(params, self.mp.dmp_weights.shape) # [num_basis, num_dof]
|
||||
return goal_pos * self.goal_scale, weight_matrix * self.weights_scale
|
||||
|
||||
def mp_rollout(self, action):
|
||||
# if self.mp.start_pos is None:
|
||||
self.mp.dmp_start_pos = self.env.init_qpos # start_pos
|
||||
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)
|
||||
@@ -0,0 +1,33 @@
|
||||
from abc import abstractmethod
|
||||
from typing import Union
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
||||
|
||||
class MPEnv(gym.Env):
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def corrected_obs_index(self):
|
||||
"""Returns boolean value for each observation entry
|
||||
whether the observation is returned by the DMP for the contextual case or not.
|
||||
This effectively allows to filter unwanted or unnecessary observations from the full step-based case.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def start_pos(self) -> Union[float, int, np.ndarray]:
|
||||
"""
|
||||
Returns the current 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
|
||||
@@ -1,32 +1,18 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
||||
from alr_envs.utils.mps.mp_environments import MPEnv
|
||||
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 = None,
|
||||
post_traj_time: float = 0.,
|
||||
policy_type: str = None,
|
||||
weights_scale: float = 1.,
|
||||
render_mode: str = None,
|
||||
**mp_kwargs
|
||||
):
|
||||
def __init__(self, env: MPEnv, num_dof: int, dt: float, duration: int = 1, post_traj_time: float = 0.,
|
||||
policy_type: str = None, weights_scale: float = 1., render_mode: str = None, **mp_kwargs):
|
||||
super().__init__(env)
|
||||
|
||||
# self.num_dof = num_dof
|
||||
# self.num_basis = num_basis
|
||||
# self.duration = duration # seconds
|
||||
|
||||
# dt = env.dt if hasattr(env, "dt") else dt
|
||||
assert dt is not None # this should never happen as MPWrapper is a base class
|
||||
self.post_traj_steps = int(post_traj_time / dt)
|
||||
|
||||
@@ -40,8 +26,11 @@ class MPWrapper(gym.Wrapper, ABC):
|
||||
self.render_mode = render_mode
|
||||
self.render_kwargs = {}
|
||||
|
||||
# TODO: not yet final
|
||||
# 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 = []
|
||||
@@ -63,7 +52,7 @@ class MPWrapper(gym.Wrapper, ABC):
|
||||
|
||||
def reset(self):
|
||||
obs = self.env.reset()
|
||||
return obs
|
||||
return obs[self.env]
|
||||
|
||||
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"""
|
||||
@@ -77,15 +66,9 @@ class MPWrapper(gym.Wrapper, ABC):
|
||||
# self._velocity = velocity
|
||||
|
||||
rewards = 0
|
||||
# infos = defaultdict(list)
|
||||
|
||||
# TODO: @Max Why do we need this configure, states should be part of the model
|
||||
# TODO: Ask Onur if the context distribution needs to be outside the environment
|
||||
# TODO: For now create a new env with each context
|
||||
# TODO: Explicitly call reset before step to obtain context from obs?
|
||||
# self.env.configure(context)
|
||||
# obs = self.env.reset()
|
||||
info = {}
|
||||
# create random obs as the reset function is called externally
|
||||
obs = self.env.observation_space.sample()
|
||||
|
||||
for t, pos_vel in enumerate(zip(trajectory, velocity)):
|
||||
ac = self.policy.get_action(pos_vel[0], pos_vel[1])
|
||||
@@ -107,18 +90,6 @@ class MPWrapper(gym.Wrapper, ABC):
|
||||
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):
|
||||
"""
|
||||
Reference in New Issue
Block a user