|
import pytest |
|
from itertools import product |
|
import torch |
|
from ding.model.template import BCQ |
|
from ding.torch_utils import is_differentiable |
|
|
|
B = 4 |
|
obs_shape = [4, (8, )] |
|
act_shape = [3, (6, )] |
|
args = list(product(*[obs_shape, act_shape])) |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestBCQ: |
|
|
|
def output_check(self, model, outputs): |
|
if isinstance(outputs, torch.Tensor): |
|
loss = outputs.sum() |
|
elif isinstance(outputs, dict): |
|
loss = sum([v.sum() for v in outputs.values()]) |
|
is_differentiable(loss, model) |
|
|
|
@pytest.mark.parametrize('obs_shape, act_shape', args) |
|
def test_BCQ(self, obs_shape, act_shape): |
|
if isinstance(obs_shape, int): |
|
inputs_obs = torch.randn(B, obs_shape) |
|
else: |
|
inputs_obs = torch.randn(B, *obs_shape) |
|
if isinstance(act_shape, int): |
|
inputs_act = torch.randn(B, act_shape) |
|
else: |
|
inputs_act = torch.randn(B, *act_shape) |
|
inputs = {'obs': inputs_obs, 'action': inputs_act} |
|
model = BCQ(obs_shape, act_shape) |
|
|
|
outputs_c = model(inputs, mode='compute_critic') |
|
assert isinstance(outputs_c, dict) |
|
if isinstance(act_shape, int): |
|
assert torch.stack(outputs_c['q_value']).shape == (2, B) |
|
else: |
|
assert torch.stack(outputs_c['q_value']).shape == (2, B) |
|
self.output_check(model.critic, torch.stack(outputs_c['q_value'])) |
|
|
|
outputs_a = model(inputs, mode='compute_actor') |
|
assert isinstance(outputs_a, dict) |
|
if isinstance(act_shape, int): |
|
assert outputs_a['action'].shape == (B, act_shape) |
|
elif len(act_shape) == 1: |
|
assert outputs_a['action'].shape == (B, *act_shape) |
|
self.output_check(model.actor, outputs_a) |
|
|
|
outputs_vae = model(inputs, mode='compute_vae') |
|
assert isinstance(outputs_vae, dict) |
|
if isinstance(act_shape, int): |
|
assert outputs_vae['recons_action'].shape == (B, act_shape) |
|
assert outputs_vae['mu'].shape == (B, act_shape * 2) |
|
assert outputs_vae['log_var'].shape == (B, act_shape * 2) |
|
assert outputs_vae['z'].shape == (B, act_shape * 2) |
|
elif len(act_shape) == 1: |
|
assert outputs_vae['recons_action'].shape == (B, *act_shape) |
|
assert outputs_vae['mu'].shape == (B, act_shape[0] * 2) |
|
assert outputs_vae['log_var'].shape == (B, act_shape[0] * 2) |
|
assert outputs_vae['z'].shape == (B, act_shape[0] * 2) |
|
if isinstance(obs_shape, int): |
|
assert outputs_vae['prediction_residual'].shape == (B, obs_shape) |
|
else: |
|
assert outputs_vae['prediction_residual'].shape == (B, *obs_shape) |
|
|
|
outputs_eval = model(inputs, mode='compute_eval') |
|
assert isinstance(outputs_eval, dict) |
|
assert isinstance(outputs_eval, dict) |
|
if isinstance(act_shape, int): |
|
assert outputs_eval['action'].shape == (B, act_shape) |
|
elif len(act_shape) == 1: |
|
assert outputs_eval['action'].shape == (B, *act_shape) |
|
|
