Transformers are Sample Efficient World Modelsを読んだメモ

論文

出典

概要

　Discrete autoencoderでRGB画像を埋め込み、自己回帰Transformerで世界モデルとして学習させる。

　Atariゲームにおける2時間分のプレイデータで、10ゲーム/26ゲームにおいて人間のスコアに達した。

メモ

  • 実環境との作用経験はモデルの学習にしか使わない
  • 以下の3ステップを繰り返す
    1. 経験の収集
    2. 世界モデルの更新
    3. 方策・価値関数の更新

Figure 1

  • 緑線E : Encoder
  • 緑先D : Decoder
  • 紫先π : 方策モデル
    • 方策モデルは、デコーダで復元したRGB画像を入力とする
    • 離散化表現の方を入力として学習しても良さそうだけど、なんで復元している？
    • Conclusionでちょっと触れられているけど、特に強い意味はなさそう
  • 青線G : 環境モデル
    • GPT的な自己回帰Transformerに突っ込むだけ
    • 離散化表現トークンは1つのタイムステップtの中に複数あるので、それを先頭から順番に予測していく

実験

　Atari 100kという、100k分までの行動が許される設定がわりとメジャー？ なものとしてあるらしい。

Table1

　MedianでSPRという手法に負けているのは少し気になるが、26種のゲームのHuman-normalizedスコアのMedianという意味ならそこまで大きな問題ではないか。

　PongやBreakoutなど画面自体の切り替わりがないゲームには強く、FrostbiteやKrullなど複数レベル・ステージを持つようなゲームには弱い。

付録部分

　パラメータの具体的な数値など

Table 3

　計算資源

For each Atari environment, we repeatedly trained IRIS with 5 different random seeds. We ran our experiments with 8 Nvidia A100 40GB GPUs. With two Atari environments running on the same GPU, training takes around 7 days, resulting in an average of 3.5 days per environment.