|
import pytest |
|
import torch |
|
from ding.torch_utils import is_differentiable |
|
from ding.model.template.vae import VanillaVAE |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_vae(): |
|
batch_size = 32 |
|
action_shape = 6 |
|
original_action_shape = 2 |
|
obs_shape = 6 |
|
hidden_size_list = [256, 256] |
|
inputs = { |
|
'action': torch.randn(batch_size, original_action_shape), |
|
'obs': torch.randn(batch_size, obs_shape), |
|
'next_obs': torch.randn(batch_size, obs_shape) |
|
} |
|
|
|
vae_model = VanillaVAE(original_action_shape, obs_shape, action_shape, hidden_size_list) |
|
outputs = vae_model(inputs) |
|
|
|
assert outputs['recons_action'].shape == (batch_size, original_action_shape) |
|
assert outputs['prediction_residual'].shape == (batch_size, obs_shape) |
|
assert isinstance(outputs['input'], dict) |
|
assert outputs['mu'].shape == (batch_size, obs_shape) |
|
assert outputs['log_var'].shape == (batch_size, obs_shape) |
|
assert outputs['z'].shape == (batch_size, action_shape) |
|
|
|
outputs_decode = vae_model.decode_with_obs(outputs['z'], inputs['obs']) |
|
assert outputs_decode['reconstruction_action'].shape == (batch_size, original_action_shape) |
|
assert outputs_decode['predition_residual'].shape == (batch_size, obs_shape) |
|
|
|
outputs['original_action'] = inputs['action'] |
|
outputs['true_residual'] = inputs['next_obs'] - inputs['obs'] |
|
vae_loss = vae_model.loss_function(outputs, kld_weight=0.01, predict_weight=0.01) |
|
is_differentiable(vae_loss['loss'], vae_model) |
|
|
