世界模型是什么？

The image of the world around us, which we carry in our head, is just a model. Nobody in his head imagines all the world, government or country. He has only selected concepts, and relationships between them, and uses those to represent the real system. —Jay Wright Forrester
以上是系统动力学之父对mental model或world model的描述。

构建方法

世界模型主要包含一个任务动力学模型，用于在给定当前的动作和状态情况下预测下一步的状态。

recurrent world model利用recurrent state space model建模世界模型，用于学习隐式动力学，从而实现智能体的规划。该模型把状态空间分为两部分，分别是隐藏层的确定性状态$h_t$和随机部分$s_t$。

iVideoGPT*利用交互式视频生成模型作为世界模型，如图1所示，状态空间模型、视频生成模型、以及交互式视频生成模型的三者区别。

图1 概念比较

RoboScape 被设计frame级别动作可控的机器人视频生成。为了提升视频生成的可行性，引入了两个辅助训练任务，一个是用于提升几何一致性的时序深度预测，另一个是用于捕获局部动态对象的运动和变形细节的自适应采样关键点的动力学学习。

用处

有助于RL实现通用能力

Dreamer-v3基于通过学习世界模型可使智能体拥有丰富的感知和想象未来的能力的想法，利用世界模型预测潜在动作的结果，critic网络预测每个网络的价值，actor神经网络选择可实现最好结果的动作。为了实现较强的表现，Dreamer-v3利用了normalization、balancing、以及transformations鲁棒性技术。这些元件从replay-buffer中并发训练，同时智能体与环境交互。为了应用各种领域，三个元件需要包含不同的信号幅度和目标函数中的robust平衡项。确切的说，世界模型主要包含
$$
\begin{aligned}
Sequence~model&:h_t=f_{\phi}(h_{t-1},z_{t-1},a_{t-1}) \\
Encoder&:z_t\sim q_{\phi}(z_t\vert h_t,x_t) \\
Dynamics~predictor&:\hat{z}_t\sim p_{\phi}(\hat{z}_t\vert h_t) \\
Reward~predictor&:\hat{r}_t\sim p_{\phi}(\hat{r}_t\vert h_t,z_t) \\
Continue~predictor&:\hat{c}_t\sim p_{\phi}(\hat{c}_t\vert h_t,z_t) \\
Decoder&:\hat{x}_t\sim p_{\phi}(\hat{x}_t\vert h_t,z_t)
\end{aligned}\tag{1}
$$
世界模型利用recurrent state-space模型进行建模。编码器用于学习随机的抽象表示，便于世界模型适用于多领域的环境。解码器为了确保隐藏层状态和编码包含足够多的未来状态信息。在actor和critic的训练，主要依赖世界模型预测的抽象轨迹表示，即：
$$
\begin{aligned}
Actor&:a_t\sim\pi_{\theta}(a_t\vert s_t) \\
critic&:v_{\psi}(R_t\vert s_t)
\end{aligned}\tag{2}
$$
式中$s_t\doteq\{h_t,z_t\}$

世界模型的损失函数：
$$
\begin{aligned}
\mathcal{L}(\phi)\doteq E_{q_{\phi}}[\sum_{t=1}^T(\beta_{pred}\mathcal{L}_{pred}(\phi)+\beta_{dyn}\mathcal{L}_{dyn}(\phi)+\beta_{rep}\mathcal{L}_{rep}(\phi))]
\end{aligned}\tag{3}
$$

$$
\begin{aligned}
\mathcal{L}_{pred}(\phi)&\doteq-logp_{\phi}(x_t\vert z_t,h_t)-logp_{\phi}(r_t\vert z_t,h_t)-logp_{\phi}(c_t\vert z_t,h_t) \\
\mathcal{L}_{dyn}(\phi)&\doteq max(1,KL[sg(q_{\phi}(z_t\vert h_t,x_t))\Vert p_{\phi}(z_t\vert h_t)]) \\
\mathcal{L}_{rep}(\phi)&\doteq max(1,KL[q_{\phi}(z_t\vert h_t,x_t)\Vert sg(p_{\phi}(z_t\vert h_t))])
\end{aligned}\tag{4}
$$
为了平衡世界模型训练的损失函数，采用free-bits和表示预测较小的系数，从而平衡表示损失和避免退化的表现。此外，为了避免大的输入和大的重构梯度，利用symlog函数对输入观测进行变换；为了避免KL-Divergence振荡的spikes，参数化编码器和动力学预测器为$1%$均匀分布和$99%$神经网络输出的类别分布。
critic loss
为了考虑更长远的奖励，critic学习预测回报的分布，即
$$
\begin{aligned}
\mathcal{L}(\psi)&\doteq-\sum_{t=1}^Tlnp_{\psi}(R_t^{\lambda}\vert s_t) \\
R_t^{\lambda}&\doteq r_t+\gamma c_t((1-\lambda)v_t+\lambda R_{t+1}^{\lambda})\\
R_T^{\lambda}&\doteq v_T
\end{aligned}\tag{5}
$$
式中$r_t$为环境的奖励，而$R_{t+1}^{\gamma}$为预测的奖励。
同时，由于不同环境的回报分布显著不同，因此参数化critic为类别部分。

actor loss
对于离散和连续动作，均利用Reinforce估计器。同时，利用熵正则化器鼓励探索
$$
\begin{aligned}
\mathcal{L}(\theta)\doteq-\sum_{t=1}^T sg((R_t^{\lambda}-v_t)/max(1,S))log\pi_{\theta}(a_t\vert s_t)+\eta H(\pi_{\theta}(a_t\vert s_t))
\end{aligned}\tag{6}
$$
式中$S\doteq EMA(Percentage(R_t^{\lambda},95)-Percentage(R_t^{\lambda},5), 0.99)$为回报缩放项，避免奖励对探索产生影响。$max(1,S)$是为了避免对小回报进行scale，从而放大噪音。

Robust Prediction
各种领域奖励信号的规模是不同的，MSE预测大的目标值会发散，而Huber和$L_1$损失会停止学习，normalization会引入非静态。为了解决该困境，利用symlog平方误差作为损失函数，即
$$
\begin{aligned}
\mathcal{L}(\theta)\doteq\frac{1}{2}(f(x,\theta)-symlog(y))^2~\hat{y}=symexp(f(x,\theta))
\end{aligned}\tag{7}
$$
为了预测负的目标值，利用bi-symmentric logarithmic，即
$$
\begin{aligned}
symlog(x)&\doteq sign(x)log(\vert x\vert+1) \\
symexp(x)&\doteq sign(x)exp(\vert x\vert-1)
\end{aligned}\tag{8}
$$
为了应对奖励和回报的噪音化目标，引入了symexp two-hot loss。对于连续值，把支持区间划分为bins，预测目标为softmax分布，值为bin位置的加权平均，从而可以输出任意连续值，即
$$
\begin{aligned}
\hat{y}\doteq softmax(f(x))^TB \\
B\doteq symexp([-20,\ldots,+20])
\end{aligned}\tag{9}
$$
确切的说，two-hot指的是$\vert B\vert$个bins中与编码标量相近的两个$k,k+1$桶的概率，那么损失函数为
$$
\begin{aligned}
\mathcal{L}(\theta)\doteq-twohot(y)^Tlog~softmax(f(x,\theta))
\end{aligned}\tag{10}
$$
最终，梯度的幅度与信号scale*进行了解耦合。

Spatial Reasoning

MindJourney通过Test-Time scaling框架把VLM与世界模型相结合，从而提高视觉语言模型的空间推理能力，即给定观测和动作下3D场景的变化。