|
import pytest |
|
from itertools import product |
|
import torch |
|
from ding.model.template import NGU |
|
from ding.torch_utils import is_differentiable |
|
|
|
B = 4 |
|
H = 4 |
|
obs_shape = [4, (8, ), (4, 64, 64)] |
|
act_shape = [4, (4, )] |
|
args = list(product(*[obs_shape, act_shape])) |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestNGU: |
|
|
|
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) |
|
|
|
@pytest.mark.parametrize('obs_shape, act_shape', args) |
|
def test_ngu(self, obs_shape, act_shape): |
|
if isinstance(obs_shape, int): |
|
inputs_obs = torch.randn(B, H, obs_shape) |
|
else: |
|
inputs_obs = torch.randn(B, H, *obs_shape) |
|
if isinstance(act_shape, int): |
|
inputs_prev_action = torch.ones(B, act_shape).long() |
|
else: |
|
inputs_prev_action = torch.ones(B, *act_shape).long() |
|
inputs_prev_reward_extrinsic = torch.randn(B, H, 1) |
|
inputs_beta = 2 * torch.ones([4, 4], dtype=torch.long) |
|
inputs = { |
|
'obs': inputs_obs, |
|
'prev_state': None, |
|
'prev_action': inputs_prev_action, |
|
'prev_reward_extrinsic': inputs_prev_reward_extrinsic, |
|
'beta': inputs_beta |
|
} |
|
|
|
model = NGU(obs_shape, act_shape, collector_env_num=3) |
|
outputs = model(inputs) |
|
assert isinstance(outputs, dict) |
|
if isinstance(act_shape, int): |
|
assert outputs['logit'].shape == (B, act_shape, act_shape) |
|
elif len(act_shape) == 1: |
|
assert outputs['logit'].shape == (B, *act_shape, *act_shape) |
|
self.output_check(model, outputs['logit']) |
|
|
|
inputs = { |
|
'obs': inputs_obs, |
|
'prev_state': None, |
|
'action': inputs_prev_action, |
|
'reward': inputs_prev_reward_extrinsic, |
|
'prev_reward_extrinsic': inputs_prev_reward_extrinsic, |
|
'beta': inputs_beta |
|
} |
|
model = NGU(obs_shape, act_shape, collector_env_num=3) |
|
outputs = model(inputs) |
|
assert isinstance(outputs, dict) |
|
if isinstance(act_shape, int): |
|
assert outputs['logit'].shape == (B, act_shape, act_shape) |
|
elif len(act_shape) == 1: |
|
assert outputs['logit'].shape == (B, *act_shape, *act_shape) |
|
self.output_check(model, outputs['logit']) |
|
|
