☰
世界模型驱动的强化学习:从想象 rollout 到 CartPole 实战
2026/10/27 10:45:45 网站建设 项目流程

做强化学习的朋友大多有过类似的纠结:算法本身不难写,难的是让它在哪儿跑。真机上撞一次墙可能损失几千块,仿真里参数调偏了结果又迁移不过去,LLM 智能体更麻烦,每一步都要真调接口、真读写文件,试错成本按 token 计费。这篇文章就围绕“世界模型”这条路展开,讲清楚它为什么能让强化学习算法在“不碰真实环境”的前提下完成训练,并给出一套可以直接跑起来的最小实现。读者对象是有 Python 和 PyTorch 基础、了解强化学习基本概念(状态、动作、奖励、策略)的开发者,读完你能掌握世界模型的核心组成、想象 rollout 的实现方式、CartPole 上的完整训练流程,以及在自动驾驶、机器人、LLM 智能体三类场景中落地时的差异和坑点。

1. 从“试错太贵”说起:世界模型解决了什么问题

1.1 强化学习在真实环境中的三个硬约束

强化学习的第一性原理是试错:智能体在环境里执行动作,拿到奖励信号,再根据信号调整策略。这套逻辑在游戏里毫无压力,一旦挪到真实世界就会撞上三堵墙。

第一堵是成本墙。机械臂抓取失败一次要人工复位,自动驾驶的极端场景几乎无法主动复现,线上系统的错误操作可能直接写脏数据。第二堵是速度墙。真实环境的时间流速是固定的,训练一个策略往往需要上百万步交互,按真实时钟走根本等不起。第三堵是安全墙。探索阶段必然包含大量“坏动作”,而这些坏动作在真实环境中可能造成不可逆的后果。

仿真环境能缓解前两堵墙,但引入了新的问题:仿真与真实之间的差距,也就是常说的 sim-to-real gap。手工搭的仿真器越精细越贵,越粗糙越不准。世界模型提供的思路是:不追求用物理公式精确重建环境,而是从数据里学一个“能预测下一步会发生什么”的模型,让智能体在这个学出来的模型里练习。图省事的说法是“给算法造一个梦”,专业一点叫 model-based reinforcement learning(基于模型的强化学习)。

1.2 世界模型的最小定义

严格来说,世界模型是一个从历史信息出发、预测环境未来演化的参数化函数。最朴素的形态是一个状态转移函数:

s_{t+1} ~ p(s_{t+1} | s_t, a_t) r_t ~ p(r_t | s_t, a_t)

只要能把这两项学准,智能体就可以在任何起点上自己“脑补”出一段轨迹,用脑补出来的奖励去更新策略。这段脑补过程在文献里常被称作 imagination rollout,本文统一叫“想象 rollout”。

需要强调的是,世界模型不等于仿真器。仿真器是人为编写的规则系统,世界模型是从数据中拟合出来的统计模型,它天然带有误差,且误差会随预测步数累积。所以世界模型的关键从来不是“预测得准不准”,而是“预测得够不够支撑决策”。一个在像素层面重建得很漂亮的模型,未必能训出好策略;一个重建很粗糙的模型,只要对奖励和终止的判断足够可靠,也可能效果不错。这个认识是后面所有工程决策的基础。

1.3 有模型与无模型:不是替代而是分工

初学者容易把 model-based 和 model-free 对立起来,实际上在工程里它们经常混用。无模型方法(DQN、PPO、SAC)样本效率低但稳定,适合交互廉价、环境简单的场景;有模型方法样本效率高,但多了一层模型误差,调参更敏感。

比较务实的组合是这样的:先用少量真实交互数据训练世界模型,再在模型里做大规模策略优化,最后用真实环境做小批量验证和微调。这样真实环境的调用次数被压缩到原来的百分之几,同时保留了真实信号作为最终裁判。后文的实战案例也遵循这个结构:离线数据训练模型 → 隐空间训练策略 → 真实环境评估。

2. 世界模型的三种形态与适用场景

2.1 状态级模型:低成本、可解释、适合控制类任务

状态级模型直接以低维观测(位置、速度、角度、角速度这类结构化数值)作为输入输出,网络通常就是几层 MLP。它的优点是训练快、可解释、调试方便,你能直接对比“预测的下一帧状态”和“真实的下一帧状态”,误差在哪个维度上一目了然。

缺点是它要求你事先拿到干净的状态表示。真实机器人上这往往意味着要先做状态估计,而状态估计本身的误差会被世界模型继续放大。所以在教学、仿真验证、以及状态本身就很干净的控制任务里,状态级模型是首选;到了图像输入的场景,它就不够用了。本文第 5 章的实战就采用这一形态,因为它的代码最短、最能把原理讲透。

