[論文レビュー] Representation Learning for Online and Offline RL in Low-rank MDPs

Masatoshi Uehara, Xuezhou Zhang+1 more|arXiv (Cornell University)|Oct 9, 2021

Advanced Bandit Algorithms Research参考文献 58被引用数 5

ひとこと要約

本稿では、低ランクマルコフ決定過程（MDP）における表現学習のための新規アルゴリズム Rep-UCB を提案する。この手法により、オンラインおよびオフライン強化学習において、サンプル効率が向上する。上位信頼区間を表現学習と統合することで、先行研究と比較して著しく改善されたサンプル複雑度 $ widetilde{O}(d^4 A^2 / ( epsilon^2 (1- gamma)^5))$ を達成する一方で、Flambe らの explore-then-commit 手法よりも単純かつデータ効率が良い。

ABSTRACT

This work studies the question of Representation Learning in RL: how can we learn a compact low-dimensional representation such that on top of the representation we can perform RL procedures such as exploration and exploitation, in a sample efficient manner. We focus on the low-rank Markov Decision Processes (MDPs) where the transition dynamics correspond to a low-rank transition matrix. Unlike prior works that assume the representation is known (e.g., linear MDPs), here we need to learn the representation for the low-rank MDP. We study both the online RL and offline RL settings. For the online setting, operating with the same computational oracles used in FLAMBE (Agarwal et.al), the state-of-art algorithm for learning representations in low-rank MDPs, we propose an algorithm REP-UCB Upper Confidence Bound driven Representation learning for RL), which significantly improves the sample complexity from $\widetilde{O}( A^9 d^7 / (ε^{10} (1-γ)^{22}))$ for FLAMBE to $\widetilde{O}( A^2 d^4 / (ε^2 (1-γ)^{5}) )$ with $d$ being the rank of the transition matrix (or dimension of the ground truth representation), $A$ being the number of actions, and $γ$ being the discounted factor. Notably, REP-UCB is simpler than FLAMBE, as it directly balances the interplay between representation learning, exploration, and exploitation, while FLAMBE is an explore-then-commit style approach and has to perform reward-free exploration step-by-step forward in time. For the offline RL setting, we develop an algorithm that leverages pessimism to learn under a partial coverage condition: our algorithm is able to compete against any policy as long as it is covered by the offline distribution.

研究の動機と目的

• 遷移ダイナミクスが低ランクであるが、表現が未知である低ランク MDP において、コンact かつ低次元の表現を学習する課題に対処すること。
• オンライン強化学習におけるサンプル効率を向上させるために、表現学習、探索、利用を統合したフレームワークで直接バランスを取ること。
• 部分的カバレッジ条件の下で、任意のポリシーがデータ分布によってカバーされている場合に性能保証が得られる、証明可能に効率的なオフライン強化学習アルゴリズムの開発。
• Flambe らの explore-then-commit 法の制限を克服し、時間ステップ間でデータを共有可能にするとともに、アルゴリズム設計を簡素化すること。

提案手法

• 低ランク MDP のオンライン強化学習において、表現学習、探索、利用を統合的に最適化する Upper Confidence Bound に基づくアルゴリズム Rep-UCB を提案する。
• 最新に学習された表現に基づく信頼区間ボーナスを用いて探索を誘導し、状態-行動空間の十分なカバレッジを保証する。
• Flambe らの計算オракル（例：MLE と ERM）を活用することで、オラクル効率を維持しつつ、サンプル複雑度を改善する。
• オフライン強化学習では、分布シフトに対処するための懐疑的アプローチ（pessimism）を適用し、オフラインデータ分布によって部分的にカバーされる任意のポリシーと競合できるようにする。
• 真の表現 $\phi^\star$ のみに依存する部分的カバレッジ条件を導入し、先行研究よりも強い一般化保証を可能にする。
• 表現誤差の分解と有限標本一般化バウンドを用い、Lemma 23 を通じて誤差バウンドをサンプル複雑度に変換することで、サンプル複雑度の結果を導出する。

実験結果

リサーチクエスチョン

• RQ1 表現が未知である低ランク MDP におけるオンライン強化学習で、Flambe らの explore-then-commit 手法と比較して、より良いサンプル複雑度を達成できるか？
• RQ2 表現学習、探索、利用をより効果的にバランスさせる統合的アルゴリズムをどのように設計できるか？
• RQ3 分布シフトが存在する状況で、任意のポリシーに対する性能保証を得るために必要な最小のカバレッジ条件は何か？
• RQ4 表現が事前に与えられていない状況でも、低ランク MDP において証明可能に効率的な学習を達成できるか？
• RQ5 提案手法のサンプル複雑度は、ランク $d$、行動数 $A$、割引因子 $\gamma$ に対してどのようにスケーリングされるか？

主な発見

• Rep-UCB はオンライン強化学習において $\widetilde{O}(d^4 A^2 / (\epsilon^2 (1-\gamma)^5))$ のサンプル複雑度を達成し、Flambe の $\widetilde{O}(A^9 d^7 / (\epsilon^{10} (1-\gamma)^{22}))$ と比較して顕著な改善を示す。
• 本手法は Flambe よりも単純である。これは、段階的で explore-then-commit 構造を避けており、時間ステップ間でのデータ共有が可能であるためである。
• オフライン強化学習において、本手法は真の表現に基づく部分的カバレッジ条件の下で、オフラインデータ分布によってカバーされる任意のポリシーと競合できる。
• 本研究で用いられる部分的カバレッジ条件は、先行研究（例：Xie ら、2021）のものよりも弱く、報酬を含む拡張特徴ではなく、$\phi^\star$ のみに依存する。
• 本手法の性能保証は、マーカフ的ポリシーに限定されず、履歴依存ポリシーにも拡張可能である。
• 理論的分析では、新規の有限標本誤差から複雑度への変換（Lemma 23）を用い、一般化誤差の対数因子を制御することで、よりタイトなサンプル複雑度バウンドを達成している。