3. 世界模型预测下一帧图像

任务说明

构建并训练一个轻量级世界模型：给定当前帧图像与动作指令，预测下一帧图像。使用仿真环境生成轨迹数据，以像素级重建误差和结构相似性作为预测质量评估指标。

环境与依赖

pip install torch torchvision numpy matplotlib scikit-image

数据生成

使用简单的移动方块环境：每帧为32×32灰度图，方块位置由动作（dx, dy∈{-1,0,1}）驱动，边界处方块不动。共生成8000条训练轨迹和2000条测试轨迹，每条轨迹长度5帧。

模型实现

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from skimage.metrics import structural_similarity as ssim

# ---------- 数据生成 ----------
def generate_trajectory(num_steps=5, img_size=32):
    """生成单条轨迹：返回 (frames, actions)"""
    frames = []
    actions = []
    x, y = np.random.randint(4, img_size-4, size=2)
    for _ in range(num_steps):
        img = np.zeros((img_size, img_size), dtype=np.float32)
        img[max(0,y-2):y+3, max(0,x-2):x+3] = 1.0
        frames.append(img)
        dx, dy = np.random.randint(-1, 2, size=2)
        actions.append([dx, dy])
        x = min(max(x+dx, 2), img_size-3)
        y = min(max(y+dy, 2), img_size-3)
    return np.stack(frames), np.array(actions, dtype=np.float32)

def make_dataset(n_traj):
    frames_list, actions_list = [], []
    for _ in range(n_traj):
        frames, actions = generate_trajectory()
        frames_list.append(frames[:-1])  # 输入帧 t=0..3
        actions_list.append(actions[:-1])
    # 输入: (N, T, 1, H, W), 动作: (N, T, 2), 目标: (N, T, 1, H, W)
    X = torch.tensor(np.stack(frames_list)[:, :, None]).float()
    A = torch.tensor(np.stack(actions_list)).float()
    Y = torch.tensor(np.stack(frames_list)[:, 1:, None]).float()
    return X, A, Y

train_X, train_A, train_Y = make_dataset(8000)
test_X, test_A, test_Y = make_dataset(2000)

# ---------- 世界模型 ----------
class WorldModel(nn.Module):
    def __init__(self, img_size=32, action_dim=2, hidden_dim=64):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((8, 8))
        )
        self.fc = nn.Linear(32*8*8 + action_dim, hidden_dim)
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim, 32*8*8),
            nn.ReLU(),
            nn.Unflatten(1, (32, 8, 8)),
            nn.Upsample(scale_factor=2, mode='bilinear'),
            nn.Conv2d(32, 16, 3, padding=1), nn.ReLU(),
            nn.Upsample(scale_factor=2, mode='bilinear'),
            nn.Conv2d(16, 1, 3, padding=1),
            nn.Sigmoid()
        )

    def forward(self, obs, act):
        # obs: (B, 1, H, W), act: (B, 2)
        z = self.encoder(obs).flatten(1)
        z = torch.cat([z, act], dim=1)
        h = torch.relu(self.fc(z))
        return self.decoder(h)

model = WorldModel()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

# ---------- 训练 ----------
def train_step(x, a, y):
    # x: (B, T, 1, H, W) -> 按时间步展开
    B, T = x.shape[0], x.shape[1]
    total_loss = 0
    for t in range(T):
        pred = model(x[:, t], a[:, t])
        total_loss += loss_fn(pred, y[:, t])
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()
    return total_loss.item() / T

model.train()
for epoch in range(5):
    epoch_loss = 0
    for i in range(0, len(train_X), 64):
        batch_x = train_X[i:i+64]
        batch_a = train_A[i:i+64]
        batch_y = train_Y[i:i+64]
        epoch_loss += train_step(batch_x, batch_a, batch_y)
    print(f"Epoch {epoch+1}, Loss: {epoch_loss / (len(train_X)//64):.4f}")

# ---------- 评估 ----------
model.eval()
mse_total = 0
ssim_total = 0
count = 0
with torch.no_grad():
    for i in range(0, len(test_X), 64):
        bx = test_X[i:i+64]
        ba = test_A[i:i+64]
        by = test_Y[i:i+64]
        B, T = bx.shape[0], bx.shape[1]
        for t in range(T):
            pred = model(bx[:, t], ba[:, t]).cpu().numpy()
            true = by[:, t].cpu().numpy()
            for j in range(B):
                p = pred[j, 0]
                g = true[j, 0]
                mse_total += np.mean((p - g) ** 2)
                ssim_total += ssim(g, p, data_range=1.0)
                count += 1

print(f"Test MSE: {mse_total/count:.6f}")
print(f"Test SSIM: {ssim_total/count:.4f}")

# ---------- 可视化示例 ----------
import matplotlib.pyplot as plt
idx = 0
fig, axes = plt.subplots(2, 4, figsize=(8, 4))
with torch.no_grad():
    for t in range(4):
        pred = model(test_X[idx:idx+1, t], test_A[idx:idx+1, t])[0, 0].cpu().numpy()
        axes[0, t].imshow(test_Y[idx, t, 0].cpu().numpy(), cmap='gray')
        axes[0, t].set_title(f'True t+{t+1}')
        axes[1, t].imshow(pred, cmap='gray')
        axes[1, t].set_title(f'Pred t+{t+1}')
        axes[0, t].axis('off'); axes[1, t].axis('off')
plt.tight_layout()
plt.savefig('prediction_visual.png', dpi=100)
print("可视化已保存至 prediction_visual.png")

运行与验证

1. 将上述代码保存为 world_model.py 并执行 python world_model.py。
2. 训练过程应显示损失逐步下降（初始约0.08，最终低于0.02）。
3. 测试输出中，MSE应低于0.01，SSIM应高于0.90。
4. 打开 prediction_visual.png，对比真实与预测帧，方块位置应基本吻合。

通过标准：测试集SSIM > 0.90 且 MSE < 0.01，同时可视化图中预测方块与真实方块位置偏差不超过2像素。