|
import torch |
|
import numpy as np |
|
import pytest |
|
|
|
from ding.model.common.head import DuelingHead, ReparameterizationHead, MultiHead, StochasticDuelingHead, EnsembleHead |
|
from ding.torch_utils import is_differentiable |
|
|
|
B = 4 |
|
T = 6 |
|
embedding_dim = 64 |
|
action_shape = 12 |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestHead: |
|
|
|
def output_check(self, model, outputs): |
|
if isinstance(outputs, torch.Tensor): |
|
loss = outputs.sum() |
|
elif isinstance(outputs, list): |
|
loss = sum([t.sum() for t in outputs]) |
|
elif isinstance(outputs, dict): |
|
loss = sum([v.sum() for v in outputs.values()]) |
|
is_differentiable(loss, model) |
|
|
|
def test_dueling(self): |
|
inputs = torch.randn(B, embedding_dim) |
|
model = DuelingHead(embedding_dim, action_shape, 3, 3) |
|
outputs = model(inputs)['logit'] |
|
self.output_check(model, outputs) |
|
assert outputs.shape == (B, action_shape) |
|
|
|
@pytest.mark.parametrize('action_shape', [1, 8]) |
|
def test_reparameterization(self, action_shape): |
|
inputs = torch.randn(B, embedding_dim) |
|
for sigma_type in ['fixed', 'independent', 'conditioned']: |
|
if sigma_type == 'fixed': |
|
model = ReparameterizationHead( |
|
embedding_dim, action_shape, sigma_type=sigma_type, fixed_sigma_value=0.5 |
|
) |
|
outputs = model(inputs) |
|
mu, sigma = outputs['mu'], outputs['sigma'] |
|
assert mu.shape == (B, action_shape) and sigma.shape == (B, action_shape) |
|
assert sigma.eq(torch.full((B, action_shape), 0.5)).all() |
|
self.output_check(model, outputs) |
|
elif sigma_type == 'independent': |
|
model = ReparameterizationHead(embedding_dim, action_shape, sigma_type=sigma_type) |
|
outputs = model(inputs) |
|
mu, sigma = outputs['mu'], outputs['sigma'] |
|
assert mu.shape == (B, action_shape) and sigma.shape == (B, action_shape) |
|
self.output_check(model, outputs) |
|
assert model.log_sigma_param.grad is not None |
|
elif sigma_type == 'conditioned': |
|
model = ReparameterizationHead(embedding_dim, action_shape, sigma_type=sigma_type) |
|
outputs = model(inputs) |
|
mu, sigma = outputs['mu'], outputs['sigma'] |
|
assert mu.shape == (B, action_shape) and sigma.shape == (B, action_shape) |
|
self.output_check(model, outputs) |
|
|
|
def test_multi_head(self): |
|
output_size_list = [2, 3, 7] |
|
head = MultiHead(DuelingHead, embedding_dim, output_size_list, activation=torch.nn.Tanh()) |
|
print(head) |
|
inputs = torch.randn(B, embedding_dim) |
|
outputs = head(inputs) |
|
assert isinstance(outputs, dict) |
|
self.output_check(head, outputs['logit']) |
|
for i, d in enumerate(output_size_list): |
|
assert outputs['logit'][i].shape == (B, d) |
|
|
|
@pytest.mark.tmp |
|
def test_stochastic_dueling(self): |
|
obs = torch.randn(B, embedding_dim) |
|
behaviour_action = torch.randn(B, action_shape).clamp(-1, 1) |
|
mu = torch.randn(B, action_shape).requires_grad_(True) |
|
sigma = torch.rand(B, action_shape).requires_grad_(True) |
|
model = StochasticDuelingHead(embedding_dim, action_shape, 3, 3) |
|
|
|
assert mu.grad is None and sigma.grad is None |
|
outputs = model(obs, behaviour_action, mu, sigma) |
|
self.output_check(model, outputs['q_value']) |
|
assert isinstance(mu.grad, torch.Tensor) |
|
print(mu.grad) |
|
assert isinstance(sigma.grad, torch.Tensor) |
|
assert outputs['q_value'].shape == (B, 1) |
|
assert outputs['v_value'].shape == (B, 1) |
|
|
|
def test_ensemble(self): |
|
inputs = torch.randn(B, embedding_dim * 3, 1) |
|
model = EnsembleHead(embedding_dim, action_shape, 3, 3, 3) |
|
outputs = model(inputs)['pred'] |
|
self.output_check(model, outputs) |
|
assert outputs.shape == (B, action_shape * 3) |
|
|
