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 | from typing import Any, Optional
from gymnasium import spaces
import torch
from torch import nn
from lit_sb3.core.preprocessing import get_action_dim
from lit_sb3.core.types import PyTorchObs
from lit_sb3.modules import create_mlp
from lit_sb3.policies import BasePolicy
class ARSPolicy(BasePolicy):
"""Policy network for ARS.
Args:
observation_space: The observation space of the environment
action_space: The action space of the environment
net_arch: Network architecture, defaults to a 2 layers MLP with 64 hidden nodes.
activation_fn: Activation function
with_bias: If set to False, the layers will not learn an additive bias
squash_output: For continuous actions, whether the output is squashed
or not using a ``tanh()`` function. If not squashed with tanh the output will instead be clipped.
"""
def __init__(
self,
observation_space: spaces.Space,
action_space: spaces.Space,
net_arch: list[int] | None = None,
activation_fn: type[nn.Module] = nn.ReLU,
with_bias: bool = True,
squash_output: bool = True,
) -> None:
super().__init__(
observation_space,
action_space,
squash_output=isinstance(action_space, spaces.Box) and squash_output,
)
if net_arch is None:
net_arch = [64, 64]
self.net_arch = net_arch
self.features_extractor = self.make_features_extractor()
self.features_dim = self.features_extractor.features_dim
self.activation_fn = activation_fn
if isinstance(action_space, spaces.Box):
action_dim = get_action_dim(action_space)
actor_net = create_mlp(
self.features_dim, action_dim, net_arch, activation_fn, with_bias=with_bias, squash_output=squash_output
)
elif isinstance(action_space, spaces.Discrete):
actor_net = create_mlp(self.features_dim, int(action_space.n), net_arch, activation_fn, with_bias=with_bias)
else:
raise NotImplementedError(
f"Error: ARS policy not implemented for action space of type {type(action_space)}."
)
self.action_net = nn.Sequential(*actor_net)
def _get_constructor_parameters(self) -> dict[str, Any]:
# data = super()._get_constructor_parameters() this adds normalize_images, which we don't support...
data = {
"observation_space": self.observation_space,
"action_space": self.action_space,
"net_arch": self.net_arch,
"activation_fn": self.activation_fn,
}
return data
def forward(self, obs: PyTorchObs) -> torch.Tensor:
features = self.extract_features(obs, self.features_extractor)
if isinstance(self.action_space, spaces.Box):
return self.action_net(features)
elif isinstance(self.action_space, spaces.Discrete):
logits = self.action_net(features)
return torch.argmax(logits, dim=1)
else:
raise NotImplementedError()
def _predict(self, observation: PyTorchObs, deterministic: bool = True) -> torch.Tensor:
# Non deterministic action does not really make sense for ARS, we ignore this parameter for now..
return self(observation)
class ARSLinearPolicy(ARSPolicy):
"""Linear policy network for ARS.
Args:
observation_space: The observation space of the environment
action_space: The action space of the environment
with_bias: With or without bias on the output
squash_output: For continuous actions, whether the output is squashed
or not using a ``tanh()`` function. If not squashed with tanh the output will instead be clipped.
"""
def __init__(
self,
observation_space: spaces.Space,
action_space: spaces.Space,
with_bias: bool = False,
squash_output: bool = False,
) -> None:
super().__init__(observation_space, action_space, net_arch=[], with_bias=with_bias, squash_output=squash_output)
MlpPolicy = ARSPolicy
LinearPolicy = ARSLinearPolicy
|