drlab 0.2.2__tar.gz → 0.2.3__tar.gz

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

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

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

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

{drlab-0.2.2 → drlab-0.2.3}/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.3"
 
 __all__ = [
     "__version__",

{drlab-0.2.2 → drlab-0.2.3}/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.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.2 → 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.2 → 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.2 → 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)
@@ -149,7 +149,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,

{drlab-0.2.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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),
         )

{drlab-0.2.2 → drlab-0.2.3}/src/drlab/runners/runner.py

@@ -34,7 +34,8 @@ class Runner:
         self.return_last_episode = return_last_episode
         self.ep_len, self.ep_ret = 0, 0.0
         self._ep_states, self._ep_actions, self._ep_rewards = [], [], []
-        self._ep_dones, self._ep_next_states, self._ep_returns = [], [], []
+        self._ep_terminated, self._ep_truncated = [], []
+        self._ep_next_states, self._ep_returns = [], []
 
     def _discounted_returns(self, rewards: list[float]) -> list[float]:
         out = [0.0] * len(rewards)
@@ -68,7 +69,8 @@ class Runner:
         states,
         actions,
         rewards,
-        dones,
+        terminated,
+        truncated,
         next_states,
         returns,
         as_numpy: bool,
@@ -76,18 +78,22 @@ class Runner:
         states = np.asarray(states, dtype=np.float32)
         actions = self._actions_to_array(actions)
         rewards = np.asarray(rewards, dtype=np.float32).reshape(-1, 1)
-        dones = np.asarray(dones, dtype=np.bool_).reshape(-1, 1)
+        terminated = np.asarray(terminated, dtype=np.bool_).reshape(-1, 1)
+        truncated = np.asarray(truncated, dtype=np.bool_).reshape(-1, 1)
         next_states = np.asarray(next_states, dtype=np.float32)
         returns = np.asarray(returns, dtype=np.float32).reshape(-1, 1)
 
         if as_numpy:
-            return TransitionBatch(states, actions, rewards, dones, next_states, returns)
+            return TransitionBatch(
+                states, actions, rewards, terminated, truncated, next_states, returns
+            )
 
         return TransitionBatch(
             states=th.from_numpy(states),
             actions=th.from_numpy(actions),
             rewards=th.from_numpy(rewards),
-            dones=th.from_numpy(dones),
+            terminated=th.from_numpy(terminated),
+            truncated=th.from_numpy(truncated),
             next_states=th.from_numpy(next_states),
             returns=th.from_numpy(returns),
         )
@@ -102,7 +108,9 @@ class Runner:
         last_episode = None
         ep_returns, ep_lengths = [], []
         ep_start = 0
-        states, actions, rewards, dones, next_states, returns = [], [], [], [], [], []
+        states, actions, rewards = [], [], []
+        terminated_flags, truncated_flags = [], []
+        next_states, returns = [], []
         run_one_episode = num_steps <= 0
         
         st = 0
@@ -113,14 +121,15 @@ class Runner:
             with th.inference_mode():
                 action_t = self.controller.choose(state_t)
             action = self._action_to_env(action_t)
-            next_state, reward, terminated, truncated, _ = self.env.step(action) 
-            done = terminated or truncated
+            next_state, reward, terminated, truncated, _ = self.env.step(action)
+            episode_end = terminated or truncated
 
             # 2. Store transition
             states.append(self.state)
             actions.append(action)
             rewards.append(reward)
-            dones.append(done)
+            terminated_flags.append(terminated)
+            truncated_flags.append(truncated)
             next_states.append(next_state)
             returns.append(0)
             self.ep_ret += reward
@@ -129,34 +138,37 @@ class Runner:
                 self._ep_states.append(self.state)
                 self._ep_actions.append(action)
                 self._ep_rewards.append(reward)
-                self._ep_dones.append(done)
+                self._ep_terminated.append(terminated)
+                self._ep_truncated.append(truncated)
                 self._ep_next_states.append(next_state)
                 self._ep_returns.append(0)
 
             # 3. Compute cumlative returns