2.2 潜空间模型:图像输入的主流路线

当观测是图像时,直接在像素空间做预测代价极高,而且大量像素细节与决策无关。于是有了潜空间模型:先用编码器把高维观测压成一个低维隐向量,动力学模型在隐空间里递推,奖励和终止信号也从隐向量上预测。Dreamer 系列工作是这条路线的代表性方案,它把观测重建、奖励预测、终止预测、以及隐空间的先验后验对齐统一在一个损失里,然后在隐空间里做 actor-critic 学习。

这条路线的好处是样本效率高,在不少连续控制任务上能用远少于无模型方法的交互量达到相当的水平。代价是实现复杂度显著上升,尤其是隐变量既要承载重建信息又要承载动力学信息,两者的平衡(通常通过 KL 项调节)很容易调崩。如果你准备走这条路,建议先在一个最简单的环境上把损失曲线跑通,再谈复杂任务。

2.3 决策相关模型:只学与决策有关的部分

第三类思路更激进:既然模型的目的是辅助决策,那就不必重建观测,只学“对决策有用的那部分”。MuZero 是典型代表,它学习的是隐状态之间的转移、即时奖励和价值,而不要求隐状态能解码回原始观测。这种设计避免了“为了重建而浪费容量”,但也带来一个问题:隐状态缺乏直观含义,调试时很难判断模型学到了什么。

在工程选择上可以这么决策:如果你需要可视化验证模型质量、需要人工审查预测结果(比如医疗、工业质检这类场景),选可解码的潜空间模型;如果只关心最终策略表现、且算力紧张,可以考虑不解码的决策型模型。

2.4 三个落地场景的差异

自动驾驶。特点是高维感知输入、长时序、极端场景稀缺。世界模型的典型用法是做“条件生成式 rollout”:给定当前帧和历史动作,生成未来若干秒的传感器画面或鸟瞰图,用来扩充稀有场景、做闭环仿真、或者直接作为策略优化的训练场。这里的核心难点是长时序一致性和多模态未来(同一个起点,可能有多种合理演化),工程上通常需要引入隐变量来建模不确定性。

机器人。特点是真机昂贵、状态维度中等、物理约束强。用小数据量训练世界模型再在模型里练策略,是降低真机磨损的直接手段。但机器人任务的接触动力学很难学,模型在接触瞬间的预测误差会非常大,所以实践中往往采用短视界想象 + 频繁真机校正的滚动方式,而不是一次性在模型里跑几千步。

LLM 智能体。这里的“世界模型”含义更宽泛:它可以是工具调用的结果预测器(执行这个动作后,文件系统/数据库/网页会变成什么样),也可以是整个沙箱环境的可回放模拟。核心价值同样是降低真实调用的成本和风险——先在预测模型里筛掉明显错误的动作序列,再把少量候选送到真实工具上执行。这类场景的评估指标通常不是重建误差,而是“动作后果预测的准确率”和“下游任务成功率”。

3. 环境准备与版本说明

3.1 运行环境与依赖

本文的示例代码在 Python 3.10 及以上、PyTorch 2.x、Gymnasium 环境下编写。由于这些库的小版本迭代较快,本文不写死具体小版本号,你按本地实际情况安装即可,重点是配置思路而不是某个精确版本。

# 文件路径:requirements.txt gymnasium>=0.29.0 numpy>=1.24.0 torch>=2.0.0

安装命令:

pip install -r requirements.txt

如果你使用 GPU 版 PyTorch,请按照官方安装指引选择对应 CUDA 版本的安装命令,不要直接用上面这行覆盖,否则可能装成 CPU 版。本文的示例规模很小,CPU 训练完全够用,不必强求 GPU。

示例使用CartPole-v1环境。它是一个经典的控制任务:小车上竖着一根杆子,智能体只能施加向左或向右的力,目标是让杆子尽量不倒。观测是 4 维向量(小车位置、小车速度、杆子角度、杆子角速度),动作是 2 个离散值,每坚持一步得 1 分。选择它的原因是状态维度低、训练秒级完成,非常适合把世界模型的每个环节看清楚。如果你本地的 Gymnasium 版本里环境 ID 有变化,用gymnasium.envs.registry查一下可用列表即可。

3.2 项目目录结构

为了让每一步都能单独运行和调试,我们把代码拆成五个文件:

