|
import torch |
|
import numpy as np |
|
import pytest |
|
from itertools import product |
|
|
|
from ding.model.template import DiscreteMAQAC, ContinuousMAQAC |
|
from ding.torch_utils import is_differentiable |
|
from ding.utils.default_helper import squeeze |
|
|
|
B = 32 |
|
agent_obs_shape = [216, 265] |
|
global_obs_shape = [264, 324] |
|
agent_num = 8 |
|
action_shape = 14 |
|
args = list(product(*[agent_obs_shape, global_obs_shape, [False, True]])) |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('agent_obs_shape, global_obs_shape, twin_critic', args) |
|
class TestDiscreteMAQAC: |
|
|
|
def output_check(self, model, outputs, action_shape): |
|
if isinstance(action_shape, tuple): |
|
loss = sum([t.sum() for t in outputs]) |
|
elif np.isscalar(action_shape): |
|
loss = outputs.sum() |
|
is_differentiable(loss, model) |
|
|
|
def test_maqac(self, agent_obs_shape, global_obs_shape, twin_critic): |
|
data = { |
|
'obs': { |
|
'agent_state': torch.randn(B, agent_num, agent_obs_shape), |
|
'global_state': torch.randn(B, agent_num, global_obs_shape), |
|
'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) |
|
} |
|
} |
|
model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=twin_critic) |
|
|
|
logit = model(data, mode='compute_actor')['logit'] |
|
value = model(data, mode='compute_critic')['q_value'] |
|
|
|
value_sum = sum(t.sum() for t in value) if twin_critic else value.sum() |
|
outputs = value_sum + logit.sum() |
|
self.output_check(model, outputs, action_shape) |
|
|
|
for p in model.parameters(): |
|
p.grad = None |
|
logit = model(data, mode='compute_actor')['logit'] |
|
self.output_check(model.actor, logit, action_shape) |
|
|
|
for p in model.parameters(): |
|
p.grad = None |
|
value = model(data, mode='compute_critic')['q_value'] |
|
if twin_critic: |
|
for v in value: |
|
assert v.shape == (B, agent_num, action_shape) |
|
else: |
|
assert value.shape == (B, agent_num, action_shape) |
|
self.output_check(model.critic, sum(t.sum() for t in value) if twin_critic else value.sum(), action_shape) |
|
|
|
|
|
B = 32 |
|
agent_obs_shape = [216, 265] |
|
global_obs_shape = [264, 324] |
|
agent_num = 8 |
|
action_shape = 14 |
|
action_space = ['regression', 'reparameterization'] |
|
args = list(product(*[agent_obs_shape, global_obs_shape, action_space, [False, True]])) |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('agent_obs_shape, global_obs_shape, action_space, twin_critic', args) |
|
class TestContinuousMAQAC: |
|
|
|
def output_check(self, model, outputs, action_shape): |
|
if isinstance(action_shape, tuple): |
|
loss = sum([t.sum() for t in outputs]) |
|
elif np.isscalar(action_shape): |
|
loss = outputs.sum() |
|
is_differentiable(loss, model) |
|
|
|
def test_continuousmaqac(self, agent_obs_shape, global_obs_shape, action_space, twin_critic): |
|
data = { |
|
'obs': { |
|
'agent_state': torch.randn(B, agent_num, agent_obs_shape), |
|
'global_state': torch.randn(B, agent_num, global_obs_shape), |
|
'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape)) |
|
}, |
|
'action': torch.randn(B, agent_num, squeeze(action_shape)) |
|
} |
|
model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, action_space, twin_critic=twin_critic) |
|
|
|
for p in model.parameters(): |
|
p.grad = None |
|
|
|
if action_space == 'regression': |
|
action = model(data['obs'], mode='compute_actor')['action'] |
|
if squeeze(action_shape) == 1: |
|
assert action.shape == (B, ) |
|
else: |
|
assert action.shape == (B, agent_num, squeeze(action_shape)) |
|
assert action.eq(action.clamp(-1, 1)).all() |
|
self.output_check(model.actor, action, action_shape) |
|
|
|
elif action_space == 'reparameterization': |
|
(mu, sigma) = model(data['obs'], mode='compute_actor')['logit'] |
|
assert mu.shape == (B, agent_num, action_shape) |
|
assert sigma.shape == (B, agent_num, action_shape) |
|
is_differentiable(mu.sum() + sigma.sum(), model.actor) |
|
|
|
for p in model.parameters(): |
|
p.grad = None |
|
value = model(data, mode='compute_critic')['q_value'] |
|
if twin_critic: |
|
for v in value: |
|
assert v.shape == (B, agent_num) |
|
else: |
|
assert value.shape == (B, agent_num) |
|
self.output_check(model.critic, sum(t.sum() for t in value) if twin_critic else value.sum(), action_shape) |
|
|
