wip

alr_envs/utils/wrapper/dmp_wrapper.py

@@ -9,8 +9,10 @@ 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 = None,
+    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 = None,
                  learn_goal: bool = False, return_to_start: bool = False, post_traj_time: float = 0.,
                  weights_scale: float = 1., goal_scale: float = 1., bandwidth_factor: float = 3.,
                  policy_type: str = None, render_mode: str = None):
@@ -35,26 +37,30 @@ class DmpWrapper(MPWrapper):
         self.learn_goal = learn_goal
         dt = env.dt if hasattr(env, "dt") else dt
         assert dt is not None
-        start_pos = start_pos if start_pos is not None else env.start_pos if hasattr(env, "start_pos") else None
+        # start_pos = start_pos if start_pos is not None else env.start_pos if hasattr(env, "start_pos") else None
         # TODO: assert start_pos is not None  # start_pos will be set in initialize, do we need this here?
-        if learn_goal:
+        # if learn_goal:
             # final_pos = np.zeros_like(start_pos)  # arbitrary, will be learned
-            final_pos = np.zeros((1, num_dof))  # arbitrary, will be learned
-        else:
-            final_pos = final_pos if final_pos is not None else start_pos if return_to_start else None
-        assert final_pos is not None
+            # final_pos = np.zeros((1, num_dof))  # arbitrary, will be learned
+        # else:
+        #     final_pos = final_pos if final_pos is not None else start_pos if return_to_start else None
+        # assert final_pos is not None
         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, render_mode,
-                         num_basis=num_basis, start_pos=start_pos, final_pos=final_pos, alpha_phase=alpha_phase,
+                         num_basis=num_basis,
+                         # start_pos=start_pos, final_pos=final_pos,
+                         alpha_phase=alpha_phase,
                          bandwidth_factor=bandwidth_factor)
 
         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., bandwidth_factor: float = 3.):
+    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., bandwidth_factor: float = 3.):
 
         phase_generator = ExpDecayPhaseGenerator(alpha_phase=alpha_phase, duration=duration)
         basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=num_basis,
@@ -66,12 +72,12 @@ class DmpWrapper(MPWrapper):
         # dmp.dmp_start_pos = start_pos.reshape((1, num_dof))
         # in a contextual environment, the start_pos may be not fixed, set in mp_rollout?
         # TODO: Should we set start_pos in init at all? It's only used after calling rollout anyway...
-        dmp.dmp_start_pos = start_pos.reshape((1, num_dof)) if start_pos is not None else np.zeros((1, num_dof))
+        # dmp.dmp_start_pos = start_pos.reshape((1, num_dof)) if start_pos is not None else np.zeros((1, num_dof))
 
-        weights = np.zeros((num_basis, num_dof))
-        goal_pos = np.zeros(num_dof) if self.learn_goal else final_pos
+        # 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)
+        # dmp.set_weights(weights, goal_pos)
         return dmp
 
     def goal_and_weights(self, params):
@@ -83,7 +89,7 @@ class DmpWrapper(MPWrapper):
             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()
+            goal_pos = self.env.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])
 
@@ -91,8 +97,8 @@ class DmpWrapper(MPWrapper):
         return goal_pos * self.goal_scale, weight_matrix * self.weights_scale
 
     def mp_rollout(self, action):
-        if self.mp.start_pos is None:
-            self.mp.start_pos = self.env.start_pos
+        # if self.mp.start_pos is None:
+        self.mp.dmp_start_pos = self.env.init_qpos  # start_pos
         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

@@ -62,7 +62,8 @@ class MPWrapper(gym.Wrapper, ABC):
         self.env.configure(context)
 
     def reset(self):
-        return self.env.reset()
+        obs = self.env.reset()
+        return 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"""