update TD returns + value losses

pyrobolearn/estimators/estimator.py

@@ -1,10 +1,13 @@
 #!/usr/bin/env python
-"""Describes the various `Estimators` (aka `Returns`) used in reinforcement learning.
+"""Describes the various `Estimators` / `Returns` used in reinforcement learning.
 
 Dependencies:
 - `pyrobolearn.storages`
 """
 
+from abc import ABCMeta
+import torch
+
 from pyrobolearn.storages import RolloutStorage
 
 
@@ -18,7 +21,40 @@ __email__ = "briandelhaisse@gmail.com"
 __status__ = "Development"
 
 
-class Estimator(object):
+class BaseReturn(object):
+    r"""Base Return / Estimator
+
+    Base return / estimator computed in RL algorithms.
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, gamma=1.):
+        """
+        Initialize the base return / estimator.
+
+        Args:
+            gamma (float): discount factor
+        """
+        self.gamma = gamma
+
+    @property
+    def gamma(self):
+        """Return the discount factor"""
+        return self._gamma
+
+    @gamma.setter
+    def gamma(self, gamma):
+        """Set the discount factor"""
+        if gamma > 1.:
+            gamma = 1.
+        elif gamma < 0.:
+            gamma = 0.
+
+        self._gamma = gamma
+
+
+class Estimator(BaseReturn):
     r"""Estimator / Return
 
     Estimator / Return used for gradient based algorithms. The gradient is given by:
@@ -46,8 +82,8 @@ class Estimator(object):
             storage (RolloutStorage): rollout storage
             gamma (float): discount factor
         """
+        super(Estimator, self).__init__(gamma)
         self.storage = storage
-        self.gamma = gamma
 
     ##############
     # Properties #
@@ -66,21 +102,6 @@ class Estimator(object):
                             "type {}".format(storage, type(storage)))
         self._storage = storage
 
-    @property
-    def gamma(self):
-        """Return the discount factor"""
-        return self._gamma
-
-    @gamma.setter
-    def gamma(self, gamma):
-        """Set the discount factor"""
-        if gamma > 1.:
-            gamma = 1.
-        elif gamma < 0.:
-            gamma = 0.
-
-        self._gamma = gamma
-
     @property
     def returns(self):
         """Return the returns tensor from the rollout storage."""
@@ -234,7 +255,7 @@ class BaselineRewardEstimator(ActionRewardEstimator):
             self.returns[t] = self.rewards[t] + self.gamma * self.returns[t + 1] - self.baseline(self.states[t])
 
 
-class StateValueEstimator(Estimator):
+class ValueEstimator(Estimator):
     r"""State Value Estimator
 
     Return the state value function:
@@ -255,14 +276,14 @@ class StateValueEstimator(Estimator):
             storage (RolloutStorage): rollout storage
             gamma (float): discount factor
         """
-        super(StateValueEstimator, self).__init__(storage=storage, gamma=gamma)
+        super(ValueEstimator, self).__init__(storage=storage, gamma=gamma)
 
     def _evaluate(self):
         """Evaluate the estimator / return."""
         self.returns[:] = self.values.clone()
 
 
-class StateActionValueEstimator(Estimator):
+class QValueEstimator(Estimator):
     r"""State Action Value Estimator
 
     Return the state action value function:
@@ -283,7 +304,7 @@ class StateActionValueEstimator(Estimator):
             storage (RolloutStorage): rollout storage
             gamma (float): discount factor
         """
-        super(StateActionValueEstimator, self).__init__(storage=storage, gamma=gamma)
+        super(QValueEstimator, self).__init__(storage=storage, gamma=gamma)
 
     def _evaluate(self):
         """Evaluate the estimator / return."""
