Merge with stochastic search branch

alr_envs/utils/dmp_async_vec_env.py

@@ -0,0 +1,173 @@
+import gym
+from gym.error import (AlreadyPendingCallError, NoAsyncCallError)
+from gym.vector.utils import concatenate, create_empty_array
+from gym.vector.async_vector_env import AsyncState
+import numpy as np
+import multiprocessing as mp
+import sys
+
+
+def _worker(index, env_fn, pipe, parent_pipe, shared_memory, error_queue):
+    assert shared_memory is None
+    env = env_fn()
+    parent_pipe.close()
+    try:
+        while True:
+            command, data = pipe.recv()
+            if command == 'reset':
+                observation = env.reset()
+                pipe.send((observation, True))
+            elif command == 'step':
+                observation, reward, done, info = env.step(data)
+                if done:
+                    observation = env.reset()
+                pipe.send(((observation, reward, done, info), True))
+            elif command == 'rollout':
+                rewards = []
+                infos = []
+                for p, c in zip(*data):
+                    reward, info = env.rollout(p, c)
+                    rewards.append(reward)
+                    infos.append(info)
+                pipe.send(((rewards, infos), (True, ) * len(rewards)))
+            elif command == 'seed':
+                env.seed(data)
+                pipe.send((None, True))
+            elif command == 'close':
+                env.close()
+                pipe.send((None, True))
+                break
+            elif command == 'idle':
+                pipe.send((None, True))
+            elif command == '_check_observation_space':
+                pipe.send((data == env.observation_space, True))
+            else:
+                raise RuntimeError('Received unknown command `{0}`. Must '
+                    'be one of {`reset`, `step`, `seed`, `close`, '
+                    '`_check_observation_space`}.'.format(command))
+    except (KeyboardInterrupt, Exception):
+        error_queue.put((index,) + sys.exc_info()[:2])
+        pipe.send((None, False))
+    finally:
+        env.close()
+
+
+class DmpAsyncVectorEnv(gym.vector.AsyncVectorEnv):
+    def __init__(self, env_fns, n_samples, observation_space=None, action_space=None,
+                 shared_memory=False, copy=True, context="spawn", daemon=True, worker=_worker):
+        super(DmpAsyncVectorEnv, self).__init__(env_fns,
+                                                observation_space=observation_space,
+                                                action_space=action_space,
+                                                shared_memory=shared_memory,
+                                                copy=copy,
+                                                context=context,
+                                                daemon=daemon,
+                                                worker=worker)
+
+        # we need to overwrite the number of samples as we may sample more than num_envs
+        self.observations = create_empty_array(self.single_observation_space,
+                                               n=n_samples,
+                                               fn=np.zeros)
+
+    def __call__(self, params, contexts=None):
+        return self.rollout(params, contexts)
+
+    def rollout_async(self, params, contexts):
+        """
+        Parameters
+        ----------
+        params : iterable of samples from `action_space`
+            List of actions.
+        """
+        self._assert_is_running()
+        if self._state != AsyncState.DEFAULT:
+            raise AlreadyPendingCallError('Calling `rollout_async` while waiting '
+                                          'for a pending call to `{0}` to complete.'.format(
+                self._state.value), self._state.value)
+
+        params = np.atleast_2d(params)
+        split_params = np.array_split(params, np.minimum(len(params), self.num_envs))
+        if contexts is None:
+            split_contexts = np.array_split([None, ] * len(params), np.minimum(len(params), self.num_envs))
+        else:
+            split_contexts = np.array_split(contexts, np.minimum(len(contexts), self.num_envs))
+
+        assert np.all([len(p) == len(c) for p, c in zip(split_params, split_contexts)])
+        for pipe, param, context in zip(self.parent_pipes, split_params, split_contexts):
+            pipe.send(('rollout', (param, context)))
+        for pipe in self.parent_pipes[len(split_params):]:
+            pipe.send(('idle', None))
+        self._state = AsyncState.WAITING_ROLLOUT
+
+    def rollout_wait(self, timeout=None):
+        """
+        Parameters
+        ----------
+        timeout : int or float, optional
+            Number of seconds before the call to `step_wait` times out. If
+            `None`, the call to `step_wait` never times out.
+
+        Returns
+        -------
+        observations : sample from `observation_space`
+            A batch of observations from the vectorized environment.
+
+        rewards : `np.ndarray` instance (dtype `np.float_`)
+            A vector of rewards from the vectorized environment.
+
+        dones : `np.ndarray` instance (dtype `np.bool_`)
+            A vector whose entries indicate whether the episode has ended.
+
+        infos : list of dict
+            A list of auxiliary diagnostic information.
+        """
+        self._assert_is_running()
+        if self._state != AsyncState.WAITING_ROLLOUT:
+            raise NoAsyncCallError('Calling `rollout_wait` without any prior call '
+                                   'to `rollout_async`.', AsyncState.WAITING_ROLLOUT.value)
+
+        if not self._poll(timeout):
+            self._state = AsyncState.DEFAULT
+            raise mp.TimeoutError('The call to `rollout_wait` has timed out after '
+                                  '{0} second{1}.'.format(timeout, 's' if timeout > 1 else ''))
+
+        results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes])
+        results = [r for r in results if r is not None]
+        self._raise_if_errors(successes)
+        self._state = AsyncState.DEFAULT
+
+        rewards, infos = [_flatten_list(r) for r in zip(*results)]
+
+        # for now, we ignore the observations and only return the rewards
+
+        # if not self.shared_memory:
+        #     self.observations = concatenate(observations_list, self.observations,
+        #                                     self.single_observation_space)
+
+        # return (deepcopy(self.observations) if self.copy else self.observations,
+        #         np.array(rewards), np.array(dones, dtype=np.bool_), infos)
+
+        return np.array(rewards), infos
+
+    def rollout(self, actions, contexts):
+        self.rollout_async(actions, contexts)
+        return self.rollout_wait()
+
+
+def _flatten_obs(obs):
+    assert isinstance(obs, (list, tuple))
+    assert len(obs) > 0
+
+    if isinstance(obs[0], dict):
+        keys = obs[0].keys()
+        return {k: np.stack([o[k] for o in obs]) for k in keys}
+    else:
+        return np.stack(obs)
+
+
+def _flatten_list(l):
+    assert isinstance(l, (list, tuple))
+    assert len(l) > 0
+    assert all([len(l_) > 0 for l_ in l])
+
+    return [l__ for l_ in l for l__ in l_]

