|
import copy |
|
|
|
import pytest |
|
import torch |
|
from easydict import EasyDict |
|
from ding.policy.cql import CQLPolicy, DiscreteCQLPolicy |
|
from ding.utils.data import offline_data_save_type |
|
from tensorboardX import SummaryWriter |
|
from ding.model.wrapper.model_wrappers import ArgmaxSampleWrapper, EpsGreedySampleWrapper, TargetNetworkWrapper |
|
import os |
|
from typing import List |
|
from collections import namedtuple |
|
from ding.utils import deep_merge_dicts |
|
|
|
obs_space = 5 |
|
action_space = 3 |
|
|
|
cfg1 = EasyDict(CQLPolicy.default_config()) |
|
cfg1.model.obs_shape = obs_space |
|
cfg1.model.action_shape = action_space |
|
|
|
cfg2 = copy.deepcopy(cfg1) |
|
cfg2.learn.auto_alpha = False |
|
cfg2.learn.log_space = False |
|
|
|
cfg3 = EasyDict(DiscreteCQLPolicy.default_config()) |
|
cfg3.model = {} |
|
cfg3.model.obs_shape = obs_space |
|
cfg3.model.action_shape = action_space |
|
|
|
cfg4 = copy.deepcopy(cfg3) |
|
cfg4.learn.auto_alpha = False |
|
|
|
|
|
def get_batch(size=8): |
|
data = {} |
|
for i in range(size): |
|
obs = torch.zeros(obs_space) |
|
data[i] = obs |
|
return data |
|
|
|
|
|
def get_transition(size=20): |
|
data = [] |
|
for i in range(size): |
|
sample = {} |
|
sample['obs'] = torch.zeros(obs_space) |
|
sample['action'] = torch.zeros(action_space) |
|
sample['done'] = False |
|
sample['next_obs'] = torch.zeros(obs_space) |
|
sample['reward'] = torch.Tensor([1.]) |
|
data.append(sample) |
|
return data |
|
|
|
|
|
def get_transition_batch(bs=1): |
|
sample = {} |
|
sample['obs'] = torch.zeros(bs, obs_space) |
|
sample['action'] = torch.zeros(bs, action_space) |
|
return sample |
|
|
|
|
|
@pytest.mark.parametrize('cfg', [cfg1, cfg2]) |
|
@pytest.mark.unittest |
|
def test_cql_continuous(cfg): |
|
policy = CQLPolicy(cfg, enable_field=['collect', 'eval', 'learn']) |
|
assert type(policy._target_model) == TargetNetworkWrapper |
|
q_value = policy._get_q_value(get_transition_batch(cfg.learn.num_actions)) |
|
assert q_value[0].shape[-1] == 1 and q_value[0].shape[-2] == cfg.learn.num_actions |
|
act, log_prob = policy._get_policy_actions(get_transition_batch(cfg.learn.num_actions)) |
|
assert list(act.shape) == [cfg.learn.num_actions * 10, action_space] |
|
sample = get_transition(size=20) |
|
out = policy._forward_learn(sample) |
|
|
|
|
|
def get_transition_discrete(size=20): |
|
data = [] |
|
for i in range(size): |
|
sample = {} |
|
sample['obs'] = torch.zeros(obs_space) |
|
sample['action'] = torch.tensor(i % action_space) |
|
sample['done'] = False |
|
sample['next_obs'] = torch.zeros(obs_space) |
|
sample['reward'] = torch.Tensor([1.]) |
|
data.append(sample) |
|
return data |
|
|
|
|
|
@pytest.mark.parametrize('cfg', [cfg3, cfg4]) |
|
@pytest.mark.unittest |
|
def test_cql_discrete(cfg): |
|
policy = DiscreteCQLPolicy(cfg, enable_field=['collect', 'eval', 'learn']) |
|
assert type(policy._learn_model) == ArgmaxSampleWrapper |
|
assert type(policy._target_model) == TargetNetworkWrapper |
|
assert type(policy._collect_model) == EpsGreedySampleWrapper |
|
sample = get_transition_batch(bs=20) |
|
samples = policy._get_train_sample(sample) |
|
assert len(samples['obs']) == 20 |
|
state = policy._state_dict_learn() |
|
policy._load_state_dict_learn(state) |
|
sample = get_transition_discrete(size=1) |
|
out = policy._forward_learn(sample) |
|
out = policy._forward_collect(get_batch(size=8), eps=0.1) |
|
|
