lots of new stuff...

2021-02-05 17:10:03 +01:00
parent cab2c249bb
commit 07195fa2dc
108 changed files with 1086 additions and 51 deletions
@@ -0,0 +1,88 @@
+from mp_lib import det_promp
+import numpy as np
+import gym
+
+
+class DetPMPEnvWrapperBase(gym.Wrapper):
+    def __init__(self,
+                 env,
+                 num_dof,
+                 num_basis,
+                 width,
+                 start_pos=None,
+                 duration=1,
+                 dt=0.01,
+                 post_traj_time=0.,
+                 policy=None,
+                 weights_scale=1):
+        super(DetPMPEnvWrapperBase, self).__init__(env)
+        self.num_dof = num_dof
+        self.num_basis = num_basis
+        self.dim = num_dof * num_basis
+        self.pmp = det_promp.DeterministicProMP(n_basis=num_basis, width=width, off=0.01)
+        weights = np.zeros(shape=(num_basis, num_dof))
+        self.pmp.set_weights(duration, weights)
+        self.weights_scale = weights_scale
+
+        self.duration = duration
+        self.dt = dt
+        self.post_traj_steps = int(post_traj_time / dt)
+
+        self.start_pos = start_pos
+
+        self.policy = policy
+
+    def __call__(self, params):
+        params = np.atleast_2d(params)
+        observations = []
+        rewards = []
+        dones = []
+        infos = []
+        for p in params:
+            observation, reward, done, info = self.rollout(p)
+            observations.append(observation)
+            rewards.append(reward)
+            dones.append(done)
+            infos.append(info)
+
+        return np.array(rewards), infos
+
+    def rollout(self, params, render=False):
+        """ This function generates a trajectory based on a DMP and then does the usual loop over reset and step"""
+        raise NotImplementedError
+
+
+class DetPMPEnvWrapperPD(DetPMPEnvWrapperBase):
+    """
+    Wrapper for gym environments which creates a trajectory in joint velocity space
+    """
+    def rollout(self, params, render=False):
+        params = np.reshape(params, newshape=(self.num_basis, self.num_dof)) * self.weights_scale
+        self.pmp.set_weights(self.duration, params)
+        t, des_pos, des_vel, des_acc = self.pmp.compute_trajectory(1/self.dt, 1.)
+        des_pos += self.start_pos[None, :]
+
+        if self.post_traj_steps > 0:
+            des_pos = np.vstack([des_pos, np.tile(des_pos[-1, :], [self.post_traj_steps, 1])])
+            des_vel = np.vstack([des_vel, np.zeros(shape=(self.post_traj_steps, self.num_dof))])
+
+        self._trajectory = des_pos
+
+        rews = []
+        infos = []
+
+        self.env.reset()
+
+        for t, pos_vel in enumerate(zip(des_pos, des_vel)):
+            ac = self.policy.get_action(self.env, pos_vel[0], pos_vel[1])
+            obs, rew, done, info = self.env.step(ac)
+            rews.append(rew)
+            infos.append(info)
+            if render:
+                self.env.render(mode="human")
+            if done:
+                break
+
+        reward = np.sum(rews)
+
+        return obs, reward, done, info
@@ -6,19 +6,31 @@ import gym


 class DmpEnvWrapperBase(gym.Wrapper):
-    def __init__(self, env, num_dof, num_basis, duration=1, dt=0.01, learn_goal=False):
+    def __init__(self,
+                 env,
+                 num_dof,
+                 num_basis,
+                 start_pos=None,
+                 final_pos=None,
+                 duration=1,
+                 alpha_phase=2,
+                 dt=0.01,
+                 learn_goal=False,
+                 post_traj_time=0.,
+                 policy=None):
        super(DmpEnvWrapperBase, self).__init__(env)
        self.num_dof = num_dof
        self.num_basis = num_basis
        self.dim = num_dof * num_basis
        if learn_goal:
            self.dim += num_dof
-        self.learn_goal = True
+        self.learn_goal = learn_goal
        self.duration = duration   # seconds
        time_steps = int(duration / dt)
        self.t = np.linspace(0, duration, time_steps)
+        self.post_traj_steps = int(post_traj_time / dt)

-        phase_generator = ExpDecayPhaseGenerator(alpha_phase=5, duration=duration)
+        phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
        basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=self.num_basis)

        self.dmp = dmps.DMP(num_dof=num_dof,
@@ -28,13 +40,18 @@ class DmpEnvWrapperBase(gym.Wrapper):
                            dt=dt
                            )