alr_envs/utils/policies.py

@@ -0,0 +1,48 @@
+from gym import Env
+
+from alr_envs.mujoco.alr_mujoco_env import AlrMujocoEnv
+
+
+class BaseController:
+    def __init__(self, env: Env):
+        self.env = env
+
+    def get_action(self, des_pos, des_vel):
+        raise NotImplementedError
+
+
+class PosController(BaseController):
+    def get_action(self, des_pos, des_vel):
+        return des_pos
+
+
+class VelController(BaseController):
+    def get_action(self, des_pos, des_vel):
+        return des_vel
+
+
+class PDController(BaseController):
+    def __init__(self, env: AlrMujocoEnv):
+        self.p_gains = env.p_gains
+        self.d_gains = env.d_gains
+        super(PDController, self).__init__(env)
+
+    def get_action(self, des_pos, des_vel):
+        # TODO: make standardized ALRenv such that all of them have current_pos/vel attributes
+        cur_pos = self.env.current_pos
+        cur_vel = self.env.current_vel
+        if len(des_pos) != len(cur_pos):
+            des_pos = self.env.extend_des_pos(des_pos)
+        if len(des_vel) != len(cur_vel):
+            des_vel = self.env.extend_des_vel(des_vel)
+        trq = self.p_gains * (des_pos - cur_pos) + self.d_gains * (des_vel - cur_vel)
+        return trq
+
+
+def get_policy_class(policy_type):
+    if policy_type == "motor":
+        return PDController
+    elif policy_type == "velocity":
+        return VelController
+    elif policy_type == "position":
+        return PosController

alr_envs/utils/utils.py

@@ -19,3 +19,31 @@ def angle_normalize(x, type="deg"):
 
     two_pi = 2 * np.pi
     return x - two_pi * np.floor((x + np.pi) / two_pi)
+
+
+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

alr_envs/utils/wrapper/detpmp_wrapper.py

@@ -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

alr_envs/utils/wrapper/dmp_wrapper.py

@@ -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)

alr_envs/utils/wrapper/mp_wrapper.py

@@ -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