def prepare_modules(model_dir, dump_dir=None):
model_dir = model_dir[:-1] if model_dir.endswith("/") else model_dir
model_dir_name = model_dir.rsplit("/", 1)[-1]
if dump_dir is not None:
current_date = datetime.now().strftime("%y-%m-%d-%H-%M-%S-%M-%f")
new_src_dir = f"/src-{current_date}/"
# @TODO: hardcoded
old_pro_dir = os.path.dirname(os.path.abspath(__file__)) + "/../../"
new_pro_dir = dump_dir + f"/{new_src_dir}/catalyst/"
shutil.copytree(old_pro_dir, new_pro_dir)
old_model_dir = os.path.abspath(model_dir)
model_dir_ = model_dir.rsplit("/", 1)[-1]
new_model_dir = dump_dir + f"/{new_src_dir}/{model_dir_}/"
shutil.copytree(old_model_dir, new_model_dir)
pyfiles = list(
map(lambda x: x.name[:-3],
pathlib.Path(model_dir).glob("*.py"))
)
modules = {}
for name in pyfiles:
module_name = f"{model_dir_name}.{name}"
module_src = model_dir + "/" + f"{name}.py"
module = import_module(module_name, module_src)
modules[name] = module
return modules
def algos_by_dir(dir):
algos = []
dirs = path.Path(dir).listdir()
for logpath in dirs:
config_path = logpath + "/config.json"
checkpoints = path.Path(logpath).glob("*.pth.tar")
for checkpoint_path in checkpoints:
args = argparse.Namespace(config=config_path)
args, config = parse_args_uargs(args, [])
config.get("algorithm", {}).pop("resume", None)
config.get("algorithm", {}).pop("load_optimizer", None)
algo_module = import_module("algo_module", args.algorithm)
trainer_kwargs = algo_module.prepare_for_trainer(config)
algorithm = trainer_kwargs["algorithm"]
algorithm.load_checkpoint(checkpoint_path, load_optimizer=False)
actor_ = algorithm.actor.eval()
name = str(algorithm.__class__).lower()
if "ensemblecritic" in name:
critics_ = [x.eval() for x in algorithm.critics]
elif "td3" in name:
critics_ = [algorithm.critic.eval(), algorithm.critic2.eval()]
else:
raise NotImplemented
history_len = trainer_kwargs["history_len"]
algos.append(
AlgoWrapper(actor=actor_,
critics=critics_,
history_len=history_len,
consensus=IN_CONSENSUS))
return algos
parser.add_argument(
"--config",
type=str,
required=True)
parser.add_argument(
"--algorithm",
type=str,
default=None)
parser.add_argument(
"--logdir",
type=str,
default=None)
args, unknown_args = parser.parse_known_args()
args, config = parse_args_uargs(args, unknown_args, dump_config=True)
algo_module = import_module("algo_module", args.algorithm)
algo_kwargs = algo_module.prepare_for_trainer(config)
redis_server = StrictRedis(port=config.get("redis", {}).get("port", 12000))
redis_prefix = config.get("redis", {}).get("prefix", "")
pprint(config["trainer"])
pprint(algo_kwargs)
trainer = Trainer(
**config["trainer"],
**algo_kwargs,
logdir=args.logdir,
redis_server=redis_server,
redis_prefix=redis_prefix)
type=float,
default=None)
parser.add_argument(
"--max-action-noise",
type=float,
default=None)
parser.add_argument(
"--max-param-noise",
type=float,
default=None)
boolean_flag(parser, "debug", default=False)
args = parser.parse_args()
args, config = parse_args_uargs(args, [])
env_module = import_module("env_module", args.environment)
algo_module = import_module("algo_module", args.algorithm)
def run_sampler(
*,
config, vis, infer,
action_noise_prob, param_noise_prob,
action_noise=None, param_noise=None,
noise_power=None, # @TODO: remove
id=None, resume=None, debug=False):
config_ = copy.deepcopy(config)
if debug:
redis_server = None
redis_prefix = None