Skip to content

Commit

Permalink
Fix outdated docs for TimeLimit max_episode_steps and add validation.
Browse files Browse the repository at this point in the history
Closes #1147.
  • Loading branch information
troiganto committed Aug 26, 2024
1 parent 94384bb commit 7fc466f
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 1 deletion.
5 changes: 4 additions & 1 deletion gymnasium/wrappers/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,8 +96,11 @@ def __init__(
Args:
env: The environment to apply the wrapper
max_episode_steps: An optional max episode steps (if ``None``, ``env.spec.max_episode_steps`` is used)
max_episode_steps: the environment step after which the episode is truncated (``elapsed >= max_episode_steps``)
"""
assert (
isinstance(max_episode_steps, int) and max_episode_steps > 0
), f"Expect the `max_episode_steps` to be positive, actually: {max_episode_steps}"
gym.utils.RecordConstructorArgs.__init__(
self, max_episode_steps=max_episode_steps
)
Expand Down
21 changes: 21 additions & 0 deletions tests/wrappers/test_time_limit.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,3 +57,24 @@ def patched_step(_action):
_, _, terminated, truncated, _ = env.step(env.action_space.sample())
assert terminated is True
assert truncated is True


def test_max_episode_steps():
env = gym.make("CartPole-v1", disable_env_checker=True)

assert env.spec.max_episode_steps == 500
assert TimeLimit(env, max_episode_steps=10).spec.max_episode_steps == 10

with pytest.raises(
AssertionError,
match="Expect the `max_episode_steps` to be positive, actually: -1",
):
TimeLimit(env, max_episode_steps=-1)

with pytest.raises(
AssertionError,
match="Expect the `max_episode_steps` to be positive, actually: None",
):
TimeLimit(env, max_episode_steps=None)

env.close()

0 comments on commit 7fc466f

Please sign in to comment.