world-model-demo/ ├── requirements.txt ├── collect_data.py # 收集离线转移数据 ├── world_model.py # 世界模型网络定义 ├── train_world_model.py # 训练世界模型 ├── imagine_policy.py # 在想象中训练策略 ├── evaluate.py # 真实环境评估 ├── data/ # 数据目录(运行时生成) └── checkpoints/ # 模型权重目录(运行时生成)

拆分的目的是让“数据—模型—策略—评估”四段流程彼此解耦。你之后换环境时,只需要改collect_data.py和网络输入输出维度,其余逻辑基本可以复用。

3.3 版本与兼容性提醒

有三个地方容易踩版本坑,先提前说明。

第一,Gymnasium 的step返回值是五元组(obs, reward, terminated, truncated, info),terminated表示任务真正结束(比如杆子倒了),truncated表示达到步数上限被强制截断。这两个信号在训练世界模型时必须区分对待,后文会详细解释原因。

第二,reset现在返回(obs, info)而不是单独的obs,老代码直接用会报解包错误。

第三,env.reset(seed=...)的随机种子机制建议显式传入,否则每次运行结果差异较大,排查问题时无法复现。

4. 核心原理拆解:一次“想象”是怎么跑起来的

4.1 编码器:把观测压成隐状态

编码器的作用是把观测映射成一个稠密的隐向量。在状态级模型里它看起来有点多余(4 维直接当输入也行),但在潜空间模型里它是刚需,因为原始观测往往高维、冗余、且包含大量与决策无关的信息。

保留编码器的另一个好处是统一接口:无论后面换成图像输入还是状态输入,策略网络的输入始终是固定维度的隐向量,策略部分代码不用改。这在工程上很有价值,意味着你可以先用低维状态把整条链路调通,再替换编码器去接图像,而不需要重写策略训练部分。

需要注意的是,编码器输出的隐向量尺度要稳定。如果隐向量数值范围飘忽不定,动力学模型会很难学,策略网络也会不稳定。实践中常用 LayerNorm 或者对隐向量做归一化,本文为了代码简洁没有加,但你在复杂任务上应该加上。

4.2 动力学模型:动作怎么注入

动力学模型接收“当前隐状态 + 动作”,输出“下一隐状态”。动作是离散的时候通常做 one-hot 编码后与隐状态拼接,连续动作则直接拼接。这个拼接看起来简单,但有个细节值得注意:动作的数值尺度要和隐状态匹配。如果动作是 [-1, 1] 的连续值而隐状态在几十的量级,直接拼接会让网络对动作不敏感,学出来的模型几乎忽略动作输入,想象 rollout 就会退化成“不管做什么,未来都一样”。

判断这个问题的方法很直接:固定同一个起点隐状态,输入不同的动作,看输出的下一隐状态差异大不大。如果几乎一样,说明动作信息没有被有效利用,需要检查拼接方式、加动作维度的缩放、或者检查训练数据里动作是否足够多样。

4.3 奖励头与终止头:策略的“裁判”

奖励头和终止头都接在预测出的下一隐状态上,分别输出标量奖励和终止概率。这两个头虽然结构简单,却决定了想象 rollout 的质量上限——因为策略优化的梯度最终来自奖励信号。

这里有个训练目标对齐的问题容易被忽略:如果训练时用z_next预测奖励,那么想象 rollout 时也必须用z_next的奖励,不能中途改用z_t。本文代码严格保持了这个一致性。看起来是小细节,但一旦不一致,模型在训练集上损失很低,想象出来的回报却完全不对,排查起来非常痛苦。

终止头用二分类交叉熵训练,输出的是 logit,使用时过 sigmoid 得到概率。是否需要按概率截断 rollout,取决于任务:如果终止是稀疏事件(比如杆子倒了),提前截断能减少无意义的想象步;如果终止很频繁,过早截断会导致想象轨迹太短、梯度信号不足。

4.4 想象 rollout 与误差累积

想象 rollout 的本质是拿模型当环境用:从一个起点隐状态开始,策略采样动作,模型递推下一状态并给出奖励,循环若干步。这段轨迹完全由模型生成,不涉及任何真实交互。

问题在于误差累积。单步预测误差是 e,滚动 T 步之后误差量级大致会按 e 的倍数放大,而且模型在训练数据分布之外的表现会急速恶化。这就是为什么想象视界不能设太长。经验上,短视界(十几到几十步)配合频繁的策略更新和真实环境校正,比一次性想象几百步要可靠得多。

还有一种缓解手段叫“分岔想象”:不是一条轨迹滚到底,而是从同一批起点分岔出多条短轨迹,用多条轨迹的平均来降低单条轨迹的方差。实现上就是把 batch 维度放大,成本可控。

