Merge pull request #155 from muupan/fix-episodic-buffer-len

toslunar · web-flow · commit 51a2762246db · 2017-10-20T16:24:53.000+09:00
Fix episodic buffer __len__
diff --git a/chainerrl/agents/pcl.py b/chainerrl/agents/pcl.py
@@ -273,6 +273,9 @@ def update_from_replay(self):
         if len(self.replay_buffer) < self.replay_start_size:
             return
 
+        if self.replay_buffer.n_episodes < self.batchsize:
+            return
+
         if self.process_idx == 0:
             self.logger.debug('update_from_replay')
 
diff --git a/chainerrl/replay_buffer.py b/chainerrl/replay_buffer.py
@@ -4,8 +4,13 @@
 from __future__ import absolute_import
 from builtins import *  # NOQA
 from future import standard_library
+from future.utils import with_metaclass
 standard_library.install_aliases()
 
+from abc import ABCMeta
+from abc import abstractmethod
+from abc import abstractproperty
+
 import numpy as np
 import six.moves.cPickle as pickle
 
@@ -14,14 +19,17 @@
 from chainerrl.misc.prioritized import PrioritizedBuffer
 
 
-class ReplayBuffer(object):
+class AbstractReplayBuffer(with_metaclass(ABCMeta, object)):
+    """Defines a common interface of replay buffer.
 
-    def __init__(self, capacity=None):
-        self.memory = RandomAccessQueue(maxlen=capacity)
+    You can append transitions to the replay buffer and later sample from it.
+    Replay buffers are typically used in experience replay.
+    """
 
+    @abstractmethod
     def append(self, state, action, reward, next_state=None, next_action=None,
                is_state_terminal=False):
-        """Append a transition to this replay buffer
+        """Append a transition to this replay buffer.
 
         Args:
             state: s_t
