drlab 0.2.3__tar.gz → 0.2.5__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {drlab-0.2.3/src/drlab.egg-info → drlab-0.2.5}/PKG-INFO +1 -1
- {drlab-0.2.3 → drlab-0.2.5}/pyproject.toml +1 -1
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/__init__.py +1 -1
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/base.py +58 -19
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/configs.py +6 -3
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/dqn.py +2 -1
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/sac.py +46 -34
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/base.py +1 -1
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/configs.py +1 -1
- {drlab-0.2.3 → drlab-0.2.5/src/drlab.egg-info}/PKG-INFO +1 -1
- {drlab-0.2.3 → drlab-0.2.5}/tests/test_learners.py +171 -20
- {drlab-0.2.3 → drlab-0.2.5}/tests/test_public_api.py +1 -1
- {drlab-0.2.3 → drlab-0.2.5}/LICENSE +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/README.md +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/setup.cfg +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/base.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/e_greedy.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/gaussian.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/greedy.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/stochastic.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/warmup.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/experiments/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/experiments/off_policy_experiment.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/experiments/on_policy_experiment.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/actor_critic.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/ppo.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/reinforce.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/replay/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/replay/replay_buffer.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/replay/transition_batch.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/runners/__init__.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab/runners/runner.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/SOURCES.txt +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/dependency_links.txt +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/requires.txt +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/top_level.txt +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/tests/test_controllers.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/tests/test_off_policy_experiment.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/tests/test_replay.py +0 -0
- {drlab-0.2.3 → drlab-0.2.5}/tests/test_runner.py +0 -0
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import torch as th
|
|
2
2
|
from copy import deepcopy
|
|
3
3
|
from abc import ABC, abstractmethod
|
|
4
|
+
from collections.abc import Iterable
|
|
4
5
|
|
|
5
6
|
from .configs import OffPolicyConfig, TargetUpdate
|
|
6
7
|
|
|
@@ -15,8 +16,7 @@ class OffPolicyLearner(ABC):
|
|
|
15
16
|
self.regularizers = config.regularizers
|
|
16
17
|
self.reg_lams = config.reg_lams
|
|
17
18
|
|
|
18
|
-
self.
|
|
19
|
-
self.grad_norm_clip = config.grad_norm_clip
|
|
19
|
+
self.clipnorm = config.clipnorm
|
|
20
20
|
|
|
21
21
|
self.use_target_model = config.use_target_model
|
|
22
22
|
self.target_update = TargetUpdate(config.target_update)
|
|
@@ -52,25 +52,59 @@ class OffPolicyLearner(ABC):
|
|
|
52
52
|
|
|
53
53
|
return target
|
|
54
54
|
|
|
55
|
-
|
|
55
|
+
@staticmethod
|
|
56
|
+
def _validate_optimizer_parameters(
|
|
57
|
+
optimizer: th.optim.Optimizer,
|
|
58
|
+
expected_parameters: Iterable[th.nn.Parameter],
|
|
59
|
+
optimizer_name: str,
|
|
60
|
+
) -> None:
|
|
61
|
+
expected_parameter_ids = sorted(
|
|
62
|
+
id(parameter) for parameter in expected_parameters
|
|
63
|
+
)
|
|
64
|
+
optimizer_parameter_ids = sorted(
|
|
65
|
+
id(parameter)
|
|
66
|
+
for group in optimizer.param_groups
|
|
67
|
+
for parameter in group["params"]
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
if optimizer_parameter_ids != expected_parameter_ids:
|
|
71
|
+
raise ValueError(
|
|
72
|
+
f"{optimizer_name} must contain exactly its model parameters."
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
def _hard_update(self, target: th.nn.Module, source: th.nn.Module):
|
|
56
76
|
target.load_state_dict(source.state_dict())
|
|
57
77
|
target.eval()
|
|
58
78
|
|
|
59
|
-
def
|
|
60
|
-
|
|
79
|
+
def _soft_update(
|
|
80
|
+
self,
|
|
81
|
+
pairs: list[tuple[th.nn.Module, th.nn.Module]],
|
|
82
|
+
) -> None:
|
|
83
|
+
target_parameters = [
|
|
84
|
+
parameter
|
|
85
|
+
for target, _ in pairs
|
|
86
|
+
for parameter in target.parameters()
|
|
87
|
+
]
|
|
88
|
+
source_parameters = [
|
|
89
|
+
parameter
|
|
90
|
+
for _, source in pairs
|
|
91
|
+
for parameter in source.parameters()
|
|
92
|
+
]
|
|
61
93
|
|
|
62
94
|
with th.no_grad():
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
95
|
+
th._foreach_mul_(target_parameters, 1.0 - self.soft_target_update_param)
|
|
96
|
+
scaled_sources = th._foreach_mul(
|
|
97
|
+
source_parameters,
|
|
98
|
+
self.soft_target_update_param,
|
|
99
|
+
)
|
|
100
|
+
th._foreach_add_(target_parameters, scaled_sources)
|
|
66
101
|
|
|
67
102
|
def update_targets(
|
|
68
103
|
self,
|
|
69
104
|
pairs: list[tuple[th.nn.Module, th.nn.Module]],
|
|
70
105
|
):
|
|
71
106
|
if self.target_update == TargetUpdate.SOFT:
|
|
72
|
-
|
|
73
|
-
self.soft_update(target, source)
|
|
107
|
+
self._soft_update(pairs)
|
|
74
108
|
return
|
|
75
109
|
|
|
76
110
|
if self.target_update == TargetUpdate.HARD:
|
|
@@ -78,7 +112,7 @@ class OffPolicyLearner(ABC):
|
|
|
78
112
|
|
|
79
113
|
if self.target_update_calls % self.target_update_interval == 0:
|
|
80
114
|
for target, source in pairs:
|
|
81
|
-
self.
|
|
115
|
+
self._hard_update(target, source)
|
|
82
116
|
|
|
83
117
|
return
|
|
84
118
|
|
|
@@ -119,16 +153,21 @@ class OffPolicyLearner(ABC):
|
|
|
119
153
|
self,
|
|
120
154
|
loss: th.Tensor,
|
|
121
155
|
optimizer: th.optim.Optimizer,
|
|
122
|
-
parameters: list[th.nn.Parameter]
|
|
123
|
-
|
|
156
|
+
parameters: list[Iterable[th.nn.Parameter]],
|
|
157
|
+
clipnorm: float | None = None,
|
|
158
|
+
) -> None:
|
|
159
|
+
parameter_groups = [tuple(group) for group in parameters]
|
|
160
|
+
parameters = [
|
|
161
|
+
parameter
|
|
162
|
+
for group in parameter_groups
|
|
163
|
+
for parameter in group
|
|
164
|
+
]
|
|
124
165
|
optimizer.zero_grad(set_to_none=True)
|
|
125
|
-
loss.backward()
|
|
166
|
+
loss.backward(inputs=parameters)
|
|
126
167
|
|
|
127
|
-
if
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
self.grad_norm_clip,
|
|
131
|
-
)
|
|
168
|
+
if clipnorm is not None:
|
|
169
|
+
for group in parameter_groups:
|
|
170
|
+
th.nn.utils.clip_grad_norm_(group, clipnorm)
|
|
132
171
|
|
|
133
172
|
optimizer.step()
|
|
134
173
|
|
|
@@ -20,8 +20,7 @@ class OffPolicyConfig:
|
|
|
20
20
|
regularizers: list[Callable[..., th.Tensor | float]] = field(default_factory=list)
|
|
21
21
|
reg_lams: list[float] = field(default_factory=list)
|
|
22
22
|
|
|
23
|
-
|
|
24
|
-
grad_norm_clip: float = 1.0
|
|
23
|
+
clipnorm: float | None = 1.0
|
|
25
24
|
|
|
26
25
|
use_target_model: bool = True
|
|
27
26
|
target_update: TargetUpdate | str = TargetUpdate.SOFT
|
|
@@ -47,6 +46,7 @@ class SACConfig(OffPolicyConfig):
|
|
|
47
46
|
criterion: th.nn.Module = field(default_factory=th.nn.MSELoss)
|
|
48
47
|
|
|
49
48
|
target_entropy: float | None = None
|
|
49
|
+
initial_alpha: float = 1.0
|
|
50
50
|
alpha_lr: float = 3e-4
|
|
51
51
|
|
|
52
52
|
min_log_std: float = -20.0
|
|
@@ -65,8 +65,11 @@ class SACConfig(OffPolicyConfig):
|
|
|
65
65
|
if self.action_dim <= 0:
|
|
66
66
|
raise ValueError("action_dim must be > 0.")
|
|
67
67
|
|
|
68
|
+
if self.initial_alpha <= 0.0:
|
|
69
|
+
raise ValueError("initial_alpha must be > 0.")
|
|
70
|
+
|
|
68
71
|
if self.alpha_lr <= 0.0:
|
|
69
72
|
raise ValueError("alpha_lr must be > 0.")
|
|
70
73
|
|
|
71
74
|
if self.min_log_std >= self.max_log_std:
|
|
72
|
-
raise ValueError("min_log_std must be < max_log_std.")
|
|
75
|
+
raise ValueError("min_log_std must be < max_log_std.")
|
|
@@ -1,5 +1,6 @@
|
|
|
1
|
+
import math
|
|
2
|
+
|
|
1
3
|
import torch as th
|
|
2
|
-
from torch.distributions import Normal
|
|
3
4
|
|
|
4
5
|
from .base import OffPolicyLearner
|
|
5
6
|
from .configs import SACConfig
|
|
@@ -14,8 +15,7 @@ class SACLearner(OffPolicyLearner):
|
|
|
14
15
|
critic1: th.nn.Module,
|
|
15
16
|
critic2: th.nn.Module,
|
|
16
17
|
actor_optimizer: th.optim.Optimizer,
|
|
17
|
-
|
|
18
|
-
critic2_optimizer: th.optim.Optimizer,
|
|
18
|
+
critic_optimizer: th.optim.Optimizer,
|
|
19
19
|
config: SACConfig,
|
|
20
20
|
):
|
|
21
21
|
super().__init__(config)
|
|
@@ -31,8 +31,21 @@ class SACLearner(OffPolicyLearner):
|
|
|
31
31
|
self.critic2_target = self.make_target_model(self.critic2)
|
|
32
32
|
|
|
33
33
|
self.actor_optimizer = actor_optimizer
|
|
34
|
-
self.
|
|
35
|
-
|
|
34
|
+
self.critic_optimizer = critic_optimizer
|
|
35
|
+
|
|
36
|
+
self._validate_optimizer_parameters(
|
|
37
|
+
self.actor_optimizer,
|
|
38
|
+
self.actor.parameters(),
|
|
39
|
+
"actor_optimizer",
|
|
40
|
+
)
|
|
41
|
+
self._validate_optimizer_parameters(
|
|
42
|
+
self.critic_optimizer,
|
|
43
|
+
[
|
|
44
|
+
*self.critic1.parameters(),
|
|
45
|
+
*self.critic2.parameters(),
|
|
46
|
+
],
|
|
47
|
+
"critic_optimizer",
|
|
48
|
+
)
|
|
36
49
|
|
|
37
50
|
self.criterion = config.criterion
|
|
38
51
|
|
|
@@ -42,7 +55,7 @@ class SACLearner(OffPolicyLearner):
|
|
|
42
55
|
self.target_entropy = -float(config.action_dim)
|
|
43
56
|
|
|
44
57
|
self.log_alpha = th.nn.Parameter(
|
|
45
|
-
th.tensor(
|
|
58
|
+
th.tensor(math.log(config.initial_alpha), device=self.device)
|
|
46
59
|
)
|
|
47
60
|
|
|
48
61
|
self.alpha_optimizer = th.optim.Adam(
|
|
@@ -54,15 +67,15 @@ class SACLearner(OffPolicyLearner):
|
|
|
54
67
|
def alpha(self) -> th.Tensor:
|
|
55
68
|
return self.log_alpha.exp()
|
|
56
69
|
|
|
57
|
-
def
|
|
58
|
-
for param in model.parameters():
|
|
59
|
-
param.requires_grad_(requires_grad)
|
|
60
|
-
|
|
61
|
-
def get_policy_dist(self, states: th.Tensor):
|
|
70
|
+
def _get_policy_params(self, states: th.Tensor) -> tuple[th.Tensor, th.Tensor]:
|
|
62
71
|
output = self.actor(states)
|
|
63
72
|
mean = output[:, : self.config.action_dim]
|
|
64
73
|
log_std = output[:, self.config.action_dim :]
|
|
65
74
|
log_std = th.clamp(log_std, self.config.min_log_std, self.config.max_log_std)
|
|
75
|
+
return mean, log_std
|
|
76
|
+
|
|
77
|
+
def get_policy_dist(self, states: th.Tensor) -> tuple[th.Tensor, th.Tensor]:
|
|
78
|
+
mean, log_std = self._get_policy_params(states)
|
|
66
79
|
std = th.exp(log_std)
|
|
67
80
|
return mean, std
|
|
68
81
|
|
|
@@ -71,19 +84,23 @@ class SACLearner(OffPolicyLearner):
|
|
|
71
84
|
obs: th.Tensor,
|
|
72
85
|
deterministic: bool = False,
|
|
73
86
|
) -> tuple[th.Tensor, th.Tensor]:
|
|
74
|
-
mean,
|
|
87
|
+
mean, log_std = self._get_policy_params(obs)
|
|
75
88
|
|
|
76
89
|
if deterministic:
|
|
77
90
|
u = mean
|
|
78
91
|
return th.tanh(u), th.zeros(obs.shape[0], 1, device=obs.device)
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
92
|
+
|
|
93
|
+
# Specialized reparameterized Gaussian sampling avoids the generic
|
|
94
|
+
# distribution object's redundant scale and normalized-noise work.
|
|
95
|
+
eps = th.randn_like(mean)
|
|
96
|
+
u = mean + th.exp(log_std) * eps
|
|
82
97
|
action = th.tanh(u)
|
|
83
98
|
|
|
84
|
-
log_prob =
|
|
85
|
-
|
|
86
|
-
|
|
99
|
+
log_prob = -0.5 * (
|
|
100
|
+
eps.square() + 2.0 * log_std + math.log(2.0 * math.pi)
|
|
101
|
+
)
|
|
102
|
+
log_prob -= th.log(1.0 - action.pow(2) + 1e-6) # Tanh correction
|
|
103
|
+
log_prob = log_prob.sum(dim=-1, keepdim=True) # Sum over action dimensions
|
|
87
104
|
|
|
88
105
|
return action, log_prob
|
|
89
106
|
|
|
@@ -123,22 +140,18 @@ class SACLearner(OffPolicyLearner):
|
|
|
123
140
|
critic2_loss = self.criterion(current_q2, targets)
|
|
124
141
|
critic_loss = critic1_loss + critic2_loss
|
|
125
142
|
self.optimize(
|
|
126
|
-
|
|
127
|
-
optimizer=self.
|
|
128
|
-
parameters=
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
parameters=self.critic2.parameters(),
|
|
143
|
+
critic_loss,
|
|
144
|
+
optimizer=self.critic_optimizer,
|
|
145
|
+
parameters=[
|
|
146
|
+
self.critic1.parameters(),
|
|
147
|
+
self.critic2.parameters(),
|
|
148
|
+
],
|
|
149
|
+
clipnorm=self.clipnorm,
|
|
134
150
|
)
|
|
135
151
|
|
|
136
152
|
# --------------------------------------------------
|
|
137
153
|
# 2. Update actor
|
|
138
154
|
# --------------------------------------------------
|
|
139
|
-
self._set_requires_grad(self.critic1, False)
|
|
140
|
-
self._set_requires_grad(self.critic2, False)
|
|
141
|
-
|
|
142
155
|
new_actions, log_probs = self.sample_action_and_log_prob(states)
|
|
143
156
|
new_action_states = th.cat([states, new_actions], dim=-1)
|
|
144
157
|
new_q1 = self.critic1(new_action_states)
|
|
@@ -158,12 +171,10 @@ class SACLearner(OffPolicyLearner):
|
|
|
158
171
|
self.optimize(
|
|
159
172
|
actor_loss,
|
|
160
173
|
optimizer=self.actor_optimizer,
|
|
161
|
-
parameters=self.actor.parameters(),
|
|
174
|
+
parameters=[self.actor.parameters()],
|
|
175
|
+
clipnorm=self.clipnorm,
|
|
162
176
|
)
|
|
163
177
|
|
|
164
|
-
self._set_requires_grad(self.critic1, True)
|
|
165
|
-
self._set_requires_grad(self.critic2, True)
|
|
166
|
-
|
|
167
178
|
# --------------------------------------------------
|
|
168
179
|
# 3. Update alpha
|
|
169
180
|
# --------------------------------------------------
|
|
@@ -174,7 +185,8 @@ class SACLearner(OffPolicyLearner):
|
|
|
174
185
|
self.optimize(
|
|
175
186
|
alpha_loss,
|
|
176
187
|
optimizer=self.alpha_optimizer,
|
|
177
|
-
parameters=[self.log_alpha],
|
|
188
|
+
parameters=[[self.log_alpha]],
|
|
189
|
+
clipnorm=self.clipnorm,
|
|
178
190
|
)
|
|
179
191
|
|
|
180
192
|
# Update target networks
|
|
@@ -1,4 +1,5 @@
|
|
|
1
1
|
import unittest
|
|
2
|
+
from unittest.mock import patch
|
|
2
3
|
|
|
3
4
|
import torch as th
|
|
4
5
|
|
|
@@ -25,6 +26,39 @@ def parameters_changed(model, before):
|
|
|
25
26
|
|
|
26
27
|
|
|
27
28
|
class LearnerSmokeTest(unittest.TestCase):
|
|
29
|
+
def test_soft_target_update_matches_original_arithmetic_exactly(self):
|
|
30
|
+
source1 = th.nn.Linear(4, 3)
|
|
31
|
+
source2 = th.nn.Linear(3, 2)
|
|
32
|
+
target1 = th.nn.Linear(4, 3)
|
|
33
|
+
target2 = th.nn.Linear(3, 2)
|
|
34
|
+
expected1 = th.nn.Linear(4, 3)
|
|
35
|
+
expected2 = th.nn.Linear(3, 2)
|
|
36
|
+
expected1.load_state_dict(target1.state_dict())
|
|
37
|
+
expected2.load_state_dict(target2.state_dict())
|
|
38
|
+
|
|
39
|
+
model = th.nn.Linear(1, 2)
|
|
40
|
+
learner = DQNLearner(
|
|
41
|
+
model,
|
|
42
|
+
th.optim.SGD(model.parameters(), lr=0.1),
|
|
43
|
+
DQNConfig(num_actions=2, soft_target_update_param=0.005),
|
|
44
|
+
)
|
|
45
|
+
|
|
46
|
+
with th.no_grad():
|
|
47
|
+
for target, source in ((expected1, source1), (expected2, source2)):
|
|
48
|
+
for target_param, source_param in zip(
|
|
49
|
+
target.parameters(), source.parameters()
|
|
50
|
+
):
|
|
51
|
+
target_param.data.mul_(0.995)
|
|
52
|
+
target_param.data.add_(0.005 * source_param.data)
|
|
53
|
+
|
|
54
|
+
learner.update_targets([(target1, source1), (target2, source2)])
|
|
55
|
+
|
|
56
|
+
for actual, expected in zip(
|
|
57
|
+
[*target1.parameters(), *target2.parameters()],
|
|
58
|
+
[*expected1.parameters(), *expected2.parameters()],
|
|
59
|
+
):
|
|
60
|
+
self.assertTrue(th.equal(actual, expected))
|
|
61
|
+
|
|
28
62
|
def test_dqn_train_step_returns_float_and_updates_parameters(self):
|
|
29
63
|
th.manual_seed(0)
|
|
30
64
|
model = th.nn.Sequential(
|
|
@@ -45,18 +79,45 @@ class LearnerSmokeTest(unittest.TestCase):
|
|
|
45
79
|
)
|
|
46
80
|
before = [param.detach().clone() for param in model.parameters()]
|
|
47
81
|
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
82
|
+
with patch(
|
|
83
|
+
"torch.nn.utils.clip_grad_norm_",
|
|
84
|
+
wraps=th.nn.utils.clip_grad_norm_,
|
|
85
|
+
) as clip_grad_norm:
|
|
86
|
+
loss = learner.train(
|
|
87
|
+
rewards=th.tensor([[1.0], [0.0], [1.0]]),
|
|
88
|
+
terminated=th.tensor([[False], [True], [False]]),
|
|
89
|
+
states=th.randn(3, 4),
|
|
90
|
+
actions=th.tensor([[0], [1], [0]], dtype=th.long),
|
|
91
|
+
next_states=th.randn(3, 4),
|
|
92
|
+
)
|
|
55
93
|
|
|
56
94
|
self.assertIsInstance(loss, float)
|
|
57
95
|
self.assertIsInstance(learner, OffPolicyLearner)
|
|
58
96
|
self.assertTrue(th.isfinite(th.tensor(loss)))
|
|
59
97
|
self.assertTrue(parameters_changed(model, before))
|
|
98
|
+
self.assertEqual(clip_grad_norm.call_count, 1)
|
|
99
|
+
|
|
100
|
+
def test_dqn_gradient_clipping_can_be_disabled(self):
|
|
101
|
+
model = th.nn.Linear(4, 2)
|
|
102
|
+
learner = DQNLearner(
|
|
103
|
+
model,
|
|
104
|
+
th.optim.SGD(model.parameters(), lr=0.05),
|
|
105
|
+
DQNConfig(
|
|
106
|
+
num_actions=2,
|
|
107
|
+
clipnorm=None,
|
|
108
|
+
),
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
with patch("torch.nn.utils.clip_grad_norm_") as clip_grad_norm:
|
|
112
|
+
learner.train(
|
|
113
|
+
rewards=th.tensor([[1.0], [0.0]]),
|
|
114
|
+
terminated=th.tensor([[False], [True]]),
|
|
115
|
+
states=th.randn(2, 4),
|
|
116
|
+
actions=th.tensor([[0], [1]], dtype=th.long),
|
|
117
|
+
next_states=th.randn(2, 4),
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
clip_grad_norm.assert_not_called()
|
|
60
121
|
|
|
61
122
|
def test_sac_train_step_returns_float_and_updates_parameters(self):
|
|
62
123
|
th.manual_seed(0)
|
|
@@ -82,27 +143,115 @@ class LearnerSmokeTest(unittest.TestCase):
|
|
|
82
143
|
critic1=critic1,
|
|
83
144
|
critic2=critic2,
|
|
84
145
|
actor_optimizer=th.optim.SGD(actor.parameters(), lr=0.05),
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
146
|
+
critic_optimizer=th.optim.SGD(
|
|
147
|
+
[*critic1.parameters(), *critic2.parameters()],
|
|
148
|
+
lr=0.05,
|
|
149
|
+
),
|
|
150
|
+
config=SACConfig(action_shape=(action_dim,), initial_alpha=0.2),
|
|
88
151
|
)
|
|
89
|
-
|
|
152
|
+
actor_before = [param.detach().clone() for param in actor.parameters()]
|
|
153
|
+
critic1_before = [param.detach().clone() for param in critic1.parameters()]
|
|
154
|
+
critic2_before = [param.detach().clone() for param in critic2.parameters()]
|
|
155
|
+
critic_grad_calls = {
|
|
156
|
+
id(parameter): 0
|
|
157
|
+
for parameter in (*critic1.parameters(), *critic2.parameters())
|
|
158
|
+
}
|
|
159
|
+
hooks = []
|
|
160
|
+
for parameter in (*critic1.parameters(), *critic2.parameters()):
|
|
161
|
+
parameter_id = id(parameter)
|
|
162
|
+
hooks.append(
|
|
163
|
+
parameter.register_hook(
|
|
164
|
+
lambda gradient, parameter_id=parameter_id: critic_grad_calls.__setitem__(
|
|
165
|
+
parameter_id,
|
|
166
|
+
critic_grad_calls[parameter_id] + 1,
|
|
167
|
+
)
|
|
168
|
+
)
|
|
169
|
+
)
|
|
90
170
|
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
171
|
+
self.assertAlmostEqual(learner.alpha.item(), 0.2)
|
|
172
|
+
|
|
173
|
+
probe_states = th.randn(3, obs_dim)
|
|
174
|
+
th.manual_seed(123)
|
|
175
|
+
sampled_actions, sampled_log_probs = learner.sample_action_and_log_prob(
|
|
176
|
+
probe_states
|
|
97
177
|
)
|
|
178
|
+
th.manual_seed(123)
|
|
179
|
+
mean, std = learner.get_policy_dist(probe_states)
|
|
180
|
+
dist = th.distributions.Normal(mean, std, validate_args=False)
|
|
181
|
+
unsquashed_actions = dist.rsample()
|
|
182
|
+
expected_actions = th.tanh(unsquashed_actions)
|
|
183
|
+
expected_log_probs = dist.log_prob(unsquashed_actions)
|
|
184
|
+
expected_log_probs -= th.log(1.0 - expected_actions.pow(2) + 1e-6)
|
|
185
|
+
expected_log_probs = expected_log_probs.sum(dim=-1, keepdim=True)
|
|
186
|
+
|
|
187
|
+
self.assertTrue(th.equal(sampled_actions, expected_actions))
|
|
188
|
+
self.assertTrue(
|
|
189
|
+
th.allclose(sampled_log_probs, expected_log_probs, atol=1e-6, rtol=1e-5)
|
|
190
|
+
)
|
|
191
|
+
|
|
192
|
+
with patch(
|
|
193
|
+
"torch.nn.utils.clip_grad_norm_",
|
|
194
|
+
wraps=th.nn.utils.clip_grad_norm_,
|
|
195
|
+
) as clip_grad_norm:
|
|
196
|
+
loss = learner.train(
|
|
197
|
+
rewards=th.tensor([[1.0], [0.0], [1.0]]),
|
|
198
|
+
terminated=th.tensor([[False], [True], [False]]),
|
|
199
|
+
states=th.randn(3, obs_dim),
|
|
200
|
+
actions=th.randn(3, action_dim).tanh(),
|
|
201
|
+
next_states=th.randn(3, obs_dim),
|
|
202
|
+
)
|
|
98
203
|
|
|
99
204
|
self.assertIsInstance(loss, float)
|
|
100
205
|
self.assertTrue(th.isfinite(th.tensor(loss)))
|
|
101
|
-
self.assertTrue(parameters_changed(actor,
|
|
206
|
+
self.assertTrue(parameters_changed(actor, actor_before))
|
|
207
|
+
self.assertTrue(parameters_changed(critic1, critic1_before))
|
|
208
|
+
self.assertTrue(parameters_changed(critic2, critic2_before))
|
|
209
|
+
self.assertTrue(all(call_count == 1 for call_count in critic_grad_calls.values()))
|
|
210
|
+
self.assertEqual(clip_grad_norm.call_count, 4)
|
|
102
211
|
self.assertEqual(
|
|
103
212
|
set(learner.last_losses),
|
|
104
213
|
{"actor", "critic", "alpha", "regularization", "total"},
|
|
105
214
|
)
|
|
215
|
+
for hook in hooks:
|
|
216
|
+
hook.remove()
|
|
217
|
+
|
|
218
|
+
def test_sac_initial_alpha_must_be_positive(self):
|
|
219
|
+
with self.assertRaisesRegex(ValueError, "initial_alpha must be > 0"):
|
|
220
|
+
SACConfig(initial_alpha=0.0)
|
|
221
|
+
|
|
222
|
+
def test_sac_validates_actor_and_critic_optimizer_parameters(self):
|
|
223
|
+
obs_dim = 4
|
|
224
|
+
action_dim = 2
|
|
225
|
+
actor = th.nn.Linear(obs_dim, 2 * action_dim)
|
|
226
|
+
critic1 = th.nn.Linear(obs_dim + action_dim, 1)
|
|
227
|
+
critic2 = th.nn.Linear(obs_dim + action_dim, 1)
|
|
228
|
+
critic_parameters = [*critic1.parameters(), *critic2.parameters()]
|
|
229
|
+
|
|
230
|
+
with self.assertRaisesRegex(
|
|
231
|
+
ValueError,
|
|
232
|
+
"actor_optimizer must contain exactly its model parameters",
|
|
233
|
+
):
|
|
234
|
+
SACLearner(
|
|
235
|
+
actor=actor,
|
|
236
|
+
critic1=critic1,
|
|
237
|
+
critic2=critic2,
|
|
238
|
+
actor_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
|
|
239
|
+
critic_optimizer=th.optim.SGD(critic_parameters, lr=0.05),
|
|
240
|
+
config=SACConfig(action_shape=(action_dim,)),
|
|
241
|
+
)
|
|
242
|
+
|
|
243
|
+
with self.assertRaisesRegex(
|
|
244
|
+
ValueError,
|
|
245
|
+
"critic_optimizer must contain exactly its model parameters",
|
|
246
|
+
):
|
|
247
|
+
SACLearner(
|
|
248
|
+
actor=actor,
|
|
249
|
+
critic1=critic1,
|
|
250
|
+
critic2=critic2,
|
|
251
|
+
actor_optimizer=th.optim.SGD(actor.parameters(), lr=0.05),
|
|
252
|
+
critic_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
|
|
253
|
+
config=SACConfig(action_shape=(action_dim,)),
|
|
254
|
+
)
|
|
106
255
|
|
|
107
256
|
def test_sac_actor_regularizer_contributes_to_loss(self):
|
|
108
257
|
th.manual_seed(0)
|
|
@@ -134,8 +283,10 @@ class LearnerSmokeTest(unittest.TestCase):
|
|
|
134
283
|
critic1=critic1,
|
|
135
284
|
critic2=critic2,
|
|
136
285
|
actor_optimizer=th.optim.SGD(actor.parameters(), lr=0.05),
|
|
137
|
-
|
|
138
|
-
|
|
286
|
+
critic_optimizer=th.optim.SGD(
|
|
287
|
+
[*critic1.parameters(), *critic2.parameters()],
|
|
288
|
+
lr=0.05,
|
|
289
|
+
),
|
|
139
290
|
config=SACConfig(
|
|
140
291
|
action_shape=(action_dim,),
|
|
141
292
|
regularizers=[actor_l2],
|
|
@@ -37,7 +37,7 @@ class PublicAPITest(unittest.TestCase):
|
|
|
37
37
|
"WarmupController",
|
|
38
38
|
}
|
|
39
39
|
|
|
40
|
-
self.assertEqual(drlab.__version__, "0.2.
|
|
40
|
+
self.assertEqual(drlab.__version__, "0.2.5")
|
|
41
41
|
self.assertTrue(expected.issubset(set(drlab.__all__)))
|
|
42
42
|
for name in expected:
|
|
43
43
|
self.assertTrue(hasattr(drlab, name), name)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|