drlab 0.2.7__tar.gz → 0.2.9__tar.gz

{drlab-0.2.7/src/drlab.egg-info → drlab-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drlab
-Version: 0.2.7
+Version: 0.2.9
 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.7 → drlab-0.2.9}/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.7 → drlab-0.2.9}/pyproject.toml

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

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

{drlab-0.2.7 → drlab-0.2.9}/src/drlab/learners/off_policy/base.py

@@ -28,6 +28,10 @@ class OffPolicyLearner(ABC):
 
         self._validate_config()
 
+    def reset_optimizers(self) -> None:
+        """Reset optimizer state owned by the learner, if any."""
+        pass
+
     def _validate_config(self) -> None:
         if len(self.regularizers) != len(self.reg_lams):
             raise ValueError("regularizers and reg_lams must have same length.")

{drlab-0.2.7 → drlab-0.2.9}/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.7 → drlab-0.2.9}/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()

{drlab-0.2.7 → drlab-0.2.9}/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.7 → drlab-0.2.9}/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.7 → drlab-0.2.9/src/drlab.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: drlab
-Version: 0.2.7
+Version: 0.2.9
 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.7 → drlab-0.2.9}/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.7 → drlab-0.2.9}/tests/test_off_policy_experiment.py

@@ -71,6 +71,11 @@ def one_transition() -> TransitionBatch:
 
 
 class OffPolicyExperimentTest(unittest.TestCase):
+    def test_base_learner_optimizer_reset_is_a_noop(self):
+        learner = CountingLearner()
+
+        self.assertIsNone(learner.reset_optimizers())
+
     def make_experiment(self, config: OffPolicyExperimentConfig) -> OffPolicyExperiment:
         env = OneStepEnv()
         learner = CountingLearner()

{drlab-0.2.7 → drlab-0.2.9}/tests/test_public_api.py

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

{drlab-0.2.7 → drlab-0.2.9}/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)