Feature/extension collector buffer (#1196)

Adds important functionality to buffer and collector. The PR is very
large but I didn't want to split it up. It's easiest to review commit by
commit, and I think various people should have a look. One can also look
file by file. Together we can do this ;)

I'll edit the description when the review is done

@Trinkle23897: pls manly have a look at the the changes in buffer
related things, and if you want also in the computation of n_step
return. I had to slightly modify one of the tests that was changing the
private `_insertion_index` leading to a malformed buffer, which now
raises an error. Ofc you are very welcome to look at the rest as well :)

@opcode81 and @maxhuettenrauch : pls have a look at the extensions in
Collector. They are untested for now, wanted to get your opinion on the
design first. Also, a quick glance at the trainer would be nice

Ah, also @Trinkle23897: I think I found a bug in the PPO implementation,
see corresponding commit

@dantp-ai : the changes to the buffer here will make the task of fixing
slicing issues easier, especially the new names and additional comments.
Would also be happy about your review, if you have time!

docs/02_notebooks/L0_overview.ipynb

@@ -15,15 +15,6 @@
     "Before we get started, we must first install Tianshou's library and Gym environment by running the commands below. This tutorials will always keep up with the latest version of Tianshou since they also serve as a test for the latest version. If you are using an older version of Tianshou, please refer to the [documentation](https://tianshou.readthedocs.io/en/latest/) of your version.\n"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# !pip install tianshou gym"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -67,7 +58,7 @@
     "import gymnasium as gym\n",
     "import torch\n",
     "\n",
-    "from tianshou.data import Collector, VectorReplayBuffer\n",
+    "from tianshou.data import Collector, CollectStats, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
     "from tianshou.policy import PPOPolicy\n",
     "from tianshou.trainer import OnpolicyTrainer\n",
@@ -114,8 +105,12 @@
     ")\n",
     "\n",
     "# collector\n",
-    "train_collector = Collector(policy, train_envs, VectorReplayBuffer(20000, len(train_envs)))\n",
-    "test_collector = Collector(policy, test_envs)\n",
+    "train_collector = Collector[CollectStats](\n",
+    "    policy,\n",
+    "    train_envs,\n",
+    "    VectorReplayBuffer(20000, len(train_envs)),\n",
+    ")\n",
+    "test_collector = Collector[CollectStats](policy, test_envs)\n",
     "\n",
     "# trainer\n",
     "train_result = OnpolicyTrainer(\n",

docs/02_notebooks/L5_Collector.ipynb

@@ -58,7 +58,7 @@
     "import gymnasium as gym\n",
     "import torch\n",
     "\n",
-    "from tianshou.data import Collector, VectorReplayBuffer\n",
+    "from tianshou.data import Collector, CollectStats, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
     "from tianshou.policy import PGPolicy\n",
     "from tianshou.utils.net.common import Net\n",
@@ -94,7 +94,7 @@
     "    action_space=env.action_space,\n",
     "    action_scaling=False,\n",
     ")\n",
-    "test_collector = Collector(policy, test_envs)"
+    "test_collector = Collector[CollectStats](policy, test_envs)"
    ]
   },
   {
@@ -187,7 +187,7 @@
     "train_envs = DummyVectorEnv([lambda: gym.make(\"CartPole-v1\") for _ in range(train_env_num)])\n",
     "replayBuffer = VectorReplayBuffer(buffer_size, train_env_num)\n",
     "\n",
-    "train_collector = Collector(policy, train_envs, replayBuffer)"
+    "train_collector = Collector[CollectStats](policy, train_envs, replayBuffer)"
    ]
   },
   {

docs/02_notebooks/L6_Trainer.ipynb

@@ -54,12 +54,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-05-06T15:34:02.969675Z",
-     "start_time": "2024-05-06T15:34:00.747309Z"
-    },
     "editable": true,
     "id": "do-xZ-8B7nVH",
     "slideshow": {
@@ -77,7 +73,7 @@
     "import gymnasium as gym\n",
     "import torch\n",
     "\n",
-    "from tianshou.data import Collector, VectorReplayBuffer\n",
+    "from tianshou.data import Collector, CollectStats, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
     "from tianshou.policy import PGPolicy\n",
     "from tianshou.trainer import OnpolicyTrainer\n",
@@ -88,13 +84,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2024-05-06T15:34:07.536452Z",
-     "start_time": "2024-05-06T15:34:03.636670Z"
-    }
-   },
+   "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
     "train_env_num = 4\n",
@@ -131,8 +122,8 @@
     "\n",
     "# Create the replay buffer and the collector\n",
     "replayBuffer = VectorReplayBuffer(buffer_size, train_env_num)\n",
