OpenAI Gym 独自Envクラス内でmlflowを使う

OpenAI Gym の独自Envで強化学習を行う場合、評価したいパラメータは独自Envクラス内に埋まってしまう。

このため mlflow を使って実験を管理する場合、独自Envクラス内に実装しないといけなくなる。

たとえば以下のような gym.Env 継承クラスの実装をおこなう。

import gym
import mlflow
from mlflow.tensorflow import autolog # tensorflow実装の場合

class HogehogeTrainEnv(gym.Env):

    def __init__(self):
        :
        # 独自環境クラス初期化処理
        self.step_count = 0
        self.total_reward = 0.0

        :
        # mlflow実行開始
        mlflow.start_run()
        # パラメータとしてenv_idを書き込む
        mlflow.log_param('env_id', 'HogehogeEnv-v0')
        autolog()

    def step(self, action):
        # 累積ステップ数
        self.step_count += 1
        :
        # 観測データ obs 更新
        # エピソード完 done 判定
        # 当該ステップの報酬値 reward 算出
        # その他情報 info の格納

        # エピソード内報酬値合計
        self.total_reward += reward
        :
        # エピソード完了時に報酬値合計を書き込む
        if done:
            ml_metric('total_reward', self.total_reward, step=self.step_count)
            self.total_reward = 0.0
        return obs, reward, done, info

    def close(self):
        :
        # mlflow 実行終了
        mlflow.end_run()

mlflow.log_params() / mlflow.log_param() は同一キーの上書きをする場合エラーを発生させることがある。
 

mlflow.log_metrics() / mlflo.log_metric() は毎step実行すると大量のデータをグラフ化できなくなるので実行タイミングを独自環境ごとに工夫する必要がある。上記スードコードの例ではエピソード完了時のみ書き込みを行っているが、total_timestepsをおおきくすればするほどログ出力が増え、トレーニング処理が増大する。

ログ出力が膨大な場合は mlflow ui するとグラフ表示できなくなることがある。tensorboard を使う場合は mlflowでは指標は書き込まずにパラメータのみ管理するなどの切り分けをして、余計なログを取らないようにする。

上記のクラスはあくまでトレーニング処理用の独自Envクラスとして設計している。評価用Envクラスと共用する場合は、評価時にもmlflowログが出力されることに注意。

上記の例では mlflow.start_run() / mlflow.end_run() を各テンプレートメソッドに埋め込んでいるが、これは Stable Baselines を使用する場合を想定している。Stable Baselinesを使う場合は autolog はどうも機能しないようなので注意。