refractoring of DMP environmets to fit gym interface better.

2021-03-26 14:05:16 +01:00
parent 6233c85904
commit 7ceadeff0a
20 changed files with 661 additions and 568 deletions
@@ -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
@@ -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)
@@ -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