Stable Baselines チュートリアル(2) / Gymラッパー、モデルの保存と読み込み

以下のColabが面白かったので、ざっくり訳してみました。

・Stable Baselines Tutorial - Gym wrappers, saving and loading models

1. はじめに

このノートブックでは、監視、正規化、ステップ数の制限、機能強化などを行うことができる「Gymラッパー」の使い方を学習します。また、モデルの保存と読み込み、およびエクスポート可能な出力ファイルの読み取り方法も確認します。

2. pipを使用して依存関係とStable Baselinesをインストール

Colabでのインストールコマンドは次の通りです。

!apt install swig
!pip install stable-baselines[mpi]==2.8.0

import gym
from stable_baselines import A2C, SAC, PPO2, TD3

3. モデルの保存と読み込み

「Stable Baselines」のモデルの保存と読み込みは簡単です。モデルのsave()とload()を直接呼び出すことができます。

import os

# 保存フォルダの生成
save_dir = "/tmp/gym/"
os.makedirs(save_dir, exist_ok=True)

model = PPO2('MlpPolicy', 'Pendulum-v0', verbose=0).learn(8000)

# モデルはPPO2_tutorial.zipの下に保存
model.save(save_dir + "/PPO2_tutorial")

# 環境からの行動をサンプリング
action = model.env.observation_space.sample()

# 保存する前に予測を確認
print("pre saved", model.predict(action, deterministic=True))

del model # 訓練されたモデルを削除して読み込みを実証

# 読み込み後に予測が同じであることを確認
loaded_model = PPO2.load(save_dir + "/PPO2_tutorial")
print("loaded", loaded_model.predict(action, deterministic=True))

「Stable Baselines」では、現在の「重み」と「訓練のハイパーパラメータ」を保存するため、非常に強力です。

つまり、実際には、ハイパーパラメータを再定義せずに、カスタムモデルを読み込むだけで、学習を続けることができます。読み込み時にモデルのクラス変数を更新することもできます。

import os
from stable_baselines.common.vec_env import DummyVecEnv

# 保存フォルダの生成
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)

# モデルはA2C_tutorial.zipの下に保存
model.save(save_dir + "/A2C_tutorial")

del model # 訓練されたモデルを削除して読み込みを実証


# モデルのロード時に詳細レベルを1に設定
loaded_model = A2C.load(save_dir + "/A2C_tutorial", verbose=1)

# 保存したハイパーパラメータを表示
print("loaded:", "gamma =", loaded_model.gamma, "n_steps =", loaded_model.n_steps)

# 環境はシリアル化できないため、新しいインスタンスを生成する必要がある
loaded_model.set_env(DummyVecEnv([lambda: gym.make('Pendulum-v0')]))

# 訓練を再開
loaded_model.learn(8000)

4. GymとVecEnvラッパー

◎Gymラッパーの構造
Gymラッパーは、Gymインターフェースに従います。reset()およびstep()があります。ラッパーは環境をラップするため、self.envを使用して各種情報にアクセスできます。これにより、元のenvを変更せずに簡単に対話できます。

事前定義された多くのラッパーがあります。完全なリストについては、Gymのドキュメントを参照してください。

class CustomWrapper(gym.Wrapper):
  """
  :param env: (gym.Env) Gym environment that will be wrapped
  """
  def __init__(self, env):ラップされるGym環境
    # 親コンストラクタを呼び出して、後でself.envにアクセスできるようにする
    super(CustomWrapper, self).__init__(env)

  def reset(self):
    """
    環境のリセット
    """
    obs = self.env.reset()
    return obs

  def step(self, action):
    """
    :param action: ([float] or int) エージェントが実行した行動
    :return: (np.ndarray, float, bool, dict) 観察、報酬、エピソード完了、情報
    """
    obs, reward, done, info = self.env.step(action)
    return obs, reward, done, info

◎最初の例：エピソードの長さを制限
ラッパーの実用的な使用例の1つは、エピソードごとにステップ数を制限する場合です。そのため、制限に達したときにエピソード完了を上書きする必要があります。また、情報辞書(step()の戻り値の情報)でその情報を渡すことをお勧めします。