4.5 潜空间坍缩与 KL 平衡

如果你采用的是带随机隐变量的模型(比如 VAE 式的编码器),会遇上一个经典问题:KL 项权重太小时,编码器几乎忽略输入,把所有观测都编码成同一个点(潜空间坍缩);KL 项权重太大时,编码器又过度关注重建,隐向量退化成压缩后的观测,动力学变得难学。

处理方式通常是两个:一是用 KL 退火,训练初期把 KL 权重设得很小,逐步升高;二是用 free bits 技巧,给每个隐维度设一个最小的 KL 目标,低于这个目标就不再施压。本文的最小实现用的是确定性编码器,避开了这个问题,但你在引入随机隐变量时必须提前规划。

5. 完整实战:在 CartPole 上搭一个最小世界模型

5.1 收集离线转移数据

第一步是攒数据。数据的多样性决定了世界模型能看到多少种状态-动作组合,用随机策略收集虽然“笨”,但覆盖面广,是性价比很高的起点。

# 文件路径:collect_data.py import os import numpy as np import gymnasium as gym def collect_transitions(episodes=300, seed=0, save_path="data/cartpole_transitions.npz"): env = gym.make("CartPole-v1") rng = np.random.default_rng(seed) obs_buf, act_buf, rew_buf, next_obs_buf, done_buf = [], [], [], [], [] for ep in range(episodes): obs, _ = env.reset(seed=seed + ep) while True: action = int(rng.integers(env.action_space.n)) next_obs, reward, terminated, truncated, _ = env.step(action) obs_buf.append(obs) act_buf.append(action) rew_buf.append(reward) next_obs_buf.append(next_obs) # 关键:只有真实终止才算 done,因步数上限被截断的不算 done_buf.append(float(terminated)) if terminated or truncated: break obs = next_obs os.makedirs(os.path.dirname(save_path), exist_ok=True) np.savez_compressed( save_path, obs=np.asarray(obs_buf, dtype=np.float32), act=np.asarray(act_buf, dtype=np.int64), rew=np.asarray(rew_buf, dtype=np.float32), next_obs=np.asarray(next_obs_buf, dtype=np.float32), done=np.asarray(done_buf, dtype=np.float32), ) print(f"共收集 {len(obs_buf)} 条转移,已保存到 {save_path}") if __name__ == "__main__": collect_transitions()

运行:

python collect_data.py

预期输出类似共收集 6000 条转移,已保存到 data/cartpole_transitions.npz,具体条数与你本地随机种子有关。

这段代码里最值得说的是done_buf那一行。如果直接把terminated or truncated存成 done,世界模型会把“跑到 500 步被强制截断”误认为“环境终止了”,于是它在所有到了 500 步的位置都会输出高终止概率。这样一来,想象 rollout 在接近 500 步时会被误截断,策略学到的时序结构就变形了。这个坑非常隐蔽,因为训练损失看起来只是略微偏高,不会报错。

5.2 实现世界模型网络

接下来定义网络。结构上包含五个部分:编码器、动力学、奖励头、终止头、解码器。解码器不是策略训练必需的,但它的重建损失提供了很强的监督信号,能帮助隐空间保留足够的信息。

# 文件路径:world_model.py import torch import torch.nn as nn import torch.nn.functional as F class WorldModel(nn.Module): """最小可用世界模型:编码器 + 动力学 + 奖励头 + 终止头 + 解码器""" def __init__(self, obs_dim=4, n_actions=2, hidden_dim=128, latent_dim=32): super().__init__() self.obs_dim = obs_dim self.n_actions = n_actions self.latent_dim = latent_dim self.encoder = nn.Sequential( nn.Linear(obs_dim, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, latent_dim), ) self.transition = nn.Sequential( nn.Linear(latent_dim + n_actions, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, latent_dim), ) self.reward_head = nn.Sequential( nn.Linear(latent_dim, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, 1), ) self.done_head = nn.Sequential( nn.Linear(latent_dim, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, 1), ) self.decoder = nn.Sequential( nn.Linear(latent_dim, hidden_dim), nn.ELU(), nn.Linear(hidden_dim, obs_dim), ) def encode(self, obs): return self.encoder(obs) def next_latent(self, latent, action): action_onehot = F.one_hot(action, num_classes=self.n_actions).float() x = torch.cat([latent, action_onehot], dim=-1) return self.transition(x) def predict_reward(self, latent): return self.reward_head(latent).squeeze(-1) def predict_done(self, latent): return self.done_head(latent).squeeze(-1) def decode(self, latent): return self.decoder(latent)

