強化学習 keras-rl2からstable-baselins3への乗り換え

はじめに

久々にkeras-rl2を使用して強化学習を行ったのですが昔作成したプログラムが動かなくなっていました。tensorflowのバージョンアップが原因のようです。とりあえず動くようには修正したのですが、keras-rl2は2年くらいメンテナンスされておらず、これを機にstable baselinesに移行することにしました。
私の環境は下記のとおりです。

• macOS: 14.0（23A344）
• Python: 3.11.5
• gym: 0.26.2
• tensorflow: 2.14.0
• keras: 2.14.0
• keras-rl2: 1.0.5

1. 確認のために使用したプログラム

下記で使用した拙作のプログラムを使用します。少し変更しているので再掲します。

環境

envs/moving.py

import time
import numpy as np
import gym


class MovingEnv(gym.Env):
    ACTION_NUM = 4
    MIN_POS = 0
    MAX_POS = 99
    MAX_STEPS = 100
    metadata = {"render.modes": ["human"]}

    def __init__(self):
        super().__init__()
        self.pos = None
        self.goal = None
        self.steps = 0
        self.action_space = gym.spaces.Discrete(self.ACTION_NUM)
        self.observation_space = gym.spaces.Box(
            low=np.array([self.MIN_POS]),
            high=np.array([self.MAX_POS]),
            dtype=np.int16,
        )

    def reset(self):
        self.pos = np.random.randint(self.MIN_POS, self.MAX_POS + 1)
        self.goal = np.random.randint(self.MIN_POS, self.MAX_POS + 1)
        self.steps = 0
        return np.array([self.goal - self.pos])

    def step(self, action):
        """
        action
        0: -10移動
        1: -1移動
        2: +1移動
        3: +10移動
        """
        if action == 0:
            next_pos = self.pos - 10
        elif action == 1:
            next_pos = self.pos - 1
        elif action == 2:
            next_pos = self.pos + 1
        elif action == 3:
            next_pos = self.pos + 10
        else:
            next_pos = self.pos

        if next_pos < self.MIN_POS:
            next_pos = self.MIN_POS
        elif next_pos > self.MAX_POS:
            next_pos = self.MAX_POS

        self.pos = next_pos
        self.steps += 1
        reward = 100 if self.pos == self.goal else -1
        done = (
            True
            if self.pos == self.goal or self.steps > self.MAX_STEPS
            else False
        )

        return np.array([self.goal - self.pos]), reward, done, {}

    def seed(self, seed=0):
        return seed

    def render(self, mode="human"):
        a = ["." for x in range(self.MIN_POS, self.MAX_POS + 1)]
        a[self.goal] = "G"
        a[self.pos] = "o"
        print(f'\r{"".join(a)}', end="")
        time.sleep(0.1)

DQNAgent作成

keras_dqn.py

from tensorflow import keras
from rl.agents.dqn import DQNAgent
from rl.policy import BoltzmannQPolicy
from rl.memory import SequentialMemory


def make_dqn(env, nb_actions):
    model = keras.models.Sequential(
        [
            keras.layers.Flatten(
                input_shape=(1,) + env.observation_space.shape
            ),
            keras.layers.Dense(16, activation="relu"),
            keras.layers.Dense(16, activation="relu"),
            keras.layers.Dense(nb_actions, activation="linear"),
        ]
    )

    memory = SequentialMemory(limit=50000, window_length=1)
    policy = BoltzmannQPolicy()
    dqn = DQNAgent(
        model=model,
        nb_actions=nb_actions,
        memory=memory,
        target_model_update=1e-2,
        policy=policy,
    )

    dqn.compile(keras.optimizers.legacy.Adam(lr=1e-3), metrics=["mae"])

    return dqn

トレーニング

keras_dqn_train.py

import gym
from tensorflow import keras
from keras_dqn import make_dqn

proj_name = "keras_dqn"


def main():
    gym.envs.registration.register(
        id="moving-v0",
        entry_point="envs.moving:MovingEnv",
    )
    env = gym.make("moving-v0")

    dqn = make_dqn(env, env.action_space.n)
    csv_logger = keras.callbacks.CSVLogger(f"{proj_name}.csv")

    dqn.fit(
        env,
        nb_steps=500000,
        visualize=False,
        verbose=0,
        callbacks=[csv_logger],
    )

    dqn.save_weights(f"{proj_name}.hdf5")


if __name__ == "__main__":
    main()

テスト

keras_dqn_test.py

import gym
from keras_dqn import make_dqn

proj_name = "keras_dqn"


def main():
    gym.envs.registration.register(
        id="moving-v0",
        entry_point="envs.moving:MovingEnv",
    )
    env = gym.make("moving-v0")

    dqn = make_dqn(env, env.action_space.n)
    dqn.load_weights(f"{proj_name}.hdf5")
    dqn.test(env, nb_episodes=10)


if __name__ == "__main__":
    main()

2. keras-rl2が動くようにする

2.1 keras-rl2でエラーが発生

tensorflow-kerasのVERSIONをインポートしようとしていますが見つからないようです。

