drlab 0.2.2__tar.gz → 0.2.4__tar.gz

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

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

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

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

{drlab-0.2.2 → 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.2"
+__version__ = "0.2.4"
 
 __all__ = [
     "__version__",

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

@@ -130,7 +130,8 @@ class OffPolicyExperiment:
             states=batch.states,
             actions=batch.actions,
             rewards=batch.rewards,
-            dones=batch.dones,
+            terminated=batch.terminated,
+            truncated=batch.truncated,
             next_states=batch.next_states,
             returns=batch.returns,
         )
@@ -142,7 +143,7 @@ class OffPolicyExperiment:
         for _ in range(self.grad_repeats):
             mb = self._make_minibatch(batch, last_episode)
             total_loss += self.learner.train(
-                mb.rewards, mb.dones, mb.states, mb.actions, mb.next_states
+                mb.rewards, mb.terminated, mb.states, mb.actions, mb.next_states
             )
 
         return total_loss / self.grad_repeats

{drlab-0.2.2 → drlab-0.2.4}/src/drlab/experiments/on_policy_experiment.py

@@ -65,7 +65,7 @@ class OnPolicyExperiment:
             # 2. Learn from batch
             loss = self.learner.train(
                 batch.rewards,
-                batch.dones,
+                batch.terminated,
                 batch.states,
                 batch.actions,
                 batch.next_states,

{drlab-0.2.2 → 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
 
@@ -88,7 +102,7 @@ class OffPolicyLearner(ABC):
         self,
         model: th.nn.Module,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
@@ -102,7 +116,7 @@ class OffPolicyLearner(ABC):
             value = regularizer(
                 model,
                 rewards,
-                dones,
+                terminated,
                 states,
                 actions,
                 next_states,
@@ -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()
@@ -139,9 +155,9 @@ class OffPolicyLearner(ABC):
     def train(
         self,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
     ) -> float:
-        pass
+        pass

{drlab-0.2.2 → 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.2 → drlab-0.2.4}/src/drlab/learners/off_policy/dqn.py

@@ -58,7 +58,7 @@ class DQNLearner(OffPolicyLearner):
     def train(
         self,
         rewards: th.Tensor,      # float32, [B,1]
-        dones: th.Tensor,        # bool or float(0/1), [B,1]
+        terminated: th.Tensor,   # bool or float(0/1), [B,1]
         states: th.Tensor,       # float32, [B, obs_dim] or [B,C,H,W]
         actions: th.Tensor,      # int64, [B,1]
         next_states: th.Tensor,  # float32, same as states
@@ -66,8 +66,10 @@ class DQNLearner(OffPolicyLearner):
         self.model.train(True)
 
         # Compute targets
-        not_done = 1.0 - dones.float()
-        targets = rewards + self.gamma * (not_done * self.next_values(next_states))
+        not_terminated = 1.0 - terminated.float()
+        targets = rewards + self.gamma * (
+            not_terminated * self.next_values(next_states)
+        )
         pred = self.current_values(states, actions)
         
         # Compute loss
@@ -75,7 +77,7 @@ class DQNLearner(OffPolicyLearner):
         reg_loss = self.regularization_loss(
             self.model,
             rewards=rewards,
-            dones=dones,
+            terminated=terminated,
             states=states,
             actions=actions,
             next_states=next_states,
@@ -87,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.2 → 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,26 +72,30 @@ 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
 
     def train(
         self,
         rewards: th.Tensor,      # float32, [B,1]
-        dones: th.Tensor,        # bool or float(0/1), [B,1]
+        terminated: th.Tensor,   # bool or float(0/1), [B,1]
         states: th.Tensor,       # float32, [B, obs_dim] or [B,C,H,W]
         actions: th.Tensor,      # float32, [B, action_dim]
         next_states: th.Tensor,  # float32, same as states
@@ -110,9 +115,9 @@ class SACLearner(OffPolicyLearner):
             target_q2 = self.critic2_target(next_action_states)
             target_q = th.min(target_q1, target_q2)
 
-            not_done = 1.0 - dones.float()
+            not_terminated = 1.0 - terminated.float()
             targets = rewards + self.config.gamma * (
-                not_done * (target_q - self.alpha.detach() * next_log_probs)
+                not_terminated * (target_q - self.alpha.detach() * next_log_probs)
             )
 
         action_states = th.cat([states, actions], dim=-1)
@@ -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)
@@ -149,7 +158,7 @@ class SACLearner(OffPolicyLearner):
         reg_loss = self.regularization_loss(
             self.actor,
             rewards=rewards,
-            dones=dones,
+            terminated=terminated,
             states=states,
             actions=actions,
             next_states=next_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.2 → drlab-0.2.4}/src/drlab/learners/on_policy/actor_critic.py

@@ -28,14 +28,14 @@ class ActorCriticLearner(OnPolicyLearner):
         self,
         returns: th.Tensor,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         values: th.Tensor,
         next_values: th.Tensor,
     ) -> th.Tensor:
         if self.config.value_targets == ValueTargets.RETURNS:
             targets = returns
         elif self.config.value_targets == ValueTargets.TD:
-            targets = rewards + self.config.gamma * (~dones * next_values)
+            targets = rewards + self.config.gamma * (~terminated * next_values)
         else:
             raise ValueError(f"Unknown value_targets: {self.config.value_targets}")
         return self.config.value_criterion(values, targets)
@@ -44,13 +44,13 @@ class ActorCriticLearner(OnPolicyLearner):
             self, 
             returns: th.Tensor, 
             rewards: th.Tensor,
-            dones: th.Tensor,
+            terminated: th.Tensor,
             values: th.Tensor,
             next_values: th.Tensor
         ) -> th.Tensor:
         advantages = None
         if self.config.advantage_bootstrap:
-            advantages = rewards + self.config.gamma * (~dones * next_values)
+            advantages = rewards + self.config.gamma * (~terminated * next_values)
         else: 
             advantages = returns
         if self.config.use_bias:
@@ -74,7 +74,7 @@ class ActorCriticLearner(OnPolicyLearner):
     def train(
         self,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
@@ -106,7 +106,7 @@ class ActorCriticLearner(OnPolicyLearner):
         advantages = self._advantages(
             returns=returns,
             rewards=rewards,
-            dones=dones,
+            terminated=terminated,
             values=values,
             next_values=next_values,
         )
@@ -119,7 +119,7 @@ class ActorCriticLearner(OnPolicyLearner):
             value_loss = self.config.value_lambda * self._value_loss(
                 returns=returns,
                 rewards=rewards,
-                dones=dones,
+                terminated=terminated,
                 values=values,
                 next_values=next_values,
             )
@@ -128,7 +128,7 @@ class ActorCriticLearner(OnPolicyLearner):
 
         reg_loss = self._regularization_loss(
             rewards=rewards,
-            dones=dones,
+            terminated=terminated,
             states=states,
             actions=actions,
             next_states=next_states,

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

@@ -50,7 +50,7 @@ class OnPolicyLearner:
     def _regularization_loss(
         self,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
@@ -59,7 +59,7 @@ class OnPolicyLearner:
             return 0.0
 
         return sum(
-            lam * reg(self.actor, rewards, dones, states, actions, next_states)
+            lam * reg(self.actor, rewards, terminated, states, actions, next_states)
             for reg, lam in zip(self.config.regularizers, self.config.reg_lams)
         )
 
@@ -67,7 +67,7 @@ class OnPolicyLearner:
         self,
         returns: th.Tensor,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         values: th.Tensor,
         next_values: th.Tensor,
     ) -> th.Tensor:
@@ -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()
 
@@ -106,7 +106,7 @@ class OnPolicyLearner:
             self, 
             returns: th.Tensor, 
             rewards: th.Tensor,
-            dones: th.Tensor,
+            terminated: th.Tensor,
             values: th.Tensor,
             next_values: th.Tensor
         ) -> th.Tensor:
@@ -115,7 +115,7 @@ class OnPolicyLearner:
     def train(
         self,
         rewards: th.Tensor,      # float32, [B,1]
-        dones: th.Tensor,        # bool or float(0/1), [B,1]
+        terminated: th.Tensor,   # bool or float(0/1), [B,1]
         states: th.Tensor,       # float32, [B, obs_dim] or [B,C,H,W]
         actions: th.Tensor,      # int64, [B,1]
         next_states: th.Tensor,  # float32, same as states

{drlab-0.2.2 → 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.2 → drlab-0.2.4}/src/drlab/learners/on_policy/ppo.py

@@ -38,7 +38,7 @@ class PPOLearner(ActorCriticLearner):
     def train(
         self,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
@@ -69,7 +69,7 @@ class PPOLearner(ActorCriticLearner):
             advantages = self._advantages(
                 returns=returns,
                 rewards=rewards,
-                dones=dones,
+                terminated=terminated,
                 values=old_values,
                 next_values=next_values,
             )
@@ -96,7 +96,7 @@ class PPOLearner(ActorCriticLearner):
                 value_loss = self.config.value_lambda * self._value_loss(
                     returns=returns,
                     rewards=rewards,
-                    dones=dones,
+                    terminated=terminated,
                     values=values,
                     next_values=next_values,
                 )
@@ -105,7 +105,7 @@ class PPOLearner(ActorCriticLearner):
 
             reg_loss = self._regularization_loss(
                 rewards=rewards,
-                dones=dones,
+                terminated=terminated,
                 states=states,
                 actions=actions,
                 next_states=next_states,

{drlab-0.2.2 → drlab-0.2.4}/src/drlab/learners/on_policy/reinforce.py

@@ -18,7 +18,7 @@ class ReinforceLearner(OnPolicyLearner):
         self,
         returns: th.Tensor,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         values: th.Tensor | None,
         next_values: th.Tensor | None,
     ) -> th.Tensor:
@@ -34,7 +34,7 @@ class ReinforceLearner(OnPolicyLearner):
     def train(
         self,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
@@ -50,7 +50,7 @@ class ReinforceLearner(OnPolicyLearner):
         advantages = self._advantages(
             returns=returns,
             rewards=rewards,
-            dones=dones,
+            terminated=terminated,
             values=None,
             next_values=None,
         )
@@ -60,7 +60,7 @@ class ReinforceLearner(OnPolicyLearner):
         entropy_loss = self._entropy_loss(logits)
         reg_loss = self._regularization_loss(
             rewards=rewards,
-            dones=dones,
+            terminated=terminated,
             states=states,
             actions=actions,
             next_states=next_states,