drlab 0.2.1__tar.gz → 0.2.2__tar.gz

{drlab-0.2.1/src/drlab.egg-info → drlab-0.2.2}/PKG-INFO

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

{drlab-0.2.1 → drlab-0.2.2}/pyproject.toml

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

{drlab-0.2.1 → drlab-0.2.2}/src/drlab/__init__.py

@@ -6,6 +6,7 @@ from drlab.controllers import (
     GaussianController,
     GreedyController,
     StochasticController,
+    WarmupController,
 )
 from drlab.experiments import (
     OnPolicyExperiment,
@@ -34,7 +35,7 @@ from drlab.learners import (
 from drlab.replay import ReplayBuffer, TransitionBatch
 from drlab.runners import Runner
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 __all__ = [
     "__version__",
@@ -68,4 +69,5 @@ __all__ = [
     "StochasticController",
     "TransitionBatch",
     "ValueTargets",
+    "WarmupController",
 ]

{drlab-0.2.1 → drlab-0.2.2}/src/drlab/controllers/__init__.py

@@ -1,8 +1,9 @@
 from .base import ContinuousActionController, Controller, DiscreteActionController
 from .greedy import GreedyController
 from .e_greedy import EpsilonGreedyController
-from .stochastic_controller import StochasticController
-from .gaussian_controller import GaussianController
+from .stochastic import StochasticController
+from .gaussian import GaussianController
+from .warmup import WarmupController
 
 __all__ = [
     "ContinuousActionController",
@@ -12,4 +13,5 @@ __all__ = [
     "GaussianController",
     "GreedyController",
     "StochasticController",
+    "WarmupController",
 ]

drlab-0.2.2/src/drlab/controllers/warmup.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+import numpy as np
+import torch as th
+
+from .base import ContinuousActionController, Controller, DiscreteActionController
+
+
+class WarmupController(ContinuousActionController, DiscreteActionController):
+    def __init__(self, controller: Controller, action_space: Any, warmup_steps: int):
+        self.controller = controller
+        self.action_space = action_space
+        self.warmup_steps = warmup_steps
+        self.steps = 0
+        self.model = controller.model
+
+        if hasattr(controller, "num_actions"):
+            self.num_actions = controller.num_actions
+        if hasattr(controller, "action_dim"):
+            self.action_dim = controller.action_dim
+
+        if warmup_steps < 0:
+            raise ValueError("warmup_steps must be >= 0.")
+
+    def choose(
+        self,
+        obs: th.Tensor,
+        **kwargs,
+    ) -> th.Tensor:
+        if self.steps >= self.warmup_steps:
+            return self.controller.choose(obs, **kwargs)
+
+        batch_size = obs.shape[0]
+        self.steps += batch_size
+        actions = [self.action_space.sample() for _ in range(batch_size)]
+        return th.as_tensor(np.asarray(actions), device=obs.device)
+
+    def probabilities(self, obs: th.Tensor, **kwargs) -> th.Tensor:
+        if self.steps >= self.warmup_steps:
+            return self.controller.probabilities(obs, **kwargs)
+
+        return th.ones((obs.shape[0], self.num_actions), device=obs.device) / self.num_actions

{drlab-0.2.1 → drlab-0.2.2}/src/drlab/experiments/off_policy_experiment.py

@@ -2,12 +2,12 @@ import numpy as np
 from tqdm import tqdm
 import gymnasium as gym
 from dataclasses import dataclass
-from typing import Callable, Tuple
+from typing import Callable
 from torch.utils.tensorboard import SummaryWriter
 
 from drlab.learners import OffPolicyLearner, SACLearner
 from drlab.runners import Runner
-from drlab.controllers import Controller
+from drlab.controllers import Controller, WarmupController
 from drlab.replay import TransitionBatch, ReplayBuffer
 
 @dataclass
@@ -24,6 +24,8 @@ class OffPolicyExperimentConfig:
     grad_repeats: int = 1
     step_callback: Callable[[int], None] | None = None
     step_callback_interval: int | None = None
+    learning_starts: int | None = None
+    warmup_steps: int = 0
 
 
 class OffPolicyExperiment:
@@ -44,9 +46,22 @@ class OffPolicyExperiment:
         self.run_steps = config.run_steps
         self.step_callback = config.step_callback
         self.step_callback_interval = config.step_callback_interval
+        self.steps = 0
+        self.learning_starts = (
+            config.batch_size if config.learning_starts is None else config.learning_starts
+        )
+        self.warmup_steps = config.warmup_steps
+
+        self._validate_config()
 
         # Init drl components
         continous_actions = isinstance(learner, SACLearner)
+        if self.warmup_steps > 0:
+            controller = WarmupController(
+                controller,
+                env.action_space,
+                warmup_steps=self.warmup_steps,
+            )
         self.runner = Runner(
             env,
             controller,
@@ -73,9 +88,17 @@ class OffPolicyExperiment:
         self.writer = SummaryWriter(log_dir=config.log_dir)
         self.experiment_name = config.experiment_name
 
+
+    def _validate_config(self):
         if self.step_callback is not None and self.step_callback_interval is None:
             raise ValueError("step_callback_interval must be set when step_callback is provided.")
 
+        if self.learning_starts < self.batch_size:
+            raise ValueError("learning_starts must be greater than or equal to batch_size.")
+
+        if self.warmup_steps < 0:
+            raise ValueError("warmup_steps must be greater than or equal to 0.")
+
     def _make_minibatch(self, batch: TransitionBatch, last_episode: TransitionBatch | None) -> TransitionBatch:
 
         # No replay buffer
@@ -97,7 +120,11 @@ class OffPolicyExperiment:
         rest = self.replay_buffer.sample(self.batch_size - ep_len)
         return rest.cat(episode)
 
-    def _learn_from_batch(self, batch: TransitionBatch, last_episode: TransitionBatch | None) -> float:
+    def _learn_from_batch(
+        self,
+        batch: TransitionBatch,
+        last_episode: TransitionBatch | None,
+    ) -> float:
         # 1) Add batch to replay buffer
         self.replay_buffer.add(
             states=batch.states,
@@ -107,7 +134,7 @@ class OffPolicyExperiment:
             next_states=batch.next_states,
             returns=batch.returns,
         )
-        if self.replay_buffer.size < self.batch_size:
+        if self.steps < self.learning_starts:
             return 0.0
 
         # 2) train repeats
@@ -121,13 +148,13 @@ class OffPolicyExperiment:
         return total_loss / self.grad_repeats
 
     def run(self):
-
-        steps = 0
+        
+        self.steps = 0
         if self.step_callback is not None:
-            self.step_callback(steps)
+            self.step_callback(self.steps)
         next_callback_step = self.step_callback_interval
         pbar = tqdm(total=self.max_steps, desc=self.experiment_name, dynamic_ncols=True, mininterval=1.0)
-        while steps < self.max_steps:
+        while self.steps < self.max_steps:
 
             # 1. Run batch of transitions
             batch, ep_returns, ep_lengths, last_episode = self.runner.run(
@@ -135,16 +162,20 @@ class OffPolicyExperiment:
                 as_numpy=True,
             )
 
+            step_inc = batch.states.shape[0]
+            self.steps += step_inc
+            pbar.update(step_inc)
+
             # 2. Learn from batch
             loss = self._learn_from_batch(batch, last_episode)
 
             # 3. Log results
-            self.writer.add_scalar("Loss", loss, steps)
+            self.writer.add_scalar("Loss", loss, self.steps)
             if ep_returns:
                 mean_return = np.mean(ep_returns)
                 mean_length = np.mean(ep_lengths)
-                self.writer.add_scalar("Episode return", mean_return, steps)
-                self.writer.add_scalar("Episode Length", mean_length, steps)
+                self.writer.add_scalar("Episode return", mean_return, self.steps)
+                self.writer.add_scalar("Episode Length", mean_length, self.steps)
 
                 pbar.set_postfix({
                     "loss": f"{loss:.3f}",
@@ -152,12 +183,7 @@ class OffPolicyExperiment:
                     "len": f"{mean_length:.1f}",
                 }, refresh=False)
 
-            # 3. Update steps & progress bar
-            step_inc = batch.states.shape[0] 
-            steps += step_inc 
-            pbar.update(step_inc)
-
-            while self.step_callback is not None and steps >= next_callback_step:
+            while self.step_callback is not None and self.steps >= next_callback_step:
                 self.step_callback(next_callback_step)
                 next_callback_step += self.step_callback_interval
         

{drlab-0.2.1 → drlab-0.2.2}/src/drlab/learners/off_policy/sac.py

@@ -145,7 +145,16 @@ class SACLearner(OffPolicyLearner):
         new_q2 = self.critic2(new_action_states)
         new_q = th.min(new_q1, new_q2)
 
-        actor_loss = (self.alpha.detach() * log_probs - new_q).mean()
+        policy_loss = (self.alpha.detach() * log_probs - new_q).mean()
+        reg_loss = self.regularization_loss(
+            self.actor,
+            rewards=rewards,
+            dones=dones,
+            states=states,
+            actions=actions,
+            next_states=next_states,
+        )
+        actor_loss = policy_loss + reg_loss
         self.optimize(
             actor_loss,
             optimizer=self.actor_optimizer,
@@ -176,9 +185,10 @@ class SACLearner(OffPolicyLearner):
 
         loss = actor_loss + critic_loss + alpha_loss
         self.last_losses = {
-            "actor": float(actor_loss.item()),
+            "actor": float(policy_loss.item()),
             "critic": float(critic_loss.item()),
             "alpha": float(alpha_loss.item()),
+            "regularization": float(reg_loss.item()),
             "total": float(loss.item()),
         }
 

{drlab-0.2.1 → drlab-0.2.2/src/drlab.egg-info}/PKG-INFO

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

{drlab-0.2.1 → drlab-0.2.2}/src/drlab.egg-info/SOURCES.txt

@@ -10,9 +10,10 @@ src/drlab.egg-info/top_level.txt
 src/drlab/controllers/__init__.py
 src/drlab/controllers/base.py
 src/drlab/controllers/e_greedy.py
-src/drlab/controllers/gaussian_controller.py
+src/drlab/controllers/gaussian.py
 src/drlab/controllers/greedy.py
-src/drlab/controllers/stochastic_controller.py
+src/drlab/controllers/stochastic.py
+src/drlab/controllers/warmup.py
 src/drlab/experiments/__init__.py
 src/drlab/experiments/off_policy_experiment.py
 src/drlab/experiments/on_policy_experiment.py
@@ -35,6 +36,7 @@ src/drlab/runners/__init__.py
 src/drlab/runners/runner.py
 tests/test_controllers.py
 tests/test_learners.py
+tests/test_off_policy_experiment.py
 tests/test_public_api.py
 tests/test_replay.py
 tests/test_runner.py

{drlab-0.2.1 → drlab-0.2.2}/tests/test_controllers.py

@@ -1,5 +1,6 @@
 import unittest
 
+import numpy as np
 import torch as th
 
 from drlab.controllers import (
@@ -9,6 +10,7 @@ from drlab.controllers import (
     GaussianController,
     GreedyController,
     StochasticController,
+    WarmupController,
 )
 
 
@@ -25,6 +27,14 @@ class FixedGaussianOutput(FixedLogits):
     pass
 
 
+class FixedActionSpace:
+    def __init__(self, action):
+        self.action = action
+
+    def sample(self):
+        return self.action
+
+
 class ControllerTest(unittest.TestCase):
     def test_discrete_controllers_share_discrete_base(self):
         model = FixedLogits([0.1, 2.0, -1.0])
@@ -115,3 +125,39 @@ class ControllerTest(unittest.TestCase):
         self.assertTrue(th.allclose(probs.sum(dim=-1), th.ones(5)))
         self.assertEqual(actions.shape, (5,))
         self.assertTrue(th.all((0 <= actions) & (actions < 3)))
+
+    def test_warmup_controller_samples_actions_before_delegating(self):
+        model = FixedLogits([1.0, 0.0, 0.0])
+        greedy = GreedyController(model, num_actions=3)
+        controller = WarmupController(greedy, FixedActionSpace(2), warmup_steps=2)
+        obs = th.zeros(1, 2)
+
+        self.assertIsInstance(controller, ContinuousActionController)
+        self.assertIsInstance(controller, DiscreteActionController)
+        self.assertTrue(th.equal(controller.choose(obs), th.tensor([2])))
+        self.assertTrue(th.equal(controller.choose(obs), th.tensor([2])))
+        self.assertTrue(th.equal(controller.choose(obs), th.tensor([0])))
+        self.assertEqual(controller.steps, 2)
+
+    def test_warmup_controller_returns_uniform_discrete_probabilities(self):
+        model = FixedLogits([1.0, 0.0, 0.0])
+        greedy = GreedyController(model, num_actions=3)
+        controller = WarmupController(greedy, FixedActionSpace(2), warmup_steps=1)
+
+        probs = controller.probabilities(th.zeros(1, 2))
+
+        self.assertTrue(th.allclose(probs, th.full((1, 3), 1 / 3)))
+
+    def test_warmup_controller_samples_continuous_actions(self):
+        model = FixedGaussianOutput([0.0, 0.0, -3.0, -3.0])
+        gaussian = GaussianController(model, action_dim=2, deterministic=True)
+        action_space = FixedActionSpace(np.asarray([0.5, -0.25], dtype=np.float32))
+        controller = WarmupController(gaussian, action_space, warmup_steps=1)
+        obs = th.zeros(2, 2)
+
+        warmup_actions = controller.choose(obs)
+        policy_actions = controller.choose(obs)
+
+        self.assertEqual(warmup_actions.shape, (2, 2))
+        self.assertTrue(th.allclose(warmup_actions[0], th.tensor([0.5, -0.25])))
+        self.assertTrue(th.allclose(policy_actions, th.zeros(2, 2)))

{drlab-0.2.1 → drlab-0.2.2}/tests/test_learners.py

@@ -101,9 +101,59 @@ class LearnerSmokeTest(unittest.TestCase):
         self.assertTrue(parameters_changed(actor, before))
         self.assertEqual(
             set(learner.last_losses),
-            {"actor", "critic", "alpha", "total"},
+            {"actor", "critic", "alpha", "regularization", "total"},
         )
 
+    def test_sac_actor_regularizer_contributes_to_loss(self):
+        th.manual_seed(0)
+        obs_dim = 4
+        action_dim = 2
+        actor = th.nn.Sequential(
+            th.nn.Linear(obs_dim, 8),
+            th.nn.ReLU(),
+            th.nn.Linear(8, 2 * action_dim),
+        )
+        critic1 = th.nn.Sequential(
+            th.nn.Linear(obs_dim + action_dim, 8),
+            th.nn.ReLU(),
+            th.nn.Linear(8, 1),
+        )
+        critic2 = th.nn.Sequential(
+            th.nn.Linear(obs_dim + action_dim, 8),
+            th.nn.ReLU(),
+            th.nn.Linear(8, 1),
+        )
+        calls = []
+
+        def actor_l2(model, rewards, dones, states, actions, next_states):
+            calls.append(model)
+            return sum(param.square().sum() for param in model.parameters())
+
+        learner = SACLearner(
+            actor=actor,
+            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),
+            config=SACConfig(
+                action_shape=(action_dim,),
+                regularizers=[actor_l2],
+                reg_lams=[0.01],
+            ),
+        )
+
+        learner.train(
+            rewards=th.tensor([[1.0], [0.0], [1.0]]),
+            dones=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.assertEqual(calls, [actor])
+        self.assertGreater(learner.last_losses["regularization"], 0.0)
+
     def test_reinforce_train_step_returns_float_and_updates_parameters(self):
         th.manual_seed(0)
         actor = th.nn.Sequential(

drlab-0.2.2/tests/test_off_policy_experiment.py

@@ -0,0 +1,154 @@
+import tempfile
+import unittest
+
+import gymnasium as gym
+import numpy as np
+import torch as th
+
+from drlab.controllers import GreedyController, WarmupController
+from drlab.experiments import OffPolicyExperiment, OffPolicyExperimentConfig
+from drlab.learners import OffPolicyConfig, OffPolicyLearner
+from drlab.replay import TransitionBatch
+
+
+class FixedLogits(th.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("logits", th.tensor([1.0, 0.0]))
+
+    def forward(self, obs):
+        return self.logits.unsqueeze(0).expand(obs.shape[0], -1)
+
+
+class OneStepEnv(gym.Env):
+    metadata = {}
+
+    def __init__(self):
+        self.observation_space = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=(1,),
+            dtype=np.float32,
+        )
+        self.action_space = gym.spaces.Discrete(2)
+
+    def reset(self, *, seed=None, options=None):
+        super().reset(seed=seed)
+        return np.asarray([0.0], dtype=np.float32), {}
+
+    def step(self, action):
+        obs = np.asarray([0.0], dtype=np.float32)
+        return obs, 0.0, True, False, {}
+
+
+class CountingLearner(OffPolicyLearner):
+    def __init__(self):
+        super().__init__(OffPolicyConfig())
+        self.train_calls = 0
+
+    def train(
+        self,
+        rewards: th.Tensor,
+        dones: th.Tensor,
+        states: th.Tensor,
+        actions: th.Tensor,
+        next_states: th.Tensor,
+    ) -> float:
+        self.train_calls += 1
+        return 7.0
+
+
+def one_transition() -> TransitionBatch:
+    return TransitionBatch(
+        states=np.asarray([[0.0]], dtype=np.float32),
+        actions=np.asarray([[0]], dtype=np.int64),
+        rewards=np.asarray([[1.0]], dtype=np.float32),
+        dones=np.asarray([[False]], dtype=np.bool_),
+        next_states=np.asarray([[1.0]], dtype=np.float32),
+        returns=np.asarray([[1.0]], dtype=np.float32),
+    )
+
+
+class OffPolicyExperimentTest(unittest.TestCase):
+    def make_experiment(self, config: OffPolicyExperimentConfig) -> OffPolicyExperiment:
+        env = OneStepEnv()
+        learner = CountingLearner()
+        controller = GreedyController(FixedLogits(), num_actions=2)
+        return OffPolicyExperiment(env, controller, learner, config)
+
+    def test_learning_starts_default_tracks_batch_size(self):
+        with tempfile.TemporaryDirectory() as log_dir:
+            experiment = self.make_experiment(
+                OffPolicyExperimentConfig(
+                    max_steps=1,
+                    batch_size=256,
+                    log_dir=log_dir,
+                )
+            )
+            try:
+                self.assertEqual(experiment.learning_starts, 256)
+            finally:
+                experiment.writer.close()
+
+    def test_learning_starts_uses_current_steps_after_increment(self):
+        with tempfile.TemporaryDirectory() as log_dir:
+            experiment = self.make_experiment(
+                OffPolicyExperimentConfig(
+                    max_steps=2,
+                    batch_size=2,
+                    replay_buffer_size=4,
+                    use_last_episode=False,
+                    log_dir=log_dir,
+                )
+            )
+            try:
+                batch = one_transition()
+
+                experiment.steps = 1
+                loss = experiment._learn_from_batch(batch, None)
+                self.assertEqual(loss, 0.0)
+                self.assertEqual(experiment.learner.train_calls, 0)
+
+                experiment.steps = 2
+                loss = experiment._learn_from_batch(batch, None)
+                self.assertEqual(loss, 7.0)
+                self.assertEqual(experiment.learner.train_calls, 1)
+            finally:
+                experiment.writer.close()
+
+    def test_warmup_steps_wraps_controller(self):
+        with tempfile.TemporaryDirectory() as log_dir:
+            experiment = self.make_experiment(
+                OffPolicyExperimentConfig(
+                    max_steps=1,
+                    warmup_steps=10,
+                    log_dir=log_dir,
+                )
+            )
+            try:
+                self.assertEqual(experiment.warmup_steps, 10)
+                self.assertIsInstance(experiment.runner.controller, WarmupController)
+                self.assertEqual(experiment.runner.controller.warmup_steps, 10)
+            finally:
+                experiment.writer.close()
+
+    def test_zero_warmup_steps_keeps_original_controller(self):
+        with tempfile.TemporaryDirectory() as log_dir:
+            experiment = self.make_experiment(
+                OffPolicyExperimentConfig(max_steps=1, log_dir=log_dir)
+            )
+            try:
+                self.assertNotIsInstance(experiment.runner.controller, WarmupController)
+            finally:
+                experiment.writer.close()
+
+    def test_negative_warmup_steps_is_rejected(self):
+        with tempfile.TemporaryDirectory() as log_dir:
+            with self.assertRaisesRegex(ValueError, "warmup_steps"):
+                self.make_experiment(
+                    OffPolicyExperimentConfig(
+                        max_steps=1,
+                        warmup_steps=-1,
+                        log_dir=log_dir,
+                    )
+                )

{drlab-0.2.1 → drlab-0.2.2}/tests/test_public_api.py

@@ -34,9 +34,10 @@ class PublicAPITest(unittest.TestCase):
             "StochasticController",
             "TransitionBatch",
             "ValueTargets",
+            "WarmupController",
         }
 
-        self.assertEqual(drlab.__version__, "0.2.1")
+        self.assertEqual(drlab.__version__, "0.2.2")
         self.assertTrue(expected.issubset(set(drlab.__all__)))
         for name in expected:
             self.assertTrue(hasattr(drlab, name), name)