Source code for torchrl.objectives.value.functional
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
import math
import warnings
from functools import wraps
from typing import Optional, Tuple, Union
import torch
__all__ = [
"generalized_advantage_estimate",
"vec_generalized_advantage_estimate",
"td0_advantage_estimate",
"td0_return_estimate",
"td1_return_estimate",
"vec_td1_return_estimate",
"td1_advantage_estimate",
"vec_td1_advantage_estimate",
"td_lambda_return_estimate",
"vec_td_lambda_return_estimate",
"td_lambda_advantage_estimate",
"vec_td_lambda_advantage_estimate",
"vtrace_advantage_estimate",
]
from torchrl.objectives.value.utils import (
_custom_conv1d,
_get_num_per_traj,
_inv_pad_sequence,
_make_gammas_tensor,
_split_and_pad_sequence,
)
SHAPE_ERR = (
"All input tensors (value, reward and done states) must share a unique shape."
)
def _transpose_time(fun):
"""Checks the time_dim argument of the function to allow for any dim.
If not -2, makes a transpose of all the multi-dim input tensors to bring
time at -2, and does the opposite transform for the outputs.
"""
ERROR = (
"The tensor shape and the time dimension are not compatible: "
"got {} and time_dim={}."
)
@wraps(fun)
def transposed_fun(*args, **kwargs):
time_dim = kwargs.pop("time_dim", -2)
def transpose_tensor(tensor):
if not isinstance(tensor, torch.Tensor) or tensor.numel() <= 1:
return tensor, False
if time_dim >= 0:
timedim = time_dim - tensor.ndim
else:
timedim = time_dim
if timedim < -tensor.ndim or timedim >= 0:
raise RuntimeError(ERROR.format(tensor.shape, timedim))
if tensor.ndim >= 2:
single_dim = False
tensor = tensor.transpose(timedim, -2)
elif tensor.ndim == 1 and timedim == -1:
single_dim = True
tensor = tensor.unsqueeze(-1)
else:
raise RuntimeError(ERROR.format(tensor.shape, timedim))
return tensor, single_dim
if time_dim != -2:
single_dim = False
if args:
args, single_dim = zip(*(transpose_tensor(arg) for arg in args))
single_dim = any(single_dim)
for k, item in list(kwargs.items()):
item, sd = transpose_tensor(item)
single_dim = single_dim or sd
kwargs[k] = item
# We don't pass time_dim because it isn't supposed to be used thereafter
out = fun(*args, **kwargs)
if isinstance(out, torch.Tensor):
out = transpose_tensor(out)[0]
if single_dim:
out = out.squeeze(-2)
return out
if single_dim:
return tuple(transpose_tensor(_out)[0].squeeze(-2) for _out in out)
return tuple(transpose_tensor(_out)[0] for _out in out)
# We don't pass time_dim because it isn't supposed to be used thereafter
out = fun(*args, **kwargs)
if isinstance(out, tuple):
for _out in out:
if _out.ndim < 2:
raise RuntimeError(ERROR.format(_out.shape, time_dim))
else:
if out.ndim < 2:
raise RuntimeError(ERROR.format(out.shape, time_dim))
return out
return transposed_fun
########################################################################
# GAE
# ---
[docs]@_transpose_time
def generalized_advantage_estimate(
gamma: float,
lmbda: float,
state_value: torch.Tensor,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
*,
time_dim: int = -2,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Generalized advantage estimate of a trajectory.
Refer to "HIGH-DIMENSIONAL CONTINUOUS CONTROL USING GENERALIZED ADVANTAGE ESTIMATION"
https://arxiv.org/pdf/1506.02438.pdf for more context.
Args:
gamma (scalar): exponential mean discount.
lmbda (scalar): trajectory discount.
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
dtype = next_state_value.dtype
device = state_value.device
not_done = (~done).int()
not_terminated = (~terminated).int()
*batch_size, time_steps, lastdim = not_done.shape
advantage = torch.empty(
*batch_size, time_steps, lastdim, device=device, dtype=dtype
)
prev_advantage = 0
g_not_terminated = gamma * not_terminated
delta = reward + (g_not_terminated * next_state_value) - state_value
discount = lmbda * gamma * not_done
for t in reversed(range(time_steps)):
prev_advantage = advantage[..., t, :] = delta[..., t, :] + (
prev_advantage * discount[..., t, :]
)
value_target = advantage + state_value
return advantage, value_target
def _geom_series_like(t, r, thr):
"""Creates a geometric series of the form [1, gammalmbda, gammalmbda**2] with the shape of `t`.
Drops all elements which are smaller than `thr`.
"""
if isinstance(r, torch.Tensor):
r = r.item()
if r == 0.0:
return torch.zeros_like(t)
elif r >= 1.0:
lim = t.numel()
else:
lim = int(math.log(thr) / math.log(r))
rs = torch.full_like(t[:lim], r)
rs[0] = 1.0
rs = rs.cumprod(0)
rs = rs.unsqueeze(-1)
return rs
def _fast_vec_gae(
reward: torch.Tensor,
state_value: torch.Tensor,
next_state_value: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor,
gamma: float,
lmbda: float,
thr: float = 1e-7,
):
"""Fast vectorized Generalized Advantage Estimate when gamma and lmbda are scalars.
In contrast to `vec_generalized_advantage_estimate` this function does not need
to allocate a big tensor of the form [B, T, T].
Args:
reward (torch.Tensor): a [*B, T, F] tensor containing rewards
state_value (torch.Tensor): a [*B, T, F] tensor containing state values (value function)
next_state_value (torch.Tensor): a [*B, T, F] tensor containing next state values (value function)
done (torch.Tensor): a [B, T] boolean tensor containing the done states.
terminated (torch.Tensor): a [B, T] boolean tensor containing the terminated states.
gamma (scalar): the gamma decay (trajectory discount)
lmbda (scalar): the lambda decay (exponential mean discount)
thr (float): threshold for the filter. Below this limit, components will ignored.
Defaults to 1e-7.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x F]``, with ``F`` feature dimensions.
"""
# _gen_num_per_traj and _split_and_pad_sequence need
# time dimension at last position
done = done.transpose(-2, -1)
terminated = terminated.transpose(-2, -1)
reward = reward.transpose(-2, -1)
state_value = state_value.transpose(-2, -1)
next_state_value = next_state_value.transpose(-2, -1)
gammalmbda = gamma * lmbda
not_terminated = (~terminated).int()
td0 = reward + not_terminated * gamma * next_state_value - state_value
num_per_traj = _get_num_per_traj(done)
td0_flat, mask = _split_and_pad_sequence(td0, num_per_traj, return_mask=True)
gammalmbdas = _geom_series_like(td0_flat[0], gammalmbda, thr=thr)
advantage = _custom_conv1d(td0_flat.unsqueeze(1), gammalmbdas)
advantage = advantage.squeeze(1)
advantage = advantage[mask].view_as(reward)
value_target = advantage + state_value
advantage = advantage.transpose(-1, -2)
value_target = value_target.transpose(-1, -2)
return advantage, value_target
[docs]@_transpose_time
def vec_generalized_advantage_estimate(
gamma: Union[float, torch.Tensor],
lmbda: Union[float, torch.Tensor],
state_value: torch.Tensor,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
*,
time_dim: int = -2,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Vectorized Generalized advantage estimate of a trajectory.
Refer to "HIGH-DIMENSIONAL CONTINUOUS CONTROL USING GENERALIZED ADVANTAGE ESTIMATION"
https://arxiv.org/pdf/1506.02438.pdf for more context.
Args:
gamma (scalar): exponential mean discount.
lmbda (scalar): trajectory discount.
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
dtype = state_value.dtype
*batch_size, time_steps, lastdim = terminated.shape
value = gamma * lmbda
if isinstance(value, torch.Tensor) and value.numel() > 1:
# create tensor while ensuring that gradients are passed
not_done = (~done).to(dtype)
gammalmbdas = not_done * value
else:
# when gamma and lmbda are scalars, use fast_vec_gae implementation
return _fast_vec_gae(
reward=reward,
state_value=state_value,
next_state_value=next_state_value,
done=done,
terminated=terminated,
gamma=gamma,
lmbda=lmbda,
)
gammalmbdas = _make_gammas_tensor(gammalmbdas, time_steps, True)
gammalmbdas = gammalmbdas.cumprod(-2)
first_below_thr = gammalmbdas < 1e-7
# if we have multiple gammas, we only want to truncate if _all_ of
# the geometric sequences fall below the threshold
first_below_thr = first_below_thr.flatten(0, 1).all(0).all(-1)
if first_below_thr.any():
first_below_thr = torch.where(first_below_thr)[0][0].item()
gammalmbdas = gammalmbdas[..., :first_below_thr, :]
not_terminated = (~terminated).to(dtype)
td0 = reward + not_terminated * gamma * next_state_value - state_value
if len(batch_size) > 1:
td0 = td0.flatten(0, len(batch_size) - 1)
elif not len(batch_size):
td0 = td0.unsqueeze(0)
td0_r = td0.transpose(-2, -1)
shapes = td0_r.shape[:2]
if lastdim != 1:
# then we flatten again the first dims and reset a singleton in between
td0_r = td0_r.flatten(0, 1).unsqueeze(1)
advantage = _custom_conv1d(td0_r, gammalmbdas)
if lastdim != 1:
advantage = advantage.squeeze(1).unflatten(0, shapes)
if len(batch_size) > 1:
advantage = advantage.unflatten(0, batch_size)
elif not len(batch_size):
advantage = advantage.squeeze(0)
advantage = advantage.transpose(-2, -1)
value_target = advantage + state_value
return advantage, value_target
########################################################################
# TD(0)
# -----
[docs]def td0_advantage_estimate(
gamma: float,
state_value: torch.Tensor,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""TD(0) advantage estimate of a trajectory.
Also known as bootstrapped Temporal Difference or one-step return.
Args:
gamma (scalar): exponential mean discount.
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
returns = td0_return_estimate(gamma, next_state_value, reward, terminated)
advantage = returns - state_value
return advantage
[docs]def td0_return_estimate(
gamma: float,
next_state_value: torch.Tensor,
reward: torch.Tensor,
terminated: torch.Tensor | None = None,
*,
done: torch.Tensor | None = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# noqa: D417
"""TD(0) discounted return estimate of a trajectory.
Also known as bootstrapped Temporal Difference or one-step return.
Args:
gamma (scalar): exponential mean discount.
next_state_value (Tensor): value function result with new_state input.
must be a [Batch x TimeSteps x 1] or [Batch x TimeSteps] tensor
reward (Tensor): reward of taking actions in the environment.
must be a [Batch x TimeSteps x 1] or [Batch x TimeSteps] tensor
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
Keyword Args:
done (Tensor): Deprecated. Use ``terminated`` instead.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if done is not None and terminated is None:
terminated = done
warnings.warn(
"done for td0_return_estimate is deprecated. Pass ``terminated`` instead."
)
if not (next_state_value.shape == reward.shape == terminated.shape):
raise RuntimeError(SHAPE_ERR)
not_terminated = (~terminated).int()
advantage = reward + gamma * not_terminated * next_state_value
return advantage
########################################################################
# TD(1)
# ----------
[docs]@_transpose_time
def td1_return_estimate(
gamma: float,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rolling_gamma: bool = None,
*,
time_dim: int = -2,
) -> torch.Tensor:
r"""TD(1) return estimate.
Args:
gamma (scalar): exponential mean discount.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (next_state_value.shape == reward.shape == done.shape == terminated.shape):
raise RuntimeError(SHAPE_ERR)
not_done = (~done).int()
not_terminated = (~terminated).int()
returns = torch.empty_like(next_state_value)
T = returns.shape[-2]
single_gamma = False
if not (isinstance(gamma, torch.Tensor) and gamma.shape == not_done.shape):
single_gamma = True
gamma = torch.full_like(next_state_value, gamma)
if rolling_gamma is None:
rolling_gamma = True
elif not rolling_gamma and single_gamma:
raise RuntimeError(
"rolling_gamma=False is expected only with time-sensitive gamma values"
)
done_but_not_terminated = (done & ~terminated).int()
if rolling_gamma:
gamma = gamma * not_terminated
g = next_state_value[..., -1, :]
for i in reversed(range(T)):
# if not done (and hence not terminated), get the bootstrapped value
# if done but not terminated, get nex_val
# if terminated, take nothing (gamma = 0)
dnt = done_but_not_terminated[..., i, :]
g = returns[..., i, :] = reward[..., i, :] + gamma[..., i, :] * (
(1 - dnt) * g + dnt * next_state_value[..., i, :]
)
else:
for k in range(T):
g = 0
_gamma = gamma[..., k, :]
nd = not_terminated
_gamma = _gamma.unsqueeze(-2) * nd
for i in reversed(range(k, T)):
dnt = done_but_not_terminated[..., i, :]
g = reward[..., i, :] + _gamma[..., i, :] * (
(1 - dnt) * g + dnt * next_state_value[..., i, :]
)
returns[..., k, :] = g
return returns
[docs]def td1_advantage_estimate(
gamma: float,
state_value: torch.Tensor,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rolling_gamma: bool = None,
time_dim: int = -2,
) -> torch.Tensor:
"""TD(1) advantage estimate.
Args:
gamma (scalar): exponential mean discount.
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
if not state_value.shape == next_state_value.shape:
raise RuntimeError("shape of state_value and next_state_value must match")
returns = td1_return_estimate(
gamma,
next_state_value,
reward,
done,
terminated=terminated,
rolling_gamma=rolling_gamma,
time_dim=time_dim,
)
advantage = returns - state_value
return advantage
[docs]@_transpose_time
def vec_td1_return_estimate(
gamma,
next_state_value,
reward,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rolling_gamma: Optional[bool] = None,
time_dim: int = -2,
):
"""Vectorized TD(1) return estimate.
Args:
gamma (scalar, Tensor): exponential mean discount. If tensor-valued,
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
return vec_td_lambda_return_estimate(
gamma=gamma,
next_state_value=next_state_value,
reward=reward,
done=done,
terminated=terminated,
rolling_gamma=rolling_gamma,
lmbda=1,
time_dim=time_dim,
)
[docs]def vec_td1_advantage_estimate(
gamma,
state_value,
next_state_value,
reward,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rolling_gamma: bool = None,
time_dim: int = -2,
):
"""Vectorized TD(1) advantage estimate.
Args:
gamma (scalar, Tensor): exponential mean discount. If tensor-valued,
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
return (
vec_td1_return_estimate(
gamma,
next_state_value,
reward,
done=done,
terminated=terminated,
rolling_gamma=rolling_gamma,
time_dim=time_dim,
)
- state_value
)
########################################################################
# TD(lambda)
# ----------
[docs]@_transpose_time
def td_lambda_return_estimate(
gamma: float,
lmbda: float,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rolling_gamma: bool = None,
*,
time_dim: int = -2,
) -> torch.Tensor:
r"""TD(:math:`\lambda`) return estimate.
Args:
gamma (scalar): exponential mean discount.
lmbda (scalar): trajectory discount.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (next_state_value.shape == reward.shape == done.shape == terminated.shape):
raise RuntimeError(SHAPE_ERR)
not_terminated = (~terminated).int()
returns = torch.empty_like(next_state_value)
next_state_value = next_state_value * not_terminated
*batch, T, lastdim = returns.shape
# if gamma is not a tensor of the same shape as other inputs, we use rolling_gamma = True
single_gamma = False
if not (isinstance(gamma, torch.Tensor) and gamma.shape == done.shape):
single_gamma = True
gamma = torch.full_like(next_state_value, gamma)
single_lambda = False
if not (isinstance(lmbda, torch.Tensor) and lmbda.shape == done.shape):
single_lambda = True
lmbda = torch.full_like(next_state_value, lmbda)
if rolling_gamma is None:
rolling_gamma = True
elif not rolling_gamma and single_gamma and single_lambda:
raise RuntimeError(
"rolling_gamma=False is expected only with time-sensitive gamma or lambda values"
)
if rolling_gamma:
g = next_state_value[..., -1, :]
for i in reversed(range(T)):
dn = done[..., i, :].int()
nv = next_state_value[..., i, :]
lmd = lmbda[..., i, :]
# if done, the bootstrapped gain is the next value, otherwise it's the
# value we computed during the previous iter
g = g * (1 - dn) + nv * dn
g = returns[..., i, :] = reward[..., i, :] + gamma[..., i, :] * (
(1 - lmd) * nv + lmd * g
)
else:
for k in range(T):
g = next_state_value[..., -1, :]
_gamma = gamma[..., k, :]
_lambda = lmbda[..., k, :]
for i in reversed(range(k, T)):
dn = done[..., i, :].int()
nv = next_state_value[..., i, :]
g = g * (1 - dn) + nv * dn
g = reward[..., i, :] + _gamma * ((1 - _lambda) * nv + _lambda * g)
returns[..., k, :] = g
return returns
[docs]def td_lambda_advantage_estimate(
gamma: float,
lmbda: float,
state_value: torch.Tensor,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rolling_gamma: bool = None,
# not a kwarg because used directly
time_dim: int = -2,
) -> torch.Tensor:
r"""TD(:math:`\lambda`) advantage estimate.
Args:
gamma (scalar): exponential mean discount.
lmbda (scalar): trajectory discount.
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
if not state_value.shape == next_state_value.shape:
raise RuntimeError("shape of state_value and next_state_value must match")
returns = td_lambda_return_estimate(
gamma,
lmbda,
next_state_value,
reward,
done,
terminated=terminated,
rolling_gamma=rolling_gamma,
time_dim=time_dim,
)
advantage = returns - state_value
return advantage
def _fast_td_lambda_return_estimate(
gamma: Union[torch.Tensor, float],
lmbda: float,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor,
thr: float = 1e-7,
):
"""Fast vectorized TD lambda return estimate.
In contrast to the generalized `vec_td_lambda_return_estimate` this function does not need
to allocate a big tensor of the form [B, T, T], but it only works with gamma/lmbda being scalars.
Args:
gamma (scalar): the gamma decay, can be a tensor with a single element (trajectory discount)
lmbda (scalar): the lambda decay (exponential mean discount)
next_state_value (torch.Tensor): a [*B, T, F] tensor containing next state values (value function)
reward (torch.Tensor): a [*B, T, F] tensor containing rewards
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for end of episode.
thr (float): threshold for the filter. Below this limit, components will ignored.
Defaults to 1e-7.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x F]``, with ``F`` feature dimensions.
"""
device = reward.device
done = done.transpose(-2, -1)
terminated = terminated.transpose(-2, -1)
reward = reward.transpose(-2, -1)
next_state_value = next_state_value.transpose(-2, -1)
# the only valid next states are those where the trajectory does not terminate
next_state_value = (~terminated).int() * next_state_value
gamma_tensor = torch.tensor([gamma], device=device)
gammalmbda = gamma_tensor * lmbda
num_per_traj = _get_num_per_traj(done)
done = done.clone()
done[..., -1] = 1
not_done = (~done).int()
t = reward + next_state_value * gamma_tensor * (1 - not_done * lmbda)
t_flat, mask = _split_and_pad_sequence(t, num_per_traj, return_mask=True)
gammalmbdas = _geom_series_like(t_flat[0], gammalmbda, thr=thr)
ret_flat = _custom_conv1d(t_flat.unsqueeze(1), gammalmbdas)
ret = ret_flat.squeeze(1)[mask]
return ret.view_as(reward).transpose(-1, -2)
[docs]@_transpose_time
def vec_td_lambda_return_estimate(
gamma,
lmbda,
next_state_value,
reward,
done,
terminated: torch.Tensor | None = None,
rolling_gamma: Optional[bool] = None,
*,
time_dim: int = -2,
):
r"""Vectorized TD(:math:`\lambda`) return estimate.
Args:
gamma (scalar, Tensor): exponential mean discount. If tensor-valued,
must be a [Batch x TimeSteps x 1] tensor.
lmbda (scalar): trajectory discount.
next_state_value (Tensor): value function result with new_state input.
must be a [Batch x TimeSteps x 1] tensor
reward (Tensor): reward of taking actions in the environment.
must be a [Batch x TimeSteps x 1] or [Batch x TimeSteps] tensor
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (next_state_value.shape == reward.shape == done.shape == terminated.shape):
raise RuntimeError(SHAPE_ERR)
gamma_thr = 1e-7
shape = next_state_value.shape
*batch, T, lastdim = shape
def _is_scalar(tensor):
return not isinstance(tensor, torch.Tensor) or tensor.numel() == 1
# There are two use-cases: if gamma/lmbda are scalars we can use the
# fast implementation, if not we must construct a gamma tensor.
if _is_scalar(gamma) and _is_scalar(lmbda):
return _fast_td_lambda_return_estimate(
gamma=gamma,
lmbda=lmbda,
next_state_value=next_state_value,
reward=reward,
done=done,
terminated=terminated,
thr=gamma_thr,
)
next_state_value = next_state_value.transpose(-2, -1).unsqueeze(-2)
if len(batch):
next_state_value = next_state_value.flatten(0, len(batch))
reward = reward.transpose(-2, -1).unsqueeze(-2)
if len(batch):
reward = reward.flatten(0, len(batch))
"""Vectorized version of td_lambda_advantage_estimate"""
device = reward.device
not_done = (~done).int()
not_terminated = (~terminated).int().transpose(-2, -1).unsqueeze(-2)
if len(batch):
not_terminated = not_terminated.flatten(0, len(batch))
next_state_value = next_state_value * not_terminated
if rolling_gamma is None:
rolling_gamma = True
if not rolling_gamma:
terminated_follows_terminated = terminated[..., 1:, :][
terminated[..., :-1, :]
].all()
if not terminated_follows_terminated:
raise NotImplementedError(
"When using rolling_gamma=False and vectorized TD(lambda) with time-dependent gamma, "
"make sure that conseducitve trajectories are separated as different batch "
"items. Propagating a gamma value across trajectories is not permitted with "
"this method. Check that you need to use rolling_gamma=False, and if so "
"consider using the non-vectorized version of the return computation or splitting "
"your trajectories."
)
if rolling_gamma:
# Make the coefficient table
gammas = _make_gammas_tensor(gamma * not_done, T, rolling_gamma)
gammas_cp = torch.cumprod(gammas, -2)
lambdas = torch.ones(T + 1, 1, device=device)
lambdas[1:] = lmbda
lambdas_cp = torch.cumprod(lambdas, -2)
lambdas = lambdas[1:]
dec = gammas_cp * lambdas_cp
gammas = _make_gammas_tensor(gamma, T, rolling_gamma)
gammas = gammas[..., 1:, :]
if gammas.ndimension() == 4 and gammas.shape[1] > 1:
gammas = gammas[:, :1]
if lambdas.ndimension() == 4 and lambdas.shape[1] > 1:
lambdas = lambdas[:, :1]
not_done = not_done.transpose(-2, -1).unsqueeze(-2)
if len(batch):
not_done = not_done.flatten(0, len(batch))
# lambdas = lambdas * not_done
v3 = (gammas * lambdas).squeeze(-1) * next_state_value * not_done
v3[..., :-1] = 0
out = _custom_conv1d(
reward
+ gammas.squeeze(-1)
* next_state_value
* (1 - lambdas.squeeze(-1) * not_done)
+ v3,
dec,
)
return out.view(*batch, lastdim, T).transpose(-2, -1)
else:
raise NotImplementedError(
"The vectorized version of TD(lambda) with rolling_gamma=False is currently not available. "
"To use this feature, use the non-vectorized version of TD(lambda). You can expect "
"good speed improvements by decorating the function with torch.compile!"
)
[docs]def vec_td_lambda_advantage_estimate(
gamma,
lmbda,
state_value,
next_state_value,
reward,
done,
terminated: torch.Tensor | None = None,
rolling_gamma: bool = None,
# not a kwarg because used directly
time_dim: int = -2,
):
r"""Vectorized TD(:math:`\lambda`) advantage estimate.
Args:
gamma (scalar, Tensor): exponential mean discount. If tensor-valued,
lmbda (scalar): trajectory discount.
state_value (Tensor): value function result with old_state input.
next_state_value (Tensor): value function result with new_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of trajectory.
terminated (Tensor): boolean flag for the end of episode. Defaults to ``done``
if not provided.
rolling_gamma (bool, optional): if ``True``, it is assumed that each gamma
if a gamma tensor is tied to a single event:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1 g2 v3 + g1 g2 g3 v4,
v2 + g2 v3 + g2 g3 v4,
v3 + g3 v4,
v4,
]
if False, it is assumed that each gamma is tied to the upcoming
trajectory:
gamma = [g1, g2, g3, g4]
value = [v1, v2, v3, v4]
return = [
v1 + g1 v2 + g1**2 v3 + g**3 v4,
v2 + g2 v3 + g2**2 v4,
v3 + g3 v4,
v4,
]
Default is True.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if terminated is None:
terminated = done
if not (
next_state_value.shape
== state_value.shape
== reward.shape
== done.shape
== terminated.shape
):
raise RuntimeError(SHAPE_ERR)
return (
vec_td_lambda_return_estimate(
gamma,
lmbda,
next_state_value,
reward,
done=done,
terminated=terminated,
rolling_gamma=rolling_gamma,
time_dim=time_dim,
)
- state_value
)
########################################################################
# V-Trace
# -----
@_transpose_time
def vtrace_advantage_estimate(
gamma: float,
log_pi: torch.Tensor,
log_mu: torch.Tensor,
state_value: torch.Tensor,
next_state_value: torch.Tensor,
reward: torch.Tensor,
done: torch.Tensor,
terminated: torch.Tensor | None = None,
rho_thresh: Union[float, torch.Tensor] = 1.0,
c_thresh: Union[float, torch.Tensor] = 1.0,
# not a kwarg because used directly
time_dim: int = -2,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Computes V-Trace off-policy actor critic targets.
Refer to "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures"
https://arxiv.org/abs/1802.01561 for more context.
Args:
gamma (scalar): exponential mean discount.
log_pi (Tensor): collection actor log probability of taking actions in the environment.
log_mu (Tensor): current actor log probability of taking actions in the environment.
state_value (Tensor): value function result with state input.
next_state_value (Tensor): value function result with next_state input.
reward (Tensor): reward of taking actions in the environment.
done (Tensor): boolean flag for end of episode.
terminated (torch.Tensor): a [B, T] boolean tensor containing the terminated states.
rho_thresh (Union[float, Tensor]): rho clipping parameter for importance weights.
c_thresh (Union[float, Tensor]): c clipping parameter for importance weights.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
All tensors (values, reward and done) must have shape
``[*Batch x TimeSteps x *F]``, with ``*F`` feature dimensions.
"""
if not (next_state_value.shape == state_value.shape == reward.shape == done.shape):
raise RuntimeError(SHAPE_ERR)
device = state_value.device
if not isinstance(rho_thresh, torch.Tensor):
rho_thresh = torch.tensor(rho_thresh, device=device)
if not isinstance(c_thresh, torch.Tensor):
c_thresh = torch.tensor(c_thresh, device=device)
c_thresh = c_thresh.to(device)
rho_thresh = rho_thresh.to(device)
not_done = (~done).int()
not_terminated = not_done if terminated is None else (~terminated).int()
*batch_size, time_steps, lastdim = not_done.shape
done_discounts = gamma * not_done
terminated_discounts = gamma * not_terminated
rho = (log_pi - log_mu).exp()
clipped_rho = rho.clamp_max(rho_thresh)
deltas = clipped_rho * (
reward + terminated_discounts * next_state_value - state_value
)
clipped_c = rho.clamp_max(c_thresh)
vs_minus_v_xs = [torch.zeros_like(next_state_value[..., -1, :])]
for i in reversed(range(time_steps)):
discount_t, c_t, delta_t = (
done_discounts[..., i, :],
clipped_c[..., i, :],
deltas[..., i, :],
)
vs_minus_v_xs.append(delta_t + discount_t * c_t * vs_minus_v_xs[-1])
vs_minus_v_xs = torch.stack(vs_minus_v_xs[1:], dim=time_dim)
vs_minus_v_xs = torch.flip(vs_minus_v_xs, dims=[time_dim])
vs = vs_minus_v_xs + state_value
vs_t_plus_1 = torch.cat(
[vs[..., 1:, :], next_state_value[..., -1:, :]], dim=time_dim
)
advantages = clipped_rho * (
reward + terminated_discounts * vs_t_plus_1 - state_value
)
return advantages, vs
########################################################################
# Reward to go
# ------------
[docs]@_transpose_time
def reward2go(
reward,
done,
gamma,
*,
time_dim: int = -2,
):
"""Compute the discounted cumulative sum of rewards given multiple trajectories and the episode ends.
Args:
reward (torch.Tensor): A tensor containing the rewards
received at each time step over multiple trajectories.
done (Tensor): boolean flag for end of episode. Differs from
truncated, where the episode did not end but was interrupted.
gamma (float, optional): The discount factor to use for computing the
discounted cumulative sum of rewards. Defaults to 1.0.
time_dim (int): dimension where the time is unrolled. Defaults to -2.
Returns:
torch.Tensor: A tensor of shape [B, T] containing the discounted cumulative
sum of rewards (reward-to-go) at each time step.
Examples:
>>> reward = torch.ones(1, 10)
>>> done = torch.zeros(1, 10, dtype=torch.bool)
>>> done[:, [3, 7]] = True
>>> reward2go(reward, done, 0.99, time_dim=-1)
tensor([[3.9404],
[2.9701],
[1.9900],
[1.0000],
[3.9404],
[2.9701],
[1.9900],
[1.0000],
[1.9900],
[1.0000]])
"""
shape = reward.shape
if shape != done.shape:
raise ValueError(
f"reward and done must share the same shape, got {reward.shape} and {done.shape}"
)
# flatten if needed
if reward.ndim > 2:
# we know time dim is at -2, let's put it at -3
rflip = reward.transpose(-2, -3)
rflip_shape = rflip.shape[-2:]
r2go = reward2go(
rflip.flatten(-2, -1), done.transpose(-2, -3).flatten(-2, -1), gamma=gamma
).unflatten(-1, rflip_shape)
return r2go.transpose(-2, -3)
# place time at dim -1
reward = reward.transpose(-2, -1)
done = done.transpose(-2, -1)
num_per_traj = _get_num_per_traj(done)
td0_flat = _split_and_pad_sequence(reward, num_per_traj)
gammas = _geom_series_like(td0_flat[0], gamma, thr=1e-7)
cumsum = _custom_conv1d(td0_flat.unsqueeze(1), gammas)
cumsum = cumsum.squeeze(1)
cumsum = _inv_pad_sequence(cumsum, num_per_traj)
cumsum = cumsum.reshape_as(reward)
cumsum = cumsum.transpose(-2, -1)
if cumsum.shape != shape:
try:
cumsum = cumsum.reshape(shape)
except RuntimeError:
raise RuntimeError(
f"Wrong shape for output reward2go: {cumsum.shape} when {shape} was expected."
)
return cumsum
