thu-ml · Trinkle23897 · Mar 17, 2022 · Mar 5, 2022 · Mar 5, 2022 · Mar 5, 2022
diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md
@@ -7,6 +7,6 @@
 - [ ] I have searched through the [issue tracker](https://github.com/thu-ml/tianshou/issues) for duplicates
 - [ ] I have mentioned version numbers, operating system and environment, where applicable:
   ```python
-  import tianshou, torch, numpy, sys
-  print(tianshou.__version__, torch.__version__, numpy.__version__, sys.version, sys.platform)
+  import tianshou, gym, torch, numpy, sys
+  print(tianshou.__version__, gym.__version__, torch.__version__, numpy.__version__, sys.version, sys.platform)
   ```
diff --git a/Makefile b/Makefile
@@ -22,10 +22,8 @@ lint:
 	flake8 ${LINT_PATHS} --count --show-source --statistics
 
 format:
-	# sort imports
 	$(call check_install, isort)
 	isort ${LINT_PATHS}
-	# reformat using yapf
 	$(call check_install, yapf)
 	yapf -ir ${LINT_PATHS}
 
@@ -57,6 +55,6 @@ doc-clean:
 
 clean: doc-clean
 
-commit-checks: format lint mypy check-docstyle spelling
+commit-checks: lint check-codestyle mypy check-docstyle spelling
 
 .PHONY: clean spelling doc mypy lint format check-codestyle check-docstyle commit-checks
diff --git a/docs/api/tianshou.trainer.rst b/docs/api/tianshou.trainer.rst
@@ -1,7 +1,49 @@
 tianshou.trainer
 ================
 
-.. automodule:: tianshou.trainer
+
+On-policy
+---------
+
+.. autoclass:: tianshou.trainer.OnpolicyTrainer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autofunction:: tianshou.trainer.onpolicy_trainer
+
+.. autoclass:: tianshou.trainer.onpolicy_trainer_iter
+
+
+Off-policy
+----------
+
+.. autoclass:: tianshou.trainer.OffpolicyTrainer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autofunction:: tianshou.trainer.offpolicy_trainer
+
+.. autoclass:: tianshou.trainer.offpolicy_trainer_iter
+
+
+Offline
+-------
+
+.. autoclass:: tianshou.trainer.OfflineTrainer
    :members:
    :undoc-members:
    :show-inheritance:
+
+.. autofunction:: tianshou.trainer.offline_trainer
+
+.. autoclass:: tianshou.trainer.offline_trainer_iter
+
+
+utils
+-----
+
+.. autofunction:: tianshou.trainer.test_episode
+
+.. autofunction:: tianshou.trainer.gather_info
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
@@ -24,12 +24,15 @@ fqf
 iqn
 qrdqn
 rl
+offpolicy
+onpolicy
 quantile
 quantiles
 dqn
 param
 async
 subprocess
+deque
 nn
 equ
 cql

diff --git a/docs/tutorials/concepts.rst b/docs/tutorials/concepts.rst
@@ -380,6 +380,26 @@ Once you have a collector and a policy, you can start writing the training metho
 
 Tianshou has three types of trainer: :func:`~tianshou.trainer.onpolicy_trainer` for on-policy algorithms such as Policy Gradient, :func:`~tianshou.trainer.offpolicy_trainer` for off-policy algorithms such as DQN, and :func:`~tianshou.trainer.offline_trainer` for offline algorithms such as BCQ. Please check out :doc:`/api/tianshou.trainer` for the usage.
 
+We also provide the corresponding iterator-based trainer classes :class:`~tianshou.trainer.OnpolicyTrainer`, :class:`~tianshou.trainer.OffpolicyTrainer`, :class:`~tianshou.trainer.OfflineTrainer` to facilitate users writing more flexible training logic:
+::
+
+    trainer = OnpolicyTrainer(...)
+    for epoch, epoch_stat, info in trainer:
+        print(f"Epoch: {epoch}")
+        print(epoch_stat)
+        print(info)
+        do_something_with_policy()
+        query_something_about_policy()
+        make_a_plot_with(epoch_stat)
+        display(info)
+
+    # or even iterate on several trainers at the same time
+
+    trainer1 = OnpolicyTrainer(...)
+    trainer2 = OnpolicyTrainer(...)
+    for result1, result2, ... in zip(trainer1, trainer2, ...):
+        compare_results(result1, result2, ...)
+
 
 .. _pseudocode:
 

diff --git a/test/continuous/test_ppo.py b/test/continuous/test_ppo.py
@@ -11,7 +11,7 @@
 from tianshou.data import Collector, VectorReplayBuffer
 from tianshou.env import DummyVectorEnv
 from tianshou.policy import PPOPolicy
-from tianshou.trainer import onpolicy_trainer
+from tianshou.trainer import OnpolicyTrainer
 from tianshou.utils import TensorboardLogger
 from tianshou.utils.net.common import ActorCritic, Net
 from tianshou.utils.net.continuous import ActorProb, Critic
@@ -157,7 +157,7 @@ def save_checkpoint_fn(epoch, env_step, gradient_step):
             print("Fail to restore policy and optim.")
 
     # trainer
-    result = onpolicy_trainer(
+    trainer = OnpolicyTrainer(
         policy,
         train_collector,
         test_collector,
@@ -173,10 +173,16 @@ def save_checkpoint_fn(epoch, env_step, gradient_step):
         resume_from_log=args.resume,
         save_checkpoint_fn=save_checkpoint_fn
     )
-    assert stop_fn(result['best_reward'])
+
+    for epoch, epoch_stat, info in trainer:
+        print(f"Epoch: {epoch}")
+        print(epoch_stat)
+        print(info)
+
+    assert stop_fn(info["best_reward"])
 
     if __name__ == '__main__':
-        pprint.pprint(result)
+        pprint.pprint(info)
         # Let's watch its performance!
         env = gym.make(args.task)
         policy.eval()

diff --git a/test/continuous/test_sac_with_il.py b/test/continuous/test_sac_with_il.py
@@ -24,7 +24,7 @@ def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('--task', type=str, default='Pendulum-v0')
     parser.add_argument('--reward-threshold', type=float, default=None)
-    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--seed', type=int, default=1)
     parser.add_argument('--buffer-size', type=int, default=20000)
     parser.add_argument('--actor-lr', type=float, default=1e-3)
     parser.add_argument('--critic-lr', type=float, default=1e-3)

diff --git a/test/continuous/test_td3.py b/test/continuous/test_td3.py
@@ -11,7 +11,7 @@
 from tianshou.env import DummyVectorEnv
 from tianshou.exploration import GaussianNoise
 from tianshou.policy import TD3Policy
-from tianshou.trainer import offpolicy_trainer
+from tianshou.trainer import OffpolicyTrainer
 from tianshou.utils import TensorboardLogger
 from tianshou.utils.net.common import Net
 from tianshou.utils.net.continuous import Actor, Critic
@@ -135,8 +135,8 @@ def save_fn(policy):
     def stop_fn(mean_rewards):
         return mean_rewards >= args.reward_threshold
 
-    # trainer
-    result = offpolicy_trainer(
+    # Iterator trainer
+    trainer = OffpolicyTrainer(
         policy,
         train_collector,
         test_collector,
@@ -148,12 +148,17 @@ def stop_fn(mean_rewards):
         update_per_step=args.update_per_step,
         stop_fn=stop_fn,
         save_fn=save_fn,
-        logger=logger
+        logger=logger,
     )
-    assert stop_fn(result['best_reward'])
+    for epoch, epoch_stat, info in trainer:
+        print(f"Epoch: {epoch}")
+        print(epoch_stat)
+        print(info)
 
-    if __name__ == '__main__':
-        pprint.pprint(result)
+    assert stop_fn(info["best_reward"])
+
+    if __name__ == "__main__":
+        pprint.pprint(info)
         # Let's watch its performance!
         env = gym.make(args.task)
         policy.eval()

diff --git a/test/offline/test_cql.py b/test/offline/test_cql.py
@@ -12,7 +12,7 @@
 from tianshou.data import Collector, VectorReplayBuffer
 from tianshou.env import DummyVectorEnv
 from tianshou.policy import CQLPolicy
-from tianshou.trainer import offline_trainer
+from tianshou.trainer import OfflineTrainer
 from tianshou.utils import TensorboardLogger
 from tianshou.utils.net.common import Net
 from tianshou.utils.net.continuous import ActorProb, Critic
@@ -195,7 +195,7 @@ def watch():
         collector.collect(n_episode=1, render=1 / 35)
 
     # trainer
-    result = offline_trainer(
+    trainer = OfflineTrainer(
         policy,
         buffer,
         test_collector,
@@ -207,11 +207,17 @@ def watch():
         stop_fn=stop_fn,
         logger=logger,
     )
-    assert stop_fn(result['best_reward'])
+
+    for epoch, epoch_stat, info in trainer:
+        print(f"Epoch: {epoch}")
+        print(epoch_stat)
+        print(info)
+
+    assert stop_fn(info["best_reward"])
 
     # Let's watch its performance!
-    if __name__ == '__main__':
-        pprint.pprint(result)
+    if __name__ == "__main__":
+        pprint.pprint(info)
         env = gym.make(args.task)
         policy.eval()
         collector = Collector(policy, env)

diff --git a/tianshou/trainer/__init__.py b/tianshou/trainer/__init__.py
@@ -1,16 +1,34 @@
 """Trainer package."""
 
-# isort:skip_file
-
-from tianshou.trainer.utils import test_episode, gather_info
-from tianshou.trainer.onpolicy import onpolicy_trainer
-from tianshou.trainer.offpolicy import offpolicy_trainer
-from tianshou.trainer.offline import offline_trainer
+from tianshou.trainer.base import BaseTrainer
+from tianshou.trainer.offline import (
+    OfflineTrainer,
+    offline_trainer,
+    offline_trainer_iter,
+)
+from tianshou.trainer.offpolicy import (
+    OffpolicyTrainer,
+    offpolicy_trainer,
+    offpolicy_trainer_iter,
+)
+from tianshou.trainer.onpolicy import (
+    OnpolicyTrainer,
+    onpolicy_trainer,
+    onpolicy_trainer_iter,
+)
+from tianshou.trainer.utils import gather_info, test_episode
 
 __all__ = [
+    "BaseTrainer",
     "offpolicy_trainer",
+    "offpolicy_trainer_iter",
+    "OffpolicyTrainer",
     "onpolicy_trainer",
+    "onpolicy_trainer_iter",
+    "OnpolicyTrainer",
     "offline_trainer",
+    "offline_trainer_iter",
+    "OfflineTrainer",
     "test_episode",
     "gather_info",
 ]