File size: 2,588 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 |
import pytest
import torch
from easydict import EasyDict
from ding.model.template import PDQN
action_args_shape_values = [1, 5]
@pytest.mark.unittest
class TestPQQN:
@pytest.mark.unittest
@pytest.mark.parametrize('action_type_shape', action_args_shape_values)
def test_dqn(self, action_type_shape):
T, B = 3, 4
obs_shape = (4, )
act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (action_type_shape, )})
if isinstance(obs_shape, int):
cont_inputs = torch.randn(B, obs_shape)
else:
cont_inputs = torch.randn(B, *obs_shape)
model = PDQN(obs_shape, act_shape)
cont_outputs = model.forward(cont_inputs, mode='compute_continuous')
assert isinstance(cont_outputs, dict)
dis_inputs = {'state': cont_inputs, 'action_args': cont_outputs['action_args']}
dis_outputs = model.forward(dis_inputs, mode='compute_discrete')
assert isinstance(dis_outputs, dict)
if isinstance(act_shape['action_type_shape'], int):
assert dis_outputs['logit'].shape == (B, act_shape.action_type_shape)
elif len(act_shape['action_type_shape']) == 1:
assert dis_outputs['logit'].shape == (B, *act_shape.action_type_shape)
else:
for i, s in enumerate(act_shape):
assert dis_outputs['logit'][i].shape == (B, s)
def test_mdqn(self):
T, B = 3, 4
obs_shape = (4, )
act_shape = EasyDict({'action_type_shape': 3, 'action_args_shape': 5})
if isinstance(obs_shape, int):
cont_inputs = torch.randn(B, obs_shape)
else:
cont_inputs = torch.randn(B, *obs_shape)
model = PDQN(
obs_shape, act_shape, multi_pass=True, action_mask=[[1, 1, 0, 0, 0], [0, 0, 1, 1, 1], [0, 0, 0, 0, 0]]
)
cont_outputs = model.forward(cont_inputs, mode='compute_continuous')
assert isinstance(cont_outputs, dict)
dis_inputs = {'state': cont_inputs, 'action_args': cont_outputs['action_args']}
dis_outputs = model.forward(dis_inputs, mode='compute_discrete')
assert isinstance(dis_outputs, dict)
if isinstance(act_shape['action_type_shape'], int):
assert dis_outputs['logit'].shape == (B, act_shape.action_type_shape)
elif len(act_shape['action_type_shape']) == 1:
assert dis_outputs['logit'].shape == (B, *act_shape.action_type_shape)
else:
for i, s in enumerate(act_shape):
assert dis_outputs['logit'][i].shape == (B, s)
|