Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add make(max_episode_steps=0) to not apply a TimeLimit wrapper #710

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions gymnasium/envs/registration.py
Original file line number Diff line number Diff line change
Expand Up @@ -645,8 +645,9 @@ def make(
Args:
id: A string for the environment id or a :class:`EnvSpec`. Optionally if using a string, a module to import can be included, e.g. ``'module:Env-v0'``.
This is equivalent to importing the module first to register the environment followed by making the environment.
max_episode_steps: Maximum length of an episode, can override the registered :class:`EnvSpec` ``max_episode_steps``.
The value is used by :class:`gymnasium.wrappers.TimeLimit`.
max_episode_steps: Maximum length of an episode, can override the registered :class:`EnvSpec` ``max_episode_steps``
with the value being passed to :class:`gymnasium.wrappers.TimeLimit`.
Using ``max_episode_steps=0`` will not apply the wrapper to the environment.
disable_env_checker: If to add :class:`gymnasium.wrappers.PassiveEnvChecker`, ``None`` will default to the
:class:`EnvSpec` ``disable_env_checker`` value otherwise use this value will be used.
kwargs: Additional arguments to pass to the environment constructor.
Expand Down Expand Up @@ -780,10 +781,11 @@ def make(
env = gym.wrappers.OrderEnforcing(env)

# Add the time limit wrapper
if max_episode_steps is not None:
env = gym.wrappers.TimeLimit(env, max_episode_steps)
elif env_spec.max_episode_steps is not None:
env = gym.wrappers.TimeLimit(env, env_spec.max_episode_steps)
if max_episode_steps != 0:
if max_episode_steps is not None:
env = gym.wrappers.TimeLimit(env, max_episode_steps)
elif env_spec.max_episode_steps is not None:
env = gym.wrappers.TimeLimit(env, env_spec.max_episode_steps)

for wrapper_spec in env_spec.additional_wrappers[num_prior_wrappers:]:
if wrapper_spec.kwargs is None:
Expand Down
12 changes: 12 additions & 0 deletions tests/envs/registration/test_make.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,18 @@ def test_max_episode_steps(register_parameter_envs):
assert env.spec.max_episode_steps == 100
assert has_wrapper(env, TimeLimit)

# Override max_episode_step to prevent applying the wrapper
for env_id in [
"CartPole-v1",
gym.spec("CartPole-v1"),
"NoMaxEpisodeStepsEnv-v0",
gym.spec("NoMaxEpisodeStepsEnv-v0"),
]:
env = gym.make(env_id, max_episode_steps=0)
assert env.spec is not None
assert env.spec.max_episode_steps is None
assert has_wrapper(env, TimeLimit) is False


@pytest.mark.parametrize(
"registration_disabled, make_disabled, if_disabled",
Expand Down