75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
import subprocess
|
|
import time
|
|
|
|
import torch
|
|
import tyro
|
|
import yaml
|
|
import os
|
|
|
|
from brittle_star_project.dataclasses import PPOArgs
|
|
from brittle_star_project.trainers.PPOTrainer import PPOTrainer
|
|
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
|
|
|
|
|
|
def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper:
|
|
if config_path is None:
|
|
return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
|
|
return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
|
|
|
|
|
|
def parse_args(log: bool = True) -> PPOArgs:
|
|
temp_args = tyro.cli(PPOArgs)
|
|
|
|
if temp_args.hyperparameter_config_path is not None:
|
|
if log:
|
|
print(f"Loading hyperparameter config from {temp_args.hyperparameter_config_path}")
|
|
|
|
with open(temp_args.hyperparameter_config_path, "r") as f:
|
|
config = yaml.safe_load(f)
|
|
if config:
|
|
# parse PPOArgs with defaults from yaml.
|
|
for key, value in config.items():
|
|
if hasattr(temp_args, key):
|
|
setattr(temp_args, key, value)
|
|
|
|
# Reparse CLI to ensure they OVERRIDE the yaml
|
|
args = tyro.cli(PPOArgs, default=temp_args)
|
|
else:
|
|
if log:
|
|
print("No hyperparameter config provided, using default config")
|
|
|
|
args = temp_args
|
|
return args
|
|
|
|
|
|
def get_git_hash() -> str:
|
|
try:
|
|
return (
|
|
subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
|
|
)
|
|
except subprocess.CalledProcessError | UnicodeDecodeError:
|
|
return "none"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
|
|
args.batch_size = args.num_envs * args.num_steps
|
|
args.minibatch_size = args.batch_size // args.num_minibatches
|
|
args.num_iterations = args.total_timesteps // args.batch_size
|
|
|
|
git_hash = get_git_hash()
|
|
run_name = f"{args.exp_name}__seed_{args.seed}__{git_hash}__{int(time.time())}"
|
|
if args.run_dir is None:
|
|
run_dir = f"runs/{run_name}"
|
|
else:
|
|
run_dir = args.run_dir
|
|
|
|
os.makedirs(run_dir, exist_ok=True)
|
|
|
|
env = make_env(args.env_config_path, args.num_envs)
|
|
|
|
torch.backends.cudnn.deterministic = args.torch_deterministic
|
|
|
|
ppo_trainer = PPOTrainer(args, env, run_dir, run_name)
|
|
ppo_trainer.train()
|