Skip to content

Commit ad8734d

Browse files
troigantotroiganto
and
troiganto
authored
Fix outdated docs for TimeLimit max_episode_steps and add validation. (#1149)
Co-authored-by: troiganto <[email protected]>
1 parent 94384bb commit ad8734d

File tree

2 files changed

+25
-1
lines changed

2 files changed

+25
-1
lines changed

Diff for: gymnasium/wrappers/common.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,11 @@ def __init__(
9696
9797
Args:
9898
env: The environment to apply the wrapper
99-
max_episode_steps: An optional max episode steps (if ``None``, ``env.spec.max_episode_steps`` is used)
99+
max_episode_steps: the environment step after which the episode is truncated (``elapsed >= max_episode_steps``)
100100
"""
101+
assert (
102+
isinstance(max_episode_steps, int) and max_episode_steps > 0
103+
), f"Expect the `max_episode_steps` to be positive, actually: {max_episode_steps}"
101104
gym.utils.RecordConstructorArgs.__init__(
102105
self, max_episode_steps=max_episode_steps
103106
)

Diff for: tests/wrappers/test_time_limit.py

+21
Original file line numberDiff line numberDiff line change
@@ -57,3 +57,24 @@ def patched_step(_action):
5757
_, _, terminated, truncated, _ = env.step(env.action_space.sample())
5858
assert terminated is True
5959
assert truncated is True
60+
61+
62+
def test_max_episode_steps():
63+
env = gym.make("CartPole-v1", disable_env_checker=True)
64+
65+
assert env.spec.max_episode_steps == 500
66+
assert TimeLimit(env, max_episode_steps=10).spec.max_episode_steps == 10
67+
68+
with pytest.raises(
69+
AssertionError,
70+
match="Expect the `max_episode_steps` to be positive, actually: -1",
71+
):
72+
TimeLimit(env, max_episode_steps=-1)
73+
74+
with pytest.raises(
75+
AssertionError,
76+
match="Expect the `max_episode_steps` to be positive, actually: None",
77+
):
78+
TimeLimit(env, max_episode_steps=None)
79+
80+
env.close()

0 commit comments

Comments
 (0)