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 | from typing import Optional
from gymnasium import spaces
import numpy as np
from lit_sb3.envs.identity_env import IdentityEnv
class InvalidActionEnvDiscrete(IdentityEnv[int]):
"""Identity env with a discrete action space. Supports action masking."""
def __init__(
self,
dim: int | None = None,
ep_length: int = 100,
n_invalid_actions: int = 0,
) -> None:
if dim is None:
dim = 1
assert n_invalid_actions < dim, f"Too many invalid actions: {n_invalid_actions} < {dim}"
space = spaces.Discrete(dim)
self.n_invalid_actions = n_invalid_actions
self.possible_actions = np.arange(space.n, dtype=int)
self.invalid_actions: list[int] = []
super().__init__(space=space, ep_length=ep_length)
def _choose_next_state(self) -> None:
self.state = self.action_space.sample()
# Randomly choose invalid actions that are not the current state
potential_invalid_actions = [i for i in self.possible_actions if i != self.state]
self.invalid_actions = np.random.choice(
potential_invalid_actions, self.n_invalid_actions, replace=False
).tolist()
def action_masks(self) -> list[bool]:
return [action not in self.invalid_actions for action in self.possible_actions]
class InvalidActionEnvMultiDiscrete(IdentityEnv[np.ndarray]):
"""Identity env with a multidiscrete action space. Supports action masking."""
action_space: spaces.MultiDiscrete
def __init__(
self,
dims: list[int] | None = None,
ep_length: int = 100,
n_invalid_actions: int = 0,
) -> None:
if dims is None:
dims = [1, 1]
if n_invalid_actions > sum(dims) - len(dims):
raise ValueError(
f"Cannot find a valid action for each dim. Set n_invalid_actions <= {sum(dims) - len(dims)}"
)
space = spaces.MultiDiscrete(dims)
self.n_invalid_actions = n_invalid_actions
self.possible_actions = np.arange(sum(dims))
self.invalid_actions: list[int] = []
super().__init__(space=space, ep_length=ep_length)
def _choose_next_state(self) -> None:
self.state = self.action_space.sample()
converted_state: list[int] = []
running_total = 0
for i in range(len(self.action_space.nvec)):
converted_state.append(running_total + self.state[i])
running_total += self.action_space.nvec[i]
# Randomly choose invalid actions that are not the current state
potential_invalid_actions = [i for i in self.possible_actions if i not in converted_state]
self.invalid_actions = np.random.choice(
potential_invalid_actions, self.n_invalid_actions, replace=False
).tolist()
def action_masks(self) -> list[bool]:
return [action not in self.invalid_actions for action in self.possible_actions]
class InvalidActionEnvMultiBinary(IdentityEnv[np.ndarray]):
"""Identity env with a multibinary action space. Supports action masking."""
def __init__(
self,
dims: int | None = None,
ep_length: int = 100,
n_invalid_actions: int = 0,
) -> None:
if dims is None:
dims = 1
if n_invalid_actions > dims:
raise ValueError(f"Cannot find a valid action for each dim. Set n_invalid_actions <= {dims}")
space = spaces.MultiBinary(dims)
self.n_dims = dims
self.n_invalid_actions = n_invalid_actions
self.possible_actions = np.arange(2 * dims)
self.invalid_actions: list[int] = []
super().__init__(space=space, ep_length=ep_length)
def _choose_next_state(self) -> None:
self.state = self.action_space.sample()
converted_state: list[int] = []
running_total = 0
for i in range(self.n_dims):
converted_state.append(running_total + self.state[i])
running_total += 2
# Randomly choose invalid actions that are not the current state
potential_invalid_actions = [i for i in self.possible_actions if i not in converted_state]
self.invalid_actions = np.random.choice(
potential_invalid_actions, self.n_invalid_actions, replace=False
).tolist()
def action_masks(self) -> list[bool]:
return [action not in self.invalid_actions for action in self.possible_actions]
|