Modernize benchmark stack and switch procgen upstream

Author: Trinkle23897

benchmark/README.md

@@ -7,13 +7,22 @@ The following results are generated from four types of machine:
 3. TPU-VM: 96 core ``Intel(R) Xeon(R) CPU @ 2.00GHz``, 2 NUMA core, TPU v3-8
 4. DGX-A100: 256 core ``AMD EPYC 7742 64-Core Processor``, 8 NUMA core, 8x A100
 
-We use `PongNoFrameskip-v4` (with environment wrappers from [OpenAI baselines](https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py)) and `Ant-v3` for Atari/Mujoco environment benchmark test with `envpool==0.6.1.post1`. Other packages' versions are all in `requirements.txt`:
+The historical numbers below were produced with `PongNoFrameskip-v4` and
+`Ant-v3` on `envpool==0.6.1.post1`. The current benchmark scripts in this
+directory use Gymnasium's `ALE/Pong-v5` and `Ant-v5`, and the baseline
+dependencies plus shared benchmark tooling are installed from
+`requirements.txt`:
 
 ```bash
 $ pip install -r requirements.txt
 ```
 
-To align with other baseline results, FPS is multiplied with `frame_skip` (4 for `PongNoFrameskip-v4` and 5 for `Ant-v3`).
+`test_gym.py` uses only the packages above. `test_envpool.py` additionally
+expects an installed EnvPool build with native modules, so install the
+EnvPool wheel you want to benchmark separately before running it.
+
+To align with other baseline results, FPS is multiplied with `frame_skip` (4
+for Atari and 5 for Mujoco).
 
 ## Highest FPS Overview
 
@@ -53,7 +62,7 @@ python3 test_gym.py --env atari --num-envs 12 --total-step 6000
 python3 test_gym.py --env mujoco --num-envs 12 --total-step 12000
 ```
 
-### Subprocess (gym.vector_env)
+### Subprocess (gymnasium.vector)
 
 Command to run:
 
@@ -66,18 +75,10 @@ python3 test_gym.py --env mujoco --async_ --num-envs 10 --total-step 50000
 
 ### Sample Factory
 
-To run with Ant-v3 in Sample Factory, add one line in `sample_factory/envs/mujoco/mujoco_utils.py`:
-
-```diff
- MUJOCO_ENVS = [
-+    MujocoSpec('mujoco_ant', 'Ant-v3'),
-     MujocoSpec('mujoco_hopper', 'Hopper-v2'),
-     MujocoSpec('mujoco_halfcheetah', 'HalfCheetah-v2'),
-     MujocoSpec('mujoco_humanoid', 'Humanoid-v2'),
- ]
-```
-
-and finally use FPS \* 5 as the result.
+Sample Factory remains a historical reference for now. The latest upstream
+release still depends on `gymnasium<1.0`, `numpy<2`, and an older `ale-py`
+build that is not available in the current devbox package mirror, so it is not
+included in `requirements.txt`.
 
 Command to run:
 
@@ -108,6 +109,9 @@ for i in num_workers:
 
 ### EnvPool
 
+Install an EnvPool wheel for the version you want to benchmark before running
+the commands below.
+
 <!--
 
 ```bash

benchmark/atari_wrappers.py

@@ -19,7 +19,7 @@
 from collections import deque
 
 import cv2
-import gym
+import gymnasium as gym
 import numpy as np
 
 
@@ -38,14 +38,14 @@ def __init__(self, env, noop_max=30):
     self.noop_action = 0
     assert env.unwrapped.get_action_meanings()[0] == "NOOP"
 
-  def reset(self):
-    self.env.reset()
+  def reset(self, *, seed=None, options=None):
+    obs, info = self.env.reset(seed=seed, options=options)
     noops = self.unwrapped.np_random.integers(1, self.noop_max + 1)
     for _ in range(noops):
-      obs, _, done, _ = self.env.step(self.noop_action)
-      if done:
-        obs = self.env.reset()
-    return obs
+      obs, _, terminated, truncated, info = self.env.step(self.noop_action)
+      if terminated or truncated:
+        obs, info = self.env.reset()
+    return obs, info
 
 
 class MaxAndSkipEnv(gym.Wrapper):
