drlab 0.2.3__tar.gz → 0.2.4__tar.gz

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

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

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

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

{drlab-0.2.3 → drlab-0.2.4}/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.3"
+__version__ = "0.2.4"
 
 __all__ = [
     "__version__",

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

@@ -1,6 +1,7 @@
 import torch as th
 from copy import deepcopy
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 
 from .configs import OffPolicyConfig, TargetUpdate
 
@@ -15,8 +16,7 @@ class OffPolicyLearner(ABC):
         self.regularizers = config.regularizers
         self.reg_lams = config.reg_lams
 
-        self.clip_grad = config.clip_grad
-        self.grad_norm_clip = config.grad_norm_clip
+        self.clipnorm = config.clipnorm
 
         self.use_target_model = config.use_target_model
         self.target_update = TargetUpdate(config.target_update)
@@ -52,25 +52,39 @@ class OffPolicyLearner(ABC):
 
         return target
 
-    def hard_update(self, target: th.nn.Module, source: th.nn.Module):
+    def _hard_update(self, target: th.nn.Module, source: th.nn.Module):
         target.load_state_dict(source.state_dict())
         target.eval()
 
-    def soft_update(self, target: th.nn.Module, source: th.nn.Module):
-        tau = self.soft_target_update_param
+    def _soft_update(
+        self,
+        pairs: list[tuple[th.nn.Module, th.nn.Module]],
+    ) -> None:
+        target_parameters = [
+            parameter
+            for target, _ in pairs
+            for parameter in target.parameters()
+        ]
+        source_parameters = [
+            parameter
+            for _, source in pairs
+            for parameter in source.parameters()
+        ]
 
         with th.no_grad():
-            for target_p, source_p in zip(target.parameters(), source.parameters()):
-                target_p.data.mul_(1.0 - tau)
-                target_p.data.add_(tau * source_p.data)
+            th._foreach_mul_(target_parameters, 1.0 - self.soft_target_update_param)
+            scaled_sources = th._foreach_mul(
+                source_parameters,
+                self.soft_target_update_param,
+            )
+            th._foreach_add_(target_parameters, scaled_sources)
 
     def update_targets(
         self,
         pairs: list[tuple[th.nn.Module, th.nn.Module]],
     ):
         if self.target_update == TargetUpdate.SOFT:
-            for target, source in pairs:
-                self.soft_update(target, source)
+            self._soft_update(pairs)
             return
 
         if self.target_update == TargetUpdate.HARD:
@@ -78,7 +92,7 @@ class OffPolicyLearner(ABC):
 
             if self.target_update_calls % self.target_update_interval == 0:
                 for target, source in pairs:
-                    self.hard_update(target, source)
+                    self._hard_update(target, source)
 
             return
 
@@ -119,15 +133,17 @@ class OffPolicyLearner(ABC):
         self,
         loss: th.Tensor,
         optimizer: th.optim.Optimizer,
-        parameters: list[th.nn.Parameter]
-    ):
+        parameters: Iterable[th.nn.Parameter],
+        clipnorm: float | None = None,
+    ) -> None:
+        parameters = tuple(parameters)
         optimizer.zero_grad(set_to_none=True)
-        loss.backward()
+        loss.backward(inputs=parameters)
 
-        if self.clip_grad:
+        if clipnorm is not None:
             th.nn.utils.clip_grad_norm_(
                 parameters,
-                self.grad_norm_clip,
+                clipnorm,
             )
 
         optimizer.step()

{drlab-0.2.3 → drlab-0.2.4}/src/drlab/learners/off_policy/configs.py

@@ -20,8 +20,7 @@ class OffPolicyConfig:
     regularizers: list[Callable[..., th.Tensor | float]] = field(default_factory=list)
     reg_lams: list[float] = field(default_factory=list)
 
-    clip_grad: bool = True
-    grad_norm_clip: float = 1.0
+    clipnorm: float | None = 1.0
 
     use_target_model: bool = True
     target_update: TargetUpdate | str = TargetUpdate.SOFT
@@ -47,6 +46,7 @@ class SACConfig(OffPolicyConfig):
     criterion: th.nn.Module = field(default_factory=th.nn.MSELoss)
 
     target_entropy: float | None = None
+    initial_alpha: float = 1.0
     alpha_lr: float = 3e-4
 
     min_log_std: float = -20.0
@@ -65,8 +65,11 @@ class SACConfig(OffPolicyConfig):
         if self.action_dim <= 0:
             raise ValueError("action_dim must be > 0.")
 
+        if self.initial_alpha <= 0.0:
+            raise ValueError("initial_alpha must be > 0.")
+
         if self.alpha_lr <= 0.0:
             raise ValueError("alpha_lr must be > 0.")
 
         if self.min_log_std >= self.max_log_std:
