def test_globaltrainingstatus(tmpdir):
path_dir = os.path.join(tmpdir, "test.json")
GlobalTrainingStatus.set_parameter_state("Category1",
StatusType.LESSON_NUM, 3)
GlobalTrainingStatus.save_state(path_dir)
with open(path_dir) as fp:
test_json = json.load(fp)
assert "Category1" in test_json
assert StatusType.LESSON_NUM.value in test_json["Category1"]
assert test_json["Category1"][StatusType.LESSON_NUM.value] == 3
assert "metadata" in test_json
GlobalTrainingStatus.load_state(path_dir)
restored_val = GlobalTrainingStatus.get_parameter_state(
"Category1", StatusType.LESSON_NUM)
assert restored_val == 3
# Test unknown categories and status types (keys)
unknown_category = GlobalTrainingStatus.get_parameter_state(
"Category3", StatusType.LESSON_NUM)
class FakeStatusType(Enum):
NOTAREALKEY = "notarealkey"
unknown_key = GlobalTrainingStatus.get_parameter_state(
"Category1", FakeStatusType.NOTAREALKEY)
assert unknown_category is None
assert unknown_key is None
def __init__(
self,
settings: Optional[Dict[str, EnvironmentParameterSettings]] = None,
run_seed: int = -1,
restore: bool = False,
):
"""
EnvironmentParameterManager manages all the environment parameters of a training
session. It determines when parameters should change and gives access to the
current sampler of each parameter.
:param settings: A dictionary from environment parameter to
EnvironmentParameterSettings.
:param run_seed: When the seed is not provided for an environment parameter,
this seed will be used instead.
:param restore: If true, the EnvironmentParameterManager will use the
GlobalTrainingStatus to try and reload the lesson status of each environment
parameter.
"""
if settings is None:
settings = {}
self._dict_settings = settings
for parameter_name in self._dict_settings.keys():
initial_lesson = GlobalTrainingStatus.get_parameter_state(
parameter_name, StatusType.LESSON_NUM)
if initial_lesson is None or not restore:
GlobalTrainingStatus.set_parameter_state(
parameter_name, StatusType.LESSON_NUM, 0)
self._smoothed_values: Dict[str, float] = defaultdict(float)
for key in self._dict_settings.keys():
self._smoothed_values[key] = 0.0
# Update the seeds of the samplers
self._set_sampler_seeds(run_seed)
def test_model_management(tmpdir):
results_path = os.path.join(tmpdir, "results")
brain_name = "Mock_brain"
final_model_path = os.path.join(results_path, brain_name)
test_checkpoint_list = [
{
"steps": 1,
"file_path": os.path.join(final_model_path, f"{brain_name}-1.nn"),
"reward": 1.312,
"creation_time": time.time(),
"auxillary_file_paths": [],
},
{
"steps": 2,
"file_path": os.path.join(final_model_path, f"{brain_name}-2.nn"),
"reward": 1.912,
"creation_time": time.time(),
"auxillary_file_paths": [],
},
{
"steps": 3,
"file_path": os.path.join(final_model_path, f"{brain_name}-3.nn"),
"reward": 2.312,
"creation_time": time.time(),
"auxillary_file_paths": [],
},
]
GlobalTrainingStatus.set_parameter_state(brain_name,
StatusType.CHECKPOINTS,
test_checkpoint_list)
new_checkpoint_4 = ModelCheckpoint(
4, os.path.join(final_model_path, f"{brain_name}-4.nn"), 2.678,
time.time())
ModelCheckpointManager.add_checkpoint(brain_name, new_checkpoint_4, 4)
assert len(ModelCheckpointManager.get_checkpoints(brain_name)) == 4
new_checkpoint_5 = ModelCheckpoint(
5, os.path.join(final_model_path, f"{brain_name}-5.nn"), 3.122,
time.time())
ModelCheckpointManager.add_checkpoint(brain_name, new_checkpoint_5, 4)
assert len(ModelCheckpointManager.get_checkpoints(brain_name)) == 4
final_model_path = f"{final_model_path}.nn"
final_model_time = time.time()
current_step = 6
final_model = ModelCheckpoint(current_step, final_model_path, 3.294,
final_model_time)
ModelCheckpointManager.track_final_checkpoint(brain_name, final_model)
assert len(ModelCheckpointManager.get_checkpoints(brain_name)) == 4
check_checkpoints = GlobalTrainingStatus.saved_state[brain_name][
StatusType.CHECKPOINTS.value]
assert check_checkpoints is not None
final_model = GlobalTrainingStatus.saved_state[
StatusType.FINAL_CHECKPOINT.value]
assert final_model is not None
def save_model(self) -> None:
"""
Forwarding call to wrapped trainers save_model.
"""
GlobalTrainingStatus.set_parameter_state(self.brain_name,
StatusType.ELO,
self.current_elo)
self.trainer.save_model()
def get_checkpoints(behavior_name: str) -> List[Dict[str, Any]]:
checkpoint_list = GlobalTrainingStatus.get_parameter_state(
behavior_name, StatusType.CHECKPOINTS)
if not checkpoint_list:
checkpoint_list = []
GlobalTrainingStatus.set_parameter_state(behavior_name,
StatusType.CHECKPOINTS,
checkpoint_list)
return checkpoint_list
def update_lessons(
self,
trainer_steps: Dict[str, int],
trainer_max_steps: Dict[str, int],
trainer_reward_buffer: Dict[str, List[float]],
) -> Tuple[bool, bool]:
"""
Given progress metrics, calculates if at least one environment parameter is
in a new lesson and if at least one environment parameter requires the env
to reset.
:param trainer_steps: A dictionary from behavior_name to the number of training
steps this behavior's trainer has performed.
:param trainer_max_steps: A dictionary from behavior_name to the maximum number
of training steps this behavior's trainer has performed.
:param trainer_reward_buffer: A dictionary from behavior_name to the list of
the most recent episode returns for this behavior's trainer.
:returns: A tuple of two booleans : (True if any lesson has changed, True if
environment needs to reset)
"""
must_reset = False
updated = False
for param_name, settings in self._dict_settings.items():
lesson_num = GlobalTrainingStatus.get_parameter_state(
param_name, StatusType.LESSON_NUM)
next_lesson_num = lesson_num + 1
lesson = settings.curriculum[lesson_num]
if (lesson.completion_criteria is not None
and len(settings.curriculum) > next_lesson_num):
behavior_to_consider = lesson.completion_criteria.behavior
if behavior_to_consider in trainer_steps:
must_increment, new_smoothing = lesson.completion_criteria.need_increment(
float(trainer_steps[behavior_to_consider]) /
float(trainer_max_steps[behavior_to_consider]),
trainer_reward_buffer[behavior_to_consider],
self._smoothed_values[param_name],
)
self._smoothed_values[param_name] = new_smoothing
if must_increment:
GlobalTrainingStatus.set_parameter_state(
param_name, StatusType.LESSON_NUM, next_lesson_num)
new_lesson_name = settings.curriculum[
next_lesson_num].name
new_lesson_value = settings.curriculum[
next_lesson_num].value
logger.info(
f"Parameter '{param_name}' has been updated to {new_lesson_value}."
+ f" Now in lesson '{new_lesson_name}'")
updated = True
if lesson.completion_criteria.require_reset:
must_reset = True
return updated, must_reset
def track_final_checkpoint(cls, behavior_name: str,
final_checkpoint: NNCheckpoint) -> None:
"""
Ensures number of checkpoints stored is within the max number of checkpoints
defined by the user and finally stores the information about the final
model (or intermediate model if training is interrupted).
:param behavior_name: Behavior name of the model.
:param final_checkpoint: Checkpoint information for the final model.
"""
final_model_dict = attr.asdict(final_checkpoint)
GlobalTrainingStatus.set_parameter_state(behavior_name,
StatusType.FINAL_CHECKPOINT,
final_model_dict)
def add_checkpoint(cls, behavior_name: str, new_checkpoint: NNCheckpoint,
keep_checkpoints: int) -> None:
"""
Make room for new checkpoint if needed and insert new checkpoint information.
:param behavior_name: Behavior name for the checkpoint.
:param new_checkpoint: The new checkpoint to be recorded.
:param keep_checkpoints: Number of checkpoints to record (user-defined).
"""
new_checkpoint_dict = attr.asdict(new_checkpoint)
checkpoints = cls.get_checkpoints(behavior_name)
checkpoints.append(new_checkpoint_dict)
cls._cleanup_extra_checkpoints(checkpoints, keep_checkpoints)
GlobalTrainingStatus.set_parameter_state(behavior_name,
StatusType.CHECKPOINTS,
checkpoints)
def advance(self, env: EnvManager) -> int:
# Get steps
with hierarchical_timer("env_step"):
num_steps = env.advance()
# Report current lesson
if self.meta_curriculum:
for brain_name, curr in self.meta_curriculum.brains_to_curricula.items(
):
if brain_name in self.trainers:
self.trainers[brain_name].stats_reporter.set_stat(
"Environment/Lesson", curr.lesson_num)
GlobalTrainingStatus.set_parameter_state(
brain_name, StatusType.LESSON_NUM, curr.lesson_num)
for trainer in self.trainers.values():
if not trainer.threaded:
with hierarchical_timer("trainer_advance"):
trainer.advance()
return num_steps