fix(core): avoid KeyError when batch_size is given without num_envs

envpool/make_test.py

@@ -282,6 +282,18 @@ def test_make_atari(self) -> None:
             env_gym.close()
             env_gymnasium.close()
 
+    def test_make_batch_size_without_num_envs(self) -> None:
+        # Regression: passing batch_size without num_envs used to raise
+        # KeyError on the bare kwargs["num_envs"] subscript in
+        # EnvRegistry._make_env_spec. batch_size=0 means "default to
+        # num_envs" (see core/env_spec.h) and must be accepted; an
+        # out-of-range batch_size must raise AssertionError, not KeyError.
+        env = envpool.make_gymnasium("CartPole-v1", batch_size=0)
+        env.close()
+        self.assertRaises(
+            AssertionError, envpool.make_gymnasium, "CartPole-v1", batch_size=2
+        )
+
     def test_make_vizdoom(self) -> None:
         try:
             with _temporary_workdir(prefix="envpool-vizdoom-smoke-"):

envpool/registration.py

@@ -235,7 +235,7 @@ def _make_env_spec(
         if "num_envs" in kwargs:
             assert kwargs["num_envs"] >= 1
         if "batch_size" in kwargs:
-            assert 0 <= kwargs["batch_size"] <= kwargs["num_envs"]
+            assert 0 <= kwargs["batch_size"] <= kwargs.get("num_envs", 1)
         if "max_num_players" in kwargs:
             assert 1 <= kwargs["max_num_players"]