Grounding Generated Videos in Feasible Plans via World Models

A video plan \(\tau_{\mathrm{vid}}\) is a sequence of visual observations generated by a conditional video generative model \(\mathcal{G}\) that provides temporally ordered visual foresight for completing a task. \(\mathcal{G}\) can be instantiated using an image-to-video (I2V) diffusion-based video model that produces temporally coherent visual transitions between the start and goal observations. We map the generated video plan into the latent state space of the action-conditioned world model using its underlying visual encoder, yielding a sequence of latent states \(z^{\mathrm{vid}}_{0:T}\).

In addition, we employ a semantic alignment loss between the optimized latent trajectory and the video plan. Specifically, the video alignment loss is defined as

\[ \mathcal{L}_{\mathrm{vid}}(z_t, z^{\mathrm{vid}}_t) = \left\| \phi(z_t) - \phi(z^{\mathrm{vid}}_t) \right\|^2, \]

where both optimized and video latent states are projected onto the unit \(\ell_2\) hypersphere by \(\phi(z) = z / \|z\|_2\). This loss penalizes angular deviation between latent embeddings while remaining invariant to their magnitude.

We formulate grounding video plans into feasible action sequences as a video-guided direct collocation problem. In direct collocation, both the latent states \(\mathcal{Z} = z_0, \ldots, z_T\) and the actions \(\mathcal{A} = a_0, \ldots, a_{T-1}\) are treated as decision variables. Our goal is to solve for a trajectory \((\mathcal{Z}^*, \mathcal{A}^*)\) that minimizes the divergence between the latent states of the optimized trajectory and the video plan \(z^{\mathrm{vid}}\), while satisfying the world model dynamics \(f_\psi\), leading to the following constrained optimization problem:

\[ \begin{aligned} \min_{\mathcal{Z}, \mathcal{A}} \quad & \lambda_{\text{v}} \sum_{t=1}^{T-1} \mathcal{L}_{\text{vid}}(z_t, z^{\text{vid}}_t) + \lambda_{\text{g}} \mathcal{L}_{\text{goal}}(z_T, z_g) + \lambda_{\text{r}} \sum_{t=0}^{T-1} \| a_t \|^2 \\ \text{s.t.} \quad & z_{t+1} = f_\psi(z_{t-H:t}, a_{t-H:t}), \quad \forall t \in \{0, \ldots, T-1\}, \\ & a_{\min} \preceq a_t \preceq a_{\max}, \quad \forall t \in \{0, \ldots, T-1\}. \end{aligned} \]

To efficiently solve the non-linear constrained optimization problem, we employ the Augmented Lagrangian Method (ALM). The loss is defined as:

\[ \mathcal{L}_\rho(\mathcal{Z}, \mathcal{A}, \Lambda) = \tilde{C}(\mathcal{Z}, \mathcal{A}) + \sum_{t=0}^{T-1} \left(\lambda_t^\top \mathcal{L}_{\mathrm{dyn}}^{t} + \frac{\rho}{2} \| \mathcal{L}_{\mathrm{dyn}}^{t} \|^2 \right), \]

where \(\tilde{C}\) denotes the latent cost objective. \(\Lambda = \{\lambda_0, \ldots, \lambda_{T-1}\}\) are the Lagrange multipliers, and \(\rho > 0\) is a scalar penalty parameter. The term \(\mathcal{L}_{\mathrm{dyn}}^{t}\) corresponds to the dynamics constraint violation at timestep \(t\), defined as:

\[ \mathcal{L}_{\mathrm{dyn}}^{t}(\mathcal{Z}, \mathcal{A}; f_\psi) = z_{t+1} - f_\psi(z_{t-H:t}, a_{t-H:t}). \]

Optimization proceeds using a standard primal–dual approach, alternating between gradient-based updates of the primal variables \((\mathcal{Z}, \mathcal{A})\) and the dual variables \(\Lambda\). In particular, we perform \(I_{\mathrm{ALM}}\) inner (primal) iterations per outer ALM step, updating the Lagrange multipliers and increasing the penalty parameter after each outer iteration.