Traceback (most recent call last):
  File "/Users/myname/rl/keras_dqn_train.py", line 3, in <module>
    from keras_dqn import make_dqn
  File "/Users/myname/rl/keras_dqn.py", line 2, in <module>
    from rl.agents.dqn import DQNAgent
  File "/Users/myname/rl/venv/lib/python3.11/site-packages/rl/agents/__init__.py", line 1, in <module>
    from .dqn import DQNAgent, NAFAgent, ContinuousDQNAgent
  File "/Users/myname/rl/venv/lib/python3.11/site-packages/rl/agents/dqn.py", line 7, in <module>
    from rl.core import Agent
  File "/Users/myname/rl/venv/lib/python3.11/site-packages/rl/core.py", line 7, in <module>
    from rl.callbacks import (
  File "/Users/myname/rl/venv/lib/python3.11/site-packages/rl/callbacks.py", line 8, in <module>
    from tensorflow.keras import __version__ as KERAS_VERSION
ImportError: cannot import name '__version__' from 'tensorflow.keras' (/Users/myname/rl/venv/lib/python3.11/site-packages/keras/api/_v2/keras/__init__.py)

応急処置として、keras-rl2のソースを修正しました。修正したのは下記の2箇所です。

callbacks.py (修正前)

  8: from tensorflow.keras import __version__ as KERAS_VERSION

callbacks.py (修正後)

  8: # from tensorflow.keras import __version__ as KERAS_VERSION

callbacks.py (修正前)

266:        if KERAS_VERSION > '2.1.3':
267:        self.progbar.update((self.step % self.interval) + 1, values=values)
268:        else:
269:            self.progbar.update((self.step % self.interval) + 1, values=values, force=True)

callbacks.py (修正後)

266:        # if KERAS_VERSION > '2.1.3':
267:        self.progbar.update((self.step % self.interval) + 1, values=values)
268:        # else:
269:        #     self.progbar.update((self.step % self.interval) + 1, values=values, force=True)

2.2 Adamに_nameという変数がない

2023-10-09 13:28:35.970778: W tensorflow/c/c_api.cc:305] Operation '{name:'dense_2_1/bias/Assign' id:154 op device:{requested: '', assigned: ''} def:{{{node dense_2_1/bias/Assign}} = AssignVariableOp[_has_manual_control_dependencies=true, dtype=DT_FLOAT, validate_shape=false](dense_2_1/bias, dense_2_1/bias/Initializer/zeros)}}' was changed by setting attribute after it was run by a session. This mutation will have no effect, and will trigger an error in the future. Either don't modify nodes after running them or create a new session.
Traceback (most recent call last):
  File "/Users/myname/rl/keras_dqn_train.py", line 30, in <module>
    main()
  File "/Users/myname/rl/keras_dqn_train.py", line 15, in main
    dqn = make_dqn(env, env.action_space.n)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/myname/rl/keras_dqn.py", line 29, in make_dqn
    dqn.compile(keras.optimizers.Adam(lr=1e-3), metrics=["mae"])
  File "/Users/myname/rl/venv/lib/python3.11/site-packages/rl/agents/dqn.py", line 175, in compile
    optimizer = AdditionalUpdatesOptimizer(optimizer, updates)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/myname/rl/venv/lib/python3.11/site-packages/rl/util.py", line 86, in __init__
    super().__init__(optimizer._name)

kerasが新しくなって、Adamから_nameという変数がなくなったようです。legacyにあるAdamを使用すると昔のAdamが使用できます。

(修正前)

    dqn.compile(keras.optimizers.Adam(lr=1e-3), metrics=["mae"])

(修正後)

    dqn.compile(keras.optimizers.legacy.Adam(lr=1e-3), metrics=["mae"])

学習させてみましたが、昔のようにうまく学習してくれません。上記の修正だけでは足りないのかも知れません。深追いはせず、stable baselinesを使用することにします。

3. stable baselinesを使ってみる

keras-rl2の次に何を使おうかと悩んだのですが、とりあえずstable baselinesというのを使用してみることにしました。

3.1 インストール

pipでインストールします。

pip install "stable-baselines3[extra]"

3.2 トレーニング

keras-rl2と比べるとだいぶ短いです。ベクトル型並行処理が効いているのかも知れません。

stable_baselines_train.py

import gym
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3 import PPO

proj_name = "stable_baselines_dqn"


def main():
    gym.envs.registration.register(
        id="moving-v0",
        entry_point="envs.moving:MovingEnv",
    )
    env = gym.make("moving-v0")

    # ベクトル型並行処理ができるので使ってみます
    env = DummyVecEnv([lambda: env])

    # 強化学習アルゴリズムとしてPPOを使用します
    # イメージデータなどを扱う際はCnnPolicyを使用するようです。ここではMlpPolicyを使用します
    model = PPO("MlpPolicy", env, verbose=1, tensorboard_log="./logs")

    model.learn(total_timesteps=100000)

    env.close()

    model.save(proj_name)


if __name__ == "__main__":
    main()

下記を実行するとTensorBoardでトレーンぐ状況を見ることができます。

tensorboard --logdir=./logs/

3.3 テスト

トレーニングで保存したモデルをロードすることでモデルを構築できます。

stable_baselines_test.py

import gym
from stable_baselines3 import PPO

proj_name = "stable_baselines_dqn"


def main():
    gym.envs.registration.register(
        id="moving-v0",
        entry_point="envs.moving:MovingEnv",
    )
    env = gym.make("moving-v0")

    model = PPO.load(proj_name)

    for _ in range(10):
        state = env.reset()
        while True:
            action, _ = model.predict(state, deterministic=True)
            state, rewards, done, info = env.step(action)
            env.render()
            if done:
                print(f"steps:{env.steps}")
                break

    env.close()


if __name__ == "__main__":
    main()

4. さいごに

これまでkeras-rl2を愛用していましたが、stable baselinesに乗り換えることにしました。使用しているフレームワークを変更するとのは勇気がいることです。stable baselinesはトレーニング時間がkeras-rl2よりだいぶ短いです。ベクトル型並行処理が効いているのかも知れません。ただGPUの使用率を見るとぜんぜん使っていないように見えます。GPUを使用するための設定が必要なのかも知れません。まずは一歩を踏み出したということで、これからいろいろ調べていきたいと思います。