Fix list index error when goal_cell and reset_cell are passed as options in PointMaze env (#164)

* Fix list index error when goal_cell and reset_cell are passed as options

* Pre-commit formatting

* Tests for reset_cell and goal_cell

Author: gauthamvasan

gymnasium_robotics/envs/maze/maze.py

@@ -218,12 +218,11 @@ def reset(
                 assert self.maze.map_length > options["goal_cell"][1]
                 assert self.maze.map_width > options["goal_cell"][0]
                 assert (
-                    self.maze.maze_map[options["goal_cell"][1], options["goal_cell"][0]]
+                    self.maze.maze_map[options["goal_cell"][1]][options["goal_cell"][0]]
                     != 1
                 ), f"Goal can't be placed in a wall cell, {options['goal_cell']}"
 
                 goal = self.maze.cell_rowcol_to_xy(options["goal_cell"])
-
             else:
                 goal = self.generate_target_goal()
 
@@ -235,8 +234,8 @@ def reset(
                 assert self.maze.map_length > options["reset_cell"][1]
                 assert self.maze.map_width > options["reset_cell"][0]
                 assert (
-                    self.maze.maze_map[
-                        options["reset_cell"][1], options["reset_cell"][0]
+                    self.maze.maze_map[options["reset_cell"][1]][
+                        options["reset_cell"][0]
                     ]
                     != 1
                 ), f"Reset can't be placed in a wall cell, {options['reset_cell']}"

gymnasium_robotics/envs/maze/maze_v4.py

@@ -295,7 +295,7 @@ def reset(
                 assert self.maze.map_length > options["goal_cell"][1]
                 assert self.maze.map_width > options["goal_cell"][0]
                 assert (
-                    self.maze.maze_map[options["goal_cell"][1], options["goal_cell"][0]]
+                    self.maze.maze_map[options["goal_cell"][1]][options["goal_cell"][0]]
                     != 1
                 ), f"Goal can't be placed in a wall cell, {options['goal_cell']}"
 
@@ -312,8 +312,8 @@ def reset(
                 assert self.maze.map_length > options["reset_cell"][1]
                 assert self.maze.map_width > options["reset_cell"][0]
                 assert (
-                    self.maze.maze_map[
-                        options["reset_cell"][1], options["reset_cell"][0]
+                    self.maze.maze_map[options["reset_cell"][1]][
+                        options["reset_cell"][0]
                     ]
                     != 1
                 ), f"Reset can't be placed in a wall cell, {options['reset_cell']}"

tests/envs/maze/test_point_maze.py

@@ -11,3 +11,31 @@ def test_reset():
         assert not info["success"]
         dist = np.linalg.norm(obs["achieved_goal"] - obs["desired_goal"])
         assert dist > 0.45, f"dist={dist} < 0.45"
+
+
+def test_reset_cell():
+    """Check that passing the reset_cell location ensures that the agent resets in the right cell."""
+    map = [
+        [1, 1, 1, 1],
+        [1, "r", "r", 1],
+        [1, "r", "g", 1],
+        [1, 1, 1, 1],
+    ]
+    env = gym.make("PointMaze_UMaze-v3", maze_map=map)
+    obs = env.reset(options={"reset_cell": [1, 2]}, seed=42)[0]
+    desired_obs = np.array([0.67929896, 0.59868401, 0, 0])
+    np.testing.assert_almost_equal(desired_obs, obs["observation"], decimal=4)
+
+
+def test_goal_cell():
+    """Check that passing the goal_cell location ensures that the goal spawns in the right cell."""
+    map = [
+        [1, 1, 1, 1],
+        [1, "r", "g", 1],
+        [1, "g", "g", 1],
+        [1, 1, 1, 1],
+    ]
+    env = gym.make("PointMaze_UMaze-v3", maze_map=map)
+    obs = env.reset(options={"goal_cell": [2, 1]}, seed=42)[0]
+    desired_goal = np.array([-0.36302198, -0.53056078])
+    np.testing.assert_almost_equal(desired_goal, obs["desired_goal"], decimal=4)