Add interface for envs controlable by a PD Controller and add more infos to mp_wrapper info return value

alr_envs/utils/mps/alr_env.py

@@ -1,14 +1,13 @@
-from abc import abstractmethod
+from abc import abstractmethod, ABC
 from typing import Union
 
 import gym
 import numpy as np
 
 
-class AlrEnv(gym.Env):
+class AlrEnv(gym.Env, ABC):
 
     @property
-    @abstractmethod
     def active_obs(self):
         """Returns boolean mask for each observation entry
         whether the observation is returned for the contextual case or not.
@@ -31,3 +30,11 @@ class AlrEnv(gym.Env):
         By default this returns the starting position.
         """
         return self.start_pos
+
+    @property
+    @abstractmethod
+    def dt(self) -> Union[float, int]:
+        """
+        Returns the time between two simulated steps of the environment
+        """
+        raise NotImplementedError()

alr_envs/utils/mps/detpmp_wrapper.py

@@ -2,29 +2,36 @@ import gym
 import numpy as np
 from mp_lib import det_promp
 
-from alr_envs.utils.mps.mp_environments import AlrEnv
+from alr_envs.utils.mps.alr_env import AlrEnv
 from alr_envs.utils.mps.mp_wrapper import MPWrapper
 
 
 class DetPMPWrapper(MPWrapper):
-    def __init__(self, env: AlrEnv, num_dof: int, num_basis: int, width: float, duration: float = 1, dt: float = 0.01,
-                 post_traj_time: float = 0., policy_type: str = None, weights_scale: float = 1.,
-                 zero_start: bool = False, zero_goal: bool = False, **mp_kwargs):
+    def __init__(self, env: AlrEnv, num_dof, num_basis, width, start_pos=None, duration=1, post_traj_time=0.,
+                 policy_type=None, weights_scale=1, zero_start=False, zero_goal=False, learn_mp_length: bool =True,
+                 **mp_kwargs):
         self.duration = duration  # seconds
 
-        dt = env.dt if hasattr(env, "dt") else dt
-        assert dt is not None
-        self.dt = dt
+        super().__init__(env=env, num_dof=num_dof, duration=duration, post_traj_time=post_traj_time,
+                         policy_type=policy_type, weights_scale=weights_scale, num_basis=num_basis,
+                         width=width, zero_start=zero_start, zero_goal=zero_goal,
+                         **mp_kwargs)
 
-        super().__init__(env, num_dof, dt, duration, post_traj_time, policy_type, weights_scale, num_basis=num_basis,
-                         width=width, zero_start=zero_start, zero_goal=zero_goal, **mp_kwargs)
+        self.learn_mp_length = learn_mp_length
+        if self.learn_mp_length:
+            parameter_space_shape = (1+num_basis*num_dof,)
+        else:
+            parameter_space_shape = (num_basis * num_dof,)
+        self.min_param = -np.inf
+        self.max_param = np.inf
+        self.parameterization_space = gym.spaces.Box(low=self.min_param, high=self.max_param,
+                                                     shape=parameter_space_shape, dtype=np.float32)
 
-        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
 
-    def initialize_mp(self, num_dof: int, duration: int, dt: float, num_basis: int = 5, width: float = None,
-                      off: float = 0.01, zero_start: bool = False, zero_goal: bool = False):
-        pmp = det_promp.DeterministicProMP(n_basis=num_basis, n_dof=num_dof, width=width, off=off,
+    def initialize_mp(self, num_dof: int, duration: int, num_basis: int = 5, width: float = None,
+                      zero_start: bool = False, zero_goal: bool = False, **kwargs):
+        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))
@@ -33,10 +40,15 @@ class DetPMPWrapper(MPWrapper):
         return pmp
 
     def mp_rollout(self, action):
