Repository: optuna/optuna

Issue:
Stable-Baselines3 RL Example
<!-- Please write a clear and concise description of the feature proposal. -->

Hello,

I was wondering if you would be interested in adding an example with Optuna + Stable-Baselines3 for hyperparameter tuning in an reinforcement learning context?

It has been used successfully in both v2 and v3 in the zoo repo:
- https://github.com/DLR-RM/rl-baselines3-zoo
- https://github.com/araffin/rl-baselines-zoo

Stable Baselines:
- https://github.com/hill-a/stable-baselines (v2, tensorflow)
- https://github.com/DLR-RM/stable-baselines3 (v3, pytorch edition)

## Motivation

Hyperparameter tuning is key to RL success and Stable Baselines (SB and SB3) is one of the most popular RL library.




Discussion:
Hi @araffin. Thanks for your suggestion and for letting us know about very interesting use cases.

I'd like to know your thoughts on the following as there is probably room for some discussion.  Namely, whether this example should be located in this repository or in one of yours, and by whom it should be created. AFAIK, there are no experts on Stable-Baselines3 in the committer team here so it might be better to have it maintained by one of your references. Is that different from what you imagined? 
>Namely, whether this example should be located in this repository or in one of yours

This is already included in our two repos in fact.
I can provide a simpler example if needed with a pull request to your repo. Our API is also fairly stable, so we should not have to update it too often.

Oh, I see. Could you point me to those examples? I'm sure other committers are interested in those as well. As you pointed out, the stability of the API is something we consider since if the example is to be added, we'll most likely run it in our daily CI jobs to make sure it's not breaking. I don't know how heavy of a task it is but as it's easier to discuss based on code, would you mind creating a pull request? 
>Could you point me to those examples? 

https://github.com/DLR-RM/rl-baselines3-zoo/blob/master/utils/hyperparams_opt.py
https://github.com/DLR-RM/rl-baselines3-zoo/blob/master/utils/callbacks.py#L10

I was writting the PR when I realized that optuna was also supporting python3.5 which is not the case for SB3 (it supports python 3.6+ to use f-strings).
A solution would be to use SB2 instead for python3.5 (which depends on tensorflow), they share the same api, so it should be a one line change.

What do you think?
This issue has not seen any recent activity.
I'm really sorry for the delayed response, noticing you already started working on an example. 

You are right about the supported Python version. Did you imagine something along conditioning imports based on the Python version used? That might be a reasonable approach to keep https://github.com/optuna/optuna/blob/master/setup.py simple. I suppose that it also depends on the user base of SB2 so would like to hear your opinion on that. FYI, we are considering dropping Py 3.5 in Optuna as well, https://github.com/optuna/optuna/issues/1067.
>FYI, we are considering dropping Py 3.5 in Optuna as well, #1067.

Good news =) Any date in mind?

