DreamerV3 训练技巧：从踩坑到收敛的实战经验

Tue, 11 Aug 2026 00:00:00 +0000

上一篇文章梳理了世界模型的四条表征路线。今天我们回到 DreamerV3 的实战主题，聊聊训练过程中的那些坑和技巧。本文基于 DreamerV3 commit e3f02248、JAX + Haiku、MuJoCo + DM Control 的本地实验环境，不同版本的参数名和配置可能有差异。

DreamerV3 是目前最开源、最成熟的世界模型实现之一。但如果你真的动手训练过，一定知道这个过程并不轻松——环境配置、超参数调优、训练不稳定、收敛慢……各种坑。

这篇文章总结我在训练 DreamerV3 过程中遇到的常见问题和解决方案，希望能帮你少走弯路。

环境配置的坑

1. MuJoCo 版本问题

DreamerV3 的官方实现基于 JAX + MuJoCo。但 MuJoCo 的版本选择很关键：

MuJoCo 2.x vs 3.x。官方代码基于 MuJoCo 2.x，但 3.x 版本有性能改进。如果用 3.x，需要修改部分 API 调用。
mujoco-py vs mujoco。mujoco-py 是旧的 Python 绑定，已经停止维护。建议直接用官方的 mujoco 包。

建议：优先使用与你当前 DreamerV3 commit 匹配的 MuJoCo 版本，不要一开始就升级整个环境。

2. JAX 的 GPU 配置

JAX 的 GPU 支持需要正确安装 CUDA 和 cuDNN。常见问题：

CUDA 版本不匹配。JAX 对 CUDA 版本有严格要求，安装前查看官方文档。
GPU 内存不足。JAX 默认会占用所有 GPU 内存。可以设置 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 限制内存使用。如果遇到编译阶段内存峰值过高，也可以设置 XLA_PYTHON_CLIENT_PREALLOCATE=false 关闭预分配。
多 GPU 问题。JAX 的多 GPU 支持和 PyTorch 不同，需要特别配置。

下图是 DreamerV3 在 RTX 5090D 上训练时的实际 GPU 占用情况，显存占用约 25.6GB（32GB 的 78%），GPU 利用率 91%： ![DreamerV3 训练时 GPU 占用情况](/images/nvidia smi.png) 图：RTX 5090D 上 DreamerV3 训练时的 GPU 显存占用（25.6GB / 32GB，利用率 91%）

训练技巧 on WorldSense 技术笔记

DreamerV3 训练技巧：从踩坑到收敛的实战经验

环境配置的坑

1. MuJoCo 版本问题

2. JAX 的 GPU 配置