|
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) |
|
|
