From 92bea93fde6ba6fe75a7cd14714f0820be700da1 Mon Sep 17 00:00:00 2001 From: Brian Delhaisse Date: Mon, 15 Apr 2019 21:40:35 +0200 Subject: [PATCH] refactor estimators/returns, losses, and storages --- pyrobolearn/algos/ppo.py | 2 +- pyrobolearn/algos/reinforce.py | 2 +- pyrobolearn/algos/sac.py | 12 +- pyrobolearn/losses/losses.py | 23 +- pyrobolearn/losses/policy_losses.py | 49 ++- pyrobolearn/losses/value_losses.py | 45 ++- pyrobolearn/returns/__init__.py | 2 +- pyrobolearn/returns/estimators.py | 597 ++++++++++++++++++++++++++++ pyrobolearn/returns/returns.py | 167 ++++++-- pyrobolearn/returns/targets.py | 67 ++-- pyrobolearn/samplers/sampler.py | 2 +- pyrobolearn/storages/er.py | 6 +- pyrobolearn/storages/storage.py | 46 +-- pyrobolearn/values/value.py | 45 ++- 14 files changed, 924 insertions(+), 141 deletions(-) create mode 100644 pyrobolearn/returns/estimators.py diff --git a/pyrobolearn/algos/ppo.py b/pyrobolearn/algos/ppo.py index 17ca5bf..6409830 100755 --- a/pyrobolearn/algos/ppo.py +++ b/pyrobolearn/algos/ppo.py @@ -184,7 +184,7 @@ class PPO(GradientRLAlgo): states, actions = policy.states, policy.actions logger.debug('create rollout storage') storage = RolloutStorage(num_steps=1000, state_shapes=states.merged_shape, - action_shapes=actions.merged_shape, num_processes=num_workers) + action_shapes=actions.merged_shape, num_trajectories=num_workers) logger.debug('create return estimator (GAE)') estimator = GAE(storage, gamma=gamma, tau=tau) logger.debug('create storage sampler') diff --git a/pyrobolearn/algos/reinforce.py b/pyrobolearn/algos/reinforce.py index d82a867..c87a13d 100755 --- a/pyrobolearn/algos/reinforce.py +++ b/pyrobolearn/algos/reinforce.py @@ -189,7 +189,7 @@ class REINFORCE(GradientRLAlgo): # create storage states, actions = policy.states, policy.actions storage = RolloutStorage(num_steps=1000, state_shapes=states.shape, action_shapes=actions.shape, - num_processes=num_workers) + num_trajectories=num_workers) sampler = StorageSampler(storage) # create estimator diff --git a/pyrobolearn/algos/sac.py b/pyrobolearn/algos/sac.py index 079853f..db93e2e 100755 --- a/pyrobolearn/algos/sac.py +++ b/pyrobolearn/algos/sac.py @@ -16,7 +16,7 @@ from pyrobolearn.exploration import ActionExploration, GaussianActionExploration from pyrobolearn.storages import ExperienceReplay from pyrobolearn.samplers import BatchRandomSampler -from pyrobolearn.returns import TDQValueReturn +from pyrobolearn.returns import ValueTarget, EntropyValueTarget from pyrobolearn.losses import MSBELoss, QLoss from pyrobolearn.optimizers import Adam @@ -260,7 +260,8 @@ class SAC(GradientRLAlgo): [3] RLKit: https://github.com/vitchyr/rlkit/blob/master/rlkit/torch/sac/sac.py """ - def __init__(self, task, approximators, gamma=0.99, lr=5e-4, polyak=0.995, capacity=10000, num_workers=1): + def __init__(self, task, approximators, gamma=0.99, lr=5e-4, polyak=0.995, alpha=0.2, capacity=10000, + num_workers=1): """ Initialize the SAC off-policy RL algorithm. @@ -271,6 +272,9 @@ class SAC(GradientRLAlgo): importance has the future rewards we get. lr (float): learning rate polyak (float): coefficient in the polyak averaging when updating the target approximators. + alpha (float): entropy regularization coefficient which controls the tradeoff between exploration and + exploitation. Higher :attr:`alpha` means more exploration, and lower :attr:`alpha` corresponds to more + exploitation. capacity (int): capacity of the experience replay storage. num_workers (int): number of processes / workers to run in parallel """ @@ -309,8 +313,8 @@ class SAC(GradientRLAlgo): exploration = ActionExploration(policy) # create targets - # q_target = - + q_target = ValueTarget(values=value_target, gamma=gamma) + v_target = EntropyValueTarget(q_values=q_values, policy=exploration, alpha=alpha) # create losses q_loss = MSBELoss(td_return=estimator) diff --git a/pyrobolearn/losses/losses.py b/pyrobolearn/losses/losses.py index b9d2609..30f8893 100644 --- a/pyrobolearn/losses/losses.py +++ b/pyrobolearn/losses/losses.py @@ -22,12 +22,13 @@ class FixedLoss(Loss): r"""Fixed Loss """ + def __init__(self, value): super(FixedLoss, self).__init__() - self.value = value + self.value = torch.tensor(value) def compute(self, batch): - return self.value * batch + return self.value class L2Loss(Loss): @@ -35,14 +36,22 @@ class L2Loss(Loss): Compute the L2 loss given by: :math:`1/2 * (y_{target} - y_{predict})^2` """ - def __init__(self, target, approximator): + + def __init__(self, target, predictor): super(L2Loss, self).__init__() - self.target = target - self.approximator = approximator + self._target = target + self._predictor = predictor def compute(self, batch): - # based on approximator check what we need - return 0.5 * (self.target(batch) - self.approximator(batch)).pow(2).mean() + if self._target in batch: + target = batch[self._target] + else: + target = self._target(batch) + if self._predictor in batch: + output = batch[self._predictor] + else: + output = self._predictor(batch) + return 0.5 * (target - output).pow(2).mean() class HuberLoss(Loss): diff --git a/pyrobolearn/losses/policy_losses.py b/pyrobolearn/losses/policy_losses.py index fabf92e..922c129 100644 --- a/pyrobolearn/losses/policy_losses.py +++ b/pyrobolearn/losses/policy_losses.py @@ -5,6 +5,7 @@ import torch from pyrobolearn.losses.loss import Loss +from pyrobolearn.returns.estimators import Estimator __author__ = "Brian Delhaisse" __copyright__ = "Copyright 2018, PyRoboLearn" @@ -26,29 +27,39 @@ class PGLoss(Loss): where :math:`\psi_t` is the associated return estimator, which can be for instance, the total reward estimator :math:`\psi_t = R(\tau)` (where :math:`\tau` represents the whole trajectory), the state action value estimator :math:`\psi_t = Q(s_t, a_t)`, or the advantage estimator :math:`\psi_t = A_t = Q(s_t, a_t) - V(s_t)`. Other - estimators are also possible. + returns are also possible. The gradient with respect to the parameters :math:`\theta` is then given by: - .. math:: g = \mathbb{E}[ \nabla_\theta \log \pi_{\theta}(a_t | s_t) + .. math:: g = \mathbb{E}[ \nabla_{\theta} \log \pi_{\theta}(a_t | s_t) ] References: [1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017 [2] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016 """ - def __init__(self): + def __init__(self, estimator): + """ + Initialize the Policy Gradient loss. + + Args: + estimator (Estimator): estimator that has been used on the rollout storage / batch. + """ super(PGLoss, self).__init__() + if not isinstance(estimator, Estimator): + raise TypeError("The given estimator is not an instance of `Estimator`, instead got: " + "{}".format(type(estimator))) + self._estimator = estimator def compute(self, batch): log_curr_pi = batch.current['action_distributions'] log_curr_pi = log_curr_pi.log_probs(batch.current['actions']) - estimator = batch['estimator'] + estimator = batch[self._estimator] loss = torch.exp(log_curr_pi) * estimator return -loss.mean() def latex(self): - return "\\mathbb{E}[ r_t(\\theta) A_t ]" + return r"\mathbb{E}[ r_t(\theta) A_t ]" class CPILoss(Loss): @@ -67,11 +78,18 @@ class CPILoss(Loss): [2] "Proximal Policy Optimization Algorithms", Schulman et al., 2017 """ - def __init__(self): + def __init__(self, estimator): """ Initialize the CPI Loss. + + Args: + estimator (Estimator): estimator that has been used on the rollout storage / batch. """ super(CPILoss, self).__init__() + if not isinstance(estimator, Estimator): + raise TypeError("The given estimator is not an instance of `Estimator`, instead got: " + "{}".format(type(estimator))) + self._estimator = estimator def compute(self, batch): # policy_distribution, old_policy_distribution, estimator): # ratio = policy_distribution / old_policy_distribution @@ -81,13 +99,13 @@ class CPILoss(Loss): log_prev_pi = log_prev_pi.log_probs(batch['actions']) ratio = torch.exp(log_curr_pi - log_prev_pi) - estimator = batch['estimator'] + estimator = batch[self._estimator] loss = ratio * estimator return -loss.mean() def latex(self): - return "\\mathbb{E}[ r_t(\\theta) A_t ]" + return r"\mathbb{E}[ r_t(\theta) A_t ]" class CLIPLoss(Loss): @@ -105,15 +123,20 @@ class CLIPLoss(Loss): [1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017 """ - def __init__(self, clip=0.2): + def __init__(self, estimator, clip=0.2): """ Initialize the loss. Args: - epsilon (float): clip parameter + estimator (Estimator): estimator that has been used on the rollout storage / batch. + clip (float): clip parameter """ super(CLIPLoss, self).__init__() self.eps = clip + if not isinstance(estimator, Estimator): + raise TypeError("The given estimator is not an instance of `Estimator`, instead got: " + "{}".format(type(estimator))) + self._estimator = estimator def compute(self, batch): # , policy_distribution, old_policy_distribution, estimator): log_curr_pi = batch.current['action_distributions'] @@ -122,13 +145,13 @@ class CLIPLoss(Loss): log_prev_pi = log_prev_pi.log_probs(batch['actions']) ratio = torch.exp(log_curr_pi - log_prev_pi) - estimator = batch['estimator'] + estimator = batch[self._estimator] loss = torch.min(ratio * estimator, torch.clamp(ratio, 1.0-self.eps, 1.0+self.eps) * estimator) return -loss.mean() def latex(self): - return "\\mathbb{E}[ \\min(r_t(\\theta) A_t, clip(r_t(\\theta), 1-\\epsilon, 1+\\epsilon) A_t) ]" + return r"\mathbb{E}[ \min(r_t(\theta) A_t, clip(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t) ]" class KLPenaltyLoss(Loss): @@ -181,7 +204,7 @@ class EntropyLoss(Loss): [3] "Proximal Policy Optimization Algorithms", Schulman et al., 2017 """ - def __init__(self): # approximator): + def __init__(self): super(EntropyLoss, self).__init__() def compute(self, batch): diff --git a/pyrobolearn/losses/value_losses.py b/pyrobolearn/losses/value_losses.py index d08cd3d..aef4b09 100644 --- a/pyrobolearn/losses/value_losses.py +++ b/pyrobolearn/losses/value_losses.py @@ -5,6 +5,8 @@ import torch from pyrobolearn.losses.loss import Loss +from pyrobolearn.policies import Policy +from pyrobolearn.values import QValue from pyrobolearn.returns import TDReturn __author__ = "Brian Delhaisse" @@ -36,14 +38,39 @@ class QLoss(Loss): """ def __init__(self, q_value, policy): + """ + Initialize the Q-loss. + + Args: + q_value (QValue): Q-value function approximator. + policy (Policy): policy. + """ super(QLoss, self).__init__() - # TODO: check that the Q-Value accepts as inputs the actions - self.q_value = q_value - self.policy = policy + + # check the given q_value + if not isinstance(q_value, QValue): + raise TypeError("Expecting the given q_value to be an instance of `QValue`, instead got: " + "{}".format(type(q_value))) + self._q_value = q_value + + # check the policy + if not isinstance(policy, Policy): + raise TypeError("Expecting the given policy to be an instance of `Policy`, instead got: " + "{}".format(type(policy))) + self._policy = policy def compute(self, batch): - actions = self.policy.predict(batch['observations']) - q_values = self.q_value(batch['observations'], actions) + """ + Compute the loss on the given batch. + + Args: + batch (Batch): batch that contains the 'states'. + + Returns: + torch.tensor: loss scalar value + """ + actions = self._policy.predict(batch['states']) + q_values = self._q_value(batch['states'], actions) return -q_values.mean() @@ -90,9 +117,13 @@ class MSBELoss(Loss): Compute the mean-squared TD return. Args: - batch: batch + batch (Batch): batch that contains the td returns. Returns: torch.Tensor: loss value. """ - return batch[self._td].pow(2).mean() + if self._td in batch: + returns = batch[self._td] + else: + returns = self._td.evaluate(batch, store=False) + return returns.pow(2).mean() diff --git a/pyrobolearn/returns/__init__.py b/pyrobolearn/returns/__init__.py index 310ae2d..89ae44c 100644 --- a/pyrobolearn/returns/__init__.py +++ b/pyrobolearn/returns/__init__.py @@ -1,6 +1,6 @@ # import the various returns -from .estimator import * +from .estimators import * # import the various targets from .targets import * diff --git a/pyrobolearn/returns/estimators.py b/pyrobolearn/returns/estimators.py new file mode 100644 index 0000000..fc8c88f --- /dev/null +++ b/pyrobolearn/returns/estimators.py @@ -0,0 +1,597 @@ +#!/usr/bin/env python +"""Describes the various `Estimators` / `Returns` used in reinforcement learning. + +- Estimators are evaluated on trajectories (i.e. RolloutStorage). +- Returns are estimated on batches of transition tuples (s, a, s', r, d). + +In both case, they add + +Dependencies: +- `pyrobolearn.storages` +""" + +from abc import ABCMeta +import collections + +import torch + +from pyrobolearn.storages import RolloutStorage +from pyrobolearn.values import Value, QValue + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class Estimator(object): + r"""Estimator / Return + + Estimator / Return used for gradient based algorithms. The gradient is given by: + + .. math:: + + g = \mathbb{E}_{s_{0:T}, a_{0:T}}[ \sum_{t=0}^{T} \psi_t \nabla_{\theta} \log \pi_{\theta}(a_t | s_t) ] + + where :math:`\psi_t` is the associated return estimator, and the expectation is over the "states and actions + sampled sequentially from the dynamics model :math:`P(s_{t+1} | s_t, a_t)` and policy :math:`\pi(a_t | s_t)`, + respectively" [1]. + + Note that a discount factor can be used for :math:`\psi_t`. If the discount factor :math:`\gamma` is smaller + than 1, then it will reduces the variance but at the cost of introducing a bias. + + Reference: + [1] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016 + """ + + def __init__(self, storage, gamma=1.): + """ + Initialize the estimator / return function. + + Args: + storage (RolloutStorage): rollout storage + gamma (float): discount factor + """ + self.gamma = gamma + self.storage = storage + + ############## + # Properties # + ############## + + @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 storage(self): + """Return the storage instance""" + return self._storage + + @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 + def returns(self): + """Return the returns tensor from the rollout storage.""" + return self.storage['self'] + + @property + def states(self): + """Return the states / observations from the rollout storage.""" + return self.storage['states'] + + @property + def rewards(self): + """Return the rewards tensor from the rollout storage.""" + return self.storage['rewards'] + + @property + def actions(self): + """Return the actions from the rollout storage.""" + return self.storage['actions'] + + @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 + + ########### + # Methods # + ########### + + def _evaluate(self): + """Evaluate the estimator on the given rollout storage. To be implemented in the child class. + + Returns: + torch.Tensor: the computed returns. + """ + raise NotImplementedError + + def evaluate(self, storage=None): + """Evaluate the estimator on the given rollout storage. + + Args: + storage (RolloutStorage): rollout storage. + + Returns: + torch.Tensor: the computed returns. + """ + # set storage + if storage is not None: + self.storage = storage + + # create returns in the storage + if self not in self.storage: + self.storage.create_new_entry(key=self, shapes=1, num_steps=self.num_steps+1) + + return self._evaluate() + + ############# + # Operators # + ############# + + def __call__(self, storage=None): + """Evaluate the estimator on the rollout storage. + + Args: + storage (RolloutStorage): rollout storage. + + Returns: + torch.Tensor: the computed returns. + """ + return self.evaluate(storage=storage) + + +class TotalRewardEstimator(Estimator): + r"""Total reward Estimator (aka (finite-horizon) discounted return) + + Return the total reward of the trajectory given by: + + .. math:: + + \psi_t = R(\tau) = \sum_{t'=0}^{T} \gamma^{t'} r_{t'} + + where :math:`\tau` represents a trajectory :math:`\tau = (s_0, a_0, s_1,..., a_{T-1}, s_T)`. + """ + + 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): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage + returns, rewards, masks = self.returns, self.rewards, self.masks + + # compute the returns + returns[-1] = rewards[-1] + for t in reversed(range(self.num_steps)): + returns[t] = rewards[t] + self.gamma * masks[t+1] * returns[t + 1] + returns[:] = returns[0] + + return returns + + +class ActionRewardEstimator(Estimator): + r"""Action Reward Estimator + + Return the accumulated reward following action :math:`a_t`: + + .. math:: + + \psi_t = \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'} + + This is based on the observation that future actions don't have any effects on previous rewards. + """ + + 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): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage + returns, rewards, masks = self.returns, self.rewards, self.masks + + # compute the returns + returns[-1] = rewards[-1] + for t in reversed(range(self.num_steps)): + returns[t] = rewards[t] + self.gamma * masks[t+1] * returns[t + 1] + + return returns + + +class BaselineRewardEstimator(ActionRewardEstimator): # TODO: check with masks + r"""Baseline Reward Estimator + + Return the accumulated reward following action a_t minus a baseline depending on the state: + + .. math:: + + \psi_t = \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'} - b(s_t) + + This baseline allows to reduce the variance, and does not introduce any biases. + The baseline is often selected to be the state value function: :math:`b(s_t) = V^{\pi}(s_t)`. + """ + + 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): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage + returns, rewards, masks, states = self.returns, self.rewards, self.masks, self.states + + # compute the returns + returns[-1] = rewards[-1] + for t in reversed(range(self.num_steps)): + state = [state[t] for state in states] + returns[t] = rewards[t] + self.gamma * masks[t+1] * returns[t + 1] - self.baseline(state) + return returns + + +class ValueEstimator(Estimator): # TODO: check with masks + r"""State Value Estimator + + Return the state value function: + + .. math:: + + \psi_t = V^{\pi, \gamma}(s_t) = \mathbb{E}_{s_{t+1:T}, a_{t:T}}[ \sum_{l=0}^{T} \gamma^{l} r_{t+l} ] + + If the discount factor :math:`\gamma` is smaller than 1, then it will reduces the variance but at the cost + of introducing a bias. + """ + + def __init__(self, storage, value, gamma=1.): + """ + Initialize the state value estimator. + + Args: + storage (RolloutStorage): rollout storage + value (Value): value function approximator. + gamma (float): discount factor + """ + super(ValueEstimator, self).__init__(storage=storage, gamma=gamma) + # check value function + if not isinstance(value, Value): + raise TypeError("Expecting the given value to be an instance of `Value`, instead got: " + "{}".format(type(value))) + self._value = value + + def _evaluate(self): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage if present. If not, evaluate the value based on the states. + if self._value in self.storage: + values = self.storage[self._value] + else: + values = self._value(self.states) + + # compute the returns + self.returns[:] = values.clone() + return self.returns + + +class QValueEstimator(Estimator): # TODO: check with masks + r"""State Action Value Estimator + + Return the state action value function: + + .. math:: + + \psi_t = Q^{\pi, \gamma}(s_t, a_t) = \mathbb{E}_{s_{t+1:T}, a_{t+1:T}}[ \sum_{l=0}^{T} \gamma^{l} r_{t+l} ] + + If the discount factor :math:`\gamma` is smaller than 1, then it will reduces the variance but at the cost + of introducing a bias. + """ + + def __init__(self, storage, q_value, gamma=1.): + """ + Initialize the station-action value estimator. + + Args: + storage (RolloutStorage): rollout storage + q_value (QValue): Q-value function approximator. + gamma (float): discount factor + """ + super(QValueEstimator, self).__init__(storage=storage, gamma=gamma) + + # check q-value + if not isinstance(q_value, QValue): + raise TypeError("Expecting the given q_value to be an instance of `QValue`, instead got: " + "{}".format(type(q_value))) + self._q_value = q_value + + def _evaluate(self): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage if present. If not, evaluate the value based on the states and actions. + if self._q_value in self.storage: + q_values = self.storage[self._q_value] + else: + q_values = self._q_value(self.states, self.actions) + + # compute the returns + self.returns[:] = q_values.clone() + return self.returns + + +class AdvantageEstimator(Estimator): # TODO: check with masks + r"""Advantage Estimator + + Return the advantage function: + + .. math:: + + \psi_t = A^{\pi,\gamma}(s_t, a_t) = Q^{\pi,\gamma}(s_t,a_t) - V^{\pi,\gamma}(s_t) + + where + + .. math:: + + Q^{\pi,\gamma}(s_t, a_t) = \mathbb{E}_{s_{t+1:T}, a_{t+1:T}}[ \sum_{l=0}^{T} \gamma^{l} r_{t+l} ] + + V^{\pi,\gamma}(s_t) = \mathbb{E}_{s_{t+1:T}, a_{t:T}}[ \sum_{l=0}^{T} \gamma^{l} r_{t+l} ] + + The advantage function represents what is the 'advantage' of taking a certain action at a certain state. + If the discount factor :math:`\gamma` is smaller than 1, then it will reduces the variance but at the cost + of introducing a bias. + + Note that: + + .. math:: + + A^{\pi,\gamma}(s_t, a_t) &= Q^{\pi,\gamma}(s_t,a_t) - V^{\pi,\gamma}(s_t) \\ + &= \mathbb{E}_{s_{t+1}}[ Q^{\pi,\gamma}(s_t,a_t) - V^{\pi,\gamma}(s_t) ] \\ + &= \mathbb{E}_{s_{t+1}}[ r_t + \gamma V^{\pi,\gamma}(s_{t+1}) - V^{\pi,\gamma}(s_t) ] \\ + &= \mathbb{E}_{s_{t+1}}[ \delta_t^{V^{\pi,\gamma}} ] + """ + + def __init__(self, storage, value, q_value, gamma=1.): + """ + Initialize the Advantage estimator. + + Args: + storage (RolloutStorage): rollout storage + value (Value): value function approximator. + q_value (QValue): Q-value function approximator. + gamma (float): discount factor + """ + super(AdvantageEstimator, self).__init__(storage=storage, gamma=gamma) + + # check value function + if not isinstance(value, Value): + raise TypeError("Expecting the given value to be an instance of `Value`, instead got: " + "{}".format(type(value))) + self._value = value + + # check q-value + if not isinstance(q_value, QValue): + raise TypeError("Expecting the given q_value to be an instance of `QValue`, instead got: " + "{}".format(type(q_value))) + self._q_value = q_value + + def _evaluate(self): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + + # get values from storage if present. If not, evaluate the value based on the states. + if self._value in self.storage: + values = self.storage[self._value] + else: + values = self._value(self.states) + + # get Q-values from storage if present. If not, evaluate the Q-values based on the states and actions. + if self._q_value in self.storage: + q_values = self.storage[self._q_value] + else: + q_values = self._q_value(self.states, self.actions) + + self.returns[:] = q_values - values + return self.returns + + +class TDResidualEstimator(Estimator): # TODO: check with masks + r"""TD Residual Estimator + + 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) + + where, + + .. math:: V^{\pi,\gamma}(s_t) = \mathbb{E}_{s_{t+1:T}, a_{t:T}}[ \sum_{l=0}^{T} \gamma^{l} r_{t+l} ] + + Note that: + + .. math:: A^{\pi,\gamma}(s_t, a_t) = \mathbb{E}_{s_{t+1}}[ \delta_t^{V^{\pi,\gamma}} ] + """ + + def __init__(self, storage, value, gamma=1.): + """ + Initialize the TD residual Estimator. + + Args: + storage (RolloutStorage): rollout storage + value (Value): value function approximator. + gamma (float): discount factor + """ + super(TDResidualEstimator, self).__init__(storage=storage, gamma=gamma) + # check value function + if not isinstance(value, Value): + raise TypeError("Expecting the given value to be an instance of `Value`, instead got: " + "{}".format(type(value))) + self._value = value + + def _evaluate(self): # , next_value): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage + returns, rewards, masks = self.returns, self.rewards, self.masks + + # get values from storage if present. If not, evaluate the value based on the states. + if self._value in self.storage: + values = self.storage[self._value] + else: + values = self._value(self.states) + + # compute the returns + # self.returns[-1] = next_value + for t in reversed(range(self.num_steps)): + returns[t] = rewards[t] + self.gamma * masks[t+1] * values[t + 1] - values[t] + + return returns + + +class GAE(Estimator): + r"""Generalized Advantage Estimator + + Return the GAE, which is the exponentially-weighted average of the discounted sum of the TD/Bellman residuals: + + .. math:: \psi_t = A^{GAE(\gamma, \tau)}_t = \sum_{l=0}^{T} (\gamma \tau)^{l} \delta_{t+l}^{V} + + Notes: + + .. math:: + + GAE(\gamma, 0) = \delta_t^{V} = r_t + \gamma V(s_{t+1}) - V(s_t) + + GAE(\gamma, 1) = \sum_{l=0}^{T} \gamma^{l} \delta_{t+l}^{V} = \sum_{l=0}^{T} \gamma^{l} r_{t+l} - V(s_t) + + where :math:`GAE(\gamma, 1)` has high variance, while :math:`GAE(\gamma, 0)` has usually lower variance. + + A compromise between bias and variance is made by controlling the open trace-decay parameter :math:`\tau`. + Good values for GAE are obtained when :math:`\gamma` and :math:`\tau` are in :math:`[0.9,0.99]`. + + References: + [1] "High-Dimensional Continuous Control using Generalized Advantage Estimation", Schulman et al., 2016 + """ + + def __init__(self, storage, value, gamma=0.98, tau=0.99): + """ + Initialize the Generalized Advantage Estimator. + + Args: + storage (RolloutStorage): rollout storage + value (Value): value function approximator. + 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. + """ + super(GAE, self).__init__(storage=storage, gamma=gamma) + self.tau = tau + # check value function + if not isinstance(value, Value): + raise TypeError("Expecting the given value to be an instance of `Value`, instead got: " + "{}".format(type(value))) + self._value = value + + def _evaluate(self): + """Evaluate the estimator / return. + + Returns: + torch.Tensor: the computed returns. + """ + # get reference to data in storage + returns, rewards, masks = self.returns, self.rewards, self.masks + + # get values from storage if present. If not, evaluate the value based on the states. + if self._value in self.storage: + values = self.storage[self._value] + else: + values = self._value(self.states) + + # compute the returns + # self.values[-1] = next_value + gae = 0 + for t in reversed(range(self.num_steps)): + delta = rewards[t] + self.gamma * values[t + 1] * masks[t + 1] - values[t] + gae = delta + self.gamma * self.tau * masks[t + 1] * gae + returns[t] = gae + values[t] + + return returns diff --git a/pyrobolearn/returns/returns.py b/pyrobolearn/returns/returns.py index b1bf2d5..d43b961 100644 --- a/pyrobolearn/returns/returns.py +++ b/pyrobolearn/returns/returns.py @@ -8,12 +8,14 @@ Dependencies: - `pyrobolearn.values` """ +from abc import ABCMeta + import torch from pyrobolearn.storages import Batch -from pyrobolearn.values import Value, QValue +from pyrobolearn.values import Value, QValue, QValueOutput from pyrobolearn.policies import Policy -from pyrobolearn.returns.estimator import BaseReturn +from pyrobolearn.returns.targets import ValueTarget, QValueTarget, QLearningTarget __author__ = "Brian Delhaisse" @@ -26,7 +28,7 @@ __email__ = "briandelhaisse@gmail.com" __status__ = "Development" -class Return(BaseReturn): +class Return(object): r"""Return Compared to the estimator that used the whole trajectory, it only uses transition tuples @@ -43,14 +45,46 @@ class Return(BaseReturn): [1] "Reinforcement Learning: an Introduction" (chap 8.13), Sutton and Barto, 2018 """ - def __init__(self, gamma=1.): + def _evaluate(self, batch): """ - Initialize the return / estimator. + Evaluate the return on the given batch. Args: - gamma (float): discount factor + batch (Batch): batch containing the transitions. + + Returns: + torch.Tensor: evaluated return. """ - super(Return, self).__init__(gamma) + raise NotImplementedError + + def evaluate(self, batch, store=True): + """ + Evaluate the return on the given batch. + + Args: + batch (Batch): batch containing the transitions. + store (bool): If True, it will save the evaluation of the target in the given batch. + + Returns: + torch.Tensor: evaluated return. + """ + output = self._evaluate(batch) + if store: # store the target in the batch if specified + batch[self] = output + return output + + def __call__(self, batch, store=True): + """ + Evaluate the return on the given batch. + + Args: + batch (Batch): batch containing the transitions. + store (bool): If True, it will save the evaluation of the target in the given batch. + + Returns: + torch.Tensor: evaluated return. + """ + return self.evaluate(batch, store=store) class TDReturn(Return): @@ -58,7 +92,33 @@ class TDReturn(Return): Return based on the one-step temporal difference TD(0). """ - pass + + __metaclass__ = ABCMeta + + def __init__(self, gamma=1.): + """ + Initialize the base return / estimator. + + Args: + gamma (float): discount factor + """ + super(TDReturn, self).__init__() + 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 TDValueReturn(TDReturn): @@ -75,27 +135,36 @@ class TDValueReturn(TDReturn): Initialize the TD state value return. Args: - value (ValueApproximator): state value function. - target_value (ValueApproximator): target state value function. + value (Value): state value function. + target_value (Value, list of Value, None): target state value function(s). If None, it will be set to be + the same as the given :attr:`value`. Note however that this can lead to unstable behaviors. 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): + # check value + if not isinstance(value, Value): + raise TypeError("Expecting the given value to be an instance of `Value`, instead got: " + "{}".format(type(value))) + self._value = value + + # check target value + if target_value is None: + # Warning: using the same value function for the target can lead to unstable behaviors. + target_value = value + self._target = ValueTarget(values=target_value, gamma=gamma) + + def _evaluate(self, batch): """Evaluate the TD return on the given batch. Args: batch (Batch): batch containing transitions. Returns: - Batch: batch + torch.Tensor: evaluated TD-return. """ - target = batch['rewards'] + self.gamma * (1 - batch['masks']) * self.target_value(batch['states']) - batch[self] = target - self.value(batch['states']) - return batch + target = self._target(batch, store=False) + return target - self._value(batch['states']) class TDQValueReturn(TDReturn): @@ -113,31 +182,37 @@ class TDQValueReturn(TDReturn): Initialize the TD state-action value return. Args: - q_value (QValueApproximator): Q-value function. + q_value (QValue): 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`. + target_qvalue (QValue, list of QValue, None): target Q-value function(s). If None, it will use the given + :attr:`q_value`. Note however that this can lead to unstable behaviors. 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): + # check Q-value + if not isinstance(q_value, QValue): + raise TypeError("Expecting the given q_value to be an instance of `QValue`, instead got: " + "{}".format(type(q_value))) + self._q_value = q_value + + # check target Q-value + if target_qvalue is None: + # Warning: using the same value function for the target can lead to unstable behaviors. + target_qvalue = q_value + self._target = QValueTarget(q_values=target_qvalue, policy=policy, gamma=gamma) + + def _evaluate(self, batch): """Evaluate the TD return on the given batch. Args: batch (Batch): batch containing transitions. Returns: - Batch: batch + torch.Tensor: evaluated TD-return. """ - action = self.policy.predict(batch['next_states']) - target = batch['rewards'] + self.gamma * (1 - batch['masks']) * self.target_qvalue(action) - batch[self] = target - self.q_value(batch['states']) - return batch + target = self._target(batch, store=False) + return target - self._q_value(batch['states'], actions) class TDQLearningReturn(TDReturn): @@ -161,26 +236,32 @@ class TDQLearningReturn(TDReturn): 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`. + q_value (QValue): Q-value function. + target_qvalue (QValue): 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): + # check Q-value + if not isinstance(q_value, QValueOutput): + raise TypeError("Expecting the given q_value to be an instance of `QValueOutput`, instead got: " + "{}".format(type(q_value))) + self._q_value = q_value + + # check target Q-value + if target_qvalue is None: + # Warning: using the same value function for the target can lead to unstable behaviors. + target_qvalue = q_value + self._target = QLearningTarget(q_values=target_qvalue, gamma=gamma) + + def _evaluate(self, batch): """Evaluate the TD return on the given batch. Args: batch (Batch): batch containing transitions. Returns: - Batch: batch + torch.Tensor: evaluated TD-return. """ - 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 + target = self._target(batch, store=False) + return target - self._q_value(batch['states']) diff --git a/pyrobolearn/returns/targets.py b/pyrobolearn/returns/targets.py index 2e173a7..a5c51ca 100644 --- a/pyrobolearn/returns/targets.py +++ b/pyrobolearn/returns/targets.py @@ -17,7 +17,7 @@ from pyrobolearn.storages import Batch from pyrobolearn.values import Value, QValue from pyrobolearn.policies import Policy from pyrobolearn.exploration import Exploration -from pyrobolearn.returns.estimator import BaseReturn + __author__ = "Brian Delhaisse" __copyright__ = "Copyright 2018, PyRoboLearn" @@ -54,31 +54,35 @@ class Target(object): targets = [] self._targets = targets - def _compute(self, batch): - """Compute the target on the given batch, and return the result. + def _evaluate(self, batch): + """Compute/evaluate the target on the given batch, and return the result. Args: batch (Batch): batch containing the transitions / trajectories. Returns: - torch.Tensor: result of evaluating the batch. + torch.Tensor: evaluated targets. """ raise NotImplementedError - def compute(self, batch): + def evaluate(self, batch, store=True): """Compute/evaluate the target on the given batch and insert the result in the given batch. Args: batch (Batch): batch containing the transitions / trajectories. + store (bool): If True, it will save the evaluation of the target in the given batch. Returns: - Batch: updated batch. + torch.Tensor: evaluated targets. """ - batch[self] = self._compute(batch) - return batch + output = self._evaluate(batch) + if store: # store the target in the batch if specified + batch[self] = output + return output - def __call__(self, batch): - return self.compute(batch) + def __call__(self, batch, store=True): + """Evaluate the target on the given batch.""" + return self.evaluate(batch) class GammaTarget(Target): @@ -133,7 +137,7 @@ class VTarget(Target): self._value = value self._flag = flag % 2 - def _compute(self, batch): + def _evaluate(self, batch): r""" Compute :math:`V(s)`. @@ -141,7 +145,7 @@ class VTarget(Target): batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ if self._flag == 0: return self._value(batch['states']) @@ -177,7 +181,7 @@ class QTarget(Target): "{}".format(type(policy))) self._policy = policy - def _compute(self, batch): + def _evaluate(self, batch): r""" Compute :math:`Q(s, a)` or :math:`Q(s', \pi(s'))`. @@ -185,7 +189,7 @@ class QTarget(Target): batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ if self._policy is None: return self._qvalue(batch['states'], batch['actions']) @@ -219,7 +223,7 @@ class PolicyTarget(Target): # TODO self._flag = flag % 2 - def _compute(self, batch): + def _evaluate(self, batch): """ Compute the log-likelihood on the action :math:`a` returned by the policy. @@ -227,7 +231,7 @@ class PolicyTarget(Target): # TODO batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ if self._flag == 0: return self._policy.predict(batch['states']) @@ -257,7 +261,7 @@ class ValueTarget(GammaTarget): raise TypeError('The {}th value is not an instance of `Value`, instead got: {}'.format(i, type(value))) self._values = values - def _compute(self, batch): + def _evaluate(self, batch): r""" Compute the value target :math:`(r + \gamma (1-d) \min_i V_{\phi_i}(s'))` @@ -265,10 +269,10 @@ class ValueTarget(GammaTarget): batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ value = torch.min(torch.cat([value(batch['next_states']) for value in self._values], dim=1), dim=1)[0] - return batch['rewards'] + self.gamma * (1 - batch['masks']) * value + return batch['rewards'] + self.gamma * batch['masks'] * value class QValueTarget(GammaTarget): @@ -302,7 +306,7 @@ class QValueTarget(GammaTarget): raise TypeError("Expecting the policy to be an instance of `Policy`, instead got: {}".format(type(policy))) self._policy = policy - def _compute(self, batch): + def _evaluate(self, batch): r""" Compute the value target :math:`(r + \gamma (1-d) \min_i Q_{\phi_i}(s',a'))`. @@ -310,12 +314,12 @@ class QValueTarget(GammaTarget): batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ - actions = self._policy.predict(batch['next_states']) - value = torch.min(torch.cat([value(batch['next_states'], actions) for value in self._q_values], dim=1), dim=1)[ - 0] - return batch['rewards'] + self.gamma * (1 - batch['masks']) * value + next_states = batch['next_states'] + actions = self._policy.predict(next_states) + value = torch.min(torch.cat([value(next_states, actions) for value in self._q_values], dim=1), dim=1)[0] + return batch['rewards'] + self.gamma * batch['masks'] * value class QLearningTarget(GammaTarget): @@ -341,7 +345,7 @@ class QLearningTarget(GammaTarget): 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): + def _evaluate(self, batch): r""" Compute the value target :math:`(r + \gamma (1-d) \min_i \max_{a'} Q_{\phi_i}(s',a'))`. @@ -349,11 +353,12 @@ class QLearningTarget(GammaTarget): batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ - q_max = [torch.max(value(batch['next_states']), dim=1, keepdim=True)[0] for value in self._q_values] + next_states = batch['next_states'] + q_max = [torch.max(q_value(next_states), dim=1, keepdim=True)[0] for q_value in self._q_values] value = torch.min(torch.cat(q_max, dim=1), dim=1)[0] - return batch['rewards'] + self.gamma * (1 - batch['masks']) * value + return batch['rewards'] + self.gamma * batch['masks'] * value class EntropyValueTarget(Target): @@ -392,7 +397,7 @@ class EntropyValueTarget(Target): self._alpha = float(alpha) - def _compute(self, batch): + def _evaluate(self, batch): r""" Compute the entropy regularized Q-value target :math:`min_i Q_{\phi_i}(s, \tilde{a}) - \alpha \log \pi_\theta(\tilde{a}|s)`, where @@ -402,7 +407,7 @@ class EntropyValueTarget(Target): batch (Batch): batch containing the transitions. Returns: - torch.Tensor: result + torch.Tensor: evaluated targets. """ actions, distribution = self._policy.predict(batch['states']) value = torch.min(torch.cat([value(batch['states'], actions) for value in self._q_values], dim=1), dim=1)[0] diff --git a/pyrobolearn/samplers/sampler.py b/pyrobolearn/samplers/sampler.py index 1790eb4..b3401a4 100644 --- a/pyrobolearn/samplers/sampler.py +++ b/pyrobolearn/samplers/sampler.py @@ -71,7 +71,7 @@ class StorageSampler(Sampler): raise ValueError("Expecting the batch size (={}) to be smaller than the size of the storage (={})" ".".format(batch_size, self.size)) sampler = torch_sampler.BatchSampler(sampler=torch_sampler.SubsetRandomSampler(range(self.size)), - batch_size=batch_size, drop_last=False) + batch_size=batch_size, drop_last=True) self.sampler = sampler ############## diff --git a/pyrobolearn/storages/er.py b/pyrobolearn/storages/er.py index 26431a0..636cd8d 100644 --- a/pyrobolearn/storages/er.py +++ b/pyrobolearn/storages/er.py @@ -212,8 +212,8 @@ class ExperienceReplay(DictStorage): # ExperienceReplayStorage(DictStorage): ########### def create_new_entry(self, key, shapes, dtype=torch.dtype): - """Create a new entry (=tensor) in the experience replay storage dictionary. The tensor will have the dimension - (num_steps, self.num_processes, *shape) for each shape in shapes, and will be initialized to zero. + """Create a new entry (=tensor) in the experience replay storage dictionary. The tensor will have the + dimension (capacity, *shape) for each shape in shapes, and will be initialized to zero. The tensor will also have the same type than the other tensors and will be sent to the correct device. Args: @@ -380,7 +380,7 @@ class ExperienceReplay(DictStorage): # ExperienceReplayStorage(DictStorage): # def __setattr__(self, key, value): # """Set the attribute using the given key and value. That is, instead of `D[key] = value`, you can do - # `D.key = value`. By default, this creates a tensor with shape (num_steps + 1, self.num_processes, 1). + # `D.key = value`. By default, this creates a tensor with shape (num_steps + 1, self.num_trajectories, 1). # # Warnings: avoid to use this. # """ diff --git a/pyrobolearn/storages/storage.py b/pyrobolearn/storages/storage.py index 374a4f9..e7f7edf 100644 --- a/pyrobolearn/storages/storage.py +++ b/pyrobolearn/storages/storage.py @@ -626,7 +626,7 @@ class RolloutStorage(DictStorage): key variable check that it is correctly present inside the storage. Nonetheless, having a dynamic rollout storage has its advantages, as you can for example store multiple value scalars from multiple value function approximators. - Also, in contrast to [1], we do not compute the returns / returns here. This is done by the `Estimator` class + Also, in contrast to [1], we do not compute the returns / estimators here. This is done by the `Estimator` class which takes as input a `RolloutStorage`. In PRL, this storage is notably used by `RLAlgo` (`Explorator`, `Evaluator`, `Updater`), `Loss`, `Estimators`, etc. @@ -635,7 +635,7 @@ class RolloutStorage(DictStorage): [1] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/storage.py """ - def __init__(self, num_steps, state_shapes, action_shapes, num_processes=1): + def __init__(self, num_steps, state_shapes, action_shapes, num_trajectories=1): # , recurrent_hidden_state_size=0): """ Initialize the rollout storage. @@ -644,7 +644,7 @@ class RolloutStorage(DictStorage): num_steps (int): number of steps in one episode state_shapes (list of tuple of int, tuple of int): each tuple represents the shape of an observation/state. action_shapes (list of tuple of int, tuple of int): each tuple represents the shape of an action. - num_processes (int): number of processes + num_trajectories (int): number of trajectories. """ # recurrent_hidden_state_size (int): size of the internal state print("\nStorage: state shape: {}".format(state_shapes)) @@ -652,9 +652,9 @@ class RolloutStorage(DictStorage): super(RolloutStorage, self).__init__() self._step = 0 self._num_steps = int(num_steps) - self._num_processes = int(num_processes) + self._num_trajectories = int(num_trajectories) self._shifts = {} # dictionary that maps the key to the time shift; this is add to the current time step - self.init(self.num_steps, state_shapes, action_shapes, self.num_processes) + self.init(self.num_steps, state_shapes, action_shapes, self.num_trajectories) ############## # Properties # @@ -666,14 +666,14 @@ class RolloutStorage(DictStorage): return self._num_steps @property - def num_processes(self): + def num_trajectories(self): """Return the number of processes used.""" - return self._num_processes + return self._num_trajectories @property def size(self): """Return the size (=number of steps * number of processes) of the rollout storage.""" - return self._num_steps * self._num_processes + return self._num_steps * self._num_trajectories @property def curr_step(self): @@ -691,7 +691,7 @@ class RolloutStorage(DictStorage): def create_new_entry(self, key, shapes, num_steps=None, dtype=torch.dtype): """Create a new entry (=tensor) in the rollout storage dictionary. The tensor will have the dimension - (num_steps, self.num_processes, *shape) for each shape in shapes, and will be initialized to zero. + (num_steps, self.num_trajectories, *shape) for each shape in shapes, and will be initialized to zero. The tensor will also have the same type than the other tensors and will be sent to the correct device. Args: @@ -722,24 +722,24 @@ class RolloutStorage(DictStorage): # if we have a list of shapes if isinstance(shapes, list): if isinstance(dtype, torch.dtype): - self[key] = [torch.zeros(num_steps, self.num_processes, *shape).to(device=self.device, dtype=dtype) + self[key] = [torch.zeros(num_steps, self.num_trajectories, *shape).to(device=self.device, dtype=dtype) for shape in shapes] else: # numpy array - self[key] = [np.zeros((num_steps, self.num_processes,) + shape, dtype=dtype) for shape in shapes] + self[key] = [np.zeros((num_steps, self.num_trajectories,) + shape, dtype=dtype) for shape in shapes] # if the 'shapes' is a tuple elif isinstance(shapes, tuple): if isinstance(dtype, torch.dtype): - self[key] = torch.zeros(num_steps, self.num_processes, *shapes).to(device=self.device, dtype=dtype) + self[key] = torch.zeros(num_steps, self.num_trajectories, *shapes).to(device=self.device, dtype=dtype) else: - self[key] = np.zeros((num_steps, self.num_processes,) + shapes, dtype=dtype) + self[key] = np.zeros((num_steps, self.num_trajectories,) + shapes, dtype=dtype) # if the 'shapes' is an int elif isinstance(shapes, int): if isinstance(dtype, torch.dtype): - self[key] = torch.zeros(num_steps, self.num_processes, shapes).to(device=self.device, dtype=dtype) + self[key] = torch.zeros(num_steps, self.num_trajectories, shapes).to(device=self.device, dtype=dtype) else: - self[key] = np.zeros((num_steps, self.num_processes, shapes), dtype=dtype) + self[key] = np.zeros((num_steps, self.num_trajectories, shapes), dtype=dtype) else: raise TypeError("Expecting the given shapes {} to be a list of tuple of int, a tuple of int, or an int, " @@ -748,7 +748,7 @@ class RolloutStorage(DictStorage): # add shift self._shifts[key] = num_steps - self.num_steps - def init(self, num_steps, state_shapes, action_shapes, num_processes=1): + def init(self, num_steps, state_shapes, action_shapes, num_trajectories=1): """ Initialize the rollout storage by allocating the appropriate tensors for the observations (states), actions, rewards, masks, and returns. @@ -759,13 +759,13 @@ class RolloutStorage(DictStorage): num_steps (int): number of time steps in the finite-horizon RL setting. state_shapes (list of tuple of int, tuple of int): each tuple represents the shape of an observation/state. action_shapes (list of tuple of int, tuple of int): each tuple represents the shape of an action. - num_processes (int): number of process. + num_trajectories (int): number of trajectories. """ # clear itself: remove all items from the DictStorage, and reset all variables self.clear() self._step = 0 self._num_steps = int(num_steps) - self._num_processes = int(num_processes) + self._num_trajectories = int(num_trajectories) # allocate space for observations / states logger.debug('creating space for states with shape: {}'.format(state_shapes)) @@ -783,9 +783,7 @@ class RolloutStorage(DictStorage): logger.debug('creating space for rewards') self.create_new_entry('rewards', shapes=1, num_steps=self.num_steps) - # allocate space for the returns and masks - logger.debug('creating space for returns') - self.create_new_entry('returns', shapes=1, num_steps=self.num_steps + 1) + # allocate space for the masks logger.debug('creating space for masks') self.create_new_entry('masks', shapes=1, num_steps=self.num_steps + 1) @@ -927,11 +925,11 @@ class RolloutStorage(DictStorage): """Return a batch of the Rollout storage in the form of a `DictStorage`. Args: - indices (list of int): indices. Each index must be between 0 and `num_steps * num_processes`. + indices (list of int): indices. Each index must be between 0 and `num_steps * num_trajectories`. Returns: DictStorage / Batch: batch containing a part of the storage. Variables such as `states`, `actions`, - `rewards`, `returns`, and others can be accessed from the object. + `rewards`, `masks`, and others can be accessed from the object. """ # In the next comments, T = number of time steps, P = number of processes, and I = number of indices batch = {} @@ -970,7 +968,7 @@ class RolloutStorage(DictStorage): # def __setattr__(self, key, value): # """Set the attribute using the given key and value. That is, instead of `D[key] = value`, you can do - # `D.key = value`. By default, this creates a tensor with shape (num_steps + 1, self.num_processes, 1). + # `D.key = value`. By default, this creates a tensor with shape (num_steps + 1, self.num_trajectories, 1). # # Warnings: avoid to use this. # """ diff --git a/pyrobolearn/values/value.py b/pyrobolearn/values/value.py index 05a27d0..3ac6d9f 100644 --- a/pyrobolearn/values/value.py +++ b/pyrobolearn/values/value.py @@ -414,17 +414,21 @@ class ParametrizedQValue(QValueApproximator): # ParametrizedValue, QValueApprox """ self.model.set_vectorized_parameters(vector=vector) - def evaluate(self, state=None, to_numpy=False): + def evaluate(self, state=None, action=None, to_numpy=False): """Compute the output of the value function. Args: state (None, State, (list of) np.array, (list of) torch.Tensor): state input data. If None, it will get - the data from the inputs that were given at the initialization. + the data from the states that were given at the initialization. + action (None, Action, (list of) np.array, (list of) torch.Tensor): input actions. If None, it will get + the data from the actions that were given at the initialization. to_numpy (bool): If True, it will convert the data (torch.Tensors) to numpy arrays. """ # if no input is given, take the provided inputs at the beginning if state is None: state = self.state + if action is None: + action = self.action # if the input is an instance of State, get the inner merged data. if isinstance(state, State): @@ -432,12 +436,17 @@ class ParametrizedQValue(QValueApproximator): # ParametrizedValue, QValueApprox if len(state) == 1: state = state[0] - self.value = self.model.predict(state, to_numpy=to_numpy, return_logits=True, set_output_data=False) + if isinstance(action, Action): + action = action.merged_data + if len(action) == 1: + action = action[0] + + self.value = self.model.predict([state, action], to_numpy=to_numpy, return_logits=True, set_output_data=False) return self.value - def __call__(self, state=None, to_numpy=False): + def __call__(self, state=None, action=None, to_numpy=False): """Predict the value.""" - return self.evaluate(state=state, to_numpy=to_numpy) + return self.evaluate(state=state, action=action, to_numpy=to_numpy) # alias @@ -483,6 +492,32 @@ class ParametrizedQValueOutput(ParametrizedQValue): raise TypeError("Expecting the action to be an int, float, Action, torch.Tensor, or np.ndarray.") self._action = action + def evaluate(self, state=None, action=None, to_numpy=False): + """Compute the output of the value function. + + Args: + state (None, State, (list of) np.array, (list of) torch.Tensor): state input data. If None, it will get + the data from the states that were given at the initialization. + action (None): this argument is discarded. + to_numpy (bool): If True, it will convert the data (torch.Tensors) to numpy arrays. + """ + # if no input is given, take the provided inputs at the beginning + if state is None: + state = self.state + + # if the input is an instance of State, get the inner merged data. + if isinstance(state, State): + state = state.merged_data + if len(state) == 1: + state = state[0] + + self.value = self.model.predict(state, to_numpy=to_numpy, return_logits=True, set_output_data=False) + return self.value + + def __call__(self, state=None, action=None, to_numpy=False): + """Predict the value.""" + return self.evaluate(state=state, to_numpy=to_numpy) + # alias QValueOutput = ParametrizedQValueOutput