Some typing improvements, add py.typed (#284)

Co-authored-by: Mark Towers <mark.m.towers@gmail.com>

.gitignore

@@ -21,4 +21,5 @@ __pycache__
 
 # Virtual environments
 .env
+.venv
 venv

docs/conf.py

@@ -18,9 +18,11 @@
 # TODO: change to minigrid version
 # from TODO import __version__ as minigrid_version
 
+from __future__ import annotations
+
 import os
 import sys
-from typing import Any, Dict
+from typing import Any
 
 project = "MiniGrid"
 copyright = "2022"
@@ -79,7 +81,7 @@ html_theme_options = {
     "dark_logo": "img/minigrid-white.svg",
     "gtag": "G-FBXJQQLXKD",
 }
-html_context: Dict[str, Any] = {}
+html_context: dict[str, Any] = {}
 html_context["conf_py_path"] = "/docs/"
 html_context["display_github"] = True
 html_context["github_user"] = "Farama-Foundation"

docs/scripts/gen_envs_display.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 
 import gymnasium

docs/scripts/gen_gifs.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 

docs/scripts/move404.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import sys
 
 if __name__ == "__main__":

docs/scripts/utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 
 

minigrid/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from gymnasium.envs.registration import register
 
 from minigrid import minigrid_env, wrappers

minigrid/benchmark.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import time
 
 import gymnasium as gym

minigrid/core/actions.py

@@ -1,4 +1,6 @@
 # Enumeration of possible actions
+from __future__ import annotations
+
 from enum import IntEnum
 
 

minigrid/core/constants.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 
 TILE_PIXELS = 32

minigrid/core/grid.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import math
-from typing import Any, List, Optional, Tuple, Type
+from typing import Any, Callable
 
 import numpy as np
 
@@ -21,7 +23,7 @@ class Grid:
     """
 
     # Static cache of pre-renderer tiles
-    tile_cache = {}
+    tile_cache: dict[tuple[Any, ...], Any] = {}
 
     def __init__(self, width: int, height: int):
         assert width >= 3
@@ -30,7 +32,7 @@ class Grid:
         self.width: int = width
         self.height: int = height
 
-        self.grid: List[Optional[WorldObj]] = [None] * width * height
+        self.grid: list[WorldObj | None] = [None] * (width * height)
 
     def __contains__(self, key: Any) -> bool:
         if isinstance(key, WorldObj):
@@ -47,25 +49,29 @@ class Grid:
                     return True
         return False
 
-    def __eq__(self, other: "Grid") -> bool:
+    def __eq__(self, other: Grid) -> bool:
         grid1 = self.encode()
         grid2 = other.encode()
         return np.array_equal(grid2, grid1)
 
-    def __ne__(self, other: "Grid") -> bool:
+    def __ne__(self, other: Grid) -> bool:
         return not self == other
 
-    def copy(self) -> "Grid":
+    def copy(self) -> Grid:
         from copy import deepcopy
 
         return deepcopy(self)
 
-    def set(self, i: int, j: int, v: Optional[WorldObj]):
-        assert 0 <= i < self.width
-        assert 0 <= j < self.height
+    def set(self, i: int, j: int, v: WorldObj | None):
+        assert (
+            0 <= i < self.width
+        ), f"column index {j} outside of grid of width {self.width}"
+        assert (
+            0 <= j < self.height
+        ), f"row index {j} outside of grid of height {self.height}"
         self.grid[j * self.width + i] = v
 
-    def get(self, i: int, j: int) -> Optional[WorldObj]:
+    def get(self, i: int, j: int) -> WorldObj | None:
         assert 0 <= i < self.width
         assert 0 <= j < self.height
         assert self.grid is not None
@@ -75,8 +81,8 @@ class Grid:
         self,
         x: int,
         y: int,
-        length: Optional[int] = None,
-        obj_type: Type[WorldObj] = Wall,
+        length: int | None = None,
+        obj_type: Callable[[], WorldObj] = Wall,
     ):
         if length is None:
             length = self.width - x
@@ -87,8 +93,8 @@ class Grid:
         self,
         x: int,
         y: int,
-        length: Optional[int] = None,
-        obj_type: Type[WorldObj] = Wall,
+        length: int | None = None,
+        obj_type: Callable[[], WorldObj] = Wall,
     ):
         if length is None:
             length = self.height - y
@@ -101,7 +107,7 @@ class Grid:
         self.vert_wall(x, y, h)
         self.vert_wall(x + w - 1, y, h)
 
-    def rotate_left(self) -> "Grid":
+    def rotate_left(self) -> Grid:
         """
         Rotate the grid to the left (counter-clockwise)
         """
@@ -115,7 +121,7 @@ class Grid:
 
         return grid
 
-    def slice(self, topX: int, topY: int, width: int, height: int) -> "Grid":
+    def slice(self, topX: int, topY: int, width: int, height: int) -> Grid:
         """
         Get a subset of the grid
         """
@@ -139,8 +145,8 @@ class Grid:
     @classmethod
     def render_tile(
         cls,
-        obj: WorldObj,
-        agent_dir: Optional[int] = None,
+        obj: WorldObj | None,
+        agent_dir: int | None = None,
         highlight: bool = False,
         tile_size: int = TILE_PIXELS,
         subdivs: int = 3,
@@ -150,7 +156,7 @@ class Grid:
         """
 
         # Hash map lookup key for the cache
-        key = (agent_dir, highlight, tile_size)
+        key: tuple[Any, ...] = (agent_dir, highlight, tile_size)
         key = obj.encode() + key if obj else key
 
         if key in cls.tile_cache:
@@ -194,9 +200,9 @@ class Grid:
     def render(
         self,
         tile_size: int,
-        agent_pos: Tuple[int, int],
-        agent_dir: Optional[int] = None,
-        highlight_mask: Optional[np.ndarray] = None,
+        agent_pos: tuple[int, int],
+        agent_dir: int | None = None,
+        highlight_mask: np.ndarray | None = None,
     ) -> np.ndarray:
         """
         Render this grid at a given scale
@@ -235,7 +241,7 @@ class Grid:
 
         return img
 
-    def encode(self, vis_mask: Optional[np.ndarray] = None) -> np.ndarray:
+    def encode(self, vis_mask: np.ndarray | None = None) -> np.ndarray:
         """
         Produce a compact numpy encoding of the grid
         """
@@ -262,7 +268,7 @@ class Grid:
         return array
 
     @staticmethod
-    def decode(array: np.ndarray) -> Tuple["Grid", np.ndarray]:
+    def decode(array: np.ndarray) -> tuple[Grid, np.ndarray]:
         """
         Decode an array grid encoding back into a grid
         """
@@ -282,7 +288,7 @@ class Grid:
 
         return grid, vis_mask
 
-    def process_vis(self, agent_pos: Tuple[int, int]) -> np.ndarray:
+    def process_vis(self, agent_pos: tuple[int, int]) -> np.ndarray:
         mask = np.zeros(shape=(self.width, self.height), dtype=bool)
 
         mask[agent_pos[0], agent_pos[1]] = True

minigrid/core/mission.py

@@ -1,4 +1,6 @@
-from typing import Any, Callable, Optional, Union
+from __future__ import annotations
+
+from typing import Any, Callable
 
 from gymnasium import spaces
 from gymnasium.utils import seeding
@@ -26,8 +28,8 @@ class MissionSpace(spaces.Space[str]):
     def __init__(
         self,
         mission_func: Callable[..., str],
-        ordered_placeholders: Optional["list[list[str]]"] = None,
-        seed: Optional[Union[int, seeding.RandomNumberGenerator]] = None,
+        ordered_placeholders: list[list[str]] | None = None,
+        seed: int | seeding.RandomNumberGenerator | None = None,
     ):
         r"""Constructor of :class:`MissionSpace` space.
 

minigrid/core/roomgrid.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple, Union
+from __future__ import annotations
 
 import numpy as np
 
@@ -8,7 +8,7 @@ from minigrid.core.world_object import Ball, Box, Door, Key, WorldObj
 from minigrid.minigrid_env import MiniGridEnv
 
 
-def reject_next_to(env: MiniGridEnv, pos: Tuple[int, int]):
+def reject_next_to(env: MiniGridEnv, pos: tuple[int, int]):
     """
     Function to filter out object positions that are right next to
     the agent's starting point
@@ -21,27 +21,27 @@ def reject_next_to(env: MiniGridEnv, pos: Tuple[int, int]):
 
 
 class Room:
-    def __init__(self, top: Tuple[int, int], size: Tuple[int, int]):
+    def __init__(self, top: tuple[int, int], size: tuple[int, int]):
         # Top-left corner and size (tuples)
         self.top = top
         self.size = size
 
         # List of door objects and door positions
         # Order of the doors is right, down, left, up
-        self.doors: List[Optional[Union[bool, Door]]] = [None] * 4
-        self.door_pos: List[Optional[Tuple[int, int]]] = [None] * 4
+        self.doors: list[bool | Door | None] = [None] * 4
+        self.door_pos: list[tuple[int, int] | None] = [None] * 4
 
         # List of rooms adjacent to this one
         # Order of the neighbors is right, down, left, up
-        self.neighbors: List[Optional[Room]] = [None] * 4
+        self.neighbors: list[Room | None] = [None] * 4
 
         # Indicates if this room is behind a locked door
         self.locked: bool = False
 
         # List of objects contained
-        self.objs: List[WorldObj] = []
+        self.objs: list[WorldObj] = []
 
-    def rand_pos(self, env: MiniGridEnv) -> Tuple[int, int]:
+    def rand_pos(self, env: MiniGridEnv) -> tuple[int, int]:
         topX, topY = self.top
         sizeX, sizeY = self.size
         return env._randPos(topX + 1, topX + sizeX - 1, topY + 1, topY + sizeY - 1)
@@ -180,7 +180,7 @@ class RoomGrid(MiniGridEnv):
 
     def place_in_room(
         self, i: int, j: int, obj: WorldObj
-    ) -> Tuple[WorldObj, Tuple[int, int]]:
+    ) -> tuple[WorldObj, tuple[int, int]]:
         """
         Add an existing object to room (i, j)
         """
@@ -199,9 +199,9 @@ class RoomGrid(MiniGridEnv):
         self,
         i: int,
         j: int,
-        kind: Optional[str] = None,
-        color: Optional[str] = None,
-    ) -> Tuple[WorldObj, Tuple[int, int]]:
+        kind: str | None = None,
+        color: str | None = None,
+    ) -> tuple[WorldObj, tuple[int, int]]:
         """
         Add a new object to room (i, j)
         """
@@ -231,10 +231,10 @@ class RoomGrid(MiniGridEnv):
         self,
         i: int,
         j: int,
-        door_idx: Optional[int] = None,
-        color: Optional[str] = None,
-        locked: Optional[bool] = None,
-    ) -> Tuple[Door, Tuple[int, int]]:
+        door_idx: int | None = None,
+        color: str | None = None,
+        locked: bool | None = None,
+    ) -> tuple[Door, tuple[int, int]]:
         """
         Add a door to a room, connecting it to a neighbor
         """
@@ -311,7 +311,7 @@ class RoomGrid(MiniGridEnv):
         neighbor.doors[(wall_idx + 2) % 4] = True
 
     def place_agent(
-        self, i: Optional[int] = None, j: Optional[int] = None, rand_dir: bool = True
+        self, i: int | None = None, j: int | None = None, rand_dir: bool = True
     ) -> np.ndarray:
         """
         Place the agent in a room
@@ -334,8 +334,8 @@ class RoomGrid(MiniGridEnv):
         return self.agent_pos
 
     def connect_all(
-        self, door_colors: List[str] = COLOR_NAMES, max_itrs: int = 5000
-    ) -> List[Door]:
+        self, door_colors: list[str] = COLOR_NAMES, max_itrs: int = 5000
+    ) -> list[Door]:
         """
         Make sure that all rooms are reachable by the agent from its
         starting position
@@ -395,11 +395,11 @@ class RoomGrid(MiniGridEnv):
 
     def add_distractors(
         self,
-        i: Optional[int] = None,
-        j: Optional[int] = None,
+        i: int | None = None,
+        j: int | None = None,
         num_distractors: int = 10,
         all_unique: bool = True,
-    ) -> List[WorldObj]:
+    ) -> list[WorldObj]:
         """
         Add random objects that can potentially distract/confuse the agent.
         """

minigrid/core/world_object.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, Optional, Tuple
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
 
 import numpy as np
 
@@ -19,6 +21,8 @@ from minigrid.utils.rendering import (
 if TYPE_CHECKING:
     from minigrid.minigrid_env import MiniGridEnv
 
+Point = Tuple[int, int]
+
 
 class WorldObj:
 
@@ -34,10 +38,10 @@ class WorldObj:
         self.contains = None
 
         # Initial position of the object
-        self.init_pos = None
+        self.init_pos: Point | None = None
 
         # Current position of the object
-        self.cur_pos = None
+        self.cur_pos: Point | None = None
 
     def can_overlap(self) -> bool:
         """Can the agent overlap with this?"""
@@ -55,16 +59,16 @@ class WorldObj:
         """Can the agent see behind this object?"""
         return True
 
-    def toggle(self, env: "MiniGridEnv", pos: Tuple[int, int]) -> bool:
+    def toggle(self, env: MiniGridEnv, pos: tuple[int, int]) -> bool:
         """Method to trigger/toggle an action this object performs"""
         return False
 
-    def encode(self) -> Tuple[int, int, int]:
+    def encode(self) -> tuple[int, int, int]:
         """Encode the a description of this object as a 3-tuple of integers"""
         return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], 0)
 
     @staticmethod
-    def decode(type_idx: int, color_idx: int, state: int) -> Optional["WorldObj"]:
+    def decode(type_idx: int, color_idx: int, state: int) -> WorldObj | None:
         """Create an object from a 3-tuple state description"""
 
         obj_type = IDX_TO_OBJECT[type_idx]
@@ -267,7 +271,7 @@ class Ball(WorldObj):
 
 
 class Box(WorldObj):
-    def __init__(self, color, contains: Optional[WorldObj] = None):
+    def __init__(self, color, contains: WorldObj | None = None):
         super().__init__("box", color)
         self.contains = contains
 

minigrid/envs/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from minigrid.envs.blockedunlockpickup import BlockedUnlockPickupEnv
 from minigrid.envs.crossing import CrossingEnv
 from minigrid.envs.distshift import DistShiftEnv

minigrid/envs/babyai/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from minigrid.envs.babyai.goto import (
     GoTo,
     GoToDoor,

minigrid/envs/babyai/core/levelgen.py

@@ -1,6 +1,8 @@
 """
 Copied and adapted from https://github.com/mila-iqia/babyai
 """
+from __future__ import annotations
+
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.roomgrid import Room
 from minigrid.envs.babyai.core.roomgrid_level import RoomGridLevel
@@ -37,7 +39,7 @@ class LevelGen(RoomGridLevel):
         implicit_unlock=True,
         action_kinds=["goto", "pickup", "open", "putnext"],
         instr_kinds=["action", "and", "seq"],
-        **kwargs
+        **kwargs,
     ):
         self.num_dists = num_dists
         self.locked_room_prob = locked_room_prob

minigrid/envs/babyai/core/roomgrid_level.py

@@ -1,7 +1,7 @@
 """
 Copied and adapted from https://github.com/mila-iqia/babyai
 """
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.roomgrid import RoomGrid
 from minigrid.envs.babyai.core.verifier import (
@@ -50,7 +50,7 @@ class RoomGridLevel(RoomGrid):
     of approximately similar difficulty.
     """
 
