def test_soft_update():
class TestModule(torch.nn.Module):
def __init__(self, vals):
super().__init__()
self.parameter = torch.nn.Parameter(torch.ones(5, 5, 5) * vals)
tm1 = TestModule(0)
tm2 = TestModule(1)
tm3 = TestModule(2)
ModelUtils.soft_update(tm1, tm3, tau=0.5)
assert torch.equal(tm3.parameter, torch.ones(5, 5, 5))
ModelUtils.soft_update(tm1, tm2, tau=1.0)
assert torch.equal(tm2.parameter, tm1.parameter)
def compare_models(module_1, module_2):
is_same = True
for key_item_1, key_item_2 in zip(module_1.state_dict().items(),
module_2.state_dict().items()):
# Compare tensors in state_dict and not the keys.
is_same = torch.equal(key_item_1[1], key_item_2[1]) and is_same
return is_same
def test_load_policy_different_hidden_units(tmp_path, vis_encode_type):
path1 = os.path.join(tmp_path, "runid1")
trainer_params = TrainerSettings()
trainer_params.network_settings = NetworkSettings(
hidden_units=12, vis_encode_type=EncoderType(vis_encode_type))
policy = create_policy_mock(trainer_params, use_visual=True)
conv_params = [
mod for mod in policy.actor.parameters() if len(mod.shape) > 2
]
model_saver = TorchModelSaver(trainer_params, path1)
model_saver.register(policy)
model_saver.initialize_or_load(policy)
policy.set_step(2000)
mock_brain_name = "MockBrain"
model_saver.save_checkpoint(mock_brain_name, 2000)
# Try load from this path
trainer_params2 = TrainerSettings()
trainer_params2.network_settings = NetworkSettings(
hidden_units=10, vis_encode_type=EncoderType(vis_encode_type))
model_saver2 = TorchModelSaver(trainer_params2, path1, load=True)
policy2 = create_policy_mock(trainer_params2, use_visual=True)
conv_params2 = [
mod for mod in policy2.actor.parameters() if len(mod.shape) > 2
]
# asserts convolutions have different parameters before load
for conv1, conv2 in zip(conv_params, conv_params2):
assert not torch.equal(conv1, conv2)
# asserts layers still have different dimensions
for mod1, mod2 in zip(policy.actor.parameters(),
policy2.actor.parameters()):
if mod1.shape[0] == 12:
assert mod2.shape[0] == 10
model_saver2.register(policy2)
model_saver2.initialize_or_load(policy2)
# asserts convolutions have same parameters after load
for conv1, conv2 in zip(conv_params, conv_params2):
assert torch.equal(conv1, conv2)
# asserts layers still have different dimensions
for mod1, mod2 in zip(policy.actor.parameters(),
policy2.actor.parameters()):
if mod1.shape[0] == 12:
assert mod2.shape[0] == 10
def test_list_to_tensor():
# Test converting pure list
unconverted_list = [[1.0, 2], [1, 3], [1, 4]]
tensor = ModelUtils.list_to_tensor(unconverted_list)
# Should be equivalent to torch.tensor conversion
assert torch.equal(tensor, torch.tensor(unconverted_list))
# Test converting pure numpy array
np_list = np.asarray(unconverted_list)
tensor = ModelUtils.list_to_tensor(np_list)
# Should be equivalent to torch.tensor conversion
assert torch.equal(tensor, torch.tensor(unconverted_list))
# Test converting list of numpy arrays
list_of_np = [np.asarray(_el) for _el in unconverted_list]
tensor = ModelUtils.list_to_tensor(list_of_np)
# Should be equivalent to torch.tensor conversion
assert torch.equal(tensor, torch.tensor(unconverted_list, dtype=torch.float32))
def test_actions_to_onehot():
all_actions = torch.tensor([[1, 0, 2], [1, 0, 2]])
action_size = [2, 1, 3]
oh_actions = ModelUtils.actions_to_onehot(all_actions, action_size)
expected_result = [
torch.tensor([[0, 1], [0, 1]], dtype=torch.float),
torch.tensor([[1], [1]], dtype=torch.float),
torch.tensor([[0, 0, 1], [0, 0, 1]], dtype=torch.float),
]
for res, exp in zip(oh_actions, expected_result):
assert torch.equal(res, exp)
def test_deterministic_sample_action():
inp_size = 4
act_size = 8
action_model, masks = create_action_model(inp_size, act_size, deterministic=True)
sample_inp = torch.ones((1, inp_size))
dists = action_model._get_dists(sample_inp, masks=masks)
agent_action1 = action_model._sample_action(dists)
agent_action2 = action_model._sample_action(dists)
agent_action3 = action_model._sample_action(dists)
assert torch.equal(agent_action1.continuous_tensor, agent_action2.continuous_tensor)
assert torch.equal(agent_action1.continuous_tensor, agent_action3.continuous_tensor)
assert torch.equal(agent_action1.discrete_tensor, agent_action2.discrete_tensor)
assert torch.equal(agent_action1.discrete_tensor, agent_action3.discrete_tensor)
action_model, masks = create_action_model(inp_size, act_size, deterministic=False)
sample_inp = torch.ones((1, inp_size))
dists = action_model._get_dists(sample_inp, masks=masks)
agent_action1 = action_model._sample_action(dists)
agent_action2 = action_model._sample_action(dists)
agent_action3 = action_model._sample_action(dists)
assert not torch.equal(
agent_action1.continuous_tensor, agent_action2.continuous_tensor
)
assert not torch.equal(
agent_action1.continuous_tensor, agent_action3.continuous_tensor
)
chance_counter = 0
if not torch.equal(agent_action1.discrete_tensor, agent_action2.discrete_tensor):
chance_counter += 1
if not torch.equal(agent_action1.discrete_tensor, agent_action3.discrete_tensor):
chance_counter += 1
if not torch.equal(agent_action2.discrete_tensor, agent_action3.discrete_tensor):
chance_counter += 1
assert chance_counter > 1
