VPG playing PongNoFrameskip-v4 from https://github.com/sgoodfriend/rl-algo-impls/tree/2067e21d62fff5db60168687e7d9e89019a8bfc0
88739bd
| import optuna | |
| from copy import deepcopy | |
| from rl_algo_impls.runner.config import Config, Hyperparams, EnvHyperparams | |
| from rl_algo_impls.runner.env import make_eval_env | |
| from rl_algo_impls.shared.policy.optimize_on_policy import sample_on_policy_hyperparams | |
| from rl_algo_impls.tuning.optimize_env import sample_env_hyperparams | |
| def sample_params( | |
| trial: optuna.Trial, | |
| base_hyperparams: Hyperparams, | |
| base_config: Config, | |
| ) -> Hyperparams: | |
| hyperparams = deepcopy(base_hyperparams) | |
| base_env_hyperparams = EnvHyperparams(**hyperparams.env_hyperparams) | |
| env = make_eval_env(base_config, base_env_hyperparams, override_n_envs=1) | |
| # env_hyperparams | |
| env_hyperparams = sample_env_hyperparams(trial, hyperparams.env_hyperparams, env) | |
| # policy_hyperparams | |
| policy_hyperparams = sample_on_policy_hyperparams( | |
| trial, hyperparams.policy_hyperparams, env | |
| ) | |
| # algo_hyperparams | |
| algo_hyperparams = hyperparams.algo_hyperparams | |
| learning_rate = trial.suggest_float("learning_rate", 1e-5, 2e-3, log=True) | |
| learning_rate_decay = trial.suggest_categorical( | |
| "learning_rate_decay", ["none", "linear"] | |
| ) | |
| n_steps_exp = trial.suggest_int("n_steps_exp", 1, 10) | |
| n_steps = 2**n_steps_exp | |
| trial.set_user_attr("n_steps", n_steps) | |
| gamma = 1.0 - trial.suggest_float("gamma_om", 1e-4, 1e-1, log=True) | |
| trial.set_user_attr("gamma", gamma) | |
| gae_lambda = 1 - trial.suggest_float("gae_lambda_om", 1e-4, 1e-1) | |
| trial.set_user_attr("gae_lambda", gae_lambda) | |
| ent_coef = trial.suggest_float("ent_coef", 1e-8, 2.5e-2, log=True) | |
| ent_coef_decay = trial.suggest_categorical("ent_coef_decay", ["none", "linear"]) | |
| vf_coef = trial.suggest_float("vf_coef", 0.1, 0.7) | |
| max_grad_norm = trial.suggest_float("max_grad_norm", 1e-1, 1e1, log=True) | |
| use_rms_prop = trial.suggest_categorical("use_rms_prop", [True, False]) | |
| normalize_advantage = trial.suggest_categorical( | |
| "normalize_advantage", [True, False] | |
| ) | |
| algo_hyperparams.update( | |
| { | |
| "learning_rate": learning_rate, | |
| "learning_rate_decay": learning_rate_decay, | |
| "n_steps": n_steps, | |
| "gamma": gamma, | |
| "gae_lambda": gae_lambda, | |
| "ent_coef": ent_coef, | |
| "ent_coef_decay": ent_coef_decay, | |
| "vf_coef": vf_coef, | |
| "max_grad_norm": max_grad_norm, | |
| "use_rms_prop": use_rms_prop, | |
| "normalize_advantage": normalize_advantage, | |
| } | |
| ) | |
| if policy_hyperparams.get("use_sde", False): | |
| sde_sample_freq = 2 ** trial.suggest_int("sde_sample_freq_exp", 0, n_steps_exp) | |
| trial.set_user_attr("sde_sample_freq", sde_sample_freq) | |
| algo_hyperparams["sde_sample_freq"] = sde_sample_freq | |
| elif "sde_sample_freq" in algo_hyperparams: | |
| del algo_hyperparams["sde_sample_freq"] | |
| env.close() | |
| return hyperparams | |