update estimators, samplers, and storages

pyrobolearn/estimators/estimator.py

@@ -6,9 +6,12 @@ Dependencies:
 """
 
 from abc import ABCMeta
+import collections
+
 import torch
 
-from pyrobolearn.storages import RolloutStorage
+from pyrobolearn.storages import RolloutStorage, Batch
+from pyrobolearn.values import Value, QValue
 
 
 __author__ = "Brian Delhaisse"
@@ -472,6 +475,121 @@ class TDReturn(Return):
     pass
 
 
+class Target(BaseReturn):
+
+    def compute(self, batch):
+        pass
+
+    def __call__(self, batch):
+        return self.compute(batch)
+
+
+class ValueTarget(Target):
+    r"""Value target.
+
+    Compute the value target given by :math:`(r + \gamma (1-d) \min_i V_{\phi_i}(s'))`, where the index `i` is in
+    the case there are multiple value function approximators given to this class.
+    """
+
+    def __init__(self, values, gamma=1.):
+        """
+        Initialize the state value target.
+
+        Args:
+            values (Value, list of Value): state value function(s).
+            gamma (float): discount factor
+        """
+        super(ValueTarget, self).__init__(gamma)
+        if not isinstance(values, collections.Iterable):
+            values = [values]
+        for i, value in enumerate(values):
+            if not isinstance(value, Value):
+                raise TypeError('The {}th value is not an instance of `Value`, instead got: {}'.format(i, type(value)))
+        self._values = values
+
+    def compute(self, batch):
+        r"""
+        Compute the value target :math:`(r + \gamma (1-d) \min_i V_{\phi_i}(s'))`
+
+        Args:
+            batch (Batch): batch containing the transitions.
+        """
+        value = torch.min(torch.cat([value(batch['states']) for value in self._values], dim=1), dim=1)[0]
+        batch[self] = batch['rewards'] + self.gamma * (1 - batch['masks']) * value
+        return batch
+
+
+class QValueTarget(Target):
+    r"""Q-Value target.
+
+    Compute the Q-value target given by :math:`(r + \gamma (1-d) \min_i Q_{\phi_i}(s',a'))`, where the index `i` is in
+    the case there are multiple Q-value function approximators given to this class.
+    """
+
+    def __init__(self, q_values, gamma=1.):
+        """
+        Initialize the state value target.
+
+        Args:
+            q_values (QValue, list of QValue): state-action value function(s).
+            gamma (float): discount factor
+        """
+        super(QValueTarget, self).__init__(gamma)
+        if not isinstance(q_values, collections.Iterable):
+            q_values = [q_values]
+        for i, value in enumerate(q_values):
+            if not isinstance(value, QValue):
+                raise TypeError('The {}th value is not an instance of `QValue`, instead got: {}'.format(i, type(value)))
+        self._q_values = q_values
+
+    def compute(self, batch):
+        r"""
+        Compute the value target :math:`(r + \gamma (1-d) \min_i Q_{\phi_i}(s',a'))`
+
+        Args:
+            batch (Batch): batch containing the transitions.
+        """
+        value = torch.min(torch.cat([value(batch['states']) for value in self._q_values], dim=1), dim=1)[0]
+        batch[self] = batch['rewards'] + self.gamma * (1 - batch['masks']) * value
+        return batch
+
+
+class QLearningTarget(Target):
+    r"""Q-Learning target.
+
+    Compute the Q-value target given by :math:`(r + \gamma (1-d) \min_i \max_{a'} Q_{\phi_i}(s',a'))`, where the
+    index `i` is in the case there are multiple Q-value function approximators given to this class.
+    """
+
+    def __init__(self, q_values, gamma=1.):
+        """
+        Initialize the state value target.
+
+        Args:
+            q_values (QValue, list of QValue): state-action value function(s).
+            gamma (float): discount factor
+        """
+        super(QLearningTarget, self).__init__(gamma)
+        if not isinstance(q_values, collections.Iterable):
+            q_values = [q_values]
+        for i, value in enumerate(q_values):
+            if not isinstance(value, QValue):
+                raise TypeError('The {}th value is not an instance of `QValue`, instead got: {}'.format(i, type(value)))
+        self._q_values = q_values
+
+    def compute(self, batch):
+        r"""
+        Compute the value target :math:`(r + \gamma (1-d) \min_i \max_{a'} Q_{\phi_i}(s',a'))`.
+
+        Args:
+            batch (Batch): batch containing the transitions.
+        """
+        q_max = [torch.max(value(batch['states']), dim=1, keepdim=True)[0] for value in self._q_values]
+        value = torch.min(torch.cat(q_max, dim=1), dim=1)[0]
+        batch[self] = batch['rewards'] + self.gamma * (1 - batch['masks']) * value
+        return batch
+
+
 class TDValueReturn(TDReturn):
     r"""TD State Value Return
 
