mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-20 14:02:03 +00:00
30 lines
1.2 KiB
Python
30 lines
1.2 KiB
Python
import gym
|
|
import warnings
|
|
|
|
|
|
class TimeLimit(gym.Wrapper):
|
|
def __init__(self, env, max_episode_steps=None):
|
|
super(TimeLimit, self).__init__(env)
|
|
warnings.warn("Gym\'s internal preprocessing wrappers are now deprecated. While they will continue to work for the foreseeable future, we strongly recommend using SuperSuit instead: https://github.com/PettingZoo-Team/SuperSuit")
|
|
if max_episode_steps is None and self.env.spec is not None:
|
|
max_episode_steps = env.spec.max_episode_steps
|
|
if self.env.spec is not None:
|
|
self.env.spec.max_episode_steps = max_episode_steps
|
|
self._max_episode_steps = max_episode_steps
|
|
self._elapsed_steps = None
|
|
|
|
def step(self, action):
|
|
assert (
|
|
self._elapsed_steps is not None
|
|
), "Cannot call env.step() before calling reset()"
|
|
observation, reward, done, info = self.env.step(action)
|
|
self._elapsed_steps += 1
|
|
if self._elapsed_steps >= self._max_episode_steps:
|
|
info["TimeLimit.truncated"] = not done
|
|
done = True
|
|
return observation, reward, done, info
|
|
|
|
def reset(self, **kwargs):
|
|
self._elapsed_steps = 0
|
|
return self.env.reset(**kwargs)
|