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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159 | from typing import Any, Generic, Optional, TypeVar, Union
import gymnasium as gym
from gymnasium import spaces
import numpy as np
from lit_sb3.core.types import GymStepReturn
T = TypeVar("T", int, np.ndarray)
class IdentityEnv[T: (int, np.ndarray)](gym.Env):
def __init__(self, dim: int | None = None, space: spaces.Space | None = None, ep_length: int = 100) -> None:
"""Identity environment for testing purposes.
Args:
dim: the size of the action and observation dimension you want
to learn. Provide at most one of ``dim`` and ``space``. If both are
None, then initialization proceeds with ``dim=1`` and ``space=None``.
space: the action and observation space. Provide at most one of
``dim`` and ``space``.
ep_length: the length of each episode in timesteps
"""
if space is None:
if dim is None:
dim = 1
space = spaces.Discrete(dim)
else:
assert dim is None, "arguments for both 'dim' and 'space' provided: at most one allowed"
self.action_space = self.observation_space = space
self.ep_length = ep_length
self.current_step = 0
self.num_resets = -1 # Becomes 0 after __init__ exits.
self.reset()
def reset(self, *, seed: int | None = None, options: dict | None = None) -> tuple[T, dict]:
if seed is not None:
super().reset(seed=seed)
self.current_step = 0
self.num_resets += 1
self._choose_next_state()
return self.state, {}
def step(self, action: T) -> tuple[T, float, bool, bool, dict[str, Any]]:
reward = self._get_reward(action)
self._choose_next_state()
self.current_step += 1
terminated = False
truncated = self.current_step >= self.ep_length
return self.state, reward, terminated, truncated, {}
def _choose_next_state(self) -> None:
self.state = self.action_space.sample()
def _get_reward(self, action: T) -> float:
return 1.0 if np.all(self.state == action) else 0.0
def render(self, mode: str = "human") -> None:
pass
class IdentityEnvBox(IdentityEnv[np.ndarray]):
def __init__(self, low: float = -1.0, high: float = 1.0, eps: float = 0.05, ep_length: int = 100) -> None:
"""Identity environment for testing purposes.
Args:
low: the lower bound of the box dim
high: the upper bound of the box dim
eps: the epsilon bound for correct value
ep_length: the length of each episode in timesteps
"""
space = spaces.Box(low=low, high=high, shape=(1,), dtype=np.float32)
super().__init__(ep_length=ep_length, space=space)
self.eps = eps
def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]:
reward = self._get_reward(action)
self._choose_next_state()
self.current_step += 1
terminated = False
truncated = self.current_step >= self.ep_length
return self.state, reward, terminated, truncated, {}
def _get_reward(self, action: np.ndarray) -> float:
return 1.0 if (self.state - self.eps) <= action <= (self.state + self.eps) else 0.0
class IdentityEnvMultiDiscrete(IdentityEnv[np.ndarray]):
def __init__(self, dim: int = 1, ep_length: int = 100) -> None:
"""Identity environment for testing purposes.
Args:
dim: the size of the dimensions you want to learn
ep_length: the length of each episode in timesteps
"""
space = spaces.MultiDiscrete([dim, dim])
super().__init__(ep_length=ep_length, space=space)
class IdentityEnvMultiBinary(IdentityEnv[np.ndarray]):
def __init__(self, dim: int = 1, ep_length: int = 100) -> None:
"""Identity environment for testing purposes.
Args:
dim: the size of the dimensions you want to learn
ep_length: the length of each episode in timesteps
"""
space = spaces.MultiBinary(dim)
super().__init__(ep_length=ep_length, space=space)
class FakeImageEnv(gym.Env):
"""Fake image environment for testing purposes, it mimics Atari games.
Args:
action_dim: Number of discrete actions
screen_height: Height of the image
screen_width: Width of the image
n_channels: Number of color channels
discrete: Create discrete action space instead of continuous
channel_first: Put channels on first axis instead of last
"""
def __init__(
self,
action_dim: int = 6,
screen_height: int = 84,
screen_width: int = 84,
n_channels: int = 1,
discrete: bool = True,
channel_first: bool = False,
) -> None:
self.observation_shape = (screen_height, screen_width, n_channels)
if channel_first:
self.observation_shape = (n_channels, screen_height, screen_width)
self.observation_space = spaces.Box(low=0, high=255, shape=self.observation_shape, dtype=np.uint8)
if discrete:
self.action_space = spaces.Discrete(action_dim)
else:
self.action_space = spaces.Box(low=-1, high=1, shape=(5,), dtype=np.float32)
self.ep_length = 10
self.current_step = 0
def reset(self, *, seed: int | None = None, options: dict | None = None) -> tuple[np.ndarray, dict]:
if seed is not None:
super().reset(seed=seed)
self.current_step = 0
return self.observation_space.sample(), {}
def step(self, action: np.ndarray | int) -> GymStepReturn:
reward = 0.0
self.current_step += 1
terminated = False
truncated = self.current_step >= self.ep_length
return self.observation_space.sample(), reward, terminated, truncated, {}
def render(self, mode: str = "human") -> None:
pass
|