[論文レビュー] Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learning

Qinjie Lin, Han Liu+1 more|arXiv (Cornell University)|Mar 14, 2022

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

SwitchTT は、スパースでエキスパートベースのトランスフォーマー構造（スイッチトランスフォーマー）と分布価値推定を用いたマルチタスクオフライン強化学習手法を提案する。これにより、サンプル効率が向上し、計算負荷が低減される。10のスパース報酬タスクにおいて、Trajectory Transformer よりも10%高いパフォーマンスを達成し、最大90%の高速化を実現。マルチタスク設定におけるスケーラビリティと価値推定の優位性が示された。

ABSTRACT

We propose SwitchTT, a multi-task extension to Trajectory Transformer but enhanced with two striking features: (i) exploiting a sparsely activated model to reduce computation cost in multi-task offline model learning and (ii) adopting a distributional trajectory value estimator that improves policy performance, especially in sparse reward settings. These two enhancements make SwitchTT suitable for solving multi-task offline reinforcement learning problems, where model capacity is critical for absorbing the vast quantities of knowledge available in the multi-task dataset. More specifically, SwitchTT exploits switch transformer model architecture for multi-task policy learning, allowing us to improve model capacity without proportional computation cost. Also, SwitchTT approximates the distribution rather than the expectation of trajectory value, mitigating the effects of the Monte-Carlo Value estimator suffering from poor sample complexity, especially in the sparse-reward setting. We evaluate our method using the suite of ten sparse-reward tasks from the gym-mini-grid environment.We show an improvement of 10% over Trajectory Transformer across 10-task learning and obtain up to 90% increase in offline model training speed. Our results also demonstrate the advantage of the switch transformer model for absorbing expert knowledge and the importance of value distribution in evaluating the trajectory.

研究の動機と目的

• 多様で異種のデータセットを用いたマルチタスクオフライン強化学習における高い計算コストと低いサンプル効率の課題に取り組む。
• マルチタスク設定でパフォーマンスを低下させる、タスク間で共有されるポリシーパラメータの制限を克服する。
• モンテカルロ推定器がサンプル複雑性に劣るために、スパース報酬環境での価値推定の改善を図る。
• 推論コストの比例的増加なしに、大規模なマルチタスクデータセットからの高容量の知識吸収を可能にする。

提案手法

• Mixture-of-Experts (MoE) レイヤーを備えたスイッチトランスフォーマー構造を採用。ゲーティングネットワークが入力をタスク固有のエキスパートネットワークにルーティングし、スパースなパラメータ活性化と計算削減を実現する。
• 3つのトランスフォーマーヘッド（行動予測、ダイナミクスモデリング、リターン・トゥ・ゴー（RTG）予測）を統合し、軌道生成のための統一されたシーケンスモデリングフレームワークを構築する。
• 標準的な期待値推定を、軌道リターンのカテゴリカル分布をモデル化する分布価値推定器に置き換えることで、不確実性を考慮した評価の向上を図る。
• 分布価値推定値をガイドとして、生成された軌道のビームサーチを実施し、計画段階で高リターン行動シーケンスを選択する。
• オフラインマルチタスクデータセット上でエンドツーエンドにモデルを訓練。状態、行動、リターン、タスク識別子の全シーケンスを入力トークンとして活用する。
• スイッチトランスフォーマーの高容量を活かし、複数のタスクにわたる多様なエキスパートデモンストレーションを吸収しながら、計算効率を維持する。

実験結果

リサーチクエスチョン

• RQ1 スパースに活性化されるトランスフォーマー構造は、モデル容量やパフォーマンスを損なわず、マルチタスクオフラインRLにおける計算コストを低減できるか？
• RQ2 期待リターンだけでなく、リターンの分布をモデル化することで、スパース報酬設定におけるよりロバストで正確な価値推定が達成できるか？
• RQ3 スパース報酬環境におけるマルチタスク設定において、SwitchTT は標準的なTrajectory Transformerと比較して、パフォーマンスとトレーニング効率で優れているか？
• RQ4 スイッチトランスフォーマー構造は、オフラインRLにおいて、多様で異種のタスク間で効果的に表現を学習・共有できるか？

主な発見

• SwitchTT は、gym-mini-grid環境における10のスパース報酬タスクで、Trajectory Transformer よりも10%高いパフォーマンスを達成した。
• エキスパート計算のスパarsityのおかげで、オフラインモデルトレーニング時間を標準的なTrajectory Transformerと比較して最大90%短縮した。
• 分布価値推定器は、低サンプル数・スパース報酬状況での価値推定精度を顕著に向上させ、モンテカルロ推定器の悪いサンプル複雑性を緩和した。
• スイッチトランスフォーマー構造は、タスク分布が異なる場合でも、より良いパラメータ共有と知識移転を実現した。
• 分布価値推定値をガイドとしたビームサーチは、より信頼性の高い軌道計画と高い累積報酬をもたらした。
• 実験的結果から、期待値ではなく分布として価値をモデル化することで、マルチタスクオフラインRLにおけるポリシーのパフォーマンスと安定性が向上することが確認された。