|
import pytest |
|
from itertools import product |
|
import numpy as np |
|
import torch |
|
from ding.rl_utils import coma_data, coma_error |
|
|
|
random_weight = torch.rand(128, 4, 8) + 1 |
|
weight_args = [None, random_weight] |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('weight, ', weight_args) |
|
def test_coma(weight): |
|
T, B, A, N = 128, 4, 8, 32 |
|
logit = torch.randn( |
|
T, |
|
B, |
|
A, |
|
N, |
|
).requires_grad_(True) |
|
action = torch.randint( |
|
0, N, size=( |
|
T, |
|
B, |
|
A, |
|
) |
|
) |
|
reward = torch.rand(T, B) |
|
q_value = torch.randn(T, B, A, N).requires_grad_(True) |
|
target_q_value = torch.randn(T, B, A, N).requires_grad_(True) |
|
mask = torch.randint(0, 2, (T, B, A)) |
|
data = coma_data(logit, action, q_value, target_q_value, reward, weight) |
|
loss = coma_error(data, 0.99, 0.95) |
|
assert all([l.shape == tuple() for l in loss]) |
|
assert logit.grad is None |
|
assert q_value.grad is None |
|
total_loss = sum(loss) |
|
total_loss.backward() |
|
assert isinstance(logit.grad, torch.Tensor) |
|
assert isinstance(q_value.grad, torch.Tensor) |
|
|
