naoyabone、5x5 Grid World を用いたメモリ付き世界モデルの実装解説を公開──GRU と RSSM による部分観測問題の解決手法
5×5 Grid World と教育用コードを用いて、部分観測環境における GRU や RSSM、Pixel MSE の課題と実装上の工夫を体系的に解説した。
リリース: 2026-08-23 · 読了 3 分記事の要約
1. 核心（What）
- 5×5のGrid World環境を題材に、世界モデルにおけるMemory（GRU）の必要性をコードベースで解説している。
- 部分観測環境（周囲3×3の視界）において、同じ画像から過去の推移によって状態が異なるalias問題を定義している。
- RSSMのprior（予測）とposterior（修正）の役割や、KL divergenceを通じた確率的状態の伝播を整理している。
- Pixel MSEにおける背景への過剰適応（shortcut）を防ぐため、Agentの25クラス分類によるposition lossの導入を示している。
2. 影響（Why）
- 世界モデルの設計判断基準: 部分的観測や長期rolloutにおける誤差蓄積を検証するミニマルな実装パターンを学べるため、自前でモデルを構築・デバッグする際のベースライン設計に直結する。
- 国内の強化学習・ロボティクス検証への影響: [国内のロボティクスシミュレーションやエージェント開発] を行う小規模チームは、大規模な環境構築の前に、alias問題やlossのshortcutを防ぐための基本設計の検証用リファレンスとして活用できる。
3. 根拠・詳細（How）
- GRU と RSSM の状態管理機構: GRUCellを用いて過去の要約メモ h_t を保持しつつ、RSSM構造により画像特徴 e_t との比較を通じて posterior から prior への分布推定をKL divergenceで整合させる設計を採用している。
- Position Loss による対策: 画像全体の平均二乗誤差（Pixel MSE）に依存すると小さなオブジェクトが無視されるため、Agentの位置を25クラスの分類問題として直接学習させる position loss を導入している。
4. 展望・課題（Next）
- 比較検証の予定: 同一データ・同一条件でのMemoryなしモデルとGRUの予測精度比較（03_memory/05_comparison）の実施が予定されている。