@@ -499,10 +617,10 @@ class TDValueReturn(TDReturn):
         """Evaluate the TD return on the given batch.
 
         Args:
-            batch (): batch containing transitions.
+            batch (Batch): batch containing transitions.
 
         Returns:
-            batch
+            Batch: batch
         """
         target = batch['rewards'] + self.gamma * (1 - batch['masks']) * self.target_value(batch['states'])
         batch[self] = target - self.value(batch['states'])
@@ -540,10 +658,10 @@ class TDQValueReturn(TDReturn):
         """Evaluate the TD return on the given batch.
 
         Args:
-            batch (): batch containing transitions.
+            batch (Batch): batch containing transitions.
 
         Returns:
-            batch
+            Batch: batch
         """
         action = self.policy.predict(batch['states'])
         target = batch['rewards'] + self.gamma * (1 - batch['masks']) * self.target_qvalue(action)
@@ -555,7 +673,7 @@ class TDQLearningReturn(TDReturn):
     r"""TD Q-Learning Value Return
 
     Compute the one-step Q-Learning, given by:
-    .. math:: (r + \gamma (1-d) max_{a'} Q_{\phi_{target}}(s',a')) - Q_{\phi}(s,a)
+    .. math:: (r + \gamma (1-d) \max_{a'} Q_{\phi_{target}}(s',a')) - Q_{\phi}(s,a)
 
     where if the actions :math:`a` are discrete, then :math:`a'` is selected such that it maximizes the Q-value, while
     if :math:`a` are continuous, :math:`a'`, with the assumption that the policy is fully differentiable, is selected
@@ -586,10 +704,10 @@ class TDQLearningReturn(TDReturn):
         """Evaluate the TD return on the given batch.
 
         Args:
-            batch (): batch containing transitions.
+            batch (Batch): batch containing transitions.
 
         Returns:
-            batch
+            Batch: batch
         """
         q_max = torch.max(self.target_qvalue(batch['states']), dim=1, keepdim=True)[0]
         target = batch['rewards'] + self.gamma * (1 - batch['masks']) * q_max

pyrobolearn/samplers/sampler.py

@@ -21,7 +21,7 @@ References:
 import torch.utils.data.sampler as torch_sampler
 import torch.utils.data.dataset as torch_dataset
 
-from pyrobolearn.storages import RolloutStorage
+from pyrobolearn.storages import Storage
 
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
@@ -34,7 +34,11 @@ __status__ = "Development"
 
 
 class Sampler(object):
-    pass
+
+    @property
+    def batch_size(self):
+        """Return the batch size."""
+        return 0
 
 
 class RandomSampler(Sampler):
@@ -82,8 +86,8 @@ class StorageSampler(Sampler):
     @storage.setter
     def storage(self, storage):
         """Set the storage instance."""
-        if not isinstance(storage, RolloutStorage):
-            raise TypeError("Expecting the storage to be an instance of `RolloutStorage`, instead got: "
+        if not isinstance(storage, Storage):
+            raise TypeError("Expecting the storage to be an instance of `Storage`, instead got: "
                             "{}".format(type(storage)))
         self._storage = storage
 
@@ -147,3 +151,19 @@ class StorageSampler(Sampler):
         """Iterate over the storage."""
         for indices in self.sampler:
             yield self.storage.get_batch(indices)
+
+
+class BatchRandomSampler(StorageSampler):
+    r"""Batch Random Sampler
+
+    """
+
+    def __init__(self, storage, num_batches=10):
+        """
+        Initialize the storage sampler.
+
+        Args:
+            storage (RolloutStorage): rollout storage.
+            num_batches (int): number of batches
+        """
+        super(BatchRandomSampler, self).__init__(storage, num_batches=num_batches)

pyrobolearn/storages/storage.py

@@ -41,6 +41,22 @@ class Storage(object):
         """Load the storage from the disk."""
         return pickle.load(open(filename, 'r'))
 
+    @property
+    def size(self):
+        """Return the size of the storage. Need to be implemented in the child class."""
+        return 0
+
+    def get_batch(self, indices):
+        """Return a batch of the storage as a `Storage` type.
+
+        Args:
+            indices (list of int): indices. Each index must be between 0 and the size of the storage.
+
+        Returns:
+            Storage: batch containing a part of the storage.
+        """
+        pass
+
     # def __repr__(self):
     #     """Return a string representing the class."""
     #     return self.__class__.__name__
@@ -221,6 +237,11 @@ class ListStorage(list, PyTorchStorage):
         args = self._to(args, device=self.device, dtype=self.dtype)
         super(ListStorage, self).__init__(args)
 
+    @property
+    def size(self):
+        """Return the size of the storage."""
+        return len(self)
+
     def insert(self, index, item):
         """Insert item before index."""
         item = self._to(item, device=self.device, dtype=self.dtype)
@@ -257,6 +278,11 @@ class FIFOQueueStorage(queue.Queue, PyTorchStorage):
         """
         super(FIFOQueueStorage, self).__init__(maxsize)
 
+    @property
+    def size(self):
+        """Return the size of the storage."""
+        return len(self)
+
     def put(self, item, block=False, timeout=None):
         """Put an item into the queue.
 
@@ -314,6 +340,11 @@ class LIFOQueueStorage(queue.LifoQueue, PyTorchStorage):
         """
         super(LIFOQueueStorage, self).__init__(maxsize)
 
+    @property
+    def size(self):
+        """Return the size of the storage."""
+        return len(self)
+
     def put(self, item, block=False, timeout=None):
         """Put an item into the queue.
 
@@ -376,6 +407,11 @@ class PriorityQueueStorage(queue.PriorityQueue, PyTorchStorage):
         super(PriorityQueueStorage, self).__init__(maxsize)
         self.ascending = ascending
 
+    @property
+    def size(self):
+        """Return the size of the storage."""
+        return len(self)
+
     def put(self, item, block=False, timeout=None):
         """Put an item into the queue.
 
@@ -472,6 +508,11 @@ class SetStorage(set, PyTorchStorage):
         args = self._to(args, device=self.device, dtype=self.dtype)
         super(SetStorage, self).__init__(args)
 
+    @property
+    def size(self):
+        """Return the size of the storage."""
+        return len(self)
+
     def add(self, item):
         """Add new item in the set."""
         item = self._to(item, device=self.device, dtype=self.dtype)
@@ -512,6 +553,19 @@ class DictStorage(dict, PyTorchStorage):
             kwargs = self._to(kwargs, device=self.device, dtype=self.dtype)
         super(DictStorage, self).__init__(kwargs)
 
+    ##############
+    # Properties #
+    ##############
+
+    @property
+    def size(self):
+        """Return the size of the storage."""
+        return len(self)
+
+    ###########
+    # Methods #
+    ###########
+
     def update(self, dictionary, **kwargs):
         """Update the current dictionary from the other given dictionary / iterable.