|
import pytest |
|
import torch |
|
from ding.torch_utils import is_differentiable |
|
from ding.model.template.wqmix import MixerStar, WQMix |
|
|
|
args = [True, False] |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_mixer_star(): |
|
agent_num, bs, embedding_dim = 4, 3, 32 |
|
agent_q = torch.randn(bs, agent_num) |
|
state_embedding = torch.randn(bs, embedding_dim) |
|
mixer_star = MixerStar(agent_num, embedding_dim, 64) |
|
total_q = mixer_star(agent_q, state_embedding) |
|
assert total_q.shape == (bs, ) |
|
loss = total_q.mean() |
|
is_differentiable(loss, mixer_star) |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('is_q_star', args) |
|
def test_wqmix(is_q_star): |
|
agent_num, bs, T = 4, 3, 8 |
|
obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 |
|
embedding_dim = 64 |
|
wqmix_model = WQMix(agent_num, obs_dim, global_obs_dim, action_dim, [128, embedding_dim], 'gru') |
|
data = { |
|
'obs': { |
|
'agent_state': torch.randn(T, bs, agent_num, obs_dim), |
|
'global_state': torch.randn(T, bs, 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 = wqmix_model(data, single_step=False, q_star=is_q_star) |
|
assert set(output.keys()) == set(['total_q', 'logit', 'next_state', 'action_mask']) |
|
assert output['total_q'].shape == (T, bs) |
|
assert output['logit'].shape == (T, bs, agent_num, action_dim) |
|
assert len(output['next_state']) == bs and all([len(n) == agent_num for n in output['next_state']]) |
|
print(output['next_state'][0][0]['h'].shape) |
|
loss = output['total_q'].sum() |
|
if is_q_star: |
|
is_differentiable(loss, [wqmix_model._q_network_star, wqmix_model._mixer_star]) |
|
else: |
|
is_differentiable(loss, [wqmix_model._q_network, wqmix_model._mixer]) |
|
data.pop('action') |
|
output = wqmix_model(data, single_step=False, q_star=is_q_star) |
|
|
