PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jun Guo, Xiaojian Ma, Yikai Wang, Min Yang, Huaping Liu, Qing Li

arXiv:2505.10075v1cs.ROcs.CV

TL;DR

Existing RGB-D world models commonly combine dynamics prediction and visual rendering in one model, motivating explicit dynamics modeling for more transparent and accurate prediction. FlowDreamer predicts 3D scene flow from RGB-D observations and robot actions, then conditions a diffusion model to generate the next visual observation, with end-to-end training. FlowDreamer achieves a 7% increase in semantic similarity and 11% on pixel quality on action-conditioned benchmarks, plus a 6% increase in manipulation success rate on VP2 evaluations.

  • Problem

    Existing RGB-D world models commonly combine dynamics prediction and visual rendering in one model, motivating explicit dynamics modeling for more transparent and accurate prediction.

    Source: §1 · pp. 1–2

  • Method

    FlowDreamer predicts 3D scene flow from RGB-D observations and robot actions, then conditions a diffusion model to generate the next visual observation, with end-to-end training.

    Source: §1 · p. 2, §3 · p. 3, §3.2 · p. 4

  • Results

    FlowDreamer achieves a 7% increase in semantic similarity and 11% on pixel quality on action-conditioned benchmarks, plus a 6% increase in manipulation success rate on VP2 evaluations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Explicit flow-based dynamics modeling improves FlowDreamer’s reported visual prediction and visual planning performance across robot manipulation domains.

    Source: §1 · p. 2, §5 · p. 8

    In Robosuite, a single RGB-D image cannot reflect post-push object velocity, and unattached objects’ scene flows can mislead prediction.

    Source: §4.2 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

FlowDreamer addresses RGB-D world modeling for robot manipulation by explicitly separating dynamics prediction from visual rendering. It uses 3D scene flow within a two-stage, end-to-end framework and reports improvements across visual prediction and planning benchmarks.

Motivation

  • Visual world models serve as learnable simulators that can help produce and evaluate robot-manipulation plans without requiring precise scene modeling and simulation.They predict future motion and render visual observations conditioned on current sensory observations and robot actions.Source: §1 · p. 1
  • Existing RGB-D world models generally combine dynamics prediction and visual rendering in one model, reducing transparency and potentially weakening future prediction.The paper hypothesizes that frame-only prediction losses can favor visual appearance over accurate dynamics modeling.Source: §1 · pp. 1–2

Proposed approach

  • FlowDreamer explicitly models dynamics with 3D scene flow, separating environment-dynamics prediction from visual rendering in a two-stage framework.A scene-flow module predicts action-induced changes, while a conditional diffusion model generates the next visual observation; both stages are trained end to end.Source: §1 · p. 2

Empirical results

  • 7% higher semantic similarity, 11% higher pixel quality, and 6% higher manipulation success rate are reported across action-conditioned and visual-planning benchmarks.The evaluations cover multiple robotic manipulation benchmarks, including RT-1, Language Table, VP2, RoboDesk, and Robosuite.Source: §1 · p. 2

2. Related Work

Robotic world models predict future states from observations and conditions, supporting planning and control, while existing approaches often omit explicit dynamical-process modeling. FlowDreamer contributes a control-conditioned RGB-D design that uses 3D scene flow as an intermediate motion representation and diffusion for future observation generation.

2.1. Generative World Models in Robotics

  • World models predict future states from current observations and conditions, supporting robot planning, learning, and reduced real-environment interaction costs.Control-based models use robot signals as conditions, whereas instruction-based models use task identifiers or natural-language descriptions.Source: §2.1 · p. 2
  • FlowDreamer is a control-based world model that conditions on robot actions, uses scene flow to simulate environmental dynamics, and generates human-interpretable images.Its two-stage design predicts scene flow before producing the next observation through diffusion.Source: §2.1 · p. 2, Figure 2 · p. 3

