thu-ml · Trinkle23897 · Aug 27, 2020 · Aug 20, 2020 · Aug 20, 2020 · Aug 20, 2020
diff --git a/README.md b/README.md
@@ -247,6 +247,8 @@ policy.load_state_dict(torch.load('dqn.pth'))
 Watch the performance with 35 FPS:
 
 ```python
+policy.eval()
+policy.set_eps(eps_test)
 collector = ts.data.Collector(policy, env)
 collector.collect(n_episode=1, render=1 / 35)
 ```

diff --git a/docs/tutorials/cheatsheet.rst b/docs/tutorials/cheatsheet.rst
@@ -96,7 +96,7 @@ This is related to `Issue 42 <https://github.com/thu-ml/tianshou/issues/42>`_.
 
 If you want to get log stat from data stream / pre-process batch-image / modify the reward with given env info, use ``preproces_fn`` in :class:`~tianshou.data.Collector`. This is a hook which will be called before the data adding into the buffer.
 
-This function receives typically 7 keys, as listed in :class:`~tianshou.data.Batch`, and returns the modified part within a dict or a Batch. For example, you can write your hook as:
+This function receives up to 7 keys ``obs``, ``act``, ``rew``, ``done``, ``obs_next``, ``info``, and ``policy``, as listed in :class:`~tianshou.data.Batch`, and returns the modified part within a :class:`~tianshou.data.Batch`. Only ``obs`` is defined at env reset, while every key is specified for normal steps. For example, you can write your hook as:
 ::
 
     import numpy as np
@@ -109,9 +109,11 @@ This function receives typically 7 keys, as listed in :class:`~tianshou.data.Bat
             self.baseline = 0
         def preprocess_fn(**kwargs):
             """change reward to zero mean"""
+            # if only obs exist -> reset
+            # if obs/act/rew/done/... exist -> normal step
             if 'rew' not in kwargs:
                 # means that it is called after env.reset(), it can only process the obs
-                return {}  # none of the variables are needed to be updated
+                return Batch()  # none of the variables are needed to be updated
             else:
                 n = len(kwargs['rew'])  # the number of envs in collector
                 if self.episode_log is None:
@@ -125,7 +127,6 @@ This function receives typically 7 keys, as listed in :class:`~tianshou.data.Bat
                         self.episode_log[i] = []
                         self.baseline = np.mean(self.main_log)
                 return Batch(rew=kwargs['rew'])
-                # you can also return with {'rew': kwargs['rew']}
 
 And finally,
 ::

diff --git a/docs/tutorials/dqn.rst b/docs/tutorials/dqn.rst
@@ -176,6 +176,8 @@ Watch the Agent's Performance
 :class:`~tianshou.data.Collector` supports rendering. Here is the example of watching the agent's performance in 35 FPS:
 ::
 
+    policy.eval()
+    policy.set_eps(0.05)
     collector = ts.data.Collector(policy, env)
     collector.collect(n_episode=1, render=1 / 35)
 

diff --git a/docs/tutorials/tictactoe.rst b/docs/tutorials/tictactoe.rst
@@ -285,6 +285,8 @@ With the above preparation, we are close to the first learned agent. The followi
         env = TicTacToeEnv(args.board_size, args.win_size)
         policy, optim = get_agents(
             args, agent_learn=agent_learn, agent_opponent=agent_opponent)
