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.
Files changed (44) hide show
  1. {drlab-0.2.3/src/drlab.egg-info → drlab-0.2.5}/PKG-INFO +1 -1
  2. {drlab-0.2.3 → drlab-0.2.5}/pyproject.toml +1 -1
  3. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/__init__.py +1 -1
  4. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/base.py +58 -19
  5. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/configs.py +6 -3
  6. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/dqn.py +2 -1
  7. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/sac.py +46 -34
  8. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/base.py +1 -1
  9. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/configs.py +1 -1
  10. {drlab-0.2.3 → drlab-0.2.5/src/drlab.egg-info}/PKG-INFO +1 -1
  11. {drlab-0.2.3 → drlab-0.2.5}/tests/test_learners.py +171 -20
  12. {drlab-0.2.3 → drlab-0.2.5}/tests/test_public_api.py +1 -1
  13. {drlab-0.2.3 → drlab-0.2.5}/LICENSE +0 -0
  14. {drlab-0.2.3 → drlab-0.2.5}/README.md +0 -0
  15. {drlab-0.2.3 → drlab-0.2.5}/setup.cfg +0 -0
  16. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/__init__.py +0 -0
  17. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/base.py +0 -0
  18. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/e_greedy.py +0 -0
  19. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/gaussian.py +0 -0
  20. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/greedy.py +0 -0
  21. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/stochastic.py +0 -0
  22. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/controllers/warmup.py +0 -0
  23. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/experiments/__init__.py +0 -0
  24. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/experiments/off_policy_experiment.py +0 -0
  25. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/experiments/on_policy_experiment.py +0 -0
  26. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/__init__.py +0 -0
  27. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/off_policy/__init__.py +0 -0
  28. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/__init__.py +0 -0
  29. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/actor_critic.py +0 -0
  30. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/ppo.py +0 -0
  31. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/learners/on_policy/reinforce.py +0 -0
  32. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/replay/__init__.py +0 -0
  33. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/replay/replay_buffer.py +0 -0
  34. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/replay/transition_batch.py +0 -0
  35. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/runners/__init__.py +0 -0
  36. {drlab-0.2.3 → drlab-0.2.5}/src/drlab/runners/runner.py +0 -0
  37. {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/SOURCES.txt +0 -0
  38. {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/dependency_links.txt +0 -0
  39. {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/requires.txt +0 -0
  40. {drlab-0.2.3 → drlab-0.2.5}/src/drlab.egg-info/top_level.txt +0 -0
  41. {drlab-0.2.3 → drlab-0.2.5}/tests/test_controllers.py +0 -0
  42. {drlab-0.2.3 → drlab-0.2.5}/tests/test_off_policy_experiment.py +0 -0
  43. {drlab-0.2.3 → drlab-0.2.5}/tests/test_replay.py +0 -0
  44. {drlab-0.2.3 → drlab-0.2.5}/tests/test_runner.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: drlab
3
- Version: 0.2.3
3
+ Version: 0.2.5
4
4
  Summary: Deep Reinforcement Learning kit for research.
5
5
  Author: Tomas Osarte
6
6
  Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "drlab"
7
- version = "0.2.3"
7
+ version = "0.2.5"
8
8
  description = "Deep Reinforcement Learning kit for research."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.11"
@@ -35,7 +35,7 @@ from drlab.learners import (
35
35
  from drlab.replay import ReplayBuffer, TransitionBatch
36
36
  from drlab.runners import Runner
37
37
 
38
- __version__ = "0.2.3"
38
+ __version__ = "0.2.5"
39
39
 
40
40
  __all__ = [
41
41
  "__version__",
@@ -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.clip_grad = config.clip_grad
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
- def hard_update(self, target: th.nn.Module, source: th.nn.Module):
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 soft_update(self, target: th.nn.Module, source: th.nn.Module):
60
- tau = self.soft_target_update_param
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
- for target_p, source_p in zip(target.parameters(), source.parameters()):
64
- target_p.data.mul_(1.0 - tau)
65
- target_p.data.add_(tau * source_p.data)
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
- for target, source in pairs:
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.hard_update(target, source)
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 self.clip_grad:
128
- th.nn.utils.clip_grad_norm_(
129
- parameters,
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
- clip_grad: bool = True
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.")
@@ -88,7 +88,8 @@ class DQNLearner(OffPolicyLearner):
88
88
  self.optimize(
89
89
  loss,
90
90
  optimizer=self.optimizer,
91
- parameters=self.model.parameters(),
91
+ parameters=[self.model.parameters()],
92
+ clipnorm=self.clipnorm,
92
93
  )
93
94
 
94
95
  if self.target_model is not None:
@@ -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
- critic1_optimizer: th.optim.Optimizer,
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.critic1_optimizer = critic1_optimizer
35
- self.critic2_optimizer = critic2_optimizer
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(0.0, device=self.device)
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 _set_requires_grad(self, model: th.nn.Module, requires_grad: bool) -> None:
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, std = self.get_policy_dist(obs)
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
- dist = Normal(mean, std)
81
- u = dist.rsample() # reparameterization trick
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 = dist.log_prob(u) # Gaussian log prob before tanh
85
- log_prob -= th.log(1.0 - action.pow(2) + 1e-6) # Tanh correction
86
- log_prob = log_prob.sum(dim=-1, keepdim=True) # Sum over action dimensions
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
- critic1_loss,
127
- optimizer=self.critic1_optimizer,
128
- parameters=self.critic1.parameters(),
129
- )
130
- self.optimize(
131
- critic2_loss,
132
- optimizer=self.critic2_optimizer,
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
@@ -87,7 +87,7 @@ class OnPolicyLearner:
87
87
  if self.config.clip_grad:
88
88
  th.nn.utils.clip_grad_norm_(
89
89
  self.parameters,
90
- self.config.grad_norm_clip,
90
+ self.config.clipnorm,
91
91
  )
92
92
  self.optimizer.step()
93
93
 
@@ -18,7 +18,7 @@ class OnPolicyConfig:
18
18
  num_actions: int = 2
19
19
 
20
20
  clip_grad: bool = True
21
- grad_norm_clip: float = 1.0
21
+ clipnorm: float = 1.0
22
22
 
23
23
  use_entropy: bool = False
24
24
  entropy_max_lambda: float = 0.0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: drlab
3
- Version: 0.2.3
3
+ Version: 0.2.5
4
4
  Summary: Deep Reinforcement Learning kit for research.
5
5
  Author: Tomas Osarte
6
6
  Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium
@@ -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
- loss = learner.train(
49
- rewards=th.tensor([[1.0], [0.0], [1.0]]),
50
- terminated=th.tensor([[False], [True], [False]]),
51
- states=th.randn(3, 4),
52
- actions=th.tensor([[0], [1], [0]], dtype=th.long),
53
- next_states=th.randn(3, 4),
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
- critic1_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
86
- critic2_optimizer=th.optim.SGD(critic2.parameters(), lr=0.05),
87
- config=SACConfig(action_shape=(action_dim,)),
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
- before = [param.detach().clone() for param in actor.parameters()]
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
- loss = learner.train(
92
- rewards=th.tensor([[1.0], [0.0], [1.0]]),
93
- terminated=th.tensor([[False], [True], [False]]),
94
- states=th.randn(3, obs_dim),
95
- actions=th.randn(3, action_dim).tanh(),
96
- next_states=th.randn(3, obs_dim),
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, before))
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
- critic1_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
138
- critic2_optimizer=th.optim.SGD(critic2.parameters(), lr=0.05),
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.3")
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