class TimeLimitWrapper(gym.Wrapper):
  """
  :param env: (gym.Env) ラップされるGym環境
  :param max_steps: (int) エピソードごとの最大ステップ数
  """
  def __init__(self, env, max_steps=100):
    # Call the parent constructor, so we can access self.env later
    super(TimeLimitWrapper, self).__init__(env)
    self.max_steps = max_steps
    # Counter of steps per episode
    self.current_step = 0

  def reset(self):
    """
    環境のリセット
    """
    # Reset the counter
    self.current_step = 0
    return self.env.reset()

  def step(self, action):
    """
    :param action: ([float] or int) エージェントが実行した行動
    :return: (np.ndarray, float, bool, dict) 観察、報酬、エピソード完了、情報
    """
    self.current_step += 1
    obs, reward, done, info = self.env.step(action)

    # エピソード完了を上書きする場合
    if self.current_step >= self.max_steps:
      done = True
      # 情報辞書を更新して、制限を超えたことを通知
      info['time_limit_reached'] = True
    return obs, reward, done, info

ラッパーをテストします。

from gym.envs.classic_control.pendulum import PendulumEnv

# gym.make()が既に環境をTimeLimitラッパーでラップされているため、環境を直接作成
env = PendulumEnv()

# 環境をラップ
env = TimeLimitWrapper(env, max_steps=100)

obs = env.reset()
done = False
n_steps = 0
while not done:
  # ランダムな行動を採る
  random_action = env.action_space.sample()
  obs, reward, done, info = env.step(random_action)
  n_steps += 1

print(n_steps, info)

実際にはGymには、ほとんどの環境で使用される「TimeLimit」（gym.wrappers.TimeLimit）という名前のラッパーが既にあります。

◎2番目の例 : 行動の正規化
通常は、観察と行動をエージェントに渡す前に「正規化」することをお勧めします。これにより、デバッグしにくい問題を防ぐことができます。

この例では、Pendulum-v0の行動空間を「正規化」して、[-2、2]ではなく[-1、1]になるようにします。

注：ここでは、連続行動、つまりgym.Box空間を扱っています。

import numpy as np

class NormalizeActionWrapper(gym.Wrapper):
  """
  :param env: (gym.Env) ラップされるGym環境
  """
  def __init__(self, env):
    # 行動空間の取得
    action_space = env.action_space
    assert isinstance(action_space, gym.spaces.Box), "This wrapper only works with continuous action space (spaces.Box)"

    # 最大/最小値の取得
    self.low, self.high = action_space.low, action_space.high

    # 行動空間を変更したため、すべての行動は[-1、1]にある
    env.action_space = gym.spaces.Box(low=-1, high=1, shape=action_space.shape, dtype=np.float32)

    # 親コンストラクタを呼び出して、後でself.envにアクセスできるようにする
    super(NormalizeActionWrapper, self).__init__(env)

  def rescale_action(self, scaled_action):
      """
      行動を[-1, 1]から[low, high]に再スケーリング
      (対称行動空間は不要）
      :param scaled_action: (np.ndarray)
      :return: (np.ndarray)
      """
      return self.low + (0.5 * (scaled_action + 1.0) * (self.high -  self.low))

  def reset(self):
    """
    環境のリセット
    """
    # カウンタのリセット
    return self.env.reset()

  def step(self, action):
    """
    :param action: ([float] or int) エージェントが採った行動
    :return: (np.ndarray, float, bool, dict) 観察、報酬、エピソード完了、情報
    """
    # 行動を[-1, 1]から元の[low, high]に再スケーリング
    rescaled_action = self.rescale_action(action)
    obs, reward, done, info = self.env.step(rescaled_action)
    return obs, reward, done, info

正規化前のランダム行動は次の通りです。

original_env = gym.make("Pendulum-v0")

print(original_env.action_space.low)
for _ in range(10):
  print(original_env.action_space.sample())

[-2.]
[-0.5033201]
[-0.14569831]
[-0.8894852]
[0.3471374]
[1.4554224]
[-1.5298725]
[0.06951643]
[-1.4717276]
[0.86743873]
[-0.4157612]

正規化後のランダム行動は次の通りです。

env = NormalizeActionWrapper(gym.make("Pendulum-v0"))

