【祝】強化学習ライブラリTF2RL v1.0到達 ~CIやドキュメントサイトなど開発・利用環境の整備~

1. はじめに

以前紹介した友人が開発しているTensorFlow 2.x 向け強化学習ライブラリTF2RLが、諸々整備してバージョン1.0に到達しました🎉

(いつの間にかスターも300超えていてすごい！)

バージョン1.0到達以降も、まだまだ様々な強化学習アルゴリズムを追加しようと開発が進んでいます。（この記事を準備している間にも、v1.1.0が公開されてます。）

インストール方法や基本となる使い方は、公式ReadMeや、前の記事を読んでいただければと思うので、この記事では割愛します。

この記事では、私もお手伝いさせてもらって整備したアルゴリズム以外の部分について紹介します。

2. マルチプラットフォームテスト (PR 97)

GitHub Actions によって、Windows/macOS/Ubuntu のマルチプラットフォームで、push や pull requestの度にユニットテストを自動で走らせれるようになりました。

特に、Windowsの開発機を友人も私も持っておらず、なかなかサポートできていなかったのですが、このActionsによってWindows上での問題があぶり出されWindowsでも無事動作できるようになりました。
（Windowsでは、multiprocessing.Pool などを利用する際に渡すことができるのは、ネストしていないモジュール直下で定義されたクラス・関数だけみたいですね。なんでも fork()システムコールの代わりに利用している pickleのシリアライズが原因みたいですが、なんでそこに制約があるのかまでは理解できていません。）

同じコードを複数のTFとPythonのバージョンの関係もあって、下のように綺麗なマトリックス戦略になっていないのですが、もうちょっとうまく記述する方法はありませんかね。

.github/workflows/test.yml(抜粋)

    strategy:
      matrix:
        os: [ubuntu-latest, macos-latest, windows-latest]
        python: ['3.7', '3.8']
        TF: ['2.2', '2.3']
        include:
          - os: ubuntu-latest
            python: '3.7'
            TF: '2.0'
          - os: ubuntu-latest
            python: '3.7'
            TF: '2.1'
          - os: macos-latest
            python: '3.7'
            TF: '2.0'
          - os: macos-latest
            python: '3.7'
            TF: '2.1'
          - os: windows-latest
            python: '3.7'
            TF: '2.0'
          - os: windows-latest
            python: '3.7'
            TF: '2.1'

3. Super-Linterでコードの自動チェック (PR 102)

GitHub Super-Linterを利用し、元々手作業で行っていたLinterでのコードのルールチェックが自動化されました。

Super-Linterで躓いたのは、カスタム設定のファイルを .github/linters/ ディレクトリに置かないといけないことでした。 カスタム設定のファイル名自体は、Super-Linterの起動時のオプションで指定できますが、あくまで上記のディレクトリにおいていないと見つからずデフォルトの設定が使用されてしまいます。（最終的にわからなくて質問しました。）

4. Sphinx を利用したドキュメントサイト構築 (PR 107)

GitHub Actionsの中で、Sphinxを利用してドキュメントを自動生成しています。

docstring から自動でクラスリファレンスを生成する autodoc とMarkdown形式文書を取り込む recommonmark を有効にしています。(テーマはRead The Docsテーマとしました。)

conf.py

from recommonmark.transform import AutoStructify


extensions = ["sphinx.ext.autodoc", "recommonmark"]
html_theme = "sphinx_rtd_theme"


def setup(app):
    app.add_transform(AutoStructify)

ちょっとだけつまづいたのは、 sphinx-apidoc コマンドで autodocの雛形となる一式をソースコードから生成するのですが、自動生成される index.rst が 手作りの index.md より優先されてしまって、うまく表示ができなかったので（しばらく悩んだあとに） index.rst を削除するコードを追加しすることで表示できるようになりました。

GitHub Pagesへの公開は、MasterブランチかつFork先ではない時に、peaceiris/actions-gh-pagesを利用して公開しています。
（どうも、gh-pages というブランチを作ってそこに生成物をpushする仕組みみたいです。なのでよくある master ブランチの doc以下を公開という設定とは異なります。）

.github/workflows/doc.yml（抜粋）

  - uses: peaceiris/actions-gh-pages@v3
    with:
      github_token: ${{ secrets.GITHUB_TOKEN }}
      publish_dir: ./public
    if: (github.ref == 'refs/heads/master') && (github.repository == 'keiohta/tf2rl')

Sphinxによって自動的に書き出されたクラスレファレンスはこちら
（クラスリファレンス以外のチュートリアルも自分で書きたいと聞いていたので、Markdownで書けるように環境は整備しましたが、まだ無いようです。。。乞うご期待。）

5. Docker と GitHub Container Registry を利用した構築済みコンテナ環境の提供 (PR 111)

ログインせずともコンテナイメージをダウンロードできるGitHub Container Registryがベータ公開されたのに併せて、構築済みコンテナ環境を利用できるようにしました。（それ以前のGitHub Packagesはログインしないとpublicなイメージもダウンロードできないという謎仕様で、フォーラムでも何人もの人がおかしいって言い続けていました。）

docker run -it ghcr.io/keiohta/tf2rl/cpu:v1.1.0 bash

また、Linuxオンリーかつマルチプロセスでの学習(ApeX)でうまく行かない例が散見されるのでexperimentalという扱いですが、NVIDIAのGPUを利用できるGPUコンテナバージョンも準備しています。
(私がMacBook Proしかなく、NVIDIA GPUの刺さったマシンを手元に持っていないためデバッグが難しく、環境をもった有識者が支援してくれると助かります。)

docker run --gpus all　-it ghcr.io/keiohta/tf2rl/nvidia:v1.1.0 bash

GPUコンテナについては、こちらの記事が非常に詳しくてためになりました。

コンテナイメージのビルドは以下のように設定してあり、tagつきでpushされたら、ビルドしてそのtag名をもったコンテナイメージをGitHub Container Registryに格納する仕組みになっています。

.github/workflows/docker.yml

name: docker

on:
  push:
    tags:
      - '**'

jobs:
  build:
    runs-on: ubuntu-latest
    env:
      DOCKER_BUILDKIT: 1
    steps:
      - uses: actions/checkout@v2
      - uses: docker/build-push-action@v1
        with:
          registry: ghcr.io
          repository: ${{ github.repository_owner }}/tf2rl/cpu
          tag_with_ref: true
          username: ${{ github.repository_owner }}
          password: ${{ secrets.GHCR_TOKEN }}
          dockerfile: Dockerfile
          always_pull: true
      - uses: docker/build-push-action@v1
        with:
          registry: ghcr.io
          repository: ${{ github.repository_owner }}/tf2rl/nvidia
          tag_with_ref: true
          username: ${{ github.repository_owner }}
          password: ${{ secrets.GHCR_TOKEN }}
          dockerfile: Dockerfile.nvidia
          always_pull: true

docker/build-push-actionは v2 があるのですが、ログインが分離されて二度手間で、tagを手動で切り出さなければならず、使い勝手が悪かったので、あえて v1 で固定しています。

また、ベータ版だからなのか、GitHub Container Registryは、GitHub Actionsの自動生成されるシークレットでは認証が通らないので、Personal Access Tokenを作って設定する必要があります。

もうひとつ、忘れがちなのはGitHub Container Registry にpushされたイメージは初めはプライベートになっています。皆が使えるようにするには、https://github.com/<ユーザー名>?tab=packages から指定のイメージを選んで、ポチポチ公開設定する必要があります。

6. Trainer をコマンドラインプログラム以外 (Jupyter Notebook等) からも実行可能に (PR 105)

TF2RLは元々コマンドラインからのスクリプト実行を前提とした作りになっており、構築したモデルを学習させる Trainer クラスは argparse と強く結合してしまっており、Jupyter Notebook上で利用するには（できなくはないが）ちょっと面倒であった。

後方互換性を維持するかつ元のコードへの変更点を最小限にするために、argparse ではなく dict が渡されたときには空の Namespace を構築してその中に dict 指定のデータを詰め込むことにした。こうすることで、わざわざコマンドラインパラメータを模擬した文字列のリストを経由しないで、パラメータを設定できる。

tf2rl/experiments/trainer.py(抜粋)

        if isinstance(args, dict):
            _args = args
            args = policy.__class__.get_argument(Trainer.get_argument())
            args = args.parse_args([])
            for k, v in _args.items():
                if hasattr(args, k):
                    setattr(args, k, v)
                else:
                    raise ValueError(f"{k} is invalid parameter.")

7. TensorFlow Probability のバージョンの自動指定 (PR 113)

最近TF2RLがTensorFlow Probability (TFP) を利用するようになったのですが、TFPは動作するTensorFlow (TF)のバージョンが決まっている割にはインストール時にいい感じにしてくれなくて、何も考えずにインストールするとバージョン非互換でエラーが発生することがありました。

そこで、TF2RLの setup.py に以下のコードブロックを追加し、ユーザーがインストール済みのTFのバージョンに応じて、インストールすべきTFPのバージョンを切り替える方式を導入しました。

setup.py（抜粋）

tf_version = "2.3"  # Default Version
compatible_tfp = {"2.3": ["tensorflow~=2.3.0",
                          "tensorflow-probability~=0.11.0"],
                  "2.2": ["tensorflow-probability~=0.10.0"],
                  "2.1": ["tensorflow-probability~=0.8.0"],
                  "2.0": ["tensorflow-probability~=0.8.0"]}

try:
    import tensorflow as tf
    tf_version = tf.version.VERSION.rsplit('.', 1)[0]
except ImportError:
    pass


install_requires = [
    "cpprb>=8.1.1",
    "setuptools>=41.0.0",
    "numpy>=1.16.0",
    "joblib",
    "scipy",
    *compatible_tfp[tf_version]
]

~= は、バージョンの最後の数字はアップデートを許容するシンタックスで、この場合はメジャーバージョンやマイナーバージョンの更新は許可しないが、バグフィックス（だと思われる）更新は許可するためにこうしてあります。

TF2.4 が出たらエラーになりますが、それは他のTF2.4変更対応とともに対処することになるでしょう。

8. 最後に

TensorFlow は書きにくいから PyTorch って思っている方も結構いるようですが、TF 2.0 以降非常に書きやすくなっていると個人的には思います。

TF2RLのユーザーがもっと増えたらいいなと思っています。

TF2RLでも利用している経験再生(Experience Replay)用のライブラリ cpprb (GitHubミラー) を開発しています。
こちらもぜひ興味を持ってもらえると嬉しいです。