-    def __init__(self, room_size=8, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, room_size=8, max_steps: int | None = None, **kwargs):
         mission_space = BabyAIMissionSpace()
 
         # If `max_steps` arg is passed it will be fixed for every episode,
@@ -64,7 +64,7 @@ class RoomGridLevel(RoomGrid):
             room_size=room_size,
             mission_space=mission_space,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     def reset(self, **kwargs):

minigrid/envs/babyai/core/verifier.py

@@ -1,6 +1,8 @@
 """
 Copied and adapted from https://github.com/mila-iqia/babyai
 """
+from __future__ import annotations
+
 import os
 from abc import ABC, abstractmethod
 

minigrid/envs/babyai/goto.py

@@ -2,6 +2,8 @@
 Copied and adapted from https://github.com/mila-iqia/babyai.
 Levels described in the Baby AI ICLR 2019 submission, with the `Go to` instruction.
 """
+from __future__ import annotations
+
 from minigrid.envs.babyai.core.levelgen import LevelGen
 from minigrid.envs.babyai.core.roomgrid_level import RejectSampling, RoomGridLevel
 from minigrid.envs.babyai.core.verifier import GoToInstr, ObjDesc
@@ -106,7 +108,7 @@ class GoTo(RoomGridLevel):
         num_cols=3,
         num_dists=18,
         doors_open=False,
-        **kwargs
+        **kwargs,
     ):
         self.num_dists = num_dists
         self.doors_open = doors_open
@@ -198,7 +200,7 @@ class GoToSeq(LevelGen):
             locked_room_prob=0,
             locations=False,
             unblocking=False,
-            **kwargs
+            **kwargs,
         )
 
 

minigrid/envs/babyai/open.py

@@ -2,7 +2,7 @@
 Copied and adapted from https://github.com/mila-iqia/babyai.
 Levels described in the Baby AI ICLR 2019 submission, with the `Open` instruction.
 """
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.envs.babyai.core.roomgrid_level import RoomGridLevel
@@ -103,8 +103,8 @@ class OpenTwoDoors(RoomGridLevel):
         first_color=None,
         second_color=None,
         strict=False,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         self.first_color = first_color
         self.second_color = second_color
