LLMでの「世界モデル（簡易版）」を自作してみました

＊＊＊＊＊＊＊＊

※LINE対応チャットボット版の
「LINEチャットボット屋」
いろんなチャットボットがあります。
ぜひ、ご覧ください！

＊＊＊＊＊＊＊＊＊＊＊＊＊＊＊

昨今のAIは文字の生成や画像、動画の生成は目を見張るものがあります。

しかし、人間のように物理世界を理解して動くのはまだまだ苦手です。

例えば、人は前から車が来たら、「どっちに来るかな。来ない方に避けよう」と思うでしょう。
AIにはそれがまだ理解できない。

単純に人間社会のさまざまな現象を学習していないからです。

これらを理解し、頭の中でシミュレーションできるのが世界モデル。

ここでは、開発環境に制限があるので、簡易的な世界モデルを実装してみようと。

そんな好奇心から始まったプロジェクトをご紹介します。

Google Colab上でTransformerを一から実装し、環境のルールを学習させ、脳内シミュレーションだけで迷路を突破する。

最初は何も知らなかったAIが、物理法則を学び、不確実性を理解し、最後にはゴールへ駆け込むまでの進化の記録を共有します。

技術的な工夫と、AIの成長プロセスをぜひご覧ください。

世界モデル（World Model）とは？

今回実装したのは、「世界モデル」と呼ばれるAIアーキテクチャです。
簡単に言えば、AIの中に「現実世界のシミュレーター（脳）」を作ることです。

従来のAI：画面を見て、反射的にボタンを押す。
世界モデル：画面を見て、「もし右へ行ったらどうなるか？」を脳内で予測（Dreaming）してから、最善の手を選ぶ。

今回は、この脳の部分にLLM（大規模語言モデル）の基盤技術であるTransformerを使用しました。
「次の単語」を予測するように、「次の画面（盤面）」を予測させるアプローチです。

第1の進化物理法則の習得と「夢」の可視化

最初のステップは、AIに「この簡易ゲーム世界（Minigrid環境）」のルールを教えること。

5×5のシンプルな部屋でエージェントをランダムに動かし、そのデータをTransformerに学習させます。

課題と解決：データの可視化

実装初期、学習は順調に進んでいる（Lossが低下している）のに、生成された予測画像が真っ黒に見えるという現象に直面。

これはMinigridのデータ形式が「RGB画像」ではなく「ID（小さな整数値）」で構成されていたためでした。

ここで「値を30倍にスケーリングして可視化する」という調整に変更。

その結果、真っ黒だった画面から、AIが描いた鮮明な「夢」が浮かび上がりました。
AIは壁の位置や、自分の動きに合わせて景色が変わる法則を、完璧に理解していたのです。

第2の進化不確実性への適応

次なる挑戦は、難易度を一気に上げた「動くモンスター（青い丸）」がいる環境への適応です。
自分（エージェント）の動きは決まっていますが、モンスターはランダムに動きます。

確率的な未来予測

ここで非常に興味深い現象が観測されました。AIが生成した数手先の予測画像が、ぼんやりと「分身」したようになったのです。
これはAIの混乱ではなく、その事象に合った適合でした。

「右に来るかもしれないし、左かもしれない」

AIは不確実な未来に対し、一点に賭けるのではなく、可能性の重ね合わせとして確率的に世界を捉えるようになりました。
この「ぼやけた影」こそが、AIがランダム性を理解した証拠でしょう。

第3の進化生存戦略から「ゴール」へ

世界を理解したAIに、いよいよゲームをプレイさせます。

手法はモデル予測制御（MPC）。
行動する前に脳内で複数のパラレルワールドをシミュレーションし、ベストな未来を選ぶ方法です。

段階的な最適化プロセス

Phase 1: 慎重な生存戦略
最初に実装したプランナーは「生存」を最優先しました。
予測画像内のエージェントの有無を判定基準にしたところ、AIは「絶対に死なない動き」をマスター。
しかし、リスクを回避するあまり、安全地帯から動こうとしない「慎重さ」も見られました。

Phase 2: ロジカルな評価への転換

より積極的な動きを引き出すため、評価関数をアップデート。
「色のピクセル数」を見る曖昧な判定から、「オブジェクトIDと座標計算」による論理的な判定へ変更。

「モンスターから2マス以上離れているか？」といった正確な状況判断が可能になり、AIの動きに迷いがなくなりました。

Phase 3 Hero Mode

最後に、ゴールへ向かう「意志」を強化。

探索空間の最適化：「ドアを開ける」などの不要な行動選択肢を排除し、移動のみに集中。
思考の深化：シミュレーション回数を200回に増加し、読みの深さも強化。
報酬設計：ゴール到達へのインセンティブを最大化し、恐怖を乗り越える動機付けを強化。

最終結果 33ステップの完全勝利

すべての調整を終えた「Hero Mode」での実行結果です。

