確率的な環境でもDecision Transformerを動かしたい！
3つの要点
✔️ 確率的な環境でRvSの成功に必要な情報統計量の性質を提示
✔️ 上記の情報統計量を推定するアルゴリズムESPERを提案
✔️ 2048などの確率的な環境で既存のRvS手法, Decision Transformerを超える性能を達成
You Can't Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments
written by Keiran Paster, Sheila McIlraith, Jimmy Ba
(Submitted on 31 May 2022 (v1), last revised 28 Nov 2022 (this version, v2))
Comments: Added experiments with Decision Transformers; Fixed error in Theorem 2.1; Updated related works; Added link for code
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
code：
本記事で使用している画像は論文中のもの, 紹介スライドのもの, またはそれを参考に作成したものを使用しております.
導入
近年, Transformerなどの生成モデルを用いて, 軌道や報酬で条件づけられた元での行動の予測タスク(教師あり学習)を解くことによって, 強化学習の問題を解く, Reinforcement learning via Supervised learning(RvS)と呼ばれる枠組みが注目を集めています. 既存のRvS手法では, 訓練時に, 状態行動の履歴に加えて軌道全体の報酬を入力として与えて行動を予測することで, モデルに軌道全体の報酬と行動の関係を学習させて, 実行時には高い報酬を入力することで高い性能を達成することを目指します. 本論文では, 環境が確率的な場合, 単に軌道の報酬で条件づけるだけでは性能の高いエージェントは学習できないということを示し, その問題に対する解決策, ESPERを提示しています.
Decision Transformer と RvS
本章では, はじめに, RvSの代表的な手法であるDecision Transformerについて紹介して, それを一般化する形でRvSを定式化します. 強化学習の基本的な記法について整理します. 強化学習はマルコフ決定過程によって定式化されます. マルコフ決定過程は状態集合$S$, 行動集合 $A$, 遷移確率 $T$, 報酬関数 $r$, 割引率$\gamma$からなります. 強化学習では, 累積割引リターン, $\sum_{t} \gamma^t r_t$(以下単にリターンと呼びます) を最大化するような行動方策 $\pi$を学習することを目的とします. あるエピソードの状態, 行動, 報酬の履歴を軌道と呼び$\tau = (s_0, a_0, r_0, \dots)$と記します. Decision Transformerでは, 与えられた相互作用のデータ$D = \{\tau_i\}_{i=1}^N$から方策を, transformerを用いて以下の形でモデル化して学習します.
$$\pi(a_t|s_0, a_0, \dots, s_t, \hat{R}) = p_D(a_t|s_0, a_0, \dots, s_t, \hat{R})$$
ここで, $p_D$はデータから誘導される分布です. ここで, $\hat{R}$はその軌道のリターンです. 各時刻 $t$ において, それまでの状態と行動の系列とリターンで条件づけた, 行動の最尤推定問題を解くことで方策を学習します. 言語タスクでのtransformerによる次単語予測との対比で考えると理解しやすいと思います. Transformerは直前までの言葉で条件づけた時に, 次に来る一番最もらしい言葉を予測するタスクを解くことで, 言語の使用法を学習します. 一方でDecition transformerは直前までの状態行動の履歴と, その軌道全体のリターンで条件づけた時に, データから見て最も尤もらしい行動を予測することで, リターンと行動の関係を学習します.
これは, 以下のように一般化することができます. まず, 軌道( $\tau$ )から決まる情報統計量 $I(\tau)$ を考えます. これは軌道のスコアのようなものでDecision transformerでは $I(\tau)$ はリターンでした. そこで, 情報統計量同士の距離を測る指標$D$を考えます. これはDecision Transformerでは条件付き尤度でした. この要素を導入した上で, RvSは, データから決まる情報統計量の値 $z$の分布が与えられた時に, $z$と$I(\tau)$の距離$D$を最小にする軌道$\tau$を導く$\pi$を学習する問題として定式化できます. 言葉だとわかりにくいと思うのでいかに数式を示します.
$$\min_\pi \mathbb{E}_{z\sim p(z), \tau \sim p_{\pi_z}[D(I(\tau), z)]}$$
本論文では, 確率的な環境の場合, $I(\tau)$ をリターンとして上記の問題を解いても, 「学習が成功しない」ということを示し, 確率的な環境で望ましい$I(\tau)$の性質を明らかにしています.
確率的な環境におけるRvS
なぜ通常のRvSは失敗するのか
なぜRvSは確率的な環境でリターンで条件付けた方策は失敗するのでしょうか. 直感的には, 得られたリターンが, 行動の結果によるものか, 確率的な遷移によるものか判断できないからです. 上のMDPの例を元に考えてみましょう[図1]. 上の遷移図はデータの分布を表しています. 報酬$r=1$について考えてみましょう, $a_1$を取った場合は50%で, $a_2$を取った場合は100%で得られることがわかります. $a_2$については行動のダイレクトな結果として, $a_1$については確率的な遷移によってこの報酬が達成されていることがわかります. このデータをもとに条件付き分布を学習したとします. 学習後の分布を$r=1$で条件づけると, $a_1, a_2$両方確率の値を取ることがわかります. 最適な行動は$a_2$ですが, それを得るために$r=1$を指定しても, $a_1$も一定の確率で取ってしまいます. これは, 確率的な遷移の結果と, 行動の結果を見分けられないような統計量であるリターンを学習に使用したからだと考えることができます.
環境のランダムネスと独立な統計量
ここでは, 確率的な環境でRvSが成功するために必要な条件を理論的に解析します. まず, 学習が成功するということの定義を以下で行います.
Definition 2.1 (Consistently Achievable). A goal $z$ is consistently achievable from state $s_0$ under policy $\pi(a|s, z)$ if $\mathbb{E}_{\tau \sim p_{\pi_z}(\tau|s_0)}[D(I(\tau), z)] = 0$.
つまり, 学習が成功するとは, 学習した方策をある統計量の値$z$で条件づけた際に誘導される軌道$\tau$が, 期待的に指定された$z$と同じ統計量$I(\tau)$を達成するということです. リターンの話で置き換えると, 学習した方策が期待的に指定したリターンを達成するということです. これを踏まえた上で, 学習が成功するために必要な統計量の性質を以下の定理で述べています.
Theorem 2.1. Let $\pi_D$ be a data collecting policy used to gather data used to train an RvS policy, and assume this policy is trained such that $\pi(a_t|s_0, a_0, \dots, s_t, z) = p_{\pi_D}(a_t|s_0, a_0, \dots, s_t, I(\tau)=z)$. Then, for any goal z such that $p_{\pi_D}(I(\tau)=z|s_0)>0$, goal $z$ is consistently achievable iff $p_{\pi_D}(s_t|s_0, a_0, \dots, a_{t-1}) = p_{\pi_D}(s_t|s_0, a_0, \dots, a_{t-1}, I(\tau))$.
この定理が要請する条件は, 統計量 $I(\tau)$ が環境のランダムネスと独立であるということです. つまり, 統計量が環境のランダムネスの情報を含んでいないということでです. こうした統計量を構成できれば, 確率的な環境においてもDecision Transformerのような既存のRvSの手法を用いて方策を学習することができるというわけです.
手法
ここでは, 環境のダイナミクスと独立な統計量 $I(\tau)$ を学習する, ESPER (Environment Stochastic Independent Representation)と彼らが呼ぶ手法を解説します. ESPERではまず, 敵対的クラスタリングによって統計量の表現 $I(\tau)$ を学習します. 次にその統計量によってクラスタを作りクラスタごとに平均のリターンを計算します. 最後にそのリターンを用いて通常のRvSを実行します[図2].
敵対的クラスタリング
敵対的クラスタリングでは, $I(\tau)$をダイナミクスの推定について敵対的に学習を進めることで, $I(\tau)$ がダイナミクスと独立となるよう促します. まず以下が学習対象のモデルです.
- Cluster model: $I(\tau) \sim p_{\theta}(I(\tau)|\tau)$
- Action predictor: $a_t \sim p_\theta(a_t|s_t, I(\tau))$
- Return predictor: $\hat{R}=f_{\psi}(I(\tau))$
- Transition predictor: $s_{t+1} \sim p_{\phi}(s_{t+1}|s_0, a_0, \dots, s_t, a_t, I(\tau))$
以下が目的関数です.
$$L(\theta) = \mathbb{E}_{I(\tau)\sim p_\theta(I(\tau)|\tau)}\left[- \underbrace{\beta_{\mathrm{act}} \log p_{\theta}(a_t|s_t, I(\tau))}_{行動を予測} + \underbrace{\beta_{\mathrm{adv}} \log p_{\phi}(s_{t+1}|s_0, a_0, \dots, s_t, a_t, I(\tau))}_{ダイナミクスの予測に対してadversarialとなることを要請}\right]$$
$$L(\phi) = \mathbb{E}_{I(\tau)\sim p_\theta(I(\tau)|\tau)}\left[\underbrace{- \log p_{\phi}(s_{t+1}|s_0, a_0, \dots, s_t, a_t, I(\tau))}_{ダイナミクスを予測}\right]$$
リターンの推定モデルは, 行動の予測に寄与する一方で, ダイナミクスの推定に対しては敵対的であることが求められていることがわかると思います.
クラスターの平均リターンを推定
次に, 学習した統計量の表現モデル$p_\theta(I(\tau)|\tau)$を用いて, 軌道の報酬を$f_{\psi}(\tau)$によって推定します. これによって, 統計量をもつ軌道の平均的なリターンをで推定します. リターンが推定できれば, ここまでに推定した$p_{\theta}(I(\tau)|\tau)$, $f_{\psi}(\tau)$を用いることで, データ中の各軌道$\tau$に対して, 確率のランダムネスに依存しない統計量 $I(\tau)$ でクラスタリングした時の平均的なリターンの推定値を得ることができます.
$$L(\psi) = \mathbb{E}_{I(\tau)\sim p_\theta(I(\tau)|\tau)}\left[||R - f_{\psi}(I(\tau))||_2^2\right]$$
RvS
最後に, 上で述べたクラスタの平均リターンを条件付け因子として, 通常のRvSを行います.
実験
本研究の実験では, 確率的なダイナミクスを含む環境として以下の3つを取り上げています.
ベンチマーク
- Gambling: 上で例証に用いた行動が3つの環境.
- Connect Four: 確率的な敵をもうけた4目並べ. 勝利すると報酬1, それ以外は報酬0としています.
- 2048: タイルの値と出現位置が確率的に決まる含むパズルゲーム. 上下左右のスライドを繰り返し, 2048を作ることを目指す. 本実験では128のタイルを作ることができれば報酬1, そうでなければ報酬0として環境を単純化しています.
ベースライン
- Return-Conditioned RvS (DT): 従来の, リターンによって条件づけられたDecision Transformer.
- CQL: 価値関数ベースのオフライン強化学習手法.
結果
図3をみてわかるように, ESPERはReturn-Conditioned RvS (DT)だけでなく, CQLに対しても優位性を持つことがわかります. この実験結果から, 確率的な環境では通常のリターンを用いるRvSの性能が低くなること, ESPERによって, 性能が大幅に改善することが確かめられ, 先程の解析と併せて, ESPERの有効性が理論的, 実験的に示されました.
まとめ
いかがだったでしょうか. 本研究では, 確率的な環境で動作するRvS手法, ESPERを提案した論文を紹介しました. RvSは大規模データとの相性が良く, 現在盛んに研究が進められています. ESPERは, ノイズなどによって遷移が決定論的に推定できないような現実問題へのRvSの適用を可能にする点で, 実社会での活用事例も期待できます. 今後も, RvSの研究の動向を追ってみてください.
この記事に関するカテゴリー
