[論文レビュー] Masked World Models for Visual Control

Younggyo Seo, Danijar Hafner+5 more|arXiv (Cornell University)|Jun 28, 2022

Advanced Vision and Imaging被引用数 10

ひとこと要約

本稿では、畳み込み特徴マスク自動エンコーダと視覚変換器を用いて微細な視覚的表現を学習する、視覚的表現とダイナミクス学習を分離したモデルベース強化学習フレームワークであるマスクドワールドモデル（MWM）を提案する。さらに、計画に使用する別個の潜在ダイナミクスモデルを採用する。タスク関連の詳細をエンコードするための補助報酬予測を導入することで、MWMは、50のメタワールド視覚的操作タスクにおいて81.7%の成功率を達成し、以前のSOTAベースライン（67.9%）を著しく上回る最先端の性能を達成した。

ABSTRACT

Visual model-based reinforcement learning (RL) has the potential to enable sample-efficient robot learning from visual observations. Yet the current approaches typically train a single model end-to-end for learning both visual representations and dynamics, making it difficult to accurately model the interaction between robots and small objects. In this work, we introduce a visual model-based RL framework that decouples visual representation learning and dynamics learning. Specifically, we train an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features, and learn a latent dynamics model that operates on the representations from the autoencoder. Moreover, to encode task-relevant information, we introduce an auxiliary reward prediction objective for the autoencoder. We continually update both autoencoder and dynamics model using online samples collected from environment interaction. We demonstrate that our decoupling approach achieves state-of-the-art performance on a variety of visual robotic tasks from Meta-world and RLBench, e.g., we achieve 81.7% success rate on 50 visual robotic manipulation tasks from Meta-world, while the baseline achieves 67.9%. Code is available on the project website: https://sites.google.com/view/mwm-rl.

研究の動機と目的

• ロボットと小さな物体の間の相互作用を視覚的制御タスクで正確にモデル化する課題に対処すること。
• エンドツーエンド学習における視覚的ワールドモデルの限界を克服すること、特に表現品質とダイナミクスの正確さのトレードオフを解消すること。
• 視覚的表現とダイナミクス学習の分離により、視覚的モデルベース強化学習におけるサンプル効率と漸近的性能を向上させること。
• 畳み込み特徴マスクを用いた自己教師付き表現学習が、ピクセルパッチマスクを上回る微細な視覚的詳細の捉え方を可能にするかどうかを検証すること。
• 補助報酬予測が、より良い下流ポリシー学習のためのタスク関連情報で視覚的表現を強化するかどうかを調査すること。

提案手法

• ピクセルパッチの代わりにマスクされた畳み込み特徴を用いて、視覚変換器ベースの自動エンコーダを訓練し、視覚的観測を再構築する。
• 自動エンコーダが抽出した表現を処理する別個の潜在ダイナミクスモデルを用いて、将来の状態を予測する。
• 自動エンコーダの学習中に補助報酬予測ヘッドを導入し、視覚的表現にタスク関連情報を取り入れる。
• 環境との相互作用中に収集したオンライン経験を用いて、自動エンコーダとダイナミクスモデルを継続的に更新する。
• 学習プロセスを分離する：まず再構築と報酬予測に基づいて自動エンコーダを更新し、その後、固定された自動エンコーダ特徴を用いてダイナミクスモデルを更新する。
• 自己教師付き事前学習中に段階的にマスク率を増加させるカリキュラム学習を適用し、安定性と特徴品質を向上させる。

実験結果

リサーチクエスチョン

• RQ1 視覚的表現学習とダイナミクス学習の分離は、視覚的モデルベース強化学習における性能向上に寄与するか？
• RQ2 ピクセルパッチの代わりに畳み込み特徴をマスクすることで、視覚的制御タスクにおける微細な視覚的詳細の捉え方が向上するか？
• RQ3 表現学習中に補助報酬予測を導入することで、下流ポリシー学習のための視覚的特徴の品質が向上するか？
• RQ4 MWMはエンドツーエンドワールドモデルや他の分離ベースラインと比較して、サンプル効率と最終性能の両面で優れているか？
• RQ5 あるタスクで事前学習された表現は、未観測の操作タスクへどの程度一般化可能か？

主な発見

• MWMは、メタワールドの50の視覚的ロボット操作タスクにおいて81.7%の成功率を達成し、以前のSOTAベースライン（67.9%）を著しく上回った。
• アブレーションスタディの結果、報酬予測を含むMWMは、それなしの同じ設定よりも優れた性能を示し、表現学習におけるタスク関連信号の有効性を確認した。
• プッシュタスクで事前学習した固定された表現は、ピックアンドプレース、プッシュバック、ドアラー開閉など未観測タスクへも良好に一般化され、MWMを初期化から学習し直した場合と同等またはそれ以上の性能を示した。
• 状態回帰分析により、報酬予測を用いて学習された表現は、体感的状態を予測する際の誤差が低くなることが確認され、特徴品質の向上が示された。
• コントラスト表現学習を用いたDreamerV2と比較して、MWMは特にピックアンドプレースのような複雑なタスクで優れた性能を示し、分離型でマスクされた特徴アプローチの利点を裏付けた。
• ImageNet分類では類似した性能を示すものの、視覚的制御タスクではピクセルパッチマスクよりも畳み込み特徴マスクが優れた性能を示し、タスク固有の利点が顕著に現れた。