refactor returns and losses

This commit is contained in:
Brian Delhaisse
2019-04-18 03:34:46 +02:00
parent 4cc9454f15
commit f5abc15008
13 changed files with 732 additions and 523 deletions
+3
View File
@@ -10,3 +10,6 @@ from .policy_losses import *
# import value losses
from .value_losses import *
# import dynamic losses
from .dynamic_losses import *
+3 -4
View File
@@ -6,9 +6,8 @@ That is, the losses that are used with dynamic models.
import torch
from pyrobolearn.losses.loss import Loss
from pyrobolearn.losses import BatchLoss
from pyrobolearn.dynamics import DynamicModel
from pyrobolearn.storages import Batch
__author__ = "Brian Delhaisse"
@@ -21,7 +20,7 @@ __email__ = "briandelhaisse@gmail.com"
__status__ = "Development"
class DynamicL2Loss(Loss):
class DynamicL2Loss(BatchLoss):
r"""Dynamic L2 loss.
This loss computes the Frobenius norm between the prediction of a dynamic model and the next states that are
@@ -42,7 +41,7 @@ class DynamicL2Loss(Loss):
"{}".format(type(dynamic_model)))
self._dynamic_model = dynamic_model
def compute(self, batch):
def _compute(self, batch):
"""
Compute the frobenius norm (i.e. L2-norm).
+7 -1
View File
@@ -32,6 +32,12 @@ class Loss(object):
__metaclass__ = ABCMeta
def __init__(self, losses=None):
"""
Initialize the loss abstract class.
Args:
losses (None, list of Loss): internal losses to compute.
"""
self.losses = losses
##############
@@ -70,7 +76,7 @@ class Loss(object):
"""Compute the loss and return the scalar value."""
pass
def latex(self):
def latex(self): # TODO: check when using operators with latex formula
"""Return a latex formula of the loss."""
pass
+165 -13
View File
@@ -7,6 +7,8 @@ Losses are evaluated on model parameters, data batches / storages, or transition
import torch
from pyrobolearn.losses.loss import Loss
from pyrobolearn.storages import Batch
__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2018, PyRoboLearn"
@@ -18,39 +20,102 @@ __email__ = "briandelhaisse@gmail.com"
__status__ = "Development"
class FixedLoss(Loss):
class BatchLoss(Loss):
r"""Loss evaluated on a batch.
"""
def _compute(self, batch):
"""Compute the loss on the given batch. This method has to be implemented in the child classes."""
raise NotImplementedError
def compute(self, batch):
"""
Compute the loss on the given batch.
Args:
batch (Batch): batch to evaluate the loss on.
Returns:
torch.Tensor: scalar loss value.
"""
# check that we are given a batch
if not isinstance(batch, Batch):
raise TypeError("Expecting the given 'batch' to be an instance of `Batch`, instead got: "
"{}".format(type(batch)))
return self._compute(batch)
class FixedLoss(BatchLoss):
r"""Fixed Loss
This is a dummy loss that returned always the same values given initially.
"""
def __init__(self, value):
"""
Initialize the fixed loss.
Args:
value (torch.Tensor, float, int, np.array, list): fixed initial values that will be returned at each call.
"""
super(FixedLoss, self).__init__()
self.value = torch.tensor(value)
def compute(self, batch):
def _compute(self, batch):
"""
Compute the fixed loss.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
return self.value
class L2Loss(Loss):
class L2Loss(BatchLoss):
r"""L2 Loss
Compute the L2 loss given by: :math:`1/2 * (y_{target} - y_{predict})^2`
"""
def __init__(self, target, predictor):
"""
Initialize the L2 loss.
Args:
target (callable): callable target that accepts a Batch instance as input. If it is not in the given batch,
it will give the batch to it.
predictor (callable): callable predictor that accepts a Batch instance as input. If it is not in the given
batch, it will give the batch to it.
"""
super(L2Loss, self).__init__()
self._target = target
self._predictor = predictor
def compute(self, batch):
def _compute(self, batch):
r"""
Compute the L2 loss: :math:`1/2 * (y_{target} - y_{predict})^2`.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
# get target data
if self._target in batch:
if self._target in batch.current:
target = batch.current[self._target]
elif self._target in batch:
target = batch[self._target]
else:
target = self._target(batch)
# get predicted data
if self._predictor in batch:
if self._predictor in batch.current:
output = batch.current[self._predictor]
elif self._predictor in batch:
output = batch[self._predictor]
else:
output = self._predictor(batch)
@@ -59,7 +124,7 @@ class L2Loss(Loss):
return 0.5 * (target - output).pow(2).mean()
class HuberLoss(Loss):
class HuberLoss(BatchLoss):
r"""Huber Loss
"In statistics, the Huber loss is a loss function used in robust regression, that is less sensitive to outliers
@@ -67,7 +132,7 @@ class HuberLoss(Loss):
This loss is given by [1]:
.. math:: {\mathcal L}(\delta) = \left\{ \begin{array}{lc} 1/2 a^2 & for |a| \leq \delta, \\
.. math:: {\mathcal L}_{\delta}(a) = \left\{ \begin{array}{lc} 1/2 a^2 & for |a| \leq \delta, \\
\delta (|a| - 1/2 \delta), & \mbox{otherwise} \left \end{array}
"This function is quadratic for small values of :math:`a`, and linear for large values, with equal values and
@@ -81,19 +146,81 @@ class HuberLoss(Loss):
[2] "Reinforcement Learning (DQN) Tutorial":
https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html
"""
def __init__(self, loss, delta=1.):
"""
Initialize the Huber loss.
Args:
loss (Loss): initial loss to smooth.
delta (float): coefficient
"""
super(HuberLoss, self).__init__()
self.loss = loss
self.delta = delta
def compute(self, batch):
def _compute(self, batch):
r"""
Compute the Huber loss.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
a = self.loss(batch)
if abs(a) <= self.delta:
return 0.5 * torch.pow(a, 2)
return self.delta * (torch.abs(a) - 0.5 * self.delta)
class KLLoss(Loss):
class PseudoHuberLoss(BatchLoss):
r"""Pseudo-Huber Loss
"The Pseudo-Huber loss function can be used as a smooth approximation of the Huber loss function. It combines the
best properties of L2 squared loss and L1 absolute loss by being strongly convex when close to the target/minimum
and less steep for extreme values. This steepness can be controlled by the :math:`\delta` value. The Pseudo-Huber
loss function ensures that derivatives are continuous for all degrees. It is defined as:
.. math:: {\mathcal L}_{\delta}(a) = \delta^2 \left( \sqrt{1 + (a/\delta)^2} - 1 \right)
As such, this function approximates :math:`a^2/2` for small values of :math:`a`, and approximates a straight line
with slope :math:`\delta` for large values of :math:`a`.
While the above is the most common form, other smooth approximations of the Huber loss function also exist." [1]
References:
[1] Huber Loss (on Wikipedia): https://en.wikipedia.org/wiki/Huber_loss#Pseudo-Huber_loss_function
"""
def __init__(self, loss, delta=1.):
"""
Initialize the Pseudo-Huber loss.
Args:
loss (Loss): initial loss to smooth.
delta (float): steepness coefficient
"""
super(PseudoHuberLoss, self).__init__()
self.loss = loss
self.delta = delta
def _compute(self, batch):
r"""
Compute the pseudo Huber loss.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
a = self.loss(batch)
return self.delta**2 * (torch.sqrt(1 + (a/self.delta)**2) - 1)
class KLLoss(BatchLoss):
r"""KL Penalty Loss
KL Penalty to minimize:
@@ -115,14 +242,21 @@ class KLLoss(Loss):
self.p = p
self.q = q
def compute(self, batch):
def _compute(self, batch):
"""
Compute :math:`KL(p||q)`.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
# TODO use the batch
return torch.distributions.kl.kl_divergence(self.p, self.q)
def latex(self):
"""Return a latex formula of the loss."""
return r"\mathbb{E}[ KL( p || q ) ]"
@@ -136,7 +270,7 @@ class KLLoss(Loss):
# pass
class HLoss(Loss):
class HLoss(BatchLoss):
r"""Entropy Loss
Entropy loss of a distribution:
@@ -147,10 +281,28 @@ class HLoss(Loss):
"""
def __init__(self, distribution):
"""
Initialize the entropy loss.
Args:
distribution (torch.distributions.Distribution): probability distribution.
"""
super(HLoss, self).__init__()
if not isinstance(distribution, torch.distributions.Distribution):
raise TypeError("Expecting the given distribution to be an instance of `torch.distributions.Distribution`, "
"instead got: {}".format(type(distribution)))
self.p = distribution
def compute(self, batch):
def _compute(self, batch):
"""
Compute the entropy loss.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
entropy = self.p.entropy().mean()
return entropy
+69 -20
View File
@@ -4,9 +4,10 @@
import torch
from pyrobolearn.losses.loss import Loss
from pyrobolearn.losses import BatchLoss
from pyrobolearn.returns.estimators import Estimator
__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2018, PyRoboLearn"
__credits__ = ["Brian Delhaisse"]
@@ -17,7 +18,7 @@ __email__ = "briandelhaisse@gmail.com"
__status__ = "Development"
class PGLoss(Loss):
class PGLoss(BatchLoss):
r"""Policy Gradient Loss
Compute the policy gradient loss which is maximized and given by:
@@ -31,7 +32,7 @@ class PGLoss(Loss):
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) \psi_t ]
References:
[1] "Proximal Policy Optimization Algorithms", Schulman et al., 2017
@@ -43,7 +44,7 @@ class PGLoss(Loss):
Initialize the Policy Gradient loss.
Args:
estimator (Estimator): estimator that has been used on the rollout storage / batch.
estimator (Estimator): estimator/return that has been used on the rollout storage / batch.
"""
super(PGLoss, self).__init__()
if not isinstance(estimator, Estimator):
@@ -51,18 +52,31 @@ class PGLoss(Loss):
"{}".format(type(estimator)))
self._estimator = estimator
def compute(self, batch):
def _compute(self, batch):
"""
Compute the PG loss on the given batch.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
# evaluate the action
log_curr_pi = batch.current['action_distributions']
log_curr_pi = log_curr_pi.log_probs(batch.current['actions'])
estimator = batch[self._estimator]
# compute loss and return it
loss = torch.exp(log_curr_pi) * estimator
return -loss.mean()
def latex(self):
"""Return a latex formula of the loss."""
return r"\mathbb{E}[ r_t(\theta) A_t ]"
class CPILoss(Loss):
class CPILoss(BatchLoss):
r"""CPI Loss
Conservative Policy Iteration objective which is maximized and defined in [1]:
@@ -91,7 +105,16 @@ class CPILoss(Loss):
"{}".format(type(estimator)))
self._estimator = estimator
def compute(self, batch): # policy_distribution, old_policy_distribution, estimator):
def _compute(self, batch):
"""
Compute the CPI loss on the given batch.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
# ratio = policy_distribution / old_policy_distribution
log_curr_pi = batch.current['action_distributions']
log_curr_pi = log_curr_pi.log_probs(batch.current['actions'])
@@ -105,10 +128,11 @@ class CPILoss(Loss):
return -loss.mean()
def latex(self):
"""Return a latex formula of the loss."""
return r"\mathbb{E}[ r_t(\theta) A_t ]"
class CLIPLoss(Loss):
class CLIPLoss(BatchLoss):
r"""CLIP Loss
Loss defined in [1] which is maximized and given by:
@@ -138,7 +162,16 @@ class CLIPLoss(Loss):
"{}".format(type(estimator)))
self._estimator = estimator
def compute(self, batch): # , policy_distribution, old_policy_distribution, estimator):
def _compute(self, batch):
"""
Compute the CLIP loss on the given batch.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
log_curr_pi = batch.current['action_distributions']
log_curr_pi = log_curr_pi.log_probs(batch.current['actions'])
log_prev_pi = batch['action_distributions']
@@ -151,10 +184,11 @@ class CLIPLoss(Loss):
return -loss.mean()
def latex(self):
"""Return a latex formula of the loss."""
return r"\mathbb{E}[ \min(r_t(\theta) A_t, clip(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t) ]"
class KLPenaltyLoss(Loss):
class KLPenaltyLoss(BatchLoss):
r"""KL Penalty Loss
KL Penalty to minimize:
@@ -164,32 +198,35 @@ class KLPenaltyLoss(Loss):
where :math:`KL(.||.)` is the KL-divergence between two probability distributions.
"""
def __init__(self): # p, q):
def __init__(self):
"""
Initialize the KL Penalty loss.
Args:
p (torch.distributions.Distribution): 1st distribution
q (torch.distributions.Distribution): 2nd distribution
"""
super(KLPenaltyLoss, self).__init__()
# self.p = p
# self.q = q
def compute(self, batch):
def _compute(self, batch):
"""
Compute :math:`KL(p||q)`.
Compute the KL divergence loss: :math:`KL(p||q)`.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
curr_pi = batch.current['action_distributions']
prev_pi = batch['action_distributions']
return torch.distributions.kl.kl_divergence(prev_pi, curr_pi)
return torch.distributions.kl.kl_divergence(prev_pi, curr_pi).mean()
def latex(self):
"""Return a latex formula of the loss."""
return r"\mathbb{E}[ KL( \pi_{\theta_{old}}(a_t | s_t) || \pi_{\theta}(a_t | s_t) ) ]"
class EntropyLoss(Loss):
class EntropyLoss(BatchLoss):
r"""Entropy Loss
Entropy loss, which is used to ensure sufficient exploration when maximized [1,2,3]:
@@ -205,9 +242,21 @@ class EntropyLoss(Loss):
"""
def __init__(self):
"""
Initialize the entropy loss.
"""
super(EntropyLoss, self).__init__()
def compute(self, batch):
def _compute(self, batch):
"""
Compute the entropy loss.
Args:
batch (Batch): batch containing the states, actions, rewards, etc.
Returns:
torch.Tensor: loss scalar value
"""
distribution = batch.current['action_distributions']
entropy = distribution.entropy().mean()
return entropy
+37 -13
View File
@@ -4,11 +4,12 @@
import torch
from pyrobolearn.losses.loss import Loss
from pyrobolearn.losses import BatchLoss
from pyrobolearn.policies import Policy
from pyrobolearn.values import QValue, Value
from pyrobolearn.returns import TDReturn, Estimator, Return
__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2018, PyRoboLearn"
__credits__ = ["Brian Delhaisse"]
@@ -19,7 +20,7 @@ __email__ = "briandelhaisse@gmail.com"
__status__ = "Development"
class ValueLoss(Loss):
class ValueL2Loss(BatchLoss):
r"""L2 loss for values
"""
@@ -28,10 +29,10 @@ class ValueLoss(Loss):
Initialize the L2 loss between the returns and values.
Args:
returns ():
returns (Estimator, Return): returns.
value (Value): value function approximator.
"""
super(ValueLoss, self).__init__()
super(ValueL2Loss, self).__init__()
# check the given returns or estimators
if not isinstance(returns, (Estimator, Return)):
@@ -45,16 +46,37 @@ class ValueLoss(Loss):
"{}".format(type(value)))
self._value = value
def compute(self, batch):
returns = batch[self._returns]
values = batch.current[self._value]
def _compute(self, batch):
"""
Compute the loss on the given batch.
Args:
batch (Batch): batch that contains the 'states'.
Returns:
torch.tensor: loss scalar value
"""
if self._returns in batch.current:
returns = batch.current[self._returns]
elif self._returns in batch:
returns = batch[self._returns]
else:
returns = self._returns.evaluate(batch, store=False)
if self._value in batch.current:
values = batch.current[self._value]
elif self._value in batch:
values = batch[self._value]
else:
values = self._value(state=batch['states'])
return 0.5 * (returns - values).pow(2).mean()
class QLoss(Loss):
class QLoss(BatchLoss):
r"""QLoss
This computes :math:`\frac{1}{|B|} \sum_{s \in B} Q_{s, \mu_{\theta}(s)}}`, where :math:`\mu_\theta` is the policy.
This computes :math:`\frac{1}{|B|} \sum_{s \in B} Q_{s, \mu_{\theta}(s)}}`, where :math:`\mu_\theta` is the
policy, and this quantity is being maximized.
"""
def __init__(self, q_value, policy):
@@ -79,7 +101,7 @@ class QLoss(Loss):
"{}".format(type(policy)))
self._policy = policy
def compute(self, batch):
def _compute(self, batch):
"""
Compute the loss on the given batch.
@@ -94,7 +116,7 @@ class QLoss(Loss):
return -q_values.mean()
class MSBELoss(Loss):
class MSBELoss(BatchLoss):
r"""Mean-squared Bellman error
The mean-squared Bellman error (MSBE) computes the Bellman error, also known as the one-step temporal difference
@@ -132,7 +154,7 @@ class MSBELoss(Loss):
"{}".format(type(td_return)))
self._td = td_return
def compute(self, batch):
def _compute(self, batch):
"""
Compute the mean-squared TD return.
@@ -142,7 +164,9 @@ class MSBELoss(Loss):
Returns:
torch.Tensor: loss value.
"""
if self._td in batch:
if self._td in batch.current:
returns = batch.current[self._td]
elif self._td in batch:
returns = batch[self._td]
else:
returns = self._td.evaluate(batch, store=False)
+3
View File
@@ -7,3 +7,6 @@ from .targets import *
# import the various returns
from .returns import *
# import evaluators
from .evaluators import *
-445
View File
@@ -1,445 +0,0 @@
#!/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, Batch
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 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:
.. 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
"""
super(Estimator, self).__init__(gamma)
self.storage = storage
##############
# Properties #
##############
@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.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.states
###########
# Methods #
###########
def _evaluate(self):
"""Evaluate the estimator.
To be implemented in the child class.
"""
raise NotImplementedError
def evaluate(self, storage=None):
"""Evaluate the estimator"""
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 (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."""
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):
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."""
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):
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."""
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 ValueEstimator(Estimator):
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, gamma=1.):
"""
Initialize the state value estimator.
Args:
storage (RolloutStorage): rollout storage
gamma (float): discount factor
"""
super(ValueEstimator, self).__init__(storage=storage, gamma=gamma)
def _evaluate(self):
"""Evaluate the estimator / return."""
self.returns[:] = self.values.clone()
class QValueEstimator(Estimator):
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, gamma=1.):
"""
Initialize the station-action value estimator.
Args:
storage (RolloutStorage): rollout storage
gamma (float): discount factor
"""
super(QValueEstimator, self).__init__(storage=storage, gamma=gamma)
def _evaluate(self):
"""Evaluate the estimator / return."""
self.returns[:] = self.action_values.clone()
class AdvantageEstimator(Estimator):
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, 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):
"""Evaluate the estimator / return."""
self.returns[:] = self.action_values - self.values
class TDResidualEstimator(Estimator):
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, 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): # , 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):
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, gamma=0.98, tau=0.99):
"""
Initialize the Generalized Advantage Estimator.
Args:
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.
"""
super(GAE, self).__init__(storage=storage, gamma=gamma)
self.tau = tau
def _evaluate(self):
"""Evaluate the estimator / return."""
# self.values[-1] = next_value
gae = 0
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]
+32 -14
View File
@@ -1,18 +1,16 @@
#!/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).
- Estimators are evaluated on trajectories and directly saved in the rollout storage unit (i.e. RolloutStorage).
- Returns are estimated on batches of transition tuples (s, a, s', r, d) and saved in the `batch.current` attribute.
In both case, they add
In both case, the computed data is saved in the storage / batch unit, which can then be later used by the loss for
instance.
Dependencies:
- `pyrobolearn.storages`
"""
from abc import ABCMeta
import collections
import torch
from pyrobolearn.storages import RolloutStorage
@@ -144,11 +142,12 @@ class Estimator(object):
"""
raise NotImplementedError
def evaluate(self, storage=None):
def evaluate(self, storage=None, store=True):
"""Evaluate the estimator on the given rollout storage.
Args:
storage (RolloutStorage): rollout storage.
store (True): This is always True.
Returns:
torch.Tensor: the computed returns.
@@ -426,15 +425,16 @@ class AdvantageEstimator(Estimator): # TODO: check with masks
&= \mathbb{E}_{s_{t+1}}[ \delta_t^{V^{\pi,\gamma}} ]
"""
def __init__(self, storage, value, q_value, gamma=1.):
def __init__(self, storage, value, q_value, gamma=1., standardize=False):
"""
Initialize the Advantage estimator.
Args:
storage (RolloutStorage): rollout storage
value (Value): value function approximator.
q_value (QValue): Q-value function approximator.
q_value (QValue, Estimator): Q-value function approximator, or estimator.
gamma (float): discount factor
standardize (bool): If True, it will standardize the advantage estimates.
"""
super(AdvantageEstimator, self).__init__(storage=storage, gamma=gamma)
@@ -445,11 +445,14 @@ class AdvantageEstimator(Estimator): # TODO: check with masks
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: "
if not isinstance(q_value, (QValue, Estimator)):
raise TypeError("Expecting the given q_value to be an instance of `QValue` or `Estimator`, instead got: "
"{}".format(type(q_value)))
self._q_value = q_value
# set if we should standardize or not
self._standardize = bool(standardize)
def _evaluate(self):
"""Evaluate the estimator / return.
@@ -463,13 +466,28 @@ class AdvantageEstimator(Estimator): # TODO: check with masks
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.
# get Q-values/estimators 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)
if isinstance(self._q_value, Estimator):
# evaluate the estimator
self._q_value.evaluate()
# get the returns
q_values = self.storage[self._q_value]
else:
q_values = self._q_value(self.states, self.actions)
self.returns[:] = q_values - values
# compute the advantage estimates
returns = q_values - values
# standardize the advantage estimates
if self._standardize:
returns = (returns - returns.mean()) / (returns.std() + 1.e-5)
# set the returns
self.returns[:] = returns
return self.returns
+335
View File
@@ -0,0 +1,335 @@
#!/usr/bin/env python
"""Computes the various targets based on value function on batches of trajectories/transitions (used in RL).
The approximators are evaluated in . targets that are evaluated are placed inside the given batch, which can then be accessed by other classes.
Dependencies:
- `pyrobolearn.storages`
- `pyrobolearn.values`
"""
from pyrobolearn.approximators import Approximator
from pyrobolearn.policies import Policy
from pyrobolearn.values import Value
from pyrobolearn.dynamics import DynamicModel
from pyrobolearn.actorcritics import ActorCritic
from pyrobolearn.exploration import Exploration # TODO change that name to Explorer instead
from pyrobolearn.storages import Batch
from pyrobolearn.returns import Return, Estimator
__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 Evaluator(object):
r"""Evaluator
Evaluator on the batches.
"""
def _evaluate(self, batch):
"""Compute/evaluate the evaluator on the given batch, and return the result.
Args:
batch (Batch): batch containing the transitions / trajectories.
Returns:
torch.Tensor: evaluated targets.
"""
raise NotImplementedError
def evaluate(self, batch, store=True):
"""Compute/evaluate the evaluator 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:
torch.Tensor: evaluated targets.
"""
# check batch type
if not isinstance(batch, Batch):
raise TypeError("Expecting the given 'batch' to be an instance of `Batch`, instead got: "
"{}".format(type(batch)))
output = self._evaluate(batch)
if store: # store the target in the batch if specified
if isinstance(output, list):
# outputs = []
for key, value in output:
batch.current[key] = value
# outputs.append(value)
# output = outputs
else:
batch.current[self] = output
return output
def __call__(self, batch, store=True):
"""Evaluate the evaluator on the given batch."""
return self.evaluate(batch)
class AdvantageEvaluator(Evaluator):
r"""Advantage evaluator
Compute :math:`\hat{A}_t = R_t - V(s_t)`.
"""
def __init__(self, returns, value, standardize=False):
r"""
Initialize the advantage evaluator.
Args:
returns (Return, Estimator): returns.
value (Value): value function.
standardize (bool): if True, it will standardize the advantage.
"""
# check returns
if not isinstance(returns, (Return, Estimator)):
raise TypeError("Expecting the given 'returns' to be an instance of `Return`, `Estimator`, instead got: "
"{}".format(returns))
self._returns = returns
# 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
# set standardize
self._standardize = bool(standardize)
def _evaluate(self, batch):
"""
Evaluate the advantage estimate.
Args:
batch (Batch): batch containing the transitions / trajectories.
Returns:
torch.Tensor: advantage estimates.
"""
# check value
if self._value in batch.current:
values = batch.current[self._value]
else:
values = self._value(batch['states'])
# check returns
if self._returns in batch.current:
returns = batch.current[self._returns]
elif self._returns in batch:
returns = batch[self._returns]
else:
returns = self._returns(batch)
# compute advantage estimates
advantages = returns - values
# standardize the advantage estimates
if self._standardize:
advantages = (advantages - advantages.mean()) / (advantages.std() + 1.e-5)
return advantages
class PolicyEvaluator(Evaluator):
r"""Policy evaluator
Evaluate a policy by computing :math:`\pi_{\theta}(a|s)` and if possible the distribution :math:`\pi(.|s)`. The
policy is evaluated on a batch.
"""
def __init__(self, policy):
"""Initialize the policy evaluator.
policy (Exploration): policy (with exploration) to evaluate.
"""
# check policy
if not isinstance(policy, Exploration):
raise TypeError("Expecting the given policy to be an instance of `Exploration`, instead got: "
"{}".format(type(policy)))
self._policy = policy
def _evaluate(self, batch):
"""Evaluate the policy on the given batch. If None, it will evaluate on the previous batch.
Args:
batch (Batch): batch containing the transitions / trajectories.
Returns:
torch.Tensor: advantage estimates.
"""
# evaluate policy
actions, action_distributions = self._policy.predict(batch['states'])
# return actions and distribution over actions
return [('actions', actions), ('action_distributions', action_distributions)]
class ValueEvaluator(Evaluator):
r"""Value evaluator
Evaluate a value by computing :math:`V_{\phi}(s)`, :math:`Q_{\phi}(s,a)`, and / or :math:`A_{\phi}(s,a)`.
The value is evaluated on a batch.
"""
def __init__(self, value):
"""Initialize the value evaluator.
value (Value): value to evaluate.
"""
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, batch):
"""Evaluate the value on the given batch. If None, it will evaluate on the previous batch.
Args:
batch (Batch): batch containing the transitions / trajectories.
"""
# evaluate value
values = self._value.evaluate(batch['states'], batch['actions'])
# return values
return [(self._value, values)]
# class ActorCriticEvaluator(Evaluator):
# r"""ActorCritic evaluator
#
# Evaluate an action by computing :math:`\pi_{\theta}(a|s)` and if possible the distribution :math:`\pi(.|s)`. It
# also evaluates the value by computing :math:`V_{\phi}(s)`, :math:`Q_{\phi}(s,a)`, and / or :math:`A_{\phi}(s,a)`.
# Both are evaluated on a batch.
# """
#
# def __init__(self, actorcritic):
# """Initialize the actorcritic evaluator.
#
# actorcritic (ActorCritic): actorcritic to evaluate.
# """
# if not isinstance(actorcritic, ActorCritic):
# raise TypeError("Expecting the given actorcritic to be an instance of `ActorCritic`, instead got: "
# "{}".format(type(actorcritic)))
# self._actorcritic = actorcritic
#
# def _evaluate(self, batch):
# """Evaluate the actorcritic on the given batch. If None, it will evaluate on the previous batch.
#
# Args:
# batch (Batch): batch containing the transitions / trajectories.
# """
# # evaluate actorcritic
# actions, action_distributions, values = self._actorcritic.evaluate(batch['states']) # , batch['actions'])
#
# # put them in the batch
# batch.current['actions'] = actions
# batch.current['action_distributions'] = action_distributions
# batch.current['values'] = values
#
# # return batch
# return batch
class DynamicModelEvaluator(Evaluator):
r"""Dynamic model evaluator
Evaluate the next state given the current state and action.
"""
def __init__(self, dynamic_model):
"""Initialize the dynamic_model evaluator.
dynamic_model (DynamicModel): dynamic_model to evaluate.
"""
# set the dynamic model
if not isinstance(dynamic_model, DynamicModel):
raise TypeError("Expecting the given dynamic_model to be an instance of `ActorCritic`, instead got: "
"{}".format(type(dynamic_model)))
self._dynamic_model = dynamic_model
def _evaluate(self, batch):
"""Evaluate the dynamic_model on the given batch. If None, it will evaluate on the previous batch.
Args:
batch (Batch): batch containing the transitions / trajectories.
"""
# evaluate dynamic_model
next_states, state_distributions = self._dynamic_model.predict(states=batch['states'], actions=batch['actions'],
deterministic=False, to_numpy=False,
set_state_data=False)
# return states and distributions over them
return [('next_states', next_states), ('state_distributions', state_distributions)]
class ApproximatorEvaluator(Evaluator):
r"""Approximators evaluator
Approximators evaluator used mostly during the update phase. Evaluate the various approximators on the given batch.
This consists:
- for policies, to compute :math:`\pi_{\theta}(a|s)` and :math:`\pi_{\theta}(.|s)` if possible.
- for value functions, to compute :math:`V_{\phi}(s)`, :math:`Q_{\phi}(s,a)`, and/or :math:`A_{\phi}`(s,a)
- for dynamic models, to compute :math:``
"""
def __init__(self, approximators):
"""
Initialize the evaluator for the approximators.
Args:
approximators ((list of) Approximator): approximators
"""
if not isinstance(approximators, list):
approximators = [approximators]
for approximator in approximators:
if not isinstance(approximator, (Approximator, Policy, Value, ActorCritic, DynamicModel, Exploration)):
raise TypeError("Expecting the approximator to be an instance of `Approximator`, `Policy`, `Value`, "
"`ActorCritic`, `DynamicModel`, or `Exploration`. Instead got: "
"{}".format(type(approximator)))
self._approximators = approximators
def _evaluate(self, batch):
"""Evaluate the various approximators.
Args:
batch (Batch): batch containing the transitions / trajectories.
"""
# sub-evaluation with the current parameter
outputs = []
for approximator in self._approximators:
if isinstance(approximator, (Policy, Exploration)):
actions, action_distributions = approximator.predict(batch['states'])
outputs.extend([('actions', actions), ('action_distributions', action_distributions)])
elif isinstance(approximator, Value):
values = approximator.evaluate(batch['states'])
outputs.extend([('values', values)])
# elif isinstance(approximator, ActorCritic):
# actions, action_distributions, values = approximator.evaluate(batch['states'], batch['actions'])
# batch.current['actions'] = actions
# batch.current['action_distributions'] = action_distributions
# batch.current['values'] = values
elif isinstance(approximator, DynamicModel):
next_states, state_distributions = approximator.predict(batch['states'], batch['actions'],
deterministic=False, to_numpy=False,
set_state_data=False)
outputs.extend([('next_states', next_states), ('state_distributions', state_distributions)])
else:
raise TypeError("Expecting the approximator to be an instance of `Policy`, `Value`, `ActorCritic`, or "
"`DynamicModel`, instead got: {}".format(type(approximator)))
return outputs
+12 -6
View File
@@ -1,7 +1,8 @@
#!/usr/bin/env python
"""Computes the various returns evaluated on batches of transitions (used in RL).
The targets that are evaluated are placed inside the given batch, which can then be accessed by other classes.
The returns that are evaluated are placed inside the given batch, which can then be accessed by other classes.
Note that most of the TD-returns are using the targets defined in `pyrobolearn/returns/targets.py`.
Dependencies:
- `pyrobolearn.storages`
@@ -68,9 +69,14 @@ class Return(object):
Returns:
torch.Tensor: evaluated return.
"""
# check batch type
if not isinstance(batch, Batch):
raise TypeError("Expecting the given 'batch' to be an instance of `Batch`, instead got: "
"{}".format(type(batch)))
output = self._evaluate(batch)
if store: # store the target in the batch if specified
batch[self] = output
batch.current[self] = output
return output
def __call__(self, batch, store=True):
@@ -135,7 +141,7 @@ class TDValueReturn(TDReturn):
Initialize the TD state value return.
Args:
value (Value): state value function.
value (Value, list of Value): state value function(s).
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
@@ -182,7 +188,7 @@ class TDQValueReturn(TDReturn):
Initialize the TD state-action value return.
Args:
q_value (QValue): Q-value function.
q_value (QValue, list of QValue): Q-value function(s).
policy (Policy): policy to compute the action a'.
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.
@@ -212,7 +218,7 @@ class TDQValueReturn(TDReturn):
torch.Tensor: evaluated TD-return.
"""
target = self._target(batch, store=False)
return target - self._q_value(batch['states'], actions)
return target - self._q_value(batch['states'], self._target.actions)
class TDQLearningReturn(TDReturn):
@@ -236,7 +242,7 @@ class TDQLearningReturn(TDReturn):
Initialize the TD Q-Learning value return.
Args:
q_value (QValue): Q-value function.
q_value (QValue, list of QValue): Q-value function(s).
target_qvalue (QValue): target Q-value function. If None, it will use the given :attr:`q_value`.
gamma (float): discount factor
"""
+9 -2
View File
@@ -75,9 +75,14 @@ class Target(object):
Returns:
torch.Tensor: evaluated targets.
"""
# check batch type
if not isinstance(batch, Batch):
raise TypeError("Expecting the given 'batch' to be an instance of `Batch`, instead got: "
"{}".format(type(batch)))
output = self._evaluate(batch)
if store: # store the target in the batch if specified
batch[self] = output
if store: # store the target in the current batch if specified
batch.current[self] = output
return output
def __call__(self, batch, store=True):
@@ -318,6 +323,7 @@ class QValueTarget(GammaTarget):
"""
next_states = batch['next_states']
actions = self._policy.predict(next_states)
self.actions = actions
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
@@ -410,5 +416,6 @@ class EntropyValueTarget(Target):
torch.Tensor: evaluated targets.
"""
actions, distribution = self._policy.predict(batch['states'])
self.actions, self.distribution = actions, distribution
value = torch.min(torch.cat([value(batch['states'], actions) for value in self._q_values], dim=1), dim=1)[0]
return value - self._alpha * distribution.log_prob(actions)
+57 -5
View File
@@ -607,14 +607,61 @@ class DictStorage(dict, PyTorchStorage):
# alias
class Batch(DictStorage):
r"""Batch storage"""
r"""Batch storage
The Batch storage contains states, actions, rewards, and masks. It might be filled later by the various estimators,
targets, and returns that are defined in `pyrobolearn/returns` folder.
The Batch is notably returned by rollout storages (used for on-policy algorithms) and experience replays (used for
off-policy algorithms) when requesting batches.
For the states, actions, rewards, and masks; these can be accessed from the batch using `batch[name]`. For the
estimators, returns, targets, and others, they can be accessed using `batch[object]`.
Finally, the Batch is created by the various storages, filled by the various returns/targets/estimators (see
`pyrobolearn/returns` folder), and given to the various losses (see `pyrobolearn/losses` folder). Thus, there is a
tight coupling between these 3 concepts. As for the original storages from which the batches are created, these
are filled by the exploration phase in RL algorithms (see `pyrobolearn/algos/explorer`).
"""
def __init__(self, kwargs=None, device=None, dtype=None):
"""
Initialize the Batch storage.
Args:
kwargs (dict): initial dictionary containing the various states, acti
device (torch.device, str, None): the device to put the data on (e.g. `torch.device("cuda:0")` or
`torch.device("cpu")`). If string, it can be 'cpu' or 'cuda'. If None, it will keep the original device
to which the tensor is allocated.
dtype (torch.dtype, None): convert the `torch.Tensor` to the specified data type. If None, it will keep
the original dtype
"""
super(Batch, self).__init__(kwargs=kwargs, device=device, dtype=dtype, update=False)
# contains the current values evaluated during the update phase of RL algorithms.
self.current = DictStorage(kwargs={}, device=device, dtype=dtype, update=False)
def get_current(self, key, default=None):
"""Try first to get the key from :attr:`current`, if not present, try to get it from the batch storage.
class RolloutStorage(DictStorage):
Args:
key (object): dictionary key
default (object): default value to return if the key is not found in `Batch.current` and `Batch`.
"""
# check in current
if key in self.current:
return self.current[key]
# check in self
if key in self:
return self[key]
# return default
return default
# def __contains__(self, key):
# """Check if the given key is in batch.current and in batch."""
# return (key in self.current) or (key in self)
class RolloutStorage(DictStorage): # TODO: think about when multiple policies: storage[policy_class]['actions']?
r"""Rollout Storage
Specific storage used in RL which stores transitions at each step `(s_t, a_t, s_{t+1}, r_t)`, and allows to
@@ -689,7 +736,12 @@ class RolloutStorage(DictStorage):
###########
def step(self, rollout_idx=0):
"""Perform one step; increment by one the current step. If it reaches the end, start from 0 again."""
"""Perform one step; increment by one the current step. If it reaches the end, start from 0 again.
Args:
rollout_idx (int, torch.tensor, np.array, list): trajectory/rollout index(ices). This index must be below
`self.num_trajectories`.
"""
# if end of storage, go at the beginning
self._step[rollout_idx] = (self._step[rollout_idx] + 1) % self.num_steps
@@ -792,7 +844,7 @@ class RolloutStorage(DictStorage):
self.create_new_entry('masks', shapes=1, num_steps=self.num_steps + 1)
# allocate space for action distribution
self.create_new_entry('distributions', shapes=[() for _ in action_shapes], num_steps=self.num_steps,
self.create_new_entry('action_distributions', shapes=[() for _ in action_shapes], num_steps=self.num_steps,
dtype=object)
# space for log probabilities on policy, distributions, scalar values from value functions,
@@ -921,7 +973,7 @@ class RolloutStorage(DictStorage):
self['masks'][t + 1][rollout_idx].copy_(self._convert_to_tensor(mask))
# insert distributions
for distribution, storage in zip(distributions, self['distributions']):
for distribution, storage in zip(distributions, self['action_distributions']):
storage[t][rollout_idx] = distribution
# add other elements