import torch


def self_critical_loss(log_probs, lengths, sampling_reward, greedy_reward):
    """
    Self critical sequence training loss for RL

    log_probs: [batch_size, time_steps]
    lengths: [batch_size], containing the lengths of the predicted polygons
    sampling_reward: [batch_size, ]
    greedy_reward: [batch_size, ]
    """
    reward = sampling_reward - greedy_reward

    loss = 0

    for i in torch.arange(reward.size(0), dtype=torch.long, device=reward.device):
        l = -1 * log_probs[:lengths[i]] * reward[i]
        # Expectation per polygon
        l = torch.mean(l)
        loss += l
    
    # mean across batches
    return loss/reward.size(0)