新しいプロジェクトである`tinyworldformer`は、低コストのマイクロコントローラであるRaspberry Pi Pico 2 W上で直接ゲームを生成および実行可能なコンパクトなワールドモデルを実証しています。このシステムは、現在の状態とアクションに基づいて次のフレームと報酬を予測するワールドモデルと、生成されたフレームに基づいてアクションを選択するアクターネットワークという、2つの小さなニューラルネットワークを利用しています。浮動小数点演算ユニット(FPU)を備えたマイクロコントローラ上でこれらのモデルを実行することで、このプロジェクトは極めて制約のあるハードウェア上で複雑な生成タスクを実行する限界を探求しています。

ワールドモデルとアクターネットワークの間の相互作用を示すアニメーション
システムアーキテクチャは、状態予測のためのワールドモデルとアクション選択のためのアクターネットワークを特徴としています。github.com

この視覚資料は、ワールドモデルがフレームを予測し、アクターネットワークがアクションを選択する二重ネットワークアーキテクチャを説明しています。

核心となる革新は、「ループ型Transformer」アーキテクチャの使用にあります。複数の異なる層を積み重ねる代わりに、このモデルは単一のTransformerブロックを同じ重みで数回適用します。このアプローチは、マイクロコントローラにとって極めて重要なメモリオーバーヘッドを大幅に削減します。Pico 2 Wでパフォーマンスを維持するために、モデルは2つのループを使用し、実行可能なフレームレートを得るためにわずかな精度の低下を許容しています。ワールドモデルはさらに8ビット量子化(QAT)を通じて最適化されており、チップのFPUを活用するためにノルムと埋め込みを`float32`に保ちつつ、重みと活性化を`int8`に圧縮しています。これにより、モデルをデバイスの520 KBのSRAM内に収めることが可能になります。

このプロジェクトは、Breakout、Flappy Bird、そして「Doom-lite」と呼ばれるレイキャストによる一人称視点など、いくつかのゲームをサポートしています。Pac-Manのようなより複雑な環境では、高い精度を達成するためにモデルはより多くのパラメータ(最大131k)とより多くのループを必要とします。注目すべき技術的課題の一つは、モデルの出力が決定論的で論理的に一貫していることを保証することでした。例えば、開発者はBreakoutのボールのようなオブジェクトが一度だけレンダリングされることを保証するために「グローバル・シングルトン」ルールを実装しました。ワールドモデルは新しいシナリオを「夢見る」ことができますが、このプロジェクトには、これらの夢の環境内でプレイすることを学習できるストリーミング強化学習エージェントも含まれており、高度な手続き型生成のための基礎を提供しています。

ワールドモデルがBreakoutのゲームを生成している様子を示すGIF
ワールドモデルはBreakoutを「夢見て」プレイ可能なバージョンを生成します。github.com

これは、マイクロコントローラ上でBreakoutを生成および実行するプロジェクトの能力を示しています。

ワールドモデルがDoom-liteの一人称視点を生成している様子を示すGIF
モデルは「Doom-lite」として知られるレイキャストによる一人称視点を生成します。github.com

これは、tinyworldformerプロジェクトによってサポートされるより複雑な「Doom-lite」環境を例示しています。

ワールドモデルがFlappy Birdのゲームを生成している様子を示すGIF
モデルはFlappy Birdのバージョンを生成し、実行します。github.com

これは、tinyworldformerシステムによってサポートされる主要なゲームの一つを示しています。