@@ -142,9 +142,7 @@ class OpenDoorsOrder(RoomGridLevel):
     Open one or two doors in the order specified.
     """
 
-    def __init__(
-        self, num_doors, debug=False, max_steps: Optional[int] = None, **kwargs
-    ):
+    def __init__(self, num_doors, debug=False, max_steps: int | None = None, **kwargs):
         assert num_doors >= 2
         self.num_doors = num_doors
         self.debug = debug

minigrid/envs/babyai/other.py

@@ -2,7 +2,7 @@
 Copied and adapted from https://github.com/mila-iqia/babyai.
 Levels described in the Baby AI ICLR 2019 submission, with different instructions than those in other files.
 """
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.envs.babyai.core.roomgrid_level import RoomGridLevel
 from minigrid.envs.babyai.core.verifier import (
@@ -56,7 +56,7 @@ class FindObjS5(RoomGridLevel):
     This level requires potentially exhaustive exploration
     """
 
-    def __init__(self, room_size=5, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, room_size=5, max_steps: int | None = None, **kwargs):
 
         if max_steps is None:
             max_steps = 20 * room_size**2
@@ -85,8 +85,8 @@ class KeyCorridor(RoomGridLevel):
         num_rows=3,
         obj_type="ball",
         room_size=6,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         self.obj_type = obj_type
 
@@ -143,7 +143,7 @@ class MoveTwoAcross(RoomGridLevel):
     """
 
     def __init__(
-        self, room_size, objs_per_room, max_steps: Optional[int] = None, **kwargs
+        self, room_size, objs_per_room, max_steps: int | None = None, **kwargs
     ):
         assert objs_per_room <= 9
         self.objs_per_room = objs_per_room

minigrid/envs/babyai/pickup.py

@@ -2,7 +2,7 @@
 Copied and adapted from https://github.com/mila-iqia/babyai.
 Levels described in the Baby AI ICLR 2019 submission, with the `Pick up` instruction.
 """
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.envs.babyai.core.levelgen import LevelGen
 from minigrid.envs.babyai.core.roomgrid_level import RejectSampling, RoomGridLevel
@@ -63,7 +63,7 @@ class PickupLoc(LevelGen):
             locked_room_prob=0,
             locations=True,
             unblocking=False,
-            **kwargs
+            **kwargs,
         )
 
 
@@ -102,7 +102,7 @@ class PickupAbove(RoomGridLevel):
     This task requires to use the compass to be solved effectively.
     """
 
-    def __init__(self, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, max_steps: int | None = None, **kwargs):
         room_size = 6
         if max_steps is None:
             max_steps = 8 * room_size**2

minigrid/envs/babyai/putnext.py

@@ -2,7 +2,7 @@
 Copied and adapted from https://github.com/mila-iqia/babyai.
 Levels described in the Baby AI ICLR 2019 submission, with the `Put Next` instruction.
 """
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.envs.babyai.core.roomgrid_level import RoomGridLevel
 from minigrid.envs.babyai.core.verifier import ObjDesc, PutNextInstr
@@ -41,8 +41,8 @@ class PutNext(RoomGridLevel):
         room_size,
         objs_per_room,
         start_carrying=False,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         assert room_size >= 4
         assert objs_per_room <= 9
@@ -86,6 +86,7 @@ class PutNext(RoomGridLevel):
 
         # If the agent starts off carrying the object
         if self.start_carrying:
+            assert self.obj_a.init_pos is not None
             self.grid.set(*self.obj_a.init_pos, None)
             self.carrying = self.obj_a
 

minigrid/envs/babyai/synth.py

@@ -4,6 +4,8 @@ Levels described in the Baby AI ICLR 2019 submission.
 The instructions are a synthesis of those from `PutNext`, `Open`, `GoTo`, and `Pickup`.
 """
 
+from __future__ import annotations
+
 from minigrid.envs.babyai.core.levelgen import LevelGen
 
 
@@ -28,7 +30,7 @@ class Synth(LevelGen):
             locations=False,
             unblocking=True,
             implicit_unlock=False,
-            **kwargs
+            **kwargs,
         )
 
 