-            if self.calculate_returns and (done or st == num_steps - 1):
+            if self.calculate_returns and (episode_end or st == num_steps - 1):
                 ep_rewards = rewards[ep_start: st + 1]
                 discounted_returns = self._discounted_returns(ep_rewards)
 
                 returns[ep_start: st + 1] = discounted_returns
-                if self.return_last_episode and done:
+                if self.return_last_episode and episode_end:
                     self._ep_returns = discounted_returns
 
             # 4. Update state
             self.state = next_state
-            if done:
+            if episode_end:
                 if self.return_last_episode:
                     last_episode = self._make_batch(
                         self._ep_states,
                         self._ep_actions,
                         self._ep_rewards,
-                        self._ep_dones,
+                        self._ep_terminated,
+                        self._ep_truncated,
                         self._ep_next_states,
                         self._ep_returns,
                         as_numpy,
                     )
                     self._ep_states, self._ep_actions, self._ep_rewards = [], [], []
-                    self._ep_dones, self._ep_next_states, self._ep_returns = [], [], []
+                    self._ep_terminated, self._ep_truncated = [], []
+                    self._ep_next_states, self._ep_returns = [], []
                 ep_start = st + 1
                 ep_returns.append(self.ep_ret)
                 ep_lengths.append(self.ep_len)
@@ -173,7 +185,8 @@ class Runner:
             states,
             actions,
             rewards,
-            dones,
+            terminated_flags,
+            truncated_flags,
             next_states,
             returns,
             as_numpy,

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

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

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

@@ -47,7 +47,7 @@ class LearnerSmokeTest(unittest.TestCase):
 
         loss = learner.train(
             rewards=th.tensor([[1.0], [0.0], [1.0]]),
-            dones=th.tensor([[False], [True], [False]]),
+            terminated=th.tensor([[False], [True], [False]]),
             states=th.randn(3, 4),
             actions=th.tensor([[0], [1], [0]], dtype=th.long),
             next_states=th.randn(3, 4),
@@ -90,7 +90,7 @@ class LearnerSmokeTest(unittest.TestCase):
 
         loss = learner.train(
             rewards=th.tensor([[1.0], [0.0], [1.0]]),
-            dones=th.tensor([[False], [True], [False]]),
+            terminated=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),
@@ -125,7 +125,7 @@ class LearnerSmokeTest(unittest.TestCase):
         )
         calls = []
 
-        def actor_l2(model, rewards, dones, states, actions, next_states):
+        def actor_l2(model, rewards, terminated, states, actions, next_states):
             calls.append(model)
             return sum(param.square().sum() for param in model.parameters())
 
