fix(core): avoid KeyError when batch_size is given without num_envs#423
Merged
Trinkle23897 merged 1 commit intoJun 17, 2026
Merged
Conversation
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
approved these changes
Jun 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
envpool.make(...)raises aKeyError: 'num_envs'wheneverbatch_sizeis supplied without an explicitnum_envs.In
EnvRegistry._make_env_spec(envpool/registration.py), thebatch_sizevalidation indexedkwargs["num_envs"]directly:…while the two seed-normalization blocks immediately above already use the safe form
kwargs.get("num_envs", 1).num_envsis never injected intokwargs— it only appears there if the user passes it — so the bare subscript crashes beforegen_configis ever reached.This is reachable through the documented public API:
batch_sizeis a first-classmake()argument (README /docs/content/python_interface.rst).batch_size=0is a valid value meaning "default tonum_envs" (core/env_spec.hturnsbatch_size == 0intonum_envs).Reproduction
Fix
Use the same default as the neighboring validation (one-line change, matches the existing
kwargs.get("num_envs", 1)pattern):Behavior after the fix:
batch_size=0(nonum_envs)KeyError✗num_envs)batch_size=1(nonum_envs)KeyError✗batch_size=2(nonum_envs, out of range)KeyError✗AssertionError(correctly rejected)num_envs=64, batch_size=16(README usage)The documented multi-env usage is unaffected; only the spurious
KeyErroris removed and out-of-range values are now rejected with a properAssertionError.Test
Added
test_make_batch_size_without_num_envstoenvpool/make_test.py, which fails (KeyError) before the fix and passes after.Validation