@@ -53,7 +55,7 @@ class SynthLoc(LevelGen):
             locations=True,
             unblocking=True,
             implicit_unlock=False,
-            **kwargs
+            **kwargs,
         )
 
 
@@ -81,7 +83,7 @@ class MiniBossLevel(LevelGen):
             room_size=5,
             num_dists=7,
             locked_room_prob=0.25,
-            **kwargs
+            **kwargs,
         )
 
 

minigrid/envs/babyai/unlock.py

@@ -2,7 +2,7 @@
 Copied and adapted from https://github.com/mila-iqia/babyai.
 Levels described in the Baby AI ICLR 2019 submission, with the `Unlock` instruction.
 """
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.world_object import Ball, Box, Key
@@ -110,7 +110,7 @@ class UnlockPickup(RoomGridLevel):
     Unlock a door, then pick up a box in another room
     """
 
-    def __init__(self, distractors=False, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, distractors=False, max_steps: int | None = None, **kwargs):
         self.distractors = distractors
         room_size = 6
         if max is None:
@@ -141,7 +141,7 @@ class BlockedUnlockPickup(RoomGridLevel):
     in another room
     """
 
-    def __init__(self, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, max_steps: int | None = None, **kwargs):
         room_size = 6
         if max_steps is None:
             max_steps = 16 * room_size**2
@@ -171,7 +171,7 @@ class UnlockToUnlock(RoomGridLevel):
     Unlock a door A that requires to unlock a door B before
     """
 
