forked from ccf-ai-infra/GPUCodeForces
59 lines
2.5 KiB
Python
59 lines
2.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
__global__ void trpo_kernel(const float* log_probs, const float* old_log_probs, const float* advantages, const float* old_probs, const float* new_probs, float* surrogate_out, float* kl_out, int batch_size, int action_dim) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx < batch_size) {
|
|
float ratio = expf(log_probs[idx] - old_log_probs[idx]);
|
|
surrogate_out[idx] = -ratio * advantages[idx];
|
|
|
|
float kl_sum = 0.0f;
|
|
for (int i = 0; i < action_dim; i++) {
|
|
float old_p = old_probs[idx * action_dim + i];
|
|
float new_p = new_probs[idx * action_dim + i];
|
|
if (old_p > 1e-8f && new_p > 1e-8f) {
|
|
kl_sum += old_p * (logf(old_p) - logf(new_p));
|
|
}
|
|
}
|
|
kl_out[idx] = kl_sum;
|
|
}
|
|
}
|
|
|
|
torch::Tensor trpo_cuda(torch::Tensor log_probs, torch::Tensor old_log_probs, torch::Tensor advantages, torch::Tensor old_probs, torch::Tensor new_probs, float max_kl) {
|
|
auto batch_size = log_probs.size(0);
|
|
auto action_dim = old_probs.size(1);
|
|
auto surrogate_out = torch::empty({batch_size}, log_probs.options());
|
|
auto kl_out = torch::empty({batch_size}, log_probs.options());
|
|
const int block_size = 256;
|
|
int num_blocks = (batch_size + block_size - 1) / block_size;
|
|
trpo_kernel<<<num_blocks, block_size>>>(log_probs.data_ptr<float>(), old_log_probs.data_ptr<float>(), advantages.data_ptr<float>(), old_probs.data_ptr<float>(), new_probs.data_ptr<float>(), surrogate_out.data_ptr<float>(), kl_out.data_ptr<float>(), batch_size, action_dim);
|
|
return surrogate_out.mean() + max_kl * kl_out.mean();
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor trpo_cuda(torch::Tensor log_probs, torch::Tensor old_log_probs, torch::Tensor advantages, torch::Tensor old_probs, torch::Tensor new_probs, float max_kl);
|
|
"""
|
|
|
|
trpo_module = load_inline(
|
|
name="trpo_loss",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["trpo_cuda"],
|
|
verbose=True
|
|
)
|
|
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
def __init__(self, max_kl):
|
|
super(ModelNew, self).__init__()
|
|
self.max_kl = max_kl
|
|
self.trpo_module = trpo_module
|
|
|
|
def forward(self, log_probs, old_log_probs, advantages, old_probs, new_probs):
|
|
return self.trpo_module.trpo_cuda(log_probs, old_log_probs, advantages, old_probs, new_probs, self.max_kl) |