|
import pytest |
|
import torch |
|
from ding.rl_utils import vtrace_data, vtrace_error_discrete_action, vtrace_error_continuous_action |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_vtrace_discrete_action(): |
|
T, B, N = 4, 8, 16 |
|
value = torch.randn(T + 1, B).requires_grad_(True) |
|
reward = torch.rand(T, B) |
|
target_output = torch.randn(T, B, N).requires_grad_(True) |
|
behaviour_output = torch.randn(T, B, N) |
|
action = torch.randint(0, N, size=(T, B)) |
|
data = vtrace_data(target_output, behaviour_output, action, value, reward, None) |
|
loss = vtrace_error_discrete_action(data, rho_clip_ratio=1.1) |
|
assert all([l.shape == tuple() for l in loss]) |
|
assert target_output.grad is None |
|
assert value.grad is None |
|
loss = sum(loss) |
|
loss.backward() |
|
assert isinstance(target_output, torch.Tensor) |
|
assert isinstance(value, torch.Tensor) |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_vtrace_continuous_action(): |
|
T, B, N = 4, 8, 16 |
|
value = torch.randn(T + 1, B).requires_grad_(True) |
|
reward = torch.rand(T, B) |
|
target_output = {} |
|
target_output['mu'] = torch.randn(T, B, N).requires_grad_(True) |
|
target_output['sigma'] = torch.exp(torch.randn(T, B, N).requires_grad_(True)) |
|
behaviour_output = {} |
|
behaviour_output['mu'] = torch.randn(T, B, N) |
|
behaviour_output['sigma'] = torch.exp(torch.randn(T, B, N)) |
|
action = torch.randn((T, B, N)) |
|
data = vtrace_data(target_output, behaviour_output, action, value, reward, None) |
|
loss = vtrace_error_continuous_action(data, rho_clip_ratio=1.1) |
|
assert all([l.shape == tuple() for l in loss]) |
|
assert target_output['mu'].grad is None |
|
assert target_output['sigma'].grad is None |
|
assert value.grad is None |
|
loss = sum(loss) |
|
loss.backward() |
|
assert isinstance(target_output['mu'], torch.Tensor) |
|
assert isinstance(target_output['sigma'], torch.Tensor) |
|
assert isinstance(value, torch.Tensor) |
|
|
