|
import torch |
|
import numpy as np |
|
import pytest |
|
from itertools import product |
|
|
|
from ding.model.template import PG |
|
from ding.torch_utils import is_differentiable |
|
from ding.utils import squeeze |
|
|
|
B = 4 |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestDiscretePG: |
|
|
|
def output_check(self, model, outputs): |
|
if isinstance(outputs, torch.Tensor): |
|
loss = outputs.sum() |
|
elif isinstance(outputs, list): |
|
loss = sum([t.sum() for t in outputs]) |
|
elif isinstance(outputs, dict): |
|
loss = sum([v.sum() for v in outputs.values()]) |
|
is_differentiable(loss, model) |
|
|
|
def test_discrete_pg(self): |
|
obs_shape = (4, 84, 84) |
|
action_shape = 5 |
|
model = PG( |
|
obs_shape, |
|
action_shape, |
|
) |
|
inputs = torch.randn(B, 4, 84, 84) |
|
|
|
outputs = model(inputs) |
|
assert isinstance(outputs, dict) |
|
assert outputs['logit'].shape == (B, action_shape) |
|
assert outputs['dist'].sample().shape == (B, ) |
|
self.output_check(model, outputs['logit']) |
|
|
|
def test_continuous_pg(self): |
|
N = 32 |
|
action_shape = (6, ) |
|
inputs = {'obs': torch.randn(B, N), 'action': torch.randn(B, squeeze(action_shape))} |
|
model = PG( |
|
obs_shape=(N, ), |
|
action_shape=action_shape, |
|
action_space='continuous', |
|
) |
|
|
|
print(model) |
|
outputs = model(inputs['obs']) |
|
assert isinstance(outputs, dict) |
|
dist = outputs['dist'] |
|
action = dist.sample() |
|
assert action.shape == (B, *action_shape) |
|
|
|
logit = outputs['logit'] |
|
mu, sigma = logit['mu'], logit['sigma'] |
|
assert mu.shape == (B, *action_shape) |
|
assert sigma.shape == (B, *action_shape) |
|
is_differentiable(mu.sum() + sigma.sum(), model) |
|
|
