GENFLOWRL: Generative Object-Centric Flow Matching for Reward Shaping in Visual Reinforcement Learning

Kelin Yu, Sheng Zhang, Harshit Soora, Furong Huang, Heng Huang, Pratap Tokekar, Ruohan Gao

Keywords: Object-centric Representation, Reinforcement Learning, Diffusion Model

TL;DR: The framework "GenFlowRL" utilizes the foundation generative model “Stable Diffusion” for generating object-centric flow and uses delta-flow matching to shape dense reward in reinforcement learning.

Abstract: Recent advances have demonstrated the potential of video generation models to guide robot learning by deriving effective robot actions through inverse dynamics. However, these methods heavily depend on the quality of generated data and struggle with fine-grained manipulation due to the lack of environment feedback. While video-based reinforcement learning improves policy robustness, it remains constrained by the artifacts of generated video and the challenges of collecting large-scale in-domain robot datasets for training diffusion models. Motivated by the above, we propose GenFlowRL, which derives shaped rewards from generated flow trained from cross-embodiment datasets. This enables learning generalizable and robust policies from expert demostrations using low-dimensional, object-centric features. Experiments on 10 manipulation tasks, both in simulation and real-world cross-embodiment evaluations, demonstrate that GenFlowRL effectively leverages manipulation features extracted from generated object-centric flow, consistently achieving superior performance across diverse and challenging scenarios.

Submission Number: 24