wip

alr_envs/utils/wrapper/dmp_wrapper.py

@@ -1,4 +1,3 @@
-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
@@ -11,9 +10,9 @@ 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,
+                 final_pos: np.ndarray = None, duration: int = 1, alpha_phase: float = 2., dt: float = None,
                  learn_goal: bool = False, post_traj_time: float = 0., policy_type: str = None,
-                 weights_scale: float = 1., goal_scale: float = 1.):
+                 weights_scale: float = 1., goal_scale: float = 1., bandwidth_factor: float = 3.):
 
         """
         This Wrapper generates a trajectory based on a DMP and will only return episodic performances.
@@ -33,20 +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
+        start_pos = env.start_pos if hasattr(env, "start_pos") else start_pos
+        assert start_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,
-                         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.):
+                      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)
+        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,
                        num_time_steps=int(duration / dt), dt=dt)

alr_envs/utils/wrapper/mp_wrapper.py

@@ -13,19 +13,20 @@ class MPWrapper(gym.Wrapper, ABC):
                  env: gym.Env,
                  num_dof: int,
                  duration: int = 1,
-                 dt: float = 0.01,
-                 # learn_goal: bool = False,
+                 dt: float = None,
                  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
+
+        # dt = env.dt if hasattr(env, "dt") else dt
+        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.mp = self.initialize_mp(num_dof, duration, dt, **mp_kwargs)
@@ -38,6 +39,26 @@ class MPWrapper(gym.Wrapper, ABC):
         self.render_mode = None
         self.render_kwargs = None
 
+    # TODO: not yet final
+    def __call__(self, params, contexts=None):
+        params = np.atleast_2d(params)
+        obs = []
+        rewards = []
+        dones = []
+        infos = []
+        for p, c in zip(params, contexts):
+            self.configure(c)
+            ob, reward, done, info = self.step(p)
+            obs.append(ob)
+            rewards.append(reward)
+            dones.append(done)
+            infos.append(info)
+
+        return obs, np.array(rewards), dones, infos
+
+    def configure(self, context):
+        self.env.configure(context)
+
     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)
@@ -53,6 +74,7 @@ class MPWrapper(gym.Wrapper, ABC):
         # infos = defaultdict(list)
 
         # TODO: @Max Why do we need this configure, states should be part of the model
+        # TODO: Ask Onur if the context distribution needs to be outside the environment
         # self.env.configure(context)
         obs = self.env.reset()
         info = {}
@@ -77,8 +99,8 @@ class MPWrapper(gym.Wrapper, ABC):
         self.render_mode = mode
         self.render_kwargs = kwargs
 
-    def __call__(self, actions):
-        return self.step(actions)
+    # def __call__(self, actions):
+    #     return self.step(actions)
         # params = np.atleast_2d(params)
         # rewards = []
         # infos = []