论文阅读笔记——TesserAct: Learning 4D Embodied World Models

TesserAct 论文
采用RGB-DN（RGB+深度+法线） 作为 4D 场景中间表示，由此建模 4D 场景，比纯 2D 视频更准确地建模 3D 几何结构。相比现有的 4D 视频生成，优化速度快，收敛好，且首次从当前帧和文本描述的具身智能体动作直接预测4D场景。
建模分布为 $p(v,d,n|v^0,d^0,n^0,T)$ 其中 v, d, n 表示预测的未来 RGB、深度图和法线图的潜在序列，$v^0$、$d^0$、$n^0$ 是 RGB 图像、深度和法线图的潜在表示以及具身智能体的文本动作。
$$\begin{array}{r}q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})=\mathcal{N}\left({\mathbf{z}}_t;\sqrt{{\alpha }_t}{\mathbf{z}}_{t-1},(1-{\alpha }_t)\mathbf{I}\right)\\ p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}^0,\mathcal{T})=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t,{\mathbf{x}}^0,\mathcal{T}),{\Sigma }_{\theta }({\mathbf{x}}_t,t)\right)\end{array}$$
通过法线积分优化深度：
在第 i 帧的 2D 图像坐标系中，像素位置 $u=(u,v{)}^T\in {\mathbf{V}}^i$，其对应的深度标量、法向量为 $d\in D^i,n=(n_x,n_y,n_z)\in N^2$。在焦距为 f，主点为 $(c_u,c_v{)}^T$ 的透视相机假设下，对数深度 $\tilde{d}=log(d)$ 应满足一下方程：${\tilde{n}}_z{\partial }_v\tilde{d}+n_x=0$ 和 ${\tilde{n}}_z{\partial }_v\tilde{d}+n_y=0$，其中 $\tilde{n}=n_x(u-c_x)+n_y(v-c_y)+n_{z}f$。由此得到：
$$\underset{d}{\min }{\iint }_{\Omega }({\tilde{n}}_z{\partial }_u\tilde{d}+n_x{)}^2+({\tilde{n}}_z{\partial }_u\tilde{d}+n_y{)}^2\mathrm{d}u\mathrm{d}v.$$
转化为迭代预测得到：
$${\tilde{d}}_{t+1}=\arg \underset{\tilde{d}}{\min }(A\tilde{d}-b{)}^{T}W({\tilde{d}}_t)(A\tilde{d}-b)\stackrel{\mathsf{def}}{=}\arg \underset{\tilde{\mathcal{D}}}{\min }{\mathcal{L}}_s(\tilde{\mathcal{D}},{\mathcal{N}}^i)$$

• 数据集：从 RLBench 中选择 20 个难度较高的任务，为每个任务从 4 个不同视角生成 1000 个实例，共生成 80k 个合成 4D 具身视频。虽然模拟器提供了度量深度信息，但缺乏表面法线数据，我们使用DSINE中的depth2normal 函数估计法线，并采用场景随机化技术增强泛化性。合成数据的多样性有限，与现实场景存在差距。纳入了现实世界视频数据集，利用 RollingDepth 为其标注仿射不变深度，使用 Temporal-Consistent Marigold-LCM-normal1 标注法线图。我们还选择了 OpenX 中的 Fractal data 和 Bridge 数据集，并纳入了人类 - 物体交互数据集 Something Something V2，以增加指令的多样性。
• 模型：使用 CogVideoX 中的 3D VAE 分别对 RGB、深度和法线视频进行编码，不额外微调 VAE。输入设计上，为每个模态引入三个单独的投影仪提取嵌入：$f_z=InputProj(z_t,z_0)$。DiT 以这些嵌入的和为输入，结合文本输入 T 和去噪步骤 t，得到隐藏状态 $h=DiT(\sum f_z,t,T)$。输出方面，保留原始 RGB 输出方法，同时为深度和法线预测引入额外模块，通过Conv3D层编码输入潜在和预测的RGB去噪输出的连接，与DiT骨干网络产生的隐藏状态结合，经过输出投影仪得到深度和法线的去噪预测。$$L={\mathbb{E}}_{{\mathbf{v}}_0,\mathcal{T},t,ϵ}\left[{\Vert [{ϵ}_{\mathbf{v}},{ϵ}_{\mathbf{d}},{ϵ}_{\mathbf{n}}]-{ϵ}_{\theta }({\mathbf{x}}_t,t,{\mathbf{x}}^0,\mathcal{T})\Vert }^2\right]$$
• 4D 场景重建：输入深度图 $D_i$ 是归一化的绝对值（[0,1]），无法直接用于 3D 重建；之前方法假设固定尺度或预测绝对深度，但重建结果粗糙；采用法线图 $N_i$ 通过法线积分优化深度，得到初步深度图 $\hat{D}$，并加空间一致性损失 $c_s$。单纯逐帧优化会导致动态场景的深度抖动，因此引入 光流（Optical Flow） 约束跨帧一致性：
  • 静态/动态区域分割：计算相邻帧光流 $F=RAFT(V)$；静态区域掩码 $M_i^s$——$||F_i||\le c$；动态区域掩码 $M_i^d$——剩余部分，$M_i^s$ 的剩余部分；背景区域掩码 $M_i^b$——动态掩码与静态掩码的交集；
  • 跨帧传播：$D_{i\to (i-1)}(u,v)=D_{i-1}(u-\Delta u,v-\Delta -v)$
  • 时间一致性损失：根据光流，可以从先前帧检索当前帧对应位置的深度，以施加一致性约束。强制动态区域和背景区域的深度与前一帧对齐：${\mathcal{L}}_c={\lambda }_{cd}\Vert {\tilde{D}}_i\circ M_i^d-D_{i\to (i-1)}\circ M_i^d{\Vert }^2+{\lambda }_{cb}\Vert {\tilde{D}}_i\circ M_i^b-D_{i\to (i-1)}\circ M_i^b{\Vert }^2$
  • 深度正则化损失：防止优化后的深度 ${\tilde{D}}_i$ 过度偏离初始预测 $D_i$：${\mathcal{L}}_r={\lambda }_{rd}\Vert {\tilde{D}}_i\circ M_i^d-D_i\circ M_i^d{\Vert }^2+{\lambda }_{rb}\Vert {\tilde{D}}_i\circ M_i^b-D_i\circ M_i^b{\Vert }^2$$$\arg \underset{\tilde{D}}{\min }\underset{\text{法线约束}}{\underbrace{{\mathcal{L}}_s(\tilde{D},N_i)}}+\underset{\text{时间约束}}{\underbrace{{\mathcal{L}}_c(\tilde{D},{\hat{D}}_{i-1},F_i,F_{i-1})}}+\underset{\text{正则化}}{\underbrace{{\mathcal{L}}_r(\tilde{D},D_i)}}$$

实验结果

凭借更好的深度和法线图，方法重建的4D点云在真实和合成数据集上均实现了最低的Chamfer距离。4D Point-E方法在RLBench上的表现优于视频扩散模型，但仍落后于我们的方法。此外，直接使用点云进行训练计算成本高昂，限制了所用的帧数。相比之下，我们的模型利用RGB-DN视频的高效表示生成更精确的4D场景，尤其在捕捉动态场景的细粒度细节方面表现出色。