drlab 0.2.6__tar.gz → 0.2.8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drlab
-Version: 0.2.6
+Version: 0.2.8
 Summary: Deep Reinforcement Learning kit for research.
 Author: Tomas Osarte
 Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium
@@ -139,11 +139,15 @@ episode returns, episode lengths, and optionally the last completed episode.
 
 `ReplayBuffer` stores transitions for off-policy learning. It supports both
 discrete and continuous action spaces through the configured action shape and
-action type.
+action type. Call `reset()` to remove all stored transitions and reuse the
+buffer.
 
 `TransitionBatch` is the tensor container passed between runners, replay
 buffers, experiments, and learners.
 
+All learners provide `reset_optimizers()` to clear optimizer state without
+changing gradients, optimizer settings, or model parameters.
+
 ## Simple Usage
 
 This example trains a DQN agent on `CartPole-v1`:

{drlab-0.2.6 → drlab-0.2.8}/README.md

@@ -102,11 +102,15 @@ episode returns, episode lengths, and optionally the last completed episode.
 
 `ReplayBuffer` stores transitions for off-policy learning. It supports both
 discrete and continuous action spaces through the configured action shape and
-action type.
+action type. Call `reset()` to remove all stored transitions and reuse the
+buffer.
 
 `TransitionBatch` is the tensor container passed between runners, replay
 buffers, experiments, and learners.
 
+All learners provide `reset_optimizers()` to clear optimizer state without
+changing gradients, optimizer settings, or model parameters.
+
 ## Simple Usage
 
 This example trains a DQN agent on `CartPole-v1`:

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

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

{drlab-0.2.6 → drlab-0.2.8}/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.6"
+__version__ = "0.2.8"
 
 __all__ = [
     "__version__",

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

@@ -28,6 +28,10 @@ class DQNLearner(OffPolicyLearner):
         if self.use_target_model:
             self.target_model = self.make_target_model(self.model)
 
+    def reset_optimizers(self) -> None:
+        """Clear the optimizer state."""
+        self.optimizer.state.clear()
+
     def _validate_dqn_config(self) -> None:
         if self.double_q and not self.use_target_model:
             raise ValueError("double_q=True requires use_target_model=True.")

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

@@ -63,6 +63,15 @@ class SACLearner(OffPolicyLearner):
             lr=config.alpha_lr,
         )
 
+    def reset_optimizers(self) -> None:
+        """Clear the state of every optimizer owned by the learner."""
+        for optimizer in (
+            self.actor_optimizer,
+            self.critic_optimizer,
+            self.alpha_optimizer,
+        ):
+            optimizer.state.clear()
+
     @property
     def alpha(self) -> th.Tensor:
         return self.log_alpha.exp()
@@ -196,12 +205,18 @@ class SACLearner(OffPolicyLearner):
         ])
 
         loss = actor_loss + critic_loss + alpha_loss
+
+        # A single batched host transfer avoids five separate device
+        # synchronizations (one per .item() call) on every training step.
+        actor_v, critic_v, alpha_v, reg_v, total_v = th.stack(
+            [policy_loss, critic_loss, alpha_loss, reg_loss, loss]
+        ).tolist()
         self.last_losses = {
-            "actor": float(policy_loss.item()),
-            "critic": float(critic_loss.item()),
-            "alpha": float(alpha_loss.item()),
-            "regularization": float(reg_loss.item()),
-            "total": float(loss.item()),
+            "actor": actor_v,
+            "critic": critic_v,
+            "alpha": alpha_v,
+            "regularization": reg_v,
+            "total": total_v,
         }
 
         return self.last_losses["total"]

{drlab-0.2.6 → drlab-0.2.8}/src/drlab/learners/on_policy/base.py

@@ -20,6 +20,10 @@ class OnPolicyLearner:
 
         self._validate_config()
 
+    def reset_optimizers(self) -> None:
+        """Clear the optimizer state."""
+        self.optimizer.state.clear()
+
     def _validate_config(self):
         if len(self.config.regularizers) != len(self.config.reg_lams):
             raise ValueError("regularizers and reg_lams must have same length.")

{drlab-0.2.6 → drlab-0.2.8}/src/drlab/replay/replay_buffer.py

@@ -37,6 +37,11 @@ class ReplayBuffer:
     def __len__(self) -> int:
         return self.size
 
+    def reset(self) -> None:
+        """Remove all transitions from the buffer."""
+        self.ptr = 0
+        self.size = 0
+
     def add(
         self,
         states: np.ndarray,       # [B, *obs_shape]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drlab
-Version: 0.2.6
+Version: 0.2.8
 Summary: Deep Reinforcement Learning kit for research.
 Author: Tomas Osarte
 Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium
@@ -139,11 +139,15 @@ episode returns, episode lengths, and optionally the last completed episode.
 
 `ReplayBuffer` stores transitions for off-policy learning. It supports both
 discrete and continuous action spaces through the configured action shape and
-action type.
+action type. Call `reset()` to remove all stored transitions and reuse the
+buffer.
 
 `TransitionBatch` is the tensor container passed between runners, replay
 buffers, experiments, and learners.
 
+All learners provide `reset_optimizers()` to clear optimizer state without
+changing gradients, optimizer settings, or model parameters.
+
 ## Simple Usage
 
 This example trains a DQN agent on `CartPole-v1`:

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

@@ -25,7 +25,88 @@ def parameters_changed(model, before):
     )
 
 
