renaming to fancy_gym
This commit is contained in:
Fabian
@@ -0,0 +1,180 @@
|
||||
import os
|
||||
|
||||
os.environ["MUJOCO_GL"] = "egl"
|
||||
|
||||
from typing import Tuple, Optional
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
from gym import spaces
|
||||
from mp_pytorch.mp.mp_interfaces import MPInterface
|
||||
|
||||
from fancy_gym.black_box.controller.base_controller import BaseController
|
||||
from fancy_gym.black_box.raw_interface_wrapper import RawInterfaceWrapper
|
||||
from fancy_gym.utils.utils import get_numpy
|
||||
|
||||
|
||||
class BlackBoxWrapper(gym.ObservationWrapper):
|
||||
|
||||
def __init__(self,
|
||||
env: RawInterfaceWrapper,
|
||||
trajectory_generator: MPInterface,
|
||||
tracking_controller: BaseController,
|
||||
duration: float,
|
||||
verbose: int = 1,
|
||||
learn_sub_trajectories: bool = False,
|
||||
replanning_schedule: Optional[callable] = None,
|
||||
reward_aggregation: callable = np.sum
|
||||
):
|
||||
"""
|
||||
gym.Wrapper for leveraging a black box approach with a trajectory generator.
|
||||
|
||||
Args:
|
||||
env: The (wrapped) environment this wrapper is applied on
|
||||
trajectory_generator: Generates the full or partial trajectory
|
||||
tracking_controller: Translates the desired trajectory to raw action sequences
|
||||
duration: Length of the trajectory of the movement primitive in seconds
|
||||
verbose: level of detail for returned values in info dict.
|
||||
learn_sub_trajectories: Transforms full episode learning into learning sub-trajectories, similar to
|
||||
step-based learning
|
||||
replanning_schedule: callable that receives
|
||||
reward_aggregation: function that takes the np.ndarray of step rewards as input and returns the trajectory
|
||||
reward, default summation over all values.
|
||||
"""
|
||||
super().__init__(env)
|
||||
|
||||
self.duration = duration
|
||||
self.learn_sub_trajectories = learn_sub_trajectories
|
||||
self.do_replanning = replanning_schedule is not None
|
||||
self.replanning_schedule = replanning_schedule or (lambda *x: False)
|
||||
self.current_traj_steps = 0
|
||||
|
||||
# trajectory generation
|
||||
self.traj_gen = trajectory_generator
|
||||
self.tracking_controller = tracking_controller
|
||||
# self.time_steps = np.linspace(0, self.duration, self.traj_steps)
|
||||
# self.traj_gen.set_mp_times(self.time_steps)
|
||||
self.traj_gen.set_duration(np.array([self.duration]), np.array([self.dt]))
|
||||
|
||||
# reward computation
|
||||
self.reward_aggregation = reward_aggregation
|
||||
|
||||
# spaces
|
||||
self.return_context_observation = not (learn_sub_trajectories or self.do_replanning)
|
||||
self.traj_gen_action_space = self._get_traj_gen_action_space()
|
||||
self.action_space = self._get_action_space()
|
||||
self.observation_space = self._get_observation_space()
|
||||
|
||||
# rendering
|
||||
self.render_kwargs = {}
|
||||
self.verbose = verbose
|
||||
|
||||
def observation(self, observation):
|
||||
# return context space if we are
|
||||
obs = observation[self.env.context_mask] if self.return_context_observation else observation
|
||||
# cast dtype because metaworld returns incorrect that throws gym error
|
||||
return obs.astype(self.observation_space.dtype)
|
||||
|
||||
def get_trajectory(self, action: np.ndarray) -> Tuple:
|
||||
clipped_params = np.clip(action, self.traj_gen_action_space.low, self.traj_gen_action_space.high)
|
||||
self.traj_gen.set_params(clipped_params)
|
||||
# TODO: is this correct for replanning? Do we need to adjust anything here?
|
||||
self.traj_gen.set_boundary_conditions(
|
||||
bc_time=np.array(0) if not self.do_replanning else np.array([self.current_traj_steps * self.dt]),
|
||||
bc_pos=self.current_pos, bc_vel=self.current_vel)
|
||||
self.traj_gen.set_duration(None if self.learn_sub_trajectories else np.array([self.duration]),
|
||||
np.array([self.dt]))
|
||||
traj_dict = self.traj_gen.get_trajs(get_pos=True, get_vel=True)
|
||||
trajectory_tensor, velocity_tensor = traj_dict['pos'], traj_dict['vel']
|
||||
|
||||
return get_numpy(trajectory_tensor), get_numpy(velocity_tensor)
|
||||
|
||||
def _get_traj_gen_action_space(self):
|
||||
"""This function can be used to set up an individual space for the parameters of the traj_gen."""
|
||||
min_action_bounds, max_action_bounds = self.traj_gen.get_param_bounds()
|
||||
action_space = gym.spaces.Box(low=min_action_bounds.numpy(), high=max_action_bounds.numpy(),
|
||||
dtype=self.env.action_space.dtype)
|
||||
return action_space
|
||||
|
||||
def _get_action_space(self):
|
||||
"""
|
||||
This function can be used to modify the action space for considering actions which are not learned via movement
|
||||
primitives. E.g. ball releasing time for the beer pong task. By default, it is the parameter space of the
|
||||
movement primitive.
|
||||
Only needs to be overwritten if the action space needs to be modified.
|
||||
"""
|
||||
try:
|
||||
return self.traj_gen_action_space
|
||||
except AttributeError:
|
||||
return self._get_traj_gen_action_space()
|
||||
|
||||
def _get_observation_space(self):
|
||||
mask = self.env.context_mask
|
||||
if not self.return_context_observation:
|
||||
# return full observation
|
||||
mask = np.ones_like(mask, dtype=bool)
|
||||
min_obs_bound = self.env.observation_space.low[mask]
|
||||
max_obs_bound = self.env.observation_space.high[mask]
|
||||
return spaces.Box(low=min_obs_bound, high=max_obs_bound, dtype=self.env.observation_space.dtype)
|
||||
|
||||
def step(self, action: np.ndarray):
|
||||
""" This function generates a trajectory based on a MP and then does the usual loop over reset and step"""
|
||||
|
||||
# TODO remove this part, right now only needed for beer pong
|
||||
mp_params, env_spec_params = self.env.episode_callback(action, self.traj_gen)
|
||||
trajectory, velocity = self.get_trajectory(mp_params)
|
||||
|
||||
trajectory_length = len(trajectory)
|
||||
rewards = np.zeros(shape=(trajectory_length,))
|
||||
if self.verbose >= 2:
|
||||
actions = np.zeros(shape=(trajectory_length,) + self.env.action_space.shape)
|
||||
observations = np.zeros(shape=(trajectory_length,) + self.env.observation_space.shape,
|
||||
dtype=self.env.observation_space.dtype)
|
||||
|
||||
infos = dict()
|
||||
done = False
|
||||
|
||||
for t, (pos, vel) in enumerate(zip(trajectory, velocity)):
|
||||
step_action = self.tracking_controller.get_action(pos, vel, self.current_pos, self.current_vel)
|
||||
c_action = np.clip(step_action, self.env.action_space.low, self.env.action_space.high)
|
||||
obs, c_reward, done, info = self.env.step(c_action)
|
||||
rewards[t] = c_reward
|
||||
|
||||
if self.verbose >= 2:
|
||||
actions[t, :] = c_action
|
||||
observations[t, :] = obs
|
||||
|
||||
for k, v in info.items():
|
||||
elems = infos.get(k, [None] * trajectory_length)
|
||||
elems[t] = v
|
||||
infos[k] = elems
|
||||
|
||||
if self.render_kwargs:
|
||||
self.env.render(**self.render_kwargs)
|
||||
|
||||
if done or self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action,
|
||||
t + 1 + self.current_traj_steps):
|
||||
break
|
||||
|
||||
infos.update({k: v[:t + 1] for k, v in infos.items()})
|
||||
self.current_traj_steps += t + 1
|
||||
|
||||
if self.verbose >= 2:
|
||||
infos['positions'] = trajectory
|
||||
infos['velocities'] = velocity
|
||||
infos['step_actions'] = actions[:t + 1]
|
||||
infos['step_observations'] = observations[:t + 1]
|
||||
infos['step_rewards'] = rewards[:t + 1]
|
||||
|
||||
infos['trajectory_length'] = t + 1
|
||||
trajectory_return = self.reward_aggregation(rewards[:t + 1])
|
||||
return self.observation(obs), trajectory_return, done, infos
|
||||
|
||||
def render(self, **kwargs):
|
||||
"""Only set render options here, such that they can be used during the rollout.
|
||||
This only needs to be called once"""
|
||||
self.render_kwargs = kwargs
|
||||
|
||||
def reset(self, *, seed: Optional[int] = None, return_info: bool = False, options: Optional[dict] = None):
|
||||
self.current_traj_steps = 0
|
||||
return super(BlackBoxWrapper, self).reset()
|
||||
@@ -0,0 +1,4 @@
|
||||
class BaseController:
|
||||
|
||||
def get_action(self, des_pos, des_vel, c_pos, c_vel):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,24 @@
|
||||
import numpy as np
|
||||
|
||||
from fancy_gym.black_box.controller.base_controller import BaseController
|
||||
|
||||
|
||||
class MetaWorldController(BaseController):
|
||||
"""
|
||||
A Metaworld Controller. Using position and velocity information from a provided environment,
|
||||
the tracking_controller calculates a response based on the desired position and velocity.
|
||||
Unlike the other Controllers, this is a special tracking_controller for MetaWorld environments.
|
||||
They use a position delta for the xyz coordinates and a raw position for the gripper opening.
|
||||
|
||||
"""
|
||||
|
||||
def get_action(self, des_pos, des_vel, c_pos, c_vel):
|
||||
gripper_pos = des_pos[-1]
|
||||
|
||||
cur_pos = c_pos[:-1]
|
||||
xyz_pos = des_pos[:-1]
|
||||
|
||||
assert xyz_pos.shape == cur_pos.shape, \
|
||||
f"Mismatch in dimension between desired position {xyz_pos.shape} and current position {cur_pos.shape}"
|
||||
trq = np.hstack([(xyz_pos - cur_pos), gripper_pos])
|
||||
return trq
|
||||
@@ -0,0 +1,28 @@
|
||||
from typing import Union, Tuple
|
||||
|
||||
from fancy_gym.black_box.controller.base_controller import BaseController
|
||||
|
||||
|
||||
class PDController(BaseController):
|
||||
"""
|
||||
A PD-Controller. Using position and velocity information from a provided environment,
|
||||
the tracking_controller calculates a response based on the desired position and velocity
|
||||
|
||||
:param env: A position environment
|
||||
:param p_gains: Factors for the proportional gains
|
||||
:param d_gains: Factors for the differential gains
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
p_gains: Union[float, Tuple] = 1,
|
||||
d_gains: Union[float, Tuple] = 0.5):
|
||||
self.p_gains = p_gains
|
||||
self.d_gains = d_gains
|
||||
|
||||
def get_action(self, des_pos, des_vel, c_pos, c_vel):
|
||||
assert des_pos.shape == c_pos.shape, \
|
||||
f"Mismatch in dimension between desired position {des_pos.shape} and current position {c_pos.shape}"
|
||||
assert des_vel.shape == c_vel.shape, \
|
||||
f"Mismatch in dimension between desired velocity {des_vel.shape} and current velocity {c_vel.shape}"
|
||||
trq = self.p_gains * (des_pos - c_pos) + self.d_gains * (des_vel - c_vel)
|
||||
return trq
|
||||
@@ -0,0 +1,9 @@
|
||||
from fancy_gym.black_box.controller.base_controller import BaseController
|
||||
|
||||
|
||||
class PosController(BaseController):
|
||||
"""
|
||||
A Position Controller. The tracking_controller calculates a response only based on the desired position.
|
||||
"""
|
||||
def get_action(self, des_pos, des_vel, c_pos, c_vel):
|
||||
return des_pos
|
||||
@@ -0,0 +1,9 @@
|
||||
from fancy_gym.black_box.controller.base_controller import BaseController
|
||||
|
||||
|
||||
class VelController(BaseController):
|
||||
"""
|
||||
A Velocity Controller. The tracking_controller calculates a response only based on the desired velocity.
|
||||
"""
|
||||
def get_action(self, des_pos, des_vel, c_pos, c_vel):
|
||||
return des_vel
|
||||
@@ -0,0 +1,17 @@
|
||||
from mp_pytorch import PhaseGenerator, NormalizedRBFBasisGenerator, ZeroStartNormalizedRBFBasisGenerator
|
||||
from mp_pytorch.basis_gn.rhytmic_basis import RhythmicBasisGenerator
|
||||
|
||||
ALL_TYPES = ["rbf", "zero_rbf", "rhythmic"]
|
||||
|
||||
|
||||
def get_basis_generator(basis_generator_type: str, phase_generator: PhaseGenerator, **kwargs):
|
||||
basis_generator_type = basis_generator_type.lower()
|
||||
if basis_generator_type == "rbf":
|
||||
return NormalizedRBFBasisGenerator(phase_generator, **kwargs)
|
||||
elif basis_generator_type == "zero_rbf":
|
||||
return ZeroStartNormalizedRBFBasisGenerator(phase_generator, **kwargs)
|
||||
elif basis_generator_type == "rhythmic":
|
||||
return RhythmicBasisGenerator(phase_generator, **kwargs)
|
||||
else:
|
||||
raise ValueError(f"Specified basis generator type {basis_generator_type} not supported, "
|
||||
f"please choose one of {ALL_TYPES}.")
|
||||
@@ -0,0 +1,21 @@
|
||||
from fancy_gym.black_box.controller.meta_world_controller import MetaWorldController
|
||||
from fancy_gym.black_box.controller.pd_controller import PDController
|
||||
from fancy_gym.black_box.controller.pos_controller import PosController
|
||||
from fancy_gym.black_box.controller.vel_controller import VelController
|
||||
|
||||
ALL_TYPES = ["motor", "velocity", "position", "metaworld"]
|
||||
|
||||
|
||||
def get_controller(controller_type: str, **kwargs):
|
||||
controller_type = controller_type.lower()
|
||||
if controller_type == "motor":
|
||||
return PDController(**kwargs)
|
||||
elif controller_type == "velocity":
|
||||
return VelController()
|
||||
elif controller_type == "position":
|
||||
return PosController()
|
||||
elif controller_type == "metaworld":
|
||||
return MetaWorldController()
|
||||
else:
|
||||
raise ValueError(f"Specified controller type {controller_type} not supported, "
|
||||
f"please choose one of {ALL_TYPES}.")
|
||||
@@ -0,0 +1,20 @@
|
||||
from mp_pytorch import LinearPhaseGenerator, ExpDecayPhaseGenerator
|
||||
from mp_pytorch.phase_gn.rhythmic_phase_generator import RhythmicPhaseGenerator
|
||||
from mp_pytorch.phase_gn.smooth_phase_generator import SmoothPhaseGenerator
|
||||
|
||||
ALL_TYPES = ["linear", "exp", "rhythmic", "smooth"]
|
||||
|
||||
|
||||
def get_phase_generator(phase_generator_type, **kwargs):
|
||||
phase_generator_type = phase_generator_type.lower()
|
||||
if phase_generator_type == "linear":
|
||||
return LinearPhaseGenerator(**kwargs)
|
||||
elif phase_generator_type == "exp":
|
||||
return ExpDecayPhaseGenerator(**kwargs)
|
||||
elif phase_generator_type == "rhythmic":
|
||||
return RhythmicPhaseGenerator(**kwargs)
|
||||
elif phase_generator_type == "smooth":
|
||||
return SmoothPhaseGenerator(**kwargs)
|
||||
else:
|
||||
raise ValueError(f"Specified phase generator type {phase_generator_type} not supported, "
|
||||
f"please choose one of {ALL_TYPES}.")
|
||||
@@ -0,0 +1,21 @@
|
||||
from mp_pytorch.basis_gn.basis_generator import BasisGenerator
|
||||
from mp_pytorch.mp.dmp import DMP
|
||||
from mp_pytorch.mp.idmp import IDMP
|
||||
from mp_pytorch.mp.promp import ProMP
|
||||
|
||||
ALL_TYPES = ["promp", "dmp", "idmp"]
|
||||
|
||||
|
||||
def get_trajectory_generator(
|
||||
trajectory_generator_type: str, action_dim: int, basis_generator: BasisGenerator, **kwargs
|
||||
):
|
||||
trajectory_generator_type = trajectory_generator_type.lower()
|
||||
if trajectory_generator_type == "promp":
|
||||
return ProMP(basis_generator, action_dim, **kwargs)
|
||||
elif trajectory_generator_type == "dmp":
|
||||
return DMP(basis_generator, action_dim, **kwargs)
|
||||
elif trajectory_generator_type == 'idmp':
|
||||
return IDMP(basis_generator, action_dim, **kwargs)
|
||||
else:
|
||||
raise ValueError(f"Specified movement primitive type {trajectory_generator_type} not supported, "
|
||||
f"please choose one of {ALL_TYPES}.")
|
||||
@@ -0,0 +1,68 @@
|
||||
from abc import abstractmethod
|
||||
from typing import Union, Tuple
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
from mp_pytorch.mp.mp_interfaces import MPInterface
|
||||
|
||||
|
||||
class RawInterfaceWrapper(gym.Wrapper):
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def context_mask(self) -> np.ndarray:
|
||||
"""
|
||||
This function defines the contexts. The contexts are defined as specific observations.
|
||||
Returns:
|
||||
bool array representing the indices of the observations
|
||||
|
||||
"""
|
||||
return np.ones(self.env.observation_space.shape[0], dtype=bool)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
"""
|
||||
Returns the current position of the action/control dimension.
|
||||
The dimensionality has to match the action/control dimension.
|
||||
This is not required when exclusively using velocity control,
|
||||
it should, however, be implemented regardless.
|
||||
E.g. The joint positions that are directly or indirectly controlled by the action.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
|
||||
"""
|
||||
Returns the current velocity of the action/control dimension.
|
||||
The dimensionality has to match the action/control dimension.
|
||||
This is not required when exclusively using position control,
|
||||
it should, however, be implemented regardless.
|
||||
E.g. The joint velocities that are directly or indirectly controlled by the action.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
def dt(self) -> float:
|
||||
"""
|
||||
Control frequency of the environment
|
||||
Returns: float
|
||||
|
||||
"""
|
||||
return self.env.dt
|
||||
|
||||
def episode_callback(self, action: np.ndarray, traj_gen: MPInterface) -> Tuple[
|
||||
np.ndarray, Union[np.ndarray, None]]:
|
||||
"""
|
||||
Used to extract the parameters for the movement primitive and other parameters from an action array which might
|
||||
include other actions like ball releasing time for the beer pong environment.
|
||||
This only needs to be overwritten if the action space is modified.
|
||||
Args:
|
||||
action: a vector instance of the whole action space, includes traj_gen parameters and additional parameters if
|
||||
specified, else only traj_gen parameters
|
||||
|
||||
Returns:
|
||||
Tuple: mp_arguments and other arguments
|
||||
"""
|
||||
return action, None
|
||||
Reference in New Issue
Block a user