current state

alr_envs/black_box/black_box_wrapper.py

@@ -0,0 +1,192 @@
+from abc import ABC
+from typing import Tuple, Union
+
+import gym
+import numpy as np
+from gym import spaces
+from mp_pytorch.mp.mp_interfaces import MPInterface
+
+from alr_envs.black_box.controller.base_controller import BaseController
+from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
+from alr_envs.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: Union[None, 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.replanning_schedule = replanning_schedule
+        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 (self.learn_sub_trajectories or replanning_schedule)
+        self.traj_gen_action_space = self.get_traj_gen_action_space()
+        self.action_space = self.get_action_space()
+        self.observation_space = spaces.Box(low=self.env.observation_space.low[self.env.context_mask],
+                                            high=self.env.observation_space.high[self.env.context_mask],
+                                            dtype=self.env.observation_space.dtype)
+
+        # rendering
+        self.render_kwargs = {}
+        self.verbose = verbose
+
+    def observation(self, observation):
+        # return context space if we are
+        return observation[self.env.context_mask] if self.return_context_observation else observation
+
+    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: Bruce said DMP, ProMP, ProDMP can have 0 bc_time for sequencing
+        # TODO Check with Bruce for replanning
+        self.traj_gen.set_boundary_conditions(
+            bc_time=np.zeros((1,)) if not self.replanning_schedule else self.current_traj_steps * self.dt,
+            bc_pos=self.current_pos, bc_vel=self.current_vel)
+        # TODO: is this correct for replanning? Do we need to adjust anything here?
+        self.traj_gen.set_duration(None if self.learn_sub_trajectories else 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()
+        mp_action_space = gym.spaces.Box(low=min_action_bounds.numpy(), high=max_action_bounds.numpy(),
+                                         dtype=np.float32)
+        return mp_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 motion
+        primitives. E.g. ball releasing time for the beer pong task. By default, it is the parameter space of the
+        motion 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 step(self, action: np.ndarray):
+        """ This function generates a trajectory based on a MP and then does the usual loop over reset and step"""
+
+        # agent to learn when to release the ball
+        mp_params, env_spec_params = self._episode_callback(action)
+        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[0], pos_vel[1], self.current_pos,
+                                                              self.current_vel)
+            step_action = self._step_callback(t, env_spec_params, step_action)  # include possible callback info
+            c_action = np.clip(step_action, self.env.action_space.low, self.env.action_space.high)
+            # print('step/clipped action ratio: ', step_action/c_action)
+            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.render(**self.render_kwargs)
+
+            if done:
+                break
+
+            if self.replanning_schedule and 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['trajectory'] = trajectory
+            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 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
+        # self.env.render(mode=self.render_mode, **self.render_kwargs)
+        self.env.render(**kwargs)
+
+    def reset(self, **kwargs):
+        self.current_traj_steps = 0
+        super(BlackBoxWrapper, self).reset(**kwargs)
+
+    def plot_trajs(self, des_trajs, des_vels):
+        import matplotlib.pyplot as plt
+        import matplotlib
+        matplotlib.use('TkAgg')
+        pos_fig = plt.figure('positions')
+        vel_fig = plt.figure('velocities')
+        for i in range(des_trajs.shape[1]):
+            plt.figure(pos_fig.number)
+            plt.subplot(des_trajs.shape[1], 1, i + 1)
+            plt.plot(np.ones(des_trajs.shape[0]) * self.current_pos[i])
+            plt.plot(des_trajs[:, i])
+
+            plt.figure(vel_fig.number)
+            plt.subplot(des_vels.shape[1], 1, i + 1)
+            plt.plot(np.ones(des_trajs.shape[0]) * self.current_vel[i])
+            plt.plot(des_vels[:, i])

alr_envs/black_box/controller/base_controller.py

@@ -0,0 +1,4 @@
+class BaseController:
+
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        raise NotImplementedError

alr_envs/black_box/controller/controller_factory.py

@@ -0,0 +1,21 @@
+from alr_envs.black_box.controller.meta_world_controller import MetaWorldController
+from alr_envs.black_box.controller.pd_controller import PDController
+from alr_envs.black_box.controller.vel_controller import VelController
+from alr_envs.black_box.controller.pos_controller import PosController
+
+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}.")

alr_envs/black_box/controller/meta_world_controller.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+from alr_envs.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.
+
+    :param env: A position environment
+    """
+
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        gripper_pos = des_pos[-1]
+
+        cur_pos = env.current_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

alr_envs/black_box/controller/pd_controller.py

@@ -0,0 +1,28 @@
+from typing import Union, Tuple
+
+from alr_envs.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

alr_envs/black_box/controller/pos_controller.py

@@ -0,0 +1,9 @@
+from alr_envs.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

alr_envs/black_box/controller/vel_controller.py

@@ -0,0 +1,9 @@
+from alr_envs.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

alr_envs/black_box/factory/basis_generator_factory.py

@@ -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}.")

alr_envs/black_box/factory/phase_generator_factory.py

@@ -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}.")

alr_envs/black_box/factory/trajectory_generator_factory.py

@@ -0,0 +1,22 @@
+from mp_pytorch.mp.dmp import DMP
+from mp_pytorch.mp.promp import ProMP
+from mp_pytorch.mp.idmp import IDMP
+
+from mp_pytorch.basis_gn.basis_generator import BasisGenerator
+
+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}.")

alr_envs/black_box/raw_interface_wrapper.py

@@ -0,0 +1,88 @@
+from typing import Union, Tuple
+
+import gym
+import numpy as np
+from abc import abstractmethod
+
+
+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 do_replanning(self, pos, vel, s, a, t):
+        # return t % 100 == 0
+        # return bool(self.replanning_model(s))
+        return False
+
+    def _episode_callback(self, action: np.ndarray) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
+        """
+        Used to extract the parameters for the motion 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
+
+    def _step_callback(self, t: int, env_spec_params: Union[np.ndarray, None], step_action: np.ndarray) -> Union[
+        np.ndarray]:
+        """
+        This function can be used to modify the step_action with additional parameters e.g. releasing the ball in the
+        Beerpong env. The parameters used should not be part of the motion primitive parameters.
+        Returns step_action by default, can be overwritten in individual mp_wrappers.
+        Args:
+            t: the current time step of the episode
+            env_spec_params: the environment specific parameter, as defined in function _episode_callback
+            (e.g. ball release time in Beer Pong)
+            step_action: the current step-based action
+
+        Returns:
+            modified step action
+        """
+        return step_action