@@ -31,13 +39,107 @@ def append(self, state, action, reward, next_state=None, next_action=None,
             next_action: a_{t+1} (can be None for off-policy algorithms)
             is_state_terminal (bool)
         """
+        raise NotImplementedError
+
+    @abstractmethod
+    def sample(self, n):
+        """Sample n unique transitions from this replay buffer.
+
+        Args:
+            n (int): Number of transitions to sample.
+        Returns:
+            Sequence of n sampled transitions.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def __len__(self):
+        """Return the number of transitions in the buffer.
+
+        Returns:
+            Number of transitions in the buffer.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def save(self, filename):
+        """Save the content of the buffer to a file.
+
+        Args:
+            filename (str): Path to a file.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def load(self, filename):
+        """Load the content of the buffer from a file.
+
+        Args:
+            filename (str): Path to a file.
+        """
+        raise NotImplementedError
+
+
+class AbstractEpisodicReplayBuffer(AbstractReplayBuffer):
+    """Defines a common interface of episodic replay buffer.
+
+    Episodic replay buffers allows you to append and sample episodes.
+    """
+
+    @abstractmethod
+    def sample_episodes(self, n_episodes, max_len=None):
+        """Sample n unique (sub)episodes from this replay buffer.
+
+        Args:
+            n (int): Number of episodes to sample.
+            max_len (int or None): Maximum length of sampled episodes. If it is
+                smaller than the length of some episode, the subsequence of the
+                episode is sampled instead. If None, full episodes are always
+                returned.
+        Returns:
+            Sequence of n sampled epiosodes, each of which is a sequence of
+            transitions.
+        """
+        raise NotImplementedError
+
+    @abstractproperty
+    def n_episodes(self):
+        """Returns the number of episodes in the buffer.
+
+        Returns:
+            Number of episodes in the buffer.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def stop_current_episode(self):
+        """Notify the buffer that the current episode is interrupted.
+
+        You may want to interrupt the current episode and start a new one
+        before observing a terminal state. This is typical in continuing envs.
+        In such cases, you need to call this method before appending a new
+        transition so that the buffer will treat it as an initial transition of
+        a new episode.
+
+        This method should not be called after an episode whose termination is
+        already notified by appending a transition with is_state_terminal=True.
+        """
+        raise NotImplementedError
+
+
+class ReplayBuffer(AbstractReplayBuffer):
+
+    def __init__(self, capacity=None):
+        self.memory = RandomAccessQueue(maxlen=capacity)
+
+    def append(self, state, action, reward, next_state=None, next_action=None,
+               is_state_terminal=False):
         experience = dict(state=state, action=action, reward=reward,
                           next_state=next_state, next_action=next_action,
                           is_state_terminal=is_state_terminal)
         self.memory.append(experience)
 
     def sample(self, n):
-        """Sample n unique samples from this replay buffer"""
         assert len(self.memory) >= n
         return self.memory.sample(n)
 
@@ -117,7 +219,6 @@ def __init__(self, capacity=None,
             self, alpha, beta0, betasteps, eps, normalize_by_max)
 
     def sample(self, n):
-        """Sample n unique samples from this replay buffer"""
         assert len(self.memory) >= n
         sampled, probabilities = self.memory.sample(n)
         weights = self.weights_from_probabilities(probabilities)
@@ -137,7 +238,7 @@ def random_subseq(seq, subseq_len):
         return seq[i:i + subseq_len]
 
 
-class EpisodicReplayBuffer(object):
+class EpisodicReplayBuffer(AbstractEpisodicReplayBuffer):
 
     def __init__(self, capacity=None):
         self.current_episode = []
@@ -147,16 +248,6 @@ def __init__(self, capacity=None):
 
     def append(self, state, action, reward, next_state=None, next_action=None,
                is_state_terminal=False, **kwargs):
-        """Append a transition to this replay buffer
-
-        Args:
-            state: s_t
-            action: a_t
-            reward: r_t
-            next_state: s_{t+1} (can be None if terminal)
-            next_action: a_{t+1} (can be None for off-policy algorithms)
-            is_state_terminal (bool)
-        """
         experience = dict(state=state, action=action, reward=reward,
                           next_state=next_state, next_action=next_action,
                           is_state_terminal=is_state_terminal,
@@ -166,12 +257,10 @@ def append(self, state, action, reward, next_state=None, next_action=None,
             self.stop_current_episode()
 
     def sample(self, n):
-        """Sample n unique samples from this replay buffer"""
         assert len(self.memory) >= n
         return self.memory.sample(n)
 
     def sample_episodes(self, n_episodes, max_len=None):
-        """Sample n unique samples from this replay buffer"""
         assert len(self.episodic_memory) >= n_episodes
         episodes = self.episodic_memory.sample(n_episodes)
         if max_len is not None:
@@ -180,6 +269,10 @@ def sample_episodes(self, n_episodes, max_len=None):
             return episodes
 
     def __len__(self):
+        return len(self.memory)
+
+    @property
+    def n_episodes(self):
         return len(self.episodic_memory)
 
     def save(self, filename):
@@ -313,6 +406,11 @@ def __init__(self, replay_buffer, update_func, batchsize, episodic_update,
     def update_if_necessary(self, iteration):
         if len(self.replay_buffer) < self.replay_start_size:
             return
+
+        if (self.episodic_update
+                and self.replay_buffer.n_episodes < self.batchsize):
+            return
+
         if iteration % self.update_interval != 0:
             return
 
diff --git a/examples/gym/train_dqn_gym.py b/examples/gym/train_dqn_gym.py
@@ -54,7 +54,7 @@ def main():
     parser.add_argument('--steps', type=int, default=10 ** 5)
     parser.add_argument('--prioritized-replay', action='store_true')
     parser.add_argument('--episodic-replay', action='store_true')
-    parser.add_argument('--replay-start-size', type=int, default=None)
+    parser.add_argument('--replay-start-size', type=int, default=1000)
     parser.add_argument('--target-update-interval', type=int, default=10 ** 2)
     parser.add_argument('--target-update-method', type=str, default='hard')
     parser.add_argument('--soft-update-tau', type=float, default=1e-2)
@@ -130,11 +130,8 @@ def make_env(for_eval):
     if args.episodic_replay:
         if args.minibatch_size is None:
             args.minibatch_size = 4
-        if args.replay_start_size is None:
-            args.replay_start_size = 10
         if args.prioritized_replay:
-            betasteps = \
-                (args.steps - timestep_limit * args.replay_start_size) \
+            betasteps = (args.steps - args.replay_start_size) \
                 // args.update_interval
             rbuf = replay_buffer.PrioritizedEpisodicReplayBuffer(
                 rbuf_capacity, betasteps=betasteps)
@@ -143,8 +140,6 @@ def make_env(for_eval):
     else:
         if args.minibatch_size is None:
             args.minibatch_size = 32
-        if args.replay_start_size is None:
-            args.replay_start_size = 1000
         if args.prioritized_replay:
             betasteps = (args.steps - args.replay_start_size) \
                 // args.update_interval
diff --git a/tests/test_replay_buffer.py b/tests/test_replay_buffer.py
@@ -113,6 +113,9 @@ def subtest_append_and_sample(self, capacity):
             for trans in transs:
                 rbuf.append(**trans)
 
+        self.assertEqual(len(rbuf), 90)
+        self.assertEqual(rbuf.n_episodes, 9)
+
         for k in [10, 30, 90]:
             s = rbuf.sample(k)
             self.assertEqual(len(s), k)
@@ -130,13 +133,13 @@ def subtest_append_and_sample(self, capacity):
 
     def test_save_and_load(self):
         for capacity in [100, None]:
-            self.subtest_append_and_sample(capacity)
+            self.subtest_save_and_load(capacity)
 
     def subtest_save_and_load(self, capacity):
 
         tempdir = tempfile.mkdtemp()
 
-        rbuf = replay_buffer.ReplayBuffer(capacity)
+        rbuf = replay_buffer.EpisodicReplayBuffer(capacity)
 
         transs = [dict(state=n, action=n+10, reward=n+20,
                        next_state=n+1, next_action=n+11,
@@ -153,12 +156,15 @@ def subtest_save_and_load(self, capacity):
         rbuf.append(**transs[4])
         rbuf.stop_current_episode()
 
+        self.assertEqual(len(rbuf), 5)
+        self.assertEqual(rbuf.n_episodes, 2)
+
         # Save
         filename = os.path.join(tempdir, 'rbuf.pkl')
         rbuf.save(filename)
 
         # Initialize rbuf
-        rbuf = replay_buffer.ReplayBuffer(capacity)
+        rbuf = replay_buffer.EpisodicReplayBuffer(capacity)
 
         # Of course it has no transition yet
         self.assertEqual(len(rbuf), 0)
@@ -168,22 +174,26 @@ def subtest_save_and_load(self, capacity):
 
         # Sampled transitions are exactly what I added!
         s5 = rbuf.sample(5)
-        self.assertEqual(len(s5) == 5)
+        self.assertEqual(len(s5), 5)
         for t in s5:
             n = t['state']
             self.assertIn(n, range(5))
             self.assertEqual(t, transs[n])
 
         # And sampled episodes are exactly what I added!
         s2e = rbuf.sample_episodes(2)
-        self.assertEqual(len(s2e) == 2)
+        self.assertEqual(len(s2e), 2)
         if s2e[0][0]['state'] == 0:
             self.assertEqual(s2e[0], [transs[0], transs[1]])
             self.assertEqual(s2e[1], [transs[2], transs[3], transs[4]])
         else:
             self.assertEqual(s2e[0], [transs[2], transs[3], transs[4]])
             self.assertEqual(s2e[1], [transs[0], transs[1]])
 
+        # Sizes are correct!
+        self.assertEqual(len(rbuf), 5)
+        self.assertEqual(rbuf.n_episodes, 2)
+
 
 class TestPrioritizedReplayBuffer(unittest.TestCase):
 
@@ -354,7 +364,9 @@ def test_append_and_sample(self):
                       for i in range(n)]
             for trans in transs:
                 rbuf.append(**trans)
-        self.assertEqual(len(rbuf), 9)
+
+        self.assertEqual(len(rbuf), 90)
+        self.assertEqual(rbuf.n_episodes, 9)
 
         for k in [10, 30, 90]:
             s = rbuf.sample(k)
