thu-ml · Trinkle23897 · Mar 2, 2021 · Mar 2, 2021 · Mar 2, 2021 · Mar 2, 2021
diff --git a/docs/api/tianshou.data.rst b/docs/api/tianshou.data.rst
@@ -5,7 +5,7 @@ tianshou.data
 Batch
 -----
 
-.. automodule:: tianshou.data.batch
+.. autoclass:: tianshou.data.Batch
    :members:
    :undoc-members:
    :show-inheritance:
@@ -14,16 +14,77 @@ Batch
 Buffer
 ------
 
-.. automodule:: tianshou.data.buffer
+ReplayBuffer
+~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.ReplayBuffer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+PrioritizedReplayBuffer
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.PrioritizedReplayBuffer
    :members:
    :undoc-members:
    :show-inheritance:
 
+ReplayBufferManager
+~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.ReplayBufferManager
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+PrioritizedReplayBufferManager
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.PrioritizedReplayBufferManager
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+VectorReplayBuffer
+~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.VectorReplayBuffer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+PrioritizedVectorReplayBuffer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.PrioritizedVectorReplayBuffer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+CachedReplayBuffer
+~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.CachedReplayBuffer
+   :members:
+   :undoc-members:
+   :show-inheritance:
 
 Collector
 ---------
 
-.. automodule:: tianshou.data.collector
+Collector
+~~~~~~~~~
+
+.. autoclass:: tianshou.data.Collector
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+AsyncCollector
+~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.data.AsyncCollector
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/docs/api/tianshou.env.rst b/docs/api/tianshou.env.rst
@@ -5,7 +5,42 @@ tianshou.env
 VectorEnv
 ---------
 
-.. automodule:: tianshou.env
+BaseVectorEnv
+~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.BaseVectorEnv
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+DummyVectorEnv
+~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.DummyVectorEnv
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+SubprocVectorEnv
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.SubprocVectorEnv
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+ShmemVectorEnv
+~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.ShmemVectorEnv
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+RayVectorEnv
+~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.RayVectorEnv
    :members:
    :undoc-members:
    :show-inheritance:
@@ -14,7 +49,34 @@ VectorEnv
 Worker
 ------
 
-.. automodule:: tianshou.env.worker
+EnvWorker
+~~~~~~~~~
+
+.. autoclass:: tianshou.env.worker.EnvWorker
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+DummyEnvWorker
+~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.worker.DummyEnvWorker
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+SubprocEnvWorker
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.worker.SubprocEnvWorker
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+RayEnvWorker
+~~~~~~~~~~~~
+
+.. autoclass:: tianshou.env.worker.RayEnvWorker
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/docs/api/tianshou.policy.rst b/docs/api/tianshou.policy.rst
@@ -1,7 +1,106 @@
 tianshou.policy
 ===============
 
-.. automodule:: tianshou.policy
+Base
+----
+
+.. autoclass:: tianshou.policy.BasePolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.RandomPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Model-free
+----------
+
+DQN Family
+~~~~~~~~~~
+
+.. autoclass:: tianshou.policy.DQNPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.C51Policy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.QRDQNPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+On-policy
+~~~~~~~~~
+
+.. autoclass:: tianshou.policy.PGPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.A2CPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.PPOPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Off-policy
+~~~~~~~~~~
+
+.. autoclass:: tianshou.policy.DDPGPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.TD3Policy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.SACPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.DiscreteSACPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Imitation
+---------
+
+.. autoclass:: tianshou.policy.ImitationPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: tianshou.policy.DiscreteBCQPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Model-based
+-----------
+
+.. autoclass:: tianshou.policy.PSRLPolicy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Multi-agent
+-----------
+
+.. autoclass:: tianshou.policy.MultiAgentPolicyManager
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/docs/api/tianshou.utils.rst b/docs/api/tianshou.utils.rst
@@ -10,16 +10,25 @@ tianshou.utils
 Pre-defined Networks
 --------------------
 
+Common
+~~~~~~
+
 .. automodule:: tianshou.utils.net.common
    :members:
    :undoc-members:
    :show-inheritance:
 
+Discrete
+~~~~~~~~
+
 .. automodule:: tianshou.utils.net.discrete
    :members:
    :undoc-members:
    :show-inheritance:
 
+Continuous
+~~~~~~~~~~
+
 .. automodule:: tianshou.utils.net.continuous
    :members:
    :undoc-members:

