colossalai

1
from typing import Optional, Union
2

3
import torch
4
import torch.nn.functional as F
5

6

7
def _compute_approx_kl(
8
    log_probs: torch.Tensor, log_probs_base: torch.Tensor, action_mask: Optional[torch.Tensor] = None
9
) -> torch.Tensor:
10
    """
11
    Compute the approximate KL divergence between two distributions.
12
    Schulman blog: http://joschu.net/blog/kl-approx.html
13

14
    Args:
15
        log_probs: Log probabilities of the new distribution.
16
        log_probs_base: Log probabilities of the base distribution.
17
        action_mask: Mask for actions.
18
    """
19

20
    log_ratio = log_probs_base - log_probs
21
    approx_kl = (log_ratio.exp() - 1) - log_ratio
22
    if action_mask is not None:
23
        approx_kl = masked_mean(approx_kl, action_mask, dim=1)
24
        return approx_kl
25
    approx_kl = approx_kl.mean(dim=1)
26
    return approx_kl
27

28

29
def compute_reward(
30
    r: Union[torch.Tensor, float],
31
    kl_coef: float,
32
    log_probs: torch.Tensor,
33
    log_probs_base: torch.Tensor,
34
    action_mask: Optional[torch.Tensor] = None,
35
) -> torch.Tensor:
36
    if kl_coef <= 0.0:
37
        return r
38
    kl = _compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask)
39
    reward = r - kl_coef * kl
40
    return reward
41

42

43
def _log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
44
    log_probs = F.log_softmax(logits, dim=-1)
45
    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
46
    return log_probs_labels.squeeze(-1)
47

48

49
def calc_action_log_probs(logits: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor:
50
    """Calculate action log probs.
51

52
    Args:
53
        output (torch.Tensor): Output tensor of Actor.forward.logits.
54
        sequences (torch.LongTensor): Input sequences.
55
        num_actions (int): Number of actions.
56

57
    Returns:
58
        torch.Tensor: Action log probs.
59
    """
60
    log_probs = _log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])
61
    return log_probs[:, -num_actions:]
62

63

64
def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
65
    tensor = tensor * mask
66
    tensor = tensor.sum(dim=dim)
67
    mask_sum = mask.sum(dim=dim)
68
    mean = tensor / (mask_sum + 1e-8)
69
    return mean
70