-            raise ValueError("min_log_std must be < max_log_std.")
+            raise ValueError("min_log_std must be < max_log_std.")

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

@@ -89,6 +89,7 @@ class DQNLearner(OffPolicyLearner):
             loss,
             optimizer=self.optimizer,
             parameters=self.model.parameters(),
+            clipnorm=self.clipnorm,
         )
 
         if self.target_model is not None:

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

@@ -1,5 +1,6 @@
+import math
+
 import torch as th
-from torch.distributions import Normal
 
 from .base import OffPolicyLearner
 from .configs import SACConfig
@@ -42,7 +43,7 @@ class SACLearner(OffPolicyLearner):
             self.target_entropy = -float(config.action_dim)
 
         self.log_alpha = th.nn.Parameter(
-            th.tensor(0.0, device=self.device)
+            th.tensor(math.log(config.initial_alpha), device=self.device)
         )
 
         self.alpha_optimizer = th.optim.Adam(
@@ -54,15 +55,15 @@ class SACLearner(OffPolicyLearner):
     def alpha(self) -> th.Tensor:
         return self.log_alpha.exp()
 
-    def _set_requires_grad(self, model: th.nn.Module, requires_grad: bool) -> None:
-        for param in model.parameters():
-            param.requires_grad_(requires_grad)
-
-    def get_policy_dist(self, states: th.Tensor):
+    def _get_policy_params(self, states: th.Tensor) -> tuple[th.Tensor, th.Tensor]:
         output = self.actor(states)
         mean = output[:, : self.config.action_dim]
         log_std = output[:, self.config.action_dim :]
         log_std = th.clamp(log_std, self.config.min_log_std, self.config.max_log_std)
+        return mean, log_std
+
+    def get_policy_dist(self, states: th.Tensor) -> tuple[th.Tensor, th.Tensor]:
+        mean, log_std = self._get_policy_params(states)
         std = th.exp(log_std)
         return mean, std
 
@@ -71,19 +72,23 @@ class SACLearner(OffPolicyLearner):
         obs: th.Tensor,
         deterministic: bool = False,
     ) -> tuple[th.Tensor, th.Tensor]:
-        mean, std = self.get_policy_dist(obs)
+        mean, log_std = self._get_policy_params(obs)
 
         if deterministic:
             u = mean
             return th.tanh(u), th.zeros(obs.shape[0], 1, device=obs.device)
-        
-        dist = Normal(mean, std)
-        u = dist.rsample()  # reparameterization trick
+
+        # Specialized reparameterized Gaussian sampling avoids the generic
+        # distribution object's redundant scale and normalized-noise work.
+        eps = th.randn_like(mean)
+        u = mean + th.exp(log_std) * eps
         action = th.tanh(u)
 
-        log_prob = dist.log_prob(u) # Gaussian log prob before tanh        
-        log_prob -= th.log(1.0 - action.pow(2) + 1e-6) # Tanh correction
-        log_prob = log_prob.sum(dim=-1, keepdim=True) # Sum over action dimensions
+        log_prob = -0.5 * (
+            eps.square() + 2.0 * log_std + math.log(2.0 * math.pi)
+        )
+        log_prob -= th.log(1.0 - action.pow(2) + 1e-6)  # Tanh correction
+        log_prob = log_prob.sum(dim=-1, keepdim=True)  # Sum over action dimensions
 
         return action, log_prob
 
@@ -122,23 +127,27 @@ class SACLearner(OffPolicyLearner):
         critic1_loss = self.criterion(current_q1, targets)
         critic2_loss = self.criterion(current_q2, targets)
         critic_loss = critic1_loss + critic2_loss
-        self.optimize(
-            critic1_loss,
-            optimizer=self.critic1_optimizer,
-            parameters=self.critic1.parameters(),
-        )
-        self.optimize(
-            critic2_loss,
-            optimizer=self.critic2_optimizer,
-            parameters=self.critic2.parameters(),
-        )
+        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.critic1.parameters(),
+                self.clipnorm,
+            )
+            th.nn.utils.clip_grad_norm_(
+                self.critic2.parameters(),
+                self.clipnorm,
+            )
+
+        self.critic1_optimizer.step()
+        self.critic2_optimizer.step()
 
         # --------------------------------------------------
         # 2. Update actor
         # --------------------------------------------------
-        self._set_requires_grad(self.critic1, False)
-        self._set_requires_grad(self.critic2, False)
-        
         new_actions, log_probs = self.sample_action_and_log_prob(states)
         new_action_states = th.cat([states, new_actions], dim=-1)
         new_q1 = self.critic1(new_action_states)