@@ -145,7 +145,7 @@ class LearnerSmokeTest(unittest.TestCase):
 
         learner.train(
             rewards=th.tensor([[1.0], [0.0], [1.0]]),
-            dones=th.tensor([[False], [True], [False]]),
+            terminated=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),
@@ -171,7 +171,7 @@ class LearnerSmokeTest(unittest.TestCase):
 
         loss = learner.train(
             rewards=th.tensor([[1.0], [0.5], [0.0]]),
-            dones=th.tensor([[False], [False], [True]]),
+            terminated=th.tensor([[False], [False], [True]]),
             states=th.randn(3, 4),
             actions=th.tensor([[0], [1], [0]], dtype=th.long),
             next_states=th.randn(3, 4),
@@ -199,7 +199,7 @@ class LearnerSmokeTest(unittest.TestCase):
 
         loss = learner.train(
             rewards=th.tensor([[1.0], [0.5], [0.0]]),
-            dones=th.tensor([[False], [False], [True]]),
+            terminated=th.tensor([[False], [False], [True]]),
             states=th.randn(3, 4),
             actions=th.tensor([[0], [1], [0]], dtype=th.long),
             next_states=th.randn(3, 4),
@@ -227,7 +227,7 @@ class LearnerSmokeTest(unittest.TestCase):
 
         loss = learner.train(
             rewards=th.tensor([[1.0], [0.5], [0.0]]),
-            dones=th.tensor([[False], [False], [True]]),
+            terminated=th.tensor([[False], [False], [True]]),
             states=th.randn(3, 4),
             actions=th.tensor([[0], [1], [0]], dtype=th.long),
             next_states=th.randn(3, 4),

{drlab-0.2.2 → drlab-0.2.3}/tests/test_off_policy_experiment.py

@@ -49,7 +49,7 @@ class CountingLearner(OffPolicyLearner):
     def train(
         self,
         rewards: th.Tensor,
-        dones: th.Tensor,
+        terminated: th.Tensor,
         states: th.Tensor,
         actions: th.Tensor,
         next_states: th.Tensor,
@@ -63,7 +63,8 @@ def one_transition() -> 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_),
+        terminated=np.asarray([[False]], dtype=np.bool_),
+        truncated=np.asarray([[False]], dtype=np.bool_),
         next_states=np.asarray([[1.0]], dtype=np.float32),
         returns=np.asarray([[1.0]], dtype=np.float32),
     )

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

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

{drlab-0.2.2 → drlab-0.2.3}/tests/test_replay.py

@@ -12,7 +12,8 @@ class TransitionBatchTest(unittest.TestCase):
             states=th.zeros(2, 3),
             actions=th.zeros(2, 1, dtype=th.long),
             rewards=th.ones(2, 1),
-            dones=th.zeros(2, 1, dtype=th.bool),
+            terminated=th.zeros(2, 1, dtype=th.bool),
+            truncated=th.zeros(2, 1, dtype=th.bool),
             next_states=th.ones(2, 3),
             returns=th.full((2, 1), 2.0),
         )
@@ -28,7 +29,8 @@ class TransitionBatchTest(unittest.TestCase):
             states=np.zeros((2, 3), dtype=np.float32),
             actions=np.zeros((2, 1), dtype=np.int64),
             rewards=np.ones((2, 1), dtype=np.float32),
-            dones=np.zeros((2, 1), dtype=np.bool_),
+            terminated=np.zeros((2, 1), dtype=np.bool_),
+            truncated=np.zeros((2, 1), dtype=np.bool_),
             next_states=np.ones((2, 3), dtype=np.float32),
             returns=np.full((2, 1), 2.0, dtype=np.float32),
         )
@@ -37,7 +39,8 @@ class TransitionBatchTest(unittest.TestCase):
 
         self.assertEqual(moved.states.dtype, th.float32)
         self.assertEqual(moved.actions.dtype, th.int64)
-        self.assertEqual(moved.dones.dtype, th.bool)
+        self.assertEqual(moved.terminated.dtype, th.bool)
+        self.assertEqual(moved.truncated.dtype, th.bool)
         self.assertTrue(th.equal(moved.returns, th.full((2, 1), 2.0)))
 
     def test_cat_concatenates_every_field(self):
@@ -45,7 +48,8 @@ class TransitionBatchTest(unittest.TestCase):
             states=th.zeros(1, 2),
             actions=th.zeros(1, 1, dtype=th.long),
             rewards=th.ones(1, 1),
-            dones=th.zeros(1, 1, dtype=th.bool),
+            terminated=th.zeros(1, 1, dtype=th.bool),
+            truncated=th.ones(1, 1, dtype=th.bool),
             next_states=th.ones(1, 2),
             returns=th.ones(1, 1),
         )
@@ -53,7 +57,8 @@ class TransitionBatchTest(unittest.TestCase):
             states=th.ones(2, 2),
             actions=th.ones(2, 1, dtype=th.long),
             rewards=th.full((2, 1), 2.0),
-            dones=th.ones(2, 1, dtype=th.bool),
+            terminated=th.ones(2, 1, dtype=th.bool),
+            truncated=th.zeros(2, 1, dtype=th.bool),
             next_states=th.full((2, 2), 3.0),
             returns=th.full((2, 1), 4.0),
         )
@@ -63,6 +68,8 @@ class TransitionBatchTest(unittest.TestCase):
         self.assertEqual(combined.states.shape, (3, 2))
         self.assertEqual(combined.actions.shape, (3, 1))
         self.assertTrue(th.equal(combined.rewards.squeeze(-1), th.tensor([1.0, 2.0, 2.0])))
+        episode_ends = combined.terminated | combined.truncated
+        self.assertTrue(th.equal(episode_ends, th.ones(3, 1, dtype=th.bool)))
 
 
 class ReplayBufferTest(unittest.TestCase):
