※この記事は，CAMPHOR- Advent Calendar 2023の2日目の記事になります．
CAMPHOR- Advent Calendar 2023についてはこちら：https://advent.camph.net
記事を書いている人
京都の技術コミュニティCAMPHOR-の運営メンバー
研究ではロボットを強化学習で制御しているM1
この記事について
この記事では，Wandbという可視化ツールを紹介し強化学習における実用例を示します．
Wandbとは
実験結果の可視化，パラメータ管理，データセット管理などができるライブラリです．企業やKaggleで使われていることもあるようで，機械学習（ML）や強化学習（RL）で使うことを想定されていそうですが，データロギング自体は昨今どの分野でもある程度行うでしょうから，分野を問わず活躍しそうです．
個人的によいと思う点は，
- ブラウザで結果を確認できる
- 簡単にSlackと連携できる
- 扱いが簡単
- 計算資源のモニタリングを勝手にやってくれる
- 動画ロギングも容易
といったところです．ブラウザで結果を確認できるので，MLやRLを書いたファイルを実行したあと外出して，出先で結果を確認できます．サーバにSSHしなくてもいいのは個人的にうれしいポイントです．
また，コードを数行差し込むだけでデータのログを作れるので，既存のコードと親和性が高いです．公式の触れ込みでは，「5行のコードでトラッキング、バージョン管理、可視化を開始」とのことですから，手軽な可視化を可能にすることが意識されているのでしょう．
Weights & Biases のアカウント作成
Wandbを利用するためには，まず下記URLよりWeights & Biasesのアカウントを作成する必要があります．
https://www.wandb.jp/
メールアドレスとパスワードを入力して次へ．
Full nameと所属・Usernameを入力します．この後は仕事で使うのか，個人利用かと聞かれるのでここではいったん個人利用としておきます．というか無料で使えるのは個人利用のみです．ただし，個人利用では他のメンバーと結果を共有できませんので，共同研究者等がいるときは注意しましょう．
アカウントを作成して最初に出てくるこの画面はそのままにしておきましょう．後で使います．
実装
では，実際に使ってみましょう．私の環境は以下の通りです．
- Ubuntu 22.04
- VSCode 1.78.2
- Python 3.8.10
- Wandb 0.16.0
まずは手始めにWandbをインストールします．Wandbを「使ってみる」ということなので，仮想環境でやっていきます．適当なディレクトリをエディタで開いて，
python3 -m venv .venv
source .venv/bin/activate
ここからコマンドラインでWandbにログインするのですが，これについては先ほどの画面の指示に従いましょう．ごく短いですしわかりやすいです．書くのが面倒
さて，ログインが終わったらコードを書いていきましょう．
import wandb
run = wandb.init(project="test")
for i in range(10):
y1 = i**2
y2 = i + 1
wandb.log({"y1": y1, "y2": y2})
これを実行したあと，先ほどの画面をリロードすると実行結果が保存されています．
左端のOverviewには
- 実行時刻
- OS
- Pythonのバージョン
- CPU・GPU数
などが記述されています．Logsにはコマンドライン出力が保存されています．どういったバグが出たかどうかを後からも確認できるので，見覚えのある他の人のバグを解決するときにも使えそうですね．Filesにはrequirements.txtが自動で生成されているので，環境を移植するときには非常にありがたいです．研究室ではテーマの引継ぎに際して環境の移植で困ることもあるでしょうから，積極的にこの辺を使っていきましょう．
次は，stable baselines3と併用してみたいと思います．
https://github.com/DLR-RM/stable-baselines3
stable-baselines3は，PyTorchで実装された強化学習アルゴリズムを利用できるライブラリです．OpenAIのGymnasiumという学習環境と合わせて使うことができます．様々な環境に対してアルゴリズムを試したり，研究するうえでコーディングのミスをしないために用いたりします．stable-baselines3をインストールします．
pip install stable-baselines3[extra]
私の環境ではstable-baselines3のバージョンは2.2.1でした．
今回は，鉛直下向きの棒を立たせるPendulumという環境でRLをやってみましょう．学習アルゴリズムはProximal Policy Optimization（John Schulman et al, 2017）というものです．
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import VecVideoRecorder
from stable_baselines3.common.env_util import make_vec_env
import wandb
from wandb.integration.sb3 import WandbCallback
cfg = {
"policy_class": "MlpPolicy",
"timesteps": 3e5,
"env": "Pendulum-v1"
}
run = wandb.init(
"sb3-pendulum",
config=cfg,
sync_tensorboard=True,
monitor_gym=True
)
env = make_vec_env(cfg["env"])
env = VecVideoRecorder(
env,
video_folder=f"video_log/{run.id}",
record_video_trigger=lambda x: x%3e4==0
)
model = PPO(
cfg["policy_class"],
env,
tensorboard_log="log/"
)
model.learn(
cfg["timesteps"],
callback=WandbCallback(verbose=2, gradient_save_freq=100)
)
引用元：https://docs.wandb.ai/guides/integrations/stable-baselines-3
cfgはコンフィグとしてwandbのログとして保存しておきたいものを書いておきます．Overviewから確認できます．学習結果は以下の通りです．
シミュレータであるMujocoが入っていれば
pip install gymnasium[mujoco] moviepy
を実行したうえで"env": "Pendulum-v1"を書き換えることで同様に学習結果の記録ができます．
"env": "Hopper-v1"
学習するとこんな感じになります．
あまりコードを書くのに慣れていなかった学部4回生の頃，強化学習アルゴリズムを自分で書くことよりも可視化のコードを書く方が苦労しました．修士1回生になってからこれを知って，なんと素晴らしいものがあるのかと研究やコーディングのモチベーションすら上がった記憶が残っています．
まとめ
- Wandbを使えば楽に可視化を実装できる
- 可視化がきれいにできるとコーディングも楽しい
参考・引用
- ほかにも可視化ツールはあります．ご自身に合うものを探してみるのもいいでしょう．
- Tensorboard
- Neptune
- Comet
- ...etc
- Policy Proximal Optimization Algorithm: https://arxiv.org/abs/1707.06347
- PPOの論文です