restructuring

alr_envs/alr/__init__.py

@@ -346,7 +346,7 @@ for _v in _versions:
         kwargs={
             "name": f"alr_envs:{_v}",
             "wrappers": [classic_control.simple_reacher.MPWrapper],
-            "mp_kwargs": {
+            "traj_gen_kwargs": {
                 "num_dof": 2 if "long" not in _v.lower() else 5,
                 "num_basis": 5,
                 "duration": 2,
@@ -386,7 +386,7 @@ register(
     kwargs={
         "name": "alr_envs:ViaPointReacher-v0",
         "wrappers": [classic_control.viapoint_reacher.MPWrapper],
-        "mp_kwargs": {
+        "traj_gen_kwargs": {
             "num_dof": 5,
             "num_basis": 5,
             "duration": 2,
@@ -424,7 +424,7 @@ for _v in _versions:
         kwargs={
             "name": f"alr_envs:HoleReacher-{_v}",
             "wrappers": [classic_control.hole_reacher.MPWrapper],
-            "mp_kwargs": {
+            "traj_gen_kwargs": {
                 "num_dof": 5,
                 "num_basis": 5,
                 "duration": 2,
@@ -467,7 +467,7 @@ for _v in _versions:
         kwargs={
             "name": f"alr_envs:{_v}",
             "wrappers": [mujoco.reacher.MPWrapper],
-            "mp_kwargs": {
+            "traj_gen_kwargs": {
                 "num_dof": 5 if "long" not in _v.lower() else 7,
                 "num_basis": 2,
                 "duration": 4,

alr_envs/alr/mujoco/ant_jump/new_mp_wrapper.py

@@ -1,12 +1,13 @@
-from alr_envs.mp.episodic_wrapper import EpisodicWrapper
+from alr_envs.mp.black_box_wrapper import BlackBoxWrapper
 from typing import Union, Tuple
 import numpy as np
 
+from alr_envs.mp.raw_interface_wrapper import RawInterfaceWrapper
 
 
-class NewMPWrapper(EpisodicWrapper):
+class NewMPWrapper(RawInterfaceWrapper):
 
-    def set_active_obs(self):
+    def get_context_mask(self):
         return np.hstack([
             [False] * 111, # ant has 111 dimensional observation space !!
             [True] # goal height

alr_envs/alr/mujoco/beerpong/new_mp_wrapper.py

@@ -1,15 +1,11 @@
-from typing import Tuple, Union
+from typing import Union, Tuple
 
 import numpy as np
 
-from alr_envs.mp.episodic_wrapper import EpisodicWrapper
+from alr_envs.mp.raw_interface_wrapper import RawInterfaceWrapper
 
 
-class NewMPWrapper(EpisodicWrapper):
-
-    # def __init__(self, replanning_model):
-    #     self.replanning_model = replanning_model
-
+class NewMPWrapper(RawInterfaceWrapper):
     @property
     def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
         return self.env.sim.data.qpos[0:7].copy()
@@ -18,7 +14,7 @@ class NewMPWrapper(EpisodicWrapper):
     def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
         return self.env.sim.data.qvel[0:7].copy()
 
-    def set_active_obs(self):
+    def get_context_mask(self):
         return np.hstack([
             [False] * 7,  # cos
             [False] * 7,  # sin
@@ -27,12 +23,7 @@ class NewMPWrapper(EpisodicWrapper):
             [False] * 3,  # cup_goal_diff_top
             [True] * 2,  # xy position of cup
             [False]  # env steps
-            ])
-
-    def do_replanning(self, pos, vel, s, a, t, last_replan_step):
-        return False
-        # const = np.arange(0, 1000, 10)
-        # return bool(self.replanning_model(s))
+        ])
 
     def _episode_callback(self, action: np.ndarray) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
         if self.mp.learn_tau:

alr_envs/alr/mujoco/hopper_jump/new_mp_wrapper.py

@@ -1,9 +1,9 @@
-from alr_envs.mp.episodic_wrapper import EpisodicWrapper
+from alr_envs.mp.black_box_wrapper import BlackBoxWrapper
 from typing import Union, Tuple
 import numpy as np
 
 
-class NewMPWrapper(EpisodicWrapper):
+class NewMPWrapper(BlackBoxWrapper):
     @property
     def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
         return self.env.sim.data.qpos[3:6].copy()
@@ -21,7 +21,7 @@ class NewMPWrapper(EpisodicWrapper):
     #     ])
 
     # Random x goal + random init pos
-    def set_active_obs(self):
+    def get_context_mask(self):
         return np.hstack([
                 [False] * (2 + int(not self.env.exclude_current_positions_from_observation)),  # position
                 [True] * 3,    # set to true if randomize initial pos
@@ -31,7 +31,7 @@ class NewMPWrapper(EpisodicWrapper):
 
 
 class NewHighCtxtMPWrapper(NewMPWrapper):
-    def set_active_obs(self):
+    def get_context_mask(self):
         return np.hstack([
             [False] * (2 + int(not self.env.exclude_current_positions_from_observation)),  # position
             [True] * 3,  # set to true if randomize initial pos

alr_envs/alr/mujoco/reacher/alr_reacher.py

@@ -149,4 +149,4 @@ if __name__ == '__main__':
         if d:
             env.reset()
 
-    env.close()
+    env.close()

alr_envs/alr/mujoco/reacher/new_mp_wrapper.py

@@ -1,9 +1,9 @@
-from alr_envs.mp.episodic_wrapper import EpisodicWrapper
+from alr_envs.mp.black_box_wrapper import BlackBoxWrapper
 from typing import Union, Tuple
 import numpy as np
 
 
-class MPWrapper(EpisodicWrapper):
+class MPWrapper(BlackBoxWrapper):
 
     @property
     def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
@@ -12,7 +12,7 @@ class MPWrapper(EpisodicWrapper):
     def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
         return self.env.sim.data.qvel.flat[:self.env.n_links]
 
-    def set_active_obs(self):
+    def get_context_mask(self):
         return np.concatenate([
             [False] * self.env.n_links,  # cos
             [False] * self.env.n_links,  # sin