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