Source code for ding.rl_utils.upgo
import torch
import torch.nn.functional as F
from ding.hpc_rl import hpc_wrapper
from .td import generalized_lambda_returns
def tb_cross_entropy(logit, label, mask=None):
"""
Overview:
Compute the cross entropy loss for label and logit, with mask support
Arguments:
- logit (:obj:`torch.Tensor`): the logit tensor, of size [T, B, N] or [T, B, N, N2]
- label (:obj:`torch.Tensor`): the label tensor, of size [T, B] or [T, B, N2]
- mask (:obj:`torch.Tensor` or :obj:`None`): the mask tensor, of size [T, B] or [T, B, N2]
Returns:
- ce (:obj:`torch.Tensor`): the computed cross entropy, of size [T, B]
Examples:
>>> T, B, N, N2 = 4, 8, 5, 7
>>> logit = torch.randn(T, B, N, N2).softmax(-1).requires_grad_(True)
>>> action = logit.argmax(-1).detach()
>>> ce = tb_cross_entropy(logit, action)
"""
assert (len(label.shape) >= 2)
T, B = label.shape[:2]
# Special 2D case
if len(label.shape) > 2:
assert len(label.shape) == 3
s, n = logit.shape[-2:]
logit = logit.reshape(-1, n)
label = label.reshape(-1)
ce = -F.cross_entropy(logit, label, reduction='none')
ce = ce.view(T * B, -1)
if mask is not None:
ce *= mask.reshape(-1, s)
ce = ce.sum(dim=1)
ce = ce.reshape(T, B)
else:
label = label.reshape(-1)
logit = logit.reshape(-1, logit.shape[-1])
ce = -F.cross_entropy(logit, label, reduction='none')
ce = ce.reshape(T, B, -1)
ce = ce.mean(dim=2)
return ce
[docs]def upgo_returns(rewards: torch.Tensor, bootstrap_values: torch.Tensor) -> torch.Tensor:
"""
Overview:
Computing UPGO return targets. Also notice there is no special handling for the terminal state.
Arguments:
- rewards (:obj:`torch.Tensor`): the returns from time step 0 to T-1, \
of size [T_traj, batchsize]
- bootstrap_values (:obj:`torch.Tensor`): estimation of the state value at step 0 to T, \
of size [T_traj+1, batchsize]
Returns:
- ret (:obj:`torch.Tensor`): Computed lambda return value for each state from 0 to T-1, \
of size [T_traj, batchsize]
Examples:
>>> T, B, N, N2 = 4, 8, 5, 7
>>> rewards = torch.randn(T, B)
>>> bootstrap_values = torch.randn(T + 1, B).requires_grad_(True)
>>> returns = upgo_returns(rewards, bootstrap_values)
"""
# UPGO can be viewed as a lambda return! The trace continues for V_t (i.e. lambda = 1.0) if r_tp1 + V_tp2 > V_tp1.
# as the lambdas[-1, :] is ignored in generalized_lambda_returns, we don't care about bootstrap_values_tp2[-1]
lambdas = (rewards + bootstrap_values[1:]) >= bootstrap_values[:-1]
lambdas = torch.cat([lambdas[1:], torch.ones_like(lambdas[-1:])], dim=0)
return generalized_lambda_returns(bootstrap_values, rewards, 1.0, lambdas)
[docs]@hpc_wrapper(
shape_fn=lambda args: args[0].shape,
namedtuple_data=True,
include_args=5,
include_kwargs=['target_output', 'rhos', 'action', 'rewards', 'bootstrap_values']
)
def upgo_loss(
target_output: torch.Tensor,
rhos: torch.Tensor,
action: torch.Tensor,
rewards: torch.Tensor,
bootstrap_values: torch.Tensor,
mask=None
) -> torch.Tensor:
"""
Overview:
Computing UPGO loss given constant gamma and lambda. There is no special handling for terminal state value,
if the last state in trajectory is the terminal, just pass a 0 as bootstrap_terminal_value.
Arguments:
- target_output (:obj:`torch.Tensor`): the output computed by the target policy network, \
of size [T_traj, batchsize, n_output]
- rhos (:obj:`torch.Tensor`): the importance sampling ratio, of size [T_traj, batchsize]
- action (:obj:`torch.Tensor`): the action taken, of size [T_traj, batchsize]
- rewards (:obj:`torch.Tensor`): the returns from time step 0 to T-1, of size [T_traj, batchsize]
- bootstrap_values (:obj:`torch.Tensor`): estimation of the state value at step 0 to T, \
of size [T_traj+1, batchsize]
Returns:
- loss (:obj:`torch.Tensor`): Computed importance sampled UPGO loss, averaged over the samples, of size []
Examples:
>>> T, B, N, N2 = 4, 8, 5, 7
>>> rhos = torch.randn(T, B)
>>> loss = upgo_loss(logit, rhos, action, rewards, bootstrap_values)
"""
# discard the value at T as it should be considered in the next slice
with torch.no_grad():
returns = upgo_returns(rewards, bootstrap_values)
advantages = rhos * (returns - bootstrap_values[:-1])
metric = tb_cross_entropy(target_output, action, mask)
assert (metric.shape == action.shape[:2])
losses = advantages * metric
return -losses.mean()