|
from unittest.mock import MagicMock, patch |
|
import pytest |
|
from mlagents.torch_utils import torch |
|
|
|
from mlagents.trainers.trainer_controller import TrainerController |
|
from mlagents.trainers.environment_parameter_manager import EnvironmentParameterManager |
|
from mlagents.trainers.ghost.controller import GhostController |
|
|
|
|
|
@pytest.fixture |
|
def basic_trainer_controller(): |
|
trainer_factory_mock = MagicMock() |
|
trainer_factory_mock.ghost_controller = GhostController() |
|
return TrainerController( |
|
trainer_factory=trainer_factory_mock, |
|
output_path="test_model_path", |
|
run_id="test_run_id", |
|
param_manager=EnvironmentParameterManager(), |
|
train=True, |
|
training_seed=99, |
|
) |
|
|
|
|
|
@patch("numpy.random.seed") |
|
@patch.object(torch, "manual_seed") |
|
def test_initialization_seed(numpy_random_seed, torch_set_seed): |
|
seed = 27 |
|
trainer_factory_mock = MagicMock() |
|
trainer_factory_mock.ghost_controller = GhostController() |
|
TrainerController( |
|
trainer_factory=trainer_factory_mock, |
|
output_path="", |
|
run_id="1", |
|
param_manager=None, |
|
train=True, |
|
training_seed=seed, |
|
) |
|
numpy_random_seed.assert_called_with(seed) |
|
torch_set_seed.assert_called_with(seed) |
|
|
|
|
|
@pytest.fixture |
|
def trainer_controller_with_start_learning_mocks(basic_trainer_controller): |
|
trainer_mock = MagicMock() |
|
trainer_mock.get_step = 0 |
|
trainer_mock.get_max_steps = 5 |
|
trainer_mock.should_still_train = True |
|
trainer_mock.parameters = {"some": "parameter"} |
|
trainer_mock.write_tensorboard_text = MagicMock() |
|
|
|
tc = basic_trainer_controller |
|
tc.trainers = {"testbrain": trainer_mock} |
|
tc.advance = MagicMock() |
|
tc.trainers["testbrain"].get_step = 0 |
|
|
|
def take_step_sideeffect(env): |
|
tc.trainers["testbrain"].get_step += 1 |
|
if ( |
|
not tc.trainers["testbrain"].get_step |
|
<= tc.trainers["testbrain"].get_max_steps |
|
): |
|
tc.trainers["testbrain"].should_still_train = False |
|
if tc.trainers["testbrain"].get_step > 10: |
|
raise KeyboardInterrupt |
|
return 1 |
|
|
|
tc.advance.side_effect = take_step_sideeffect |
|
|
|
tc._save_models = MagicMock() |
|
return tc, trainer_mock |
|
|
|
|
|
def test_start_learning_trains_forever_if_no_train_model( |
|
trainer_controller_with_start_learning_mocks, |
|
): |
|
tc, trainer_mock = trainer_controller_with_start_learning_mocks |
|
tc.train_model = False |
|
|
|
env_mock = MagicMock() |
|
env_mock.close = MagicMock() |
|
env_mock.reset = MagicMock() |
|
env_mock.training_behaviors = MagicMock() |
|
|
|
tc.start_learning(env_mock) |
|
env_mock.reset.assert_called_once() |
|
assert tc.advance.call_count == 11 |
|
tc._save_models.assert_not_called() |
|
|
|
|
|
def test_start_learning_trains_until_max_steps_then_saves( |
|
trainer_controller_with_start_learning_mocks, |
|
): |
|
tc, trainer_mock = trainer_controller_with_start_learning_mocks |
|
|
|
brain_info_mock = MagicMock() |
|
env_mock = MagicMock() |
|
env_mock.close = MagicMock() |
|
env_mock.reset = MagicMock(return_value=brain_info_mock) |
|
env_mock.training_behaviors = MagicMock() |
|
|
|
tc.start_learning(env_mock) |
|
env_mock.reset.assert_called_once() |
|
assert tc.advance.call_count == trainer_mock.get_max_steps + 1 |
|
tc._save_models.assert_called_once() |
|
|
|
|
|
@pytest.fixture |
|
def trainer_controller_with_take_step_mocks(basic_trainer_controller): |
|
trainer_mock = MagicMock() |
|
trainer_mock.get_step = 0 |
|
trainer_mock.get_max_steps = 5 |
|
trainer_mock.parameters = {"some": "parameter"} |
|
trainer_mock.write_tensorboard_text = MagicMock() |
|
|
|
tc = basic_trainer_controller |
|
tc.trainers = {"testbrain": trainer_mock} |
|
tc.managers = {"testbrain": MagicMock()} |
|
|
|
return tc, trainer_mock |
|
|
|
|
|
def test_advance_adds_experiences_to_trainer_and_trains( |
|
trainer_controller_with_take_step_mocks, |
|
): |
|
tc, trainer_mock = trainer_controller_with_take_step_mocks |
|
|
|
brain_name = "testbrain" |
|
|
|
env_mock = MagicMock() |
|
|
|
tc.brain_name_to_identifier[brain_name].add(brain_name) |
|
|
|
tc.advance(env_mock) |
|
|
|
env_mock.reset.assert_not_called() |
|
env_mock.get_steps.assert_called_once() |
|
env_mock.process_steps.assert_called_once() |
|
|
|
|
|
|
