forked from ccf-ai-infra/GPUCodeForces
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, max_kl):
|
|
super(Model, self).__init__()
|
|
self.max_kl = max_kl
|
|
|
|
def forward(self, log_probs: torch.Tensor, old_log_probs: torch.Tensor, advantages: torch.Tensor,
|
|
old_probs: torch.Tensor, new_probs: torch.Tensor) -> torch.Tensor:
|
|
ratio = torch.exp(log_probs - old_log_probs)
|
|
surrogate_loss = -(ratio * advantages).mean()
|
|
kl_div = (old_probs * (torch.log(old_probs) - torch.log(new_probs))).sum(dim=-1).mean()
|
|
loss = surrogate_loss + self.max_kl * kl_div
|
|
return loss
|
|
|
|
|
|
batch_size = 32
|
|
action_dim = 4
|
|
|
|
|
|
def get_inputs():
|
|
log_probs = torch.randn(batch_size)
|
|
old_log_probs = torch.randn(batch_size)
|
|
advantages = torch.randn(batch_size)
|
|
old_probs = torch.softmax(torch.randn(batch_size, action_dim), dim=-1)
|
|
new_probs = torch.softmax(torch.randn(batch_size, action_dim), dim=-1)
|
|
return [log_probs, old_log_probs, advantages, old_probs, new_probs]
|
|
|
|
|
|
def get_init_inputs():
|
|
max_kl = torch.tensor(0.01)
|
|
return [max_kl] |