+        policy.eval()
+        policy.set_eps(args.eps_test)
         collector = Collector(policy, env)
         result = collector.collect(n_episode=1, render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')

diff --git a/examples/atari/pong_dqn.py b/examples/atari/pong_dqn.py
diff --git a/examples/atari/atari.py → examples/atari/runnable/atari.py b/examples/atari/atari.py → examples/atari/runnable/atari.py
diff --git a/examples/atari/pong_a2c.py → examples/atari/runnable/pong_a2c.py b/examples/atari/pong_a2c.py → examples/atari/runnable/pong_a2c.py
diff --git a/examples/atari/pong_ppo.py → examples/atari/runnable/pong_ppo.py b/examples/atari/pong_ppo.py → examples/atari/runnable/pong_ppo.py
diff --git a/examples/box2d/acrobot_dualdqn.py b/examples/box2d/acrobot_dualdqn.py
@@ -102,9 +102,12 @@ def test_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        policy.set_eps(args.eps_test)
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/box2d/bipedal_hardcore_sac.py b/examples/box2d/bipedal_hardcore_sac.py
@@ -44,6 +44,7 @@ def get_args():
 
 class EnvWrapper(object):
     """Env wrapper for reward scale, action repeat and action noise"""
+
     def __init__(self, task, action_repeat=3,
                  reward_scale=5, act_noise=0.3):
         self._env = gym.make(task)
@@ -71,19 +72,20 @@ def step(self, action):
 def test_sac_bipedal(args=get_args()):
     torch.set_num_threads(1)  # we just need only one thread for NN
 
+    env = EnvWrapper(args.task)
+
     def IsStop(reward):
-        return reward >= 300 * 5
+        return reward >= env.spec.reward_threshold
 
-    env = EnvWrapper(args.task)
     args.state_shape = env.observation_space.shape or env.observation_space.n
     args.action_shape = env.action_space.shape or env.action_space.n
     args.max_action = env.action_space.high[0]
 
     train_envs = SubprocVectorEnv(
         [lambda: EnvWrapper(args.task) for _ in range(args.training_num)])
     # test_envs = gym.make(args.task)
-    test_envs = SubprocVectorEnv(
-        [lambda: EnvWrapper(args.task) for _ in range(args.test_num)])
+    test_envs = SubprocVectorEnv([lambda: EnvWrapper(args.task, reward_scale=1)
+                                  for _ in range(args.test_num)])
 
     # seed
     np.random.seed(args.seed)
@@ -138,9 +140,11 @@ def save_fn(policy):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = EnvWrapper(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=16, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/box2d/lunarlander_dqn.py b/examples/box2d/lunarlander_dqn.py
@@ -99,9 +99,12 @@ def test_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        policy.set_eps(args.eps_test)
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/box2d/sac_mcc.py → examples/box2d/mcc_sac.py b/examples/box2d/sac_mcc.py → examples/box2d/mcc_sac.py
@@ -112,9 +112,11 @@ def stop_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/mujoco/ant_v2_ddpg.py b/examples/mujoco/ant_v2_ddpg.py
@@ -88,9 +88,11 @@ def stop_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/mujoco/ant_v2_sac.py b/examples/mujoco/ant_v2_sac.py
@@ -98,9 +98,11 @@ def stop_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/mujoco/ant_v2_td3.py b/examples/mujoco/ant_v2_td3.py
@@ -98,9 +98,11 @@ def stop_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/mujoco/halfcheetahBullet_v0_sac.py b/examples/mujoco/halfcheetahBullet_v0_sac.py
@@ -104,9 +104,11 @@ def stop_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_episode=1, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/examples/mujoco/point_maze_td3.py b/examples/mujoco/point_maze_td3.py
@@ -104,9 +104,11 @@ def stop_fn(x):
     if __name__ == '__main__':
         pprint.pprint(result)
         # Let's watch its performance!
-        env = gym.make(args.task)
-        collector = Collector(policy, env)
-        result = collector.collect(n_step=1000, render=args.render)
+        policy.eval()
+        test_envs.seed(args.seed)
+        test_collector.reset()
+        result = test_collector.collect(n_episode=[1] * args.test_num,
+                                        render=args.render)
         print(f'Final reward: {result["rew"]}, length: {result["len"]}')
 
 

diff --git a/setup.py b/setup.py
@@ -3,18 +3,10 @@
 
 from setuptools import setup, find_packages
 
-import re
-from os import path
-
-here = path.abspath(path.dirname(__file__))
-
-# Get the version string
-with open(path.join(here, 'tianshou', '__init__.py')) as f:
-    version = re.search(r'__version__ = \'(.*?)\'', f.read()).group(1)
 
 setup(
     name='tianshou',
-    version=version,
+    version='0.2.6',
     description='A Library for Deep Reinforcement Learning',
     long_description=open('README.md', encoding='utf8').read(),
     long_description_content_type='text/markdown',

diff --git a/test/base/env.py b/test/base/env.py
@@ -21,6 +21,8 @@ def __init__(self, size, sleep=0, dict_state=False, recurse_state=False,
         self.recurse_state = recurse_state
         self.ma_rew = ma_rew
         self._md_action = multidiscrete_action
+        # how many steps this env has stepped
+        self.steps = 0
         if dict_state:
             self.observation_space = Dict(
                 {"index": Box(shape=(1, ), low=0, high=size - 1),
@@ -74,6 +76,7 @@ def _get_state(self):
             return np.array([self.index], dtype=np.float32)
 
     def step(self, action):
+        self.steps += 1
         if self._md_action:
             action = action[0]
         if self.done: