|
import pytest |
|
from itertools import product |
|
|
|
import torch |
|
|
|
from ding.model.template import ProcedureCloningMCTS, ProcedureCloningBFS |
|
|
|
B = 4 |
|
T = 15 |
|
obs_shape = [(64, 64, 3)] |
|
action_dim = [9] |
|
obs_embeddings = 256 |
|
args = list(product(*[obs_shape, action_dim])) |
|
|
|
|
|
@pytest.mark.unittest |
|
@pytest.mark.parametrize('obs_shape, action_dim', args) |
|
class TestProcedureCloning: |
|
|
|
def test_procedure_cloning_mcts(self, obs_shape, action_dim): |
|
inputs = { |
|
'states': torch.randn(B, *obs_shape), |
|
'goals': torch.randn(B, *obs_shape), |
|
'actions': torch.randn(B, T, action_dim) |
|
} |
|
model = ProcedureCloningMCTS(obs_shape=obs_shape, action_dim=action_dim) |
|
goal_preds, action_preds = model(inputs['states'], inputs['goals'], inputs['actions']) |
|
assert goal_preds.shape == (B, obs_embeddings) |
|
assert action_preds.shape == (B, T + 1, action_dim) |
|
|
|
def test_procedure_cloning_bfs(self, obs_shape, action_dim): |
|
o_shape = (obs_shape[2], obs_shape[0], obs_shape[1]) |
|
model = ProcedureCloningBFS(obs_shape=o_shape, action_shape=action_dim) |
|
|
|
inputs = torch.randn(B, *obs_shape) |
|
map_preds = model(inputs) |
|
assert map_preds['logit'].shape == (B, obs_shape[0], obs_shape[1], action_dim + 1) |
|
|
