自作中のシンプルな分散強化学習フレームワークの紹介

今までいろいろな強化学習アルゴリズムを実装してきましたが、以下の問題を毎回感じていました。

1. アルゴリズム毎に環境と学習のための実装が必要
   　(gym等は環境側のインタフェースまでは提供されているが、学習を含めた実装は提供されていない)
2. アルゴリズムの学習フローがそもそも複雑になりやすい（デバッグしづらい）
3. 分散学習まで考えるとなお大変
4. アルゴリズムと環境が分離しきれず結構依存する
   1. Atariの画像の前処理とその入力に依存したニューラルネットワークの設計
   2. アクションが離散値と連続値の場合によるアルゴリズム側の違い
   3. よくある数フレーム入力やフレームスキップ、実行できないアクション等を加味した場合の変更
   4. 学習フローの差異（AlphaZeroやモデルベース強化学習など）
5. ログや可視化等の作成

既存のフレームワークを試したりしましたが自作アルゴリズムを作ろうとした場合、フレームワークの理解から勉強が必要になりかなりハードルが高いような気がしました。
そこでなるべくハードルが低くなるようなフレームワークを自作しようと思い作成しました。

このフレームワークの特徴は以下です。

1. カスタマイズ可能な環境の提供
2. カスタマイズ可能な強化学習アルゴリズムの提供
3. 環境とアルゴリズム間のインタフェースの自動調整
4. 分散強化学習のサポート
5. 有名な強化学習アルゴリズムの提供
6. （新しいアルゴリズムへの対応）

本記事はフレームワークの説明をすると共に実装しているアルゴリズムも説明していこうかと思います。
各項目については別記事にして少しずつ作成する予定です。

Github

各項目

• 概要（この記事）
• 最初に
  • 強化学習の基礎
• フレームワークの使い方
  • オリジナル環境の作成方法
  • オリジナルアルゴリズムの作成方法
  • KaggleのConnectXへの投稿方法を解説（スコア900越えコード付き）
  • クラウドサービスを利用した分散強化学習（無料編）
  • クラウドサービスを利用した分散強化学習（kubernetes編）
  • クラウドサービスを利用した分散強化学習（GKE/有料編）
  • エミュレータ連携（Gym×BizHawk×強化学習）
• 技術関係
  • 分散強化学習について
  • 部分観測マルコフ決定過程について
  • Off-Policyにおける安全な学習方法(Retrace)
• 価値ベース
  • 価値ベースアルゴリズムの基礎（Q学習）
  • DQN
  • Rainbow
  • R2D2
  • Agent57
• 方策ベース
  • 方策ベースアルゴリズムの基礎（方策勾配法/REINFORCE）
  • TRPO(過去記事、フレームワーク未実装)
  • PPO
  • DDPG/TD3(過去記事、フレームワーク実装あり)
  • SAC(過去記事、フレームワーク実装あり)
• Alphaシリーズ
  • モンテカルロ木探索(MCTS)
  • AlphaZero
  • MuZero
  • StochasticMuZero
• WorldModelsシリーズ
  • WorldModels
  • PlaNet
  • Dreamer
  • DreamerV2
  • DreamerV3
• その他の強化学習アルゴリズム
  • Go-Explore
  • SND
• オリジナル
  • SearchDynaQ（テーブルベースアルゴリズム）
• 最適化アルゴリズム
  • 最適化アルゴリズムを実装していくぞ（概要）
• フレームワークとは関係ない機械学習関係の記事
  • 重み付き乱択アルゴリズムをまとめて計測してみた
  • Softmax+多クラス分類の感覚をつかむために基礎からちゃんと調べてみた（交差エントロピー損失）
  • BitNet b1.58(BitLinear)を実装してMNISTで検証してみた(Tensorflow/Torch)
  • 大規模言語モデルと強化学習：強化学習にLLMを組み込んで実装してみた（ローカルLLM）
• Tensorflow/ディープラーニング関係
  • 0から理解する誤差逆伝播法（理論からコードの実装まで）
  • 基本的な出力層の活性化関数と損失関数の組み合わせまとめ（TF2.0）
  • LSTM(RNN)で可変長な時系列と隠れ状態について調べてみた(Tensorflow2.0)
  • 乱数の固定方法(random/numpy/tensorflow/gym)
  • 混合密度ネットワーク（Tensorflow2.0）
  • tf.GradientTapeで正則化項が反映されているか調査してみた（過学習の公式チュートリアルを添えて）
  • 離散空間のReparameterizationTrick（GumbelMaxTrick/StraightGradients/VQ-VAE）
  • ELBO（変分下限）を最大化するという意味をゼロから理解する（尤度、最尤推定）
• その他
  • 中心極限定理が本当か試してみた
  • フォルダでほぼ完結する開発環境を0から整えて実装してみた（Python+VSCocde）
  • 機械学習用の環境を0から構築した(windows11 + VSCode + venv + Tensorflow + GPU)
  • Python3.11 + TF2.12 + GPUの環境を構築した(Windows11 + WSL + VSCode)

1.インストール

Github からの pip install を想定しています。(PyPIへの登録予定は今のところありません)
また、ダウンロードだけでも使えます。詳細はGitHubを見てください。

インストールする一例は以下です。

git clone https://github.com/pocokhc/simple_distributed_rl.git
cd simple_distributed_rl
pip install .

2.使い方

import numpy as np
import srl

# ---  1. 使うアルゴリズムをimport
from srl.algorithms import ql

def main():

    # --- 2. 各configを設定
    env_config = srl.EnvConfig("Grid")  # 環境config
    rl_config = ql.Config()     # アルゴリズムconfig

    # --- 3. Runnerを作成
    runner = srl.Runner(env_config, rl_config)

    # --- 4. 学習
    if True:
        # 逐次学習の場合
        runner.train(timeout=20)
    else:
        # 分散学習の場合
        runner.train_mp(timeout=20)
    
    # 評価
    rewards = runner.evaluate(max_episodes=10)
    print(f"10回テストした結果の平均報酬: {np.mean(rewards)}")

    # 可視化
    # Terminal上で1エピソードを見る
    runner.render_terminal()

    # 1エピソードをアニメーション化('opencv-python','pillow','matplotlib','pygame'が別途必要)
    runner.animation_save_gif("_Grid.gif")

if __name__ == '__main__':
    main()

• アニメーション

• Terminal上での結果

10回テストした結果の平均報酬: 0.7960000045597553
### 0, action 3, rewards [0.], next 0
env   None
work0 None
......
.   G.
. . X.
.P   .
......


 ←  : 0.10151
 ↓  : 0.07913
 →  : 0.02751
*↑  : 0.36339

### 1, action 3, rewards [-0.04], next 0
env   {}
work0 {}
......
.   G.
.P. X.
.    .
......


 ←  : 0.17283
 ↓  : 0.11623
 →  : 0.17698
*↑  : 0.48977

### 2, action 2, rewards [-0.04], next 0
env   {}
work0 {}
......
.P  G.
. . X.
.    .
......


 ←  : 0.23222
 ↓  : 0.18907
*→  : 0.61112
 ↑  : 0.24942

### 3, action 2, rewards [-0.04], next 0
env   {}
work0 {}
......
. P G.
. . X.
.    .
......


 ←  : 0.25572
 ↓  : 0.31604
*→  : 0.76530
 ↑  : 0.32385

### 4, action 2, rewards [-0.04], next 0
env   {}
work0 {}
......
.  PG.
. . X.
.    .
......


 ←  : 0.32941
 ↓  : 0.26414
*→  : 0.92709
 ↑  : 0.44327

### 5, action 2, rewards [-0.04], next 0
env   {}
work0 {}
......
.  PG.
. . X.
.    .
......


 ←  : 0.32941
 ↓  : 0.26414
*→  : 0.92709
 ↑  : 0.44327

### 6, action 2, rewards [1.], done(env), next 0
env   {}
work0 {}
......
.   P.
. . X.
.    .
......

おわりに

少しずつ開発しているので温かく見守っていただければ幸いです。