@@ -339,9 +360,9 @@ class AdvantageEstimator(Estimator):
 class TDResidualEstimator(Estimator):
     r"""TD Residual Estimator
 
-    Return the TD residual, given by:
+    Return the temporal difference (TD) residual, given by:
 
-    .. math:: \psi_t = \delta_t^{V^{\pi,\gamma}} = r_t + \gamma V^{\pi,\gamma}(s_{t+1}) - V^{\pi,\gamma}(s_t)
+    .. math:: \psi_t = \delta_t^{V^{\pi,\gamma}} = (r_t + \gamma V^{\pi,\gamma}(s_{t+1})) - V^{\pi,\gamma}(s_t)
 
     where,
 
@@ -414,3 +435,163 @@ class GAE(Estimator):
             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]
+
+
+class Return(BaseReturn):
+    r"""Return
+
+    Compared to the estimator that used the whole trajectory, it only uses transition tuples
+    :math:`(s_t, a_t, s_{t+1}, r_t, d)`, where is :math:`s_t` is the state at time :math:`t`, :math:`a_t` is the action
+    outputted by the policy in response to the state :math:`s_t`, :math:`s_{t+1}` is the next state returned by the
+    environment due to the policy's action :math:`a_t` and the current state :math:`s_t`, :math:`r_t` is the reward
+    signal returned by the environment, and :math:`d` is a boolean value that specifies if the task is over or not
+    (i.e. if it has failed or succeeded).
+
+    That is, estimators are estimated on the Monte-Carlo trajectories, while returns are estimated on temporal
+    difference errors [1].
+
+    References:
+        [1] "Reinforcement Learning: an Introduction" (chap 8.13), Sutton and Barto, 2018
+    """
+
+    def __init__(self, gamma=1.):
+        """
+        Initialize the return / estimator.
+
+        Args:
+            gamma (float): discount factor
+        """
+        super(Return, self).__init__(gamma)
+
+
+class TDReturn(Return):
+    r"""TD Return
+
+    Return based on the one-step temporal difference TD(0).
+    """
+    pass
+
+
+class TDValueReturn(TDReturn):
+    r"""TD State Value Return
+
+    Compute the state value one-step temporal difference TD(0), given by:
+    .. math:: \delta_t^{V} = (r + \gamma (1-d) V_{\phi_{target}}(s')) - V_{\phi}(s)
+
+    where :math:`V_{\phi_{target}}` is the target value function.
+    """
+
+    def __init__(self, value, target_value=None, gamma=1.):
+        """
+        Initialize the TD state value return.
+
+        Args:
+            value (ValueApproximator): state value function.
+            target_value (ValueApproximator): target state value function.
+            gamma (float): discount factor
+        """
+        super(TDValueReturn, self).__init__(gamma)
+        self.value = value
+        if target_value is None:
+            self.target_value = self.value
+
+    def evaluate(self, batch):
+        """Evaluate the TD return on the given batch.
+
+        Args:
+            batch (): batch containing transitions.
+
+        Returns:
+            batch
+        """
+        target = batch['rewards'] + self.gamma * (1 - batch['masks']) * self.target_value(batch['states'])
+        batch[self] = target - self.value(batch['states'])
+        return batch
+
+
+class TDQValueReturn(TDReturn):
+    r"""TD Action Value Return
+
+    Compute the action value one-step temporal difference TD(0), given by:
+    .. math:: (r + \gamma (1-d) Q_{\phi_{target}}(s',a')) - Q_{\phi}(s,a)
+    where :math:`a'` is the action performed by the policy given the state :math:`s'`.
+
+    This is also known as the Sarsa (an on-policy TD control) algorithm.
+    """
+
+    def __init__(self, q_value, policy, target_qvalue=None, gamma=1.):
+        """
+        Initialize the TD state-action value return.
+
+        Args:
+            q_value (QValueApproximator): Q-value function.
+            policy (Policy): policy to compute the action a'.
+            target_qvalue (QValueApproximator, None): target Q-value function. If None, it will use the given
+                :attr:`q_value`.
+            gamma (float): discount factor
+        """
+        super(TDQValueReturn, self).__init__(gamma)
+        self.q_value = q_value
+        self.policy = policy
+        if target_qvalue is None:
+            self.target_qvalue = self.q_value
+
+    def evaluate(self, batch):
+        """Evaluate the TD return on the given batch.
+
+        Args:
+            batch (): batch containing transitions.
+
+        Returns:
+            batch
+        """
+        action = self.policy.predict(batch['states'])
+        target = batch['rewards'] + self.gamma * (1 - batch['masks']) * self.target_qvalue(action)
+        batch[self] = target - self.q_value(batch['states'])
+        return batch
+
+
+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)
+
+    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
+    such that it maximizes locally the Q-value. That is, in the latter case, the action :math:`a'` is first computed
+    using the policy :math:`\pi(a'|s')`, then by taking the gradient of the Q-value with respect to the action, we
+    increment the initial action by :math:`a' = a' + \alpha \grad_{a} Q_{\phi_{target}}(s',a)`. This increment step
+    can be computed for few iterations such that it locally maximizes the :math:` Q_{\phi_{target}}`.
+
+    This is known as the Q-Learning (an off-policy TD control) algorithm.
+    """
+
+    def __init__(self, q_value, target_qvalue=None, gamma=1.):
+        """
+        Initialize the TD Q-Learning value return.
+
+        Args:
+            q_value (QValueApproximator): Q-value function.
+            target_qvalue (QValueApproximator): target Q-value function. If None, it will use the given
+                :attr:`q_value`.
+            gamma (float): discount factor
+        """
+        super(TDQLearningReturn, self).__init__(gamma)
+        self.q_value = q_value
+        if target_qvalue is None:
+            self.target_qvalue = self.q_value
+
+    def evaluate(self, batch):
+        """Evaluate the TD return on the given batch.
+
+        Args:
+            batch (): batch containing transitions.
+
+        Returns:
+            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
+        batch[self] = target - self.q_value(batch['states'])
+        return batch

pyrobolearn/losses/losses.py

@@ -7,6 +7,7 @@ Losses are evaluated on model parameters, data batches / storages, or transition
 import torch
 
 from pyrobolearn.losses.loss import Loss
+from pyrobolearn.estimators.estimator import TDReturn
 
 __author__ = "Brian Delhaisse"
 __copyright__ = "Copyright 2018, PyRoboLearn"
@@ -57,6 +58,24 @@ class ValueLoss(Loss):
         return 0.5 * (returns - values).pow(2).mean()
 
 
+class QLoss(Loss):
+    r"""QLoss
+
+    This computes :math:`\frac{1}{|B|} \sum_{s \in B} Q_{s, \mu_{\theta}(s)}}`, where :math:`\mu_\theta` is the policy.
+    """
+
+    def __init__(self, q_value, policy):
+        super(QLoss, self).__init__()
+        # TODO: check that the Q-Value accepts as inputs the actions
+        self.q_value = q_value
+        self.policy = policy
+
+    def compute(self, batch):
+        actions = self.policy.predict(batch['observations'])
+        q_values = self.q_value(batch['observations'], actions)
+        return -q_values.mean()
+
+
 class HuberLoss(Loss):
     r"""Huber Loss
 
@@ -300,11 +319,30 @@ class MSBELoss(Loss):
         [2] "Reinforcement Learning: An Introduction", Sutton and Barto, 2018
     """
 
-    def __init__(self):
+    def __init__(self, td_return):
+        """
+        Initialize the mean-squared Bellman error (MSBE).
+
+        Args:
+            td_return (TDReturn): Temporal difference return.
+        """
         super(MSBELoss, self).__init__()
+        if not isinstance(td_return, TDReturn):
+            raise TypeError("Expecting the given 'td_return' to be an instance of `TDReturn`, instead got: "
+                            "{}".format(type(td_return)))
+        self._td = td_return
 
     def compute(self, batch):
-        pass
+        """
+        Compute the mean-squared TD return.
+
+        Args:
+            batch:
+
+        Returns:
+            torch.Tensor: loss value.
+        """
+        return batch[self._td].pow(2).mean()
 
 
 # Tests