Rolling

tests/test_envs.py

@@ -6,11 +6,13 @@ import mujoco_maze
 
 @pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
 def test_ant_maze(maze_id):
+    if "Rolling" in maze_id:
+        return
     for i in range(2):
         env = gym.make(f"Ant{maze_id}-v{i}")
         s0 = env.reset()
         s, _, _, _ = env.step(env.action_space.sample())
-        if not env.unwrapped._top_down_view and not env.unwrapped._observe_blocks:
+        if not env.unwrapped.has_extended_obs:
             assert s0.shape == (30,)
             assert s.shape == (30,)
 
@@ -21,9 +23,12 @@ def test_point_maze(maze_id):
         env = gym.make(f"Point{maze_id}-v{i}")
         s0 = env.reset()
         s, r, _, _ = env.step(env.action_space.sample())
-        if not env.unwrapped._top_down_view and not env.unwrapped._observe_blocks:
+        if not env.unwrapped.has_extended_obs:
             assert s0.shape == (7,)
             assert s.shape == (7,)
+        if env.unwrapped._observe_balls:
+            assert s0.shape == (10,)
+            assert s.shape == (10,)
         if i == 0:
             assert r != 0.0
         else:
@@ -33,42 +38,35 @@ def test_point_maze(maze_id):
 
 @pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
 def test_reacher_maze(maze_id):
-    for inhibited in ["Fall", "Push", "Block"]:
+    for inhibited in ["Fall", "Push", "Block", "Rolling"]:
         if inhibited in maze_id:
             return
     for i in range(2):
         env = gym.make(f"Reacher{maze_id}-v{i}")
         s0 = env.reset()
         s, _, _, _ = env.step(env.action_space.sample())
-        if not env.unwrapped._top_down_view and not env.unwrapped._observe_blocks:
+        if not env.unwrapped.has_extended_obs:
             assert s0.shape == (9,)
             assert s.shape == (9,)
 
 
 @pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
 def test_swimmer_maze(maze_id):
-    for inhibited in ["Fall", "Push", "Block"]:
+    for inhibited in ["Fall", "Push", "Block", "Rolling"]:
         if inhibited in maze_id:
             return
     for i in range(2):
         env = gym.make(f"Swimmer{maze_id}-v{i}")
         s0 = env.reset()
         s, _, _, _ = env.step(env.action_space.sample())
-        if not env.unwrapped._top_down_view and not env.unwrapped._observe_blocks:
+        if not env.unwrapped.has_extended_obs:
             assert s0.shape == (11,)
             assert s.shape == (11,)
 
 
 @pytest.mark.parametrize("v", [0, 1])
 def test_maze_args(v):
-    env = gym.make(f"PointTRoom-v{v}", task_kwargs={"goals": [(-2.0, -3.0)]})
-    assert env.reset().shape == (7,)
-    s, _, _, _ = env.step(env.action_space.sample())
-    assert s.shape == (7,)
-
-
-def test_getting_movable(v):
-    env = gym.make("PointBlockMaze-v1")
+    env = gym.make(f"PointTRoom-v{v}", task_kwargs={"goal": (-2.0, -3.0)})
     assert env.reset().shape == (7,)
     s, _, _, _ = env.step(env.action_space.sample())
     assert s.shape == (7,)