From fc41b5e7c3468c44a9c80bdbbc197cf6d5d6bc7d Mon Sep 17 00:00:00 2001
From: Michael Joseph Rosenthal <rosenthalm93@gmail.com>
Date: Sat, 12 Nov 2022 09:58:51 -0800
Subject: [PATCH] Some typing improvements, add py.typed (#284)

Co-authored-by: Mark Towers <mark.m.towers@gmail.com>
---
 .gitignore                                  |  1 +
 docs/conf.py                                |  6 +-
 docs/scripts/gen_envs_display.py            |  2 +
 docs/scripts/gen_gifs.py                    |  2 +
 docs/scripts/move404.py                     |  2 +
 docs/scripts/utils.py                       |  2 +
 minigrid/__init__.py                        |  2 +
 minigrid/benchmark.py                       |  2 +
 minigrid/core/actions.py                    |  2 +
 minigrid/core/constants.py                  |  2 +
 minigrid/core/grid.py                       | 56 +++++++++--------
 minigrid/core/mission.py                    |  8 ++-
 minigrid/core/roomgrid.py                   | 44 +++++++-------
 minigrid/core/world_object.py               | 18 +++---
 minigrid/envs/__init__.py                   |  2 +
 minigrid/envs/babyai/__init__.py            |  2 +
 minigrid/envs/babyai/core/levelgen.py       |  4 +-
 minigrid/envs/babyai/core/roomgrid_level.py |  6 +-
 minigrid/envs/babyai/core/verifier.py       |  2 +
 minigrid/envs/babyai/goto.py                |  6 +-
 minigrid/envs/babyai/open.py                | 10 ++-
 minigrid/envs/babyai/other.py               | 10 +--
 minigrid/envs/babyai/pickup.py              |  6 +-
 minigrid/envs/babyai/putnext.py             |  7 ++-
 minigrid/envs/babyai/synth.py               |  8 ++-
 minigrid/envs/babyai/unlock.py              |  8 +--
 minigrid/envs/blockedunlockpickup.py        |  4 +-
 minigrid/envs/crossing.py                   |  9 +--
 minigrid/envs/distshift.py                  |  8 +--
 minigrid/envs/doorkey.py                    |  4 +-
 minigrid/envs/dynamicobstacles.py           |  9 +--
 minigrid/envs/empty.py                      |  8 +--
 minigrid/envs/fetch.py                      |  4 +-
 minigrid/envs/fourrooms.py                  |  4 +-
 minigrid/envs/gotodoor.py                   |  4 +-
 minigrid/envs/gotoobject.py                 |  4 +-
 minigrid/envs/keycorridor.py                |  4 +-
 minigrid/envs/lavagap.py                    |  6 +-
 minigrid/envs/lockedroom.py                 |  4 +-
 minigrid/envs/memory.py                     |  6 +-
 minigrid/envs/multiroom.py                  |  8 +--
 minigrid/envs/obstructedmaze.py             |  4 +-
 minigrid/envs/playground.py                 |  4 +-
 minigrid/envs/putnear.py                    |  4 +-
 minigrid/envs/redbluedoors.py               |  4 +-
 minigrid/envs/unlock.py                     |  6 +-
 minigrid/envs/unlockpickup.py               |  4 +-
 minigrid/manual_control.py                  |  2 +
 minigrid/minigrid_env.py                    | 67 ++++++++++++---------
 minigrid/py.typed                           |  0
 minigrid/utils/rendering.py                 |  2 +
 minigrid/wrappers.py                        |  2 +
 pyproject.toml                              |  5 ++
 setup.py                                    |  2 +
 tests/test_envs.py                          |  2 +
 tests/test_scripts.py                       |  2 +
 tests/test_wrappers.py                      |  2 +
 tests/utils.py                              |  2 +
 58 files changed, 251 insertions(+), 169 deletions(-)
 create mode 100644 minigrid/py.typed

diff --git a/.gitignore b/.gitignore
index 8de9f0234..a2e664972 100644
--- a/.gitignore
+++ b/.gitignore
@@ -21,4 +21,5 @@ __pycache__
 
 # Virtual environments
 .env
+.venv
 venv
diff --git a/docs/conf.py b/docs/conf.py
index 65abeb898..a5de26ea9 100644
--- a/docs/conf.py
+++ b/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 @@
     "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"
diff --git a/docs/scripts/gen_envs_display.py b/docs/scripts/gen_envs_display.py
index 122695529..9d0c76067 100644
--- a/docs/scripts/gen_envs_display.py
+++ b/docs/scripts/gen_envs_display.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 
 import gymnasium
diff --git a/docs/scripts/gen_gifs.py b/docs/scripts/gen_gifs.py
index c0ae276b3..e0712e487 100644
--- a/docs/scripts/gen_gifs.py
+++ b/docs/scripts/gen_gifs.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 
diff --git a/docs/scripts/move404.py b/docs/scripts/move404.py
index b53bae55f..ef6a18037 100644
--- a/docs/scripts/move404.py
+++ b/docs/scripts/move404.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import sys
 
 if __name__ == "__main__":
diff --git a/docs/scripts/utils.py b/docs/scripts/utils.py
index 0521eaca0..a182f3035 100644
--- a/docs/scripts/utils.py
+++ b/docs/scripts/utils.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 
 
diff --git a/minigrid/__init__.py b/minigrid/__init__.py
index 6a1f71da3..ea5b42da1 100644
--- a/minigrid/__init__.py
+++ b/minigrid/__init__.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from gymnasium.envs.registration import register
 
 from minigrid import minigrid_env, wrappers
diff --git a/minigrid/benchmark.py b/minigrid/benchmark.py
index d1ed1e58b..422a9bf86 100755
--- a/minigrid/benchmark.py
+++ b/minigrid/benchmark.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import time
 
 import gymnasium as gym
diff --git a/minigrid/core/actions.py b/minigrid/core/actions.py
index 24d3a94fb..b440438bc 100644
--- a/minigrid/core/actions.py
+++ b/minigrid/core/actions.py
@@ -1,4 +1,6 @@
 # Enumeration of possible actions
+from __future__ import annotations
+
 from enum import IntEnum
 
 
diff --git a/minigrid/core/constants.py b/minigrid/core/constants.py
index e3e524e6e..9763c63b0 100644
--- a/minigrid/core/constants.py
+++ b/minigrid/core/constants.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 
 TILE_PIXELS = 32
diff --git a/minigrid/core/grid.py b/minigrid/core/grid.py
index 3941f9a17..ccbd85211 100644
--- a/minigrid/core/grid.py
+++ b/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 @@ def __init__(self, width: int, height: int):
         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 @@ def __contains__(self, key: Any) -> bool:
                     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 @@ def horz_wall(
         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 @@ def vert_wall(
         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 @@ def wall_rect(self, x: int, y: int, w: int, h: int):
         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 @@ def rotate_left(self) -> "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 @@ def slice(self, topX: int, topY: int, width: int, height: int) -> "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 @@ def render_tile(
         """
 
         # 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 @@ def render_tile(
     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 @@ def render(
 
         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 @@ def encode(self, vis_mask: Optional[np.ndarray] = None) -> np.ndarray:
         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 @@ def decode(array: np.ndarray) -> Tuple["Grid", np.ndarray]:
 
         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
diff --git a/minigrid/core/mission.py b/minigrid/core/mission.py
index 40810e483..bca9c11d3 100644
--- a/minigrid/core/mission.py
+++ b/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.
 
diff --git a/minigrid/core/roomgrid.py b/minigrid/core/roomgrid.py
index 5923743f5..c0ba6fb60 100644
--- a/minigrid/core/roomgrid.py
+++ b/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.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 @@ def _gen_grid(self, width, height):
 
     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 @@ def add_object(
         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 @@ def add_door(
         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 @@ def remove_wall(self, i: int, j: int, wall_idx: int):
         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 @@ def place_agent(
         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 @@ def find_reach():
 
     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.
         """
diff --git a/minigrid/core/world_object.py b/minigrid/core/world_object.py
index 9b3394e0d..592be953a 100644
--- a/minigrid/core/world_object.py
+++ b/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 @@
 if TYPE_CHECKING:
     from minigrid.minigrid_env import MiniGridEnv
 
+Point = Tuple[int, int]
+
 
 class WorldObj:
 
@@ -34,10 +38,10 @@ def __init__(self, type: str, color: str):
         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 @@ def see_behind(self) -> bool:
         """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 @@ def render(self, img):
 
 
 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
 
diff --git a/minigrid/envs/__init__.py b/minigrid/envs/__init__.py
index 16794bc26..5401e7157 100644
--- a/minigrid/envs/__init__.py
+++ b/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
diff --git a/minigrid/envs/babyai/__init__.py b/minigrid/envs/babyai/__init__.py
index e9d8534dc..fb39fbf48 100644
--- a/minigrid/envs/babyai/__init__.py
+++ b/minigrid/envs/babyai/__init__.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from minigrid.envs.babyai.goto import (
     GoTo,
     GoToDoor,
diff --git a/minigrid/envs/babyai/core/levelgen.py b/minigrid/envs/babyai/core/levelgen.py
index fb011a7e9..f99c2e32c 100644
--- a/minigrid/envs/babyai/core/levelgen.py
+++ b/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 @@ def __init__(
         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
diff --git a/minigrid/envs/babyai/core/roomgrid_level.py b/minigrid/envs/babyai/core/roomgrid_level.py
index 0b1c38db2..7de73c6d3 100644
--- a/minigrid/envs/babyai/core/roomgrid_level.py
+++ b/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 @@ def __init__(self, room_size=8, max_steps: Optional[int] = None, **kwargs):
             room_size=room_size,
             mission_space=mission_space,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     def reset(self, **kwargs):
diff --git a/minigrid/envs/babyai/core/verifier.py b/minigrid/envs/babyai/core/verifier.py
index 129cce5cf..260ff949e 100644
--- a/minigrid/envs/babyai/core/verifier.py
+++ b/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
 
diff --git a/minigrid/envs/babyai/goto.py b/minigrid/envs/babyai/goto.py
index ff8e57228..b771d1b8f 100644
--- a/minigrid/envs/babyai/goto.py
+++ b/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 @@ def __init__(
         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 @@ def __init__(self, room_size=8, num_rows=3, num_cols=3, num_dists=18, **kwargs):
             locked_room_prob=0,
             locations=False,
             unblocking=False,
-            **kwargs
+            **kwargs,
         )
 
 
diff --git a/minigrid/envs/babyai/open.py b/minigrid/envs/babyai/open.py
index 3ae6d9bb3..640aff037 100644
--- a/minigrid/envs/babyai/open.py
+++ b/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 @@ def __init__(
         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
diff --git a/minigrid/envs/babyai/other.py b/minigrid/envs/babyai/other.py
index 6b990d50b..d1d8fdc1b 100644
--- a/minigrid/envs/babyai/other.py
+++ b/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 @@ def __init__(
         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
diff --git a/minigrid/envs/babyai/pickup.py b/minigrid/envs/babyai/pickup.py
index 96fdb89f6..a9ff7629a 100644
--- a/minigrid/envs/babyai/pickup.py
+++ b/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 @@ def __init__(self, **kwargs):
             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
diff --git a/minigrid/envs/babyai/putnext.py b/minigrid/envs/babyai/putnext.py
index 36b4a1072..204c3645f 100644
--- a/minigrid/envs/babyai/putnext.py
+++ b/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 @@ def __init__(
         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 @@ def reset(self, **kwargs):
 
         # 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
 
diff --git a/minigrid/envs/babyai/synth.py b/minigrid/envs/babyai/synth.py
index 710f8a69d..181ab96a1 100644
--- a/minigrid/envs/babyai/synth.py
+++ b/minigrid/envs/babyai/synth.py
@@ -4,6 +4,8 @@
 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 @@ def __init__(self, room_size=8, num_rows=3, num_cols=3, num_dists=18, **kwargs):
             locations=False,
             unblocking=True,
             implicit_unlock=False,
-            **kwargs
+            **kwargs,
         )
 
 
@@ -53,7 +55,7 @@ def __init__(self, **kwargs):
             locations=True,
             unblocking=True,
             implicit_unlock=False,
-            **kwargs
+            **kwargs,
         )
 
 
@@ -81,7 +83,7 @@ def __init__(self, **kwargs):
             room_size=5,
             num_dists=7,
             locked_room_prob=0.25,
-            **kwargs
+            **kwargs,
         )
 
 
diff --git a/minigrid/envs/babyai/unlock.py b/minigrid/envs/babyai/unlock.py
index 341ac23b8..e0f1bbd76 100644
--- a/minigrid/envs/babyai/unlock.py
+++ b/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
diff --git a/minigrid/envs/blockedunlockpickup.py b/minigrid/envs/blockedunlockpickup.py
index ec3aa5328..942a8964a 100644
--- a/minigrid/envs/blockedunlockpickup.py
+++ b/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"]],
diff --git a/minigrid/envs/crossing.py b/minigrid/envs/crossing.py
index 042646d39..c0acf6088 100644
--- a/minigrid/envs/crossing.py
+++ b/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 @@ def __init__(
         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 @@ def __init__(
             grid_size=size,
             see_through_walls=False,  # Set this to True for maximum speed
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/distshift.py b/minigrid/envs/distshift.py
index b2a4dcba6..30579c1ae 100644
--- a/minigrid/envs/distshift.py
+++ b/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 @@ def __init__(
         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 @@ def __init__(
             # Set this to True for maximum speed
             see_through_walls=True,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/doorkey.py b/minigrid/envs/doorkey.py
index e948b5880..221fe1753 100644
--- a/minigrid/envs/doorkey.py
+++ b/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)
diff --git a/minigrid/envs/dynamicobstacles.py b/minigrid/envs/dynamicobstacles.py
index 49309a687..817468b7c 100644
--- a/minigrid/envs/dynamicobstacles.py
+++ b/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 @@ def __init__(
         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 @@ def __init__(
             # 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)
diff --git a/minigrid/envs/empty.py b/minigrid/envs/empty.py
index d2c689bd2..6a0389e1d 100644
--- a/minigrid/envs/empty.py
+++ b/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 @@ def __init__(
         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 @@ def __init__(
             # Set this to True for maximum speed
             see_through_walls=True,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/fetch.py b/minigrid/envs/fetch.py
index 8cf29ff4c..3bc3408c6 100644
--- a/minigrid/envs/fetch.py
+++ b/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"]
 
diff --git a/minigrid/envs/fourrooms.py b/minigrid/envs/fourrooms.py
index be3acd4f0..dcb87fddd 100644
--- a/minigrid/envs/fourrooms.py
+++ b/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 @@ def __init__(self, agent_pos=None, goal_pos=None, max_steps=100, **kwargs):
             width=self.size,
             height=self.size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/gotodoor.py b/minigrid/envs/gotodoor.py
index 513df3868..deb288f4c 100644
--- a/minigrid/envs/gotodoor.py
+++ b/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(
diff --git a/minigrid/envs/gotoobject.py b/minigrid/envs/gotoobject.py
index 4bee7c317..8bd8b841f 100644
--- a/minigrid/envs/gotoobject.py
+++ b/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
diff --git a/minigrid/envs/keycorridor.py b/minigrid/envs/keycorridor.py
index 2e0a47442..3cf3fe4e8 100644
--- a/minigrid/envs/keycorridor.py
+++ b/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 @@ def __init__(
         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
diff --git a/minigrid/envs/lavagap.py b/minigrid/envs/lavagap.py
index b1fe8beb4..d709a649b 100644
--- a/minigrid/envs/lavagap.py
+++ b/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 @@ def __init__(
             # Set this to True for maximum speed
             see_through_walls=False,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/lockedroom.py b/minigrid/envs/lockedroom.py
index 439e86387..8c742da32 100644
--- a/minigrid/envs/lockedroom.py
+++ b/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:
diff --git a/minigrid/envs/memory.py b/minigrid/envs/memory.py
index f82c0b73f..0ef729d81 100644
--- a/minigrid/envs/memory.py
+++ b/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 @@ def __init__(
             # Set this to True for maximum speed
             see_through_walls=False,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/multiroom.py b/minigrid/envs/multiroom.py
index 2a45d0982..bc4c9145d 100644
--- a/minigrid/envs/multiroom.py
+++ b/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 @@ def __init__(
         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 @@ def __init__(
             width=self.size,
             height=self.size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/obstructedmaze.py b/minigrid/envs/obstructedmaze.py
index 42c0a6c5c..fcbcfe4f0 100644
--- a/minigrid/envs/obstructedmaze.py
+++ b/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 @@ def __init__(
         num_rows,
         num_cols,
         num_rooms_visited,
-        max_steps: Optional[int] = None,
+        max_steps: int | None = None,
         **kwargs,
     ):
         room_size = 6
diff --git a/minigrid/envs/playground.py b/minigrid/envs/playground.py
index a1b0d520c..0a0341cac 100644
--- a/minigrid/envs/playground.py
+++ b/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 @@ def __init__(self, max_steps=100, **kwargs):
             width=self.size,
             height=self.size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/putnear.py b/minigrid/envs/putnear.py
index 87dcc7338..f1a2fbc95 100644
--- a/minigrid/envs/putnear.py
+++ b/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"]
diff --git a/minigrid/envs/redbluedoors.py b/minigrid/envs/redbluedoors.py
index 22f4fe94b..ca4bc8bd8 100644
--- a/minigrid/envs/redbluedoors.py
+++ b/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)
 
diff --git a/minigrid/envs/unlock.py b/minigrid/envs/unlock.py
index 849469cdd..0fc53c1cf 100644
--- a/minigrid/envs/unlock.py
+++ b/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 @@ def __init__(self, max_steps: Optional[int] = None, **kwargs):
             num_cols=2,
             room_size=room_size,
             max_steps=max_steps,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
diff --git a/minigrid/envs/unlockpickup.py b/minigrid/envs/unlockpickup.py
index 5d90b18d5..33aa2fbfe 100644
--- a/minigrid/envs/unlockpickup.py
+++ b/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,
diff --git a/minigrid/manual_control.py b/minigrid/manual_control.py
index 7b3d84484..852952bcc 100755
--- a/minigrid/manual_control.py
+++ b/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
diff --git a/minigrid/minigrid_env.py b/minigrid/minigrid_env.py
index 887f15e3d..6014eb0d8 100755
--- a/minigrid/minigrid_env.py
+++ b/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 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 Actions(IntEnum):
     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 @@ def __init__(
             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 @@ def __init__(
         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 @@ def __str__(self):
     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 @@ def _rand_elem(self, iterable):
         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 @@ def _rand_subset(self, iterable, num_elems):
         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 @@ def _rand_subset(self, iterable, num_elems):
 
         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 @@ def place_obj(self, obj, top=None, size=None, reject_fn=None, max_tries=math.inf
 
             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 @@ def place_obj(self, obj, top=None, size=None, reject_fn=None, max_tries=math.inf
 
         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 @@ def dir_vec(self):
         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
diff --git a/minigrid/py.typed b/minigrid/py.typed
new file mode 100644
index 000000000..e69de29bb
diff --git a/minigrid/utils/rendering.py b/minigrid/utils/rendering.py
index f147745eb..0f299a08d 100644
--- a/minigrid/utils/rendering.py
+++ b/minigrid/utils/rendering.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 
 import numpy as np
diff --git a/minigrid/wrappers.py b/minigrid/wrappers.py
index 681ff207a..3f3f2205c 100644
--- a/minigrid/wrappers.py
+++ b/minigrid/wrappers.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 import operator
 from functools import reduce
diff --git a/pyproject.toml b/pyproject.toml
index 5ad2a1e07..72adb6350 100644
--- a/pyproject.toml
+++ b/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
\ No newline at end of file
diff --git a/setup.py b/setup.py
index b752e8110..f1b140614 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,7 @@
 """Setups up the Minigrid module."""
 
+from __future__ import annotations
+
 from setuptools import find_packages, setup
 
 
diff --git a/tests/test_envs.py b/tests/test_envs.py
index e2af1f7b7..35b736c2f 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pickle
 import warnings
 
diff --git a/tests/test_scripts.py b/tests/test_scripts.py
index 8c582d71e..d50e2ca6f 100644
--- a/tests/test_scripts.py
+++ b/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
diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py
index 229add089..cf2318184 100644
--- a/tests/test_wrappers.py
+++ b/tests/test_wrappers.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 
 import gymnasium as gym
diff --git a/tests/utils.py b/tests/utils.py
index a4503e688..db0bd0725 100644
--- a/tests/utils.py
+++ b/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