File size: 6,121 Bytes
05c9ac2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
import pytest
import io
import os
import yaml
from unittest.mock import patch
from mlagents.trainers.trainer import TrainerFactory
from mlagents.trainers.cli_utils import load_config, _load_config
from mlagents.trainers.ppo.trainer import PPOTrainer
from mlagents.trainers.exception import TrainerConfigError, UnityTrainerException
from mlagents.trainers.settings import RunOptions
from mlagents.trainers.tests.dummy_config import ppo_dummy_config
from mlagents.trainers.environment_parameter_manager import EnvironmentParameterManager
from mlagents.trainers.directory_utils import (
validate_existing_directories,
setup_init_path,
)
@pytest.fixture
def dummy_config():
return RunOptions(behaviors={"testbrain": ppo_dummy_config()})
@patch("mlagents_envs.base_env.BehaviorSpec")
def test_initialize_ppo_trainer(BehaviorSpecMock, dummy_config):
brain_name = "testbrain"
training_behaviors = {"testbrain": BehaviorSpecMock()}
output_path = "results_dir"
train_model = True
load_model = False
seed = 11
expected_reward_buff_cap = 1
base_config = dummy_config.behaviors
expected_config = ppo_dummy_config()
def mock_constructor(
self,
brain,
reward_buff_cap,
trainer_settings,
training,
load,
p_seed,
artifact_path,
):
assert brain == brain_name
assert trainer_settings == expected_config
assert reward_buff_cap == expected_reward_buff_cap
assert training == train_model
assert load == load_model
assert p_seed == seed
assert artifact_path == os.path.join(output_path, brain_name)
with patch.object(PPOTrainer, "__init__", mock_constructor):
trainer_factory = TrainerFactory(
trainer_config=base_config,
output_path=output_path,
train_model=train_model,
load_model=load_model,
seed=seed,
param_manager=EnvironmentParameterManager(),
)
trainers = {}
for brain_name in training_behaviors.keys():
trainers[brain_name] = trainer_factory.generate(brain_name)
assert "testbrain" in trainers
assert isinstance(trainers["testbrain"], PPOTrainer)
def test_handles_no_config_provided():
"""
Make sure the trainer setup handles no configs provided at all.
"""
brain_name = "testbrain"
no_default_config = RunOptions().behaviors
# Pretend this was created without a YAML file
no_default_config.set_config_specified(False)
trainer_factory = TrainerFactory(
trainer_config=no_default_config,
output_path="output_path",
train_model=True,
load_model=False,
seed=42,
param_manager=EnvironmentParameterManager(),
)
trainer_factory.generate(brain_name)
def test_load_config_missing_file():
with pytest.raises(TrainerConfigError):
load_config("thisFileDefinitelyDoesNotExist.yaml")
def test_load_config_valid_yaml():
file_contents = """
this:
- is fine
"""
fp = io.StringIO(file_contents)
res = _load_config(fp)
assert res == {"this": ["is fine"]}
def test_load_config_invalid_yaml():
file_contents = """
you:
- will
- not
- parse
"""
with pytest.raises(TrainerConfigError):
fp = io.StringIO(file_contents)
_load_config(fp)
def test_existing_directories(tmp_path):
output_path = os.path.join(tmp_path, "runid")
# Test fresh new unused path - should do nothing.
validate_existing_directories(output_path, False, False)
# Test resume with fresh path - should throw an exception.
with pytest.raises(UnityTrainerException):
validate_existing_directories(output_path, True, False)
# make a directory
os.mkdir(output_path)
# Test try to train w.o. force, should complain
with pytest.raises(UnityTrainerException):
validate_existing_directories(output_path, False, False)
# Test try to train w/ resume - should work
validate_existing_directories(output_path, True, False)
# Test try to train w/ force - should work
validate_existing_directories(output_path, False, True)
# Test initialize option
init_path = os.path.join(tmp_path, "runid2")
with pytest.raises(UnityTrainerException):
validate_existing_directories(output_path, False, True, init_path)
os.mkdir(init_path)
# Should pass since the directory exists now.
validate_existing_directories(output_path, False, True, init_path)
@pytest.mark.parametrize("dir_exists", [True, False])
def test_setup_init_path(tmpdir, dir_exists):
"""
:return:
"""
test_yaml = """
behaviors:
BigWallJump:
init_path: BigWallJump-6540981.pt #full path
trainer_type: ppo
MediumWallJump:
init_path: {}/test_setup_init_path_results/test_run_id/MediumWallJump/checkpoint.pt
trainer_type: ppo
SmallWallJump:
trainer_type: ppo
checkpoint_settings:
run_id: test_run_id
initialize_from: test_run_id
""".format(
tmpdir
)
run_options = RunOptions.from_dict(yaml.safe_load(test_yaml))
if dir_exists:
init_path = tmpdir.mkdir("test_setup_init_path_results").mkdir("test_run_id")
big = init_path.mkdir("BigWallJump").join("BigWallJump-6540981.pt")
big.write("content")
med = init_path.mkdir("MediumWallJump").join("checkpoint.pt")
med.write("content")
small = init_path.mkdir("SmallWallJump").join("checkpoint.pt")
small.write("content")
setup_init_path(run_options.behaviors, init_path)
assert run_options.behaviors["BigWallJump"].init_path == big
assert run_options.behaviors["MediumWallJump"].init_path == med
assert run_options.behaviors["SmallWallJump"].init_path == small
else:
# don't make dirs and fail
with pytest.raises(UnityTrainerException):
setup_init_path(
run_options.behaviors, run_options.checkpoint_settings.maybe_init_path
)
|