mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-19 05:25:54 +00:00
47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from gym import monitoring
|
|
from gym import Wrapper
|
|
from gym.wrappers.time_limit import TimeLimit
|
|
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class _Monitor(Wrapper):
|
|
def __init__(self, env, directory, video_callable=None, force=False, resume=False,
|
|
write_upon_reset=False, uid=None, mode=None):
|
|
super(_Monitor, self).__init__(env)
|
|
self._monitor = monitoring.MonitorManager(env)
|
|
self._monitor.start(directory, video_callable, force, resume,
|
|
write_upon_reset, uid, mode)
|
|
|
|
def _step(self, action):
|
|
self._monitor._before_step(action)
|
|
observation, reward, done, info = self.env.step(action)
|
|
done = self._monitor._after_step(observation, reward, done, info)
|
|
|
|
return observation, reward, done, info
|
|
|
|
def _reset(self):
|
|
self._monitor._before_reset()
|
|
observation = self.env.reset()
|
|
self._monitor._after_reset(observation)
|
|
|
|
return observation
|
|
|
|
def _close(self):
|
|
super(_Monitor, self)._close()
|
|
|
|
# _monitor will not be set if super(Monitor, self).__init__ raises, this check prevents a confusing error message
|
|
if getattr(self, '_monitor', None):
|
|
self._monitor.close()
|
|
|
|
def set_monitor_mode(self, mode):
|
|
logger.info("Setting the monitor mode is deprecated and will be removed soon")
|
|
self._monitor._set_mode(mode)
|
|
|
|
def Monitor(env, directory, video_callable=None, force=False, resume=False,
|
|
write_upon_reset=False, uid=None, mode=None):
|
|
|
|
return _Monitor(TimeLimit(env), directory, video_callable, force, resume,
|
|
write_upon_reset, uid, mode)
|