几点说明。激活函数用 ELU 而不是 ReLU,是因为 ELU 在负值区有梯度,对动力学这种需要精细拟合的网络更友好。隐向量维度取 32,比原始 4 维观测大,这是有意为之:隐空间需要额外容纳“动力学所需的历史信息”,因为我们的模型是单步马尔可夫的,无法看到历史,只能靠隐空间隐式编码趋势(比如速度)。如果维度压得太小,模型会学不准。

next_latent里的 one_hot 要求action是整型张量且形状为(batch,),如果你传入的是(batch, 1),one_hot 会得到三维张量,拼接时会报维度错误。这是新手最常遇到的错误之一。

5.3 训练世界模型

训练目标是多任务联合:重建损失、奖励损失、终止损失加权求和。三个损失的量级不同,需要调权重,本文给的经验写法是重建和奖励各占 1,终止占 0.5。

# 文件路径:train_world_model.py import os import numpy as np import torch import torch.nn.functional as F from world_model import WorldModel def load_data(path="data/cartpole_transitions.npz"): data = np.load(path) return {k: torch.from_numpy(data[k]) for k in data.files} def train(epochs=60, batch_size=256, lr=3e-4, seed=0, save_path="checkpoints/world_model.pt"): torch.manual_seed(seed) np.random.seed(seed) data = load_data() n = data["obs"].shape[0] model = WorldModel() optimizer = torch.optim.Adam(model.parameters(), lr=lr) for epoch in range(epochs): perm = torch.randperm(n) total_loss = total_rec = total_rew = total_done = 0.0 steps = 0 for i in range(0, n, batch_size): idx = perm[i:i + batch_size] obs = data["obs"][idx] act = data["act"][idx] rew = data["rew"][idx] next_obs = data["next_obs"][idx] done = data["done"][idx] z = model.encode(obs) z_next = model.next_latent(z, act) rec_loss = F.mse_loss(model.decode(z_next), next_obs) rew_loss = F.mse_loss(model.predict_reward(z_next), rew) done_logit = model.predict_done(z_next) done_loss = F.binary_cross_entropy_with_logits(done_logit, done) loss = rec_loss + rew_loss + 0.5 * done_loss optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0) optimizer.step() total_loss += loss.item() total_rec += rec_loss.item() total_rew += rew_loss.item() total_done += done_loss.item() steps += 1 if (epoch + 1) % 10 == 0: print(f"epoch {epoch + 1:3d} | loss {total_loss / steps:.4f} " f"| rec {total_rec / steps:.4f} " f"| rew {total_rew / steps:.4f} " f"| done {total_done / steps:.4f}") os.makedirs(os.path.dirname(save_path), exist_ok=True) torch.save(model.state_dict(), save_path) print(f"世界模型已保存到 {save_path}") return model if __name__ == "__main__": train()

运行:

python train_world_model.py

你大概会看到重建损失在 1e-4 到 1e-3 量级,奖励损失非常小(因为 CartPole 每步奖励恒为 1,模型很快就能学到),终止损失在 0.1 上下波动。梯度裁剪那行是为了防止早期动力学发散,在更复杂的任务上这一步几乎是必须的。

有一点需要注意:load_data里把所有字段都转成了张量,如果数据量很大(比如图像数据),一次性载入内存会爆掉,那时候要改用分块读取或者内存映射。本文的数据量很小,不必纠结。

5.4 在想象中训练策略

策略网络的输入是隐状态,输出是动作 logits。训练算法用最朴素的 REINFORCE,配合批次均值作为 baseline。选它是因为公式直观、代码短,便于看清“梯度来自想象奖励”这件事。

