164 lines
5 KiB
Python
164 lines
5 KiB
Python
import args
|
|
import os
|
|
import pathlib
|
|
|
|
import torch as T
|
|
import torch.nn as nn
|
|
|
|
from typing import Callable
|
|
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.callbacks import CheckpointCallback
|
|
from stable_baselines3.common.vec_env.vec_monitor import VecMonitor
|
|
|
|
from godot_rl.core.utils import can_import
|
|
from godot_rl.wrappers.onnx.stable_baselines_export import export_ppo_model_as_onnx
|
|
from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv
|
|
|
|
if can_import("ray"):
|
|
print("WARNING: SB3 and ray[rllib] are not compatible.")
|
|
|
|
args, extras = args.parse_args()
|
|
|
|
def handle_onnx_export():
|
|
'''
|
|
Enforces the onnx and zip extentions when saving models.
|
|
This avoids potential conflicts in case of identical names and extentions
|
|
'''
|
|
if args.onnx_export_path is not None:
|
|
path_onnx = pathlib.Path(args.onnx_export_path).with_suffix(".onnx")
|
|
print(f"Exporting onnx to: {os.path.abspath(path_onnx)}")
|
|
export_ppo_model_as_onnx(model, str(path_onnx))
|
|
|
|
def handle_model_save():
|
|
if args.save_model_path is not None:
|
|
zip_save_path = pathlib.Path(Args.save_model_path).with_suffix(".zip")
|
|
print(f"Saving model to: {os.path.abspath(zip_save_path)}")
|
|
model.save(zip_save_path)
|
|
|
|
def close_env():
|
|
try:
|
|
print("Closing env...")
|
|
env.close()
|
|
except Exception as e:
|
|
print(f"Exception while closing env: {e}")
|
|
|
|
path_checkpoint = os.path.join(args.exper_dir, f"{args.exper_name}_checkpoints")
|
|
abs_path_checkpoint = os.path.abspath(path_checkpoint)
|
|
|
|
if args.save_checkpoint_frequency is not None and os.path.isdir(path_checkpoint):
|
|
raise RuntimeError(
|
|
f"{abs_path_checkpoint} already exists."
|
|
"Use a different directory or different name."
|
|
"If you want to override previous checkpoints you have to delete them manually."
|
|
)
|
|
|
|
if args.inference and args.resume_model_path is None:
|
|
raise parser.error(
|
|
"Using --inference requires --resume_model_path to be set."
|
|
)
|
|
|
|
if args.env_path is None and args.viz:
|
|
print("Info: using --viz without --env_path set has no effect.")
|
|
print("\nIn editor training will always render.")
|
|
|
|
env = StableBaselinesGodotEnv(
|
|
env_path=args.env_path,
|
|
show_window=args.viz,
|
|
seed=args.seed,
|
|
n_parallel=args.n_parallel,
|
|
speedup=args.speedup
|
|
)
|
|
env = VecMonitor(env)
|
|
|
|
# LR schedule code snippet from:
|
|
# https://stable-baselines3.readthedocs.io/en/master/guide/examples.html#learning-rate-schedule
|
|
def linear_schedule(initial_value: float) -> Callable[[float], float]:
|
|
"""
|
|
Linear learning rate schedule.
|
|
|
|
:param initial_value: Initial learning rate.
|
|
:return: schedule that computes
|
|
current learning rate depending on remaining progress
|
|
"""
|
|
|
|
def func(progress_remaining: float) -> float:
|
|
"""
|
|
Progress will decrease from 1 (beginning) to 0.
|
|
|
|
:param progress_remaining:
|
|
:return: current learning rate
|
|
"""
|
|
return progress_remaining * initial_value
|
|
|
|
return func
|
|
|
|
policy_kwargs = dict(
|
|
activation_fn=nn.Tanh,
|
|
net_arch = dict(
|
|
pi=[256],
|
|
vf = [2048, 2048]#, 4096, 4096, 4096, 4096]
|
|
),
|
|
optimizer_class = T.optim.Adam,
|
|
optimizer_kwargs = dict(
|
|
betas=(0.9, 0.9),
|
|
eps=1e-5
|
|
),
|
|
)
|
|
|
|
if args.resume_model_path is None:
|
|
if not args.linear_lr_schedule:
|
|
learning_rate = 0.0003
|
|
else:
|
|
linear_schedule(0.0003)
|
|
|
|
model: PPO = PPO(
|
|
# 'MultiInputPolicy' serves as an alias for MultiInputActorCriticPolicy
|
|
"MultiInputPolicy",
|
|
env,
|
|
batch_size=64,
|
|
ent_coef=0.01,
|
|
verbose=2,
|
|
n_steps=256,
|
|
tensorboard_log=args.exper_dir,
|
|
learning_rate=learning_rate,
|
|
policy_kwargs=policy_kwargs,
|
|
# optimizer_kwargs=optimizer_kwargs,
|
|
)
|
|
else:
|
|
path_zip = pathlib.Path(args.resume_model_path)
|
|
print(f"Loading model: {os.path.abspath(pathzip)}")
|
|
model: PPO = PPO.load(
|
|
path_zip,
|
|
env=env,
|
|
tensorboard_log=args.exper_dir
|
|
)
|
|
|
|
if args.inference:
|
|
obs = env.reset()
|
|
for i in range(args.timesteps):
|
|
action, _state = model.predict(obs, deterministic=True)
|
|
obs, reward, done, info = env.step(action)
|
|
else:
|
|
learn_arguments = dict(total_timesteps=args.timesteps, tb_log_name=args.exper_name)
|
|
if args.save_checkpoint_frequency:
|
|
print("Checkpoint saving enabled.")
|
|
print(f"\nCheckpoints will be saved to {abs_path_checkpoint}")
|
|
checkpoint_callback = CheckpointCallback(
|
|
save_freq=(args.save_checkpoint_frequency // env.num_envs),
|
|
save_path=path_checkpoint,
|
|
name_prefix=args.exper_name
|
|
)
|
|
learn_arguments["callback"] = checkpoint_callback
|
|
try:
|
|
model.learn(**learn_arguments)
|
|
except KeyboardInterrupt:
|
|
print(
|
|
"""
|
|
Training interrupted by user. Will save if --save_model_path was set and/or export if --onnx_export was set.
|
|
"""
|
|
)
|
|
|
|
close_env()
|
|
handle_onnx_export()
|
|
handle_model_save()
|