diff --git a/pyrobolearn/estimators/estimator.py b/pyrobolearn/estimators/estimator.py index d88b20d..35d5630 100644 --- a/pyrobolearn/estimators/estimator.py +++ b/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 diff --git a/pyrobolearn/losses/losses.py b/pyrobolearn/losses/losses.py index 05d3b8c..bd03fbb 100644 --- a/pyrobolearn/losses/losses.py +++ b/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