做强化学习的朋友大多有过类似的纠结算法本身不难写难的是让它在哪儿跑。真机上撞一次墙可能损失几千块仿真里参数调偏了结果又迁移不过去LLM 智能体更麻烦每一步都要真调接口、真读写文件试错成本按 token 计费。这篇文章就围绕“世界模型”这条路展开讲清楚它为什么能让强化学习算法在“不碰真实环境”的前提下完成训练并给出一套可以直接跑起来的最小实现。读者对象是有 Python 和 PyTorch 基础、了解强化学习基本概念状态、动作、奖励、策略的开发者读完你能掌握世界模型的核心组成、想象 rollout 的实现方式、CartPole 上的完整训练流程以及在自动驾驶、机器人、LLM 智能体三类场景中落地时的差异和坑点。1. 从“试错太贵”说起世界模型解决了什么问题1.1 强化学习在真实环境中的三个硬约束强化学习的第一性原理是试错智能体在环境里执行动作拿到奖励信号再根据信号调整策略。这套逻辑在游戏里毫无压力一旦挪到真实世界就会撞上三堵墙。第一堵是成本墙。机械臂抓取失败一次要人工复位自动驾驶的极端场景几乎无法主动复现线上系统的错误操作可能直接写脏数据。第二堵是速度墙。真实环境的时间流速是固定的训练一个策略往往需要上百万步交互按真实时钟走根本等不起。第三堵是安全墙。探索阶段必然包含大量“坏动作”而这些坏动作在真实环境中可能造成不可逆的后果。仿真环境能缓解前两堵墙但引入了新的问题仿真与真实之间的差距也就是常说的 sim-to-real gap。手工搭的仿真器越精细越贵越粗糙越不准。世界模型提供的思路是不追求用物理公式精确重建环境而是从数据里学一个“能预测下一步会发生什么”的模型让智能体在这个学出来的模型里练习。图省事的说法是“给算法造一个梦”专业一点叫 model-based reinforcement learning基于模型的强化学习。1.2 世界模型的最小定义严格来说世界模型是一个从历史信息出发、预测环境未来演化的参数化函数。最朴素的形态是一个状态转移函数s_{t1} ~ p(s_{t1} | 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 gymnasium0.29.0 numpy1.24.0 torch2.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(episodes300, seed0, save_pathdata/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(seedseed 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_okTrue) np.savez_compressed( save_path, obsnp.asarray(obs_buf, dtypenp.float32), actnp.asarray(act_buf, dtypenp.int64), rewnp.asarray(rew_buf, dtypenp.float32), next_obsnp.asarray(next_obs_buf, dtypenp.float32), donenp.asarray(done_buf, dtypenp.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_dim4, n_actions2, hidden_dim128, latent_dim32): 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_classesself.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(pathdata/cartpole_transitions.npz): data np.load(path) return {k: torch.from_numpy(data[k]) for k in data.files} def train(epochs60, batch_size256, lr3e-4, seed0, save_pathcheckpoints/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(), lrlr) 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(fepoch {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_okTrue) 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_dim32, n_actions2, hidden_dim64): 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(logitslogits) 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(horizon20, iterations400, batch_size64, lr3e-3, gamma0.99, seed0, save_pathcheckpoints/policy.pt): torch.manual_seed(seed) np.random.seed(seed) model WorldModel() model.load_state_dict(torch.load(checkpoints/world_model.pt, map_locationcpu)) 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_dimmodel.latent_dim) optimizer torch.optim.Adam(policy.parameters(), lrlr) 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, dim1) # (batch, horizon) logp torch.stack(logps, dim1) 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(dim0, keepdimTrue) 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(dim1).mean().item() print(fiter {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, episodes20, seed123): env gym.make(CartPole-v1) model.eval() policy.eval() returns [] for ep in range(episodes): obs, _ env.reset(seedseed ep) total 0.0 done False while not done: obs_t torch.as_tensor(obs, dtypetorch.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(episodes20, seed123): env gym.make(CartPole-v1) rng np.random.default_rng(seed) returns [] for ep in range(episodes): obs, _ env.reset(seedseed 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_locationcpu)) policy Policy(latent_dimmodel.latent_dim) policy.load_state_dict(torch.load(checkpoints/policy.pt, map_locationcpu)) 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.py5.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_grad7. 工程最佳实践7.1 数据质量优先于模型结构在最小规模的世界模型上把数据量翻倍往往比把网络加深一层更有效。收集数据时要有意识地覆盖状态空间随机策略负责广度任务相关策略负责深度。另外建议保留原始转移数据而不只是训练好的模型权重因为后续换任务、加正则、改损失时原始数据是唯一的真相来源。数据里还要记录元信息比如收集时的策略版本、环境参数、随机种子。等到模型效果不对时你会非常需要这些信息去定位是数据变了还是代码变了。7.2 用不确定性约束策略前文提到的模型集成不是理论上的讲究而是工程上的必需品。具体做法是训练 3 到 5 个世界模型每个模型用不同的随机初始化或者对训练数据做 bootstrap 采样。想象 rollout 时让所有模型都预测一遍取均值作为奖励估计取标准差作为不确定性惩罚项# 核心片段集成模型的不确定性惩罚 def ensemble_reward(models, latent, action, penalty1.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, dim0) # (n_models, batch) mean_reward rewards.mean(dim0) std_reward rewards.std(dim0) # 不确定性越大惩罚越多 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看看训练会崩成什么样这个反例比正例更能加深理解最后尝试把编码器去掉直接用原始观测当隐状态对比前后效果差异你会对“隐空间到底提供了什么”有更具体的感受。真正把世界模型用起来的难点从来不是网络结构而是数据覆盖、误差控制和真实验证这三个环节的配合。把最小闭环跑顺比一上来就堆复杂结构要划算得多。