Raspberry Pi Pico 2 W上でループ型ワールドモデルとRLアクターを実行するTinyWorldFormer

tinyworldformerはSheepSebによるオープンソースプロジェクトで、Raspberry Pi Pico 2 W上で2つの小さなニューラルネットワークを実行します。1つはゲームフレームを生成するワールドモデル、もう1つはそれをプレイする事前学習済みアクターです。どちらのネットワークもチップ上で実際のゲームを見ることはありません。アクターはワールドモデルが生成したフレームのみを読み取り、ワールドモデルはアクターからの行動のみを受け取ります。

アーキテクチャ

• ワールドモデル： int8量子化されたループ型トランスフォーマー。26k～131kパラメータ、フレームあたり2～8M乗加算。入力は現在のフレーム＋行動、出力は次のフレーム＋報酬。
• アクター： 2つの隠れ層を持つMLP、float32。9k～16kパラメータ、フレームあたり9k～16k乗加算。入力はフレーム、出力は行動。
• どちらもRP2350のコア1で実行され、コア0はUSBとディスプレイを処理します。
• 事前学習済みアクターは、ボード上でゼロから学習する新しいstream-x強化学習エージェントに差し替え可能です。ボタンを押すと人間のプレイヤーに操作が移ります。

ゲームとベンチマーク結果

各ゲームはsrc/tinyworldformer/games/内の数百行のNumPyコードで、トークンIDの小さなグリッドにレンダリングされます。1フレームにゲーム状態全体が保持されるため、隠れ状態はフレーム内に描画されます（ボールの進行方向はトークンの一部であり、Doom-liteにはミニマップがあります）。

• Breakout — 20×16フレーム、26kパラメータ、4.3M MAC、99.99%の完全一致フレーム。
• Doom-lite — 8×8迷路の32×32レイキャストビュー、55kパラメータ、4.2M MAC、83.6%の完全一致フレーム（1つの陰影付き壁ピクセルでもフレーム全体が不一致とカウントされます）。
• Flappy Bird — 24×24フレーム、46kパラメータ、2.0M MAC、100%の完全一致フレーム。
• Pong — 20×32フレーム、46kパラメータ、2.3M MAC、99.45%の完全一致フレーム。
• Pac-Man lite — 48×48フレーム、23×21迷路、249個のペレット、131kパラメータ、8.2M MAC、86.4%の完全一致フレーム。

「完全一致フレーム」とは、予測フレームのすべてのセルが実際のゲームと一致するホールドアウト遷移の割合です（ランダムイベント直後のフレームは除外されます）。

ワールドモデルの仕組み

フレームは正方形のパッチに分割されます（Breakoutでは2×2セル、Doom-lite、Flappy Bird、Pongでは4×4、Pac-Manでは6×6）。これによりゲームあたり36～80トークンになります。各パッチ埋め込みは、セルごとに1つの学習済みベクトルを、パッチ内の位置ごとに別のテーブルで合計し、さらに位置ベクトルと行動ベクトルを加算します。モデルは単一のトランスフォーマーブロック（アテンション＋ReLU MLP）で、1組の重みを複数回適用するループ型トランスフォーマーであり、マイクロコントローラーに収まるほどメモリを抑えます。

実際のゲームを実行する

uv run python -m tinyworldformer.games.play --game pacman

breakout、doomlite、flappy、またはpongに置き換えてください。リポジトリにはtinyworldformer-rs、ドキュメント、テスト、ブラウザデモも含まれています。

エッジ推論や小さなワールドモデルに取り組む開発者にとっての実用的なポイントは、メモリと計算のトレードオフです。1つのトランスフォーマーブロックの重みを再利用することで、26kパラメータのモデルがフレームあたり数百万MACで動作し、Pico 2 Wの制限内に収まります。