+def populate_optimizer_state(optimizer):
+    optimizer.zero_grad(set_to_none=True)
+    loss = sum(
+        parameter.square().sum()
+        for group in optimizer.param_groups
+        for parameter in group["params"]
+    )
+    loss.backward()
+    optimizer.step()
+
+
+def assert_optimizer_state_reset(test_case, optimizer):
+    test_case.assertEqual(len(optimizer.state), 0)
+    for group in optimizer.param_groups:
+        for parameter in group["params"]:
+            test_case.assertIsNotNone(parameter.grad)
+
+
 class LearnerSmokeTest(unittest.TestCase):
+    def test_all_on_policy_learners_reset_their_optimizer(self):
+        cases = (
+            (ReinforceLearner, ReinforceConfig(num_actions=2), 2),
+            (ActorCriticLearner, ActorCriticConfig(num_actions=2), 3),
+            (PPOLearner, PPOConfig(num_actions=2), 3),
+        )
+
+        for learner_type, config, output_size in cases:
+            with self.subTest(learner=learner_type.__name__):
+                actor = th.nn.Linear(4, output_size)
+                optimizer = th.optim.Adam(actor.parameters())
+                learner = learner_type(actor, optimizer, config)
+                populate_optimizer_state(optimizer)
+
+                learner.reset_optimizers()
+
+                assert_optimizer_state_reset(self, optimizer)
+
+    def test_dqn_resets_its_optimizer(self):
+        model = th.nn.Linear(4, 2)
+        optimizer = th.optim.Adam(model.parameters())
+        learner = DQNLearner(
+            model,
+            optimizer,
+            DQNConfig(num_actions=2),
+        )
+        populate_optimizer_state(optimizer)
+
+        learner.reset_optimizers()
+
+        assert_optimizer_state_reset(self, optimizer)
+
+    def test_sac_resets_all_three_optimizers(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)
+        actor_optimizer = th.optim.Adam(actor.parameters())
+        critic_optimizer = th.optim.Adam(
+            [*critic1.parameters(), *critic2.parameters()]
+        )
+        learner = SACLearner(
+            actor=actor,
+            critic1=critic1,
+            critic2=critic2,
+            actor_optimizer=actor_optimizer,
+            critic_optimizer=critic_optimizer,
+            config=SACConfig(action_shape=(action_dim,)),
+        )
+        optimizers = (
+            actor_optimizer,
+            critic_optimizer,
+            learner.alpha_optimizer,
+        )
+        for optimizer in optimizers:
+            populate_optimizer_state(optimizer)
+
+        learner.reset_optimizers()
+
+        for optimizer in optimizers:
+            assert_optimizer_state_reset(self, optimizer)
+
     def test_soft_target_update_matches_original_arithmetic_exactly(self):
         source1 = th.nn.Linear(4, 3)
         source2 = th.nn.Linear(3, 2)

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

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

{drlab-0.2.6 → drlab-0.2.8}/tests/test_replay.py

@@ -134,3 +134,37 @@ class ReplayBufferTest(unittest.TestCase):
 
         with self.assertRaisesRegex(ValueError, "empty buffer"):
             buffer.sample(1)
+
+    def test_reset_empties_buffer_and_allows_reuse(self):
+        buffer = ReplayBuffer(capacity=3, obs_shape=(1,), device="cpu")
+        buffer.add(
+            states=np.asarray([[0], [1]], dtype=np.float32),
+            actions=np.asarray([0, 1]),
+            rewards=np.asarray([1.0, 2.0]),
+            terminated=np.asarray([False, True]),
+            truncated=np.asarray([False, False]),
+            next_states=np.asarray([[1], [2]], dtype=np.float32),
+            returns=np.asarray([3.0, 2.0]),
+        )
+
+        buffer.reset()
+
+        self.assertEqual(len(buffer), 0)
+        self.assertEqual(buffer.ptr, 0)
+        self.assertEqual(buffer.get_all().states.shape, (0, 1))
+        with self.assertRaisesRegex(ValueError, "empty buffer"):
+            buffer.sample(1)
+
+        buffer.add(
+            states=np.asarray([[4]], dtype=np.float32),
+            actions=np.asarray([1]),
+            rewards=np.asarray([5.0]),
+            terminated=np.asarray([False]),
+            truncated=np.asarray([True]),
+            next_states=np.asarray([[5]], dtype=np.float32),
+            returns=np.asarray([5.0]),
+        )
+
+        self.assertEqual(len(buffer), 1)
+        self.assertEqual(buffer.ptr, 1)
+        self.assertEqual(buffer.get_all().states.item(), 4.0)