|
from typing import Set, Dict, Any, TextIO |
|
import os |
|
import yaml |
|
from mlagents.trainers.exception import TrainerConfigError |
|
from mlagents_envs.environment import UnityEnvironment |
|
import argparse |
|
from mlagents_envs import logging_util |
|
|
|
logger = logging_util.get_logger(__name__) |
|
|
|
|
|
class RaiseRemovedWarning(argparse.Action): |
|
""" |
|
Internal custom Action to raise warning when argument is called. |
|
""" |
|
|
|
def __init__(self, nargs=0, **kwargs): |
|
super().__init__(nargs=nargs, **kwargs) |
|
|
|
def __call__(self, arg_parser, namespace, values, option_string=None): |
|
logger.warning(f"The command line argument {option_string} was removed.") |
|
|
|
|
|
class DetectDefault(argparse.Action): |
|
""" |
|
Internal custom Action to help detect arguments that aren't default. |
|
""" |
|
|
|
non_default_args: Set[str] = set() |
|
|
|
def __call__(self, arg_parser, namespace, values, option_string=None): |
|
setattr(namespace, self.dest, values) |
|
DetectDefault.non_default_args.add(self.dest) |
|
|
|
|
|
class DetectDefaultStoreTrue(DetectDefault): |
|
""" |
|
Internal class to help detect arguments that aren't default. |
|
Used for store_true arguments. |
|
""" |
|
|
|
def __init__(self, nargs=0, **kwargs): |
|
super().__init__(nargs=nargs, **kwargs) |
|
|
|
def __call__(self, arg_parser, namespace, values, option_string=None): |
|
super().__call__(arg_parser, namespace, True, option_string) |
|
|
|
|
|
class StoreConfigFile(argparse.Action): |
|
""" |
|
Custom Action to store the config file location not as part of the CLI args. |
|
This is because we want to maintain an equivalence between the config file's |
|
contents and the args themselves. |
|
""" |
|
|
|
trainer_config_path: str |
|
|
|
def __call__(self, arg_parser, namespace, values, option_string=None): |
|
delattr(namespace, self.dest) |
|
StoreConfigFile.trainer_config_path = values |
|
|
|
|
|
def _create_parser() -> argparse.ArgumentParser: |
|
argparser = argparse.ArgumentParser( |
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter |
|
) |
|
argparser.add_argument( |
|
"trainer_config_path", action=StoreConfigFile, nargs="?", default=None |
|
) |
|
argparser.add_argument( |
|
"--env", |
|
default=None, |
|
dest="env_path", |
|
help="Path to the Unity executable to train", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--load", |
|
default=False, |
|
dest="load_model", |
|
action=DetectDefaultStoreTrue, |
|
help=argparse.SUPPRESS, |
|
) |
|
argparser.add_argument( |
|
"--resume", |
|
default=False, |
|
dest="resume", |
|
action=DetectDefaultStoreTrue, |
|
help="Whether to resume training from a checkpoint. Specify a --run-id to use this option. " |
|
"If set, the training code loads an already trained model to initialize the neural network " |
|
"before resuming training. This option is only valid when the models exist, and have the same " |
|
"behavior names as the current agents in your scene.", |
|
) |
|
argparser.add_argument( |
|
"--deterministic", |
|
default=False, |
|
dest="deterministic", |
|
action=DetectDefaultStoreTrue, |
|
help="Whether to select actions deterministically in policy. `dist.mean` for continuous action " |
|
"space, and `dist.argmax` for deterministic action space ", |
|
) |
|
argparser.add_argument( |
|
"--force", |
|
default=False, |
|
dest="force", |
|
action=DetectDefaultStoreTrue, |
|
help="Whether to force-overwrite this run-id's existing summary and model data. (Without " |
|
"this flag, attempting to train a model with a run-id that has been used before will throw " |
|
"an error.", |
|
) |
|
argparser.add_argument( |
|
"--run-id", |
|
default="ppo", |
|
help="The identifier for the training run. This identifier is used to name the " |
|
"subdirectories in which the trained model and summary statistics are saved as well " |
|
"as the saved model itself. If you use TensorBoard to view the training statistics, " |
|
"always set a unique run-id for each training run. (The statistics for all runs with the " |
|
"same id are combined as if they were produced by a the same session.)", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--initialize-from", |
|
metavar="RUN_ID", |
|
default=None, |
|
help="Specify a previously saved run ID from which to initialize the model from. " |
|
"This can be used, for instance, to fine-tune an existing model on a new environment. " |
|
"Note that the previously saved models must have the same behavior parameters as your " |
|
"current environment.", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--seed", |
|
default=-1, |
|
type=int, |
|
help="A number to use as a seed for the random number generator used by the training code", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--train", |
|
default=False, |
|
dest="train_model", |
|
action=DetectDefaultStoreTrue, |
|
help=argparse.SUPPRESS, |
|
) |
|
argparser.add_argument( |
|
"--inference", |
|
default=False, |
|
dest="inference", |
|
action=DetectDefaultStoreTrue, |
|
help="Whether to run in Python inference mode (i.e. no training). Use with --resume to load " |
|
"a model trained with an existing run ID.", |
|
) |
|
argparser.add_argument( |
|
"--base-port", |
|
default=UnityEnvironment.BASE_ENVIRONMENT_PORT, |
|
type=int, |
|
help="The starting port for environment communication. Each concurrent Unity environment " |
|
"instance will get assigned a port sequentially, starting from the base-port. Each instance " |
|
"will use the port (base_port + worker_id), where the worker_id is sequential IDs given to " |
|
"each instance from 0 to (num_envs - 1). Note that when training using the Editor rather " |
|
"than an executable, the base port will be ignored.", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--num-envs", |
|
default=1, |
|
type=int, |
|
help="The number of concurrent Unity environment instances to collect experiences " |
|
"from when training", |
|
action=DetectDefault, |
|
) |
|
|
|
argparser.add_argument( |
|
"--num-areas", |
|
default=1, |
|
type=int, |
|
help="The number of parallel training areas in each Unity environment instance.", |
|
action=DetectDefault, |
|
) |
|
|
|
argparser.add_argument( |
|
"--debug", |
|
default=False, |
|
action=DetectDefaultStoreTrue, |
|
help="Whether to enable debug-level logging for some parts of the code", |
|
) |
|
argparser.add_argument( |
|
"--env-args", |
|
default=None, |
|
nargs=argparse.REMAINDER, |
|
help="Arguments passed to the Unity executable. Be aware that the standalone build will also " |
|
"process these as Unity Command Line Arguments. You should choose different argument names if " |
|
"you want to create environment-specific arguments. All arguments after this flag will be " |
|
"passed to the executable.", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--max-lifetime-restarts", |
|
default=10, |
|
help="The max number of times a single Unity executable can crash over its lifetime before ml-agents exits. " |
|
"Can be set to -1 if no limit is desired.", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--restarts-rate-limit-n", |
|
default=1, |
|
help="The maximum number of times a single Unity executable can crash over a period of time (period set in " |
|
"restarts-rate-limit-period-s). Can be set to -1 to not use rate limiting with restarts.", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--restarts-rate-limit-period-s", |
|
default=60, |
|
help="The period of time --restarts-rate-limit-n applies to.", |
|
action=DetectDefault, |
|
) |
|
argparser.add_argument( |
|
"--torch", |
|
default=False, |
|
action=RaiseRemovedWarning, |
|
help="(Removed) Use the PyTorch framework.", |
|
) |
|
argparser.add_argument( |
|
"--tensorflow", |
|
default=False, |
|
action=RaiseRemovedWarning, |
|
help="(Removed) Use the TensorFlow framework.", |
|
) |
|
argparser.add_argument( |
|
"--results-dir", |
|
default="results", |
|
action=DetectDefault, |
|
help="Results base directory", |
|
) |
|
|
|
eng_conf = argparser.add_argument_group(title="Engine Configuration") |
|
eng_conf.add_argument( |
|
"--width", |
|
default=84, |
|
type=int, |
|
help="The width of the executable window of the environment(s) in pixels " |
|
"(ignored for editor training).", |
|
action=DetectDefault, |
|
) |
|
eng_conf.add_argument( |
|
"--height", |
|
default=84, |
|
type=int, |
|
help="The height of the executable window of the environment(s) in pixels " |
|
"(ignored for editor training)", |
|
action=DetectDefault, |
|
) |
|
eng_conf.add_argument( |
|
"--quality-level", |
|
default=5, |
|
type=int, |
|
help="The quality level of the environment(s). Equivalent to calling " |
|
"QualitySettings.SetQualityLevel in Unity.", |
|
action=DetectDefault, |
|
) |
|
eng_conf.add_argument( |
|
"--time-scale", |
|
default=20, |
|
type=float, |
|
help="The time scale of the Unity environment(s). Equivalent to setting " |
|
"Time.timeScale in Unity.", |
|
action=DetectDefault, |
|
) |
|
eng_conf.add_argument( |
|
"--target-frame-rate", |
|
default=-1, |
|
type=int, |
|
help="The target frame rate of the Unity environment(s). Equivalent to setting " |
|
"Application.targetFrameRate in Unity.", |
|
action=DetectDefault, |
|
) |
|
eng_conf.add_argument( |
|
"--capture-frame-rate", |
|
default=60, |
|
type=int, |
|
help="The capture frame rate of the Unity environment(s). Equivalent to setting " |
|
"Time.captureFramerate in Unity.", |
|
action=DetectDefault, |
|
) |
|
eng_conf.add_argument( |
|
"--no-graphics", |
|
default=False, |
|
action=DetectDefaultStoreTrue, |
|
help="Whether to run the Unity executable in no-graphics mode (i.e. without initializing " |
|
"the graphics driver. Use this only if your agents don't use visual observations.", |
|
) |
|
|
|
torch_conf = argparser.add_argument_group(title="Torch Configuration") |
|
torch_conf.add_argument( |
|
"--torch-device", |
|
default=None, |
|
dest="device", |
|
action=DetectDefault, |
|
help='Settings for the default torch.device used in training, for example, "cpu", "cuda", or "cuda:0"', |
|
) |
|
return argparser |
|
|
|
|
|
def load_config(config_path: str) -> Dict[str, Any]: |
|
try: |
|
with open(config_path) as data_file: |
|
return _load_config(data_file) |
|
except OSError: |
|
abs_path = os.path.abspath(config_path) |
|
raise TrainerConfigError(f"Config file could not be found at {abs_path}.") |
|
except UnicodeDecodeError: |
|
raise TrainerConfigError( |
|
f"There was an error decoding Config file from {config_path}. " |
|
f"Make sure your file is save using UTF-8" |
|
) |
|
|
|
|
|
def _load_config(fp: TextIO) -> Dict[str, Any]: |
|
""" |
|
Load the yaml config from the file-like object. |
|
""" |
|
try: |
|
return yaml.safe_load(fp) |
|
except yaml.parser.ParserError as e: |
|
raise TrainerConfigError( |
|
"Error parsing yaml file. Please check for formatting errors. " |
|
"A tool such as http://www.yamllint.com/ can be helpful with this." |
|
) from e |
|
|
|
|
|
parser = _create_parser() |
|
|
