Fixing moving objects & rendering (#458)

Farama-Foundation · Apr 3, 2024 · a5b32cc · a5b32cc
1 parent c822f28
commit a5b32cc
Show file tree

Hide file tree

Showing 53 changed files with 175 additions and 200 deletions.
diff --git a/docs/_static/img/metaworld-text.svg b/docs/_static/img/metaworld-text.svg
diff --git a/metaworld/envs/mujoco/sawyer_xyz/sawyer_xyz_env.py b/metaworld/envs/mujoco/sawyer_xyz/sawyer_xyz_env.py
@@ -26,13 +26,22 @@ class SawyerMocapBase(mjenv_gym):
         "render_fps": 80,
     }
 
-    def __init__(self, model_name, frame_skip=5, render_mode=None):
+    def __init__(
+        self,
+        model_name,
+        frame_skip=5,
+        render_mode=None,
+        camera_name=None,
+        camera_id=None,
+    ):
         mjenv_gym.__init__(
             self,
             model_name,
             frame_skip=frame_skip,
             observation_space=self.sawyer_observation_space,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
         self.reset_mocap_welds()
         self.frame_skip = frame_skip
@@ -113,6 +122,8 @@ def __init__(
         action_scale=1.0 / 100,
         action_rot_scale=1.0,
         render_mode=None,
+        camera_id=None,
+        camera_name=None,
     ):
         self.action_scale = action_scale
         self.action_rot_scale = action_rot_scale
@@ -140,7 +151,13 @@ def __init__(
 
         self._partially_observable = True
 
-        super().__init__(model_name, frame_skip=frame_skip, render_mode=render_mode)
+        super().__init__(
+            model_name,
+            frame_skip=frame_skip,
+            render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
+        )
 
         mujoco.mj_forward(
             self.model, self.data
@@ -484,7 +501,7 @@ def step(self, action):
                     "unscaled_reward": 0.0,
                 },
             )
-
+        mujoco.mj_forward(self.model, self.data)
         self._last_stable_obs = self._get_obs()
 
         self._last_stable_obs = np.clip(
@@ -524,6 +541,7 @@ def evaluate_state(self, obs, action):
     def reset(self, seed=None, options=None):
         self.curr_path_length = 0
         obs, info = super().reset()
+        mujoco.mj_forward(self.model, self.data)
         self._prev_obs = obs[:18].copy()
         obs[18:36] = self._prev_obs
         obs = np.float64(obs)

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_assembly_peg_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_assembly_peg_v2.py
@@ -13,7 +13,7 @@
 class SawyerNutAssemblyEnvV2(SawyerXYZEnv):
     WRENCH_HANDLE_LENGTH = 0.02
 
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (0, 0.6, 0.02)
@@ -26,9 +26,9 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
-        if tasks is not None:
-            self.tasks = tasks
 
         self.init_config = {
             "obj_init_angle": 0.3,

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_basketball_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_basketball_v2.py
@@ -14,7 +14,7 @@ class SawyerBasketballEnvV2(SawyerXYZEnv):
     PAD_SUCCESS_MARGIN = 0.06
     TARGET_RADIUS = 0.08
 
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.1, 0.6, 0.0299)
@@ -27,11 +27,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_angle": 0.3,
             "obj_init_pos": np.array([0, 0.6, 0.03], dtype=np.float32),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_bin_picking_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_bin_picking_v2.py
@@ -23,7 +23,7 @@ class SawyerBinPickingEnvV2(SawyerXYZEnv):
         - (11/23/20) Updated reward function to new pick-place style
     """
 
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.07)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.21, 0.65, 0.02)
@@ -37,9 +37,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
-        if tasks is not None:
-            self.tasks = tasks
+
         self.init_config = {
             "obj_init_angle": 0.3,
             "obj_init_pos": np.array([-0.12, 0.7, 0.02]),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_box_close_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_box_close_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerBoxCloseEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.05, 0.5, 0.02)
@@ -24,9 +24,9 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
-        if tasks is not None:
-            self.tasks = tasks
 
         self.init_config = {
             "obj_init_angle": 0.3,

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_topdown_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_topdown_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerButtonPressTopdownEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.1, 0.8, 0.115)
@@ -22,9 +22,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
-        if tasks is not None:
-            self.tasks = tasks
+
         self.init_config = {
             "obj_init_pos": np.array([0, 0.8, 0.115], dtype=np.float32),
             "hand_init_pos": np.array([0, 0.4, 0.2], dtype=np.float32),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_topdown_wall_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_topdown_wall_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerButtonPressTopdownWallEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.1, 0.8, 0.115)
@@ -22,11 +22,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_pos": np.array([0, 0.8, 0.115], dtype=np.float32),
             "hand_init_pos": np.array([0, 0.4, 0.2], dtype=np.float32),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerButtonPressEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.1, 0.85, 0.115)
@@ -22,11 +22,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_pos": np.array([0.0, 0.9, 0.115], dtype=np.float32),
             "hand_init_pos": np.array([0, 0.4, 0.2], dtype=np.float32),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_wall_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_button_press_wall_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerButtonPressWallEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.05, 0.85, 0.1149)
@@ -22,11 +22,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_pos": np.array([0.0, 0.9, 0.115], dtype=np.float32),
             "hand_init_pos": np.array([0, 0.4, 0.2], dtype=np.float32),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_coffee_button_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_coffee_button_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerCoffeeButtonEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         self.max_dist = 0.03
 
         hand_low = (-0.5, 0.4, 0.05)
@@ -28,11 +28,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_pos": np.array([0, 0.9, 0.28]),
             "obj_init_angle": 0.3,

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_coffee_pull_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_coffee_pull_v2.py
@@ -12,7 +12,7 @@
 
 
 class SawyerCoffeePullEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.05, 0.7, -0.001)
@@ -25,11 +25,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_pos": np.array([0, 0.75, 0.0]),
             "obj_init_angle": 0.3,

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_coffee_push_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_coffee_push_v2.py
@@ -12,7 +12,7 @@
 
 
 class SawyerCoffeePushEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.1, 0.55, -0.001)
@@ -25,10 +25,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
         self.init_config = {
             "obj_init_angle": 0.3,
             "obj_init_pos": np.array([0.0, 0.6, 0.0]),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_dial_turn_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_dial_turn_v2.py
@@ -13,7 +13,7 @@
 class SawyerDialTurnEnvV2(SawyerXYZEnv):
     TARGET_RADIUS = 0.07
 
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (-0.1, 0.7, 0.0)
@@ -26,11 +26,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_pos": np.array([0, 0.7, 0.0]),
             "hand_init_pos": np.array([0, 0.6, 0.2], dtype=np.float32),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_disassemble_peg_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_disassemble_peg_v2.py
@@ -13,7 +13,7 @@
 class SawyerNutDisassembleEnvV2(SawyerXYZEnv):
     WRENCH_HANDLE_LENGTH = 0.02
 
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         hand_low = (-0.5, 0.40, 0.05)
         hand_high = (0.5, 1, 0.5)
         obj_low = (0.0, 0.6, 0.025)
@@ -26,11 +26,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_angle": 0.3,
             "obj_init_pos": np.array([0, 0.7, 0.025]),

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_door_close_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_door_close_v2.py
@@ -11,7 +11,7 @@
 
 
 class SawyerDoorCloseEnvV2(SawyerXYZEnv):
-    def __init__(self, tasks=None, render_mode=None):
+    def __init__(self, render_mode=None, camera_name=None, camera_id=None):
         goal_low = (0.2, 0.65, 0.1499)
         goal_high = (0.3, 0.75, 0.1501)
         hand_low = (-0.5, 0.40, 0.05)
@@ -24,11 +24,10 @@ def __init__(self, tasks=None, render_mode=None):
             hand_low=hand_low,
             hand_high=hand_high,
             render_mode=render_mode,
+            camera_name=camera_name,
+            camera_id=camera_id,
         )
 
-        if tasks is not None:
-            self.tasks = tasks
-
         self.init_config = {
             "obj_init_angle": 0.3,
             "obj_init_pos": np.array([0.1, 0.95, 0.15], dtype=np.float32),