diff --git a/docs/tutorials/concepts.rst b/docs/tutorials/concepts.rst
@@ -341,7 +341,30 @@ The :class:`~tianshou.data.Collector` enables the policy to interact with differ
 
 :meth:`~tianshou.data.Collector.collect` is the main method of Collector: it let the policy perform a specified number of step ``n_step`` or episode ``n_episode`` and store the data in the replay buffer, then return the statistics of the collected data such as episode's total reward.
 
-The general explanation is listed in :ref:`pseudocode`. Other usages of collector are listed in :class:`~tianshou.data.Collector` documentation.
+The general explanation is listed in :ref:`pseudocode`. Other usages of collector are listed in :class:`~tianshou.data.Collector` documentation. Here are some example usages:
+::
+
+    policy = PGPolicy(...)  # or other policies if you wish
+    env = gym.make("CartPole-v0")
+
+    replay_buffer = ReplayBuffer(size=10000)
+
+    # here we set up a collector with a single environment
+    collector = Collector(policy, env, buffer=replay_buffer)
+
+    # the collector supports vectorized environments as well
+    vec_buffer = VectorReplayBuffer(total_size=10000, buffer_num=3)
+    # buffer_num should be equal to (suggested) or larger than #envs
+    envs = DummyVectorEnv([lambda: gym.make("CartPole-v0") for _ in range(3)])
+    collector = Collector(policy, envs, buffer=vec_buffer)
+
+    # collect 3 episodes
+    collector.collect(n_episode=3)
+    # collect at least 2 steps
+    collector.collect(n_step=2)
+    # collect episodes with visual rendering ("render" is the sleep time between
+    # rendering consecutive frames)
+    collector.collect(n_episode=1, render=0.03)
 
 There is also another type of collector :class:`~tianshou.data.AsyncCollector` which supports asynchronous environment setting (for those taking a long time to step). However, AsyncCollector only supports **at least** ``n_step`` or ``n_episode`` collection due to the property of asynchronous environments.
 

diff --git a/tianshou/__init__.py b/tianshou/__init__.py
@@ -1,7 +1,7 @@
 from tianshou import data, env, utils, policy, trainer, exploration
 
 
-__version__ = "0.3.2"
+__version__ = "0.4.0"
 
 __all__ = [
     "env",

diff --git a/tianshou/data/__init__.py b/tianshou/data/__init__.py
@@ -1,15 +1,13 @@
 from tianshou.data.batch import Batch
 from tianshou.data.utils.converter import to_numpy, to_torch, to_torch_as
 from tianshou.data.utils.segtree import SegmentTree
-from tianshou.data.buffer import (
-    ReplayBuffer,
-    PrioritizedReplayBuffer,
-    ReplayBufferManager,
-    PrioritizedReplayBufferManager,
-    VectorReplayBuffer,
-    PrioritizedVectorReplayBuffer,
-    CachedReplayBuffer,
-)
+from tianshou.data.buffer.base import ReplayBuffer
+from tianshou.data.buffer.prio import PrioritizedReplayBuffer
+from tianshou.data.buffer.manager import ReplayBufferManager
+from tianshou.data.buffer.manager import PrioritizedReplayBufferManager
+from tianshou.data.buffer.vecbuf import VectorReplayBuffer
+from tianshou.data.buffer.vecbuf import PrioritizedVectorReplayBuffer
+from tianshou.data.buffer.cached import CachedReplayBuffer
 from tianshou.data.collector import Collector, AsyncCollector
 
 __all__ = [

diff --git a/tianshou/data/batch.py b/tianshou/data/batch.py
@@ -146,6 +146,19 @@ def _parse_value(v: Any) -> Optional[Union["Batch", np.ndarray, torch.Tensor]]:
         return v
 
 
+def _alloc_by_keys_diff(
+    meta: "Batch", batch: "Batch", size: int, stack: bool = True
+) -> None:
+    for key in batch.keys():
+        if key in meta.keys():
+            if isinstance(meta[key], Batch) and isinstance(batch[key], Batch):
+                _alloc_by_keys_diff(meta[key], batch[key], size, stack)
+            elif isinstance(meta[key], Batch) and meta[key].is_empty():
+                meta[key] = _create_value(batch[key], size, stack)
+        else:
+            meta[key] = _create_value(batch[key], size, stack)
+
+
 class Batch:
     """The internal data structure in Tianshou.