2.2. Dynamics Modeling

  • Most existing world models directly predict future states without explicitly modeling dynamics, although explicit dynamics can improve interpretability and prediction accuracy.Prior motion representations include SE(3) transformations, which are limited to rigid bodies and require object-centric awareness with additional segmentation.Source: §2.2 · p. 2
  • FlowDreamer’s two stages first predict 3D scene flow from RGB-D observations and robot actions, then generate the next observation with a denoising U-Net and diffusion.The framework uses the predicted flow as motion information for high-resolution future-observation generation.Source: Figure 2 · p. 3, §2.2 · p. 3
  • Flow is a flexible motion representation defined as each observation-space point’s displacement between timesteps, supporting non-rigid motion.In RGB-D space, 3D scene flow can be obtained by integrating optical flow from consecutive frames with corresponding depth information.Source: §2.2 · p. 3

3. FlowDreamer: A RGB-D World Model

FlowDreamer separates explicit 3D dynamics prediction from RGB future-frame generation in a two-stage RGB-D world model. It predicts scene flow from RGB-D observations and actions, then uses diffusion conditioned on that flow to generate future observations.

3. FlowDreamer: A RGB-D World Model

  • FlowDreamer receives the current RGB-D observation and robot action, predicts 3D scene flow, and generates the next RGB observation through a diffusion denoising network.The two stages are trained jointly despite the modularized design.Source: §3 · p. 3, §3.3 · p. 4

3.1. Background: Latent Diffusion Models

  • Latent diffusion progressively denoises a noisy latent representation, starting from Gaussian noise and iteratively recovering the clean latent.The model can accelerate inference by setting σ_k = 0 and using a subsequence of diffusion steps.Source: §3.1 · pp. 3–4

3.2. Dynamics Prediction

  • RGB-D pixels are projected into 3D using camera intrinsics and depth, while scene flow represents the displacement of corresponding 3D points across consecutive frames.Simulation provides flow from object transformations, whereas real-world data can use RAFT-3D or estimated depth.Source: §3.2 · p. 4
  • The dynamics module predicts scene flow from the current RGB-D observation and action using a conditional U-Net with action features integrated through cross-attention.Its flow objective is the mean squared error between predicted and estimated scene flow.Source: §3.2 · p. 4

3.3. Future Generation

  • The future-generation module fine-tunes Stable Diffusion to predict the next RGB observation from the current RGB-D input, action, and predicted scene flow.RGB observations are encoded into latent space, while depth and flow are downsampled and concatenated with the latent representation.Source: §3.3 · p. 4
  • The model does not directly predict the next depth map; instead, a pretrained depth estimator with a DPT head derives metric depth from the generated RGB frame.This choice reflects the diffusion decoder’s lack of specialization for metric depths.Source: §3.3 · p. 4
  • End-to-end training combines diffusion and flow objectives as L_total = L_diff + αL_flow, with α controlling their relative weight.The joint objective trains the dynamics prediction module and denoising U-Net together.Source: §3.3 · p. 4

4. Experiments

Experiments evaluate FlowDreamer across video prediction, visual planning, and flow-analysis settings. The results show strong video-prediction performance, improved planning relative to Vanilla, and evidence that predicted scene flow supports future-frame generation.

4.1. Video Prediction

  • FlowDreamer achieves the best performance on most RT-1 and Language Table metrics, including semantic similarity, pixel similarity, and media quality.The evaluation uses PSNR, SSIM, LPIPS, FID, FVD, DINOv2 L2, and CLIP score; SepTrain is generally second-best and close to FlowDreamer.Source: §4.1 · pp. 5–6, Table 1 · p. 5, Table 2 · p. 5

4.2. Visual Planning

  • FlowDreamer always outperforms Vanilla on VP2 visual planning and achieves comparable average performance to video-prediction baselines using two-frame context.It performs best on some tasks, including Green button and Upright block, but fails on Open slide.Source: §4.2 · pp. 7–8, Figure 4 · p. 7
  • Robosuite push tasks are harder because objects retain velocity after being pushed, while scene flow from unattached objects can mislead prediction from a single RGB-D image.The qualitative results show accurate flow and goal prediction on RoboDesk but less precise flow on Robosuite push tasks.Source: Figure 5 · p. 7, §4.2 · p. 8