@@ -159,11 +168,9 @@ class SACLearner(OffPolicyLearner):
             actor_loss,
             optimizer=self.actor_optimizer,
             parameters=self.actor.parameters(),
+            clipnorm=self.clipnorm,
         )
 
-        self._set_requires_grad(self.critic1, True)
-        self._set_requires_grad(self.critic2, True)
-
         # --------------------------------------------------
         # 3. Update alpha
         # --------------------------------------------------
@@ -175,6 +182,7 @@ class SACLearner(OffPolicyLearner):
             alpha_loss,
             optimizer=self.alpha_optimizer,
             parameters=[self.log_alpha],
+            clipnorm=self.clipnorm,
         )
 
         # Update target networks

{drlab-0.2.3 → drlab-0.2.4}/src/drlab/learners/on_policy/base.py

@@ -87,7 +87,7 @@ class OnPolicyLearner:
         if self.config.clip_grad:
             th.nn.utils.clip_grad_norm_(
                 self.parameters,
-                self.config.grad_norm_clip,
+                self.config.clipnorm,
             )
         self.optimizer.step()
 

{drlab-0.2.3 → drlab-0.2.4}/src/drlab/learners/on_policy/configs.py

@@ -18,7 +18,7 @@ class OnPolicyConfig:
     num_actions: int = 2
 
     clip_grad: bool = True
-    grad_norm_clip: float = 1.0
+    clipnorm: float = 1.0
 
     use_entropy: bool = False
     entropy_max_lambda: float = 0.0

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

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

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

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import patch
 
 import torch as th
 
@@ -25,6 +26,39 @@ def parameters_changed(model, before):
 
 
 class LearnerSmokeTest(unittest.TestCase):
+    def test_soft_target_update_matches_original_arithmetic_exactly(self):
+        source1 = th.nn.Linear(4, 3)
+        source2 = th.nn.Linear(3, 2)
+        target1 = th.nn.Linear(4, 3)
+        target2 = th.nn.Linear(3, 2)
+        expected1 = th.nn.Linear(4, 3)
+        expected2 = th.nn.Linear(3, 2)
+        expected1.load_state_dict(target1.state_dict())
+        expected2.load_state_dict(target2.state_dict())
+
+        model = th.nn.Linear(1, 2)
+        learner = DQNLearner(
+            model,
+            th.optim.SGD(model.parameters(), lr=0.1),
+            DQNConfig(num_actions=2, soft_target_update_param=0.005),
+        )
+
+        with th.no_grad():
+            for target, source in ((expected1, source1), (expected2, source2)):
+                for target_param, source_param in zip(
+                    target.parameters(), source.parameters()
+                ):
+                    target_param.data.mul_(0.995)
+                    target_param.data.add_(0.005 * source_param.data)
+
+        learner.update_targets([(target1, source1), (target2, source2)])
+
+        for actual, expected in zip(
+            [*target1.parameters(), *target2.parameters()],
+            [*expected1.parameters(), *expected2.parameters()],
+        ):
+            self.assertTrue(th.equal(actual, expected))
+
     def test_dqn_train_step_returns_float_and_updates_parameters(self):
         th.manual_seed(0)
         model = th.nn.Sequential(
@@ -45,18 +79,45 @@ class LearnerSmokeTest(unittest.TestCase):
         )
         before = [param.detach().clone() for param in model.parameters()]
 
-        loss = learner.train(
-            rewards=th.tensor([[1.0], [0.0], [1.0]]),
-            terminated=th.tensor([[False], [True], [False]]),
-            states=th.randn(3, 4),
-            actions=th.tensor([[0], [1], [0]], dtype=th.long),
-            next_states=th.randn(3, 4),
-        )
+        with patch(
+            "torch.nn.utils.clip_grad_norm_",
+            wraps=th.nn.utils.clip_grad_norm_,
+        ) as clip_grad_norm:
+            loss = learner.train(
+                rewards=th.tensor([[1.0], [0.0], [1.0]]),
+                terminated=th.tensor([[False], [True], [False]]),
+                states=th.randn(3, 4),
+                actions=th.tensor([[0], [1], [0]], dtype=th.long),
+                next_states=th.randn(3, 4),
+            )
 
         self.assertIsInstance(loss, float)
         self.assertIsInstance(learner, OffPolicyLearner)
         self.assertTrue(th.isfinite(th.tensor(loss)))
         self.assertTrue(parameters_changed(model, before))
