Moved the tests around

ml-agents/mlagents/trainers/tests/torch/saver/test_saver.py

+import pytest
+from unittest import mock
+import os
+
+import numpy as np
+import torch
+from mlagents.trainers.policy.torch_policy import TorchPolicy
+from mlagents.trainers.ppo.optimizer_torch import TorchPPOOptimizer
+from mlagents.trainers.saver.torch_saver import TorchSaver
+from mlagents.trainers.settings import TrainerSettings
+from mlagents.trainers.tests import mock_brain as mb
+from mlagents.trainers.tests.torch.test_policy import create_policy_mock
+
+
+def test_register(tmp_path):
+    trainer_params = TrainerSettings()
+    saver = TorchSaver(trainer_params, tmp_path)
+
+    opt = mock.Mock(spec=TorchPPOOptimizer)
+    opt.get_modules = mock.Mock(return_value={})
+    saver.register(opt)
+    assert saver.policy is None
+
+    trainer_params = TrainerSettings()
+    policy = create_policy_mock(trainer_params)
+    opt.get_modules = mock.Mock(return_value={})
+    saver.register(policy)
+    assert saver.policy is not None
+
+
+def test_load_save(tmp_path):
+    path1 = os.path.join(tmp_path, "runid1")
+    path2 = os.path.join(tmp_path, "runid2")
+    trainer_params = TrainerSettings()
+    policy = create_policy_mock(trainer_params)
+    saver = TorchSaver(trainer_params, path1)
+    saver.register(policy)
+    saver.initialize_or_load(policy)
+    policy.set_step(2000)
+
+    mock_brain_name = "MockBrain"
+    saver.save_checkpoint(mock_brain_name, 2000)
+    assert len(os.listdir(tmp_path)) > 0
+
+    # Try load from this path
+    saver2 = TorchSaver(trainer_params, path1, load=True)
+    policy2 = create_policy_mock(trainer_params)
+    saver2.register(policy2)
+    saver2.initialize_or_load(policy2)
+    _compare_two_policies(policy, policy2)
+    assert policy2.get_current_step() == 2000
+
+    # Try initialize from path 1
+    trainer_params.init_path = path1
+    saver3 = TorchSaver(trainer_params, path2)
+    policy3 = create_policy_mock(trainer_params)
+    saver3.register(policy3)
+    saver3.initialize_or_load(policy3)
+    _compare_two_policies(policy2, policy3)
+    # Assert that the steps are 0.
+    assert policy3.get_current_step() == 0
+
+
+# TorchPolicy.evalute() returns log_probs instead of all_log_probs like tf does.
+# resulting in indeterministic results for testing.
+# So here use sample_actions instead.
+def _compare_two_policies(policy1: TorchPolicy, policy2: TorchPolicy) -> None:
+    """
+    Make sure two policies have the same output for the same input.
+    """
+    decision_step, _ = mb.create_steps_from_behavior_spec(
+        policy1.behavior_spec, num_agents=1
+    )
+    vec_vis_obs, masks = policy1._split_decision_step(decision_step)
+    vec_obs = [torch.as_tensor(vec_vis_obs.vector_observations)]
+    vis_obs = [torch.as_tensor(vis_ob) for vis_ob in vec_vis_obs.visual_observations]
+    memories = torch.as_tensor(
+        policy1.retrieve_memories(list(decision_step.agent_id))
+    ).unsqueeze(0)
+
+    with torch.no_grad():
+        _, log_probs1, _, _, _ = policy1.sample_actions(
+            vec_obs, vis_obs, masks=masks, memories=memories, all_log_probs=True
+        )
+        _, log_probs2, _, _, _ = policy2.sample_actions(
+            vec_obs, vis_obs, masks=masks, memories=memories, all_log_probs=True
+        )
+
+    np.testing.assert_array_equal(log_probs1, log_probs2)
+
+
+@pytest.mark.parametrize("discrete", [True, False], ids=["discrete", "continuous"])
+@pytest.mark.parametrize("visual", [True, False], ids=["visual", "vector"])
+@pytest.mark.parametrize("rnn", [True, False], ids=["rnn", "no_rnn"])
+def test_checkpoint_conversion(tmpdir, rnn, visual, discrete):
+    dummy_config = TrainerSettings()
+    model_path = os.path.join(tmpdir, "Mock_Brain")
+    policy = create_policy_mock(
+        dummy_config, use_rnn=rnn, use_discrete=discrete, use_visual=visual
+    )
+    trainer_params = TrainerSettings()
+    saver = TorchSaver(trainer_params, model_path)
+    saver.register(policy)
+    saver.save_checkpoint("Mock_Brain", 100)
+    assert os.path.isfile(model_path + "/Mock_Brain-100.onnx")

ml-agents/mlagents/trainers/tests/torch/saver/test_saver_reward_providers.py

+import pytest
+import os
+
+from mlagents.trainers.ppo.optimizer_torch import TorchPPOOptimizer
+from mlagents.trainers.sac.optimizer_torch import TorchSACOptimizer
+from mlagents.trainers.saver.torch_saver import TorchSaver
+from mlagents.trainers.settings import (
+    TrainerSettings,
+    RewardSignalType,
+    CuriositySettings,
+    GAILSettings,
+    PPOSettings,
+    SACSettings,
+)
+from mlagents.trainers.tests.torch.test_policy import create_policy_mock
+
+DEMO_PATH = (
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir, os.pardir)
+    + "/test.demo"
+)
+
+
+@pytest.mark.parametrize(
+    "optimizer",
+    [(TorchPPOOptimizer, PPOSettings), (TorchSACOptimizer, SACSettings)],
+    ids=["ppo", "sac"],
+)
+def test_reward_provider_save(tmp_path, optimizer):
+    OptimizerClass, HyperparametersClass = optimizer
+
+    trainer_settings = TrainerSettings()
+    trainer_settings.hyperparameters = HyperparametersClass()
+    trainer_settings.reward_signals = {
+        RewardSignalType.CURIOSITY: CuriositySettings(),
+        RewardSignalType.GAIL: GAILSettings(demo_path=DEMO_PATH),
+    }
+    policy = create_policy_mock(trainer_settings, use_discrete=False)
+    optimizer = OptimizerClass(policy, trainer_settings)
+
+    # save at path 1
+    path1 = os.path.join(tmp_path, "runid1")
+    saver = TorchSaver(trainer_settings, path1)
+    saver.register(policy)
+    saver.register(optimizer)
+    saver.initialize_or_load(policy)
+    policy.set_step(2000)
+    saver.save_checkpoint("MockBrain", 2000)
+
+    # create a new optimizer and policy
+    optimizer2 = OptimizerClass(policy, trainer_settings)
+    policy2 = create_policy_mock(trainer_settings, use_discrete=False)
+
+    # load weights
+    saver2 = TorchSaver(trainer_settings, path1, load=True)
+    saver2.register(policy2)
+    saver2.register(optimizer2)
+    saver2.initialize_or_load(policy2)
+    saver2.initialize_or_load()  # This is to load the optimizers
+
+    # assert the models have the same weights
+    module_dict_1 = optimizer.get_modules()
+    module_dict_2 = optimizer2.get_modules()
+    assert "Module:GAIL" in module_dict_1
+    assert "Module:GAIL" in module_dict_2
+    for name, module1 in module_dict_1.items():
+        assert name in module_dict_2
+        module2 = module_dict_2[name]
+        if hasattr(module1, "parameters"):
+            for param1, param2 in zip(module1.parameters(), module2.parameters()):
+                print(param1.data)
+                print(param2.data)
+                assert param1.data.ne(param2.data).sum() == 0