4.3. Additional Analysis on Flow Prediction

  • Reversing the predicted flow worsens future-frame prediction, indicating that stage-one scene flow provides auxiliary information beyond the robot action.With incorrect reversed flow, the diffusion model does not produce contrary actions but cannot reach the goal state effectively.Source: §4.3 · p. 8, Figure 6 · p. 8
  • Flow error correlates highly with PSNR and SSIM but only weakly with DINOv2 L2 and CLIP score.The authors infer that semantic metrics do not reliably distinguish relatively minor prediction errors from the ground truth.Source: §4.3 · p. 8, Table 3 · p. 8

5. Conclusion

FlowDreamer is an action-conditioned RGB-D world model that uses flow-based motion representations and separates dynamics learning from visual rendering. Across four benchmarks spanning video prediction and visual planning, it outperforms other RGB-D world models.

5. Conclusion

  • FlowDreamer uses 3D scene flow as a motion representation and a separate dynamics prediction module to learn environment dynamics.Despite its modularized design, the model is jointly trained end-to-end.Source: §5 · p. 8
  • Experiments on four benchmarks covering video prediction and visual planning demonstrate FlowDreamer’s superiority over other RGB-D world models.Source: §5 · p. 8

A. Implementation Details

The implementation uses conditional networks for action-conditioned dynamics prediction and a diffusion-based rendering stack trained with fixed encoder and decoder weights. Training uses distributed mixed-precision optimization and 20-step PNDM inference.

A. Implementation Details

  • A conditional 2D U-Net with four downsampling and upsampling layers serves as FlowDreamer’s dynamics prediction module.Action conditions are integrated with image features through cross-attention.Source: A. Implementation Details · p. 13
  • A conditional MinkowskiNet baseline replaces batch normalization with FiLM layers to introduce action conditioning in a point-cloud model.Source: A. Implementation Details · p. 13
  • The denoising U-Net and image encoder/decoder use Stable Diffusion 2.1 components, while only the denoising network is fine-tuned.Source: A. Implementation Details · p. 13
  • Models train for 60k steps with AdamW, learning rate 1e-4, mixed precision, and batch size 16 per GPU across eight GPUs.Source: A. Implementation Details · p. 13
  • Inference uses a PNDM scheduler with 20 sampling steps because additional steps do not improve future-prediction accuracy.Source: A. Implementation Details · p. 13

B. Data Collection

The data pipeline combines simulator trajectories and estimated real-world scene flow to train and evaluate FlowDreamer across robotic manipulation environments. Simulation provides ground-truth flow, whereas RT-1 uses depth and flow estimators to construct training targets.

B. Data Collection

  • The dataset collection covers four simulation environments and the real-world RT-1 dataset, with trajectory and sample statistics summarized in Table 4.An episode is a complete task trajectory, while a sample is a consecutive frame pair with its robot action.Source: B. Data Collection · p. 13, Table 4 · p. 13
  • Simulation trajectories come from CogACT, an RRT* oracle, or VP2-provided trajectories, with failed SimplerEnv trajectories discarded.RoboDesk observations are resized and cropped for planning rewards, while training and prediction use the original observations.Source: B. Data Collection · p. 13
  • For simulated rigid and articulated objects, ground-truth scene flow is computed from point coordinates transformed by consecutive 4 × 4 pose matrices.The flow is the difference between the transformed positions at consecutive timesteps.Source: B. Data Collection · p. 13
  • On RT-1, Video Depth Anything provides temporally coherent relative depth, Depth Anything v2 provides metric depth, and RAFT-3D estimates frame-paired 3D scene flow.Depth maps are aligned using scale and shift parameters, and RAFT supplies a flow mask for refinement.Source: Figure 8 · p. 13, B. Data Collection · pp. 13–14
  • Table 5 reports RT-1 video-prediction metrics grouped into semantic similarity, pixel similarity, and media quality.Source: Table 5 · p. 14

C. Extended Experimental Results

On the harder real-world RT-1 dataset, FlowDreamer remains better than the single-stage Vanilla model, although the performance gap narrows relative to simulation. Estimated scene flow introduces noisy supervision that degrades flow prediction, yet FlowDreamer still preserves background consistency better than Vanilla.

C. Extended Experimental Results

D. Limitations and Future Works

FlowDreamer’s current limitations concern training-data coverage, limited temporal context, and relatively slow diffusion inference. Future work is suggested to improve these boundaries while retaining explicit dynamics modeling.

Related papers

Loading 2505.10075v1…