@@ -75,7 +82,8 @@ class ReplayBufferTest(unittest.TestCase):
             states=states,
             actions=np.asarray([0, 1, 0]),
             rewards=np.asarray([1.0, 2.0, 3.0]),
-            dones=np.asarray([False, False, True]),
+            terminated=np.asarray([False, False, True]),
+            truncated=np.asarray([False, True, False]),
             next_states=next_states,
             returns=np.asarray([6.0, 5.0, 3.0]),
         )
@@ -87,7 +95,9 @@ class ReplayBufferTest(unittest.TestCase):
         self.assertIsInstance(all_data, TransitionBatch)
         self.assertEqual(all_data.states.shape, (3, 2))
         self.assertEqual(all_data.actions.dtype, th.int64)
-        self.assertEqual(all_data.dones.dtype, th.bool)
+        self.assertEqual(all_data.terminated.dtype, th.bool)
+        self.assertEqual(all_data.truncated.dtype, th.bool)
+        self.assertTrue(all_data.truncated[1].item())
         self.assertIsInstance(sample, TransitionBatch)
         self.assertEqual(sample.states.shape, (2, 2))
 
@@ -98,7 +108,8 @@ class ReplayBufferTest(unittest.TestCase):
             states=np.asarray([[0], [1]], dtype=np.float32),
             actions=np.asarray([0, 1]),
             rewards=np.asarray([0.0, 1.0]),
-            dones=np.asarray([False, False]),
+            terminated=np.asarray([False, False]),
+            truncated=np.asarray([False, False]),
             next_states=np.asarray([[1], [2]], dtype=np.float32),
             returns=np.asarray([1.0, 1.0]),
         )
@@ -106,7 +117,8 @@ class ReplayBufferTest(unittest.TestCase):
             states=np.asarray([[2], [3]], dtype=np.float32),
             actions=np.asarray([0, 1]),
             rewards=np.asarray([2.0, 3.0]),
-            dones=np.asarray([False, True]),
+            terminated=np.asarray([False, True]),
+            truncated=np.asarray([False, False]),
             next_states=np.asarray([[3], [4]], dtype=np.float32),
             returns=np.asarray([5.0, 3.0]),
         )

{drlab-0.2.2 → drlab-0.2.3}/tests/test_runner.py

@@ -48,6 +48,12 @@ class ThreeStepEnv(gym.Env):
         return obs, reward, terminated, truncated, {}
 
 
+class TruncatingEnv(ThreeStepEnv):
+    def step(self, action):
+        obs, reward, _, _, info = super().step(action)
+        return obs, reward, False, self.t >= 3, info
+
+
 class RunnerTest(unittest.TestCase):
     def test_run_zero_collects_one_complete_episode_with_returns(self):
         env = ThreeStepEnv()
@@ -65,6 +71,23 @@ class RunnerTest(unittest.TestCase):
         self.assertEqual(ep_lengths, [3])
         self.assertTrue(th.equal(batch.returns.squeeze(-1), th.tensor([6.0, 4.0, 2.0])))
         self.assertTrue(th.equal(last_episode.returns.squeeze(-1), th.tensor([6.0, 4.0, 2.0])))
+        self.assertTrue(batch.terminated[-1].item())
+        self.assertFalse(batch.truncated.any().item())
+
+    def test_truncation_ends_episode_without_marking_transition_terminal(self):
+        env = TruncatingEnv()
+        model = FixedLogits([0.0, 1.0])
+        controller = GreedyController(model, num_actions=2)
+        runner = Runner(env, controller, calculate_returns=False)
+
+        batch, ep_returns, ep_lengths, last_episode = runner.run(0)
+
+        self.assertEqual(ep_returns, [6.0])
+        self.assertEqual(ep_lengths, [3])
+        self.assertIsNotNone(last_episode)
+        self.assertFalse(batch.terminated.any().item())
+        self.assertTrue(batch.truncated[-1].item())
+        self.assertTrue((batch.terminated | batch.truncated)[-1].item())
 
     def test_run_positive_steps_can_return_partial_batch(self):
         env = ThreeStepEnv()