Part VI · 第 24 章 Model-based 方法 · 24.3Part VI · Ch. 24 Model-Based Methods · 24.3

Dreamer 与 TD-MPCDreamer & TD-MPC

机器人在真实世界里的每一次尝试都很贵。世界模型的办法是:先从经验里学一个能"做梦"的模型,再把大部分学习搬进梦里。本节讲当今最有影响的两条路线。Dreamer 用循环状态空间模型(RSSM)把图像压成一个既有记忆又有随机性的隐状态,然后完全在想象出来的轨迹上训练 actor-critic;TD-MPC 不重建图像,只学一个能预测奖励和价值的隐空间,每一步都用 MPPI 在里面规划一小段,再用学到的价值函数接上远处的未来。五个实验都在浏览器里实时训练:一个会分岔的模型、在想象中学会控制关节的 Dreamer、被 U 形陷阱困住又靠价值函数逃出来的规划器、被干扰物骗走的重建式隐变量,以及让同一组超参数适用于任何奖励尺度的 symlog 与 two-hot。 Every trial a robot makes in the real world is expensive. The world-model answer is to learn, from experience, a model that can "dream", and then move most of the learning into the dream. This lesson covers today's two most influential lines of work. Dreamer compresses images into a latent state with both memory and randomness, the recurrent state-space model (RSSM), and trains an actor-critic purely on imagined trajectories. TD-MPC reconstructs nothing: it learns a latent space that predicts rewards and values, plans a short stretch in it with MPPI at every step, and lets a learned value function account for the distant future. Five experiments train live in your browser: a model that has to predict a fork, a Dreamer that learns to control a joint in imagination, a planner trapped in a U-shaped dead end until a value function gets it out, a reconstruction latent that a distractor leads astray, and symlog and two-hot targets that let one set of hyperparameters work at any reward scale.

