Use top_down_view in Push and Fall

tests/test_envs.py

@@ -8,18 +8,22 @@ import mujoco_maze
 def test_ant_maze(maze_id):
     for i in range(2):
         env = gym.make(f"Ant{maze_id}-v{i}")
-        assert env.reset().shape == (30,)
+        s0 = env.reset()
         s, _, _, _ = env.step(env.action_space.sample())
-        assert s.shape == (30,)
+        if not env.unwrapped._top_down_view:
+            assert s0.shape == (30,)
+            assert s.shape == (30,)
 
 
 @pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
 def test_point_maze(maze_id):
     for i in range(2):
         env = gym.make(f"Point{maze_id}-v{i}")
-        assert env.reset().shape == (7,)
+        s0 = env.reset()
         s, _, _, _ = env.step(env.action_space.sample())
-        assert s.shape == (7,)
+        if not env.unwrapped._top_down_view:
+            assert s0.shape == (7,)
+            assert s.shape == (7,)
 
 
 @pytest.mark.parametrize("v", [0, 1])