print(env.action_space.low)

for _ in range(10):
  print(env.action_space.sample())

[-1.]
[0.13084263]
[-0.6334403]
[-0.7103045]
[-0.02388744]
[-0.28877452]
[0.8808639]
[0.5306505]
[0.49732724]
[0.8074395]
[-0.83315516]

◎RLアルゴリズムでテスト
「Stable Baselines」のMonitorラッパーを使用して、訓練統計（平均エピソード報酬、平均エピソード長）をモニターできるようにします。

from stable_baselines.bench import Monitor
from stable_baselines.common.vec_env import DummyVecEnv

env = Monitor(gym.make('Pendulum-v0'), filename=None, allow_early_resets=True)
env = DummyVecEnv([lambda: env])

model = A2C("MlpPolicy", env, verbose=1).learn(int(1000))

NormalizeActionWrapperを使用します。

normalized_env = Monitor(gym.make('Pendulum-v0'), filename=None, allow_early_resets=True)

# 複数のラッパーを連続して使用できます
normalized_env = NormalizeActionWrapper(normalized_env)
normalized_env = DummyVecEnv([lambda: normalized_env])

model_2 = A2C("MlpPolicy", normalized_env, verbose=1).learn(int(1000))

5. VecEnvラッパー

「OpenAI Gym」で「Gymラッパー」が提供されているように、「Stable Baselines」は」「VecEnvラッパー」を提供しています。「VecEnvラッパー」には、次のようなものがあります。

・VecNormalize：観測値とリターンを正規化するために移動平均と標準偏差を計算します。
・VecFrameStack：連続した複数の観測をスタックします。観測に時間を統合するのに便利です（例：Atariゲームの連続フレーム）。

詳細はドキュメントを参照してください。

注：「VecNormalizeラッパー」を使用する場合、実行中の平均値と標準値をモデルとともに保存する必要があります。そうしないと、エージェントを再度ロードしたときに適切な結果が得られません。「RL Zoo」を使用する場合、これは自動的に行われます。

from stable_baselines.common.vec_env import VecNormalize, VecFrameStack

env = DummyVecEnv([lambda: gym.make("Pendulum-v0")])
normalized_vec_env = VecNormalize(env)

obs = normalized_vec_env.reset()
for _ in range(10):
  action = [normalized_vec_env.action_space.sample()]
  obs, reward, _, _ = normalized_vec_env.step(action)
  print(obs, reward)

6. 【演習】 Monitorラッパーを所有するコード

ラッパーがどのように機能するか、何ができるかがわかったので、演習に入ります。ここでの目標は、訓練の進行状況を監視するラッパーを作成し、エピソードの報酬（1つのエピソードの報酬の合計）とエピソードの長さ（最後のエピソードのステップ数）の両方を保存することです。

エピソードが終了するたびに、情報辞書を使用してこれらの値を返します。

class MyMonitorWrapper(gym.Wrapper):
  """
  :param env: (gym.Env) ラップされるGym環境
  """
  def __init__(self, env):
    # 親コンストラクタを呼び出して、後でself.envにアクセスできるようにする
    super(MyMonitorWrapper, self).__init__(env)
    # === YOU CODE HERE ===#
    # エピソードの長さとエピソードの報酬を保存するために使用される変数を初期化

    # ====================== #

  def reset(self):
    """
    Reset the environment
    """
    obs = self.env.reset()
    # === YOU CODE HERE ===#
    # 変数をリセット

    # ====================== #
    return obs

  def step(self, action):
    """
    :param action: ([float] or int) エージェントが採った行動
    :return: (np.ndarray, float, bool, dict) 観察、報酬、エピソード完了、情報
    """
    obs, reward, done, info = self.env.step(action)
    # === YOU CODE HERE ===#
    # 現在のエピソードの報酬とエピソードの長さを更新

    # ====================== #

    if done:
      # === YOU CODE HERE ===#
      # エピソードの長さとエピソードの報酬を情報辞書に保存

      # ====================== #
    return obs, reward, done, info

ラッパーをテストします。

# LunarLanderを使用するには、box2d box2d-kengz（pip）とswig（apt-get）をインストールする必要がある
!pip install box2d box2d-kengz