+        self.assertEqual(clip_grad_norm.call_count, 1)
+
+    def test_dqn_gradient_clipping_can_be_disabled(self):
+        model = th.nn.Linear(4, 2)
+        learner = DQNLearner(
+            model,
+            th.optim.SGD(model.parameters(), lr=0.05),
+            DQNConfig(
+                num_actions=2,
+                clipnorm=None,
+            ),
+        )
+
+        with patch("torch.nn.utils.clip_grad_norm_") as clip_grad_norm:
+            learner.train(
+                rewards=th.tensor([[1.0], [0.0]]),
+                terminated=th.tensor([[False], [True]]),
+                states=th.randn(2, 4),
+                actions=th.tensor([[0], [1]], dtype=th.long),
+                next_states=th.randn(2, 4),
+            )
+
+        clip_grad_norm.assert_not_called()
 
     def test_sac_train_step_returns_float_and_updates_parameters(self):
         th.manual_seed(0)
@@ -84,25 +145,77 @@ class LearnerSmokeTest(unittest.TestCase):
             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),
-            config=SACConfig(action_shape=(action_dim,)),
+            config=SACConfig(action_shape=(action_dim,), initial_alpha=0.2),
         )
-        before = [param.detach().clone() for param in actor.parameters()]
+        actor_before = [param.detach().clone() for param in actor.parameters()]
+        critic1_before = [param.detach().clone() for param in critic1.parameters()]
+        critic2_before = [param.detach().clone() for param in critic2.parameters()]
+        critic_grad_calls = {
+            id(parameter): 0
+            for parameter in (*critic1.parameters(), *critic2.parameters())
+        }
+        hooks = []
+        for parameter in (*critic1.parameters(), *critic2.parameters()):
+            parameter_id = id(parameter)
+            hooks.append(
+                parameter.register_hook(
+                    lambda gradient, parameter_id=parameter_id: critic_grad_calls.__setitem__(
+                        parameter_id,
+                        critic_grad_calls[parameter_id] + 1,
+                    )
+                )
+            )
 
-        loss = learner.train(
-            rewards=th.tensor([[1.0], [0.0], [1.0]]),
-            terminated=th.tensor([[False], [True], [False]]),
-            states=th.randn(3, obs_dim),
-            actions=th.randn(3, action_dim).tanh(),
-            next_states=th.randn(3, obs_dim),
+        self.assertAlmostEqual(learner.alpha.item(), 0.2)
+
+        probe_states = th.randn(3, obs_dim)
+        th.manual_seed(123)
+        sampled_actions, sampled_log_probs = learner.sample_action_and_log_prob(
+            probe_states
         )
+        th.manual_seed(123)
+        mean, std = learner.get_policy_dist(probe_states)
+        dist = th.distributions.Normal(mean, std, validate_args=False)
+        unsquashed_actions = dist.rsample()
+        expected_actions = th.tanh(unsquashed_actions)
+        expected_log_probs = dist.log_prob(unsquashed_actions)
+        expected_log_probs -= th.log(1.0 - expected_actions.pow(2) + 1e-6)
+        expected_log_probs = expected_log_probs.sum(dim=-1, keepdim=True)
+
+        self.assertTrue(th.equal(sampled_actions, expected_actions))
+        self.assertTrue(
+            th.allclose(sampled_log_probs, expected_log_probs, atol=1e-6, rtol=1e-5)
+        )
+
+        with patch(
+            "torch.nn.utils.clip_grad_norm_",
+            wraps=th.nn.utils.clip_grad_norm_,
+        ) as clip_grad_norm:
+            loss = learner.train(
+                rewards=th.tensor([[1.0], [0.0], [1.0]]),
+                terminated=th.tensor([[False], [True], [False]]),
+                states=th.randn(3, obs_dim),
+                actions=th.randn(3, action_dim).tanh(),
+                next_states=th.randn(3, obs_dim),
+            )
 
         self.assertIsInstance(loss, float)
         self.assertTrue(th.isfinite(th.tensor(loss)))
-        self.assertTrue(parameters_changed(actor, before))
+        self.assertTrue(parameters_changed(actor, actor_before))
+        self.assertTrue(parameters_changed(critic1, critic1_before))
+        self.assertTrue(parameters_changed(critic2, critic2_before))
+        self.assertTrue(all(call_count == 1 for call_count in critic_grad_calls.values()))
+        self.assertEqual(clip_grad_norm.call_count, 4)
         self.assertEqual(
             set(learner.last_losses),
             {"actor", "critic", "alpha", "regularization", "total"},
         )
+        for hook in hooks:
+            hook.remove()
+
+    def test_sac_initial_alpha_must_be_positive(self):
+        with self.assertRaisesRegex(ValueError, "initial_alpha must be > 0"):
+            SACConfig(initial_alpha=0.0)
 
     def test_sac_actor_regularizer_contributes_to_loss(self):
         th.manual_seed(0)

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

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