SheepSebは、2つの小型ニューラルネットワークでゲームの生成とプレイを同時に行う「tinyworldformer」をRaspberry Pi Pico 2 W上で動作させ、GitHubで公開しました。ワールドモデルは26k〜131kパラメータで、Flappy Birdでは2ループ時に完全一致フレーム率100%を記録しています。
何が新しいか
tinyworldformerは、同じ重みのトランスフォーマーブロックを1〜4回繰り返すループ型ワールドモデルと、生成フレームから行動を選ぶアクターネットワークで構成されます。ワールドモデルは現在のフレームとアクションから次のフレームと報酬を予測し、アクターは生成されたフレームだけを見てアクションを選びます。Pico上では両ネットワークともゲームの実装を直接参照せず、ゲームの「夢」を進行させます。ゲーム全体の状態は小さなトークン化フレームに描画され、Breakoutではボールの向きをトークンに含め、Doom-liteにはミニマップを描きます。
開発者は、重みの再利用はメモリ面で安価なため、ループ型ワールドモデルならマイコンで動くほど小さくできるかを試したと述べています。Pico 2 Wを選んだ理由は、安価でFPUを備えていることです。
仕様・数値
| 項目 | ワールドモデル | アクターネットワーク |
|---|---|---|
| 構成 | ループ型トランスフォーマー、int8 | 2隠れ層のMLP、float32 |
| パラメータ数 | 26k〜131k | 9k〜16k |
| フレーム当たり計算量 | 2〜8M multiply-adds | 9k〜16k multiply-adds |
| 学習 | 記録した実ゲームの遷移 | Stream AC(λ)、実ゲームまたは夢の中 |
ゲームごとの構成と完全一致フレーム率(2ループ)は次のとおりです。完全一致フレーム率は、保留した遷移のうち予測フレームの全セルが実ゲームと一致した割合で、ランダムイベント直後のフレームは除外しています。
| ゲーム | フレーム | パラメータ | MACs | 完全一致率 |
|---|---|---|---|---|
| Breakout | 20×16 | 26k | 4.3M | 99.99% |
| Doom-lite | 32×32 | 55k | 4.2M | 83.6% |
| Flappy Bird | 24×24 | 46k | 2.0M | 100% |
| Pong | 20×32 | 46k | 2.3M | 99.45% |
| Pac-Man lite | 48×48 | 131k | 8.2M | 86.4% |
Doom-liteの3Dビューは768セルあり、陰影付きの壁ピクセルが1つずれるだけでフレーム全体が不一致になるため、原文はゲーム状態自体はほぼ常に正しいと説明しています。ループ回数を1→2→3→4と増やすと、Doom-liteの一致率は59.8%→83.6%→85.9%→84.7%と推移します。Pico上ではフレームレートを保つため2ループで動かします。
学習データはゲームごとに1〜3 million examplesで、ワールドモデルはセル単位のクロスエントロピー損失で30k steps学習します。最初の3〜4k stepsをfloatで回した後、int8丸めを有効にするQATに切り替えます。量子化では全行列をint8化し、出力行ごとに1つのスケールを使います。
Pico 2 Wでの実行と試し方
- 両ネットワークはRP2350のcore 1で動き、core 0がUSBとディスプレイを処理します。
- 事前学習済みアクターは、ボード上でゼロから学習するstream-xエージェントに置き換えられます。Pico上のポリシー/価値ネットワークは隠れ層32ユニットです。
- ボタンを押すと操作をユーザーに移せます。
- Rust実装は
tinyworldformer-rsにあり、int8重みでPico向けファームウェアを動かします。int8重みはcargo testでPythonの整数仕様とビット単位で一致することを確認しています。 - 実ゲームはターミナルで
uv run python -m tinyworldformer.games.play --game pacmanで遊べます。ゲーム名にはbreakout、doomlite、flappy、pongも指定できます。 - Doom-liteのノイズ拡張は
train --noise 0.3で実行します。
夢の中の抜け穴対策
夢の中だけで学習したDoom-liteエージェントは、実際の最大値56に対し夢の中で210点を出し、実ゲームでは28点にとどまりました。開発者は、エージェントで実ゲームを記録してワールドモデルを再学習するラウンドを重ね、エージェントが見つけた抜け穴を一つずつ塞いだと述べています。ホスト上では、Rustのint8 Flappy Birdの夢の中で新規エージェントが25k stepsで平均6本、100kで121本、250kで300本以上のパイプを通過しました。リポジトリのGitHubスター数は84、フォーク数は3です。
まだコメントはありません。