drlab 0.2.1__tar.gz → 0.2.3__tar.gz

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

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

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

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

{drlab-0.2.1 → drlab-0.2.3}/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.3"
 
 __all__ = [
     "__version__",
@@ -68,4 +69,5 @@ __all__ = [
     "StochasticController",
     "TransitionBatch",
     "ValueTargets",
+    "WarmupController",
 ]

{drlab-0.2.1 → drlab-0.2.3}/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.3/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.3}/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,17 +120,22 @@ 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,
             actions=batch.actions,
             rewards=batch.rewards,
-            dones=batch.dones,
+            terminated=batch.terminated,
+            truncated=batch.truncated,
             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
@@ -115,19 +143,19 @@ 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
 
     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 +163,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 +184,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.3}/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.1 → drlab-0.2.3}/src/drlab/learners/off_policy/base.py

@@ -88,7 +88,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 +102,7 @@ class OffPolicyLearner(ABC):
             value = regularizer(
                 model,
                 rewards,
-                dones,
+                terminated,
                 states,
                 actions,
                 next_states,
@@ -139,9 +139,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.1 → drlab-0.2.3}/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,

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

@@ -90,7 +90,7 @@ class SACLearner(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,      # float32, [B, action_dim]
         next_states: th.Tensor,  # float32, same as states
@@ -110,9 +110,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)
@@ -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,
+            terminated=terminated,
+            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.3}/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.1 → drlab-0.2.3}/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:
@@ -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.1 → drlab-0.2.3}/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.1 → drlab-0.2.3}/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,

{drlab-0.2.1 → drlab-0.2.3}/src/drlab/replay/replay_buffer.py

@@ -27,7 +27,8 @@ class ReplayBuffer:
         self.actions = np.zeros((capacity, *action_shape), dtype=action_dtype)
 
         self.rewards = np.zeros((capacity, 1), dtype=np.float32)
-        self.dones = np.zeros((capacity, 1), dtype=np.bool_)
+        self.terminated = np.zeros((capacity, 1), dtype=np.bool_)
+        self.truncated = np.zeros((capacity, 1), dtype=np.bool_)
         self.returns = np.zeros((capacity, 1), dtype=np.float32)
 
         self.ptr = 0
@@ -41,7 +42,8 @@ class ReplayBuffer:
         states: np.ndarray,       # [B, *obs_shape]
         actions: np.ndarray,      # [B] or [B,1]
         rewards: np.ndarray,      # [B] or [B,1]
-        dones: np.ndarray,        # [B] or [B,1] (bool)
+        terminated: np.ndarray,   # [B] or [B,1] (bool)
+        truncated: np.ndarray,    # [B] or [B,1] (bool)
         next_states: np.ndarray,  # [B, *obs_shape]
         returns: np.ndarray       # [B, 1]
 
@@ -56,7 +58,8 @@ class ReplayBuffer:
         else:
             actions = np.asarray(actions, dtype=np.int64).reshape(B, 1)
         rewards = np.asarray(rewards, dtype=np.float32).reshape(B, 1)
-        dones = np.asarray(dones, dtype=np.bool_).reshape(B, 1)
+        terminated = np.asarray(terminated, dtype=np.bool_).reshape(B, 1)
+        truncated = np.asarray(truncated, dtype=np.bool_).reshape(B, 1)
         returns = np.asarray(returns, dtype=np.float32).reshape(B, 1)
 
         idx = (self.ptr + np.arange(B)) % self.capacity
@@ -65,7 +68,8 @@ class ReplayBuffer:
         self.next_states[idx] = next_states
         self.actions[idx] = actions
         self.rewards[idx] = rewards
-        self.dones[idx] = dones
+        self.terminated[idx] = terminated
+        self.truncated[idx] = truncated
         self.returns[idx] = returns
 
         self.ptr = (self.ptr + B) % self.capacity
@@ -76,9 +80,12 @@ class ReplayBuffer:
         next_states = th.from_numpy(self.next_states[:self.size]).to(self.device)
         actions = th.from_numpy(self.actions[:self.size]).to(self.device)
         rewards = th.from_numpy(self.rewards[:self.size]).to(self.device)
-        dones = th.from_numpy(self.dones[:self.size]).to(self.device)
+        terminated = th.from_numpy(self.terminated[:self.size]).to(self.device)
+        truncated = th.from_numpy(self.truncated[:self.size]).to(self.device)
         returns = th.from_numpy(self.returns[:self.size]).to(self.device)
-        return TransitionBatch(states, actions, rewards, dones, next_states, returns)
+        return TransitionBatch(
+            states, actions, rewards, terminated, truncated, next_states, returns
+        )
 
     def sample(self, batch_size: int) -> TransitionBatch:
         if self.size == 0:
@@ -90,7 +97,10 @@ class ReplayBuffer:
         next_states = th.from_numpy(self.next_states[idx]).to(self.device)
         actions = th.from_numpy(self.actions[idx]).to(self.device)               # long already
         rewards = th.from_numpy(self.rewards[idx]).to(self.device)               # float32
-        dones = th.from_numpy(self.dones[idx]).to(self.device)                   # bool
+        terminated = th.from_numpy(self.terminated[idx]).to(self.device)         # bool
+        truncated = th.from_numpy(self.truncated[idx]).to(self.device)           # bool
         returns = th.from_numpy(self.returns[idx]).to(self.device)               # float32
 
-        return TransitionBatch(states, actions, rewards, dones, next_states, returns)
+        return TransitionBatch(
+            states, actions, rewards, terminated, truncated, next_states, returns
+        )

{drlab-0.2.1 → drlab-0.2.3}/src/drlab/replay/transition_batch.py

@@ -8,7 +8,8 @@ class TransitionBatch:
     states: th.Tensor | np.ndarray       # [B, obs_dim]
     actions: th.Tensor | np.ndarray      # [B, 1] (int64/float32)
     rewards: th.Tensor | np.ndarray      # [B, 1] (float)
-    dones: th.Tensor | np.ndarray        # [B, 1] (bool)
+    terminated: th.Tensor | np.ndarray   # [B, 1] (bool), disables bootstrapping
+    truncated: th.Tensor | np.ndarray    # [B, 1] (bool), ends an episode only
     next_states: th.Tensor | np.ndarray  # [B, obs_dim]
     returns: th.Tensor | np.ndarray      # [B, 1] (float)
 
@@ -24,7 +25,8 @@ class TransitionBatch:
             states=move(self.states),
             actions=move(self.actions),
             rewards=move(self.rewards),
-            dones=move(self.dones),
+            terminated=move(self.terminated),
+            truncated=move(self.truncated),
             next_states=move(self.next_states),
             returns=move(self.returns),
         )
@@ -34,7 +36,8 @@ class TransitionBatch:
             states=th.cat([self.states, other.states], dim=0),
             actions=th.cat([self.actions, other.actions], dim=0),
             rewards=th.cat([self.rewards, other.rewards], dim=0),
-            dones=th.cat([self.dones, other.dones], dim=0),
+            terminated=th.cat([self.terminated, other.terminated], dim=0),
+            truncated=th.cat([self.truncated, other.truncated], dim=0),
             next_states=th.cat([self.next_states, other.next_states], dim=0),
             returns=th.cat([self.returns, other.returns], dim=0),
         )