制限時間が100ステップある中、AIはわずか33ステップでゴールに到達。

生成されたGIFアニメーションには、迫りくる青いモンスターを最小限の動きで華麗にかわし、最短ルートで緑色のゴールへ突き進むエージェントの姿が記録されていました。

それはもはや単純なゲームの攻略ではなく、自ら思考し、目的を達成する真の姿です。

まとめ

今回のプロジェクトを通じて、特別なスーパーコンピュータがなくても、「世界を理解し、自ら行動計画を立てるAI」を簡易的に実装できることがわかりました。

見る：状態をトークンとして認識する
学ぶ：物理法則と不確実性を学習する
夢見る：未来をシミュレーションする
叶える：最適な行動を選択し、現実を変える

この4つのステップは、ゲーム攻略だけでなく、ロボティクスや自動運転など、あらゆるAIエージェントの基本となる技術です。

今回はごく単純なゲームで学ぶ世界モデルでしたが、これの規模を拡大して、さらなる発展系の「世界モデル」も作っていくつもりです。

LLMの中で「世界モデル」が十分に認識できてくれば、AGIへの進化もまた一歩近くなるでしょう。

今回使用したPythonコード

＊表示のバグでインテントが崩れています。
コピーしてお気に入りのAIに直してもらってください。

全部載せるとアレなので、初期の一部です。
ご参考までに。


!pip install gymnasium minigrid torch numpy matplotlib

import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym
import minigrid
from minigrid.wrappers import ImgObsWrapper
import numpy as np
import matplotlib.pyplot as plt
import random

# --- 設定 ---
ENV_ID = "MiniGrid-Empty-5x5-v0" # シンプルな5x5の部屋
SEQ_LEN = 10 # 学習するシーケンス長
BATCH_SIZE = 32
EMBED_DIM = 64
N_HEAD = 4
N_LAYER = 2
LR = 1e-3
EPOCHS = 500
DATA_COUNT = 1000 # 収集するステップ数

# デバイス設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# --- 環境とトークン化の準備 ---
# グリッドの各セル(3チャンネルの値)を1つの整数トークンに変換する簡易ハッシュ関数
def state_to_tokens(state):
# state shape: (H, W, 3) -> 平坦化
return state.flatten().tolist()

# 逆変換（可視化用）
def tokens_to_state(tokens):
# 5x5グリッド, 3チャンネル (Minigridの仕様に合わせる)
return np.array(tokens).reshape(7, 7, 3).astype(np.uint8)

# 環境テストと次元確認
env = gym.make(ENV_ID, render_mode="rgb_array")
env = ImgObsWrapper(env)
obs, _ = env.reset()
obs_dim = len(state_to_tokens(obs)) # トークン列の長さ (7*7*3 = 147)
vocab_size = 256 # ピクセル値やIDは0-255の範囲
act_vocab_size = env.action_space.n

print(f"Observation Token Length: {obs_dim}")
print(f"Action Space: {act_vocab_size}")

# --- データセット収集 (Data Collection) ---
print("Collecting trajectories...")
dataset = []
buffer = []

obs, _ = env.reset()
for _ in range(DATA_COUNT):
action = env.action_space.sample() # ランダムに行動
next_obs, reward, terminated, truncated, _ = env.step(action)

# データを保存: [現在の状態トークン列, 行動]
tokens = state_to_tokens(obs)
dataset.append((tokens, action, state_to_tokens(next_obs)))

obs = next_obs
if terminated or truncated:
obs, _ = env.reset()

print(f"Collected {len(dataset)} transitions.")

# --- Transformer 世界モデル定義 ---
class WorldModel(nn.Module):
def __init__(self, obs_len, vocab_size, act_vocab_size, embed_dim, n_head, n_layer):
super().__init__()
self.obs_len = obs_len

# 埋め込み層
self.token_embedding = nn.Embedding(vocab_size, embed_dim)
self.action_embedding = nn.Embedding(act_vocab_size, embed_dim)
self.pos_embedding = nn.Parameter(torch.zeros(1, obs_len + 1, embed_dim)) # +1はAction用

# Transformer Decoder
encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=n_head, batch_first=True)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layer)

# 出力ヘッド: 次の状態の各トークンを予測
self.head = nn.Linear(embed_dim, vocab_size)

def forward(self, obs_tokens, action):
# obs_tokens: (Batch, Obs_Len)
# action: (Batch, 1)

# 状態の埋め込み
obs_emb = self.token_embedding(obs_tokens) # (B, Obs_Len, Dim)

# 行動の埋め込み
act_emb = self.action_embedding(action).unsqueeze(1) # (B, 1, Dim)

# 入力シーケンス結合: [状態..., 行動]
x = torch.cat([obs_emb, act_emb], dim=1) # (B, Obs_Len + 1, Dim)

# 位置エンコーディング加算
x = x + self.pos_embedding[:, :x.size(1), :]