# 文件路径:imagine_policy.py import numpy as np import torch import torch.nn as nn from world_model import WorldModel class Policy(nn.Module): def __init__(self, latent_dim=32, n_actions=2, hidden_dim=64): super().__init__() self.net = nn.Sequential( nn.Linear(latent_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, n_actions), ) def forward(self, latent): return self.net(latent) def imagine_rollout(model, policy, start_latent, horizon): """在隐空间中展开 horizon 步,收集奖励与对数概率""" latents, actions, rewards, logps = [], [], [], [] latent = start_latent for _ in range(horizon): logits = policy(latent) dist = torch.distributions.Categorical(logits=logits) action = dist.sample() logp = dist.log_prob(action) latent = model.next_latent(latent, action) reward = model.predict_reward(latent) latents.append(latent) actions.append(action) rewards.append(reward) logps.append(logp) return latents, actions, rewards, logps def train_policy(horizon=20, iterations=400, batch_size=64, lr=3e-3, gamma=0.99, seed=0, save_path="checkpoints/policy.pt"): torch.manual_seed(seed) np.random.seed(seed) model = WorldModel() model.load_state_dict(torch.load("checkpoints/world_model.pt", map_location="cpu")) model.eval() for p in model.parameters(): p.requires_grad_(False) # 冻结世界模型,只训练策略 data = np.load("data/cartpole_transitions.npz") start_obs = torch.from_numpy(data["obs"]) n = start_obs.shape[0] policy = Policy(latent_dim=model.latent_dim) optimizer = torch.optim.Adam(policy.parameters(), lr=lr) for it in range(iterations): idx = torch.randint(0, n, (batch_size,)) with torch.no_grad(): start_latent = model.encode(start_obs[idx]) _, _, rewards, logps = imagine_rollout(model, policy, start_latent, horizon) R = torch.stack(rewards, dim=1) # (batch, horizon) logp = torch.stack(logps, dim=1) returns = torch.zeros_like(R) running = torch.zeros(R.shape[0]) for t in reversed(range(horizon)): running = R[:, t] + gamma * running returns[:, t] = running advantage = returns - returns.mean(dim=0, keepdim=True) loss = -(logp * advantage.detach()).mean() optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(policy.parameters(), 10.0) optimizer.step() if (it + 1) % 50 == 0: avg_return = R.sum(dim=1).mean().item() print(f"iter {it + 1:4d} | loss {loss.item():.4f} " f"| 想象平均回报 {avg_return:.3f}") torch.save(policy.state_dict(), save_path) print(f"策略已保存到 {save_path}") return policy if __name__ == "__main__": train_policy()

运行:

python imagine_policy.py

这里有三个设计决策值得解释。

第一,想象起点从训练数据的真实观测里采样,而不是从固定状态开始。这是因为策略最终要在真实分布上工作,起点分布必须匹配。如果你固定从一个状态起跑,策略会过拟合到那条特定的想象轨迹上。

第二,世界模型被冻结。如果不冻结,策略优化的梯度会回传到世界模型,把它往“给出高奖励”的方向带偏,也就是模型被策略“骗”了。这是 model-based RL 里非常经典的失败模式。

第三,advantage 按时间步做了去均值。这只是最简单的方差缩减,比不做好很多,但远不是最优。如果你想进一步提升稳定性,可以引入一个价值网络作为学习到的 baseline,或者改用 actor-critic 结构。

5.5 真实环境评估

最后一步,把学到的策略放到真实环境里跑,看它到底行不行。评估时用编码器把真实观测编成隐状态再喂给策略,这样每一步都能用真实信息校正,避免隐状态漂移。

# 文件路径:evaluate.py import gymnasium as gym import numpy as np import torch from world_model import WorldModel from imagine_policy import Policy def evaluate_policy(model, policy, episodes=20, seed=123): env = gym.make("CartPole-v1") model.eval() policy.eval() returns = [] for ep in range(episodes): obs, _ = env.reset(seed=seed + ep) total = 0.0 done = False while not done: obs_t = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0) with torch.no_grad(): latent = model.encode(obs_t) action = int(torch.argmax(policy(latent), dim=-1).item()) obs, reward, terminated, truncated, _ = env.step(action) total += reward done = terminated or truncated returns.append(total) return float(np.mean(returns)), float(np.std(returns)) def evaluate_random(episodes=20, seed=123): env = gym.make("CartPole-v1") rng = np.random.default_rng(seed) returns = [] for ep in range(episodes): obs, _ = env.reset(seed=seed + ep) total = 0.0 done = False while not done: action = int(rng.integers(env.action_space.n)) obs, reward, terminated, truncated, _ = env.step(action) total += reward done = terminated or truncated returns.append(total) return float(np.mean(returns)), float(np.std(returns)) if __name__ == "__main__": model = WorldModel() model.load_state_dict(torch.load("checkpoints/world_model.pt", map_location="cpu")) policy = Policy(latent_dim=model.latent_dim) policy.load_state_dict(torch.load("checkpoints/policy.pt", map_location="cpu")) rand_mean, rand_std = evaluate_random() model_mean, model_std = evaluate_policy(model, policy) print(f"随机策略 : 平均回报 {rand_mean:.1f} ± {rand_std:.1f}") print(f"世界模型策略 : 平均回报 {model_mean:.1f} ± {model_std:.1f}")

运行:

python evaluate.py

5.6 结果解读与调参方向

判断整套流程是否真的跑通了,看两件事:随机策略的基线明显低于世界模型策略,且世界模型策略的表现明显高于随机。如果两者接近,说明策略没学到东西,问题大概率出在想象 rollout 或者损失权重上。具体数值随随机种子波动,不必追求复现某个精确数字,关键是看趋势对不对。