I think the best option would be to wait for Optuna to drop python 3.5, this will make the PR much simpler ;)
That'd be one option. On the other hand, we are quite interested in taking a look at your examples. How about creating examples that only support Python 3.6 and up? We already have such examples and it's fine as well (we'll skip those conditionally based on Python versions from the CI).

The reason is that we don't have a fixed date yet for the drop. 🙁 

Resolving patch:
diff --git a/examples/rl/sb3_simple.py b/examples/rl/sb3_simple.py
new file mode 100644
--- /dev/null
+++ b/examples/rl/sb3_simple.py
@@ -0,0 +1,179 @@
+""" Optuna example that optimizes the hyperparameters of
+a reinforcement learning agent using A2C implementation from Stable-Baselines3
+on a OpenAI Gym environment.
+
+This is a simplified version of what can be found in https://github.com/DLR-RM/rl-baselines3-zoo.
+
+You can run this example as follows:
+    $ python sb3_simple.py
+
+"""
+from typing import Any
+from typing import Dict
+
+import gym
+from stable_baselines3 import A2C
+from stable_baselines3.common.callbacks import EvalCallback
+import torch
+import torch.nn as nn
+
+import optuna
+from optuna.pruners import MedianPruner
+from optuna.samplers import TPESampler
+
+
+N_TRIALS = 100
+N_JOBS = 2
+N_STARTUP_TRIALS = 5
+N_EVALUATIONS = 2
+N_TIMESTEPS = int(2e4)
+EVAL_FREQ = int(N_TIMESTEPS / N_EVALUATIONS)
+N_EVAL_EPISODES = 3
+
+ENV_ID = "CartPole-v1"
+
+DEFAULT_HYPERPARAMS = {
+    "policy": "MlpPolicy",
+    "env": ENV_ID,
+}
+
+
+def sample_a2c_params(trial: optuna.Trial) -> Dict[str, Any]:
+    """Sampler for A2C hyperparameters."""
+    gamma = 1.0 - trial.suggest_float("gamma", 0.0001, 0.1, log=True)
+    max_grad_norm = trial.suggest_float("max_grad_norm", 0.3, 5.0, log=True)
+    gae_lambda = 1.0 - trial.suggest_float("gae_lambda", 0.001, 0.2, log=True)
+    n_steps = 2 ** trial.suggest_int("exponent_n_steps", 3, 10)
+    learning_rate = trial.suggest_float("lr", 1e-5, 1, log=True)
+    ent_coef = trial.suggest_float("ent_coef", 0.00000001, 0.1, log=True)
+    ortho_init = trial.suggest_categorical("ortho_init", [False, True])
+    net_arch = trial.suggest_categorical("net_arch", ["tiny", "small"])
+    activation_fn = trial.suggest_categorical("activation_fn", ["tanh", "relu"])
+
+    # Display true values
+    trial.set_user_attr("gamma_", gamma)
+    trial.set_user_attr("gae_lambda_", gae_lambda)
+    trial.set_user_attr("n_steps", n_steps)
+
+    net_arch = [
+        {"pi": [64], "vf": [64]} if net_arch == "tiny" else {"pi": [64, 64], "vf": [64, 64]}
+    ]
+
+    activation_fn = {"tanh": nn.Tanh, "relu": nn.ReLU,}[activation_fn]
+
+    return {
+        "n_steps": n_steps,
+        "gamma": gamma,
+        "gae_lambda": gae_lambda,
+        "learning_rate": learning_rate,
+        "ent_coef": ent_coef,
+        "max_grad_norm": max_grad_norm,
+        "policy_kwargs": {
+            "net_arch": net_arch,
+            "activation_fn": activation_fn,
+            "ortho_init": ortho_init,
+        },
+    }
+
+
+class TrialEvalCallback(EvalCallback):
+    """Callback used for evaluating and reporting a trial."""
+
+    def __init__(
+        self,
+        eval_env: gym.Env,
+        trial: optuna.Trial,
+        n_eval_episodes: int = 5,
+        eval_freq: int = 10000,
+        deterministic: bool = True,
+        verbose: int = 0,
+    ):
+
+        super().__init__(
+            eval_env=eval_env,
+            n_eval_episodes=n_eval_episodes,
+            eval_freq=eval_freq,
+            deterministic=deterministic,
+            verbose=verbose,
+        )
+        self.trial = trial
+        self.eval_idx = 0
+        self.is_pruned = False
+
+    def _on_step(self) -> bool:
+        if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0:
+            super()._on_step()
+            self.eval_idx += 1
+            self.trial.report(self.last_mean_reward, self.eval_idx)
+            # Prune trial if need
+            if self.trial.should_prune():
+                self.is_pruned = True
+                return False
+        return True
+
+
+def objective(trial: optuna.Trial) -> float:
+
+    kwargs = DEFAULT_HYPERPARAMS.copy()
+    # Sample hyperparameters
+    kwargs.update(sample_a2c_params(trial))
+    # Create the RL model
+    model = A2C(**kwargs)
+    # Create env used for evaluation
+    eval_env = gym.make(ENV_ID)
+    # Create the callback that will periodically evaluate
+    # and report the performance
+    eval_callback = TrialEvalCallback(
+        eval_env, trial, n_eval_episodes=N_EVAL_EPISODES, eval_freq=EVAL_FREQ, deterministic=True
+    )
+
+    nan_encountered = False
+    try:
+        model.learn(N_TIMESTEPS, callback=eval_callback)
+    except AssertionError as e:
+        # Sometimes, random hyperparams can generate NaN
+        print(e)
+        nan_encountered = True
+    finally:
+        # Free memory
+        model.env.close()
+        eval_env.close()
+
+    # Tell the optimizer that the trial failed
+    if nan_encountered:
+        return float("nan")
+
+    if eval_callback.is_pruned:
+        raise optuna.exceptions.TrialPruned()
+
+    return eval_callback.last_mean_reward
+
+
+if __name__ == "__main__":
+    # Set pytorch num threads to 1 for faster training
+    torch.set_num_threads(1)
+
+    sampler = TPESampler(n_startup_trials=N_STARTUP_TRIALS)
+    # Do not prune before 1/3 of the max budget is used
+    pruner = MedianPruner(n_startup_trials=N_STARTUP_TRIALS, n_warmup_steps=N_EVALUATIONS // 3)
+
+    study = optuna.create_study(sampler=sampler, pruner=pruner, direction="maximize")
+    try:
+        study.optimize(objective, n_trials=N_TRIALS, n_jobs=N_JOBS, timeout=600)
+    except KeyboardInterrupt:
+        pass
+
+    print("Number of finished trials: ", len(study.trials))
+
+    print("Best trial:")
+    trial = study.best_trial
+
+    print("  Value: ", trial.value)
+
+    print("  Params: ")
+    for key, value in trial.params.items():
+        print("    {}: {}".format(key, value))
+
+    print("  User attrs:")
+    for key, value in trial.user_attrs.items():
+        print("    {}: {}".format(key, value))
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -84,6 +84,7 @@ def get_extras_require() -> Dict[str, List[str]]:
             "torchvision==0.5.0" if sys.platform == "darwin" else "torchvision==0.5.0+cpu",
             "xgboost",
         ]
+        + (["stable-baselines3>=0.7.0"] if (3, 5) < sys.version_info[:2] else [])
         + (
             ["allennlp<1", "fastai<2", "pytorch_lightning>=0.7.1"]
             if (3, 5) < sys.version_info[:2] < (3, 8)


Test patch: