Source code for ding.rl_utils.vtrace
import torch
import torch.nn.functional as F
from torch.distributions import Categorical, Independent, Normal
from collections import namedtuple
from .isw import compute_importance_weights
from ding.hpc_rl import hpc_wrapper
[docs]def vtrace_nstep_return(clipped_rhos, clipped_cs, reward, bootstrap_values, gamma=0.99, lambda_=0.95):
"""
Overview:
Computation of vtrace return.
Returns:
- vtrace_return (:obj:`torch.FloatTensor`): the vtrace loss item, all of them are differentiable 0-dim tensor
Shapes:
- clipped_rhos (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep, B is batch size
- clipped_cs (:obj:`torch.FloatTensor`): :math:`(T, B)`
- reward (:obj:`torch.FloatTensor`): :math:`(T, B)`
- bootstrap_values (:obj:`torch.FloatTensor`): :math:`(T+1, B)`
- vtrace_return (:obj:`torch.FloatTensor`): :math:`(T, B)`
"""
deltas = clipped_rhos * (reward + gamma * bootstrap_values[1:] - bootstrap_values[:-1])
factor = gamma * lambda_
result = bootstrap_values[:-1].clone()
vtrace_item = 0.
for t in reversed(range(reward.size()[0])):
vtrace_item = deltas[t] + factor * clipped_cs[t] * vtrace_item
result[t] += vtrace_item
return result
[docs]def vtrace_advantage(clipped_pg_rhos, reward, return_, bootstrap_values, gamma):
"""
Overview:
Computation of vtrace advantage.
Returns:
- vtrace_advantage (:obj:`namedtuple`): the vtrace loss item, all of them are the differentiable 0-dim tensor
Shapes:
- clipped_pg_rhos (:obj:`torch.FloatTensor`): :math:`(T, B)`, where T is timestep, B is batch size
- reward (:obj:`torch.FloatTensor`): :math:`(T, B)`
- return (:obj:`torch.FloatTensor`): :math:`(T, B)`
- bootstrap_values (:obj:`torch.FloatTensor`): :math:`(T, B)`
- vtrace_advantage (:obj:`torch.FloatTensor`): :math:`(T, B)`
"""
return clipped_pg_rhos * (reward + gamma * return_ - bootstrap_values)
vtrace_data = namedtuple('vtrace_data', ['target_output', 'behaviour_output', 'action', 'value', 'reward', 'weight'])
vtrace_loss = namedtuple('vtrace_loss', ['policy_loss', 'value_loss', 'entropy_loss'])
def shape_fn_vtrace_discrete_action(args, kwargs):
r"""
Overview:
Return shape of vtrace for hpc
Returns:
shape: [T, B, N]
"""
if len(args) <= 0:
tmp = kwargs['data'].target_output.shape
else:
tmp = args[0].target_output.shape
return tmp
[docs]@hpc_wrapper(
shape_fn=shape_fn_vtrace_discrete_action,
namedtuple_data=True,
include_args=[0, 1, 2, 3, 4, 5],
include_kwargs=['data', 'gamma', 'lambda_', 'rho_clip_ratio', 'c_clip_ratio', 'rho_pg_clip_ratio']
)
def vtrace_error_discrete_action(
data: namedtuple,
gamma: float = 0.99,
lambda_: float = 0.95,
rho_clip_ratio: float = 1.0,
c_clip_ratio: float = 1.0,
rho_pg_clip_ratio: float = 1.0
):
"""
Overview:
Implementation of vtrace(IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner\
Architectures), (arXiv:1802.01561)
Arguments:
- data (:obj:`namedtuple`): input data with fields shown in ``vtrace_data``
- target_output (:obj:`torch.Tensor`): the output taking the action by the current policy network,\
usually this output is network output logit
- behaviour_output (:obj:`torch.Tensor`): the output taking the action by the behaviour policy network,\
usually this output is network output logit, which is used to produce the trajectory(collector)
- action (:obj:`torch.Tensor`): the chosen action(index for the discrete action space) in trajectory,\
i.e.: behaviour_action
- gamma: (:obj:`float`): the future discount factor, defaults to 0.95
- lambda: (:obj:`float`): mix factor between 1-step (lambda_=0) and n-step, defaults to 1.0
- rho_clip_ratio (:obj:`float`): the clipping threshold for importance weights (rho) when calculating\
the baseline targets (vs)
- c_clip_ratio (:obj:`float`): the clipping threshold for importance weights (c) when calculating\
the baseline targets (vs)
- rho_pg_clip_ratio (:obj:`float`): the clipping threshold for importance weights (rho) when calculating\
the policy gradient advantage
Returns:
- trace_loss (:obj:`namedtuple`): the vtrace loss item, all of them are the differentiable 0-dim tensor
Shapes:
- target_output (:obj:`torch.FloatTensor`): :math:`(T, B, N)`, where T is timestep, B is batch size and\
N is action dim
- behaviour_output (:obj:`torch.FloatTensor`): :math:`(T, B, N)`
- action (:obj:`torch.LongTensor`): :math:`(T, B)`
- value (:obj:`torch.FloatTensor`): :math:`(T+1, B)`
- reward (:obj:`torch.LongTensor`): :math:`(T, B)`
- weight (:obj:`torch.LongTensor`): :math:`(T, B)`
Examples:
>>> T, B, N = 4, 8, 16
>>> value = torch.randn(T + 1, B).requires_grad_(True)
>>> reward = torch.rand(T, B)
>>> target_output = torch.randn(T, B, N).requires_grad_(True)
>>> behaviour_output = torch.randn(T, B, N)
>>> action = torch.randint(0, N, size=(T, B))
>>> data = vtrace_data(target_output, behaviour_output, action, value, reward, None)
>>> loss = vtrace_error_discrete_action(data, rho_clip_ratio=1.1)
"""
target_output, behaviour_output, action, value, reward, weight = data
with torch.no_grad():
IS = compute_importance_weights(target_output, behaviour_output, action, 'discrete')
rhos = torch.clamp(IS, max=rho_clip_ratio)
cs = torch.clamp(IS, max=c_clip_ratio)
return_ = vtrace_nstep_return(rhos, cs, reward, value, gamma, lambda_)
pg_rhos = torch.clamp(IS, max=rho_pg_clip_ratio)
return_t_plus_1 = torch.cat([return_[1:], value[-1:]], 0)
adv = vtrace_advantage(pg_rhos, reward, return_t_plus_1, value[:-1], gamma)
if weight is None:
weight = torch.ones_like(reward)
dist_target = Categorical(logits=target_output)
pg_loss = -(dist_target.log_prob(action) * adv * weight).mean()
value_loss = (F.mse_loss(value[:-1], return_, reduction='none') * weight).mean()
entropy_loss = (dist_target.entropy() * weight).mean()
return vtrace_loss(pg_loss, value_loss, entropy_loss)
[docs]def vtrace_error_continuous_action(
data: namedtuple,
gamma: float = 0.99,
lambda_: float = 0.95,
rho_clip_ratio: float = 1.0,
c_clip_ratio: float = 1.0,
rho_pg_clip_ratio: float = 1.0
):
"""
Overview:
Implementation of vtrace(IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner\
Architectures), (arXiv:1802.01561)
Arguments:
- data (:obj:`namedtuple`): input data with fields shown in ``vtrace_data``
- target_output (:obj:`dict{key:torch.Tensor}`): the output taking the action \
by the current policy network, usually this output is network output, \
which represents the distribution by reparameterization trick.
- behaviour_output (:obj:`dict{key:torch.Tensor}`): the output taking the action \
by the behaviour policy network, usually this output is network output logit, \
which represents the distribution by reparameterization trick.
- action (:obj:`torch.Tensor`): the chosen action(index for the discrete action space) in trajectory, \
i.e.: behaviour_action
- gamma: (:obj:`float`): the future discount factor, defaults to 0.95
- lambda: (:obj:`float`): mix factor between 1-step (lambda_=0) and n-step, defaults to 1.0
- rho_clip_ratio (:obj:`float`): the clipping threshold for importance weights (rho) when calculating\
the baseline targets (vs)
- c_clip_ratio (:obj:`float`): the clipping threshold for importance weights (c) when calculating\
the baseline targets (vs)
- rho_pg_clip_ratio (:obj:`float`): the clipping threshold for importance weights (rho) when calculating\
the policy gradient advantage
Returns:
- trace_loss (:obj:`namedtuple`): the vtrace loss item, all of them are the differentiable 0-dim tensor
Shapes:
- target_output (:obj:`dict{key:torch.FloatTensor}`): :math:`(T, B, N)`, \
where T is timestep, B is batch size and \
N is action dim. The keys are usually parameters of reparameterization trick.
- behaviour_output (:obj:`dict{key:torch.FloatTensor}`): :math:`(T, B, N)`
- action (:obj:`torch.LongTensor`): :math:`(T, B)`
- value (:obj:`torch.FloatTensor`): :math:`(T+1, B)`
- reward (:obj:`torch.LongTensor`): :math:`(T, B)`
- weight (:obj:`torch.LongTensor`): :math:`(T, B)`
Examples:
>>> T, B, N = 4, 8, 16
>>> value = torch.randn(T + 1, B).requires_grad_(True)
>>> reward = torch.rand(T, B)
>>> target_output = dict(
>>> 'mu': torch.randn(T, B, N).requires_grad_(True),
>>> 'sigma': torch.exp(torch.randn(T, B, N).requires_grad_(True)),
>>> )
>>> behaviour_output = dict(
>>> 'mu': torch.randn(T, B, N),
>>> 'sigma': torch.exp(torch.randn(T, B, N)),
>>> )
>>> action = torch.randn((T, B, N))
>>> data = vtrace_data(target_output, behaviour_output, action, value, reward, None)
>>> loss = vtrace_error_continuous_action(data, rho_clip_ratio=1.1)
"""
target_output, behaviour_output, action, value, reward, weight = data
with torch.no_grad():
IS = compute_importance_weights(target_output, behaviour_output, action, 'continuous')
rhos = torch.clamp(IS, max=rho_clip_ratio)
cs = torch.clamp(IS, max=c_clip_ratio)
return_ = vtrace_nstep_return(rhos, cs, reward, value, gamma, lambda_)
pg_rhos = torch.clamp(IS, max=rho_pg_clip_ratio)
return_t_plus_1 = torch.cat([return_[1:], value[-1:]], 0)
adv = vtrace_advantage(pg_rhos, reward, return_t_plus_1, value[:-1], gamma)
if weight is None:
weight = torch.ones_like(reward)
dist_target = Independent(Normal(loc=target_output['mu'], scale=target_output['sigma']), 1)
pg_loss = -(dist_target.log_prob(action) * adv * weight).mean()
value_loss = (F.mse_loss(value[:-1], return_, reduction='none') * weight).mean()
entropy_loss = (dist_target.entropy() * weight).mean()
return vtrace_loss(pg_loss, value_loss, entropy_loss)