概要
Open AIのBaselinesの強化学習を扱い安くしたしたライブラリです。Baselinesは研究として利用する上では問題なく利用できますが、例えば画像から特徴量を学習するときなどはコメントアウトされている部分をコメント消すような修正が必要となったりと地獄を見ることがあります。こちらのStable Baselinesは地獄に落ちないためのライブラリとなります。（ベイマックスがかわいいです）
最新のPytorchを利用したStable Baselines3がこちら。
学習/推論
例えばPPOの実装はこのような簡単なコードで学習と推論まで実装することができます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
# from https://github.com/hill-a/stable-baselines import gym from stable_baselines.common.policies import MlpPolicy from stable_baselines import PPO2 env = gym.make('CartPole-v1') model = PPO2(MlpPolicy, env, verbose=1) # Train the agent model.learn(total_timesteps=10000) # Enjoy trained agent obs = env.reset() for i in range(1000): action, _states = model.predict(obs, deterministic=False) obs, reward, done, info = env.step(action) env.render() if done: obs = env.reset() env.close() |
評価
以下がtutorialを参照した評価関数です。env.stepを実行して得たrewardの平均で評価できるとわかります。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
def evaluate(model, num_episodes=100): """ Evaluate a RL agent :param model: (BaseRLModel object) the RL Agent :param num_episodes: (int) number of episodes to evaluate it :return: (float) Mean reward for the last num_episodes """ # This function will only work for a single Environment env = model.get_env() all_episode_rewards = [] for i in range(num_episodes): episode_rewards = [] done = False obs = env.reset() while not done: # _states are only useful when using LSTM policies action, _states = model.predict(obs) # here, action, rewards and dones are arrays # because we are using vectorized env obs, reward, done, info = env.step(action) episode_rewards.append(reward) all_episode_rewards.append(sum(episode_rewards)) mean_episode_reward = np.mean(all_episode_rewards) print("Mean reward:", mean_episode_reward, "Num episodes:", num_episodes) return mean_episode_reward |
このような関数を自作しなくてもstable baselinesに評価関数evaluate_policyが既に存在します。学習済みのmodelを入れることでmodelを評価可能です。
|
1 2 3 4 5 |
from stable_baselines3.common.evaluation import evaluate_policy mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=100) print(f"mean_reward:{mean_reward:.2f} +/- {std_reward:.2f}") |
モデルの保存/読み込み
医科のように簡単に学習済みモデルの保存と読み込みができます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import os from stable_baselines3.common.vec_env import DummyVecEnv # Create save dir save_dir = "/tmp/gym/" os.makedirs(save_dir, exist_ok=True) model = A2C('MlpPolicy', 'Pendulum-v0', verbose=0, gamma=0.9, n_steps=20).learn(8000) # The model will be saved under A2C_tutorial.zip model.save(save_dir + "/A2C_tutorial") del model # delete trained model to demonstrate loading # load the model, and when loading set verbose to 1 loaded_model = A2C.load(save_dir + "/A2C_tutorial", verbose=1) # show the save hyperparameters print("loaded:", "gamma =", loaded_model.gamma, "n_steps =", loaded_model.n_steps) # as the environment is not serializable, we need to set a new instance of the environment loaded_model.set_env(DummyVecEnv([lambda: gym.make('Pendulum-v0')])) # and continue training loaded_model.learn(8000) |
独自環境の生成
Gymによって生成された環境も簡単にカスタムすることができます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
class CustomWrapper(gym.Wrapper): """ :param env: (gym.Env) Gym environment that will be wrapped """ def __init__(self, env): # Call the parent constructor, so we can access self.env later super(CustomWrapper, self).__init__(env) def reset(self): """ Reset the environment """ obs = self.env.reset() return obs def step(self, action): """ :param action: ([float] or int) Action taken by the agent :return: (np.ndarray, float, bool, dict) observation, reward, is the episode over?, additional informations """ obs, reward, done, info = self.env.step(action) return obs, reward, done, info |
学習データと推論データ
実際に利用する時に少し戸惑うのは、学習データはenvで学習したとしてどこに推論データを入れればいいのかというところです。上述のコードでは環境に利用するデータは同じで推論前に一度resetしています。この部分を推論対象のenvにすり替えれば学習ができるはずです。
|
1 2 |
obs = env.reset() |
本ページのコードは学習と推論を同じファイルに記載していますが、実際は学習/モデル保存とモデル読み込み/推論で別のファイルに作ればよいかなと思います。