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

TL;DR

FlowDreamer introduces a novel RGB-D world model for robot manipulation that explicitly uses 3D scene flow as motion representation, improving prediction accuracy and planning success in various benchmarks.

Contribution

It is the first to incorporate explicit 3D scene flow with diffusion models in RGB-D world modeling for robot manipulation.

Findings

Achieves 7% better semantic similarity

Improves pixel quality by 11%

Increases success rate by 6% in manipulation tasks

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.