@@ -65,15 +65,15 @@ def step(self, action):
 
     Repeat action, sum reward, and max over last observations.
     """
-    obs_list, total_reward, done = [], 0., False
+    obs_list, total_reward = [], 0.
     for _ in range(self._skip):
-      obs, reward, done, info = self.env.step(action)
+      obs, reward, terminated, truncated, info = self.env.step(action)
       obs_list.append(obs)
       total_reward += reward
-      if done:
+      if terminated or truncated:
         break
     max_frame = np.max(obs_list[-2:], axis=0)
-    return max_frame, total_reward, done, info
+    return max_frame, total_reward, terminated, truncated, info
 
 
 class EpisodicLifeEnv(gym.Wrapper):
@@ -90,32 +90,34 @@ def __init__(self, env):
     self.was_real_done = True
 
   def step(self, action):
-    obs, reward, done, info = self.env.step(action)
-    self.was_real_done = done
+    obs, reward, terminated, truncated, info = self.env.step(action)
+    self.was_real_done = terminated or truncated
     # check current lives, make loss of life terminal, then update lives to
     # handle bonus lives
     lives = self.env.unwrapped.ale.lives()
     if 0 < lives < self.lives:
       # for Qbert sometimes we stay in lives == 0 condition for a few
       # frames, so its important to keep lives > 0, so that we only reset
       # once the environment is actually done.
-      done = True
+      terminated = True
     self.lives = lives
-    return obs, reward, done, info
+    return obs, reward, terminated, truncated, info
 
-  def reset(self):
+  def reset(self, *, seed=None, options=None):
     """Calls the Gym environment reset, only when lives are exhausted.
 
     This way all states are still reachable even though lives are episodic,
     and the learner need not know about any of this behind-the-scenes.
     """
     if self.was_real_done:
-      obs = self.env.reset()
+      obs, info = self.env.reset(seed=seed, options=options)
     else:
       # no-op step to advance from terminal/lost life state
-      obs = self.env.step(0)[0]
+      obs, _, terminated, truncated, info = self.env.step(0)
+      if terminated or truncated:
+        obs, info = self.env.reset(seed=seed, options=options)
     self.lives = self.env.unwrapped.ale.lives()
-    return obs
+    return obs, info
 
 
 class FireResetEnv(gym.Wrapper):
@@ -131,9 +133,12 @@ def __init__(self, env):
     assert env.unwrapped.get_action_meanings()[1] == "FIRE"
     assert len(env.unwrapped.get_action_meanings()) >= 3
 
-  def reset(self):
-    self.env.reset()
-    return self.env.step(1)[0]
+  def reset(self, *, seed=None, options=None):
+    self.env.reset(seed=seed, options=options)
+    obs, _, terminated, truncated, info = self.env.step(1)
+    if terminated or truncated:
+      obs, info = self.env.reset()
+    return obs, info
 
 
 class WarpFrame(gym.ObservationWrapper):
@@ -214,16 +219,16 @@ def __init__(self, env, n_frames):
       dtype=env.observation_space.dtype
     )
 
-  def reset(self):
-    obs = self.env.reset()
+  def reset(self, *, seed=None, options=None):
+    obs, info = self.env.reset(seed=seed, options=options)
     for _ in range(self.n_frames):
       self.frames.append(obs)
-    return self._get_ob()
+    return self._get_ob(), info
 
   def step(self, action):
-    obs, reward, done, info = self.env.step(action)
+    obs, reward, terminated, truncated, info = self.env.step(action)
     self.frames.append(obs)
-    return self._get_ob(), reward, done, info
+    return self._get_ob(), reward, terminated, truncated, info
 
   def _get_ob(self):
     # the original wrapper use `LazyFrames` but since we use np buffer,
@@ -251,7 +256,11 @@ def wrap_deepmind(
   :param bool warp_frame: wrap the grayscale + resize observation wrapper.
   :return: the wrapped atari environment.
   """
-  assert "NoFrameskip" in env.spec.id
+  assert env.spec is not None
+  assert (
+    "NoFrameskip" in env.spec.id or
+    (env.spec.id.startswith("ALE/") and env.spec.id.endswith("-v5"))
+  )
   env = NoopResetEnv(env, noop_max=30)
   env = MaxAndSkipEnv(env, skip=4)
   if episode_life:

benchmark/requirements.txt

@@ -1,10 +1,8 @@
-gym[accept-rom-license]==0.23.1
-ale-py==0.7.5
+gymnasium[atari,box2d,mujoco]==1.2.3
+gym>=0.26.2
+ale-py==0.11.2
 mujoco==3.6.0
-envpool==0.6.1.post1
-sample-factory==1.123.0
-mujoco_py==2.1.2.14
 tqdm
-opencv-python-headless
+opencv-python-headless>=4.11.0.86,<4.13
 dm_control==1.0.38
 packaging