# Transformer処理
x = self.transformer(x)

# 次の状態予測 (行動トークン部分を使って次の全状態を予測するのは簡易版として、
# ここではTransformerの出力全体から次の状態を回帰させる簡易構成をとります)
# 本来はAutoregressiveに1トークンずつ予測しますが、今回は一括予測で簡略化します。

# 入力の各位置の情報を使って、「次の状態の各位置」を予測するよう学習させるため、
# 状態部分の出力を使います。
logits = self.head(x[:, :self.obs_len, :]) # (B, Obs_Len, Vocab)
return logits

model = WorldModel(obs_dim, vocab_size, act_vocab_size, EMBED_DIM, N_HEAD, N_LAYER).to(device)
optimizer = optim.Adam(model.parameters(), lr=LR)
criterion = nn.CrossEntropyLoss()

# --- 学習ループ ---
print("Starting training...")
loss_history = []

for epoch in range(EPOCHS):
# バッチ作成
batch = random.sample(dataset, BATCH_SIZE)

obs_batch = torch.tensor([b[0] for b in batch], dtype=torch.long).to(device)
act_batch = torch.tensor([b[1] for b in batch], dtype=torch.long).to(device)
next_obs_batch = torch.tensor([b[2] for b in batch], dtype=torch.long).to(device)

optimizer.zero_grad()

# 予測
pred_logits = model(obs_batch, act_batch) # (B, Obs_Len, Vocab)

# Loss計算 (Flattenして計算)
loss = criterion(pred_logits.reshape(-1, vocab_size), next_obs_batch.reshape(-1))

loss.backward()
optimizer.step()

loss_history.append(loss.item())

if epoch % 50 == 0:
print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

plt.plot(loss_history)
plt.title("World Model Training Loss")
plt.show()

# --- ドリーミング (Dreaming/Simulation) ---
# 実際の環境を使わず、学習したモデルだけで未来を予測するテスト
print("Dreaming future steps...")

# 初期状態
env.reset()
start_obs, _, _, _, _ = env.step(env.action_space.sample())
current_tokens = torch.tensor([state_to_tokens(start_obs)], dtype=torch.long).to(device)

fig, axes = plt.subplots(1, 6, figsize=(15, 3))
axes[0].imshow(start_obs)
axes[0].set_title("Start (Real)")
axes[0].axis('off')

# 5ステップ先まで「夢」を見る
for i in range(5):
# ランダムな行動を決定（ポリシーがあればそれを使う）
action = torch.tensor([env.action_space.sample()], dtype=torch.long).to(device)

# 世界モデルで次の状態を予測
with torch.no_grad():
pred_logits = model(current_tokens, action)
# 最も確度の高いトークンを選択 (Greedy decoding)
next_tokens = torch.argmax(pred_logits, dim=2)

# 予測された状態を画像に戻して表示
img = tokens_to_state(next_tokens.cpu().numpy()[0])
axes[i+1].imshow(img)
axes[i+1].set_title(f"Dream Step {i+1}")
axes[i+1].axis('off')

# 現在の状態を更新（現実の観測は使わない！）
current_tokens = next_tokens

plt.show()
print("Visualization: 左端が実際の初期状態。右へ続く画像は、モデルが脳内で想像した未来の遷移です。")

# --- ドリーミング (Dreaming/Simulation) ---
print("Dreaming future steps...")

# 初期状態
env.reset()
start_obs, _, _, _, _ = env.step(env.action_space.sample())
current_tokens = torch.tensor([state_to_tokens(start_obs)], dtype=torch.long).to(device)

fig, axes = plt.subplots(1, 6, figsize=(15, 3))

# ★修正ポイント: 画像を30倍して表示（IDを見える明るさにするため）
axes[0].imshow(start_obs * 30)
axes[0].set_title("Start (Real)")
axes[0].axis('off')

# 5ステップ先まで「夢」を見る
for i in range(5):
action = torch.tensor([env.action_space.sample()], dtype=torch.long).to(device)

with torch.no_grad():
pred_logits = model(current_tokens, action)
next_tokens = torch.argmax(pred_logits, dim=2)

img = tokens_to_state(next_tokens.cpu().numpy()[0])

# ★修正ポイント: ここも30倍します
axes[i+1].imshow(img * 30)
axes[i+1].set_title(f"Dream Step {i+1}")
axes[i+1].axis('off')

current_tokens = next_tokens

plt.show()
print("Visualization: 明るく補正しました。赤いブロック（エージェント）などが動くのが見えるはずです。")

＊＊＊＊＊＊＊＊＊＊＊＊＊＊＊＊

最近のデジタルアート作品を掲載！

X　旧ツイッターもやってます。
https://x.com/ison1232

インスタグラムはこちら
https://www.instagram.com/nanahati555/

＊＊＊＊＊＊＊＊＊＊＊＊＊＊

Follow me!