如果效果不理想,按下面顺序调:

先看世界模型的单步预测质量。可以单独写个小脚本,从数据里取一批真实观测和动作,对比decode(next_latent(z, a))和真实next_obs,看四个维度的误差分别多大。如果角度维度误差特别大,说明模型对关键的动力学量学得不好。

再检查想象 rollout 的视界。视界从 20 逐步降到 8 或升到 40,观察真实评估分数的变化。这能帮你找到“模型可信”和“梯度信号足够”之间的平衡点。

最后调策略学习率和迭代次数。REINFORCE 类的方法方差大,学习率过低会几乎不动,过高会来回震荡。把学习率在 1e-3 到 1e-2 之间扫一遍通常就够了。

6. 常见问题与排查清单

6.1 模型损失很低,但策略完全学不会

现象是世界模型训练日志看着很漂亮,重建误差很小,但在想象里训出来的策略在真实环境里表现和随机策略差不多。

最可能的原因是编码器学会了“恒等映射”:直接把观测抄进隐向量,动力学模型则退化成“近似复制上一帧”。在这种情况下,重建任务确实很容易完成(因为隐空间几乎保留了全部信息),但模型对动作的响应很微弱,想象 rollout 出来的轨迹与动作几乎无关,策略自然学不到东西。

排查方法是固定同一个起点隐状态,输入不同动作,比较输出的下一隐状态差异。如果差异极小,就是这个原因。解决思路有两条:一是提高动作注入的表达能力,比如把动作先过一层独立的小网络再拼接;二是给动力学加更强的约束,比如加大奖励预测的权重,逼模型必须依赖动作才能预测准确。

6.2 想象回报虚高,真实回报崩盘

想象 rollout 的累计回报一路涨到接近视界上限,但真实环境一跑就倒。这是典型的模型利用(model exploitation):策略找到了模型的预测盲区,在那个区域模型总是给出高奖励,但真实环境里这些动作是灾难性的。

根本原因是策略会主动往模型误差大的地方钻。缓解手段主要有三类,实践中通常组合使用。第一类是限制想象视界,缩短策略“钻空子”的空间。第二类是模型集成,训练多个世界模型(不同初始化、不同数据子集),用它们预测的方差来惩罚策略在不确定区域的动作。第三类是训练过程中定期用真实数据重新训练或微调世界模型,把策略探索出来的新状态补进训练集。

6.3 数据分布覆盖不足导致模型失效

如果收集数据时只用了随机策略,模型在“杆子快倒”这类极端状态上的样本会很少,预测精度很差。策略学会了在这些状态下该怎么做,模型却预测不准,两边对不上。

解决思路有两种。一种是从数据层面解决,用混合策略收集:随机策略收集基础覆盖,加上一些简单启发式策略(比如总是往杆子倾斜方向推)收集关键区域的数据。另一种是从训练层面解决,按状态难度做重采样,让模型在稀有状态上多训练几轮。

6.4 排查清单

问题现象常见原因排查思路
重建损失下降但动作无响应编码器恒等映射,动作被忽略固定起点换动作,比较输出隐状态差异
想象回报虚高策略利用模型盲区缩短视界、加模型集成、定期重训模型
终止头误判truncated 被当作终止检查 done 标签是否只取 terminated
隐状态数值爆炸缺少归一化或学习率过高加 LayerNorm、降学习率、梯度裁剪
策略训练曲线剧烈震荡REINFORCE 方差过大加大 batch、加 baseline、衰减学习率
真实评估远低于想象隐状态漂移评估时每步用编码器重新编码观测
模型对极端状态预测差数据覆盖不足混合策略收集、重采样
训练速度异常慢每步都在建计算图冻结模型参数、加 no_grad

7. 工程最佳实践

7.1 数据质量优先于模型结构

在最小规模的世界模型上,把数据量翻倍往往比把网络加深一层更有效。收集数据时要有意识地覆盖状态空间:随机策略负责广度,任务相关策略负责深度。另外建议保留原始转移数据而不只是训练好的模型权重,因为后续换任务、加正则、改损失时,原始数据是唯一的真相来源。

数据里还要记录元信息,比如收集时的策略版本、环境参数、随机种子。等到模型效果不对时,你会非常需要这些信息去定位是数据变了还是代码变了。

7.2 用不确定性约束策略

