adjusted classic control environments to new interface

alr_envs/utils/mps/detpmp_wrapper.py

@@ -7,22 +7,22 @@ from alr_envs.utils.mps.mp_wrapper import MPWrapper
 
 
 class DetPMPWrapper(MPWrapper):
-    def __init__(self, env: MPEnv, num_dof: int, num_basis: int, width: int, start_pos=None, duration: int = 1,
-                 dt: float = 0.01, post_traj_time: float = 0., policy_type: str = None, weights_scale: float = 1.,
+    def __init__(self, env: MPEnv, num_dof: int, num_basis: int, width: int, duration: int = 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):
-        # self.duration = duration  # seconds
+        self.duration = duration  # seconds
 
         super().__init__(env, num_dof, dt, duration, 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, **mp_kwargs)
+                         width=width, zero_start=zero_start, zero_goal=zero_goal, **mp_kwargs)
+
+        self.dt = dt
 
         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):
+                      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)
 

alr_envs/utils/mps/dmp_wrapper.py

@@ -63,7 +63,7 @@ class DmpWrapper(MPWrapper):
             goal_pos = params[0, -self.mp.num_dimensions:]  # [num_dof]
             params = params[:, :-self.mp.num_dimensions]  # [1,num_dof]
         else:
-            goal_pos = self.env.goal_pos  # self.mp.dmp_goal_pos.flatten()
+            goal_pos = self.env.goal_pos
             assert goal_pos is not None
 
         weight_matrix = np.reshape(params, self.mp.dmp_weights.shape)  # [num_basis, num_dof]

alr_envs/utils/mps/mp_environments.py

@@ -9,7 +9,7 @@ class MPEnv(gym.Env):
 
     @property
     @abstractmethod
-    def corrected_obs_index(self):
+    def active_obs(self):
         """Returns boolean value for each observation entry
         whether the observation is returned by the DMP for the contextual case or not.
         This effectively allows to filter unwanted or unnecessary observations from the full step-based case.

alr_envs/utils/mps/mp_wrapper.py

@@ -13,6 +13,12 @@ class MPWrapper(gym.Wrapper, ABC):
                  policy_type: str = None, weights_scale: float = 1., render_mode: str = None, **mp_kwargs):
         super().__init__(env)
 
+        # adjust observation space to reduce version
+        obs_sp = self.env.observation_space
+        self.observation_space = gym.spaces.Box(low=obs_sp.low[self.env.active_obs],
+                                                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)
 
@@ -51,8 +57,7 @@ class MPWrapper(gym.Wrapper, ABC):
         self.env.configure(context)
 
     def reset(self):
-        obs = self.env.reset()
-        return obs[self.env]
+        return self.env.reset()[self.env.active_obs]
 
     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"""
@@ -82,7 +87,7 @@ class MPWrapper(gym.Wrapper, ABC):
                 break
 
         done = True
-        return obs, rewards, done, info
+        return obs[self.env.active_obs], rewards, done, info
 
     def render(self, mode='human', **kwargs):
         """Only set render options here, such that they can be used during the rollout.