Latent Dynamics Models

Latent Dynamics Model は、high-dimensional な observation を latent state に圧縮し、その latent space 上で dynamics を学習する world model です。Pixel 空間で future を直接予測するよりも、planning や control に必要な information を compact に扱えます。

なぜ latent なのか

Image や video は高次元です。Pixel を直接予測すると、背景 texture や照明変化のような task に不要な detail に capacity を使ってしまいます。Latent dynamics では、observation を task-relevant な state に圧縮します。

$$z_t = e_\phi(o_t)$$

Dynamics model は、現在 state と action から次の state を予測します。

$$z_{t+1} \sim p_\theta(z_{t+1} \mid z_t, a_t)$$

Recurrent State-Space Model

Dreamer 系では、recurrent state-space model (RSSM) がよく使われます。Deterministic hidden state と stochastic latent state を組み合わせます。

この構成により、uncertainty と temporal memory の両方を扱えます。

Training objective

Latent dynamics model は、主に次を同時に学習します。

• Observation reconstruction
• Reward prediction
• Continuation / termination prediction
• Latent consistency
• Policy / value learning

VAE 的な objective を使い、posterior と prior の KL term を含むことが多いです。

Planning

Latent dynamics model の利点は、model 内で future を想像できることです。

Pixel を実際に decode しなくても、latent state と reward を予測すれば planning ができます。

数式で見る latent dynamics の学習

Latent dynamics model は、高次元 observation $o_t$ を latent $\mathbf{z}_t$ に圧縮し、その latent の時間発展を学習します。

$$\mathbf{z}_t=E_\phi(o_t), \qquad \hat{\mathbf{z}}_{t+1}=F_\theta(\mathbf{z}_t,a_t)$$

Prediction loss は latent 空間で書けます。

$$\mathcal{L}_{dyn}=\sum_t\left\|F_\theta(\mathbf{z}_t,a_t)-\mathrm{sg}(E_\phi(o_{t+1}))\right\|^2$$

ここで、$\mathrm{sg}$ は stop-gradient です。この式の気持ちは、「次の画像そのものを pixel で予測するのではなく、制御に必要な抽象表現の次状態を予測する」というものです。Pixel prediction は細部に引っ張られやすい一方、latent prediction は planning に必要な dynamics を学びやすくなります。

関連ページ

• Dreamer
• World Models Overview
• JEPA Overview

主なソース

• World Models: https://worldmodels.github.io/
• Dreamer paper: https://arxiv.org/abs/1912.01603
• DreamerV3: https://arxiv.org/abs/2301.04104