Merge branch 'master' into feature/random-agent

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