前文提到的模型集成不是理论上的讲究,而是工程上的必需品。具体做法是训练 3 到 5 个世界模型,每个模型用不同的随机初始化,或者对训练数据做 bootstrap 采样。想象 rollout 时,让所有模型都预测一遍,取均值作为奖励估计,取标准差作为不确定性惩罚项:

# 核心片段:集成模型的不确定性惩罚 def ensemble_reward(models, latent, action, penalty=1.0): rewards = [] next_latents = [] for m in models: z_next = m.next_latent(latent, action) next_latents.append(z_next) rewards.append(m.predict_reward(z_next)) rewards = torch.stack(rewards, dim=0) # (n_models, batch) mean_reward = rewards.mean(dim=0) std_reward = rewards.std(dim=0) # 不确定性越大,惩罚越多 adjusted_reward = mean_reward - penalty * std_reward return adjusted_reward, next_latents[0]

这段代码需要配合多个已训练的世界模型使用,penalty一般从 0.1 到 1.0 之间试。惩罚太大会让策略过于保守,只在模型确信的区域活动,探索能力下降。

7.3 控制想象视界与重规划频率

视界不是越长越好。太短,梯度信号不足以覆盖任务的时序结构;太长,误差累积会压过真实信号。一个实用的做法是让视界与任务的关键决策间隔挂钩:如果关键决策每隔十几步才有一次,视界就设在那个量级附近。

另外建议采用“滚动重规划”而不是一次性长跑:每想象一小段,就用真实观测重置起点,再继续想象。这样虽然真实交互次数增加了一点,但换来的是隐状态不漂移,整体可靠性高很多。机器人这类任务基本都用这种方式。

7.4 生产环境的安全边界

如果这套东西要碰真实系统,有几条底线必须守住。所有真实环境的操作必须经过明确授权,在测试环境验证通过后再进入生产。涉及写操作、删除操作、控制指令的,必须加人工确认或熔断机制,不能让策略直接接管。世界模型可以作为“预演”层,先预测动作后果,过滤掉预测风险高的动作,但最终决策仍要有兜底规则。

数据层面也要注意,用真实生产数据训练模型前,先确认数据脱敏和权限边界,不要把敏感信息带进训练集或日志。模型文件本身也要当作受控资产管理,记录训练数据来源和版本。

8. 学习路线与延伸方向

8.1 从最小实现走向完整方案

本文的实现刻意做减法:确定性隐状态、单步动力学、单一模型、简单策略梯度。把它跑通之后,下一步可以按这几个方向逐个加码。

把确定性编码器换成随机隐变量编码器,引入 KL 项,你就基本走到了 VAE 式世界模型的门口。把隐空间上的 REINFORCE 换成 actor-critic,加入价值网络和 lambda-return,稳定性会明显提升。把单一模型换成集成,加入不确定性惩罚,就有了对抗模型利用的基本武器。这三步走完,结构上已经很接近 Dreamer 系列的核心思想。

8.2 三个场景的入门切入点

如果目标是自动驾驶,建议从开源的驾驶仿真环境入手,先做“给定前若干帧,预测未来若干帧”的纯预测任务,把时序一致性和多模态建模摸清楚,再谈闭环策略训练。观测是高维图像,编码器通常要用卷积或者视觉 Transformer,训练成本会显著高于本文示例。

如果目标是机器人,建议从状态级模型开始,选一个接触较少的任务(比如推到目标位置),先验证“模型里训策略、真机上验证”的闭环。接触密集的任务要额外处理接触瞬间的模型误差,常用做法是在接触阶段缩短想象步长、提高真机校正频率。

如果目标是 LLM 智能体,世界模型的形态会完全不同:它更像一个“动作后果预测器 + 可回放沙箱”。可以先从最简单的做起,记录真实执行的动作和对应的环境状态变化,训练一个小模型预测“执行这个动作后关键状态会怎么变”,在智能体决策前用它筛掉明显不合理的候选动作。这条路的技术门槛比视觉世界模型低,但在降低真实调用成本上的收益很直接。

8.3 动手建议

建议你按这个顺序动手:先原样跑通本文的五个脚本,确认随机基线和世界模型策略的差距;然后把horizon从 20 改到 5 和 50,观察曲线怎么变,体会误差累积的威力;接着故意把done_buf改成包含truncated,看看训练会崩成什么样,这个反例比正例更能加深理解;最后尝试把编码器去掉,直接用原始观测当隐状态,对比前后效果差异,你会对“隐空间到底提供了什么”有更具体的感受。

真正把世界模型用起来的难点从来不是网络结构,而是数据覆盖、误差控制和真实验证这三个环节的配合。把最小闭环跑顺,比一上来就堆复杂结构要划算得多。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询