Add observation space range

dfki-ric · Aug 20, 2024 · 61bb91f · 61bb91f
1 parent fe56df8
commit 61bb91f
Show file tree

Hide file tree

Showing 2 changed files with 171 additions and 9 deletions.
diff --git a/deformable_gym/robots/mj_robot.py b/deformable_gym/robots/mj_robot.py
@@ -39,6 +39,7 @@ def __init__(self, name: str) -> None:
         self.nq = self.model.nq
         self.dof = self.model.nv
         self.nact = self.model.nu
+        self.joint_range = self.model.jnt_range.copy()
         self.ctrl_range = self.model.actuator_ctrlrange.copy()
 
     @property

diff --git a/tests/envs/mujoco/test_mj_shadow_grasp.py b/tests/envs/mujoco/test_mj_shadow_grasp.py
@@ -1,5 +1,3 @@
-from copy import deepcopy
-
 import numpy as np
 import pytest
 from gymnasium.spaces import Box
@@ -8,8 +6,156 @@
 from deformable_gym.envs.mujoco.shadow_grasp import ShadowHandGrasp
 
 SEED = 42
-OBS_SPACE = Box(low=-np.inf, high=np.inf, shape=(30,), dtype=np.float64)
-
+OBS_SPACE = Box(
+    low=np.array(
+        [
+            -0.5,
+            -0.5,
+            -0.5,
+            -1.570796,
+            -1.570796,
+            -1.570796,
+            -0.523599,
+            -0.698132,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            0.0,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            -1.0472,
+            0.0,
+            -0.20944,
+            -0.698132,
+            -0.261799,
+        ]
+    ),
+    high=np.array(
+        [
+            0.5,
+            0.5,
+            0.5,
+            1.570796,
+            1.570796,
+            1.570796,
+            0.174533,
+            0.488692,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            0.785398,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            1.0472,
+            1.22173,
+            0.20944,
+            0.698132,
+            1.5708,
+        ]
+    ),
+    shape=(30,),
+    dtype=np.float64,
+)
+OBS_SPACE_W_OBJ_POS = Box(
+    low=np.array(
+        [
+            -0.5,
+            -0.5,
+            -0.5,
+            -1.570796,
+            -1.570796,
+            -1.570796,
+            -0.523599,
+            -0.698132,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            0.0,
+            -0.349066,
+            -0.261799,
+            0.0,
+            0.0,
+            -1.0472,
+            0.0,
+            -0.20944,
+            -0.698132,
+            -0.261799,
+            -np.inf,
+            -np.inf,
+            -np.inf,
+        ]
+    ),
+    high=np.array(
+        [
+            0.5,
+            0.5,
+            0.5,
+            1.570796,
+            1.570796,
+            1.570796,
+            0.174533,
+            0.488692,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            0.785398,
+            0.349066,
+            1.5708,
+            1.5708,
+            1.5708,
+            1.0472,
+            1.22173,
+            0.20944,
+            0.698132,
+            1.5708,
+            np.inf,
+            np.inf,
+            np.inf,
+        ]
+    ),
+    shape=(33,),
+    dtype=np.float64,
+)
 ACTION_SPACE = Box(
     low=np.array(
         [
@@ -85,10 +231,27 @@ def env():
     return env
 
 
-def test_observation_space(env: ShadowHandGrasp):
+@pytest.fixture
+def env_w_obj_pos():
+    env = ShadowHandGrasp(
+        obj_name="insole_fixed", gui=False, observable_object_pos=True
+    )
+    env.action_space.seed(SEED)
+    return env
+
+
+def test_observation_space(
+    env: ShadowHandGrasp, env_w_obj_pos: ShadowHandGrasp
+):
     assert_array_equal(env.observation_space.low, OBS_SPACE.low)
     assert_array_equal(env.observation_space.high, OBS_SPACE.high)
     assert env.observation_space.shape[0] == OBS_SPACE.shape[0]
+    assert_array_equal(
+        env_w_obj_pos.observation_space.low, OBS_SPACE_W_OBJ_POS.low
+    )
+    assert (
+        env_w_obj_pos.observation_space.shape[0] == OBS_SPACE_W_OBJ_POS.shape[0]
+    )
 
 
 def test_action_space(env: ShadowHandGrasp):
@@ -97,12 +260,10 @@ def test_action_space(env: ShadowHandGrasp):
     assert env.action_space.shape[0] == ACTION_SPACE.shape[0]
 
 
-def test_reset(env: ShadowHandGrasp):
+def test_reset(env: ShadowHandGrasp, env_w_obj_pos: ShadowHandGrasp):
     obs, info = env.reset(seed=SEED)
     assert obs.shape[0] == 30
-    env2 = deepcopy(env)
-    env2.observable_object_pos = True
-    obs, info = env2.reset(seed=SEED)
+    obs, info = env_w_obj_pos.reset(seed=SEED)
     assert obs.shape[0] == 33