Feature/improvements batch

docs/spelling_wordlist.txt

@@ -270,4 +270,4 @@ v_s
 v_s_
 obs
 obs_next
-
+dtype

examples/inverse/irl_gail.py

@@ -4,7 +4,7 @@
 import datetime
 import os
 import pprint
-from typing import SupportsFloat
+from typing import SupportsFloat, cast
 
 import d4rl
 import gymnasium as gym
@@ -16,6 +16,7 @@
 from torch.utils.tensorboard import SummaryWriter
 
 from tianshou.data import Batch, Collector, ReplayBuffer, VectorReplayBuffer
+from tianshou.data.types import RolloutBatchProtocol
 from tianshou.env import SubprocVectorEnv, VectorEnvNormObs
 from tianshou.policy import GAILPolicy
 from tianshou.policy.base import BasePolicy
@@ -185,12 +186,15 @@ def dist(loc_scale: tuple[torch.Tensor, torch.Tensor]) -> Distribution:
 
     for i in range(dataset_size):
         expert_buffer.add(
-            Batch(
-                obs=dataset["observations"][i],
-                act=dataset["actions"][i],
-                rew=dataset["rewards"][i],
-                done=dataset["terminals"][i],
-                obs_next=dataset["next_observations"][i],
+            cast(
+                RolloutBatchProtocol,
+                Batch(
+                    obs=dataset["observations"][i],
+                    act=dataset["actions"][i],
+                    rew=dataset["rewards"][i],
+                    done=dataset["terminals"][i],
+                    obs_next=dataset["next_observations"][i],
+                ),
             ),
         )
     print("dataset loaded")

test/base/env.py

@@ -147,14 +147,16 @@ def step(self, action: np.ndarray | int):  # type: ignore[no-untyped-def]  # cf.
         if self.index == self.size:
             self.terminated = True
             return self._get_state(), self._get_reward(), self.terminated, False, {}
+
+        info_dict = {"key": 1, "env": self}
         if action == 0:
             self.index = max(self.index - 1, 0)
             return (
                 self._get_state(),
                 self._get_reward(),
                 self.terminated,
                 False,
-                {"key": 1, "env": self} if self.dict_state else {},
+                info_dict,
             )
         if action == 1:
             self.index += 1
@@ -164,7 +166,7 @@ def step(self, action: np.ndarray | int):  # type: ignore[no-untyped-def]  # cf.
                 self._get_reward(),
                 self.terminated,
                 False,
-                {"key": 1, "env": self},
+                info_dict,
             )
         return None