-    def __init__(self, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, max_steps: int | None = None, **kwargs):
         room_size = 6
         if max_steps is None:
             max_steps = 30 * room_size**2

minigrid/envs/blockedunlockpickup.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.mission import MissionSpace
@@ -64,7 +64,7 @@ class BlockedUnlockPickupEnv(RoomGrid):
 
     """
 
-    def __init__(self, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, max_steps: int | None = None, **kwargs):
         mission_space = MissionSpace(
             mission_func=self._gen_mission,
             ordered_placeholders=[COLOR_NAMES, ["box", "key"]],

minigrid/envs/crossing.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 import itertools as itt
-from typing import Optional
 
 import numpy as np
 
@@ -88,8 +89,8 @@ class CrossingEnv(MiniGridEnv):
         size=9,
         num_crossings=1,
         obstacle_type=Lava,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         self.num_crossings = num_crossings
         self.obstacle_type = obstacle_type
@@ -107,7 +108,7 @@ class CrossingEnv(MiniGridEnv):
             grid_size=size,
             see_through_walls=False,  # Set this to True for maximum speed
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/distshift.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
@@ -69,8 +69,8 @@ class DistShiftEnv(MiniGridEnv):
         agent_start_pos=(1, 1),
         agent_start_dir=0,
         strip2_row=2,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         self.agent_start_pos = agent_start_pos
         self.agent_start_dir = agent_start_dir
@@ -89,7 +89,7 @@ class DistShiftEnv(MiniGridEnv):
             # Set this to True for maximum speed
             see_through_walls=True,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/doorkey.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
@@ -60,7 +60,7 @@ class DoorKeyEnv(MiniGridEnv):
 
     """
 