-        params = np.reshape(action, newshape=(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.learn_mp_length:
+            duration = max(1, self.duration*abs(action[0]))
+            params = np.reshape(action[1:], (self.mp.n_basis, -1)) * self.weights_scale # TODO: Fix Bug when zero_start is true
+        else:
+            duration = self.duration
+            params = np.reshape(action, (self.mp.n_basis, -1)) * self.weights_scale # TODO: Fix Bug when zero_start is true
+        self.mp.set_weights(1., params)
+        _, des_pos, des_vel, _ = self.mp.compute_trajectory(frequency=max(1, duration))
         if self.mp.zero_start:
-            des_pos += self.env.start_pos[None, :]
+            des_pos += self.start_pos
 
         return des_pos, des_vel

alr_envs/utils/mps/dmp_wrapper.py

@@ -4,7 +4,7 @@ from mp_lib import dmps
 from mp_lib.basis import DMPBasisGenerator
 from mp_lib.phase import ExpDecayPhaseGenerator
 
-from alr_envs.utils.mps.mp_environments import AlrEnv
+from alr_envs.utils.mps.alr_env import AlrEnv
 from alr_envs.utils.mps.mp_wrapper import MPWrapper
 
 
@@ -32,26 +32,26 @@ class DmpWrapper(MPWrapper):
             goal_scale:
         """
         self.learn_goal = learn_goal
-        dt = env.dt if hasattr(env, "dt") else dt
-        assert dt is not None
+
         self.t = np.linspace(0, duration, int(duration / dt))
         self.goal_scale = goal_scale
 
-        super().__init__(env, num_dof, dt, duration, post_traj_time, policy_type, weights_scale, render_mode,
+        super().__init__(env=env, num_dof=num_dof, duration=duration, post_traj_time=post_traj_time,
+                         policy_type=policy_type, weights_scale=weights_scale, render_mode=render_mode,
                          num_basis=num_basis, alpha_phase=alpha_phase, bandwidth_factor=bandwidth_factor)
 
         action_bounds = np.inf * np.ones((np.prod(self.mp.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, alpha_phase: float = 2.,
-                      bandwidth_factor: int = 3):
+    def initialize_mp(self, num_dof: int, duration: int, num_basis: int, alpha_phase: float = 2.,
+                      bandwidth_factor: int = 3, **kwargs):
 
         phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
         basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=num_basis,
                                             basis_bandwidth_factor=bandwidth_factor)
 
         dmp = dmps.DMP(num_dof=num_dof, basis_generator=basis_generator, phase_generator=phase_generator,
-                       duration=duration, dt=dt)
+                       dt=self.dt)
 
         return dmp
 

alr_envs/utils/mps/mp_wrapper.py

@@ -1,16 +1,36 @@
 from abc import ABC, abstractmethod
+from typing import Union
 
 import gym
 import numpy as np
 
-from alr_envs.utils.mps.mp_environments import AlrEnv
-from alr_envs.utils.policies import get_policy_class
+from alr_envs.utils.mps.alr_env import AlrEnv
+from alr_envs.utils.policies import get_policy_class, BaseController
 
 
 class MPWrapper(gym.Wrapper, ABC):
+    """
+    Base class for movement primitive based gym.Wrapper implementations.
 
-    def __init__(self, env: AlrEnv, num_dof: int, dt: float, duration: float = 1, post_traj_time: float = 0.,
-                 policy_type: str = None, weights_scale: float = 1., render_mode: str = None, **mp_kwargs):
+    :param env: The (wrapped) environment this wrapper is applied on
+    :param num_dof: Dimension of the action space of the wrapped env
+    :param duration: Number of timesteps in the trajectory of the movement primitive
+    :param post_traj_time: Time for which the last position of the trajectory is fed to the environment to continue
+    simulation
+    :param policy_type: Type or object defining the policy that is used to generate action based on the trajectory
+    :param weight_scale: Scaling parameter for the actions given to this wrapper
+    :param render_mode: Equivalent to gym render mode
+    """
+    def __init__(self,
+                 env: AlrEnv,
+                 num_dof: int,
+                 duration: int = 1,
+                 post_traj_time: float = 0.,
+                 policy_type: Union[str, BaseController] = None,
+                 weights_scale: float = 1.,
+                 render_mode: str = None,
+                 **mp_kwargs
+                 ):
         super().__init__(env)
 
         # adjust observation space to reduce version
@@ -19,14 +39,15 @@ class MPWrapper(gym.Wrapper, ABC):
                                                 high=obs_sp.high[self.env.active_obs],
                                                 dtype=obs_sp.dtype)
 
-        assert dt is not None  # this should never happen as MPWrapper is a base class
-        self.post_traj_steps = int(post_traj_time / dt)
+        self.post_traj_steps = int(post_traj_time / env.dt)
 
-        self.mp = self.initialize_mp(num_dof, duration, dt, **mp_kwargs)
+        self.mp = self.initialize_mp(num_dof=num_dof, duration=duration, **mp_kwargs)
         self.weights_scale = weights_scale
 
-        policy_class = get_policy_class(policy_type)
-        self.policy = policy_class(env)
+        if type(policy_type) is str:
+            self.policy = get_policy_class(policy_type, env, mp_kwargs)
+        else:
+            self.policy = policy_type
 
         # rendering
         self.render_mode = render_mode
@@ -62,29 +83,30 @@ class MPWrapper(gym.Wrapper, ABC):
 
         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.mp.n_dof))])
-
-        # self._trajectory = trajectory
-        # self._velocity = velocity
-
-        rewards = 0
-        info = {}
-        # create random obs as the reset function is called externally
-        obs = self.env.observation_space.sample()
+            velocity = np.vstack([velocity, np.zeros(shape=(self.post_traj_steps, self.mp.num_dimensions))])
 
+        trajectory_length = len(trajectory)
+        actions = np.zeros(shape=(trajectory_length, self.mp.num_dimensions))
+        observations= np.zeros(shape=(trajectory_length,) + self.env.observation_space.shape)
+        rewards = np.zeros(shape=(trajectory_length,))
+        trajectory_return = 0
+        infos = dict(step_infos =[])
         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
-            # TODO return all dicts?
-            # [infos[k].append(v) for k, v in info.items()]
+            actions[t,:] = self.policy.get_action(pos_vel[0], pos_vel[1])
+            observations[t,:], rewards[t], done, info = self.env.step(actions[t,:])
+            trajectory_return += rewards[t]
+            infos['step_infos'].append(info)
             if self.render_mode:
                 self.env.render(mode=self.render_mode, **self.render_kwargs)
             if done:
                 break
 
+        infos['step_actions'] = actions[:t+1]
+        infos['step_observations'] = observations[:t+1]
+        infos['step_rewards'] = rewards[:t+1]
+        infos['trajectory_length'] = t+1
         done = True
-        return obs[self.env.active_obs], rewards, done, info
+        return observations[t][self.env.active_obs], trajectory_return, done, infos
 
     def render(self, mode='human', **kwargs):
         """Only set render options here, such that they can be used during the rollout.
@@ -102,7 +124,7 @@ class MPWrapper(gym.Wrapper, ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def initialize_mp(self, num_dof: int, duration: float, dt: float, **kwargs):
+    def initialize_mp(self, num_dof: int, duration: float, **kwargs):
         """
         Create respective instance of MP
         Returns: