Repository: HumanCompatibleAI/seals (language: python)

Issue:
Migrate to gymnasium



Discussion / hints:

Fix patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 20ec919..b998576 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -8,7 +8,7 @@ orbs:
 
 defaults: &defaults
   docker:
-    - image: humancompatibleai/seals:base
+    - image: humancompatibleai/seals:base-alpha
       auth:
         username: $DOCKERHUB_USERNAME
         password: $DOCKERHUB_PASSWORD
@@ -51,7 +51,7 @@ commands:
       # released that you want to upgrade to, without mandating the newer version in setup.py.
       - restore_cache:
           keys:
-            - v2-dependencies-{{ checksum "setup.py" }}
+            - v2-dependencies-{{ checksum "setup.py" }}-{{ checksum "ci/build_venv.sh" }}
 
       # Create virtual environment and install dependencies using `ci/build_venv.sh`.
       # `mujoco_py` needs a MuJoCo key, so download that first.
@@ -64,7 +64,7 @@ commands:
       - save_cache:
           paths:
             - /venv
-          key: v2-dependencies-{{ checksum "setup.py" }}
+          key: v2-dependencies-{{ checksum "setup.py" }}-{{ checksum "ci/build_venv.sh" }}
 
       # Install seals.
       # Note we install the source distribution, not in developer mode (`pip install -e`).
diff --git a/Dockerfile b/Dockerfile
index a190ca8..01cb4aa 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,6 @@
 # base stage contains just binary dependencies.
 # This is used in the CI build.
-FROM nvidia/cuda:10.0-runtime-ubuntu18.04 AS base
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu20.04 AS base
 ARG DEBIAN_FRONTEND=noninteractive
 
 RUN    apt-get update -q \
@@ -9,6 +9,7 @@ RUN    apt-get update -q \
     curl \
     ffmpeg \
     git \
+    ssh \
     libgl1-mesa-dev \
     libgl1-mesa-glx \
     libglew-dev \
