drlab 0.2.4__tar.gz → 0.2.6__tar.gz

{drlab-0.2.4/src/drlab.egg-info → drlab-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drlab
-Version: 0.2.4
+Version: 0.2.6
 Summary: Deep Reinforcement Learning kit for research.
 Author: Tomas Osarte
 Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium

{drlab-0.2.4 → drlab-0.2.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "drlab"
-version = "0.2.4"
+version = "0.2.6"
 description = "Deep Reinforcement Learning kit for research."
 readme = "README.md"
 requires-python = ">=3.11"

{drlab-0.2.4 → drlab-0.2.6}/src/drlab/__init__.py

@@ -35,7 +35,7 @@ from drlab.learners import (
 from drlab.replay import ReplayBuffer, TransitionBatch
 from drlab.runners import Runner
 
-__version__ = "0.2.4"
+__version__ = "0.2.6"
 
 __all__ = [
     "__version__",

{drlab-0.2.4 → drlab-0.2.6}/src/drlab/learners/off_policy/base.py

@@ -52,6 +52,26 @@ class OffPolicyLearner(ABC):
 
         return target
 
+    @staticmethod
+    def _validate_optimizer_parameters(
+        optimizer: th.optim.Optimizer,
+        expected_parameters: Iterable[th.nn.Parameter],
+        optimizer_name: str,
+    ) -> None:
+        expected_parameter_ids = sorted(
+            id(parameter) for parameter in expected_parameters
+        )
+        optimizer_parameter_ids = sorted(
+            id(parameter)
+            for group in optimizer.param_groups
+            for parameter in group["params"]
+        )
+
+        if optimizer_parameter_ids != expected_parameter_ids:
+            raise ValueError(
+                f"{optimizer_name} must contain exactly its model parameters."
+            )
+
     def _hard_update(self, target: th.nn.Module, source: th.nn.Module):
         target.load_state_dict(source.state_dict())
         target.eval()
@@ -133,18 +153,21 @@ class OffPolicyLearner(ABC):
         self,
         loss: th.Tensor,
         optimizer: th.optim.Optimizer,
-        parameters: Iterable[th.nn.Parameter],
+        parameters: list[Iterable[th.nn.Parameter]],
         clipnorm: float | None = None,
     ) -> None:
-        parameters = tuple(parameters)
+        parameter_groups = [tuple(group) for group in parameters]
+        parameters = [
+            parameter
+            for group in parameter_groups
+            for parameter in group
+        ]
         optimizer.zero_grad(set_to_none=True)
         loss.backward(inputs=parameters)
 
         if clipnorm is not None:
-            th.nn.utils.clip_grad_norm_(
-                parameters,
-                clipnorm,
-            )
+            for group in parameter_groups:
+                th.nn.utils.clip_grad_norm_(group, clipnorm)
 
         optimizer.step()
 

{drlab-0.2.4 → drlab-0.2.6}/src/drlab/learners/off_policy/dqn.py

@@ -88,7 +88,7 @@ class DQNLearner(OffPolicyLearner):
         self.optimize(
             loss,
             optimizer=self.optimizer,
-            parameters=self.model.parameters(),
+            parameters=[self.model.parameters()],
             clipnorm=self.clipnorm,
         )
 

{drlab-0.2.4 → drlab-0.2.6}/src/drlab/learners/off_policy/sac.py

@@ -15,8 +15,7 @@ class SACLearner(OffPolicyLearner):
         critic1: th.nn.Module,
         critic2: th.nn.Module,
         actor_optimizer: th.optim.Optimizer,
-        critic1_optimizer: th.optim.Optimizer,
-        critic2_optimizer: th.optim.Optimizer,
+        critic_optimizer: th.optim.Optimizer,
         config: SACConfig,
     ):
         super().__init__(config)
@@ -32,8 +31,21 @@ class SACLearner(OffPolicyLearner):
         self.critic2_target = self.make_target_model(self.critic2)
 
         self.actor_optimizer = actor_optimizer
-        self.critic1_optimizer = critic1_optimizer
-        self.critic2_optimizer = critic2_optimizer
+        self.critic_optimizer = critic_optimizer
+
+        self._validate_optimizer_parameters(
+            self.actor_optimizer,
+            self.actor.parameters(),
+            "actor_optimizer",
+        )
+        self._validate_optimizer_parameters(
+            self.critic_optimizer,
+            [
+                *self.critic1.parameters(),
+                *self.critic2.parameters(),
+            ],
+            "critic_optimizer",
+        )
 
         self.criterion = config.criterion
 
@@ -127,23 +139,15 @@ class SACLearner(OffPolicyLearner):
         critic1_loss = self.criterion(current_q1, targets)
         critic2_loss = self.criterion(current_q2, targets)
         critic_loss = critic1_loss + critic2_loss
-        self.critic1_optimizer.zero_grad(set_to_none=True)
-        self.critic2_optimizer.zero_grad(set_to_none=True)
-        critic_loss.backward()
-
-        if self.clipnorm is not None:
-            # Clip each critic independently to preserve the previous behavior.
-            th.nn.utils.clip_grad_norm_(
+        self.optimize(
+            critic_loss,
+            optimizer=self.critic_optimizer,
+            parameters=[
                 self.critic1.parameters(),
-                self.clipnorm,
-            )
-            th.nn.utils.clip_grad_norm_(
                 self.critic2.parameters(),
-                self.clipnorm,
-            )
-
-        self.critic1_optimizer.step()
-        self.critic2_optimizer.step()
+            ],
+            clipnorm=self.clipnorm,
+        )
 
         # --------------------------------------------------
         # 2. Update actor
@@ -167,7 +171,7 @@ class SACLearner(OffPolicyLearner):
         self.optimize(
             actor_loss,
             optimizer=self.actor_optimizer,
-            parameters=self.actor.parameters(),
+            parameters=[self.actor.parameters()],
             clipnorm=self.clipnorm,
         )
 
@@ -181,7 +185,7 @@ class SACLearner(OffPolicyLearner):
         self.optimize(
             alpha_loss,
             optimizer=self.alpha_optimizer,
-            parameters=[self.log_alpha],
+            parameters=[[self.log_alpha]],
             clipnorm=self.clipnorm,
         )
 

{drlab-0.2.4 → drlab-0.2.6/src/drlab.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drlab
-Version: 0.2.4
+Version: 0.2.6
 Summary: Deep Reinforcement Learning kit for research.
 Author: Tomas Osarte
 Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium

{drlab-0.2.4 → drlab-0.2.6}/tests/test_learners.py

@@ -143,8 +143,10 @@ class LearnerSmokeTest(unittest.TestCase):
             critic1=critic1,
             critic2=critic2,
             actor_optimizer=th.optim.SGD(actor.parameters(), lr=0.05),
-            critic1_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
-            critic2_optimizer=th.optim.SGD(critic2.parameters(), lr=0.05),
+            critic_optimizer=th.optim.SGD(
+                [*critic1.parameters(), *critic2.parameters()],
+                lr=0.05,
+            ),
             config=SACConfig(action_shape=(action_dim,), initial_alpha=0.2),
         )
         actor_before = [param.detach().clone() for param in actor.parameters()]
@@ -217,6 +219,40 @@ class LearnerSmokeTest(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "initial_alpha must be > 0"):
             SACConfig(initial_alpha=0.0)
 
+    def test_sac_validates_actor_and_critic_optimizer_parameters(self):
+        obs_dim = 4
+        action_dim = 2
+        actor = th.nn.Linear(obs_dim, 2 * action_dim)
+        critic1 = th.nn.Linear(obs_dim + action_dim, 1)
+        critic2 = th.nn.Linear(obs_dim + action_dim, 1)
+        critic_parameters = [*critic1.parameters(), *critic2.parameters()]
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "actor_optimizer must contain exactly its model parameters",
+        ):
+            SACLearner(
+                actor=actor,
+                critic1=critic1,
+                critic2=critic2,
+                actor_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
+                critic_optimizer=th.optim.SGD(critic_parameters, lr=0.05),
+                config=SACConfig(action_shape=(action_dim,)),
+            )
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "critic_optimizer must contain exactly its model parameters",
+        ):
+            SACLearner(
+                actor=actor,
+                critic1=critic1,
+                critic2=critic2,
+                actor_optimizer=th.optim.SGD(actor.parameters(), lr=0.05),
+                critic_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
+                config=SACConfig(action_shape=(action_dim,)),
+            )
+
     def test_sac_actor_regularizer_contributes_to_loss(self):
         th.manual_seed(0)
         obs_dim = 4
@@ -247,8 +283,10 @@ class LearnerSmokeTest(unittest.TestCase):
             critic1=critic1,
             critic2=critic2,
             actor_optimizer=th.optim.SGD(actor.parameters(), lr=0.05),
-            critic1_optimizer=th.optim.SGD(critic1.parameters(), lr=0.05),
-            critic2_optimizer=th.optim.SGD(critic2.parameters(), lr=0.05),
+            critic_optimizer=th.optim.SGD(
+                [*critic1.parameters(), *critic2.parameters()],
+                lr=0.05,
+            ),
             config=SACConfig(
                 action_shape=(action_dim,),
                 regularizers=[actor_l2],

{drlab-0.2.4 → drlab-0.2.6}/tests/test_public_api.py

@@ -37,7 +37,7 @@ class PublicAPITest(unittest.TestCase):
             "WarmupController",
         }
 
-        self.assertEqual(drlab.__version__, "0.2.4")
+        self.assertEqual(drlab.__version__, "0.2.5")
         self.assertTrue(expected.issubset(set(drlab.__all__)))
         for name in expected:
             self.assertTrue(hasattr(drlab, name), name)