Cold Diffusion on the Replay Buffer: Learning to Plan from Known Good States

Zidan Wang, Takeru Oba, Takuma Yoneda, Rui Shen, Matthew Walter, Bradly C. Stadie

Keywords: Imitation Learning, Reinforcement Learning, Diffusion, Cold Diffusion, Planning, Safety

TL;DR: By using cold diffusion over a Replay Buffer, we can learn to imitate while avoiding infeasible states.

Abstract: Learning from demonstrations (LfD) has successfully trained robots to exhibit remarkable generalization capabilities. However, many powerful imitation techniques do not prioritize the feasibility of the robot behaviors they generate. In this work, we explore the feasibility of plans produced by LfD. As in prior work, we employ a temporal diffusion model with fixed start and goal states to facilitate imitation through in-painting. Unlike previous studies, we apply cold diffusion to ensure the optimization process is directed through the agent's replay buffer of previously visited states. This routing approach increases the likelihood that the final trajectories will predominantly occupy the feasible region of the robot's state space. We test this method in simulated robotic environments with obstacles and observe a significant improvement in the agent's ability to avoid these obstacles during planning.

Student First Author: yes

Instructions: I have read the instructions for authors (https://corl2023.org/instructions-for-authors/)

Code: https://github.com/zidanwang2025/cold_diffusion_on_replay_buffer