env = gym.make("LunarLander-v2")
# === YOU CODE HERE ===#
# 環境をラップ

# 環境をリセット

# 環境内でランダム行動を採り、各エピソードの終了後に正しい値を返すことを確認

# ====================== #

7. 【ラッパーボーナス】 状態空間の変更：固定長のエピソードのラッパー

from gym.wrappers import TimeLimit

class TimeFeatureWrapper(gym.Wrapper):
    """
    固定長エピソードの観測スペースに残り時間を追加。
    https://arxiv.org/abs/1712.00378およびhttps://github.com/aravindr93/mjrl/issues/13を参照
    :param env: (gym.Env)
    :param max_steps: (int) エピソードの最大ステップ数
        TimeLimitオブジェクトにラップされていない場合。
    :param test_mode: (bool) テストモードでは、時間機能は一定で、ゼロに等しくなる。
        これにより、エージェントがこの機能に過剰に適合しなかったことを確認し、
        事前に定義された決定的な一連の行動を学習できる。
    """
    def __init__(self, env, max_steps=1000, test_mode=False):
        assert isinstance(env.observation_space, gym.spaces.Box)
        # 観察に時間特徴を追加
        low, high = env.observation_space.low, env.observation_space.high
        low, high= np.concatenate((low, [0])), np.concatenate((high, [1.]))
        env.observation_space = gym.spaces.Box(low=low, high=high, dtype=np.float32)

        super(TimeFeatureWrapper, self).__init__(env)

        if isinstance(env, TimeLimit):
            self._max_steps = env._max_episode_steps
        else:
            self._max_steps = max_steps
        self._current_step = 0
        self._test_mode = test_mode

    def reset(self):
        self._current_step = 0
        return self._get_obs(self.env.reset())

    def step(self, action):
        self._current_step += 1
        obs, reward, done, info = self.env.step(action)
        return self._get_obs(obs), reward, done, info

    def _get_obs(self, obs):
        """
        時間機能を現在の観察値に連結

        :param obs: (np.ndarray)
        :return: (np.ndarray)
        """
        # 残り時間はより一般的
        time_feature = 1 - (self._current_step / self._max_steps)
        if self._test_mode:
            time_feature = 1.0
        # オプション: concatenate [time_feature, time_feature ** 2]
        return np.concatenate((obs, [time_feature]))

8. モデルの保存ファイル 形式

モデルの保存ファイル形式は、「Stable-Baselines（> 2.7.0）」で最近改良されました。これは、zipアーカイブされた「JSONダンプ」と「NumPy配列のzipアーカイブ」です。

saved_model.zip/
├── data              クラスパラメータのJSONファイル（辞書）
├── parameter_list    モデルパラメータのJSONファイルとその順序（リスト）
├── parameters        numpy.savez（numpy配列のzipファイル） ...
   ├── ...           zipアーカイブであるため、このオブジェクトも開いて参照することが可能
       ├── ...

保存ファイル 形式を確認します。

# 保存フォルダの生成
save_dir = "/tmp/gym/"
os.makedirs(save_dir, exist_ok=True)

model = PPO2('MlpPolicy', 'Pendulum-v0', verbose=0).learn(8000)
model.save(save_dir + "/PPO2_tutorial")

!ls /tmp/gym/PPO2_tutorial*

/tmp/gym/PPO2_tutorial.zip

import zipfile

archive = zipfile.ZipFile("/tmp/gym/PPO2_tutorial.zip", 'r')
for f in archive.filelist:
  print(f.filename)

data
parameters
parameter_list

9. 保存したモデルのエクスポート

最後に、Tensorflow JSまたはJavaにエクスポートしたい人のための参考資料をいくつか紹介します。
https://stable-baselines.readthedocs.io/en/master/guide/export.html

10. 参照

・Github repo: https://github.com/araffin/rl-tutorial-jnrr19
・Stable-Baselines: https://github.com/hill-a/stable-baselines
・Documentation: https://stable-baselines.readthedocs.io/en/master/
・RL Baselines zoo: https://github.com/araffin/rl-baselines-zoo

OpenAI Gym / Baselines 深層学習・強化学習 人工知能プログラミング 実践入門