Upgrade gymnasium to 1.0.0

.github/workflows/tests.yaml

@@ -10,7 +10,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04, windows-latest]
     runs-on: ${{ matrix.os }}
@@ -34,8 +34,8 @@ jobs:
       - name: Install jax
         if: runner.os == 'Linux' || runner.os == 'macOS'
         run: poetry install -E "pytest jax"
-      - name: Run gymnasium migration dependencies
-        run: poetry run pip install "stable_baselines3==2.0.0a1"
+      # - name: Run gymnasium migration dependencies
+      #   run: poetry run pip install "stable_baselines3==2.0.0a1"
       - name: Run gymnasium tests
         run: poetry run pytest tests/test_classic_control_gymnasium.py
       - name: Run core tests with jax
@@ -53,7 +53,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04, windows-latest]
     runs-on: ${{ matrix.os }}
@@ -77,8 +77,8 @@ jobs:
       - name: Install jax
         if: runner.os == 'Linux' || runner.os == 'macOS'
         run: poetry install -E "pytest atari jax"
-      - name: Run gymnasium migration dependencies
-        run: poetry run pip install "stable_baselines3==2.0.0a1" "gymnasium[atari,accept-rom-license]==0.28.1"  "ale-py==0.8.1" 
+      # - name: Run gymnasium migration dependencies
+      #   run: poetry run pip install "stable_baselines3==2.0.0a1" "gymnasium[atari,accept-rom-license]==0.28.1"  "ale-py==0.8.1" 
       - name: Run gymnasium tests
         run: poetry run pytest tests/test_atari_gymnasium.py
       - name: Run gymnasium tests with jax
@@ -89,7 +89,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04, windows-latest]
     runs-on: ${{ matrix.os }}
@@ -115,7 +115,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04]
     runs-on: ${{ matrix.os }}
@@ -134,8 +134,8 @@ jobs:
         run: poetry install -E "pytest mujoco dm_control jax"
       - name: Downgrade setuptools
         run: poetry run pip install setuptools==59.5.0
-      - name: Run gymnasium migration dependencies
-        run: poetry run pip install "stable_baselines3==2.0.0a1"
+      # - name: Run gymnasium migration dependencies
+      #   run: poetry run pip install "stable_baselines3==2.0.0a1"
       - name: install mujoco dependencies
         run: |
           sudo apt-get update && sudo apt-get -y install libgl1-mesa-glx libosmesa6 libglfw3
@@ -175,7 +175,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04]
     runs-on: ${{ matrix.os }}
@@ -193,7 +193,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04]
     runs-on: ${{ matrix.os }}
@@ -219,7 +219,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04]
     runs-on: ${{ matrix.os }}
@@ -245,7 +245,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10"]
         poetry-version: ["1.7"]
         os: [ubuntu-22.04]
     runs-on: ${{ matrix.os }}

.pre-commit-config.yaml

@@ -37,50 +37,50 @@ repos:
         args:
           - --ignore-words-list=nd,reacher,thist,ths,magent,ba
           - --skip=docs/css/termynal.css,docs/js/termynal.js,docs/get-started/CleanRL_Huggingface_Integration_Demo.ipynb
-  - repo: https://github.com/python-poetry/poetry
-    rev: 1.3.2
+  - repo: https://github.com/python-poetry/poetry-plugin-export
+    rev: 1.9.0
     hooks:
       - id: poetry-export
