1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104 | from typing import Any, SupportsFloat
import gymnasium as gym
from gymnasium import spaces
from gymnasium.core import ActType
import numpy as np
TimeFeatureObs = np.ndarray | dict[str, np.ndarray]
class TimeFeatureWrapper(gym.Wrapper[TimeFeatureObs, ActType, TimeFeatureObs, ActType]):
"""Add remaining, normalized time to observation space for fixed length episodes.
See https://arxiv.org/abs/1712.00378 and https://github.com/aravindr93/mjrl/issues/13.
.. note::
Only ``gym.spaces.Box`` and ``gym.spaces.Dict`` (``gym.GoalEnv``) 1D observation spaces
are supported for now.
Args:
env: Gym env to wrap.
max_steps: Max number of steps of an episode
if it is not wrapped in a ``TimeLimit`` object.
test_mode: In test mode, the time feature is constant,
equal to zero. This allow to check that the agent did not overfit this feature,
learning a deterministic pre-defined sequence of actions.
"""
def __init__(self, env: gym.Env, max_steps: int = 1000, test_mode: bool = False) -> None:
assert isinstance(env.observation_space, (spaces.Box, spaces.Dict)), (
"`TimeFeatureWrapper` only supports `gym.spaces.Box` and `spaces.Dict` (`gym.GoalEnv`) observation spaces."
)
# Add a time feature to the observation
if isinstance(env.observation_space, spaces.Dict):
assert "observation" in env.observation_space.spaces, "No `observation` key in the observation space"
obs_space = env.observation_space.spaces["observation"]
else:
obs_space = env.observation_space
assert isinstance(obs_space, gym.spaces.Box), (
"`TimeFeatureWrapper` only supports `gym.spaces.Box` observation space."
)
assert len(obs_space.shape) == 1, "Only 1D observation spaces are supported"
low, high = obs_space.low, obs_space.high
low, high = np.concatenate((low, [0.0])), np.concatenate((high, [1.0]))
self.dtype = obs_space.dtype
low, high = low.astype(self.dtype), high.astype(self.dtype)
if isinstance(env.observation_space, spaces.Dict):
env.observation_space.spaces["observation"] = spaces.Box(
low=low,
high=high,
dtype=self.dtype, # type: ignore[arg-type]
)
else:
env.observation_space = spaces.Box(low=low, high=high, dtype=self.dtype) # type: ignore[arg-type]
super().__init__(env)
# Try to infer the max number of steps per episode
try:
self._max_steps = env.spec.max_episode_steps # type: ignore[union-attr]
except AttributeError:
self._max_steps = None
# Fallback to provided value
if self._max_steps is None:
self._max_steps = max_steps
self._current_step = 0
self._test_mode = test_mode
def reset(self, **kwargs) -> tuple[TimeFeatureObs, dict[str, Any]]:
self._current_step = 0
obs, info = self.env.reset(**kwargs)
return self._get_obs(obs), info
def step(self, action: ActType) -> tuple[TimeFeatureObs, SupportsFloat, bool, bool, dict[str, Any]]:
self._current_step += 1
obs, reward, done, truncated, info = self.env.step(action)
return self._get_obs(obs), reward, done, truncated, info
def _get_obs(self, obs: np.ndarray | dict[str, np.ndarray]) -> np.ndarray | dict[str, np.ndarray]:
"""Concatenate the time feature to the current observation.
Args:
obs:
obs: The current observation.
"""
# for mypy
assert self._max_steps is not None
# Remaining time is more general
time_feature = 1 - (self._current_step / self._max_steps)
if self._test_mode:
time_feature = 1.0
time_feature = np.array(time_feature, dtype=self.dtype)
if isinstance(obs, dict):
obs["observation"] = np.append(obs["observation"], time_feature)
return obs
return np.append(obs, time_feature)
|