FlowDreamer: A RGB-D World Model with Flow-based Motion Representations for Robot Manipulation

TL;DR

FlowDreamer addresses the challenge of predicting future RGB-D observations in robot manipulation by decoupling dynamics prediction from visual rendering and introducing 3D scene flow as an explicit motion representation. It proposes a two-stage architecture where Stage 1 explicitly predicts 3D scene flow $\hat{f}_{t\rightarrow t+1}$ from $(I_t, D_t, a_t)$, and Stage 2 uses a latent diffusion model conditioned on the current observation, depth, action, and predicted flow to synthesize $I_{t+1}$, with end-to-end training via $\mathcal{L}_{\text{diff}}$ and $\mathcal{L}_{\text{flow}}$. The approach yields improvements in semantic similarity, pixel fidelity, and task success across RT-1, Language Table, RoboDesk, and Robosuite benchmarks, illustrating the value of explicit dynamics supervision for both video prediction and visual planning. By leveraging 3D scene flow and latent diffusion, FlowDreamer provides interpretable dynamics and high-fidelity future frames, with potential implications for more reliable real-world robotic perception and control, albeit with challenges in real-world data collection and diffusion inference speed.

Abstract

This paper investigates training better visual world models for robot manipulation, i.e., models that can predict future visual observations by conditioning on past frames and robot actions. Specifically, we consider world models that operate on RGB-D frames (RGB-D world models). As opposed to canonical approaches that handle dynamics prediction mostly implicitly and reconcile it with visual rendering in a single model, we introduce FlowDreamer, which adopts 3D scene flow as explicit motion representations. FlowDreamer first predicts 3D scene flow from past frame and action conditions with a U-Net, and then a diffusion model will predict the future frame utilizing the scene flow. FlowDreamer is trained end-to-end despite its modularized nature. We conduct experiments on 4 different benchmarks, covering both video prediction and visual planning tasks. The results demonstrate that FlowDreamer achieves better performance compared to other baseline RGB-D world models by 7% on semantic similarity, 11% on pixel quality, and 6% on success rate in various robot manipulation domains.

Figures (9)

• Figure 1: Proposed RGB-D world model with flow-based motion representations. FlowDreamer adopts a two-stage prediction framework, which explicitly predict scene flow as motion representations. FlowDreamer achieves better results on future frame prediction and visual planning tasks in various robot manipulation domains.
• Figure 2: Overview of FlowDreamer. At stage 1, FlowDreamer receives the RGB-D frame and the robot action as input to explicitly predict the scene flow as motion representations. At stage 2, FlowDreamer leverages a denoising U-Net to generate high-resolution next-step future observation via diffusion.
• Figure 3: Qualitative results on the SimplerEnv RT-1 and Language Table benchmark. We show the predicted frames and the scene flows except for Vanilla, where only RGB frames are being predicted. The R, G, and B channel values in the flow visualization represent the components of the 3D scene flow along the $x$, $y$, and $z$ directions, respectively, normalized by the maximum value of the scene flow.
• Figure 4: Visual planning results on the VP$^2$ benchmark. We report the mean and the min/max performance of different methods over multiple runs with different random seeds. On the right, "Average" means the average success rate over all reported tasks.
• Figure 5: Qualitative results on the Robodesk and Robosuite dataset. The trajectory comes from the validation set, which is split from the original training trajectories and is not used for training. For our method, we show the predicted RGB images and scene flows.