File size: 2,559 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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'])
|