|
import pytest |
|
import torch |
|
from ding.torch_utils import is_differentiable |
|
from ding.model.template import MADQN |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_madqn(): |
|
agent_num, bs, T = 4, 3, 8 |
|
obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 |
|
embedding_dim = 64 |
|
madqn_model = MADQN( |
|
agent_num=agent_num, |
|
obs_shape=obs_dim, |
|
action_shape=action_dim, |
|
hidden_size_list=[embedding_dim, embedding_dim], |
|
global_obs_shape=global_obs_dim |
|
) |
|
data = { |
|
'obs': { |
|
'agent_state': torch.randn(T, bs, agent_num, obs_dim), |
|
'global_state': torch.randn(T, bs, agent_num, global_obs_dim), |
|
'action_mask': torch.randint(0, 2, size=(T, bs, agent_num, action_dim)) |
|
}, |
|
'prev_state': [[None for _ in range(agent_num)] for _ in range(bs)], |
|
'action': torch.randint(0, action_dim, size=(T, bs, agent_num)) |
|
} |
|
output = madqn_model(data, cooperation=True, single_step=False) |
|
assert output['total_q'].shape == (T, bs) |
|
assert len(output['next_state']) == bs and all([len(n) == agent_num for n in output['next_state']]) |
|
|