-        name: poetry-export requirements.txt
+        name: poetry export requirements.txt
         args: ["--without-hashes", "-o", "requirements/requirements.txt"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-atari.txt
+        name: poetry export requirements-atari.txt
         args: ["--without-hashes", "-o", "requirements/requirements-atari.txt", "-E", "atari"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-mujoco.txt
+        name: poetry export requirements-mujoco.txt
         args: ["--without-hashes", "-o", "requirements/requirements-mujoco.txt", "-E", "mujoco"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-dm_control.txt
+        name: poetry export requirements-dm_control.txt
         args: ["--without-hashes", "-o", "requirements/requirements-dm_control.txt", "-E", "dm_control"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-procgen.txt
+        name: poetry export requirements-procgen.txt
         args: ["--without-hashes", "-o", "requirements/requirements-procgen.txt", "-E", "procgen"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-envpool.txt
+        name: poetry export requirements-envpool.txt
         args: ["--without-hashes", "-o", "requirements/requirements-envpool.txt", "-E", "envpool"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-pettingzoo.txt
+        name: poetry export requirements-pettingzoo.txt
         args: ["--without-hashes", "-o", "requirements/requirements-pettingzoo.txt", "-E", "pettingzoo"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-jax.txt
+        name: poetry export requirements-jax.txt
         args: ["--without-hashes", "-o", "requirements/requirements-jax.txt", "-E", "jax"]
         stages: [manual]
       - id: poetry-export
         name: poetry-export requirements-optuna.txt
         args: ["--without-hashes", "-o", "requirements/requirements-optuna.txt", "-E", "optuna"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-docs.txt
+        name: poetry export requirements-docs.txt
         args: ["--without-hashes", "-o", "requirements/requirements-docs.txt", "-E", "docs"]
         stages: [manual]
       - id: poetry-export
-        name: poetry-export requirements-cloud.txt
+        name: poetry export requirements-cloud.txt
         args: ["--without-hashes", "-o", "requirements/requirements-cloud.txt", "-E", "cloud"]
         stages: [manual]

cleanrl/c51.py

@@ -13,6 +13,8 @@
 from stable_baselines3.common.buffers import ReplayBuffer
 from torch.utils.tensorboard import SummaryWriter
 
+from cleanrl_utils.same_model_vector_env import SameModelSyncVectorEnv
+
 
 @dataclass
 class Args:
@@ -160,7 +162,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
     device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
 
     # env setup
-    envs = gym.vector.SyncVectorEnv(
+    envs = SameModelSyncVectorEnv(
         [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
     )
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
@@ -195,11 +197,12 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
 
         # TRY NOT TO MODIFY: record rewards for plotting purposes
         if "final_info" in infos:
-            for info in infos["final_info"]:
-                if info and "episode" in info:
-                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
-                    writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
-                    writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
+            for i in range(envs.num_envs):
+                if infos["final_info"]["_episode"][i]:
+                    print(f"global_step={global_step}, episodic_return={infos['final_info']['episode']['r'][i]}")
+                    writer.add_scalar("charts/episodic_return", infos["final_info"]["episode"]["r"][i], global_step)
+                    writer.add_scalar("charts/episodic_length", infos["final_info"]["episode"]["l"][i], global_step)
+                    break
 
         # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
         real_next_obs = next_obs.copy()

cleanrl/c51_atari.py

@@ -4,6 +4,7 @@
 import time
 from dataclasses import dataclass
 
+import ale_py
 import gymnasium as gym
 import numpy as np
 import torch
@@ -12,14 +13,17 @@
 import tyro
 from stable_baselines3.common.atari_wrappers import (
     ClipRewardEnv,
-    EpisodicLifeEnv,
     FireResetEnv,
     MaxAndSkipEnv,
     NoopResetEnv,
 )
 from stable_baselines3.common.buffers import ReplayBuffer
 from torch.utils.tensorboard import SummaryWriter
 
+from cleanrl_utils.same_model_vector_env import SameModelSyncVectorEnv
+
+gym.register_envs(ale_py)
+
 
 @dataclass
 class Args:
@@ -92,13 +96,14 @@ def thunk():
 
         env = NoopResetEnv(env, noop_max=30)
         env = MaxAndSkipEnv(env, skip=4)
-        env = EpisodicLifeEnv(env)
+        # EpisodicLifeEnv conflicts with gym v1.0.0's RecordEpisodeStatistics and will be fixed later.
+        # env = EpisodicLifeEnv(env)
         if "FIRE" in env.unwrapped.get_action_meanings():
             env = FireResetEnv(env)
         env = ClipRewardEnv(env)
         env = gym.wrappers.ResizeObservation(env, (84, 84))
-        env = gym.wrappers.GrayScaleObservation(env)
-        env = gym.wrappers.FrameStack(env, 4)
+        env = gym.wrappers.GrayscaleObservation(env)
+        env = gym.wrappers.FrameStackObservation(env, 4)
 
         env.action_space.seed(seed)
         return env
@@ -182,7 +187,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
     device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
 
     # env setup
-    envs = gym.vector.SyncVectorEnv(
+    envs = SameModelSyncVectorEnv(
         [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
     )
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
@@ -218,11 +223,12 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
 
         # TRY NOT TO MODIFY: record rewards for plotting purposes
         if "final_info" in infos:
-            for info in infos["final_info"]:
-                if info and "episode" in info:
-                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
-                    writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
-                    writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
+            for i in range(envs.num_envs):
+                if infos["final_info"]["_episode"][i]:
+                    print(f"global_step={global_step}, episodic_return={infos['final_info']['episode']['r'][i]}")
+                    writer.add_scalar("charts/episodic_return", infos["final_info"]["episode"]["r"][i], global_step)
+                    writer.add_scalar("charts/episodic_length", infos["final_info"]["episode"]["l"][i], global_step)
+                    break
 
         # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
         real_next_obs = next_obs.copy()

cleanrl/c51_atari_jax.py

@@ -7,7 +7,7 @@
 os.environ[
     "XLA_PYTHON_CLIENT_MEM_FRACTION"
 ] = "0.7"  # see https://github.com/google/jax/discussions/6332#discussioncomment-1279991
-
+import ale_py
 import flax
 import flax.linen as nn
 import gymnasium as gym
@@ -19,14 +19,17 @@
 from flax.training.train_state import TrainState
 from stable_baselines3.common.atari_wrappers import (
     ClipRewardEnv,
-    EpisodicLifeEnv,
     FireResetEnv,
     MaxAndSkipEnv,
     NoopResetEnv,
 )
 from stable_baselines3.common.buffers import ReplayBuffer
 from torch.utils.tensorboard import SummaryWriter
 
+from cleanrl_utils.same_model_vector_env import SameModelSyncVectorEnv
+
+gym.register_envs(ale_py)
+
 
 @dataclass
 class Args:
@@ -95,13 +98,14 @@ def thunk():
 
         env = NoopResetEnv(env, noop_max=30)
         env = MaxAndSkipEnv(env, skip=4)
-        env = EpisodicLifeEnv(env)
+        # EpisodicLifeEnv conflicts with gym v1.0.0's RecordEpisodeStatistics and will be fixed later.
+        # env = EpisodicLifeEnv(env)
         if "FIRE" in env.unwrapped.get_action_meanings():
             env = FireResetEnv(env)
         env = ClipRewardEnv(env)
         env = gym.wrappers.ResizeObservation(env, (84, 84))
-        env = gym.wrappers.GrayScaleObservation(env)
-        env = gym.wrappers.FrameStack(env, 4)
+        env = gym.wrappers.GrayscaleObservation(env)
+        env = gym.wrappers.FrameStackObservation(env, 4)
 
         env.action_space.seed(seed)
         return env
@@ -181,7 +185,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
     key, q_key = jax.random.split(key, 2)
 
     # env setup
-    envs = gym.vector.SyncVectorEnv(
+    envs = SameModelSyncVectorEnv(
         [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
     )
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
@@ -278,11 +282,12 @@ def get_action(q_state, obs):
 
         # TRY NOT TO MODIFY: record rewards for plotting purposes
         if "final_info" in infos:
-            for info in infos["final_info"]:
-                if info and "episode" in info:
-                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
-                    writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
-                    writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
+            for i in range(envs.num_envs):
+                if infos["final_info"]["_episode"][i]:
+                    print(f"global_step={global_step}, episodic_return={infos['final_info']['episode']['r'][i]}")
+                    writer.add_scalar("charts/episodic_return", infos["final_info"]["episode"]["r"][i], global_step)
+                    writer.add_scalar("charts/episodic_length", infos["final_info"]["episode"]["l"][i], global_step)
+                    break
 
         # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
         real_next_obs = next_obs.copy()

cleanrl/c51_jax.py

@@ -16,6 +16,8 @@
 from stable_baselines3.common.buffers import ReplayBuffer
 from torch.utils.tensorboard import SummaryWriter
 
+from cleanrl_utils.same_model_vector_env import SameModelSyncVectorEnv
+
 
 @dataclass
 class Args:
@@ -153,7 +155,7 @@ def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
     key, q_key = jax.random.split(key, 2)
 
     # env setup
-    envs = gym.vector.SyncVectorEnv(
+    envs = SameModelSyncVectorEnv(
         [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
     )
     assert isinstance(envs.single_action_space, gym.spaces.Discrete), "only discrete action space is supported"
@@ -241,11 +243,12 @@ def loss(q_params, observations, actions, target_pmfs):
 
         # TRY NOT TO MODIFY: record rewards for plotting purposes
         if "final_info" in infos:
-            for info in infos["final_info"]:
-                if info and "episode" in info:
-                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
-                    writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
-                    writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
+            for i in range(envs.num_envs):
+                if infos["final_info"]["_episode"][i]:
+                    print(f"global_step={global_step}, episodic_return={infos['final_info']['episode']['r'][i]}")
+                    writer.add_scalar("charts/episodic_return", infos["final_info"]["episode"]["r"][i], global_step)
+                    writer.add_scalar("charts/episodic_length", infos["final_info"]["episode"]["l"][i], global_step)
+                    break
 
         # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
         real_next_obs = next_obs.copy()