Policy Gradient Loss function for PyTorch

policy_gradient_loss.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch._jit_internal import weak_module, weak_script_method

@weak_module
class PolicyGradientLoss(nn.Module):
    """
    Multiplies an unreduced CrossEntropyLoss by a `q` vector.
    """

    def __init__(self):
        super(PolicyGradientLoss, self).__init__()
        self.cross_entropy_loss = nn.CrossEntropyLoss(reduction='none')

    @weak_script_method
    def forward(self, input_, target, q):
        cel = self.cross_entropy_loss.forward(input_, target)
        return torch.mean(cel * q)

Note: I use input_ to appease my editor, since input is a Python built-in. See https://stackoverflow.com/a/20670757/109618