Skip to content

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

Merged
Trinkle23897 merged 1 commit into
sail-sg:mainfrom
hobostay:fix/make-batch-size-without-num-envs
Jun 17, 2026
Merged

fix(core): avoid KeyError when batch_size is given without num_envs#423
Trinkle23897 merged 1 commit into
sail-sg:mainfrom
hobostay:fix/make-batch-size-without-num-envs

Conversation

@hobostay

Copy link
Copy Markdown
Contributor

Problem

envpool.make(...) raises a KeyError: 'num_envs' whenever batch_size is supplied without an explicit num_envs.

In EnvRegistry._make_env_spec (envpool/registration.py), the batch_size validation indexed kwargs["num_envs"] directly:

if "batch_size" in kwargs:
    assert 0 <= kwargs["batch_size"] <= kwargs["num_envs"]   # KeyError

…while the two seed-normalization blocks immediately above already use the safe form kwargs.get("num_envs", 1). num_envs is never injected into kwargs — it only appears there if the user passes it — so the bare subscript crashes before gen_config is ever reached.

This is reachable through the documented public API:

  • batch_size is a first-class make() argument (README / docs/content/python_interface.rst).
  • batch_size=0 is a valid value meaning "default to num_envs" (core/env_spec.h turns batch_size == 0 into num_envs).

Reproduction

import envpool
envpool.make("Pong-v5", env_type="gymnasium", batch_size=0)
# KeyError: 'num_envs'   <- before this PR

Fix

Use the same default as the neighboring validation (one-line change, matches the existing kwargs.get("num_envs", 1) pattern):

if "batch_size" in kwargs:
    assert 0 <= kwargs["batch_size"] <= kwargs.get("num_envs", 1)

Behavior after the fix:

call before after
batch_size=0 (no num_envs) KeyError OK (defaults to num_envs)
batch_size=1 (no num_envs) KeyError OK
batch_size=2 (no num_envs, out of range) KeyError AssertionError (correctly rejected)
num_envs=64, batch_size=16 (README usage) OK OK (unchanged)

The documented multi-env usage is unaffected; only the spurious KeyError is removed and out-of-range values are now rejected with a proper AssertionError.

Test

Added test_make_batch_size_without_num_envs to envpool/make_test.py, which fails (KeyError) before the fix and passes after.

Validation

make ruff py-format   # pass
make mypy             # changed files clean (under py311 target)

In `EnvRegistry._make_env_spec`, the `batch_size` validation indexed
`kwargs["num_envs"]` directly, while the neighboring seed-normalization
blocks use `kwargs.get("num_envs", 1)`. When a user passes `batch_size`
without an explicit `num_envs`, the validation raised
`KeyError: 'num_envs'` before reaching `gen_config`.

This is reachable through the documented public API: `batch_size` is a
first-class make() argument (see README and docs/content/python_interface.rst),
and `batch_size=0` is a valid value meaning "default to num_envs"
(core/env_spec.h turns `batch_size == 0` into `num_envs`).

Use `kwargs.get("num_envs", 1)` to match the surrounding validation, so
valid values are accepted and out-of-range `batch_size` is still rejected
as an `AssertionError` rather than an opaque `KeyError`.

Co-Authored-By: Claude <noreply@anthropic.com>
@Trinkle23897 Trinkle23897 merged commit 3446f28 into sail-sg:main Jun 17, 2026
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants