update estimators

pyrobolearn/estimators/README.md

@@ -1,6 +1,6 @@
-## Estimators
+## Estimators / Returns
 
-This folder provides the various estimators used in reinforcement learning (in the evaluation step of RL algorithms), such as:
+This folder provides the various estimators / returns used in reinforcement learning (in the evaluation step of RL algorithms), such as:
 - Total reward estimator
 - Action reward estimator
 - Baseline reward estimator

pyrobolearn/estimators/estimator.py

@@ -1,10 +1,13 @@
 #!/usr/bin/env python
-"""Describes the various `Estimators` used in reinforcement learning.
+"""Describes the various `Estimators` (aka `Returns`) used in reinforcement learning.
 
 Dependencies:
 - `pyrobolearn.storages`
 """
 
+from pyrobolearn.storages import RolloutStorage
+
+
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
 __credits__ = ["Brian Delhaisse"]
@@ -16,9 +19,9 @@ __status__ = "Development"
 
 
 class Estimator(object):
-    r"""Estimator / Returns
+    r"""Estimator / Return
 
-    Estimator used for gradient based algorithms. The gradient is given by:
+    Estimator / Return used for gradient based algorithms. The gradient is given by:
 
     .. math::
 
@@ -37,10 +40,10 @@ class Estimator(object):
 
     def __init__(self, storage, gamma=1.):
         """
-        Initialize the storage.
+        Initialize the estimator / return function.
 
         Args:
-            storage: storage
+            storage (RolloutStorage): rollout storage
             gamma (float): discount factor
         """
         self.storage = storage
@@ -58,6 +61,9 @@ class Estimator(object):
     @storage.setter
     def storage(self, storage):
         """Set the storage instance"""
+        if not isinstance(storage, RolloutStorage):
+            raise TypeError("Expecting the given storage to be an instance of `RolloutStorage`, instead got: {} with "
+                            "type {}".format(storage, type(storage)))
         self._storage = storage
 
     @property
@@ -75,32 +81,68 @@ class Estimator(object):
 
         self._gamma = gamma
 
+    @property
+    def returns(self):
+        """Return the returns tensor from the rollout storage."""
+        return self.storage.returns
+
+    @property
+    def rewards(self):
+        """Return the rewards tensor from the rollout storage."""
+        return self.storage.rewards
+
+    @property
+    def masks(self):
+        """Return the masks tensor from the rollout storage."""
+        return self.storage.masks
+
+    @property
+    def values(self):
+        """Return the value tensor V(s) from the rollout storage."""
+        return self.storage.values
+
+    @property
+    def action_values(self):
+        """Return the action value tensor Q(s,a) from the rollout storage."""
+        return self.storage.action_values
+
+    @property
+    def num_steps(self):
+        """Return the total number of steps in the storage."""
+        return self.storage.num_steps
+
+    @property
+    def states(self):
+        """Return the states / observations from the rollout storage."""
+        return self.storage.observations
+
     ###########
     # Methods #
     ###########
 
-    def _evaluate(self, storage):
+    def _evaluate(self):
         """Evaluate the estimator.
         To be implemented in the child class.
         """
         raise NotImplementedError
 
-    def evaluate(self, storage):
+    def evaluate(self, storage=None):
         """Evaluate the estimator"""
-        if storage is None:
-            storage = self.storage
-        return self.evaluate(storage)
+        if storage is not None:
+            self.storage = storage
+        return self._evaluate()
 
     #############
     # Operators #
     #############
 
     def __call__(self):
+        """Evaluate the estimator."""
         self.evaluate()
 
 
 class TotalRewardEstimator(Estimator):
-    r"""Total reward Estimator
+    r"""Total reward Estimator (aka finite-horizon undiscounted return)
 
     Return the total reward of the trajectory given by:
 
@@ -112,13 +154,21 @@ class TotalRewardEstimator(Estimator):
     """
 
     def __init__(self, storage, gamma=1.):
+        """
+        Initialize the total reward estimator (also known as finite-horizon undiscounted return)
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            gamma (float): discount factor
+        """
         super(TotalRewardEstimator, self).__init__(storage=storage, gamma=gamma)
 
-    def _evaluate(self, storage):
-        storage.returns[-1] = storage.rewards[-1]
-        for t in reversed(range(storage.num_steps)):
-            storage.returns[t] = storage.rewards[t] + self.gamma * storage.returns[t + 1]
-        storage.returns[:] = storage.returns[0]
+    def _evaluate(self):
+        """Evaluate the estimator / return."""
+        self.returns[-1] = self.rewards[-1]
+        for t in reversed(range(self.num_steps)):
+            self.returns[t] = self.rewards[t] + self.gamma * self.returns[t + 1]
+        self.returns[:] = self.returns[0]
 
 
 class ActionRewardEstimator(Estimator):
@@ -134,12 +184,20 @@ class ActionRewardEstimator(Estimator):
     """
 
     def __init__(self, storage, gamma=1.):
+        """
+        Initialize the action reward estimator.
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            gamma (float): discount factor
+        """
         super(ActionRewardEstimator, self).__init__(storage=storage, gamma=gamma)
 
-    def _evaluate(self, storage):
-        storage.returns[-1] = storage.rewards[-1]
-        for t in reversed(range(storage.num_steps)):
-            storage.returns[t] = storage.rewards[t] + self.gamma * storage.returns[t + 1]
+    def _evaluate(self):
+        """Evaluate the estimator / return."""
+        self.returns[-1] = self.rewards[-1]
+        for t in reversed(range(self.num_steps)):
+            self.returns[t] = self.rewards[t] + self.gamma * self.returns[t + 1]
 
 
 class BaselineRewardEstimator(ActionRewardEstimator):
@@ -156,14 +214,24 @@ class BaselineRewardEstimator(ActionRewardEstimator):
     """
 
     def __init__(self, storage, baseline, gamma=1.):
+        """
+        Initialize the TD residual Estimator.
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            baseline (callable): the baseline function which predicts a scalar given the state.
+            gamma (float): discount factor
+        """
         super(BaselineRewardEstimator, self).__init__(storage=storage, gamma=gamma)
+        if not callable(baseline):
+            raise TypeError("Expecting the given baseline to be callable.")
         self.baseline = baseline
 
-    def _evaluate(self, storage):
-        storage.returns[-1] = storage.rewards[-1]
-        for t in reversed(range(storage.num_steps)):
-            storage.returns[t] = storage.rewards[t] + self.gamma * storage.returns[t + 1] \
-                                - self.baseline(storage.states[t])
+    def _evaluate(self):
+        """Evaluate the estimator / return."""
+        self.returns[-1] = self.rewards[-1]
+        for t in reversed(range(self.num_steps)):
+            self.returns[t] = self.rewards[t] + self.gamma * self.returns[t + 1] - self.baseline(self.states[t])
 
 
 class StateValueEstimator(Estimator):
@@ -180,10 +248,18 @@ class StateValueEstimator(Estimator):
     """
 
     def __init__(self, storage, gamma=1.):
+        """
+        Initialize the state value estimator.
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            gamma (float): discount factor
+        """
         super(StateValueEstimator, self).__init__(storage=storage, gamma=gamma)
 
-    def _evaluate(self, storage):
-        storage.returns = storage.values.clone()
+    def _evaluate(self):
+        """Evaluate the estimator / return."""
+        self.returns[:] = self.values.clone()
 
 
 class StateActionValueEstimator(Estimator):
@@ -200,10 +276,18 @@ class StateActionValueEstimator(Estimator):
     """
 
     def __init__(self, storage, gamma=1.):
+        """
+        Initialize the station-action value estimator.
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            gamma (float): discount factor
+        """
         super(StateActionValueEstimator, self).__init__(storage=storage, gamma=gamma)
 
-    def _evaluate(self, storage):
-        storage.returns = storage.action_values.clone()
+    def _evaluate(self):
+        """Evaluate the estimator / return."""
+        self.returns[:] = self.action_values.clone()
 
 
 class AdvantageEstimator(Estimator):
@@ -238,10 +322,18 @@ class AdvantageEstimator(Estimator):
     """
 
     def __init__(self, storage, gamma=1.):
+        """
+        Initialize the Advantage estimator.
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            gamma (float): discount factor
+        """
         super(AdvantageEstimator, self).__init__(storage=storage, gamma=gamma)
 
-    def _evaluate(self, storage):
-        storage.returns = storage.action_values - storage.values
+    def _evaluate(self):
+        """Evaluate the estimator / return."""
+        self.returns[:] = self.action_values - self.values
 
 
 class TDResidualEstimator(Estimator):
@@ -261,12 +353,20 @@ class TDResidualEstimator(Estimator):
     """
 
     def __init__(self, storage, gamma=1.):
+        """
+        Initialize the TD residual Estimator.
+
+        Args:
+            storage (RolloutStorage): rollout storage
+            gamma (float): discount factor
+        """
         super(TDResidualEstimator, self).__init__(storage=storage, gamma=gamma)
 
-    def _evaluate(self, storage=None):  # , next_value):
-        # storage.returns[-1] = next_value
-        for t in reversed(range(storage.num_steps)):
-            storage.returns[t] = storage.rewards[t] + self.gamma * storage.values[t + 1] - storage.values[t]
+    def _evaluate(self):  # , next_value):
+        """Evaluate the estimator / return."""
+        # self.returns[-1] = next_value
+        for t in reversed(range(self.num_steps)):
+            self.returns[t] = self.rewards[t] + self.gamma * self.values[t + 1] - self.values[t]
 
 
 class GAE(Estimator):
@@ -298,7 +398,7 @@ class GAE(Estimator):
         Initialize the Generalized Advantage Estimator.
 
         Args:
-            storage:
+            storage (RolloutStorage): rollout storage
             gamma (float): discount factor
             tau (float): trace-decay parameter (which is a bias-variance tradeoff). If :math:`\tau=1`, this results
                 in a Monte Carlo method, while :math:`\tau=0` results in a one-step TD methods.
@@ -306,10 +406,11 @@ class GAE(Estimator):
         super(GAE, self).__init__(storage=storage, gamma=gamma)
         self.tau = tau
 
-    def _evaluate(self, storage):  # , next_value):
-        # storage.values[-1] = next_value
+    def _evaluate(self):  # , next_value):
+        """Evaluate the estimator / return."""
+        # self.values[-1] = next_value
         gae = 0
-        for t in reversed(range(storage.num_steps)):
-            delta = storage.rewards[t] + self.gamma * storage.values[t + 1] * storage.masks[t + 1] - storage.values[t]
-            gae = delta + self.gamma * self.tau * storage.masks[t + 1] * gae
-            storage.returns[t] = gae + storage.values[t]
+        for t in reversed(range(self.num_steps)):
+            delta = self.rewards[t] + self.gamma * self.values[t + 1] * self.masks[t + 1] - self.values[t]
+            gae = delta + self.gamma * self.tau * self.masks[t + 1] * gae
+            self.returns[t] = gae + self.values[t]