benchmark/test_envpool.py

@@ -13,6 +13,11 @@
 # limitations under the License.
 """EnvPool python benchmark script.
 
+This script expects an installed EnvPool build with native modules available.
+Baseline dependencies for the other benchmark scripts live in
+``benchmark/requirements.txt``; install the EnvPool wheel you want to
+benchmark separately before running this script.
+
 Single Python Thread
 ====================
 
@@ -56,18 +61,22 @@
     choices=["atari", "mujoco", "vizdoom", "box2d"],
   )
   parser.add_argument("--num-envs", type=int, default=645)
-  parser.add_argument("--batch-size", type=int, default=248)
+  parser.add_argument("--batch-size", type=int, default=None)
   # num_threads == 0 means to let envpool itself determine
   parser.add_argument("--num-threads", type=int, default=0)
   # thread_affinity_offset == -1 means no thread affinity
   parser.add_argument("--thread-affinity-offset", type=int, default=0)
   parser.add_argument("--total-step", type=int, default=50000)
   parser.add_argument("--seed", type=int, default=0)
   args = parser.parse_args()
+  if args.batch_size is None:
+    args.batch_size = min(248, args.num_envs)
+  elif args.batch_size > args.num_envs:
+    raise ValueError("--batch-size must be less than or equal to --num-envs")
   print(args)
   task_id = {
     "atari": "Pong-v5",
-    "mujoco": "Ant-v3",
+    "mujoco": "Ant-v5",
     "vizdoom": "HealthGathering-v1",
     "box2d": "LunarLander-v2",
   }[args.env]

benchmark/test_gym.py

@@ -15,57 +15,74 @@
 import argparse
 import time
 
-import gym
+import ale_py
+import gymnasium as gym
 import tqdm
 from atari_wrappers import wrap_deepmind
 
 
+def make_vector_env(num_envs, async_, make_env):
+  if async_:
+    vector_env_cls = gym.vector.AsyncVectorEnv
+  else:
+    vector_env_cls = gym.vector.SyncVectorEnv
+  return vector_env_cls([make_env for _ in range(num_envs)])
+
+
 def run(env, num_envs, total_step, async_):
   if env == "atari":
-    task_id = "PongNoFrameskip-v4"
+    gym.register_envs(ale_py)
+    task_id = "ALE/Pong-v5"
     frame_skip = 4
+    make_kwargs = {"frameskip": 1}
     if num_envs == 1:
       env = wrap_deepmind(
-        gym.make(task_id),
+        gym.make(task_id, **make_kwargs),
         episode_life=False,
         clip_rewards=False,
         frame_stack=4,
       )
     else:
-      env = gym.vector.make(
-        task_id, num_envs, async_, lambda e:
-        wrap_deepmind(e, episode_life=False, clip_rewards=False, frame_stack=4)
+      env = make_vector_env(
+        num_envs,
+        async_,
+        lambda: wrap_deepmind(
+          gym.make(task_id, **make_kwargs),
+          episode_life=False,
+          clip_rewards=False,
+          frame_stack=4,
+        ),
       )
   elif env == "mujoco":
-    task_id = "Ant-v3"
+    task_id = "Ant-v5"
     frame_skip = 5
     if num_envs == 1:
       env = gym.make(task_id)
     else:
-      env = gym.vector.make(task_id, num_envs, async_)
+      env = make_vector_env(num_envs, async_, lambda: gym.make(task_id))
   elif env == "box2d":
-    task_id = "LunarLander-v2"
+    task_id = "LunarLander-v3"
     frame_skip = 1
     if num_envs == 1:
       env = gym.make(task_id)
     else:
-      env = gym.vector.make(task_id, num_envs, async_)
+      env = make_vector_env(num_envs, async_, lambda: gym.make(task_id))
   else:
     raise NotImplementedError(f"Unknown env {env}")
-  env.seed(0)
-  env.reset()
+  env.reset(seed=0)
   action = env.action_space.sample()
-  done = False
+  terminated = truncated = False
   t = time.time()
   for _ in tqdm.trange(total_step):
     if num_envs == 1:
-      if done:
-        done = False
+      if terminated or truncated:
+        terminated = truncated = False
         env.reset()
       else:
-        done = env.step(action)[2]
+        _, _, terminated, truncated, _ = env.step(action)
     else:
       env.step(action)
+  env.close()
   print(f"FPS = {frame_skip * total_step * num_envs / (time.time() - t):.2f}")