-    "test_collector = Collector(policy, test_envs)\n",
-    "train_collector = Collector(policy, train_envs, replayBuffer)"
+    "test_collector = Collector[CollectStats](policy, test_envs)\n",
+    "train_collector = Collector[CollectStats](policy, train_envs, replayBuffer)"
    ]
   },
   {

docs/02_notebooks/L7_Experiment.ipynb

@@ -71,7 +71,7 @@
     "import gymnasium as gym\n",
     "import torch\n",
     "\n",
-    "from tianshou.data import Collector, VectorReplayBuffer\n",
+    "from tianshou.data import Collector, CollectStats, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
     "from tianshou.policy import PPOPolicy\n",
     "from tianshou.trainer import OnpolicyTrainer\n",
@@ -202,12 +202,12 @@
    },
    "outputs": [],
    "source": [
-    "train_collector = Collector(\n",
+    "train_collector = Collector[CollectStats](\n",
     "    policy=policy,\n",
     "    env=train_envs,\n",
     "    buffer=VectorReplayBuffer(20000, len(train_envs)),\n",
     ")\n",
-    "test_collector = Collector(policy=policy, env=test_envs)"
+    "test_collector = Collector[CollectStats](policy=policy, env=test_envs)"
    ]
   },
   {

docs/spelling_wordlist.txt

@@ -282,3 +282,9 @@ autocompletion
 codebase
 indexable
 sliceable
+gaussian
+logprob
+monte
+carlo
+subclass
+subclassing

examples/atari/atari_c51.py

@@ -9,7 +9,7 @@
 from atari_network import C51
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import C51Policy
 from tianshou.policy.base import BasePolicy
@@ -112,8 +112,8 @@ def test_c51(args: argparse.Namespace = get_args()) -> None:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -173,7 +173,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_dqn.py

@@ -9,7 +9,7 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import DQNPolicy
 from tianshou.policy.base import BasePolicy
@@ -148,8 +148,8 @@ def main(args: argparse.Namespace = get_args()) -> None:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -215,7 +215,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_fqf.py

@@ -9,7 +9,7 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import FQFPolicy
 from tianshou.policy.base import BasePolicy
@@ -125,8 +125,8 @@ def test_fqf(args: argparse.Namespace = get_args()) -> None:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -186,7 +186,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_iqn.py

@@ -9,7 +9,7 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import IQNPolicy
 from tianshou.policy.base import BasePolicy
@@ -122,8 +122,8 @@ def test_iqn(args: argparse.Namespace = get_args()) -> None:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -183,7 +183,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_ppo.py

@@ -11,7 +11,7 @@
 from torch.distributions import Categorical
 from torch.optim.lr_scheduler import LambdaLR
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import ICMPolicy, PPOPolicy
 from tianshou.policy.base import BasePolicy
@@ -190,8 +190,8 @@ def dist(logits: torch.Tensor) -> Categorical:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -243,7 +243,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_qrdqn.py

@@ -9,7 +9,7 @@
 from atari_network import QRDQN
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import QRDQNPolicy
 from tianshou.policy.base import BasePolicy
@@ -116,8 +116,8 @@ def test_qrdqn(args: argparse.Namespace = get_args()) -> None:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -177,7 +177,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_rainbow.py

@@ -9,7 +9,12 @@
 from atari_network import Rainbow
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, PrioritizedVectorReplayBuffer, VectorReplayBuffer
+from tianshou.data import (
+    Collector,
+    CollectStats,
+    PrioritizedVectorReplayBuffer,
+    VectorReplayBuffer,
+)
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import C51Policy, RainbowPolicy
 from tianshou.policy.base import BasePolicy
@@ -142,8 +147,8 @@ def test_rainbow(args: argparse.Namespace = get_args()) -> None:
             weight_norm=not args.no_weight_norm,
         )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -213,7 +218,7 @@ def watch() -> None:
                 alpha=args.alpha,
                 beta=args.beta,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size

examples/atari/atari_sac.py

@@ -9,7 +9,7 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.data import Collector, CollectStats, VectorReplayBuffer
 from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import DiscreteSACPolicy, ICMPolicy
 from tianshou.policy.base import BasePolicy
@@ -173,8 +173,8 @@ def test_discrete_sac(args: argparse.Namespace = get_args()) -> None:
         stack_num=args.frames_stack,
     )
     # collector
-    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
-    test_collector = Collector(policy, test_envs, exploration_noise=True)
+    train_collector = Collector[CollectStats](policy, train_envs, buffer, exploration_noise=True)
+    test_collector = Collector[CollectStats](policy, test_envs, exploration_noise=True)
 
     # log
     now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
@@ -226,7 +226,7 @@ def watch() -> None:
                 save_only_last_obs=True,
                 stack_num=args.frames_stack,
             )
-            collector = Collector(policy, test_envs, buffer, exploration_noise=True)
+            collector = Collector[CollectStats](policy, test_envs, buffer, exploration_noise=True)
             result = collector.collect(n_step=args.buffer_size)
             print(f"Save buffer into {args.save_buffer_name}")
             # Unfortunately, pickle will cause oom with 1M buffer size