Fix mypy issues in tests and examples (thu-ml#1077)

Closes thu-ml#952 - `SamplingConfig` supports `batch_size=None`. thu-ml#1077 - tests and examples are covered by `mypy`. thu-ml#1077 - `NetBase` is more used, stricter typing by making it generic. thu-ml#1077 - `utils.net.common.Recurrent` now receives and returns a `RecurrentStateBatch` instead of a dict. thu-ml#1077 --------- Co-authored-by: Michael Panchenko <m.panchenko@appliedai.de>
ZhengLi1314 · Apr 15, 2024 · 3dc5426 · 3dc5426
1 parent a70b532
commit 3dc5426
Show file tree

Hide file tree

Showing 106 changed files with 1,265 additions and 903 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -41,7 +41,7 @@ repos:
         pass_filenames: false
       - id: mypy
         name: mypy
-        entry: poetry run mypy tianshou
+        entry: poetry run mypy tianshou examples test
         # filenames should not be passed as they would collide with the config in pyproject.toml
         pass_filenames: false
         files: '^tianshou(/[^/]*)*/[^/]*\.py$'

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@
 - `Collector`s can now be closed, and their reset is more granular. #1063
 - Trainers can control whether collectors should be reset prior to training. #1063
 - Convenience constructor for `CollectStats` called `with_autogenerated_stats`. #1063
+- `SamplingConfig` supports `batch_size=None`. #1077
 
 ### Internal Improvements
 - `Collector`s rely less on state, the few stateful things are stored explicitly instead of through a `.data` attribute. #1063
@@ -20,6 +21,8 @@ instead of just `nn.Module`. #1032
 - Simplified `PGPolicy` forward by unifying the `dist_fn` interface (see associated breaking change). #1032
 - Use `.mode` of distribution instead of relying on knowledge of the distribution type. #1032
 - Exception no longer raised on `len` of empty `Batch`. #1084
+- tests and examples are covered by `mypy`. #1077
+- `NetBase` is more used, stricter typing by making it generic. #1077
 
 ### Breaking Changes
 
@@ -30,10 +33,10 @@ expicitly or pass `reset_before_collect=True` . #1063
 - Fixed `iter(Batch(...)` which now behaves the same way as `Batch(...).__iter__()`. Can be considered a bugfix. #1063
 - Changed interface of `dist_fn` in `PGPolicy` and all subclasses to take a single argument in both
 continuous and discrete cases. #1032
+- `utils.net.common.Recurrent` now receives and returns a `RecurrentStateBatch` instead of a dict. #1077
 
 ### Tests
 - Fixed env seeding it test_sac_with_il.py so that the test doesn't fail randomly. #1081
 
 
 Started after v1.0.0
-
diff --git a/docs/02_notebooks/L0_overview.ipynb b/docs/02_notebooks/L0_overview.ipynb
@@ -17,14 +17,14 @@
   },
   {
    "cell_type": "code",
-   "outputs": [],
-   "source": [
-    "# !pip install tianshou gym"
-   ],
+   "execution_count": null,
    "metadata": {
     "collapsed": false
    },
-   "execution_count": 0
+   "outputs": [],
+   "source": [
+    "# !pip install tianshou gym"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -71,7 +71,7 @@
     "\n",
     "from tianshou.data import Collector, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
-    "from tianshou.policy import BasePolicy, PPOPolicy\n",
+    "from tianshou.policy import PPOPolicy\n",
     "from tianshou.trainer import OnpolicyTrainer\n",
     "from tianshou.utils.net.common import ActorCritic, Net\n",
     "from tianshou.utils.net.discrete import Actor, Critic\n",
@@ -106,8 +106,7 @@
     "\n",
     "# PPO policy\n",
     "dist = torch.distributions.Categorical\n",
-    "policy: BasePolicy\n",
-    "policy = PPOPolicy(\n",
+    "policy: PPOPolicy = PPOPolicy(\n",
     "    actor=actor,\n",
     "    critic=critic,\n",
     "    optim=optim,\n",

diff --git a/docs/02_notebooks/L5_Collector.ipynb b/docs/02_notebooks/L5_Collector.ipynb
@@ -60,7 +60,7 @@
     "\n",
     "from tianshou.data import Collector, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
-    "from tianshou.policy import BasePolicy, PGPolicy\n",
+    "from tianshou.policy import PGPolicy\n",
     "from tianshou.utils.net.common import Net\n",
     "from tianshou.utils.net.discrete import Actor"
    ]
@@ -87,8 +87,7 @@
     "actor = Actor(net, env.action_space.n)\n",
     "optim = torch.optim.Adam(actor.parameters(), lr=0.0003)\n",
     "\n",
-    "policy: BasePolicy\n",
-    "policy = PGPolicy(\n",
+    "policy: PGPolicy = PGPolicy(\n",
     "    actor=actor,\n",
     "    optim=optim,\n",
     "    dist_fn=torch.distributions.Categorical,\n",

diff --git a/docs/02_notebooks/L6_Trainer.ipynb b/docs/02_notebooks/L6_Trainer.ipynb
@@ -75,7 +75,7 @@
     "\n",
     "from tianshou.data import Collector, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
-    "from tianshou.policy import BasePolicy, PGPolicy\n",
+    "from tianshou.policy import PGPolicy\n",
     "from tianshou.trainer import OnpolicyTrainer\n",
     "from tianshou.utils.net.common import Net\n",
     "from tianshou.utils.net.discrete import Actor"
@@ -110,9 +110,8 @@
     "actor = Actor(net, env.action_space.n)\n",
     "optim = torch.optim.Adam(actor.parameters(), lr=0.001)\n",
     "\n",
-    "policy: BasePolicy\n",
     "# We choose to use REINFORCE algorithm, also known as Policy Gradient\n",
-    "policy = PGPolicy(\n",
+    "policy: PGPolicy = PGPolicy(\n",
     "    actor=actor,\n",
     "    optim=optim,\n",
     "    dist_fn=torch.distributions.Categorical,\n",

diff --git a/docs/02_notebooks/L7_Experiment.ipynb b/docs/02_notebooks/L7_Experiment.ipynb
@@ -73,7 +73,7 @@
     "\n",
     "from tianshou.data import Collector, VectorReplayBuffer\n",
     "from tianshou.env import DummyVectorEnv\n",
-    "from tianshou.policy import BasePolicy, PPOPolicy\n",
+    "from tianshou.policy import PPOPolicy\n",
     "from tianshou.trainer import OnpolicyTrainer\n",
     "from tianshou.utils.net.common import ActorCritic, Net\n",
     "from tianshou.utils.net.discrete import Actor, Critic\n",
@@ -164,8 +164,7 @@
    "outputs": [],
    "source": [
     "dist = torch.distributions.Categorical\n",
-    "policy: BasePolicy\n",
-    "policy = PPOPolicy(\n",
+    "policy: PPOPolicy = PPOPolicy(\n",
     "    actor=actor,\n",
     "    critic=critic,\n",
     "    optim=optim,\n",

diff --git a/examples/atari/atari_c51.py b/examples/atari/atari_c51.py
@@ -9,8 +9,8 @@
 from atari_network import C51
 from atari_wrapper import make_atari_env
 
-from examples.common import logger_factory
 from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import C51Policy
 from tianshou.policy.base import BasePolicy
 from tianshou.trainer import OffpolicyTrainer
@@ -122,6 +122,7 @@ def test_c51(args: argparse.Namespace = get_args()) -> None:
     log_path = os.path.join(args.logdir, log_name)
 
     # logger
+    logger_factory = LoggerFactoryDefault()
     if args.logger == "wandb":
         logger_factory.logger_type = "wandb"
         logger_factory.wandb_project = args.wandb_project
@@ -182,8 +183,7 @@ def watch() -> None:
             print("Testing agent ...")
             test_collector.reset()
             result = test_collector.collect(n_episode=args.test_num, render=args.render)
-        rew = result.returns_stat.mean
-        print(f"Mean reward (over {result['n/ep']} episodes): {rew}")
+        result.pprint_asdict()
 
     if args.watch:
         watch()

diff --git a/examples/atari/atari_dqn.py b/examples/atari/atari_dqn.py
@@ -9,8 +9,8 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from examples.common import logger_factory
 from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import DQNPolicy
 from tianshou.policy.base import BasePolicy
 from tianshou.policy.modelbased.icm import ICMPolicy
@@ -104,7 +104,8 @@ def test_dqn(args: argparse.Namespace = get_args()) -> None:
     net = DQN(*args.state_shape, args.action_shape, args.device).to(args.device)
     optim = torch.optim.Adam(net.parameters(), lr=args.lr)
     # define policy
-    policy: DQNPolicy = DQNPolicy(
+    policy: DQNPolicy | ICMPolicy
+    policy = DQNPolicy(
         model=net,
         optim=optim,
         action_space=env.action_space,
@@ -157,6 +158,7 @@ def test_dqn(args: argparse.Namespace = get_args()) -> None:
     log_path = os.path.join(args.logdir, log_name)
 
     # logger
+    logger_factory = LoggerFactoryDefault()
     if args.logger == "wandb":
         logger_factory.logger_type = "wandb"
         logger_factory.wandb_project = args.wandb_project
@@ -223,8 +225,7 @@ def watch() -> None:
             print("Testing agent ...")
             test_collector.reset()
             result = test_collector.collect(n_episode=args.test_num, render=args.render)
-        rew = result.returns_stat.mean
-        print(f"Mean reward (over {result['n/ep']} episodes): {rew}")
+        result.pprint_asdict()
 
     if args.watch:
         watch()

diff --git a/examples/atari/atari_fqf.py b/examples/atari/atari_fqf.py
@@ -9,8 +9,8 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from examples.common import logger_factory
 from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import FQFPolicy
 from tianshou.policy.base import BasePolicy
 from tianshou.trainer import OffpolicyTrainer
@@ -135,6 +135,7 @@ def test_fqf(args: argparse.Namespace = get_args()) -> None:
     log_path = os.path.join(args.logdir, log_name)
 
     # logger
+    logger_factory = LoggerFactoryDefault()
     if args.logger == "wandb":
         logger_factory.logger_type = "wandb"
         logger_factory.wandb_project = args.wandb_project
@@ -195,8 +196,7 @@ def watch() -> None:
             print("Testing agent ...")
             test_collector.reset()
             result = test_collector.collect(n_episode=args.test_num, render=args.render)
-        rew = result.returns_stat.mean
-        print(f"Mean reward (over {result['n/ep']} episodes): {rew}")
+        result.pprint_asdict()
 
     if args.watch:
         watch()

diff --git a/examples/atari/atari_iqn.py b/examples/atari/atari_iqn.py
@@ -9,8 +9,8 @@
 from atari_network import DQN
 from atari_wrapper import make_atari_env
 
-from examples.common import logger_factory
 from tianshou.data import Collector, VectorReplayBuffer
+from tianshou.highlevel.logger import LoggerFactoryDefault
 from tianshou.policy import IQNPolicy
 from tianshou.policy.base import BasePolicy
 from tianshou.trainer import OffpolicyTrainer
@@ -132,6 +132,7 @@ def test_iqn(args: argparse.Namespace = get_args()) -> None:
     log_path = os.path.join(args.logdir, log_name)
 
     # logger
+    logger_factory = LoggerFactoryDefault()
     if args.logger == "wandb":
         logger_factory.logger_type = "wandb"
         logger_factory.wandb_project = args.wandb_project
@@ -192,8 +193,7 @@ def watch() -> None:
             print("Testing agent ...")
             test_collector.reset()
             result = test_collector.collect(n_episode=args.test_num, render=args.render)
-        rew = result.returns_stat.mean
-        print(f"Mean reward (over {result['n/ep']} episodes): {rew}")
+        result.pprint_asdict()
 
     if args.watch:
         watch()