-    def __init__(self, size=8, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=8, max_steps: int | None = None, **kwargs):
         if max_steps is None:
             max_steps = 10 * size**2
         mission_space = MissionSpace(mission_func=self._gen_mission)

minigrid/envs/dynamicobstacles.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 from operator import add
-from typing import Optional
 
 from gymnasium.spaces import Discrete
 
@@ -74,8 +75,8 @@ class DynamicObstaclesEnv(MiniGridEnv):
         agent_start_pos=(1, 1),
         agent_start_dir=0,
         n_obstacles=4,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         self.agent_start_pos = agent_start_pos
         self.agent_start_dir = agent_start_dir
@@ -97,7 +98,7 @@ class DynamicObstaclesEnv(MiniGridEnv):
             # Set this to True for maximum speed
             see_through_walls=True,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
         # Allow only 3 actions permitted: left, right, forward
         self.action_space = Discrete(self.actions.forward + 1)

minigrid/envs/empty.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
@@ -70,8 +70,8 @@ class EmptyEnv(MiniGridEnv):
         size=8,
         agent_start_pos=(1, 1),
         agent_start_dir=0,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         self.agent_start_pos = agent_start_pos
         self.agent_start_dir = agent_start_dir
@@ -87,7 +87,7 @@ class EmptyEnv(MiniGridEnv):
             # Set this to True for maximum speed
             see_through_walls=True,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/fetch.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
@@ -71,7 +71,7 @@ class FetchEnv(MiniGridEnv):
 
     """
 
-    def __init__(self, size=8, numObjs=3, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=8, numObjs=3, max_steps: int | None = None, **kwargs):
         self.numObjs = numObjs
         self.obj_types = ["key", "ball"]
 

minigrid/envs/fourrooms.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
 from minigrid.core.world_object import Goal
@@ -67,7 +69,7 @@ class FourRoomsEnv(MiniGridEnv):
             width=self.size,
             height=self.size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/gotodoor.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
@@ -63,7 +63,7 @@ class GoToDoorEnv(MiniGridEnv):
 
     """
 
-    def __init__(self, size=5, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=5, max_steps: int | None = None, **kwargs):
         assert size >= 5
         self.size = size
         mission_space = MissionSpace(

minigrid/envs/gotoobject.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
@@ -13,7 +13,7 @@ class GoToObjectEnv(MiniGridEnv):
     named using an English text string
     """
 
-    def __init__(self, size=6, numObjs=2, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=6, numObjs=2, max_steps: int | None = None, **kwargs):
 
         self.numObjs = numObjs
         self.size = size

minigrid/envs/keycorridor.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.mission import MissionSpace
@@ -77,7 +77,7 @@ class KeyCorridorEnv(RoomGrid):
         num_rows=3,
         obj_type="ball",
         room_size=6,
-        max_steps: Optional[int] = None,
+        max_steps: int | None = None,
         **kwargs,
     ):
         self.obj_type = obj_type

minigrid/envs/lavagap.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 import numpy as np
 
@@ -67,7 +67,7 @@ class LavaGapEnv(MiniGridEnv):
     """
 
     def __init__(
-        self, size, obstacle_type=Lava, max_steps: Optional[int] = None, **kwargs
+        self, size, obstacle_type=Lava, max_steps: int | None = None, **kwargs
     ):
         self.obstacle_type = obstacle_type
         self.size = size
@@ -87,7 +87,7 @@ class LavaGapEnv(MiniGridEnv):
             # Set this to True for maximum speed
             see_through_walls=False,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/lockedroom.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
@@ -76,7 +76,7 @@ class LockedRoomEnv(MiniGridEnv):
 
     """
 
-    def __init__(self, size=19, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=19, max_steps: int | None = None, **kwargs):
         self.size = size
 
         if max_steps is None:

minigrid/envs/memory.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 import numpy as np
 
@@ -68,7 +68,7 @@ class MemoryEnv(MiniGridEnv):
     """
 
     def __init__(
-        self, size=8, random_length=False, max_steps: Optional[int] = None, **kwargs
+        self, size=8, random_length=False, max_steps: int | None = None, **kwargs
     ):
         self.size = size
         self.random_length = random_length
@@ -84,7 +84,7 @@ class MemoryEnv(MiniGridEnv):
             # Set this to True for maximum speed
             see_through_walls=False,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/multiroom.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
@@ -77,8 +77,8 @@ class MultiRoomEnv(MiniGridEnv):
         minNumRooms,
         maxNumRooms,
         maxRoomSize=10,
-        max_steps: Optional[int] = None,
-        **kwargs
+        max_steps: int | None = None,
+        **kwargs,
     ):
         assert minNumRooms > 0
         assert maxNumRooms >= minNumRooms
