Stable baselines 3

Training

In this example, we showcase how the built-in stable baselines 3 (sb3) extension can be used to smoothly leverage the large associated collection of algorithm and policies.

import math
import pathlib
import random
import shutil

from stable_baselines3.common.callbacks import (
    EvalCallback,
    StopTrainingOnRewardThreshold,
)
from stable_baselines3.common.logger import configure

from amaze import Maze, Robot, Simulation, Sign, amaze_main, StartLocation
from amaze.extensions.sb3 import (
    make_vec_maze_env,
    env_method,
    load_sb3_controller,
    PPO,
    TensorboardCallback,
    sb3_controller,
    CV2QTGuard,
)

As usual, we start by importing the necessary packages and we define some global configuration options. Note that, in addition to the traditional amaze classes, we also import extension-specific items (detailed below).

    train_mazes = Maze.BuildData.from_string(train_maze).all_rotations()
    eval_mazes = [d.where(seed=test_seed) for d in train_mazes]

    train_env = make_vec_maze_env(train_mazes, robot, SEED)
    eval_env = make_vec_maze_env(eval_mazes, robot, SEED, log_trajectory=True)

The training function is defined much more shortly than in the hand-written q-learning case thanks to the added functionalities of stable baselines 3 and added wrappers. While, creating mazes and robots should be familiar by now, we see a new extension-specific function make_vec_maze_env() used to create Vectorized Environments (VecEnv)

    optimal_reward = sum(env_method(eval_env, "optimal_reward")) / len(eval_mazes)

We also, sometimes, need access to the underlying environments (regular mazes) as illustrated below. There we collect the average optimal reward by calling optimal_reward() on every maze used for intermediate performance evaluation thanks to env_method().

    tb_callback = TensorboardCallback(
        log_trajectory_every=1, max_timestep=BUDGET  # Eval callback (below)
    )
    eval_callback = EvalCallback(
        eval_env,
        best_model_save_path=FOLDER,
        log_path=FOLDER,
        eval_freq=max(100, BUDGET // (10 * len(eval_mazes))),
        verbose=1,
        n_eval_episodes=len(eval_mazes),
        callback_after_eval=tb_callback,
        callback_on_new_best=StopTrainingOnRewardThreshold(
            reward_threshold=optimal_reward, verbose=1
        ),
    )

Next we create a TensorboardCallback, an illustrative built-in callback that uses Tensorboard to provide an overview of the training process. In addition to logging numerical data such as the average rewards it also automatically generates trajectory images whenever the EventCallback is triggered. The following lines define such an object, in a traditional SB3 fashion, while adding our own tensorboard callback and also using the optimal reward to stop as soon as the agent is behaving optimally.

    model = sb3_controller(
        robot,
        PPO,
        policy="MlpPolicy",
        env=train_env,
        seed=SEED,
        learning_rate=1e-3,
    )

    print("== Starting", "=" * 68)
    model.set_logger(configure(FOLDER, ["csv", "tensorboard"]))
    model.learn(BUDGET, callback=eval_callback, progress_bar=not is_test)

    tb_callback.log_step(final=True)

Finally, we create the sb3 model, using the dedicated wrapper sb3_controller(), by providing the robot data and the of underlying model type (one of compatible_models()) and, afterwards, the usual parameters. Then after setting up the logger and letting the training process run its course, we perform a final step of the callback to render the final trajectories.

Using

def evaluate(is_test):
    model = load_sb3_controller(BEST)

Once the training process is complete, we evaluate the resulting agent’s generalization capability in the same manner as in Full example: Training. The only difference is the use of the dedicated loading function load_sb3_controller() which is a verbose alias to load(). The reminder of this function being the same, we refer the reader to the previous example (Generalization), if needed.

Full listing for examples/extensions/sb3:main

def main(is_test=False):
    folder = pathlib.Path(FOLDER)
    if folder.exists():
        shutil.rmtree(folder)
    folder.mkdir(parents=True, exist_ok=False)

    train_maze = train(is_test)
    evaluate(is_test)

    with CV2QTGuard(platform=False):
        amaze_main(
            f"--controller {BEST} --extension sb3 --maze {train_maze}"
            f" --auto-quit --no-restore-config --width 1000"
        )

Finally, the main should also be familiar from the previous example. One thing to note, however, is that, due to incompatibilities between the current opencv and PyQT5 libraries, one should use CV2QTGuard when combining stable baselines 3 with the native Qt5 components.