Stable Baselines - TensorBoardによる学習状況の監視

TensorBoard は TensorFlow のデータを可視化するツールです。

これを使うことにより、学習状況をより詳細に確認することができます。

(Ubuntu 19.10で動作確認しています。)

TensofBoardのログを出力する

TensorBoard 用のログを出力するためには、モデルを作成する際の引数 tensorboard_log に出力先ディレクトリを指定します(15行目)。

[ソース]

tensorboard_watch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import gym
import os
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2

# ログ出力先のディレクトリを作成
log_dir = './logs/'
os.makedirs(log_dir, exist_ok=True)

# 環境を作成
env = gym.make('CartPole-v1')
env = DummyVecEnv([lambda: env])

# モデルを作成
model = PPO2('MlpPolicy', env, verbose=1, tensorboard_log=log_dir)

# モデルの学習
model.learn(total_timesteps=100000)

# モデルのテスト
state = env.reset()
for i in range(200):
    # 環境の描画
    env.render()
    # モデルの推論
    action, _ = model.predict(state)
    # 1ステップ実行
    state, rewards, done, info = env.step(action)
    # エピソード完了
    if done:
        break
# 環境のクローズ
env.close()

上記コードを実行します。

[コマンド]

1
python3.7 tensorboard_watch.py

TensofBoardの実行

先ほど実行したログを TensorBoard で表示するために下記のコマンド実行します。

[コマンド]

1
tensorboard --logdir=./logs/.

正常に実行されると、コンソールにURLが表示されるのでそれをブラウザで開きます。

実行結果


学習による平均報酬のグラフやその他いろいろな情報を確認することができます。