|
import pytest |
|
from itertools import product |
|
import numpy as np |
|
import torch |
|
from ding.rl_utils import a2c_data, a2c_error, a2c_error_continuous |
|
|
|
random_weight = torch.rand(4) + 1 |
|
weight_args = [None, random_weight] |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('weight, ', weight_args) |
|
def test_a2c(weight): |
|
B, N = 4, 32 |
|
logit = torch.randn(B, N).requires_grad_(True) |
|
action = torch.randint(0, N, size=(B, )) |
|
value = torch.randn(B).requires_grad_(True) |
|
adv = torch.rand(B) |
|
return_ = torch.randn(B) * 2 |
|
data = a2c_data(logit, action, value, adv, return_, weight) |
|
loss = a2c_error(data) |
|
assert all([l.shape == tuple() for l in loss]) |
|
assert logit.grad is None |
|
assert value.grad is None |
|
total_loss = sum(loss) |
|
total_loss.backward() |
|
assert isinstance(logit.grad, torch.Tensor) |
|
assert isinstance(value.grad, torch.Tensor) |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('weight, ', weight_args) |
|
def test_a2c_continuous(weight): |
|
B, N = 4, 32 |
|
logit = { |
|
"mu": torch.randn(B, N).requires_grad_(True), |
|
"sigma": torch.exp(torch.randn(B, N)).requires_grad_(True), |
|
} |
|
action = torch.randn(B, N).requires_grad_(True) |
|
value = torch.randn(B).requires_grad_(True) |
|
adv = torch.rand(B) |
|
return_ = torch.randn(B) * 2 |
|
data = a2c_data(logit, action, value, adv, return_, weight) |
|
loss = a2c_error_continuous(data) |
|
assert all([l.shape == tuple() for l in loss]) |
|
assert logit["mu"].grad is None |
|
assert logit["sigma"].grad is None |
|
assert value.grad is None |
|
total_loss = sum(loss) |
|
total_loss.backward() |
|
assert isinstance(logit["mu"].grad, torch.Tensor) |
|
assert isinstance(logit['sigma'].grad, torch.Tensor) |
|
assert isinstance(value.grad, torch.Tensor) |
|
|