@@ -102,7 +102,7 @@ class MultiRoomEnv(MiniGridEnv):
             width=self.size,
             height=self.size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/obstructedmaze.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES, DIR_TO_VEC
 from minigrid.core.mission import MissionSpace
@@ -76,7 +76,7 @@ class ObstructedMazeEnv(RoomGrid):
         num_rows,
         num_cols,
         num_rooms_visited,
-        max_steps: Optional[int] = None,
+        max_steps: int | None = None,
         **kwargs,
     ):
         room_size = 6

minigrid/envs/playground.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
@@ -19,7 +21,7 @@ class PlaygroundEnv(MiniGridEnv):
             width=self.size,
             height=self.size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/putnear.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.grid import Grid
@@ -67,7 +67,7 @@ class PutNearEnv(MiniGridEnv):
 
     """
 
-    def __init__(self, size=6, numObjs=2, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=6, numObjs=2, max_steps: int | None = None, **kwargs):
         self.size = size
         self.numObjs = numObjs
         self.obj_types = ["key", "ball", "box"]

minigrid/envs/redbluedoors.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
@@ -59,7 +59,7 @@ class RedBlueDoorEnv(MiniGridEnv):
 
     """
 
-    def __init__(self, size=8, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, size=8, max_steps: int | None = None, **kwargs):
         self.size = size
         mission_space = MissionSpace(mission_func=self._gen_mission)
 

minigrid/envs/unlock.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.mission import MissionSpace
 from minigrid.core.roomgrid import RoomGrid
@@ -53,7 +53,7 @@ class UnlockEnv(RoomGrid):
 
     """
 
-    def __init__(self, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, max_steps: int | None = None, **kwargs):
         room_size = 6
         mission_space = MissionSpace(mission_func=self._gen_mission)
 
@@ -66,7 +66,7 @@ class UnlockEnv(RoomGrid):
             num_cols=2,
             room_size=room_size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod

minigrid/envs/unlockpickup.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 from minigrid.core.constants import COLOR_NAMES
 from minigrid.core.mission import MissionSpace
@@ -57,7 +57,7 @@ class UnlockPickupEnv(RoomGrid):
 
     """
 