约 55 分钟~55 min 前置:24.1 学习动力学模型 · 17.2 RNN 与状态空间模型 · 19.1 VAE · 20.4 Actor-Critic · 13.4 MPCPrereqs: 24.1 Learning dynamics models · 17.2 RNNs & SSMs · 19.1 VAEs · 20.4 Actor-critic · 13.4 MPC 5 个交互5 interactives

    1为什么机器人需要它Why robots need this

    20.4 节的 SAC 在一个单关节手臂上要 1000–1750 个环境步才学会到达目标(误差低于 8°),PPO 要几千到上万步。在仿真里这不算什么;在真机上,这意味着几十分钟到几小时的电机磨损和看护,复杂一点的任务要几百万步,根本付不起。24.1 节和 24.2 节已经给出了基于模型的答案:先学 \(s'=f(s,a)\),再在模型里规划(PETS)或生成假数据(MBPO)。真实机器人还要再过两道坎:

    • 机器人看到的是图像,不是状态。相机给出的是几万个像素,速度、被遮挡的物体都不在单张图里。我们需要一个从观测序列里学出来的紧凑状态,既要记住过去(17.2 节的 RNN),又要承认未来有不确定性(19.1 节的隐变量)。这就是 Dreamer 的循环状态空间模型 RSSM。
    • 模型越往后预测越不准。24.1 节看到了误差累积。所以两条路线都只在模型里走很短的一段(Dreamer 想象 15 步,TD-MPC 规划 3–5 步),更远的未来交给一个用时序差分(TD)学出来的价值函数(20.2、20.4 节)。

    围绕这两点,形成了两个家族:

    • Dreamer(Hafner 等):PlaNet(ICML 2019)学一个 RSSM 并在隐空间里用 CEM 规划;Dreamer(ICLR 2020)改为在想象中训练 actor-critic,部署时只需跑一次策略网络;DreamerV2(ICLR 2021)换成离散隐变量,在 Atari 上达到人类水平;DreamerV3(2023 年预印本,2025 年发表于 Nature)用一组固定超参数覆盖 150 多个任务,并且第一个从零开始(不用人类数据)在 Minecraft 里挖到钻石。DayDreamer(CoRL 2022)把 Dreamer 直接搬上真实机器人。
    • TD-MPC(Hansen、Wang、Su,ICML 2022)与 TD-MPC2(Hansen、Su、Wang,ICLR 2024):隐空间模型不重建图像,只要求它能预测奖励、价值和自己的下一个隐状态;每一步用 MPPI(6.11 节、13.4 节)在隐空间里规划一个短视界,末端接上学到的 Q 值。

    On a single-joint arm, the SAC of Lesson 20.4 needs 1000 to 1750 environment steps to learn to reach a target (error below 8°), and PPO thousands to tens of thousands. In simulation that is nothing; on a real robot it means tens of minutes to hours of motor wear and supervision, and harder tasks need millions of steps, which nobody can afford. Lesson 24.1 and Lesson 24.2 gave the model-based answer: learn \(s'=f(s,a)\), then plan in the model (PETS) or generate synthetic data from it (MBPO). Real robots add two hurdles:

    • Robots see images, not states. A camera delivers tens of thousands of pixels, and velocities or occluded objects are not in any single frame. We need a compact state learned from the observation sequence that remembers the past (the RNNs of Lesson 17.2) and admits that the future is uncertain (the latent variables of Lesson 19.1). That is Dreamer's recurrent state-space model, the RSSM.
    • Models get worse the further they predict. Lesson 24.1 showed the compounding error. So both lines of work go only a short way inside the model (Dreamer imagines 15 steps, TD-MPC plans 3 to 5) and leave the distant future to a value function learned by temporal differences (Lessons 20.2 and 20.4).

    Two families grew around these two points:

    • Dreamer (Hafner et al.): PlaNet (ICML 2019) learned an RSSM and planned in its latent space with CEM; Dreamer (ICLR 2020) instead trains an actor-critic in imagination, so that acting needs only one pass of a policy network; DreamerV2 (ICLR 2021) switched to discrete latents and reached human-level Atari; DreamerV3 (preprint 2023, published in Nature in 2025) covers more than 150 tasks with one fixed set of hyperparameters and was the first algorithm to collect diamonds in Minecraft from scratch, without human data. DayDreamer (CoRL 2022) took Dreamer straight onto real robots.
    • TD-MPC (Hansen, Wang & Su, ICML 2022) and TD-MPC2 (Hansen, Su & Wang, ICLR 2024): the latent model reconstructs no images; it only has to predict rewards, values and its own next latent state. At every step MPPI (Lesson 6.11, Lesson 13.4) plans a short horizon in latent space and adds a learned Q-value at the end.

    2直觉:一个会做梦的滤波器Intuition: a filter that can dream

    把 RSSM 想成 4.8 节卡尔曼滤波的学习版。每一步它维护两部分状态:确定性的 \(\mathbf h_t\)(一个 GRU 的隐藏状态,17.2 节)负责记忆,随机的 \(\mathbf z_t\) 负责"这一步可能发生什么"。它有两种工作方式:

    • 看着走(滤波):有观测 \(o_t\) 时,用后验 \(q_\theta(\mathbf z_t\mid\mathbf h_t,o_t)\) 推断隐状态,就像卡尔曼滤波的"更新"步。
    • 闭着眼走(做梦):没有观测时,只用先验 \(p_\theta(\hat{\mathbf z}_t\mid\mathbf h_t)\) 采样往前推,就像只有"预测"步。Dreamer 的想象、PlaNet 的规划都在这一模式里进行。

    Think of the RSSM as a learned version of the Kalman filter of Lesson 4.8. At every step it keeps a state in two parts: a deterministic \(\mathbf h_t\) (the hidden state of a GRU, Lesson 17.2) that carries memory, and a stochastic \(\mathbf z_t\) that captures "what might happen at this step". It runs in two modes:

    • Eyes open (filtering): when an observation \(o_t\) arrives, the posterior \(q_\theta(\mathbf z_t\mid\mathbf h_t,o_t)\) infers the latent, like the Kalman filter's update step.
    • Eyes closed (dreaming): without observations, the model samples from the prior \(p_\theta(\hat{\mathbf z}_t\mid\mathbf h_t)\) and rolls forward, like a filter that only predicts. Dreamer's imagination and PlaNet's planning both run in this mode.
    $$ \mathbf h_t=f_\theta(\mathbf h_{t-1},\mathbf z_{t-1},a_{t-1}),\qquad \hl{prior}{\hat{\mathbf z}_t\sim p_\theta(\hat{\mathbf z}_t\mid\mathbf h_t)},\qquad \mathbf z_t\sim q_\theta(\mathbf z_t\mid\mathbf h_t,o_t),\qquad \hat o_t\sim p_\theta(o_t\mid\mathbf h_t,\mathbf z_t),\qquad \hat r_t\sim p_\theta(r_t\mid\mathbf h_t,\mathbf z_t). $$

    为什么不只用一个确定性的 \(\mathbf h_t\)?因为未来常常分岔。下图是一个最小的例子:夹爪把一个销子放到一道窄脊上,销子的横向位置 \(x\) 每 0.1 秒更新一次。离开脊顶时它滑向较近的谷底(\(\pm0.7\));正好落在脊顶(\(|x|<0.12\))时,它以各一半的概率倒向左边或右边。两个一步模型同时在线训练:

    • 确定性模型 \(\hl{det}{x'=f(x)}\),用均方误差回归。在脊顶上它只能学到两种结果的平均,也就是"留在原地附近"。
    • 隐变量模型:先验 \(p(z\mid x)\)、后验 \(q(z\mid x,x')\)、解码器 \(x'=g(x,z)\),用 19.1 节的 ELBO 训练,KL 项的权重是 \(\beta\)。这就是一步的 RSSM(这里状态 \(x\) 直接可见,所以不需要 \(\mathbf h\))。

    训练之后,两个模型都从脊顶 \(x_0=0\) 出发"做梦" 10 步:随机模型每一步从先验采样 \(z\),确定性模型只有一条路。

    Why not use a deterministic \(\mathbf h_t\) alone? Because the future often forks. The figure shows a minimal example: a gripper sets a peg down on a narrow ridge, and the peg's lateral position \(x\) is updated every 0.1 s. Off the crest it slides toward the nearer valley (at \(\pm0.7\)); placed right on the crest (\(|x|<0.12\)) it tips left or right with probability one half each. Two one-step models train live side by side:

    • A deterministic model \(\hl{det}{x'=f(x)}\), regressed with mean squared error. On the crest it can only learn the average of the two outcomes, that is, "stay roughly where you are".
    • A latent-variable model: a prior \(p(z\mid x)\), a posterior \(q(z\mid x,x')\) and a decoder \(x'=g(x,z)\), trained with the ELBO of Lesson 19.1 with weight \(\beta\) on the KL term. This is a one-step RSSM (the state \(x\) is observed directly here, so no \(\mathbf h\) is needed).

    After training, both models "dream" 10 steps from the crest, \(x_0=0\): the stochastic model samples \(z\) from the prior at every step, the deterministic model has a single path.

    FIG 1一个输入,两种未来:确定性模型 vs 随机隐变量One input, two futures: deterministic vs stochastic latent
    从脊顶出发的 10 步:真实 vs 梦10 STEPS FROM THE CREST · real vs dreamed
    一步预测 x′ 的分布(x = 0)ONE-STEP x′ DISTRIBUTION (x = 0)
    左:横轴是步数,纵轴是销子位置 \(x\);浅紫色带是脊顶,两条绿色虚线是谷底。绿色细线是 60 条真实的未来,琥珀色细线是随机模型的 60 个梦(统计用 200 个),红色粗线是确定性模型唯一的梦。右:从 \(x=0\) 走一步的结果,绿色柱是真实分布(两个尖峰),琥珀色柱是随机模型从先验采样的 400 个预测,红色竖线是确定性模型的预测。改变 \(\beta\) 会从头重新训练(同一个随机种子)。 Left: steps on the horizontal axis, peg position \(x\) on the vertical; the light purple band is the crest and the two green dashed lines are the valleys. Thin green lines are 60 real futures, thin amber lines are 60 of the stochastic model's dreams (the statistics use 200), and the thick red line is the deterministic model's only dream. Right: one step from \(x=0\); green bars are the true distribution (two spikes), amber bars are 400 predictions sampled from the stochastic model's prior, and the red line is the deterministic model's prediction. Changing \(\beta\) restarts training (same random seed).
    试一试Try this
    • 用默认的 \(\beta=2\) 训练至少 500 步,直到左右两个谷底各收到至少 30% 的梦。红线这时在哪里?With the default \(\beta=2\), train for at least 500 steps until each valley receives at least 30% of the dreams. Where is the red line meanwhile?
    • 把 \(\beta\) 设为 20,训练至少 500 步,直到 KL(读数里的滑动平均)低于 0.01 nat,而且至少一个谷底收到的梦少于 15%。隐变量还在起作用吗?Set \(\beta\) to 20 and train for at least 500 steps until the KL (the running average in the readout) is below 0.01 nat and at least one valley receives fewer than 15% of the dreams. Is the latent still doing anything?
    • 把 \(\beta\) 设为 0.1,训练至少 500 步,直到 KL 高于 1 nat 而某一个谷底收到的梦少于 10%。重建误差很小,梦却是偏的:为什么?Set \(\beta\) to 0.1 and train for at least 500 steps until the KL is above 1 nat while one valley receives fewer than 10% of the dreams. The reconstruction error is small, yet the dreams are lopsided: why?
    核心洞见Key insight

    在种子 1 下,\(\beta=2\) 训练 500 步后,左右谷底分别收到约 51% 和 45% 的梦,KL 约 0.5–0.6 nat;确定性模型的唯一一条梦则随训练在左右之间摇摆,最后停在哪一边纯属偶然。能分岔的梦需要随机的隐状态,而随机隐状态需要先验和后验对得上。\(\beta=20\) 时 KL 被压到约 0.001 nat,后验不再携带信息,解码器学会忽略 \(z\)(后验坍缩,19.1 节),模型退化成确定性的;\(\beta=0.1\) 时后验自由地用 \(z\) 编码答案,KL 涨到 1.3–2 nat,先验却跟不上,从先验采到的 \(z\) 几乎都解码成"向右"(约 99%)。这正是 DreamerV2 提出 KL 平衡的原因:让 KL 主要去训练先验靠近后验,而不是把后验压向一个还没学好的先验(第 5 节)。

    With seed 1, after 500 steps at \(\beta=2\) the left and right valleys receive about 51% and 45% of the dreams, with a KL of about 0.5–0.6 nat; the deterministic model's single dream swings between left and right as training goes on, and where it ends up is pure chance. A dream that can fork needs a stochastic latent, and a stochastic latent needs its prior and posterior to agree. At \(\beta=20\) the KL is squeezed to about 0.001 nat, the posterior carries no information, the decoder learns to ignore \(z\) (posterior collapse, Lesson 19.1) and the model degenerates into a deterministic one. At \(\beta=0.1\) the posterior freely encodes the answer in \(z\) and the KL grows to 1.3–2 nat, but the prior cannot keep up, and almost every \(z\) sampled from the prior decodes to "right" (about 99%). This is exactly why DreamerV2 introduced KL balancing: let the KL term mostly train the prior toward the posterior, instead of pulling the posterior toward a prior that has not been learned yet (Section 5).

    3动手玩:Dreamer,在想象中学习Play: Dreamer, learning in imagination

    有了能做梦的世界模型,Dreamer 的训练循环只有三件事,交替进行:

    1. 采数据:用当前策略(加探索噪声)在真实环境里跑一个回合,存进重放缓冲区。
    2. 学世界模型:从缓冲区里取一段段真实序列,训练 RSSM(重建、奖励、KL)。
    3. 在想象中学行为:从缓冲区里的真实状态出发,用先验和当前策略想象 \(H\) 步,在这些梦里的轨迹上训练评论家和演员。真实环境一步都不用。

    评论家 \(v_\psi\) 回归的目标是 \(\lambda\)-回报,它把想象出的奖励和末端的价值混在一起(与 20.4 节 GAE 的 \(\lambda\) 是同一个想法):

    With a world model that can dream, Dreamer's training loop alternates three things:

    1. Collect data: run one episode in the real environment with the current policy (plus exploration noise) and store it in a replay buffer.
    2. Learn the world model: train the RSSM (reconstruction, reward, KL) on real sequences drawn from the buffer.
    3. Learn behaviour in imagination: starting from real states in the buffer, imagine \(H\) steps with the prior and the current policy, and train the critic and the actor on these dreamed trajectories. Not a single real step is used.

    The critic \(v_\psi\) regresses the \(\lambda\)-return, which mixes imagined rewards with values further along (the same idea as the \(\lambda\) of GAE in Lesson 20.4):

    $$ \hl{lamret}{G^\lambda_t}=\hat r_{t+1}+\gamma\Big[(1-\lambda)\,v_\psi(\hat s_{t+1})+\lambda\,G^\lambda_{t+1}\Big],\qquad G^\lambda_{\hl{Hh}{H}}=v_\psi(\hat s_{\hl{Hh}{H}}),\qquad \max_\phi\ \mathbb E\Big[\textstyle\sum_{t=0}^{\hl{Hh}{H}-1}\hl{lamret}{G^\lambda_t}\Big]. $$

    这里 \(\hat s_t=(\mathbf h_t,\hat{\mathbf z}_t)\) 是想象出的模型状态,\(a_t\sim\pi_\phi(a_t\mid\hat s_t)\)。第一版 Dreamer 的演员直接对模型求导:整条想象轨迹都是可微的网络(转移、奖励、价值、策略),所以 \(\sum_t G^\lambda_t\) 对策略参数 \(\phi\) 的梯度可以沿着轨迹反向传播回去(17.2 节的 BPTT),这叫"解析的价值梯度"。DreamerV3 改用 REINFORCE 式的梯度(第 5 节)。

    下图把这个循环放在 20.4 节的同一个任务上:受重力作用的单关节手臂,2 秒内转到随机目标并保持住,评估用同样的 24 组(起点,目标)。为了在浏览器里跑得动,世界模型直接作用在关节角和角速度上(不是从图像学出的隐状态),是一个 4→32→32→2 的 tanh 网络,预测一步的变化;奖励函数已知。演员和评论家都是 3→32→32→1 的网络。每个真实回合(40 步)之后做 40 次更新,每次从缓冲区取 8 个起点各想象 \(H\) 步,\(\gamma=0.97\),\(\lambda=0.95\),演员的梯度穿过学到的模型。

    Here \(\hat s_t=(\mathbf h_t,\hat{\mathbf z}_t)\) is an imagined model state and \(a_t\sim\pi_\phi(a_t\mid\hat s_t)\). The original Dreamer trains the actor by differentiating through the model: the whole imagined trajectory is made of differentiable networks (transition, reward, value, policy), so the gradient of \(\sum_t G^\lambda_t\) with respect to the policy parameters \(\phi\) can be backpropagated along the trajectory (the BPTT of Lesson 17.2). These are "analytic value gradients". DreamerV3 uses a REINFORCE-style gradient instead (Section 5).

    The figure runs this loop on the task of Lesson 20.4: a single arm joint under gravity must reach a random target within 2 seconds and hold it, evaluated on the same 24 (start, target) pairs. To run in a browser, the world model acts on the joint angle and velocity directly (not on a latent learned from images); it is a 4→32→32→2 tanh network that predicts the one-step change, and the reward function is known. Actor and critic are 3→32→32→1 networks. After every real episode (40 steps) come 40 updates, each imagining \(H\) steps from 8 start states drawn from the buffer, with \(\gamma=0.97\), \(\lambda=0.95\), and the actor's gradient flowing through the learned model.

    FIG 2Dreamer 实时训练:真实一步,想象一百步Dreamer training live: one real step, a hundred imagined
    最近一个真实回合(含探索噪声)LATEST REAL EPISODE (with exploration noise)
    评估误差(°,对数)· 真实 vs 想象EVAL ERROR (°, log) · real vs imagined
    左:琥珀色是手臂,绿色虚线是目标角度(从竖直向下量起),重放的是最近一个真实回合。右上:琥珀色是这次训练的评估误差,紫色虚线是上一次训练(换设置或重置之前),横轴是真实环境步数。右下:同一个起点和目标下,当前演员在真实物理里的关节角(绿色)和在学到的模型里想象的关节角(青色虚线),横轴是 40 步,阴影带是想象视界 \(H\),水平虚线是目标。 Left: the amber line is the arm and the green dashed line the target angle (measured from hanging straight down), replaying the latest real episode. Top right: amber is this run's evaluation error, the purple dashed line the previous run (before the last setting change or reset), against real environment steps. Bottom right: for one fixed start and target, the joint angle under the current actor in the real physics (green) and imagined in the learned model (cyan dashed) over 40 steps; the shaded band is the imagination horizon \(H\) and the horizontal dashed line is the target.
    试一试Try this
    • 用默认设置(\(H=15\),评论家开)按"训练",直到评估误差低于 3°。用了多少真实步、多少想象步?With the defaults (\(H=15\), critic on) press "Train" until the evaluation error drops below 3°. How many real steps and how many imagined steps did it take?
    • 选 \(H=1\) 并关掉评论家,训练至少 5 个回合:误差一直停在 30° 以上。演员在想象里到底能看到什么?Choose \(H=1\) with the critic off and train for at least 5 episodes: the error stays above 30°. What can the actor actually see in its imagination?
    • 保持 \(H=1\),打开评论家,训练到误差低于 25°。和上一个任务相比,是谁把一步之外的后果告诉了演员?Keep \(H=1\), turn the critic on and train until the error is below 25°. Compared with the previous task, who tells the actor about consequences beyond one step?
    核心洞见Key insight

    在种子 1 下,默认设置第 3 个回合(120 个真实步)误差就降到 7.8°,之后在 5°–17° 之间起伏,第 17 个回合(680 步)低于 3°,约 26 个回合后稳定在 1° 左右。每个真实回合配 \(40\times8\times15=4800\) 个想象步,是真实步数的 120 倍。作为对照,20.4 节里 SAC 要 1000–1750 步、PPO(\(N=16\))要 4480 步才第一次低于 8°。世界模型用计算换样本。评论家的作用要看视界:\(H=1\) 且没有评论家时,演员只看见一步的奖励 \(r(s_t,a_t)\),其中唯一依赖动作的是力矩惩罚,于是它学会"什么都不做",误差停在 36.6°;打开评论家后,价值函数替它看见了一步之外,第 21 个回合降到 21.8°(但过程很不稳定)。在这个任务里 15 步已经覆盖了整个动作,关掉评论家反而更快(第 4 个回合就低于 3°):评论家也是一个学出来的模型,早期的误差同样会被演员利用。

    With seed 1, the default run reaches 7.8° at episode 3 (120 real steps), wanders between 5° and 17°, drops below 3° at episode 17 (680 steps) and settles near 1° after about 26 episodes. Every real episode comes with \(40\times8\times15=4800\) imagined steps, 120 times the real ones. For comparison, in Lesson 20.4 SAC needed 1000–1750 steps and PPO (\(N=16\)) 4480 steps just to get below 8° for the first time. A world model trades computation for samples. What the critic does depends on the horizon: at \(H=1\) without a critic, the actor sees only the one-step reward \(r(s_t,a_t)\), whose only action-dependent part is the torque penalty, so it learns to do nothing and the error stays at 36.6°; with the critic on, the value function sees beyond the single step for it and the error reaches 21.8° at episode 21 (on a rocky path). On this task 15 steps already cover the whole motion, and switching the critic off is actually faster (below 3° at episode 4): the critic is one more learned model, and the actor exploits its early errors too.

    4第二个视角:TD-MPC,规划一小段,剩下的交给价值A second view: TD-MPC, plan a little, leave the rest to the value

    Dreamer 把规划"蒸馏"进一个策略网络;TD-MPC 反过来,每一步都真的规划。它的模型 TOLD(task-oriented latent dynamics,面向任务的隐动力学)有五个部件:编码器 \(\mathbf z_t=e_\theta(o_t)\)(论文记作 \(h_\theta\),这里换个字母,免得和 RSSM 的 \(\mathbf h_t\) 混淆)、隐动力学 \(\mathbf z_{t+1}=d_\theta(\mathbf z_t,a_t)\)、奖励 \(R_\theta(\mathbf z_t,a_t)\)、价值 \(Q_\theta(\mathbf z_t,a_t)\) 和一个策略 \(\pi_\theta(\mathbf z_t)\),全部是确定性的 MLP。规划时,MPPI 采样一批长度为 \(H\) 的动作序列,用模型推演,按下面的分数打分,取前 \(k\) 名加权平均,迭代几轮,执行第一个动作:

    Dreamer distils planning into a policy network; TD-MPC does the opposite and really plans at every step. Its model, TOLD (task-oriented latent dynamics), has five parts: an encoder \(\mathbf z_t=e_\theta(o_t)\) (written \(h_\theta\) in the paper; we use another letter so it does not clash with the RSSM's \(\mathbf h_t\)), latent dynamics \(\mathbf z_{t+1}=d_\theta(\mathbf z_t,a_t)\), a reward \(R_\theta(\mathbf z_t,a_t)\), a value \(Q_\theta(\mathbf z_t,a_t)\) and a policy \(\pi_\theta(\mathbf z_t)\), all deterministic MLPs. To plan, MPPI samples a batch of action sequences of length \(H\), rolls them through the model, scores them as below, takes a weighted average of the top \(k\), iterates a few times, and executes the first action:

    $$ \phi(a_{t:t+H})=\sum_{i=0}^{\hl{plan}{H}-1}\gamma^{i}R_\theta(\mathbf z_{t+i},a_{t+i})\;+\;\hl{Vterm}{\gamma^{H}Q_\theta(\mathbf z_{t+H},a_{t+H})},\qquad \mathbf z_{t+i+1}=d_\theta(\mathbf z_{t+i},a_{t+i}),\quad a_{t+H}=\pi_\theta(\mathbf z_{t+H}). $$

    没有末端那一项,这就是 PETS 那样的纯模型规划(24.2 节):只看得见 \(H\) 步。下图把它放进一个房间:移动机器人从左边出发,目标在右边,中间有一个开口朝向机器人的 U 形墙;每步最多移动 0.08,奖励是到目标距离的负数,\(\gamma=0.98\),每回合最多 70 步。直奔目标会一头扎进 U 形底部,而那里离目标只有 0.63,是一个很深的局部最优;正确的路要先绕开,大约 25 步。规划器是 TD-MPC 式的 MPPI:96 条样本、前 12 名、4 轮迭代、温度 0.5,上一步的解平移后作为下一步的初值。为了把规划单独拿出来看,这里用的是房间的真实动力学,价值是 40×40 网格上的一张表,用同步的 TD 备份 \(V(s)\leftarrow\max_a[r+\gamma V(s')]\)(16 个方向)训练;TD-MPC 用的是网络和学到的隐空间。

    Without the last term this is pure model-based planning like PETS (Lesson 24.2): it sees only \(H\) steps. The figure puts it in a room: a mobile robot starts on the left, the goal is on the right, and in between stands a U-shaped wall whose opening faces the robot. The robot moves at most 0.08 per step, the reward is minus the distance to the goal, \(\gamma=0.98\), and an episode lasts at most 70 steps. Heading straight for the goal drives the robot into the bottom of the U, only 0.63 from the goal: a deep local optimum. The right path goes around first and takes about 25 steps. The planner is TD-MPC-style MPPI: 96 samples, top 12, 4 iterations, temperature 0.5, and last step's solution shifted forward as the next initial guess. To isolate the planning, it uses the room's true dynamics, and the value is a table on a 40×40 grid trained by synchronous TD backups \(V(s)\leftarrow\max_a[r+\gamma V(s')]\) (over 16 directions); TD-MPC uses networks and a learned latent space.

    FIG 3短视界 + 末端价值:从 U 形陷阱里走出来Short horizon + terminal value: out of the U-shaped trap
    房间、价值表与 MPPI 样本ROOM · VALUE TABLE · MPPI SAMPLES
    底色是价值表:越青 = 价值越高(训练前全为 0,不画)。红色粗线是墙,琥珀色圆点是机器人,绿色圆圈是目标(半径 0.1 内算到达)。蓝色细线是最后一轮 MPPI 样本中的一半(每隔一条画一条),琥珀色线是前 12 名,琥珀色粗线是当前计划,绿色线是已走过的路径。"TD 备份"按钮开始或暂停价值训练(每次同步更新全部 1600 格),读数里的"残差"是最近一次备份中最大的变化量。 The background is the value table: more cyan = higher value (all 0 before training, not drawn). Thick red lines are walls, the amber dot is the robot and the green circle is the goal (within 0.1 counts as reached). Thin blue lines are every other MPPI sample of the last iteration, amber lines the top 12, the thick amber line the current plan, and the green line the path travelled so far. "TD backups" starts or pauses value training (each backup updates all 1600 cells at once); the readout's "residual" is the largest change in the latest backup.
    试一试Try this
    • 关掉末端价值,把视界调到 \(H\ge20\),按"运行"直到回合结束:机器人被困住了。再试更长的视界,最长 30。Turn the terminal value off, set the horizon to \(H\ge20\) and press "Run" until the episode ends: the robot is trapped. Try longer horizons, up to 30.
    • 按"TD 备份"把价值表训练到残差低于 0.01,打开末端价值,把视界设为 \(H\le3\),运行一个回合:机器人到达目标。Press "TD backups" and train the value table until the residual is below 0.01, turn the terminal value on, set the horizon to \(H\le3\) and run an episode: the robot reaches the goal.
    • 按"重置价值",只做不超过 15 次 TD 备份就暂停,打开末端价值再运行一个回合:机器人又被困住了。看看底色:价值信息传到哪里了?Press "Reset value", run at most 15 TD backups and pause, then run an episode with the terminal value on: the robot is trapped again. Look at the background: how far has the value information spread?
    核心洞见Key insight

    没有末端价值时,\(H\) 从 1 到 30 机器人全部停在离目标 0.63 的 U 形底部:要发现绕路更好,规划器得在 60 维的动作序列空间里偶然采到一条完整的绕路,96 条样本远远不够。价值表经过 80 次备份收敛(残差 < 0.01)后,\(H=1\) 到 10 都能在 23–26 步到达目标。价值函数把"远处"压缩成了一个数,规划只需负责附近。反过来,视界不是越长越好:用收敛的价值,\(H=15\) 要 43 步,\(H=25\) 要 62 步,\(H=30\) 又被困住,因为同样的样本数要搜索的维度翻了几倍。价值也必须先学好:15 次备份之后价值信息只从目标传出去约 15 步的距离,U 形开口处仍是一片 0(乐观的初值),规划器又被骗进陷阱;25 次之后 \(H\le10\) 已经够用。这就是 TD-MPC 的分工:TD 学习负责长远,MPC 负责眼前,而 TD-MPC2 默认只规划 \(H=3\) 步。

    Without a terminal value the robot ends up in the bottom of the U, 0.63 from the goal, for every \(H\) from 1 to 30: to discover that going around is better, the planner would have to sample a complete detour by chance in a 60-dimensional space of action sequences, and 96 samples are far too few. Once the value table has converged after 80 backups (residual < 0.01), every \(H\) from 1 to 10 reaches the goal in 23–26 steps. The value function compresses "far away" into one number, so planning only has to handle what is near. Conversely, a longer horizon is not better: with the converged value, \(H=15\) takes 43 steps, \(H=25\) takes 62, and \(H=30\) is trapped again, because the same number of samples now searches a space several times larger. And the value must be learned first: after 15 backups the value information has spread only about 15 steps out from the goal, the mouth of the U is still all 0 (an optimistic initial value), and the planner is lured into the trap again; after 25, \(H\le10\) already suffices. That is TD-MPC's division of labour: TD learning handles the long run and MPC the near term, and TD-MPC2 plans only \(H=3\) steps by default.

    要不要解码器?Do you need a decoder?

    Dreamer 的隐状态要能重建图像,这是一个很丰富的学习信号:每个像素都在提供梯度。代价是它在乎的是"图像里什么占的像素多",而不是"什么和任务有关"。TD-MPC 去掉了解码器,只用三个损失训练隐空间:预测奖励、预测价值(TD 目标),以及隐状态一致性:模型预测的下一个隐状态 \(d_\theta(\mathbf z_t,a_t)\) 要接近慢速移动平均的编码器 \(e_{\theta^-}\) 对真实下一帧的编码:

    Dreamer's latent must be able to reconstruct the image, a very rich learning signal: every pixel contributes gradient. The price is that it cares about "what covers many pixels" rather than "what matters for the task". TD-MPC drops the decoder and trains its latent space with three losses only: predict the reward, predict the value (TD target), and latent consistency: the model's predicted next latent \(d_\theta(\mathbf z_t,a_t)\) must match the encoding of the real next frame by a slowly moving average of the encoder, \(e_{\theta^-}\):

    $$ \mathcal L_{\text{recon}}=\hl{dec}{\big\|\hat o_t-o_t\big\|^2}+\big(R(\mathbf z_t)-r_t\big)^2,\qquad \mathcal L_{\text{TD-MPC}}=\hl{cons}{\big\|d_\theta(\mathbf z_t,a_t)-e_{\theta^-}(o_{t+1})\big\|^2}+\big(R_\theta(\mathbf z_t,a_t)-r_t\big)^2+\big(Q_\theta(\mathbf z_t,a_t)-y_t\big)^2. $$

    下图比较两种隐空间,都是 2 维。观测是 10×10 的图像:一个小圆斑是机器人(它在动作下移动,奖励是它到左下角目标的距离,取负),一个更大的圆斑是干扰物,比如镜头前走过的人,每一步跳到一个随机位置,和任务无关也无法预测。重建式模型训练编码器、解码器和奖励头;无解码器模型训练编码器、残差式隐动力学(一致性损失,目标编码器取移动平均 \(\tau=0.02\))和奖励头(同时预测当前和下一步的奖励)。我们不评估它们的损失,而是问:从 2 维隐码里能读回多少信息?用 5 近邻回归在 300 张固定的测试图上算 \(R^2\)(1 = 全部读回,0 或负数 = 读不回)。这里没有价值项;在 TD-MPC 里它和奖励项起同样的作用。

    The figure compares two latent spaces, both 2-dimensional. Observations are 10×10 images: a small blob is the robot (it moves under the actions, and the reward is minus its distance to a goal in the lower-left corner), and a larger blob is a distractor, say a person walking past the camera, who jumps to a random place every step, irrelevant to the task and unpredictable. The reconstruction model trains an encoder, a decoder and a reward head; the decoder-free model trains an encoder, residual latent dynamics (the consistency loss, with a target encoder averaged at \(\tau=0.02\)) and a reward head (predicting both the current and the next reward). Instead of judging their losses, we ask how much can be read back from the 2-D codes: \(R^2\) of 5-nearest-neighbour regression on 300 fixed test images (1 = everything recovered, 0 or negative = nothing). There is no value term here; in TD-MPC it plays the same role as the reward term.

    FIG 4隐空间里留下了什么:重建 vs 无解码器What a latent keeps: reconstruction vs decoder-free
    观测与重建(10×10)OBSERVATION · RECONSTRUCTION (10×10)
    300 张测试图的 2 维隐码2-D CODES OF 300 TEST IMAGES
    左:上面是一张测试观测(绿色小圆圈标出机器人的真实位置,紫色大圆圈标出干扰物,红色十字是目标),下面是重建式模型的重建。右:两个模型的隐码散点,每个点是一张测试图,青色点越不透明 = 机器人离目标越近(奖励越高)。如果隐空间编码了奖励,不透明度会沿某个方向渐变;如果只编码了干扰物,不透明的点和透明的点会混在一起。两个模型并排训练,每帧各走相同的步数;\(R^2\) 每 200 步重算一次。 Left: one test observation on top (the small green circle marks the robot's true position, the large purple circle the distractor, the red cross the goal) and the reconstruction model's reconstruction below. Right: the two models' codes, one dot per test image; the more opaque the cyan dot, the closer the robot is to the goal (the higher the reward). If a latent encodes the reward, the opacity varies along some direction; if it encodes only the distractor, opaque and transparent dots are mixed up. Both models train side by side with the same number of steps per frame; \(R^2\) is recomputed every 200 steps.
    试一试Try this
    • 保持干扰物开启,训练至少 1000 步,直到:重建式隐码读回干扰物位置的 \(R^2>0.8\)、读回奖励的 \(R^2<0.2\),而无解码器隐码读回奖励的 \(R^2>0.5\)。With the distractor on, train for at least 1000 steps until the reconstruction code recovers the distractor's position with \(R^2>0.8\) but the reward with \(R^2<0.2\), while the decoder-free code recovers the reward with \(R^2>0.5\).
    • 关掉干扰物,训练至少 400 步,直到重建式隐码读回机器人位置的 \(R^2>0.8\)。和无解码器的比,谁对机器人位置记得更全?Switch the distractor off and train for at least 400 steps until the reconstruction code recovers the robot's position with \(R^2>0.8\). Compared with the decoder-free code, which one remembers the robot's position more completely?
    核心洞见Key insight

    有干扰物时,训练 1200 步后重建式隐码读回干扰物的 \(R^2\) 是 0.95,读回奖励的 \(R^2\) 约为 0(−0.02):干扰物占的像素是机器人的 4 倍多,两个维度全被它占去了,尽管模型也在预测奖励。无解码器的隐码读回奖励的 \(R^2\) 是 0.64(2400 步时 0.75),干扰物约为 0:既然干扰物无法预测、和奖励无关,一致性和奖励损失都不需要它。没有干扰物时情况反过来:重建式在 200 步时读回机器人位置的 \(R^2\) 就有 0.83(2400 步 0.95);无解码器的读回奖励几乎完美(0.99),但机器人位置只有 0.6–0.75,因为它只需记住"离目标多远"。重建保留看得见的一切,无解码器只保留预测奖励和价值所需的东西。前者在像素都相关时学得又快又全,后者在杂乱、有干扰的真实画面里更稳,但信号更弱,学得也慢。

    With the distractor, after 1200 steps the reconstruction code recovers the distractor with \(R^2\) 0.95 and the reward with \(R^2\) about 0 (−0.02): the distractor covers more than 4 times as many pixels as the robot and takes over both dimensions, even though the model also predicts rewards. The decoder-free code recovers the reward with \(R^2\) 0.64 (0.75 at 2400 steps) and the distractor not at all: it cannot be predicted and has nothing to do with the reward, so neither the consistency loss nor the reward loss needs it. Without the distractor the picture flips: after only 200 steps the reconstruction code recovers the robot's position with \(R^2\) 0.83 (0.95 at 2400); the decoder-free code recovers the reward almost perfectly (0.99) but the robot's position only at 0.6–0.75, since all it needs is "how far from the goal". Reconstruction keeps everything visible; decoder-free keeps only what predicts rewards and values. The former learns fast and completely when all the pixels matter; the latter holds up better in cluttered real scenes with distractions, but its signal is weaker and it learns more slowly.

    一组超参数走天下:symlog 与 two-hotOne set of hyperparameters for everything: symlog and two-hot

    DreamerV3 和 TD-MPC2 的一个共同目标是:不调超参数,直接用到新任务上。最大的障碍是尺度:一个任务每步奖励约为 1,另一个约为 1000,回报和价值的量级可以差好几个数量级。两者的办法相同:评论家(以及奖励头)不直接回归原始的数值。DreamerV3 用

    A shared goal of DreamerV3 and TD-MPC2 is to work on new tasks without tuning hyperparameters. The biggest obstacle is scale: one task gives rewards of about 1 per step, another about 1000, and returns and values can differ by orders of magnitude. Both use the same remedy: the critic (and the reward head) does not regress raw numbers. DreamerV3 uses

    $$ \hl{symlog}{\operatorname{symlog}(x)=\operatorname{sign}(x)\ln(|x|+1)},\qquad \operatorname{symexp}(y)=\operatorname{sign}(y)\big(e^{|y|}-1\big),\qquad \hl{twohot}{\hat y=\textstyle\sum_i p_i\,b_i,\quad \mathcal L=-\sum_i t_i\log p_i}. $$

    symlog 把大数压扁、在 0 附近近似恒等;网络预测 symlog 值,用的时候再 symexp 回去。two-hot 把回归变成分类:在一组 symexp 间隔的格点 \(b_i\) 上输出 softmax 概率 \(p_i\),预测值是期望;目标 \(y\) 落在相邻两格之间,就按距离把权重分给这两格(目标向量 \(t\) 只有两个非零元,因此叫 two-hot),用交叉熵训练。交叉熵的梯度大小与 \(y\) 的量级无关。下图让三种头去拟合同一条"回报曲线",它乘上尺度 \(c\)(范围约 \(-1.01c\) 到 \(-0.04c\)),用同一个 Adam 学习率 \(3\times10^{-3}\)、同样的 1→32→32 网络;two-hot 用 41 个格点,覆盖 \(\pm\operatorname{symexp}(10)\approx\pm22026\)。

    symlog squashes large numbers and is close to the identity near 0; the network predicts symlog values, which are mapped back with symexp when used. Two-hot turns regression into classification: the network outputs softmax probabilities \(p_i\) over a grid of symexp-spaced bins \(b_i\), and the prediction is their expectation; a target \(y\) between two neighbouring bins splits its weight between them by distance (the target vector \(t\) has only two non-zero entries, hence "two-hot"), and training uses cross-entropy, whose gradient size does not depend on the magnitude of \(y\). In the figure three heads fit the same "return curve" multiplied by a scale \(c\) (it ranges from about \(-1.01c\) to \(-0.04c\)), with the same Adam learning rate \(3\times10^{-3}\) and the same 1→32→32 network; two-hot uses 41 bins covering \(\pm\operatorname{symexp}(10)\approx\pm22026\).

    FIG 5同一个学习率,三种尺度:MSE、symlog、two-hotOne learning rate, three scales: MSE, symlog, two-hot
    目标与三种预测(÷ c)TARGET AND PREDICTIONS (÷ c)
    相对误差(对数)vs 步数RELATIVE ERROR (log) vs STEPS
    左:横轴是状态 \(x\),纵轴是数值除以尺度 \(c\);虚线是目标,红色是原始 MSE 头,蓝色是 symlog 头,绿色是 two-hot 头(超出画面的部分被截掉)。右:相对误差 = 均方根误差 ÷ 目标的标准差(1 等于只会猜平均值),对数刻度。换尺度会从头训练三个头(同一个种子)。 Left: state \(x\) on the horizontal axis, value divided by the scale \(c\) on the vertical; the dashed line is the target, red the raw MSE head, blue the symlog head and green the two-hot head (parts outside the frame are cut off). Right: relative error = RMSE ÷ standard deviation of the targets (1 means no better than guessing the mean), on a log scale. Changing the scale restarts all three heads (same seed).
    试一试Try this
    • 选尺度 \(c=10^4\),训练至少 1500 步:原始 MSE 的相对误差仍大于 1,而 two-hot 低于 0.1。再回到 \(c=1\) 比较三者。Choose the scale \(c=10^4\) and train for at least 1500 steps: the raw MSE head's relative error is still above 1 while two-hot is below 0.1. Then go back to \(c=1\) and compare the three.
    核心洞见Key insight

    \(c=1\) 时三者都能学好(2000 步后 MSE 0.011、symlog 0.037、two-hot 0.029),原始 MSE 甚至最好。\(c=10^4\) 时 1500 步后 MSE 是 2.73,比只猜平均值还差;symlog 0.16,two-hot 0.075。原因不是梯度太大:Adam 会把梯度归一化,每一步参数最多移动约一个学习率,而输出要到达 \(-10^4\) 量级,参数得走很远的路。symlog 和 two-hot 把任何尺度都变成了网络习惯的 \(O(1)\) 数值。代价是精度:symlog 在压缩后的空间里平均误差,换回原尺度会放大;two-hot 的精度受格点间距限制。DreamerV3 对奖励头和评论家用 two-hot(格点按指数间隔),对向量观测用 symlog;TD-MPC2 对奖励和价值用对数变换空间里的离散回归(101 个格点)。

    At \(c=1\) all three learn well (after 2000 steps MSE 0.011, symlog 0.037, two-hot 0.029), and raw MSE is even the best. At \(c=10^4\), after 1500 steps MSE is at 2.73, worse than guessing the mean, while symlog is at 0.16 and two-hot at 0.075. The cause is not huge gradients: Adam normalizes the gradient, so each step moves a parameter by at most about one learning rate, and reaching outputs of order \(-10^4\) takes the parameters a very long way. symlog and two-hot turn any scale into the \(O(1)\) numbers networks are used to. The price is precision: symlog averages errors in the squashed space, which grow when mapped back; two-hot is limited by the bin spacing. DreamerV3 uses two-hot for the reward head and the critic (with exponentially spaced bins) and symlog for vector observations; TD-MPC2 uses discrete regression in a log-transformed space for rewards and values (101 bins).

    5再看数学The math

    RSSM 的训练目标:按时间展开的 ELBOTraining the RSSM: an ELBO unrolled in time

    记模型状态 \(s_t=(\mathbf h_t,\mathbf z_t)\)。对一段真实序列 \(o_{1:T},a_{1:T},r_{1:T}\),世界模型最小化

    Write the model state as \(s_t=(\mathbf h_t,\mathbf z_t)\). On a real sequence \(o_{1:T},a_{1:T},r_{1:T}\) the world model minimizes

    $$ \mathcal L(\theta)=\sum_{t=1}^{T}\Big[\hl{dec}{-\ln p_\theta(o_t\mid s_t)}-\ln p_\theta(r_t\mid s_t)+\beta\,\mathrm{KL}\big(q_\theta(\mathbf z_t\mid\mathbf h_t,o_t)\,\big\|\,\hl{prior}{p_\theta(\mathbf z_t\mid\mathbf h_t)}\big)\Big],\qquad \mathbf z_t\sim q_\theta,\ \ \mathbf h_t=f_\theta(\mathbf h_{t-1},\mathbf z_{t-1},a_{t-1}). $$

    这正是 19.1 节的 ELBO(取负号),只是每一步的先验不再是固定的 \(\mathcal N(0,I)\),而是由过去决定的 \(p_\theta(\mathbf z_t\mid\mathbf h_t)\),所以 KL 项同时做两件事:把后验拉向先验(正则化表示),把先验拉向后验(训练"做梦"的转移模型)。FIG 1 显示这两件事会互相拖累。DreamerV2 的 KL 平衡把它拆开、分别加权:\(\alpha\,\mathrm{KL}(\mathrm{sg}(q)\|p)+(1-\alpha)\,\mathrm{KL}(q\|\mathrm{sg}(p))\),\(\alpha=0.8\),sg 表示停止梯度;DreamerV3 写成"动力学损失"和"表示损失"两项,权重 1 和 0.1,并加上 free bits:每一项低于 1 nat 时截断为 1,不再优化,防止 KL 被压到 0(FIG 1 中 \(\beta=20\) 的后验坍缩)。DreamerV2 起,\(\mathbf z_t\) 是 32 个各有 32 类的类别变量,用直通(straight-through)估计反向传播;DreamerV3 再把每个类别分布混入 1% 的均匀分布,避免概率为 0。

    This is the (negated) ELBO of Lesson 19.1, except that each step's prior is no longer a fixed \(\mathcal N(0,I)\) but the past-dependent \(p_\theta(\mathbf z_t\mid\mathbf h_t)\). The KL term therefore does two jobs at once: it pulls the posterior toward the prior (regularizing the representation) and the prior toward the posterior (training the transition model that dreams). FIG 1 shows how the two get in each other's way. DreamerV2's KL balancing splits them and weights them separately, \(\alpha\,\mathrm{KL}(\mathrm{sg}(q)\|p)+(1-\alpha)\,\mathrm{KL}(q\|\mathrm{sg}(p))\) with \(\alpha=0.8\), where sg stops the gradient; DreamerV3 writes them as a "dynamics loss" and a "representation loss" weighted 1 and 0.1, and adds free bits: each term is clipped below 1 nat and no longer optimized there, which keeps the KL from being squeezed to 0 (the posterior collapse of FIG 1 at \(\beta=20\)). From DreamerV2 on, \(\mathbf z_t\) is a set of 32 categorical variables with 32 classes each, trained through straight-through gradients; DreamerV3 also mixes 1% of a uniform distribution into every categorical so that no probability is 0.

    在想象中学行为Behaviour learning in imagination

    从缓冲区里真实序列的后验状态出发,用先验和策略想象 \(\hat s_{t+1}\sim p_\theta(\cdot\mid\hat s_t,a_t)\),\(a_t\sim\pi_\phi(\cdot\mid\hat s_t)\),再算 \(\lambda\)-回报(第 3 节的递推)。评论家最小化 \(\sum_t\tfrac12\big(v_\psi(\hat s_t)-\mathrm{sg}(G^\lambda_t)\big)^2\)(DreamerV3 用 two-hot 交叉熵代替平方误差,并把评论家向自己的移动平均正则化)。演员的两种梯度:

    Starting from posterior states of real sequences in the buffer, imagine \(\hat s_{t+1}\sim p_\theta(\cdot\mid\hat s_t,a_t)\) with \(a_t\sim\pi_\phi(\cdot\mid\hat s_t)\), then compute \(\lambda\)-returns (the recursion of Section 3). The critic minimizes \(\sum_t\tfrac12\big(v_\psi(\hat s_t)-\mathrm{sg}(G^\lambda_t)\big)^2\) (DreamerV3 replaces the squared error by a two-hot cross-entropy and regularizes the critic toward a moving average of itself). The actor has two kinds of gradient:

    $$ \text{V1:}\ \ \max_\phi\ \mathbb E\Big[\sum_{t}\hl{lamret}{G^\lambda_t}\Big],\qquad \text{V3:}\ \ \mathcal L(\phi)=-\sum_t\mathrm{sg}\!\Big(\frac{\hl{lamret}{G^\lambda_t}-v_\psi(\hat s_t)}{\max(1,S)}\Big)\ln\pi_\phi(a_t\mid\hat s_t)-\eta\,\mathcal H\big[\pi_\phi(\cdot\mid\hat s_t)\big]. $$

    V1 的梯度是"解析的":沿想象轨迹反向传播,方差小,但要求模型可微而且导数可信(FIG 2 就是这样做的,梯度穿过学到的模型)。V3 对离散和连续动作都用 REINFORCE(20.3 节)加基线 \(v_\psi\),不需要模型的导数;为了让一个固定的熵系数 \(\eta=3\times10^{-4}\) 适用于所有任务,回报除以 \(\max(1,S)\),其中 \(S\) 是一批回报的第 95 与第 5 百分位数之差的移动平均:只缩小大回报,不放大小回报(否则稀疏奖励下会放大噪声)。V1 与 V3 的想象视界都是 \(H=15\),\(\lambda=0.95\)。

    V1's gradient is "analytic": backpropagated along the imagined trajectory, low-variance, but it needs a differentiable model whose derivatives can be trusted (FIG 2 does this, with gradients flowing through the learned model). V3 uses REINFORCE (Lesson 20.3) with the baseline \(v_\psi\) for both discrete and continuous actions and needs no model derivatives; so that one fixed entropy coefficient \(\eta=3\times10^{-4}\) works for every task, returns are divided by \(\max(1,S)\), where \(S\) is a moving average of the gap between the 95th and the 5th percentile of a batch of returns: large returns are scaled down but small ones are never scaled up (which would amplify noise under sparse rewards). Both V1 and V3 imagine \(H=15\) steps with \(\lambda=0.95\).

    TD-MPC:训练 TOLD 与规划TD-MPC: training TOLD, and planning

    从缓冲区取一段长度为 \(H\) 的真实片段,只编码第一帧 \(\mathbf z_t=e_\theta(o_t)\),之后的隐状态都由模型自己推出来 \(\mathbf z_{i+1}=d_\theta(\mathbf z_i,a_i)\),梯度沿时间反向传播:

    Take a real segment of length \(H\) from the buffer, encode only its first frame \(\mathbf z_t=e_\theta(o_t)\), let the model produce all later latents itself, \(\mathbf z_{i+1}=d_\theta(\mathbf z_i,a_i)\), and backpropagate through time:

    $$ \mathcal J(\theta)=\sum_{i=t}^{t+H}\rho^{\,i-t}\Big[c_1\big\|R_\theta(\mathbf z_i,a_i)-r_i\big\|^2+c_2\big\|Q_\theta(\mathbf z_i,a_i)-\big(r_i+\gamma\,Q_{\theta^-}(\mathbf z_{i+1},\pi_\theta(\mathbf z_{i+1}))\big)\big\|^2+c_3\,\hl{cons}{\big\|d_\theta(\mathbf z_i,a_i)-e_{\theta^-}(o_{i+1})\big\|^2}\Big],\qquad \mathcal J_\pi=-\sum_{i}\rho^{\,i-t}Q_\theta\big(\mathbf z_i,\pi_\theta(\mathrm{sg}(\mathbf z_i))\big). $$

    \(\theta^-\) 是参数的指数移动平均(目标网络,20.2 节),\(\rho<1\) 让近处的预测权重更大(论文记作 \(\lambda\),这里换成 \(\rho\),免得和 \(\lambda\)-回报混淆)。价值项就是 20.2 节的 TD 目标,只是"下一状态的最优动作"由策略 \(\pi_\theta\) 近似,因为在 TD 目标里每次都规划太贵了。规划时(TD-MPC 原文:\(H=5\)、512 条样本另加 5% 来自 \(\pi_\theta\) 的样本、6 轮迭代,取前 64 名),每轮按 \(\phi\) 选出前 \(k\) 名,按 \(\Omega_i\propto\exp(\phi_i)\) 的某个温度加权,重新估计每一步动作的均值和标准差;最后执行第一个动作,下一步把解平移一格作为初值。本页 FIG 3 用 \(\Omega_i=\exp\big((\phi_i-\phi_{\max})/0.5\big)\)。和 6.11 节的 MPPI 相比,唯一本质的变化是分数里多了末端价值,以及部分样本来自策略网络。

    \(\theta^-\) is an exponential moving average of the parameters (a target network, Lesson 20.2), and \(\rho<1\) weights near-term predictions more (the paper calls it \(\lambda\); we write \(\rho\) to avoid confusion with the \(\lambda\)-return). The value term is the TD target of Lesson 20.2, with the best next action approximated by the policy \(\pi_\theta\), because planning inside every TD target would be far too expensive. To plan (original TD-MPC: \(H=5\), 512 samples plus 5% from \(\pi_\theta\), 6 iterations, top 64), each iteration selects the top \(k\) by \(\phi\), weights them by \(\Omega_i\propto\exp(\phi_i)\) at some temperature, and refits the mean and standard deviation of every action in the sequence; finally the first action is executed and the solution, shifted by one step, initializes the next plan. FIG 3 on this page uses \(\Omega_i=\exp\big((\phi_i-\phi_{\max})/0.5\big)\). Compared with the MPPI of Lesson 6.11, the only essential changes are the terminal value in the score and the samples that come from the policy network.

    TD-MPC2 保留了这套框架,改的是让它在不调参的情况下变大变稳:所有 MLP 加 LayerNorm 和 Mish 激活;隐状态用 SimNorm 归一化(分成若干组,每组做 softmax,投影到单纯形上);5 个 Q 函数的集成,TD 目标取随机两个中的较小值(20.4 节双评论家的推广);奖励和价值用对数变换空间里的离散回归(101 个格点,即 FIG 5 的 two-hot 思想);策略先验用最大熵目标训练(20.4 节的 SAC);多任务时加一个可学习的任务嵌入。默认规划视界 \(H=3\)、512 条样本(其中 24 条来自策略先验)、前 64 名、6 轮迭代。

    TD-MPC2 keeps this framework and changes what is needed to make it bigger and more stable without tuning: LayerNorm and Mish activations in every MLP; the latent normalized by SimNorm (split into groups, each passed through a softmax, i.e. projected onto a simplex); an ensemble of 5 Q-functions whose TD target takes the minimum of two chosen at random (a generalization of the twin critics of Lesson 20.4); discrete regression of rewards and values in a log-transformed space (101 bins, the two-hot idea of FIG 5); a policy prior trained with a maximum-entropy objective (the SAC of Lesson 20.4); and a learnable task embedding for multi-task training. The default planner uses \(H=3\), 512 samples (24 of them from the policy prior), the top 64 and 6 iterations.

    6机器人中的它:真实机器人与两条路线的取舍In robotics: real robots and the trade-offs between the two lines

    Dreamer (V1–V3)TD-MPC / TD-MPC2
    隐状态Latent stateRSSM:GRU 的确定性部分 + 随机部分(V2 起为类别变量)RSSM: deterministic GRU part + stochastic part (categorical from V2 on)确定性 MLP 的隐向量(TD-MPC2 用 SimNorm);没有循环记忆Latent vector from deterministic MLPs (SimNorm in TD-MPC2); no recurrent memory
    学隐空间的信号What shapes the latent重建观测 + 奖励 + KLReconstruction + reward + KL奖励 + 价值(TD)+ 隐状态一致性,不重建Reward + value (TD) + latent consistency; no reconstruction
    怎样行动How it acts策略网络一次前向,部署便宜One forward pass of the policy network; cheap to deploy每一步在隐空间里做 MPPI,需要更多计算,但能在测试时纠正策略的错误MPPI in latent space at every step: more computation, but it can correct the policy at test time
    长远The long run想象 15 步 + \(\lambda\)-回报中的价值15 imagined steps + the value inside the \(\lambda\)-return规划 3–5 步 + 末端 Q3–5 planned steps + a terminal Q
    典型强项Typical strengths部分可观测、需要记忆的任务;离散和连续动作;从图像学习Partially observed tasks that need memory; discrete and continuous actions; learning from images高维连续控制(灵巧手、肌肉骨骼、人形);多任务大模型High-dimensional continuous control (dexterous hands, musculoskeletal models, humanoids); large multi-task models
    真实系统里In real systems

    DayDreamer(Wu、Escontrela、Hafner、Goldberg、Abbeel,CoRL 2022)不用仿真器,直接在真实机器人上跑 Dreamer:Unitree A1 四足机器人从背朝下躺着开始,在 1 小时内学会翻身、站起来并行走,中途不需要人工复位;之后用棍子推它,它在 10 分钟的继续学习内学会扛住轻推、被推倒后快速翻身站起。在 UR5 和 XArm 两台机械臂上,它从相机图像和稀疏奖励学会多物体的抓放,表现接近人类遥操作;Sphero 小球机器人只看图像学会导航到目标。TD-MPC 系列:TD-MPC2 用同一组超参数在 DMControl、Meta-World、ManiSkill2、MyoSuite 的 104 个连续控制任务上超过基线,并训练了一个 3.17 亿参数的单一智能体完成 80 个任务;Feng、Hansen 等(CoRL 2023)先用离线数据预训练 TD-MPC 的世界模型,再在真实的 xArm 上微调,在一个带有未见过干扰物的视觉抓取任务里,只用 20 次试验把成功率从 22% 提到 67%。DreamerV3 的 Minecraft 钻石实验说明了同一套方法能处理极长的任务链。软件:两条路线都有作者的开源实现(DreamerV3 基于 JAX,TD-MPC2 基于 PyTorch)。往后看:把模型换成大规模视频预测器的"世界模型"是 Part VII 的话题;只用离线数据学习时的分布偏移见 20.5 节;把学到的模型放进经典 MPC(梯度法、线性化、残差)见 24.4 节。

    DayDreamer (Wu, Escontrela, Hafner, Goldberg & Abbeel, CoRL 2022) ran Dreamer directly on real robots without simulators: a Unitree A1 quadruped starting on its back learned to roll over, stand up and walk within 1 hour, without manual resets; when then pushed with a pole, it learned within 10 minutes of further training to withstand light pushes and to roll back onto its feet quickly after hard ones. On two arms, a UR5 and an XArm, it learned multi-object pick and place from camera images and sparse rewards, approaching human teleoperators; a Sphero robot learned to navigate to a goal from images alone. The TD-MPC line: with one set of hyperparameters, TD-MPC2 beat its baselines on 104 continuous-control tasks across DMControl, Meta-World, ManiSkill2 and MyoSuite, and trained a single 317-million-parameter agent on 80 tasks; Feng, Hansen et al. (CoRL 2023) pretrained a TD-MPC world model on offline data and finetuned it on a real xArm, raising the success rate of a visual pick task with unseen distractors from 22% to 67% in only 20 trials. DreamerV3's Minecraft diamonds show that the same recipe can handle extremely long chains of subgoals. Software: both lines have open-source implementations by their authors (DreamerV3 in JAX, TD-MPC2 in PyTorch). Looking ahead: "world models" built from large video predictors are a topic for Part VII; the distribution shift of learning from offline data alone is the subject of Lesson 20.5; plugging learned models into classical MPC (gradient-based, linearized, residual) is Lesson 24.4.

    7常见误区Common pitfalls

    误区 1:忘了优化器会利用模型的错误Pitfall 1: forgetting that the optimizer exploits model errors

    演员和规划器都在模型里最大化回报,哪里模型过于乐观,它们就往哪里去。FIG 2 里打开评论家时误差先降后升,就是演员在追逐一个还没学好的价值函数的梯度。对策:想象和规划只走短视界;不停地采新的真实数据,让模型在策略去的地方变准;用集成估计不确定性并惩罚它(24.2 节的 PETS、MBPO,21.3 节);评估时始终看真实环境里的回报,而不是模型自己预测的回报。

    Both the actor and the planner maximize return inside the model, and they go wherever the model is too optimistic. In FIG 2 with the critic on, the error first drops and then rises again: the actor is chasing the gradient of a value function that has not been learned yet. Remedies: imagine and plan only short horizons; keep collecting fresh real data so the model becomes accurate where the policy goes; estimate uncertainty with ensembles and penalize it (PETS and MBPO in Lesson 24.2, Lesson 21.3); and always evaluate the return in the real environment, never the return the model predicts for itself.

    误区 2:只数真实步数,不数计算Pitfall 2: counting real steps and ignoring computation

    "样本效率高 10 倍"说的是真实交互。FIG 2 里每个真实步配 120 个想象步,每个想象步还要过三个网络并反向传播;TD-MPC 每个控制周期要推演几千条隐轨迹。在真机上这通常值得(电机时间比 GPU 时间贵),但控制频率受限于规划速度,比较算法时要同时报告真实步数和墙钟时间。

    "Ten times more sample-efficient" refers to real interaction. In FIG 2 every real step comes with 120 imagined ones, each passing through three networks forwards and backwards; TD-MPC rolls out thousands of latent trajectories in every control cycle. On real hardware this is usually worth it (robot time costs more than GPU time), but the control rate is limited by planning speed, and comparisons should report both real steps and wall-clock time.

    误区 3:把 KL 当成越小越好的正则项Pitfall 3: treating the KL as a regularizer that should be as small as possible

    在世界模型里,KL 衡量的是"观测带来了多少先验预料不到的信息"。KL 接近 0(FIG 1 的 \(\beta=20\))通常不是好事:后验坍缩,模型不再用随机隐变量,梦变成确定性的;KL 很大而先验跟不上(\(\beta=0.1\))则梦是偏的。看重建误差判断不了这两种失败,要直接检查从先验出发的开环预测。KL 平衡和 free bits 就是为此设计的。

    In a world model the KL measures "how much information the observation brings that the prior did not expect". A KL near 0 (FIG 1 at \(\beta=20\)) is usually bad news: the posterior has collapsed, the model no longer uses its stochastic latent and its dreams become deterministic; a large KL with a prior that cannot keep up (\(\beta=0.1\)) gives lopsided dreams. The reconstruction error reveals neither failure; check open-loop predictions from the prior directly. KL balancing and free bits were designed for exactly this.

    误区 4:以为视界越长越好,或者价值可以晚点再学Pitfall 4: assuming longer horizons are better, or that the value can wait

    FIG 3 里没有价值时 \(H=30\) 也逃不出陷阱,有价值时 \(H=30\) 反而又被困住:模型误差随视界累积(24.1 节),采样规划的搜索维度随视界线性增长。另一方面,乐观初始化(全 0)的价值在没学好时会把规划器骗进陷阱。TD-MPC 用短视界,并且价值和模型从第一步起就一起训练。

    In FIG 3 \(H=30\) cannot escape the trap without a value, and with the value \(H=30\) gets trapped again: model errors accumulate with the horizon (Lesson 24.1) and the search space of sampling-based planning grows linearly with it. On the other hand, an optimistically initialized value (all 0) that is not learned yet lures the planner into the trap. TD-MPC uses short horizons and trains the value together with the model from the very first step.

    误区 5:以为一种隐空间对所有任务都最好Pitfall 5: assuming one kind of latent is best for every task

    重建式隐空间会把容量花在大面积、但和任务无关的东西上(FIG 4 的干扰物:光照、背景、走过的人);无解码器的隐空间只记住预测这个奖励所需的信息,换一个奖励(比如目标换了位置)它可能已经把需要的信息扔掉了(FIG 4 无干扰时只记住了"离目标多远")。也别忘了尺度:同一组超参数要跨任务使用,奖励和价值头就需要 symlog 或 two-hot 这样的尺度无关的目标(FIG 5)。

    A reconstruction latent spends its capacity on large but task-irrelevant things (FIG 4's distractor: lighting, background, a passer-by); a decoder-free latent keeps only what predicts this reward, and for a different reward (say, the goal moves) it may already have thrown away what is needed (without the distractor, FIG 4's decoder-free code remembered only "how far from the goal"). And mind the scale: to reuse one set of hyperparameters across tasks, the reward and value heads need scale-free targets such as symlog or two-hot (FIG 5).

    8小练习Check yourself

    1. RSSM 为什么除了确定性的 \(\mathbf h_t\) 还要一个随机的 \(\mathbf z_t\)?1. Why does the RSSM need a stochastic \(\mathbf z_t\) besides the deterministic \(\mathbf h_t\)?
    选项 B。确定性模型对分岔的未来只能预测平均或者随便选一支(FIG 1 的红线);随机的 \(\mathbf z_t\) 从先验采样,梦才能落到两个谷底。GRU 当然能记住多步信息,这正是 \(\mathbf h_t\) 的作用(选项 A 错);重建误差主要由后验决定,与此无关(选项 C);奖励头照样需要(选项 D)。Option B. A deterministic model can only predict the average of a forking future or pick one branch arbitrarily (FIG 1's red line); a stochastic \(\mathbf z_t\) sampled from the prior lets the dreams reach both valleys. A GRU can certainly remember over many steps, which is exactly the job of \(\mathbf h_t\) (option A is wrong); the reconstruction error is mainly a matter of the posterior and unrelated (option C); a reward head is needed anyway (option D).
    2. Dreamer 训练演员和评论家时,用的轨迹从哪里来?2. Where do the trajectories come from on which Dreamer trains its actor and critic?
    选项 D。真实数据只用来训练世界模型;行为学习完全在想象中进行,起点是真实经验里的状态,长度只有 \(H\)(15 步),更远的部分由 \(\lambda\)-回报里的价值负责。选项 A 是无模型方法;选项 B 会失去样本效率;选项 C 的长想象会累积模型误差。Option D. Real data trains only the world model; behaviour learning happens entirely in imagination, starting from states in real experience and lasting only \(H\) steps (15), with the value inside the \(\lambda\)-return covering what lies beyond. Option A is a model-free method; option B would throw away the sample efficiency; option C's long imagination would accumulate model errors.
    3. TD-MPC2 默认只规划 \(H=3\) 步,却能完成需要几百步的任务。最主要的原因是?3. TD-MPC2 plans only \(H=3\) steps by default, yet solves tasks that take hundreds of steps. What is the main reason?
    选项 A。FIG 3 中没有价值时任何视界都被困住,有了收敛的价值 \(H=1\) 到 10 都能到达。模型不会"在 3 步后完全准确"(选项 B),采样也搜索不了整条长轨迹(FIG 3 中 \(H=30\) 失败,选项 C),TD-MPC 根本不重建图像(选项 D)。Option A. In FIG 3 every horizon is trapped without a value, and with the converged value every \(H\) from 1 to 10 reaches the goal. No model is "perfectly accurate after 3 steps" (option B), sampling cannot search a whole long trajectory (\(H=30\) fails in FIG 3, option C), and TD-MPC does not reconstruct images at all (option D).
    4. 一台机械臂的相机视野里有一台不停播放视频的显示器。哪种世界模型的隐空间最可能被它带偏?4. A robot arm's camera view includes a monitor that keeps playing videos. Whose latent space is most likely to be led astray by it?
    选项 C。重建损失按像素求和,显示器占很多像素、变化又大,隐空间会优先编码它(FIG 4 中重建式隐码读回干扰物的 \(R^2\) 为 0.95、读回奖励约为 0)。不可预测又与奖励无关的内容对奖励、价值和一致性损失都没有帮助,所以选项 A、B 的隐空间不会去编码它。这也是 TD-MPC 系列选择不重建的理由之一。Option C. The reconstruction loss sums over pixels; the monitor covers many pixels and changes a lot, so the latent encodes it first (in FIG 4 the reconstruction code recovers the distractor with \(R^2\) 0.95 and the reward with about 0). Content that is unpredictable and unrelated to the reward helps none of the reward, value or consistency losses, so the latents of options A and B have no reason to encode it. This is one of the reasons the TD-MPC line chose not to reconstruct.
    5. 在 FIG 5 中,尺度 \(c=10^4\) 时原始 MSE 头用 Adam 学得极慢。最直接的原因是?5. In FIG 5 at scale \(c=10^4\), the raw MSE head learns extremely slowly with Adam. What is the most direct reason?
    选项 B。Adam 用梯度的二阶矩归一化,更新量的大小约为学习率,与梯度量级无关,所以不会因为梯度大而发散(选项 A 错),但也不会因为误差大而走得更快。线性输出层要产生 \(-10^4\) 的值,权重或偏置得变得很大,以 \(3\times10^{-3}\) 一步需要成千上万步。隐藏层输入在 \([-1,1]\) 内,不会因为目标尺度而饱和(选项 C);目标只是乘了一个常数(选项 D)。symlog 和 two-hot 把输出拉回 \(O(1)\)。Option B. Adam normalizes by the gradient's second moment, so each update has a size of about the learning rate whatever the gradient's magnitude: it does not diverge because gradients are large (option A is wrong), but it does not move faster when the error is large either. For the linear output layer to produce values around \(-10^4\), weights or biases must become very large, which at \(3\times10^{-3}\) per step takes many thousands of steps. The hidden layers' inputs are in \([-1,1]\) and do not saturate because of the target scale (option C); the target is just multiplied by a constant (option D). symlog and two-hot bring the outputs back to \(O(1)\).

    9延伸阅读Further reading

    • R. S. Sutton, Dyna, an Integrated Architecture for Learning, Planning, and Reacting, SIGART Bulletin, 1991 — "用模型生成经验来学习"的源头The origin of "learning from experience generated by a model"
    • D. Ha & J. Schmidhuber, Recurrent World Models Facilitate Policy Evolution, NeurIPS 2018 — "World Models":VAE + RNN 的世界模型,在梦里训练控制器"World Models": a VAE + RNN world model, with a controller trained inside the dream
    • D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee & J. Davidson, Learning Latent Dynamics for Planning from Pixels, ICML 2019 — PlaNet:RSSM 与隐空间 CEM 规划PlaNet: the RSSM and CEM planning in latent space
    • D. Hafner, T. Lillicrap, J. Ba & M. Norouzi, Dream to Control: Learning Behaviors by Latent Imagination, ICLR 2020 — Dreamer:想象中的 actor-critic 与解析价值梯度Dreamer: actor-critic in imagination with analytic value gradients
    • D. Hafner, T. Lillicrap, M. Norouzi & J. Ba, Mastering Atari with Discrete World Models, ICLR 2021 — DreamerV2:类别隐变量与 KL 平衡DreamerV2: categorical latents and KL balancing
    • D. Hafner, J. Pasukonis, J. Ba & T. Lillicrap, Mastering Diverse Control Tasks through World Models, Nature 640, 2025 (arXiv:2301.04104, 2023, Mastering Diverse Domains through World Models) — DreamerV3
    • P. Wu, A. Escontrela, D. Hafner, K. Goldberg & P. Abbeel, DayDreamer: World Models for Physical Robot Learning, CoRL 2022
    • J. Schrittwieser et al., Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model, Nature, 2020 — MuZero:不重建观测、只预测奖励价值策略的模型MuZero: a model that predicts rewards, values and policies without reconstructing observations
    • N. Hansen, X. Wang & H. Su, Temporal Difference Learning for Model Predictive Control, ICML 2022 — TD-MPC
    • N. Hansen, H. Su & X. Wang, TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024
    • Y. Feng, N. Hansen, Z. Xiong, C. Rajagopalan & X. Wang, Finetuning Offline World Models in the Real World, CoRL 2023
    • G. Williams, A. Aldrich & E. A. Theodorou, Model Predictive Path Integral Control: From Theory to Parallel Computation, Journal of Guidance, Control, and Dynamics, 2017 — MPPI