diff --git a/README.md b/README.md
index eacef50..8b02209 100644
--- a/README.md
+++ b/README.md
@@ -39,7 +39,7 @@ All *seals* environments are available in the Gym registry. Simply import it and
 would with your usual RL or specification learning algroithm:
 
 ```python
-import gym
+import gymnasium as gym
 import seals
 
 env = gym.make('seals/CartPole-v0')
diff --git a/ci/Xdummy-entrypoint.py b/ci/Xdummy-entrypoint.py
index f1abd5c..1a45388 100755
--- a/ci/Xdummy-entrypoint.py
+++ b/ci/Xdummy-entrypoint.py
@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/python3
 
 # Adapted from https://github.com/openai/mujoco-py/blob/master/vendor/Xdummy-entrypoint
 # Copyright OpenAI; MIT License
diff --git a/ci/build_venv.sh b/ci/build_venv.sh
index 672ecea..841ed47 100755
--- a/ci/build_venv.sh
+++ b/ci/build_venv.sh
@@ -9,4 +9,5 @@ fi
 
 virtualenv -p python3.8 ${venv}
 source ${venv}/bin/activate
+pip install --upgrade pip  # Ensure we have the newest pip
 pip install .[cpu,docs,mujoco,test]
diff --git a/pyproject.toml b/pyproject.toml
index f9bcbe8..d1ddaf0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,8 +6,5 @@ build-backend = "setuptools.build_meta"
 target-version = ["py38"]
 
 [[tool.mypy.overrides]]
-module = [
-    "gym.*",
-    "setuptools_scm.*",
-]
+module = ["gym.*", "setuptools_scm.*"]
 ignore_missing_imports = true
diff --git a/setup.py b/setup.py
index 84ba62a..f2cee15 100644
--- a/setup.py
+++ b/setup.py
@@ -88,14 +88,14 @@ def get_readme() -> str:
 
 ATARI_REQUIRE = [
     "opencv-python",
-    "ale-py==0.7.4",
+    "ale-py~=0.8.1",
     "pillow",
     "autorom[accept-rom-license]~=0.4.2",
+    "shimmy[atari] >=0.1.0,<1.0",
 ]
 TESTS_REQUIRE = [
-    # remove pin once https://github.com/nedbat/coveragepy/issues/881 fixed
     "black",
-    "coverage==4.5.4",
+    "coverage~=4.5.4",
     "codecov",
     "codespell",
     "darglint>=1.5.6",
@@ -115,11 +115,8 @@ def get_readme() -> str:
     "pytest-xdist",
     "pytype",
     "stable-baselines3>=0.9.0",
-    # TODO(adam): remove pyglet pin once Gym upgraded to >0.21
-    # Workaround for https://github.com/openai/gym/issues/2986
-    # Discussed in https://github.com/HumanCompatibleAI/imitation/pull/603
-    "pyglet==1.5.27",
     "setuptools_scm~=7.0.5",
+    "gymnasium[classic-control,mujoco]",
     *ATARI_REQUIRE,
 ]
 DOCS_REQUIRE = [
@@ -140,16 +137,14 @@ def get_readme() -> str:
     packages=find_packages("src"),
     package_dir={"": "src"},
     package_data={"seals": ["py.typed"]},
-    install_requires=["gym", "numpy"],
+    install_requires=["gymnasium", "numpy"],
     tests_require=TESTS_REQUIRE,
     extras_require={
         # recommended packages for development
         "dev": ["ipdb", "jupyter", *TESTS_REQUIRE, *DOCS_REQUIRE],
         "docs": DOCS_REQUIRE,
         "test": TESTS_REQUIRE,
-        # We'd like to specify `gym[mujoco]`, but this is a no-op when Gym is already
-        # installed. See https://github.com/pypa/pip/issues/4957 for issue.
-        "mujoco": ["mujoco_py>=1.50, <2.0", "imageio"],
+        "mujoco": ["gymnasium[mujoco]"],
         "atari": ATARI_REQUIRE,
     },
     url="https://github.com/HumanCompatibleAI/benchmark-environments",
diff --git a/src/seals/__init__.py b/src/seals/__init__.py
index 7f017e4..5dd0f50 100644
--- a/src/seals/__init__.py
+++ b/src/seals/__init__.py
@@ -2,7 +2,7 @@
 
 from importlib import metadata
 
-import gym
+import gymnasium as gym
 
 from seals import atari, util
 import seals.diagnostics  # noqa: F401
@@ -31,12 +31,12 @@
 
 for env_base in ["Ant", "HalfCheetah", "Hopper", "Humanoid", "Swimmer", "Walker2d"]:
     gym.register(
-        id=f"seals/{env_base}-v0",
+        id=f"seals/{env_base}-v1",
         entry_point=f"seals.mujoco:{env_base}Env",
-        max_episode_steps=util.get_gym_max_episode_steps(f"{env_base}-v3"),
+        max_episode_steps=util.get_gym_max_episode_steps(f"{env_base}-v4"),
     )
 
 # Atari
 
-GYM_ATARI_ENV_SPECS = list(filter(atari._supported_atari_env, gym.envs.registry.all()))
+GYM_ATARI_ENV_SPECS = list(filter(atari._supported_atari_env, gym.registry.values()))
 atari.register_atari_envs(GYM_ATARI_ENV_SPECS)
diff --git a/src/seals/atari.py b/src/seals/atari.py
index a2d896c..3782c84 100644
--- a/src/seals/atari.py
+++ b/src/seals/atari.py
@@ -2,7 +2,8 @@
 
 from typing import Dict, Iterable, Optional
 
-import gym
+import gymnasium as gym
+from gymnasium.envs.registration import EnvSpec
 
 from seals.util import (
     AutoResetWrapper,
@@ -35,9 +36,9 @@ def _get_score_region(atari_env_id: str) -> Optional[MaskedRegionSpecifier]:
     return SCORE_REGIONS.get(basename)
 
 
-def make_atari_env(atari_env_id: str, masked: bool) -> gym.Env:
+def make_atari_env(atari_env_id: str, masked: bool, *args, **kwargs) -> gym.Env:
     """Fixed-length, optionally masked-score variant of a given Atari environment."""
-    env = AutoResetWrapper(gym.make(atari_env_id))
+    env: gym.Env = AutoResetWrapper(gym.make(atari_env_id, *args, **kwargs))
 
     if masked:
         score_region = _get_score_region(atari_env_id)
@@ -59,15 +60,15 @@ def _not_ram_or_det(env_id: str) -> bool:
     after_slash = slash_separated[-1]
     hyphen_separated = after_slash.split("-")
     assert len(hyphen_separated) > 1
-    not_ram = not ("ram" in hyphen_separated[1])
-    not_deterministic = not ("Deterministic" in env_id)
+    not_ram = "ram" not in hyphen_separated[1]
+    not_deterministic = "Deterministic" not in env_id
     return not_ram and not_deterministic
 
 
-def _supported_atari_env(gym_spec: gym.envs.registration.EnvSpec) -> bool:
+def _supported_atari_env(gym_spec: EnvSpec) -> bool:
     """Checks if a gym Atari environment is one of the ones we will support."""
-    is_atari = gym_spec.entry_point == "gym.envs.atari:AtariEnv"
-    v5_and_plain = gym_spec.id.endswith("-v5") and not ("NoFrameskip" in gym_spec.id)
+    is_atari = gym_spec.entry_point == "shimmy.atari_env:AtariEnv"
+    v5_and_plain = gym_spec.id.endswith("-v5") and "NoFrameskip" not in gym_spec.id
     v4_and_no_frameskip = gym_spec.id.endswith("-v4") and "NoFrameskip" in gym_spec.id
     return (
         is_atari
@@ -76,7 +77,7 @@ def _supported_atari_env(gym_spec: gym.envs.registration.EnvSpec) -> bool:
     )
 
 
-def _seals_name(gym_spec: gym.envs.registration.EnvSpec, masked: bool) -> str:
+def _seals_name(gym_spec: EnvSpec, masked: bool) -> str:
     """Makes a Gym ID for an Atari environment in the seals namespace."""
     slash_separated = gym_spec.id.split("/")
     name = "seals/" + slash_separated[-1]
@@ -88,7 +89,7 @@ def _seals_name(gym_spec: gym.envs.registration.EnvSpec, masked: bool) -> str:
 
 
 def register_atari_envs(
-    gym_atari_env_specs: Iterable[gym.envs.registration.EnvSpec],
+    gym_atari_env_specs: Iterable[EnvSpec],
 ) -> None:
     """Register masked and unmasked wrapped gym Atari environments."""
 
diff --git a/src/seals/base_envs.py b/src/seals/base_envs.py
index c7b29c6..a0e42a2 100644
--- a/src/seals/base_envs.py
+++ b/src/seals/base_envs.py
@@ -1,20 +1,27 @@
 """Base environment classes."""
 
 import abc
-from typing import Generic, Optional, Sequence, Tuple, TypeVar
+from typing import Any, Dict, Generic, Optional, Tuple, TypeVar, Union
 
-import gym
-from gym import spaces
+import gymnasium as gym
+from gymnasium import spaces
 import numpy as np
+import numpy.typing as npt
 
 from seals import util
 
-State = TypeVar("State")
-Observation = TypeVar("Observation")
-Action = TypeVar("Action")
+# Note: we redefine the type vars from gymnasium.core here, because pytype does not
+# recognize them as valid type vars if we import them from gymnasium.core.
+StateType = TypeVar("StateType")
+ActType = TypeVar("ActType")
+ObsType = TypeVar("ObsType")
 
 
-class ResettablePOMDP(gym.Env, abc.ABC, Generic[State, Observation, Action]):
+class ResettablePOMDP(
+    gym.Env[ObsType, ActType],
+    abc.ABC,
+    Generic[StateType, ObsType, ActType],
+):
     """ABC for POMDPs that are resettable.
 
     Specifically, these environments provide oracle access to sample from
@@ -23,69 +30,36 @@ class ResettablePOMDP(gym.Env, abc.ABC, Generic[State, Observation, Action]):
     meet these criteria.
     """
 
-    _state_space: gym.Space
-    _observation_space: gym.Space
-    _action_space: gym.Space
-    _cur_state: Optional[State]
-    _n_actions_taken: Optional[int]
-
-    def __init__(
-        self,
-        *,
-        state_space: gym.Space,
-        observation_space: gym.Space,
-        action_space: gym.Space,
-    ):
-        """Build resettable (PO)MDP.
+    state_space: spaces.Space[StateType]
 
-        Args:
-            state_space: gym.Space containing possible states.
-            observation_space: gym.Space containing possible observations.
-            action_space: gym.Space containing possible actions.
-        """
-        self._state_space = state_space
-        self._observation_space = observation_space
-        self._action_space = action_space
+    _cur_state: Optional[StateType]
+    _n_actions_taken: Optional[int]
 
+    def __init__(self):
+        """Build resettable (PO)MDP."""
         self._cur_state = None
         self._n_actions_taken = None
-        self.seed()
 
     @abc.abstractmethod
-    def initial_state(self) -> State:
+    def initial_state(self) -> StateType:
         """Samples from the initial state distribution."""
 
     @abc.abstractmethod
-    def transition(self, state: State, action: Action) -> State:
+    def transition(self, state: StateType, action: ActType) -> StateType:
         """Samples from transition distribution."""
 
     @abc.abstractmethod
-    def reward(self, state: State, action: Action, new_state: State) -> float:
+    def reward(self, state: StateType, action: ActType, new_state: StateType) -> float:
         """Computes reward for a given transition."""
 
     @abc.abstractmethod
-    def terminal(self, state: State, step: int) -> bool:
+    def terminal(self, state: StateType, step: int) -> bool:
         """Is the state terminal?"""
 
     @abc.abstractmethod
-    def obs_from_state(self, state: State) -> Observation:
+    def obs_from_state(self, state: StateType) -> ObsType:
         """Sample observation for given state."""
 
-    @property
-    def state_space(self) -> gym.Space:
-        """State space. Often same as observation_space, but differs in POMDPs."""
-        return self._state_space
-
-    @property
-    def observation_space(self) -> gym.Space:
-        """Observation space. Return type of reset() and component of step()."""
-        return self._observation_space
-
-    @property
-    def action_space(self) -> gym.Space:
-        """Action space. Parameter type of step()."""
-        return self._action_space
-
     @property
     def n_actions_taken(self) -> int:
         """Number of steps taken so far."""
@@ -93,37 +67,39 @@ def n_actions_taken(self) -> int:
         return self._n_actions_taken
 
     @property
-    def state(self) -> State:
+    def state(self) -> StateType:
         """Current state."""
         assert self._cur_state is not None
         return self._cur_state
 
     @state.setter
-    def state(self, state: State):
-        """Set current state."""
+    def state(self, state: StateType):
+        """Set the current state."""
         if state not in self.state_space:
             raise ValueError(f"{state} not in {self.state_space}")
         self._cur_state = state
 
-    def seed(self, seed=None) -> Sequence[int]:
-        """Set random seed."""
-        if seed is None:
-            # Gym API wants list of seeds to be returned for some reason, so
-            # generate a seed explicitly in this case
-            seed = np.random.randint(0, 1 << 31)
-        self.rand_state = np.random.RandomState(seed)
-        return [seed]
-
-    def reset(self) -> Observation:
-        """Reset episode and return initial observation."""
+    def reset(
+        self,
+        *,
+        seed: Optional[int] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[ObsType, Dict[str, Any]]:
+        """Reset the episode and return initial observation."""
+        if options is not None:
+            raise NotImplementedError("Options not supported.")
+
+        super().reset(seed=seed)
         self.state = self.initial_state()
         self._n_actions_taken = 0
-        return self.obs_from_state(self.state)
+        obs = self.obs_from_state(self.state)
+        info: Dict[str, Any] = dict()
+        return obs, info
 
-    def step(self, action: Action) -> Tuple[Observation, float, bool, dict]:
+    def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, dict]:
         """Transition state using given action."""
         if self._cur_state is None or self._n_actions_taken is None:
-            raise ValueError("Need to call reset() before first step()")
+            raise RuntimeError("Need to call reset() before first step()")
         if action not in self.action_space:
             raise ValueError(f"{action} not in {self.action_space}")
 
@@ -133,16 +109,22 @@ def step(self, action: Action) -> Tuple[Observation, float, bool, dict]:
         assert obs in self.observation_space
         reward = self.reward(old_state, action, self.state)
         self._n_actions_taken += 1
-        done = self.terminal(self.state, self.n_actions_taken)
+        terminated = self.terminal(self.state, self.n_actions_taken)
+        truncated = False
 
         infos = {"old_state": old_state, "new_state": self._cur_state}
-        return obs, reward, done, infos
+        return obs, reward, terminated, truncated, infos
 
 
-class ExposePOMDPStateWrapper(gym.Wrapper, Generic[State, Observation, Action]):
+class ExposePOMDPStateWrapper(
+    gym.Wrapper[StateType, ActType, ObsType, ActType],
+    Generic[StateType, ObsType, ActType],
+):
     """A wrapper that exposes the current state of the POMDP as the observation."""
 
-    def __init__(self, env: ResettablePOMDP[State, Observation, Action]) -> None:
+    env: ResettablePOMDP[StateType, ObsType, ActType]
+
+    def __init__(self, env: ResettablePOMDP[StateType, ObsType, ActType]) -> None:
         """Build wrapper.
 
         Args:
@@ -151,51 +133,48 @@ def __init__(self, env: ResettablePOMDP[State, Observation, Action]) -> None:
         super().__init__(env)
         self._observation_space = env.state_space
 
-    def reset(self) -> State:
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[StateType, Dict[str, Any]]:
         """Reset environment and return initial state."""
-        self.env.reset()
-        return self.env.state
+        _, info = self.env.reset(seed=seed, options=options)
+        return self.env.state, info
 
-    def step(self, action) -> Tuple[State, float, bool, dict]:
+    def step(self, action) -> Tuple[StateType, float, bool, bool, dict]:
         """Transition state using given action."""
-        obs, reward, done, info = self.env.step(action)
-        return self.env.state, reward, done, info
+        _, reward, terminated, truncated, info = self.env.step(action)
+        return self.env.state, reward, terminated, truncated, info
 
 
 class ResettableMDP(
-    ResettablePOMDP[State, State, Action],
+    ResettablePOMDP[StateType, StateType, ActType],
     abc.ABC,
-    Generic[State, Action],
+    Generic[StateType, ActType],
 ):
     """ABC for MDPs that are resettable."""
 
-    def __init__(
-        self,
-        *,
-        state_space: gym.Space,
-        action_space: gym.Space,
-    ):
-        """Build resettable MDP.
-
-        Args:
-            state_space: gym.Space containing possible states.
-            action_space: gym.Space containing possible actions.
-        """
-        super().__init__(
-            state_space=state_space,
-            observation_space=state_space,
-            action_space=action_space,
-        )
+    @property
+    def observation_space(self):
+        """Observation space."""
+        return self.state_space
 
-    def obs_from_state(self, state: State) -> State:
+    def obs_from_state(self, state: StateType) -> StateType:
         """Identity since observation == state in an MDP."""
         return state
 
 
+DiscreteSpaceInt = np.int64
+
+
 # TODO(juan) this does not implement the .render() method,
 #  so in theory it should not be instantiated directly.
 #  Not sure why this is not raising an error?
-class BaseTabularModelPOMDP(ResettablePOMDP[int, Observation, int]):
+class BaseTabularModelPOMDP(
+    ResettablePOMDP[DiscreteSpaceInt, ObsType, DiscreteSpaceInt],
+    Generic[ObsType],
+):
     """Base class for tabular environments with known dynamics.
 
     This is the general class that also allows subclassing for creating
@@ -236,6 +215,8 @@ def __init__(
             ValueError: `transition_matrix`, `reward_matrix` or
                 `initial_state_dist` have shapes different to specified above.
         """
+        super().__init__()
+
         # The following matrices should conform to the shapes below:
 
         # transition matrix: n_states x n_actions x n_states
@@ -278,43 +259,44 @@ def __init__(
         self.horizon = horizon
         self.initial_state_dist = initial_state_dist
 
-        super().__init__(
-            state_space=self._construct_state_space(),
-            action_space=self._construct_action_space(),
-            observation_space=self._construct_observation_space(),
-        )
-
-    def _construct_state_space(self) -> gym.Space:
-        return spaces.Discrete(self.state_dim)
-
-    def _construct_action_space(self) -> gym.Space:
-        return spaces.Discrete(self.action_dim)
-
-    @abc.abstractmethod
-    def _construct_observation_space(self) -> gym.Space:
-        pass  # pragma: no cover
+        self.state_space = spaces.Discrete(self.state_dim)
+        self.action_space = spaces.Discrete(self.action_dim)
 
-    def initial_state(self) -> int:
+    def initial_state(self) -> DiscreteSpaceInt:
         """Samples from the initial state distribution."""
-        return util.sample_distribution(
-            self.initial_state_dist,
-            random=self.rand_state,
+        return DiscreteSpaceInt(
+            util.sample_distribution(
+                self.initial_state_dist,
+                random=self.np_random,
+            ),
         )
 
-    def transition(self, state: int, action: int) -> int:
+    def transition(
+        self,
+        state: DiscreteSpaceInt,
+        action: DiscreteSpaceInt,
+    ) -> DiscreteSpaceInt:
         """Samples from transition distribution."""
-        return util.sample_distribution(
-            self.transition_matrix[state, action],
-            random=self.rand_state,
+        return DiscreteSpaceInt(
+            util.sample_distribution(
+                self.transition_matrix[state, action],
+                random=self.np_random,
+            ),
         )
 
-    def reward(self, state: int, action: int, new_state: int) -> float:
+    def reward(
+        self,
+        state: DiscreteSpaceInt,
+        action: DiscreteSpaceInt,
+        new_state: DiscreteSpaceInt,
+    ) -> float:
         """Computes reward for a given transition."""
         inputs = (state, action, new_state)[: len(self.reward_matrix.shape)]
         return self.reward_matrix[inputs]
 
-    def terminal(self, state: int, n_actions_taken: int) -> bool:
+    def terminal(self, state: DiscreteSpaceInt, n_actions_taken: int) -> bool:
         """Checks if state is terminal."""
+        del state
         return self.horizon is not None and n_actions_taken >= self.horizon
 
     @property
@@ -337,7 +319,13 @@ def action_dim(self) -> int:
         return self.transition_matrix.shape[1]
 
 
-class TabularModelPOMDP(BaseTabularModelPOMDP[np.ndarray]):
+ObsEntryType = TypeVar(
+    "ObsEntryType",
+    bound=Union[np.floating, np.integer],
+)
+
+
+class TabularModelPOMDP(BaseTabularModelPOMDP[np.ndarray], Generic[ObsEntryType]):
     """Tabular model POMDP.
 
     This class is specifically for environments where observation != state,
@@ -349,13 +337,14 @@ class TabularModelPOMDP(BaseTabularModelPOMDP[np.ndarray]):
     a vector with self.obs_dim entries.
     """
 
-    observation_matrix: np.ndarray
+    observation_matrix: npt.NDArray[ObsEntryType]
+    observation_space: spaces.Box
 
     def __init__(
         self,
         *,
         transition_matrix: np.ndarray,
-        observation_matrix: np.ndarray,
+        observation_matrix: npt.NDArray[ObsEntryType],
         reward_matrix: np.ndarray,
         horizon: Optional[int] = None,
         initial_state_dist: Optional[np.ndarray] = None,
@@ -377,7 +366,6 @@ def __init__(
                 f"observation_matrix.shape[0]: {observation_matrix.shape[0]}",
             )
 
-    def _construct_observation_space(self) -> gym.Space:
         min_val: float
         max_val: float
         try:
@@ -386,14 +374,14 @@ def _construct_observation_space(self) -> gym.Space:
         except ValueError:
             min_val = -np.inf
             max_val = np.inf
-        return spaces.Box(
+        self.observation_space = spaces.Box(
             low=min_val,
             high=max_val,
             shape=(self.obs_dim,),
-            dtype=self.obs_dtype,
+            dtype=self.obs_dtype,  # type: ignore
         )
 
-    def obs_from_state(self, state: int) -> np.ndarray:
+    def obs_from_state(self, state: DiscreteSpaceInt) -> npt.NDArray[ObsEntryType]:
         """Computes observation from state."""
         # Copy so it can't be mutated in-place (updates will be reflected in
         # self.observation_matrix!)
@@ -407,12 +395,12 @@ def obs_dim(self) -> int:
         return self.observation_matrix.shape[1]
 
     @property
-    def obs_dtype(self) -> int:
+    def obs_dtype(self) -> np.dtype:
         """Data type of observation vectors (e.g. np.float32)."""
         return self.observation_matrix.dtype
 
 
-class TabularModelMDP(BaseTabularModelPOMDP[int]):
+class TabularModelMDP(BaseTabularModelPOMDP[DiscreteSpaceInt]):
     """Tabular model MDP.
 
     A tabular model MDP is a tabular MDP where the transition and reward
@@ -443,10 +431,8 @@ def __init__(
             horizon=horizon,
             initial_state_dist=initial_state_dist,
         )
+        self.observation_space = self.state_space
 
-    def obs_from_state(self, state: int) -> int:
+    def obs_from_state(self, state: DiscreteSpaceInt) -> DiscreteSpaceInt:
         """Identity since observation == state in an MDP."""
         return state
-
-    def _construct_observation_space(self) -> gym.Space:
-        return self._construct_state_space()
diff --git a/src/seals/classic_control.py b/src/seals/classic_control.py
index 9854095..bf856f9 100644
--- a/src/seals/classic_control.py
+++ b/src/seals/classic_control.py
@@ -1,27 +1,30 @@
 """Adaptation of classic Gym environments for specification learning algorithms."""
 
+from typing import Any, Dict, Optional
 import warnings
 
-from gym import spaces
-import gym.envs.classic_control
+import gymnasium as gym
+from gymnasium import spaces
+from gymnasium.envs import classic_control
 import numpy as np
 
 from seals import util
 
 
-class FixedHorizonCartPole(gym.envs.classic_control.CartPoleEnv):
+class FixedHorizonCartPole(classic_control.CartPoleEnv):
     """Fixed-length variant of CartPole-v1.
 
-    Reward is 1.0 whenever the CartPole is an "ok" state (i.e. the pole is upright
-    and the cart is on the screen). Otherwise reward is 0.0.
+    Reward is 1.0 whenever the CartPole is an "ok" state (i.e., the pole is upright
+    and the cart is on the screen). Otherwise, reward is 0.0.
 
-    Done is always False. (Though note that by default, this environment is wrapped
-    in `TimeLimit` with max steps 500.)
+    Terminated is always False.
+    By default, this environment is wrapped in 'TimeLimit' with max steps 500,
+    which sets `truncated` to true after that many steps.
     """
 
-    def __init__(self):
-        """Builds FixedHorizonCartPole, modifying observation_space from Gym parent."""
-        super().__init__()
+    def __init__(self, *args, **kwargs):
+        """Builds FixedHorizonCartPole, modifying observation_space from gym parent."""
+        super().__init__(*args, **kwargs)
 
         high = [
             np.finfo(np.float32).max,  # x axis
@@ -32,14 +35,20 @@ def __init__(self):
         high = np.array(high)
         self.observation_space = spaces.Box(-high, high, dtype=np.float32)
 
-    def reset(self):
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ):
         """Reset for FixedHorizonCartPole."""
-        return super().reset().astype(np.float32)
+        observation, info = super().reset(seed=seed, options=options)
+        return observation.astype(np.float32), info
 
     def step(self, action):
         """Step function for FixedHorizonCartPole."""
         with warnings.catch_warnings():
-            # Filter out CartPoleEnv warning for calling step() beyond done=True.
+            # Filter out CartPoleEnv warning for calling step() beyond
+            # terminated=True or truncated=True
             warnings.filterwarnings("ignore", ".*You are calling.*")
             super().step(action)
 
@@ -55,10 +64,10 @@ def step(self, action):
         )
 
         rew = 1.0 if state_ok else 0.0
-        return np.array(self.state, dtype=np.float32), rew, False, {}
+        return np.array(self.state, dtype=np.float32), rew, False, False, {}
 
 
-def mountain_car():
+def mountain_car(*args, **kwargs):
     """Fixed-length variant of MountainCar-v0.
 
     In the event of early episode completion (i.e., the car reaches the
@@ -67,7 +76,7 @@ def mountain_car():
 
     Done is always returned on timestep 200 only.
     """
-    env = util.make_env_no_wrappers("MountainCar-v0")
+    env = gym.make("MountainCar-v0", *args, **kwargs)
     env = util.ObsCastWrapper(env, dtype=np.float32)
     env = util.AbsorbAfterDoneWrapper(env)
     return env
diff --git a/src/seals/diagnostics/__init__.py b/src/seals/diagnostics/__init__.py
index 74be80f..b3c9b28 100644
--- a/src/seals/diagnostics/__init__.py
+++ b/src/seals/diagnostics/__init__.py
@@ -1,6 +1,6 @@
 """Simple diagnostic environments."""
 
-import gym
+import gymnasium as gym
 
 gym.register(
     id="seals/Branching-v0",
diff --git a/src/seals/diagnostics/early_term.py b/src/seals/diagnostics/early_term.py
index 8c672bb..87287bd 100644
--- a/src/seals/diagnostics/early_term.py
+++ b/src/seals/diagnostics/early_term.py
@@ -53,7 +53,7 @@ def __init__(self, is_reward_positive: bool = True):
             reward_matrix=reward_matrix,
         )
 
-    def terminal(self, state: int, n_actions_taken: int) -> bool:
+    def terminal(self, state: base_envs.DiscreteSpaceInt, n_actions_taken: int) -> bool:
         """Returns True if (and only if) in state 2."""
         return bool(state == 2)
 
diff --git a/src/seals/diagnostics/init_shift.py b/src/seals/diagnostics/init_shift.py
index 086f561..1e47477 100644
--- a/src/seals/diagnostics/init_shift.py
+++ b/src/seals/diagnostics/init_shift.py
@@ -27,7 +27,7 @@ class InitShiftEnv(base_envs.TabularModelMDP):
     disambiguate this case.
     """
 
-    def __init__(self, initial_state: int):
+    def __init__(self, initial_state: base_envs.DiscreteSpaceInt):
         """Constructs environment.
 
         Args:
@@ -63,7 +63,7 @@ def __init__(self, initial_state: int):
             reward_matrix=reward_matrix,
         )
 
-    def initial_state(self) -> int:
+    def initial_state(self) -> base_envs.DiscreteSpaceInt:
         """Returns initial state defined in constructor."""
         return self._initial_state
 
diff --git a/src/seals/diagnostics/largest_sum.py b/src/seals/diagnostics/largest_sum.py
index d55021f..7a08dad 100644
--- a/src/seals/diagnostics/largest_sum.py
+++ b/src/seals/diagnostics/largest_sum.py
@@ -1,6 +1,6 @@
 """Environment testing scalability to high-dimensional tasks."""
 
-from gym import spaces
+from gymnasium import spaces
 import numpy as np
 
 from seals import base_envs
@@ -23,12 +23,10 @@ def __init__(self, length: int = 50):
         Args:
             length: dimensionality of state space vector.
         """
+        super().__init__()
         self._length = length
-        state_space = spaces.Box(low=0.0, high=1.0, shape=(length,))
-        super().__init__(
-            state_space=state_space,
-            action_space=spaces.Discrete(2),
-        )
+        self.state_space = spaces.Box(low=0.0, high=1.0, shape=(length,))
+        self.action_space = spaces.Discrete(2)
 
     def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
         """Always returns True, since this task should have a 1-timestep horizon."""
@@ -36,7 +34,7 @@ def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
 
     def initial_state(self) -> np.ndarray:
         """Returns vector sampled uniformly in [0, 1]**L."""
-        init_state = self.rand_state.rand(self._length)
+        init_state = self.np_random.random((self._length,))
         return init_state.astype(self.observation_space.dtype)
 
     def reward(self, state: np.ndarray, act: int, next_state: np.ndarray) -> float:
diff --git a/src/seals/diagnostics/noisy_obs.py b/src/seals/diagnostics/noisy_obs.py
index 8588416..92b3429 100644
--- a/src/seals/diagnostics/noisy_obs.py
+++ b/src/seals/diagnostics/noisy_obs.py
@@ -1,6 +1,6 @@
 """Environment testing for robustness to noise."""
 
-from gym import spaces
+from gymnasium import spaces
 import numpy as np
 
 from seals import base_envs, util
@@ -23,6 +23,8 @@ def __init__(self, *, size: int = 5, noise_length: int = 20):
             size: width and height of gridworld.
             noise_length: dimension of noise vector in observation.
         """
+        super().__init__()
+
         self._size = size
         self._noise_length = noise_length
         self._goal = np.array([self._size // 2, self._size // 2])
@@ -34,14 +36,12 @@ def __init__(self, *, size: int = 5, noise_length: int = 20):
             ([size - 1, size - 1], np.full(self._noise_length, np.inf)),  # type: ignore
         )
 
-        super().__init__(
-            state_space=spaces.MultiDiscrete([size, size]),
-            action_space=spaces.Discrete(5),
-            observation_space=spaces.Box(
-                low=obs_box_low,
-                high=obs_box_high,
-                dtype=np.float32,
-            ),
+        self.state_space = spaces.MultiDiscrete([size, size])
+        self.action_space = spaces.Discrete(5)
+        self.observation_space = spaces.Box(
+            low=obs_box_low,
+            high=obs_box_high,
+            dtype=np.float32,
         )
 
     def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
@@ -52,7 +52,7 @@ def initial_state(self) -> np.ndarray:
         """Returns one of the grid's corners."""
         n = self._size
         corners = np.array([[0, 0], [n - 1, 0], [0, n - 1], [n - 1, n - 1]])
-        return corners[self.rand_state.randint(4)]
+        return corners[self.np_random.integers(4)]
 
     def reward(self, state: np.ndarray, action: int, new_state: np.ndarray) -> float:
         """Returns  +1.0 reward if state is the goal and 0.0 otherwise."""
@@ -69,5 +69,5 @@ def transition(self, state: np.ndarray, action: int) -> np.ndarray:
 
     def obs_from_state(self, state: np.ndarray) -> np.ndarray:
         """Returns (x, y) concatenated with Gaussian noise."""
-        noise_vector = self.rand_state.randn(self._noise_length)
+        noise_vector = self.np_random.normal(size=self._noise_length)
         return np.concatenate([state, noise_vector]).astype(np.float32)
diff --git a/src/seals/diagnostics/parabola.py b/src/seals/diagnostics/parabola.py
index 07c9b8b..7b0c677 100644
--- a/src/seals/diagnostics/parabola.py
+++ b/src/seals/diagnostics/parabola.py
@@ -1,6 +1,6 @@
 """Environment testing for generalization in continuous spaces."""
 
-from gym import spaces
+from gymnasium import spaces
 import numpy as np
 
 from seals import base_envs
@@ -24,16 +24,15 @@ def __init__(self, x_step: float = 0.05, bounds: float = 5):
             bounds: limits coordinates, useful for keeping rewards in
                 a small bounded range.
         """
+        super().__init__()
         self._x_step = x_step
         self._bounds = bounds
 
         state_high = np.array([bounds, bounds, 1.0, 1.0, 1.0])
         state_low = (-1) * state_high
 
-        super().__init__(
-            state_space=spaces.Box(low=state_low, high=state_high),
-            action_space=spaces.Box(low=(-2) * bounds, high=2 * bounds, shape=()),
-        )
+        self.state_space = spaces.Box(low=state_low, high=state_high)
+        self.action_space = spaces.Box(low=(-2) * bounds, high=2 * bounds, shape=())
 
     def terminal(self, state: int, n_actions_taken: int) -> bool:
         """Always returns False."""
@@ -41,7 +40,7 @@ def terminal(self, state: int, n_actions_taken: int) -> bool:
 
     def initial_state(self) -> np.ndarray:
         """Get state by sampling a random parabola."""
-        a, b, c = -1 + 2 * self.rand_state.rand(3)
+        a, b, c = -1 + 2 * self.np_random.random((3,))
         x, y = 0, c
         return np.array([x, y, a, b, c], dtype=self.state_space.dtype)
 
diff --git a/src/seals/diagnostics/proc_goal.py b/src/seals/diagnostics/proc_goal.py
index e534b72..5136db8 100644
--- a/src/seals/diagnostics/proc_goal.py
+++ b/src/seals/diagnostics/proc_goal.py
@@ -1,6 +1,6 @@
 """Large gridworld with random agent and goal position."""
 
-from gym import spaces
+from gymnasium import spaces
 import numpy as np
 
 from seals import base_envs, util
@@ -28,13 +28,11 @@ def __init__(self, bounds: int = 100, distance: int = 10):
                 generalization harder.
             distance: initial distance between agent and goal.
         """
+        super().__init__()
         self._bounds = bounds
         self._distance = distance
-
-        super().__init__(
-            state_space=spaces.Box(low=-np.inf, high=np.inf, shape=(4,)),
-            action_space=spaces.Discrete(5),
-        )
+        self.state_space = spaces.Box(low=-np.inf, high=np.inf, shape=(4,))
+        self.action_space = spaces.Discrete(5)
 
     def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
         """Always returns False."""
@@ -42,11 +40,11 @@ def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
 
     def initial_state(self) -> np.ndarray:
         """Samples random agent position and random goal."""
-        pos = self.rand_state.randint(low=-self._bounds, high=self._bounds, size=(2,))
+        pos = self.np_random.integers(low=-self._bounds, high=self._bounds, size=(2,))
 
-        x_dist = self.rand_state.randint(self._distance)
+        x_dist = self.np_random.integers(self._distance)
         y_dist = self._distance - x_dist
-        random_signs = 2 * self.rand_state.randint(2, size=2) - 1
+        random_signs = 2 * self.np_random.integers(2, size=2) - 1
         goal = pos + random_signs * (x_dist, y_dist)
 
         return np.concatenate([pos, goal]).astype(self.observation_space.dtype)
diff --git a/src/seals/diagnostics/random_trans.py b/src/seals/diagnostics/random_trans.py
index 09eb6b4..6071e88 100644
--- a/src/seals/diagnostics/random_trans.py
+++ b/src/seals/diagnostics/random_trans.py
@@ -46,7 +46,7 @@ def __init__(
         """
         # this generator is ONLY for constructing the MDP, not for controlling
         # random outcomes during rollouts
-        rand_gen = np.random.RandomState(generator_seed)
+        rand_gen = np.random.default_rng(generator_seed)
 
         if random_obs:
             if obs_dim is None:
@@ -72,7 +72,8 @@ def __init__(
             n_states=n_states,
             rand_state=rand_gen,
         )
-        self.reward_weights = rand_gen.randn(observation_matrix.shape[-1])
+
+        self.reward_weights = rand_gen.normal(size=(observation_matrix.shape[-1],))
         reward_matrix = observation_matrix @ self.reward_weights
         super().__init__(
             transition_matrix=transition_matrix,
@@ -87,7 +88,7 @@ def make_random_trans_mat(
         n_states,
         n_actions,
         max_branch_factor,
-        rand_state: Optional[np.random.RandomState] = None,
+        rand_state: Optional[np.random.Generator] = None,
     ) -> np.ndarray:
         """Make a 'random' transition matrix.
 
@@ -110,14 +111,14 @@ def make_random_trans_mat(
             of transitioning to `next_s` after taking action `a` in state `s`.
         """
         if rand_state is None:
-            rand_state = np.random.RandomState()
+            rand_state = np.random.default_rng()
         assert rand_state is not None
         out_mat = np.zeros((n_states, n_actions, n_states), dtype="float32")
         for start_state in range(n_states):
             for action in range(n_actions):
                 # uniformly sample a number of successors in [1,max_branch_factor]
                 # for this action
-                successors = rand_state.randint(1, max_branch_factor + 1)
+                successors = rand_state.integers(1, max_branch_factor + 1)
                 next_states = rand_state.choice(
                     n_states,
                     size=(successors,),
@@ -133,7 +134,7 @@ def make_random_trans_mat(
     def make_random_state_dist(
         n_avail: int,
         n_states: int,
-        rand_state: Optional[np.random.RandomState] = None,
+        rand_state: Optional[np.random.Generator] = None,
     ) -> np.ndarray:
         """Make a random initial state distribution over n_states.
 
@@ -152,7 +153,7 @@ def make_random_state_dist(
             ValueError: If `n_avail` is not in the range `(0, n_states]`.
         """  # noqa: DAR402
         if rand_state is None:
-            rand_state = np.random.RandomState()
+            rand_state = np.random.default_rng()
         assert rand_state is not None
         assert 0 < n_avail <= n_states
         init_dist = np.zeros((n_states,))
@@ -168,7 +169,7 @@ def make_obs_mat(
         n_states: int,
         is_random: bool,
         obs_dim: Optional[int] = None,
-        rand_state: Optional[np.random.RandomState] = None,
+        rand_state: Optional[np.random.Generator] = None,
     ) -> np.ndarray:
         """Makes an observation matrix with a single observation for each state.
 
@@ -179,7 +180,7 @@ def make_obs_mat(
                         If `False`, are unique one-hot vectors for each state.
             obs_dim (int or NoneType): Must be `None` if `is_random == False`.
                     Otherwise, this must be set to the size of the random vectors.
-            rand_state (np.random.RandomState): Random number generator.
+            rand_state (np.random.Generator): Random number generator.
 
         Returns:
             A matrix of shape `(n_states, obs_dim if is_random else n_states)`.
@@ -188,7 +189,7 @@ def make_obs_mat(
             ValueError: If ``is_random == False`` and ``obs_dim is not None``.
         """
         if rand_state is None:
-            rand_state = np.random.RandomState()
+            rand_state = np.random.default_rng()
         assert rand_state is not None
         if is_random:
             if obs_dim is None:
diff --git a/src/seals/diagnostics/sort.py b/src/seals/diagnostics/sort.py
index d1c4313..6f564fc 100644
--- a/src/seals/diagnostics/sort.py
+++ b/src/seals/diagnostics/sort.py
@@ -1,6 +1,6 @@
 """Environment to sort a list using swap actions."""
 
-from gym import spaces
+from gymnasium import spaces
 import numpy as np
 
 from seals import base_envs
@@ -21,10 +21,9 @@ def __init__(self, length: int = 4):
         """
         self._length = length
 
-        super().__init__(
-            state_space=spaces.Box(low=0, high=1.0, shape=(length,)),
-            action_space=spaces.MultiDiscrete([length, length]),
-        )
+        super().__init__()
+        self.state_space = spaces.Box(low=0, high=1.0, shape=(length,))
+        self.action_space = spaces.MultiDiscrete([length, length])
 
     def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
         """Always returns False."""
@@ -32,7 +31,7 @@ def terminal(self, state: np.ndarray, n_actions_taken: int) -> bool:
 
     def initial_state(self):
         """Sample random vector uniformly in [0, 1]**L."""
-        sample = self.rand_state.random(size=self._length)
+        sample = self.np_random.random(size=self._length)
         return sample.astype(self.state_space.dtype)
 
     def reward(
@@ -42,6 +41,7 @@ def reward(
         new_state: np.ndarray,
     ) -> float:
         """Rewards fully sorted lists, and new correct positions."""
+        del action
         # This is not meant to be a potential shaping in the formal sense,
         # as it changes the trajectory returns (since we do not return
         # a fixed-potential state at termination).
diff --git a/src/seals/mujoco.py b/src/seals/mujoco.py
index 5cd2c5f..3ffb5f3 100644
--- a/src/seals/mujoco.py
+++ b/src/seals/mujoco.py
@@ -2,13 +2,13 @@
 
 import functools
 
-from gym.envs.mujoco import (
-    ant_v3,
-    half_cheetah_v3,
-    hopper_v3,
-    humanoid_v3,
-    swimmer_v3,
-    walker2d_v3,
+from gymnasium.envs.mujoco import (
+    ant_v4,
+    half_cheetah_v4,
+    hopper_v4,
+    humanoid_v4,
+    swimmer_v4,
+    walker2d_v4,
 )
 
 
@@ -27,33 +27,33 @@ def _no_early_termination(cls):
 
 @_include_position_in_observation
 @_no_early_termination
-class AntEnv(ant_v3.AntEnv):
+class AntEnv(ant_v4.AntEnv):
     """Ant with position observation and no early termination."""
 
 
 @_include_position_in_observation
-class HalfCheetahEnv(half_cheetah_v3.HalfCheetahEnv):
+class HalfCheetahEnv(half_cheetah_v4.HalfCheetahEnv):
     """HalfCheetah with position observation. Naturally does not terminate early."""
 
 
 @_include_position_in_observation
 @_no_early_termination
-class HopperEnv(hopper_v3.HopperEnv):
+class HopperEnv(hopper_v4.HopperEnv):
     """Hopper with position observation and no early termination."""
 
 
 @_include_position_in_observation
 @_no_early_termination
-class HumanoidEnv(humanoid_v3.HumanoidEnv):
+class HumanoidEnv(humanoid_v4.HumanoidEnv):
     """Humanoid with position observation and no early termination."""
 
 
 @_include_position_in_observation
-class SwimmerEnv(swimmer_v3.SwimmerEnv):
+class SwimmerEnv(swimmer_v4.SwimmerEnv):
     """Swimmer with position observation. Naturally does not terminate early."""
 
 
 @_include_position_in_observation
 @_no_early_termination
-class Walker2dEnv(walker2d_v3.Walker2dEnv):
+class Walker2dEnv(walker2d_v4.Walker2dEnv):
     """Walker2d with position observation and no early termination."""
diff --git a/src/seals/util.py b/src/seals/util.py
index 60c0323..6299bdc 100644
--- a/src/seals/util.py
+++ b/src/seals/util.py
@@ -1,14 +1,36 @@
 """Miscellaneous utilities."""
 
 from dataclasses import dataclass
-from typing import List, Optional, Sequence, Tuple, Union
-
-import gym
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Sequence,
+    SupportsFloat,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import gymnasium as gym
 import numpy as np
+import numpy.typing as npt
+
+# Note: we redefine the type vars from gymnasium.core here, because pytype does not
+# recognize them as valid type vars if we import them from gymnasium.core.
+WrapperObsType = TypeVar("WrapperObsType")
+WrapperActType = TypeVar("WrapperActType")
+ObsType = TypeVar("ObsType")
+ActType = TypeVar("ActType")
 
 
-class AutoResetWrapper(gym.Wrapper):
-    """Hides done=True and auto-resets at the end of each episode.
+class AutoResetWrapper(
+    gym.Wrapper,
+    Generic[WrapperObsType, WrapperActType, ObsType, ActType],
+):
+    """Hides terminated truncated and auto-resets at the end of each episode.
 
     Depending on the flag 'discard_terminal_observation', either discards the terminal
     observation or pads with an additional 'reset transition'. The former is the default
@@ -37,8 +59,13 @@ def __init__(self, env, discard_terminal_observation=True, reset_reward=0.0):
         self.reset_reward = reset_reward
         self.previous_done = False  # Whether the previous step returned done=True.
 
-    def step(self, action):
-        """When done=True, returns done=False, then reset depending on flag.
+    def step(
+        self,
+        action: WrapperActType,
+    ) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
+        """When terminated or truncated, resets the environment.
+
+        Always returns False for terminated and truncated.
 
         Depending on whether we are discarding the terminal observation,
         either resets the environment and discards,
@@ -50,8 +77,13 @@ def step(self, action):
         else:
             return self._step_pad(action)
 
-    def _step_pad(self, action):
-        """When done=True, return done=False instead and return the terminal obs.
+    def _step_pad(
+        self,
+        action: WrapperActType,
+    ) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
+        """When terminated or truncated, resets the environment.
+
+        Always returns False for terminated and truncated.
 
         The agent will then usually be asked to perform an action based on
         the terminal observation. In the next step, this final action will be ignored
@@ -67,26 +99,31 @@ def _step_pad(self, action):
         """
         if self.previous_done:
             self.previous_done = False
+            reset_obs, reset_info_dict = self.env.reset()
+            info = {"reset_info_dict": reset_info_dict}
             # This transition will only reset the environment, the action is ignored.
-            return self.env.reset(), self.reset_reward, False, {}
+            return reset_obs, self.reset_reward, False, False, info
 
-        obs, rew, done, info = self.env.step(action)
-        if done:
-            self.previous_done = True
-        return obs, rew, False, info
+        obs, rew, terminated, truncated, info = self.env.step(action)
+        self.previous_done = terminated or truncated
+        return obs, rew, False, False, info
 
-    def _step_discard(self, action):
-        """When done=True, returns done=False instead and automatically resets.
+    def _step_discard(
+        self,
+        action: WrapperActType,
+    ) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
+        """When terminated or truncated, return False for both and automatically reset.
 
         When an automatic reset happens, the observation from reset is returned,
         and the overridden observation is stored in
         `info["terminal_observation"]`.
         """
-        obs, rew, done, info = self.env.step(action)
-        if done:
+        obs, rew, terminated, truncated, info = self.env.step(action)
+        if terminated or truncated:
             info["terminal_observation"] = obs
-            obs = self.env.reset()
-        return obs, rew, False, info
+            obs, reset_info_dict = self.env.reset()
+            info["reset_info_dict"] = reset_info_dict
+        return obs, rew, False, False, info
 
 
 @dataclass
@@ -100,7 +137,10 @@ class BoxRegion:
 MaskedRegionSpecifier = List[BoxRegion]
 
 
-class MaskScoreWrapper(gym.Wrapper):
+class MaskScoreWrapper(
+    gym.Wrapper[npt.NDArray, ActType, npt.NDArray, ActType],
+    Generic[ActType],
+):
     """Mask a list of box-shaped regions in the observation to hide reward info.
 
     Intended for environments whose observations are raw pixels (like Atari
@@ -130,27 +170,32 @@ def __init__(
         super().__init__(env)
         self.fill_value = np.array(fill_value, env.observation_space.dtype)
 
+        if env.observation_space.shape is None:
+            raise ValueError("Observation space must have a shape.")  # pragma: no cover
         self.mask = np.ones(env.observation_space.shape, dtype=bool)
         for r in score_regions:
             if r.x[0] >= r.x[1] or r.y[0] >= r.y[1]:
                 raise ValueError('Invalid region: "x" and "y" must be increasing.')
             self.mask[r.x[0] : r.x[1], r.y[0] : r.y[1]] = 0
 
-    def _mask_obs(self, obs):
+    def _mask_obs(self, obs) -> npt.NDArray:
         return np.where(self.mask, obs, self.fill_value)
 
-    def step(self, action):
-        """Returns (obs, rew, done, info) with masked obs."""
-        obs, rew, done, info = self.env.step(action)
-        return self._mask_obs(obs), rew, done, info
+    def step(
+        self,
+        action: ActType,
+    ) -> Tuple[npt.NDArray, SupportsFloat, bool, bool, Dict[str, Any]]:
+        """Returns (obs, rew, terminated, truncated, info) with masked obs."""
+        obs, rew, terminated, truncated, info = self.env.step(action)
+        return self._mask_obs(obs), rew, terminated, truncated, info
 
     def reset(self, **kwargs):
         """Returns masked reset observation."""
-        obs = self.env.reset(**kwargs)
-        return self._mask_obs(obs)
+        obs, info = self.env.reset(**kwargs)
+        return self._mask_obs(obs), info
 
 
-class ObsCastWrapper(gym.Wrapper):
+class ObsCastWrapper(gym.ObservationWrapper):
     """Cast observations to specified dtype.
 
     Some external environments return observations of a different type than the
@@ -169,21 +214,17 @@ def __init__(self, env: gym.Env, dtype: np.dtype):
         super().__init__(env)
         self.dtype = dtype
 
-    def reset(self):
-        """Returns reset observation, cast to self.dtype."""
-        return super().reset().astype(self.dtype)
-
-    def step(self, action):
-        """Returns (obs, rew, done, info) with obs cast to self.dtype."""
-        obs, rew, done, info = super().step(action)
-        return obs.astype(self.dtype), rew, done, info
+    def observation(self, obs):
+        """Returns observation cast to self.dtype."""
+        return obs.astype(self.dtype)
 
 
 class AbsorbAfterDoneWrapper(gym.Wrapper):
     """Transition into absorbing state instead of episode termination.
 
-    When the environment being wrapped returns `done=True`, we return an absorbing
-    observation. This wrapper always returns `done=False`.
+    When the environment being wrapped returns `terminated=True` or `truncated=True`,
+    we return an absorbing observation.
+    This wrapper always returns `terminated=False` and `truncated=False`.
 
     A convenient way to add absorbing states to environments like MountainCar.
     """
@@ -217,28 +258,28 @@ def reset(self, *args, **kwargs):
     def step(self, action):
         """Advance the environment by one step.
 
-        This wrapped `step()` always returns done=False.
+        This wrapped `step()` always returns terminated=False and truncated=False.
 
-        After the first done is returned by the underlying Env, we enter an artificial
-        absorb state.
+        After the first time either terminated or truncated is returned by the
+        underlying Env, we enter an artificial absorb state.
 
         In this artificial absorb state, we stop calling
         `self.env.step(action)` (i.e. the `action` argument is entirely ignored) and
-        we return fixed values for obs, rew, done, and info. The values of `obs` and
-        `rew` depend on initialization arguments. `info` is always an empty dictionary.
+        we return fixed values for obs, rew, terminated, truncated, and info.
+        The values of `obs` and `rew` depend on initialization arguments.
+        `info` is always an empty dictionary.
         """
         if not self.at_absorb_state:
-            inner_obs, inner_rew, done, inner_info = self.env.step(action)
-            if done:
+            obs, rew, terminated, truncated, info = self.env.step(action)
+            if terminated or truncated:
                 # Initialize the artificial absorb state, which we will repeatedly use
                 # starting on the next call to `step()`.
                 self.at_absorb_state = True
 
                 if self.absorb_obs_default is None:
-                    self.absorb_obs_this_episode = inner_obs
+                    self.absorb_obs_this_episode = obs
                 else:
                     self.absorb_obs_this_episode = self.absorb_obs_default
-            obs, rew, info = inner_obs, inner_rew, inner_info
         else:
             assert self.absorb_obs_this_episode is not None
             assert self.absorb_reward is not None
@@ -246,25 +287,17 @@ def step(self, action):
             rew = self.absorb_reward
             info = {}
 
-        return obs, rew, False, info
-
-
-def make_env_no_wrappers(env_name: str, **kwargs) -> gym.Env:
-    """Gym sometimes wraps envs in TimeLimit before returning from gym.make().
-
-    This helper method builds directly from spec to avoid this wrapper.
-    """
-    return gym.envs.registry.env_specs[env_name].make(**kwargs)
+        return obs, rew, False, False, info
 
 
 def get_gym_max_episode_steps(env_name: str) -> Optional[int]:
     """Get the `max_episode_steps` attribute associated with a gym Spec."""
-    return gym.envs.registry.env_specs[env_name].max_episode_steps
+    return gym.spec(env_name).max_episode_steps
 
 
 def sample_distribution(
     p: np.ndarray,
-    random: np.random.RandomState,
+    random: np.random.Generator,
 ) -> int:
     """Samples an integer with probabilities given by p."""
     return random.choice(np.arange(len(p)), p=p)


Test patch:
diff --git a/src/seals/testing/envs.py b/src/seals/testing/envs.py
index a927a50..4d6dc0c 100644
--- a/src/seals/testing/envs.py
+++ b/src/seals/testing/envs.py
@@ -12,17 +12,18 @@
     Iterator,
     List,
     Mapping,
-    Optional,
     Sequence,
+    SupportsFloat,
     Tuple,
 )
 
-import gym
+import gymnasium as gym
 import numpy as np
 
-Step = Tuple[Any, Optional[float], bool, Mapping[str, Any]]
+Step = Tuple[Any, SupportsFloat, bool, bool, Mapping[str, Any]]
 Rollout = Sequence[Step]
-"""A sequence of 4-tuples (obs, rew, done, info) as returned by `get_rollout`."""
+"""A sequence of 5-tuples (obs, rew, terminated, truncated, info) as returned by
+`get_rollout`."""
 
 
 def make_env_fixture(
@@ -88,16 +89,17 @@ def matches_list(env_name: str, patterns: Iterable[str]) -> bool:
 
 
 def get_rollout(env: gym.Env, actions: Iterable[Any]) -> Rollout:
-    """Performs sequence of actions `actions` in `env`.
+    """Performs a sequence of actions `actions` in `env`.
 
     Args:
-      env: the environment to rollout in.
+      env: the environment to roll out in.
       actions: the actions to perform.
 
     Returns:
-      A sequence of 4-tuples (obs, rew, done, info).
+      A sequence of 5-tuples (obs, rew, terminated, truncated, info).
     """
-    ret: List[Step] = [(env.reset(), None, False, {})]
+    obs, info = env.reset()
+    ret: List[Step] = [(obs, 0, False, False, info)]
     for act in actions:
         ret.append(env.step(act))
     return ret
@@ -110,11 +112,12 @@ def assert_equal_rollout(rollout_a: Rollout, rollout_b: Rollout) -> None:
         AssertionError if they are not equal.
     """
     for step_a, step_b in zip(rollout_a, rollout_b):
-        ob_a, rew_a, done_a, info_a = step_a
-        ob_b, rew_b, done_b, info_b = step_b
+        ob_a, rew_a, terminated_a, truncated_a, info_a = step_a
+        ob_b, rew_b, terminated_b, truncated_b, info_b = step_b
         np.testing.assert_equal(ob_a, ob_b)
         assert rew_a == rew_b
-        assert done_a == done_b
+        assert terminated_a == terminated_b
+        assert truncated_a == truncated_b
         np.testing.assert_equal(info_a, info_b)
 
 
@@ -155,13 +158,10 @@ def test_seed(
     env.action_space.seed(0)
     actions = [env.action_space.sample() for _ in range(rollout_len)]
     # With the same seed, should always get the same result
-    seeds = env.seed(42)
-    # output of env.seed should be a list, but atari environments return a tuple.
-    assert isinstance(seeds, (list, tuple))
-    assert len(seeds) > 0
+    env.reset(seed=42)
     rollout_a = get_rollout(env, actions)
 
-    env.seed(42)
+    env.reset(seed=42)
     rollout_b = get_rollout(env, actions)
 
     assert_equal_rollout(rollout_a, rollout_b)
@@ -171,9 +171,9 @@ def test_seed(
     # seeds should produce the same starting state.
     def different_seeds_same_rollout(seed1, seed2):
         new_actions = [env.action_space.sample() for _ in range(rollout_len)]
-        env.seed(seed1)
+        env.reset(seed=seed1)
         new_rollout_1 = get_rollout(env, new_actions)
-        env.seed(seed2)
+        env.reset(seed=seed2)
         new_rollout_2 = get_rollout(env, new_actions)
         return has_same_observations(new_rollout_1, new_rollout_2)
 
@@ -195,21 +195,23 @@ def _check_obs(obs: np.ndarray, obs_space: gym.Space) -> None:
 def _sample_and_check(env: gym.Env, obs_space: gym.Space) -> bool:
     """Sample from env and check return value is of valid type."""
     act = env.action_space.sample()
-    obs, rew, done, info = env.step(act)
+    obs, rew, terminated, truncated, info = env.step(act)
     _check_obs(obs, obs_space)
     assert isinstance(rew, float)
-    assert isinstance(done, bool)
+    assert isinstance(terminated, bool)
+    assert isinstance(truncated, bool)
     assert isinstance(info, dict)
-    return done
+    return terminated or truncated
 
 
-def _is_mujoco_env(env: gym.Env) -> bool:
+def is_mujoco_env(env: gym.Env) -> bool:
+    """True if `env` is a MuJoCo environment."""
     return hasattr(env, "sim") and hasattr(env, "model")
 
 
 def test_rollout_schema(
     env: gym.Env,
-    steps_after_done: int = 10,
+    steps_after_terminated: int = 10,
     max_steps: int = 10_000,
     check_episode_ends: bool = True,
 ) -> None:
@@ -217,20 +219,21 @@ def test_rollout_schema(
 
     Args:
         env: The environment to test.
-        steps_after_done: The number of steps to take after `done` is True, the nominal
-            episode termination. This is an abuse of the Gym API, but we would like the
-            environments to handle this case gracefully.
-        max_steps: Test fails if we do not get `done` after this many timesteps.
+        steps_after_terminated: The number of steps to take after `terminated` is True,
+            the nominal episode termination. This is an abuse of the Gym API,
+            but we would like the environments to handle this case gracefully.
+        max_steps: Test fails if we do not get `terminated` after this many timesteps.
         check_episode_ends: Check that episode ends after `max_steps` steps, and that
-            steps taken after `done` is true are of the correct type.
+            steps taken after `terminated` is true are of the correct type.
 
     Raises:
         AssertionError if test fails.
     """
     obs_space = env.observation_space
-    obs = env.reset()
+    obs, _ = env.reset(seed=0)
     _check_obs(obs, obs_space)
 
+    done = False
     for _ in range(max_steps):
         done = _sample_and_check(env, obs_space)
         if done:
@@ -239,7 +242,7 @@ def test_rollout_schema(
     if check_episode_ends:
         assert done, "did not get to end of episode"
 
-        for _ in range(steps_after_done):
+        for _ in range(steps_after_terminated):
             _sample_and_check(env, obs_space)
 
 
@@ -257,7 +260,7 @@ def test_premature_step(env: gym.Env, skip_fn, raises_fn) -> None:
     Raises:
         AssertionError if test fails.
     """
-    if _is_mujoco_env(env):  # pragma: no cover
+    if is_mujoco_env(env):  # pragma: no cover
         # We can't use isinstance since importing mujoco_py will fail on
         # machines without MuJoCo installed
         skip_fn("MuJoCo environments cannot perform this check.")
@@ -267,49 +270,6 @@ def test_premature_step(env: gym.Env, skip_fn, raises_fn) -> None:
         env.step(act)
 
 
-def test_render(env: gym.Env, raises_fn) -> None:
-    """Test that render() supports the modes declared.
-
-    Example usage in pytest:
-        test_render(env, raises_fn=pytest.raises)
-
-    Args:
-        env: The environment to test.
-        raises_fn: Context manager to check NotImplementedError is thrown when
-            environment metadata indicates modes are supported.
-
-    Raises:
-        AssertionError: if test fails. This occurs if:
-            (a) `env.render(mode=mode)` fails for any mode declared supported
-            in `env.metadata["render.modes"]`; (b) env.render() *succeeds* when
-            `env.metadata["render.modes"]` is empty; (c) `env.render(mode="rgb_array")`
-            returns different values at the same time step.
-    """
-    env.reset()  # make sure environment is in consistent state
-
-    render_modes = env.metadata["render.modes"]
-    if not render_modes:
-        # No modes supported -- render() should fail.
-        with raises_fn(NotImplementedError):
-            env.render()
-    else:
-        for mode in render_modes:
-            env.render(mode=mode)
-
-        # WARNING(adam): there seems to be a memory leak with Gym 0.17.3
-        # & MuJoCoPy 1.50.1.68. `MujocoEnv.close()` does not call `finish()`
-        # on the viewer (commented out) so the resources are not released.
-        # For now this is OK, but may bite if we end up testing a lot of
-        # MuJoCo environments.
-        is_mujoco = _is_mujoco_env(env)
-        if "rgb_array" in render_modes and not is_mujoco:
-            # Render should not change without calling `step()`.
-            # MuJoCo rendering fails this check, ignore -- not much we can do.
-            resa = env.render(mode="rgb_array")
-            resb = env.render(mode="rgb_array")
-            assert np.allclose(resa, resb)
-
-
 class CountingEnv(gym.Env):
     """At timestep `t` of each episode, has `t == obs == reward / 10`.
 
@@ -335,10 +295,10 @@ def __init__(self, episode_length: int = 5):
         self.episode_length = episode_length
         self.timestep = None
 
-    def reset(self):
+    def reset(self, seed=None, options={}):
         """Reset method for CountingEnv."""
         t, self.timestep = 0, 1
-        return np.array(t, dtype=self.observation_space.dtype)
+        return np.array(t, dtype=self.observation_space.dtype), {}
 
     def step(self, action):
         """Step method for CountingEnv."""
@@ -352,5 +312,5 @@ def step(self, action):
         t, self.timestep = self.timestep, self.timestep + 1
         obs = np.array(t, dtype=self.observation_space.dtype)
         rew = t * 10.0
-        done = t == self.episode_length
-        return obs, rew, done, {}
+        terminated = t == self.episode_length
+        return obs, rew, terminated, False, {}
diff --git a/tests/test_base_env.py b/tests/test_base_env.py
index 59a3e32..9f089f3 100644
--- a/tests/test_base_env.py
+++ b/tests/test_base_env.py
@@ -4,7 +4,7 @@
 so the tests in this file focus on features unique to classes in `base_envs`.
 """
 
-import gym
+import gymnasium as gym
 import numpy as np
 import pytest
 
@@ -17,11 +17,12 @@ class NewEnv(base_envs.TabularModelMDP):
 
     def __init__(self):
         """Build environment."""
+        np.random.seed(0)
         nS = 3
         nA = 2
-        transition_matrix = np.random.rand(nS, nA, nS)
+        transition_matrix = np.random.random((nS, nA, nS))
         transition_matrix /= transition_matrix.sum(axis=2)[:, :, None]
-        reward_matrix = np.random.rand(nS)
+        reward_matrix = np.random.random((nS,))
         super().__init__(
             transition_matrix=transition_matrix,
             reward_matrix=reward_matrix,
@@ -36,7 +37,7 @@ def test_base_envs():
 
     envs.test_premature_step(env, skip_fn=pytest.skip, raises_fn=pytest.raises)
 
-    env.reset()
+    env.reset(seed=0)
     assert env.n_actions_taken == 0
     env.step(env.action_space.sample())
     assert env.n_actions_taken == 1
@@ -51,6 +52,9 @@ def test_base_envs():
     with pytest.raises(ValueError, match=r".*not in.*"):
         env.state = bad_state  # type: ignore
 
+    with pytest.raises(NotImplementedError, match=r"Options not supported.*"):
+        env.reset(options={"option": "value"})
+
 
 def test_tabular_env_validation():
     """Test input validation for base_envs.TabularModelEnv."""
@@ -87,7 +91,7 @@ def test_tabular_env_validation():
         transition_matrix=np.zeros((3, 1, 3)),
         reward_matrix=np.zeros((3,)),
     )
-    env.reset()
+    env.reset(seed=0)
     with pytest.raises(ValueError, match=r".*not in.*"):
         env.step(4)
 
@@ -98,12 +102,12 @@ def test_expose_pomdp_state_wrapper():
     wrapped_env = base_envs.ExposePOMDPStateWrapper(env)
 
     assert wrapped_env.observation_space == env.state_space
-    state = wrapped_env.reset()
+    state, _ = wrapped_env.reset(seed=0)
     assert state == env.state
     assert state in env.state_space
 
     action = env.action_space.sample()
-    next_state, reward, done, info = wrapped_env.step(action)
+    next_state, reward, terminated, truncated, info = wrapped_env.step(action)
     assert next_state == env.state
     assert next_state in env.state_space
 
diff --git a/tests/test_envs.py b/tests/test_envs.py
index b2e86d4..2cbce74 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -1,19 +1,19 @@
 """Smoke tests for all environments."""
 
-from typing import List
+from typing import List, Union
 
-import gym
-from gym.envs import registration
+import gymnasium as gym
+from gymnasium.envs import registration
+import numpy as np
 import pytest
 
 import seals  # noqa: F401 required for env registration
 from seals.atari import SCORE_REGIONS, _get_score_region, _seals_name, make_atari_env
 from seals.testing import envs
+from seals.testing.envs import is_mujoco_env
 
 ENV_NAMES: List[str] = [
-    env_spec.id
-    for env_spec in registration.registry.all()
-    if env_spec.id.startswith("seals/")
+    env_id for env_id in registration.registry.keys() if env_id.startswith("seals/")
 ]
 
 
@@ -120,6 +120,7 @@ def test_seed(self, env: gym.Env, env_name: str):
                 "seals/KingKong-Unmasked-v5",
                 "seals/Koolaid-Unmasked-v5",
                 "seals/NameThisGame-Unmasked-v5",
+                "seals/Casino-Unmasked-v5",
             ]
             rollout_len = 100 if env_name not in slow_random_envs else 400
             num_seeds = 2 if env_name in ATARI_NO_FRAMESKIP_ENVS else 10
@@ -149,6 +150,31 @@ def test_rollout_schema(self, env: gym.Env, env_name: str):
         else:
             envs.test_rollout_schema(env)
 
-    def test_render(self, env: gym.Env):
-        """Tests `render()` supports modes specified in environment metadata."""
-        envs.test_render(env, raises_fn=pytest.raises)
+    def test_render_modes(self, env_name: str):
+        """Tests that all render modes specifeid in the metadata work.
+
+        Note: we only check that no exception is thrown.
+        There is no test to see if something reasonable is rendered.
+        """
+        for mode in gym.make(env_name).metadata["render_modes"]:
+            # GIVEN
+            env = gym.make(env_name, render_mode=mode)
+            env.reset(seed=0)
+
+            # WHEN
+            if mode == "rgb_array" and not is_mujoco_env(env):
+                # The render should not change without calling `step()`.
+                # MuJoCo rendering fails this check, ignore -- not much we can do.
+                r1: Union[np.ndarray, List[np.ndarray], None] = env.render()
+                r2: Union[np.ndarray, List[np.ndarray], None] = env.render()
+                assert r1 is not None
+                assert r2 is not None
+                assert np.allclose(r1, r2)
+            else:
+                env.render()
+
+            # THEN
+            # no error raised
+
+            # CLEANUP
+            env.close()
diff --git a/tests/test_mujoco_rl.py b/tests/test_mujoco_rl.py
index b0e362c..13a6304 100644
--- a/tests/test_mujoco_rl.py
+++ b/tests/test_mujoco_rl.py
@@ -2,7 +2,7 @@
 
 from typing import Tuple
 
-import gym
+import gymnasium as gym
 import pytest
 import stable_baselines3
 from stable_baselines3.common import evaluation
@@ -35,9 +35,9 @@ def test_fixed_env_model_as_good_as_gym_env_model(env_base: str):  # pragma: no
     """Compare original and modified MuJoCo v3 envs."""
     train_timesteps = 200000
 
-    gym_reward, _ = _eval_env(f"{env_base}-v3", total_timesteps=train_timesteps)
+    gym_reward, _ = _eval_env(f"{env_base}-v4", total_timesteps=train_timesteps)
     fixed_reward, _ = _eval_env(
-        f"seals/{env_base}-v0",
+        f"seals/{env_base}-v1",
         total_timesteps=train_timesteps,
     )
 
diff --git a/tests/test_util.py b/tests/test_util.py
index 5829ebb..b1ca318 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -2,7 +2,7 @@
 
 import collections
 
-import gym
+import gymnasium as gym
 import numpy as np
 import pytest
 
@@ -22,11 +22,11 @@ def test_mask_score_wrapper_enforces_spec():
 def test_sample_distribution():
     """Test util.sample_distribution."""
     distr_size = 5
-    distr = np.random.rand(distr_size)
+    distr = np.random.random((distr_size,))
     distr /= distr.sum()
 
     n_samples = 1000
-    rng = np.random.RandomState()
+    rng = np.random.default_rng(0)
     sample_count = collections.Counter(
         util.sample_distribution(distr, rng) for _ in range(n_samples)
     )
@@ -39,8 +39,8 @@ def test_sample_distribution():
 
     # Same seed gives same samples
     assert all(
-        util.sample_distribution(distr, random=np.random.RandomState(seed))
-        == util.sample_distribution(distr, random=np.random.RandomState(seed))
+        util.sample_distribution(distr, random=np.random.default_rng(seed))
+        == util.sample_distribution(distr, random=np.random.default_rng(seed))
         for seed in range(20)
     )
 
diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py
index 06fac22..136db64 100644
--- a/tests/test_wrappers.py
+++ b/tests/test_wrappers.py
@@ -13,9 +13,9 @@ def test_auto_reset_wrapper_pad(episode_length=3, n_steps=100, n_manual_reset=2)
     AutoResetWrapper that pads trajectory with an extra transition containing the
     terminal observations.
     Also check that calls to .reset() do not interfere with automatic resets.
-    Due to the padding the number of steps counted inside the environment and the number
-    of steps performed outside the environment, i.e. the number of actions performed,
-    will differ. This test checks that this difference is consistent.
+    Due to the padding, the number of steps counted inside the environment and the
+    number of steps performed outside the environment, i.e., the number of actions
+    performed, will differ. This test checks that this difference is consistent.
     """
     env = util.AutoResetWrapper(
         envs.CountingEnv(episode_length=episode_length),
@@ -23,7 +23,7 @@ def test_auto_reset_wrapper_pad(episode_length=3, n_steps=100, n_manual_reset=2)
     )
 
     for _ in range(n_manual_reset):
-        obs = env.reset()
+        obs, info = env.reset()
         assert obs == 0
 
         # We count the number of episodes, so we can sanity check the padding.
@@ -31,10 +31,11 @@ def test_auto_reset_wrapper_pad(episode_length=3, n_steps=100, n_manual_reset=2)
         next_episode_end = episode_length
         for t in range(1, n_steps + 1):
             act = env.action_space.sample()
-            obs, rew, done, info = env.step(act)
+            obs, rew, terminated, truncated, info = env.step(act)
 
-            # AutoResetWrapper overrides all done signals.
-            assert done is False
+            # AutoResetWrapper overrides all terminated and truncated signals.
+            assert terminated is False
+            assert truncated is False
 
             if t == next_episode_end:
                 # Unlike the AutoResetWrapper that discards terminal observations,
@@ -75,16 +76,17 @@ def test_auto_reset_wrapper_discard(episode_length=3, n_steps=100, n_manual_rese
     )
 
     for _ in range(n_manual_reset):
-        obs = env.reset()
+        obs, info = env.reset()
         assert obs == 0
 
         for t in range(1, n_steps + 1):
             act = env.action_space.sample()
-            obs, rew, done, info = env.step(act)
+            obs, rew, terminated, truncated, info = env.step(act)
             expected_obs = t % episode_length
 
             assert obs == expected_obs
-            assert done is False
+            assert terminated is False
+            assert truncated is False
 
             if expected_obs == 0:  # End of episode
                 assert info.get("terminal_observation", None) == episode_length
@@ -113,8 +115,9 @@ def test_absorb_repeat_custom_state(
         env.reset()
         for t in range(1, n_steps + 1):
             act = env.action_space.sample()
-            obs, rew, done, _ = env.step(act)
-            assert done is False
+            obs, rew, terminated, truncated, _ = env.step(act)
+            assert terminated is False
+            assert truncated is False
             if t > episode_length:
                 expected_obs = absorb_obs
                 expected_rew = absorb_reward
@@ -134,8 +137,9 @@ def test_absorb_repeat_final_state(episode_length=6, n_steps=100, n_manual_reset
         env.reset()
         for t in range(1, n_steps + 1):
             act = env.action_space.sample()
-            obs, rew, done, _ = env.step(act)
-            assert done is False
+            obs, rew, terminated, truncated, _ = env.step(act)
+            assert terminated is False
+            assert truncated is False
             if t > episode_length:
                 expected_obs = episode_length
                 expected_rew = -1
@@ -153,16 +157,19 @@ def test_obs_cast(dtype: np.dtype, episode_length: int = 5):
     Test uses CountingEnv with small integers, which can be represented in
     all the specified dtypes without any loss of precision.
     """
-    env = envs.CountingEnv(episode_length=episode_length)
-    env = util.ObsCastWrapper(env, dtype)
+    env = util.ObsCastWrapper(
+        envs.CountingEnv(episode_length=episode_length),
+        dtype,
+    )
 
-    obs = env.reset()
+    obs, _ = env.reset()
     assert obs.dtype == dtype
     assert obs == 0
     for t in range(1, episode_length + 1):
         act = env.action_space.sample()
-        obs, rew, done, _ = env.step(act)
-        assert done == (t == episode_length)
+        obs, rew, terminated, truncated, _ = env.step(act)
+        assert terminated == (t == episode_length)
+        assert truncated is False
         assert obs.dtype == dtype
         assert obs == t
         assert rew == t * 10.0