-    def __init__(self, max_steps: Optional[int] = None, **kwargs):
+    def __init__(self, max_steps: int | None = None, **kwargs):
         room_size = 6
         mission_space = MissionSpace(
             mission_func=self._gen_mission,

minigrid/manual_control.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import gymnasium as gym
 
 from minigrid.minigrid_env import MiniGridEnv

minigrid/minigrid_env.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import hashlib
 import math
 from abc import abstractmethod
 from enum import IntEnum
-from typing import Optional
+from typing import Iterable, TypeVar
 
 import gymnasium as gym
 import numpy as np
@@ -11,8 +13,11 @@ from gymnasium import spaces
 from minigrid.core.constants import COLOR_NAMES, DIR_TO_VEC, TILE_PIXELS
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
+from minigrid.core.world_object import Point, WorldObj
 from minigrid.utils.window import Window
 
+T = TypeVar("T")
+
 
 class MiniGridEnv(gym.Env):
     """
@@ -43,13 +48,13 @@ class MiniGridEnv(gym.Env):
     def __init__(
         self,
         mission_space: MissionSpace,
-        grid_size: int = None,
-        width: int = None,
-        height: int = None,
+        grid_size: int | None = None,
+        width: int | None = None,
+        height: int | None = None,
         max_steps: int = 100,
         see_through_walls: bool = False,
         agent_view_size: int = 7,
-        render_mode: Optional[str] = None,
+        render_mode: str | None = None,
         highlight: bool = True,
         tile_size: int = TILE_PIXELS,
         agent_pov: bool = False,
@@ -62,6 +67,7 @@ class MiniGridEnv(gym.Env):
             assert width is None and height is None
             width = grid_size
             height = grid_size
+        assert width is not None and height is not None
 
         # Action enumeration for this environment
         self.actions = MiniGridEnv.Actions
@@ -107,7 +113,7 @@ class MiniGridEnv(gym.Env):
         self.see_through_walls = see_through_walls
 
         # Current position and direction of the agent
-        self.agent_pos: np.ndarray = None
+        self.agent_pos: np.ndarray | tuple[int, int] = None
         self.agent_dir: int = None
 
         # Current grid and mission and carrying
@@ -228,35 +234,35 @@ class MiniGridEnv(gym.Env):
     def _gen_grid(self, width, height):
         pass
 
-    def _reward(self):
+    def _reward(self) -> float:
         """
         Compute the reward to be given upon success
         """
 
         return 1 - 0.9 * (self.step_count / self.max_steps)
 
-    def _rand_int(self, low, high):
+    def _rand_int(self, low: int, high: int) -> int:
         """
         Generate random integer in [low,high[
         """
 
         return self.np_random.integers(low, high)
 
-    def _rand_float(self, low, high):
+    def _rand_float(self, low: float, high: float) -> float:
         """
         Generate random float in [low,high[
         """
 
         return self.np_random.uniform(low, high)
 
-    def _rand_bool(self):
+    def _rand_bool(self) -> bool:
         """
         Generate random boolean value
         """
 
         return self.np_random.integers(0, 2) == 0
 
-    def _rand_elem(self, iterable):
+    def _rand_elem(self, iterable: Iterable[T]) -> T:
         """
         Pick a random element in a list
         """
@@ -265,7 +271,7 @@ class MiniGridEnv(gym.Env):
         idx = self._rand_int(0, len(lst))
         return lst[idx]
 
-    def _rand_subset(self, iterable, num_elems):
+    def _rand_subset(self, iterable: Iterable[T], num_elems: int) -> list[T]:
         """
         Sample a random subset of distinct elements of a list
         """
@@ -273,7 +279,7 @@ class MiniGridEnv(gym.Env):
         lst = list(iterable)
         assert num_elems <= len(lst)
 
-        out = []
+        out: list[T] = []
 
         while len(out) < num_elems:
             elem = self._rand_elem(lst)
@@ -282,24 +288,33 @@ class MiniGridEnv(gym.Env):
 
         return out
 
-    def _rand_color(self):
+    def _rand_color(self) -> str:
         """
         Generate a random color name (string)
         """
 
         return self._rand_elem(COLOR_NAMES)
 
-    def _rand_pos(self, xLow, xHigh, yLow, yHigh):
+    def _rand_pos(
+        self, x_low: int, x_high: int, y_low: int, y_high: int
+    ) -> tuple[int, int]:
         """
         Generate a random (x,y) position tuple
         """
 
         return (
-            self.np_random.integers(xLow, xHigh),
-            self.np_random.integers(yLow, yHigh),
+            self.np_random.integers(x_low, x_high),
+            self.np_random.integers(y_low, y_high),
         )
 
-    def place_obj(self, obj, top=None, size=None, reject_fn=None, max_tries=math.inf):
+    def place_obj(
+        self,
+        obj: WorldObj | None,
+        top: Point = None,
+        size: tuple[int, int] = None,
+        reject_fn=None,
+        max_tries=math.inf,
+    ):
         """
         Place an object at an empty position in the grid
 
@@ -326,15 +341,11 @@ class MiniGridEnv(gym.Env):
 
             num_tries += 1
 
-            pos = np.array(
-                (
-                    self._rand_int(top[0], min(top[0] + size[0], self.grid.width)),
-                    self._rand_int(top[1], min(top[1] + size[1], self.grid.height)),
-                )
+            pos = (
+                self._rand_int(top[0], min(top[0] + size[0], self.grid.width)),
+                self._rand_int(top[1], min(top[1] + size[1], self.grid.height)),
             )
 
-            pos = tuple(pos)
-
             # Don't place the object on top of another object
             if self.grid.get(*pos) is not None:
                 continue
@@ -357,7 +368,7 @@ class MiniGridEnv(gym.Env):
 
         return pos
 
-    def put_obj(self, obj, i, j):
+    def put_obj(self, obj: WorldObj, i: int, j: int):
         """
         Put an object at a specific position in the grid
         """
@@ -387,7 +398,9 @@ class MiniGridEnv(gym.Env):
         of forward movement.
         """
 
-        assert self.agent_dir >= 0 and self.agent_dir < 4
+        assert (
+            self.agent_dir >= 0 and self.agent_dir < 4
+        ), f"Invalid agent_dir: {self.agent_dir} is not within range(0, 4)"
         return DIR_TO_VEC[self.agent_dir]
 
     @property

minigrid/utils/rendering.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 
 import numpy as np

minigrid/wrappers.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 import operator
 from functools import reduce

pyproject.toml

@@ -35,3 +35,8 @@ reportPrivateImportUsage = "none"
 
 [tool.pytest.ini_options]
 filterwarnings = ['ignore:.*step API.*:DeprecationWarning'] # TODO: to be removed when old step API is removed
+
+[tool.isort]
+profile = "black"
+add_imports = [ "from __future__ import annotations" ]
+append_only = true

setup.py

@@ -1,5 +1,7 @@
 """Setups up the Minigrid module."""
 
+from __future__ import annotations
+
 from setuptools import find_packages, setup
 
 

tests/test_envs.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pickle
 import warnings
 

tests/test_scripts.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import gymnasium as gym
 import numpy as np
 from pytest_mock import MockerFixture

tests/test_wrappers.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 
 import gymnasium as gym

tests/utils.py

@@ -1,4 +1,6 @@
 """Finds all the specs that we can test with"""
+from __future__ import annotations
+
 import gymnasium as gym
 import numpy as np