-        self.dmp.dmp_start_pos = env.start_pos.reshape((1, num_dof))
+        self.dmp.dmp_start_pos = start_pos.reshape((1, num_dof))

        dmp_weights = np.zeros((num_basis, num_dof))
-        dmp_goal_pos = np.zeros(num_dof)
+        if learn_goal:
+            dmp_goal_pos = np.zeros(num_dof)
+        else:
+            dmp_goal_pos = final_pos

        self.dmp.set_weights(dmp_weights, dmp_goal_pos)

+        self.policy = policy
+
    def __call__(self, params):
        params = np.atleast_2d(params)
        observations = []
@@ -48,7 +65,7 @@ class DmpEnvWrapperBase(gym.Wrapper):
            dones.append(done)
            infos.append(info)

-        return np.array(rewards)
+        return np.array(rewards), infos

    def goal_and_weights(self, params):
        if len(params.shape) > 1:
@@ -71,7 +88,7 @@ class DmpEnvWrapperBase(gym.Wrapper):
        raise NotImplementedError


-class DmpEnvWrapperAngle(DmpEnvWrapperBase):
+class DmpEnvWrapperPos(DmpEnvWrapperBase):
    """
    Wrapper for gym environments which creates a trajectory in joint angle space
    """
@@ -80,7 +97,12 @@ class DmpEnvWrapperAngle(DmpEnvWrapperBase):
        if hasattr(self.env, "weight_matrix_scale"):
            weight_matrix = weight_matrix * self.env.weight_matrix_scale
        self.dmp.set_weights(weight_matrix, goal_pos)
-        trajectory, velocities = self.dmp.reference_trajectory(self.t)
+        trajectory, _ = self.dmp.reference_trajectory(self.t)
+
+        if self.post_traj_steps > 0:
+            trajectory = np.vstack([trajectory, np.tile(trajectory[-1, :], [self.post_traj_steps, 1])])
+
+        self._trajectory = trajectory

        rews = []

@@ -95,8 +117,6 @@ class DmpEnvWrapperAngle(DmpEnvWrapperBase):
                break

        reward = np.sum(rews)
-        # done = True
-        info = {}

        return obs, reward, done, info

@@ -110,7 +130,7 @@ class DmpEnvWrapperVel(DmpEnvWrapperBase):
        if hasattr(self.env, "weight_matrix_scale"):
            weight_matrix = weight_matrix * self.env.weight_matrix_scale
        self.dmp.set_weights(weight_matrix, goal_pos)
-        trajectory, velocities = self.dmp.reference_trajectory(self.t)
+        _, velocities = self.dmp.reference_trajectory(self.t)

        rews = []
        infos = []
@@ -129,3 +149,41 @@ class DmpEnvWrapperVel(DmpEnvWrapperBase):
        reward = np.sum(rews)

        return obs, reward, done, info
+
+
+class DmpEnvWrapperPD(DmpEnvWrapperBase):
+    """
+    Wrapper for gym environments which creates a trajectory in joint velocity space
+    """
+    def rollout(self, action, render=False):
+        goal_pos, weight_matrix = self.goal_and_weights(action)
+        if hasattr(self.env, "weight_matrix_scale"):
+            weight_matrix = weight_matrix * self.env.weight_matrix_scale
+        self.dmp.set_weights(weight_matrix, goal_pos)
+        trajectory, velocity = self.dmp.reference_trajectory(self.t)
+
+        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.num_dof))])
+
+        self._trajectory = trajectory
+        self._velocity = velocity
+
+        rews = []
+        infos = []
+
+        self.env.reset()
+
+        for t, pos_vel in enumerate(zip(trajectory, velocity)):
+            ac = self.policy.get_action(self.env, pos_vel[0], pos_vel[1])
+            obs, rew, done, info = self.env.step(ac)
+            rews.append(rew)
+            infos.append(info)
+            if render:
+                self.env.render(mode="human")
+            if done:
+                break
+
+        reward = np.sum(rews)
+
+        return obs, reward, done, info
@@ -0,0 +1,15 @@
+class PDController:
+    def __init__(self, p_gains, d_gains):
+        self.p_gains = p_gains
+        self.d_gains = d_gains
+
+    def get_action(self, env, des_pos, des_vel):
+        # TODO: make standardized ALRenv such that all of them have current_pos/vel attributes
+        cur_pos = env.current_pos
+        cur_vel = env.current_vel
+        if len(des_pos) != len(cur_pos):
+            des_pos = env.extend_des_pos(des_pos)
+        if len(des_vel) != len(cur_vel):
+            des_vel = env.extend_des_vel(des_vel)
+        trq = self.p_gains * (des_pos - cur_pos) + self.d_gains * (des_vel - cur_vel)
+        return trq