|
from copy import deepcopy |
|
import pytest |
|
import torch |
|
from easydict import EasyDict |
|
from ding.model.wrapper.model_wrappers import BaseModelWrapper, MultinomialSampleWrapper |
|
from ding.policy import PPOSTDIMPolicy |
|
|
|
obs_shape = 4 |
|
action_shape = 2 |
|
|
|
cfg1 = EasyDict(PPOSTDIMPolicy.default_config()) |
|
cfg1.model.obs_shape = obs_shape |
|
cfg1.model.action_shape = action_shape |
|
|
|
cfg2 = deepcopy(cfg1) |
|
cfg2.action_space = "continuous" |
|
|
|
|
|
def get_transition_discrete(size=64): |
|
data = [] |
|
for i in range(size): |
|
sample = {} |
|
sample['obs'] = torch.rand(obs_shape) |
|
sample['next_obs'] = torch.rand(obs_shape) |
|
sample['action'] = torch.tensor([0], dtype=torch.long) |
|
sample['value'] = torch.rand(1) |
|
sample['logit'] = torch.rand(size=(action_shape, )) |
|
sample['done'] = False |
|
sample['reward'] = torch.rand(1) |
|
data.append(sample) |
|
return data |
|
|
|
|
|
@pytest.mark.parametrize('cfg', [cfg1]) |
|
@pytest.mark.unittest |
|
def test_stdim(cfg): |
|
policy = PPOSTDIMPolicy(cfg, enable_field=['collect', 'eval', 'learn']) |
|
assert type(policy._learn_model) == BaseModelWrapper |
|
assert type(policy._collect_model) == MultinomialSampleWrapper |
|
sample = get_transition_discrete(size=64) |
|
state = policy._state_dict_learn() |
|
policy._load_state_dict_learn(state) |
|
sample = get_transition_discrete(size=64) |
|
out = policy._forward_learn(sample) |
|
|
