Imitationによる報酬関数の可視化

概要

Imitationを使って、報酬関数の可視化を行いました。

https://imitation.readthedocs.io/

Center for Human-Compatible AI, ver 1.0

背景

GAILは、生成的敵対ネットワーク（Generative Adversarial Networks, GANs）の概念を応用して、専門家のデモンストレーションから、専門家の振る舞い(policy)を模倣するGeneratorと、専門家の行動とエージェントの行動を区別するDiscriminatorとの間で敵対的な学習を行います。

逆強化学習はこのDiscriminatorからは、逆強化学習の枠組みの中で報酬関数が学習されます。が、この報酬関数の実装例が見つからず、変なところで苦戦したため、メモします。

勉強はじめでよくわからないところが多いため、間違っているところありましたらご助言いただけるとなによりです。

方法

ここでは、Stable baselines3と、Imitationのライブラリを使っています。状態の理解が楽なため、問題にはOpenAI gymのCartPole-v0を使います。

CartpoleはCartの上に立ったPoleを倒さないように、Cartを右左に動かすゲームです。直感的に理解がしやすいですが、基本的な方程式の説明は、以下などにもあります。

https://brain.cc.kogakuin.ac.jp/~kanamaru/NN/CPRL/index-j.html

Imitationのチュートリアルのところは省略します。

エキスパートデータの作成（sealsから）
GAILで学習（チュートリアルまま）
Policy, Reward_netの取得（念のため保存）
エキスパートデータでの報酬の推定
シミュレーションでの報酬の推定
(追加) AIRLで同様

報酬関数

肝心の報酬関数の呼び出しは簡単で、

reward_net.predict(obs_tensor, act_tensor, next_obs_tensor, done_tensor)

はじめChatGPTに聞きながらやっていたところ、

reward_net(obs_tensor, act_tensor, next_obs_tensor, done_tensor)

と言われ、そこからテンソルの次元が合わないエラーがでて、あれこれやって、一時的にreward_net.use_done = True、にしたら次元があうけど、それでいいのか？と悩んで、中身見て、一つ一つ理解していったら、、、。と無駄に時間をかけました。

やはりPytorch勉強しなくては。

def estimate_reward(obs, act, next_obs, done, reward_net):
    # 観測値、行動、次の観測値、および終了フラグを適切なテンソル形式に変換
    act = np.array([act])
    done = np.array([done])
   
    obs_tensor = torch.as_tensor(np.copy(obs), dtype=torch.float32).unsqueeze(0)
    act_tensor = torch.as_tensor(np.copy(act), dtype=torch.float32).unsqueeze(0)
    next_obs_tensor = torch.as_tensor(np.copy(next_obs), dtype=torch.float32).unsqueeze(0)
    done_tensor = torch.as_tensor(np.copy(done), dtype=torch.float32).unsqueeze(0)

    # reward_netに入力して報酬を推定
    with torch.no_grad():  # 勾配計算を行わない
        reward = reward_net.predict(obs_tensor, act_tensor, next_obs_tensor, done_tensor)
    return reward.item()

可視化

いきなりですが、報酬を可視化します。

CartPoleの状態は4次元で、

Cartの位置：　-4.8～4.8
Cartの速度：　-Inf～Inf
Poleの角度：　-24～24[rad]
Poleの角速度：　-Inf～Inf

なのですが、カートの位置は意味がないように思えるため無視して、Volocityが0.1より大きいか小さいかと、更に実際のactで分け、残りの2次元で可視化します。*1

Poleが左に倒れそうなときは、Cartも左に動かしてやるのが正解な動かし方です（たぶん）。

各図の第3象限（左下）がその状態で、cとdを比較するとcの報酬が大きく、すなわちact=0で、左に移動するのが好ましくでています。同様に第1象限はPoleが右に倒れそうなときなので、右に動かしてやるのが正解で、bは正しくaより大きい報酬となっています。

ただし、本来物理的な意味を考えると、正負で反対の結果になるはずなので、bとc、aとdは同じような結果になるはずですが、なぜか角速度が負で速度が正の場合（両者は強い相関あり）は報酬が小さくなっています。エキスパートデータに偏りがあるかとチェックしましたが、そうでもないよう。要確認。AIRLを試したところ、概ね解決しました。

シミュレーション

以上は、エキスパートデータを使っていたので、Actに対する報酬は実は違う状態を比較していたため、シミュレーションし毎回非推奨側の行動の報酬も推定してみます。最大Stepを20回にして、計10000step。

図はPolicyが推奨した行動に対する推定報酬と、非推奨行動の推定報酬。

全然、ちゃんと報酬推定できてないじゃん（笑）。推奨行動の報酬が大きい割合が50％と、完全にランダムじゃん。

と思ったのですが、当然Stepに制限をかけないと、いつまでも動き続けるので、減衰なしの累積報酬和の利得を見てみると、

ちゃんと、推奨行動側の利得がほぼ100％で大きくなっています。よかった。

自分なりの理解を書くと、棒があんまり倒れていなく、どのような行動でもそんなすぐに終了に繋がらないような場合、報酬の推定は重要でない（重要なのは利得）。CartPoleは実際にプレイするとわかるのですが、actに対してVelocityが飛び飛びでめちゃくちゃ動くので、中立付近ではどちらでもよいのでしょうか？

一方で、非推奨行動の報酬が非常に小さい場合、それは回避する。その結果、利得は大きくなり、行動も続けられる。

ただし、より広い範囲をみる目的で、40％程度の確率でランダムに行動するようにしたところ、報酬和で推奨行動が必ずしもよくなっていないことも増えてきたため、要確認です。

追記(AIRL)

GAILは実装例を見ていたこともあり馴染みがあったのに対し*2、AIRLはそんな変わらないとたかをくくってやったことがなかったのですが、試しにやってみたところ、報酬関数が状態の正負において対称に近いかたちになって、GAILよりもそれっぽい形になりました。（Step数他はチュートリアルと同じ）

ただ、1回ごとの推奨行動と非推奨行動の報酬は、GAILと同様に50%にわかれて、学習できているのか、いないのかよくわからないのは、かわらず。

LogPとの比較

そもそも、逆強化学習をしなくても、方策のみから状態価値と行動の対数尤度は簡単に出力できます。

AIRLでの行動の対数尤度は下図ですが、ほとんどrewardと違いがわかりません。省略しますが、状態価値も傾向はかわらず。

教科書でも読んで勉強しろよという感じですが、逆強化学習は強化学習の一部として教科書にはちょろっと1p程度触れられる程度。

当然、価値関数は長期的な報酬和であることはわかっていますが、感覚的に、エキスパートデータのみから、詳細な報酬を推定できるわけない気もするので、行動価値関数とだいたい一緒になるとしたら、まーそうだろうな。

だとすると、逆強化学習の価値は、あくまでも強化学習をおこなうための報酬関数を推定してくれることのみ、にあるのでしょうか？

雑記、疑問点

そもそも報酬の推定に、next_obsが必要なのってなんでだろうか？それが強化学習と言われればそれまでだが、逆強化学習の観点からすれば不自然でないか？next_obsが必要だと、環境とのやり取りができない場合に、実際に行動した側のnext_obsしか得られなく、もう一声。
reward_netは、gail_trainer._reward_netにあるようですが、先頭に1つのアンダースコアついてるけど、保存から再読込してもよいんだろうか？

*1:Velocity=0付近は混合していて、図が少し見にくかったため

*2:https://tech.morikatron.ai/entry/2020/10/12/100000

NezumiNoKuni’s blog

日々気になったことをメモしてます。ポケモン、データ分析、育児他を予定してます。