Shortcuts

Source code for ding.rl_utils.ppg

from typing import Tuple
from collections import namedtuple
import torch
import torch.nn.functional as F

ppg_data = namedtuple('ppg_data', ['logit_new', 'logit_old', 'action', 'value_new', 'value_old', 'return_', 'weight'])
ppg_joint_loss = namedtuple('ppg_joint_loss', ['auxiliary_loss', 'behavioral_cloning_loss'])


[docs]def ppg_joint_error( data: namedtuple, clip_ratio: float = 0.2, use_value_clip: bool = True, ) -> Tuple[namedtuple, namedtuple]: ''' Overview: Get PPG joint loss Arguments: - data (:obj:`namedtuple`): ppg input data with fieids shown in ``ppg_data`` - clip_ratio (:obj:`float`): clip value for ratio - use_value_clip (:obj:`bool`): whether use value clip Returns: - ppg_joint_loss (:obj:`namedtuple`): the ppg loss item, all of them are the differentiable 0-dim tensor Shapes: - logit_new (:obj:`torch.FloatTensor`): :math:`(B, N)`, where B is batch size and N is action dim - logit_old (:obj:`torch.FloatTensor`): :math:`(B, N)` - action (:obj:`torch.LongTensor`): :math:`(B,)` - value_new (:obj:`torch.FloatTensor`): :math:`(B, 1)` - value_old (:obj:`torch.FloatTensor`): :math:`(B, 1)` - return (:obj:`torch.FloatTensor`): :math:`(B, 1)` - weight (:obj:`torch.FloatTensor`): :math:`(B,)` - auxiliary_loss (:obj:`torch.FloatTensor`): :math:`()`, 0-dim tensor - behavioral_cloning_loss (:obj:`torch.FloatTensor`): :math:`()` Examples: >>> action_dim = 4 >>> data = ppg_data( >>> logit_new=torch.randn(3, action_dim), >>> logit_old=torch.randn(3, action_dim), >>> action=torch.randint(0, action_dim, (3,)), >>> value_new=torch.randn(3, 1), >>> value_old=torch.randn(3, 1), >>> return_=torch.randn(3, 1), >>> weight=torch.ones(3), >>> ) >>> loss = ppg_joint_error(data, 0.99, 0.99) ''' logit_new, logit_old, action, value_new, value_old, return_, weight = data if weight is None: weight = torch.ones_like(return_) # auxiliary_loss if use_value_clip: value_clip = value_old + (value_new - value_old).clamp(-clip_ratio, clip_ratio) v1 = (return_ - value_new).pow(2) v2 = (return_ - value_clip).pow(2) auxiliary_loss = 0.5 * (torch.max(v1, v2) * weight).mean() else: auxiliary_loss = 0.5 * ((return_ - value_new).pow(2) * weight).mean() dist_new = torch.distributions.categorical.Categorical(logits=logit_new) dist_old = torch.distributions.categorical.Categorical(logits=logit_old) logp_new = dist_new.log_prob(action) logp_old = dist_old.log_prob(action) # behavioral cloning loss behavioral_cloning_loss = F.kl_div(logp_new, logp_old, reduction='batchmean') return ppg_joint_loss(auxiliary_loss, behavioral_cloning_loss)