mp wrapper fixes

alr_envs/alr/classic_control/viapoint_reacher/mp_wrapper.py

@@ -7,6 +7,7 @@ from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
 
 class MPWrapper(RawInterfaceWrapper):
 
+    @property
     def context_mask(self):
         return np.hstack([
             [self.env.random_start] * self.env.n_links,  # cos
@@ -24,7 +25,3 @@ class MPWrapper(RawInterfaceWrapper):
     @property
     def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
         return self.env.current_vel
-
-    @property
-    def dt(self) -> Union[float, int]:
-        return self.env.dt

alr_envs/alr/classic_control/viapoint_reacher/viapoint_reacher.py

@@ -1,8 +1,9 @@
-from typing import Iterable, Union
+from typing import Iterable, Union, Tuple, Optional
 
 import gym
 import matplotlib.pyplot as plt
 import numpy as np
+from gym.core import ObsType
 from gym.utils import seeding
 
 from alr_envs.alr.classic_control.base_reacher.base_reacher_direct import BaseReacherDirectEnv
@@ -40,7 +41,8 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
     # def start_pos(self):
     #     return self._start_pos
 
-    def reset(self):
+    def reset(self, *, seed: Optional[int] = None, return_info: bool = False,
+              options: Optional[dict] = None, ) -> Union[ObsType, Tuple[ObsType, dict]]:
         self._generate_goal()
         return super().reset()
 
@@ -183,8 +185,10 @@ class ViaPointReacherEnv(BaseReacherDirectEnv):
 
                 plt.pause(0.01)
 
+
 if __name__ == "__main__":
     import time
+
     env = ViaPointReacherEnv(5)
     env.reset()