87 lines
3.8 KiB
Python
87 lines
3.8 KiB
Python
"""Wrapper that autoreset environments when `terminated=True` or `truncated=True`."""
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from typing import TYPE_CHECKING
|
|
|
|
import gymnasium as gym
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from gymnasium.envs.registration import EnvSpec
|
|
|
|
|
|
class AutoResetWrapper(gym.Wrapper, gym.utils.RecordConstructorArgs):
|
|
"""A class for providing an automatic reset functionality for gymnasium environments when calling :meth:`self.step`.
|
|
|
|
When calling step causes :meth:`Env.step` to return `terminated=True` or `truncated=True`, :meth:`Env.reset` is called,
|
|
and the return format of :meth:`self.step` is as follows: ``(new_obs, final_reward, final_terminated, final_truncated, info)``
|
|
with new step API and ``(new_obs, final_reward, final_done, info)`` with the old step API.
|
|
|
|
- ``new_obs`` is the first observation after calling :meth:`self.env.reset`
|
|
- ``final_reward`` is the reward after calling :meth:`self.env.step`, prior to calling :meth:`self.env.reset`.
|
|
- ``final_terminated`` is the terminated value before calling :meth:`self.env.reset`.
|
|
- ``final_truncated`` is the truncated value before calling :meth:`self.env.reset`. Both `final_terminated` and `final_truncated` cannot be False.
|
|
- ``info`` is a dict containing all the keys from the info dict returned by the call to :meth:`self.env.reset`,
|
|
with an additional key "final_observation" containing the observation returned by the last call to :meth:`self.env.step`
|
|
and "final_info" containing the info dict returned by the last call to :meth:`self.env.step`.
|
|
|
|
Warning:
|
|
When using this wrapper to collect rollouts, note that when :meth:`Env.step` returns `terminated` or `truncated`, a
|
|
new observation from after calling :meth:`Env.reset` is returned by :meth:`Env.step` alongside the
|
|
final reward, terminated and truncated state from the previous episode.
|
|
If you need the final state from the previous episode, you need to retrieve it via the
|
|
"final_observation" key in the info dict.
|
|
Make sure you know what you're doing if you use this wrapper!
|
|
"""
|
|
|
|
def __init__(self, env: gym.Env):
|
|
"""A class for providing an automatic reset functionality for gymnasium environments when calling :meth:`self.step`.
|
|
|
|
Args:
|
|
env (gym.Env): The environment to apply the wrapper
|
|
"""
|
|
gym.utils.RecordConstructorArgs.__init__(self)
|
|
gym.Wrapper.__init__(self, env)
|
|
|
|
def step(self, action):
|
|
"""Steps through the environment with action and resets the environment if a terminated or truncated signal is encountered.
|
|
|
|
Args:
|
|
action: The action to take
|
|
|
|
Returns:
|
|
The autoreset environment :meth:`step`
|
|
"""
|
|
obs, reward, terminated, truncated, info = self.env.step(action)
|
|
if terminated or truncated:
|
|
new_obs, new_info = self.env.reset()
|
|
assert (
|
|
"final_observation" not in new_info
|
|
), 'info dict cannot contain key "final_observation" '
|
|
assert (
|
|
"final_info" not in new_info
|
|
), 'info dict cannot contain key "final_info" '
|
|
|
|
new_info["final_observation"] = obs
|
|
new_info["final_info"] = info
|
|
|
|
obs = new_obs
|
|
info = new_info
|
|
|
|
return obs, reward, terminated, truncated, info
|
|
|
|
@property
|
|
def spec(self) -> EnvSpec | None:
|
|
"""Modifies the environment spec to specify the `autoreset=True`."""
|
|
if self._cached_spec is not None:
|
|
return self._cached_spec
|
|
|
|
env_spec = self.env.spec
|
|
if env_spec is not None:
|
|
env_spec = deepcopy(env_spec)
|
|
env_spec.autoreset = True
|
|
|
|
self._cached_spec = env_spec
|
|
return env_spec
|