|
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 |
|
|
|
|
|
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 |
|
|
|
|
|
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 |
|
|
|
|
|
@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) |
|
|
|
|
|
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) |
|
|
