Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Contrastive Diffuser: Planning Towards High Return States via Contrastive Learning

Yixiang Shan, Zhengbang Zhu, Ting Long, Qifan Liang, Yi Chang, Weinan Zhang, Liang Yin

TL;DR

CDiffuser tackles offline reinforcement learning when high-return trajectories are scarce by leveraging abundant low-return trajectories through a contrastive mechanism. It combines a diffusion-based Planning Module with a Contrastive Module that uses high-return states as positives and low-return states as negatives to steer generated trajectories toward favorable regions. The approach introduces SR and SRD sampling strategies and a plan-specific contrastive loss to enforce state-level discrimination, jointly optimizing trajectory generation and a return predictor. Empirical results across 14 D4RL benchmarks show strong improvements over diffusion-based and non-diffusion baselines, with ablations validating the value of the contrastive component and its plan-execution consistency benefits.

Abstract

The performance of offline reinforcement learning (RL) is sensitive to the proportion of high-return trajectories in the offline dataset. However, in many simulation environments and real-world scenarios, there are large ratios of low-return trajectories rather than high-return trajectories, which makes learning an efficient policy challenging. In this paper, we propose a method called Contrastive Diffuser (CDiffuser) to make full use of low-return trajectories and improve the performance of offline RL algorithms. Specifically, CDiffuser groups the states of trajectories in the offline dataset into high-return states and low-return states and treats them as positive and negative samples correspondingly. Then, it designs a contrastive mechanism to pull the trajectory of an agent toward high-return states and push them away from low-return states. Through the contrast mechanism, trajectories with low returns can serve as negative examples for policy learning, guiding the agent to avoid areas associated with low returns and achieve better performance. Experiments on 14 commonly used D4RL benchmarks demonstrate the effectiveness of our proposed method. Our code is publicly available at \url{https://anonymous.4open.science/r/CDiffuser}.

Contrastive Diffuser: Planning Towards High Return States via Contrastive Learning

TL;DR

CDiffuser tackles offline reinforcement learning when high-return trajectories are scarce by leveraging abundant low-return trajectories through a contrastive mechanism. It combines a diffusion-based Planning Module with a Contrastive Module that uses high-return states as positives and low-return states as negatives to steer generated trajectories toward favorable regions. The approach introduces SR and SRD sampling strategies and a plan-specific contrastive loss to enforce state-level discrimination, jointly optimizing trajectory generation and a return predictor. Empirical results across 14 D4RL benchmarks show strong improvements over diffusion-based and non-diffusion baselines, with ablations validating the value of the contrastive component and its plan-execution consistency benefits.

Abstract

The performance of offline reinforcement learning (RL) is sensitive to the proportion of high-return trajectories in the offline dataset. However, in many simulation environments and real-world scenarios, there are large ratios of low-return trajectories rather than high-return trajectories, which makes learning an efficient policy challenging. In this paper, we propose a method called Contrastive Diffuser (CDiffuser) to make full use of low-return trajectories and improve the performance of offline RL algorithms. Specifically, CDiffuser groups the states of trajectories in the offline dataset into high-return states and low-return states and treats them as positive and negative samples correspondingly. Then, it designs a contrastive mechanism to pull the trajectory of an agent toward high-return states and push them away from low-return states. Through the contrast mechanism, trajectories with low returns can serve as negative examples for policy learning, guiding the agent to avoid areas associated with low returns and achieve better performance. Experiments on 14 commonly used D4RL benchmarks demonstrate the effectiveness of our proposed method. Our code is publicly available at \url{https://anonymous.4open.science/r/CDiffuser}.
Paper Structure (32 sections, 19 equations, 10 figures, 4 tables, 2 algorithms)

This paper contains 32 sections, 19 equations, 10 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Denoising Probabilistic Models
  4. Contrastive Learning
  5. Offline RL Problem Definition
  6. Methodology
  7. Planning Module
  8. Contrastive Module
  9. Sample Positive and Negative States
  10. Constrain the trajectory with contrastive learning
  11. Model Learning
  12. Experiments
  13. Experiment Settings
  14. Benchmark Results
  15. Study on the limited number of high-return trajectories.
  16. ...and 17 more sections

Figures (10)

  • Figure 1: (a) The probability density of trajectories’ returns in Maze2d; (b) The learned trajectory when high-return trajectories are abundant; (c) The learned trajectory when the number of high-return trajectories is limited; (d) The contrastive learning applied by previous RL models; (e) The example of our solution.
  • Figure 2: The overall framework of CDiffuser. CDiffuser is composed of two modules: the Planning Module and the Contrastive Module. The Planning Module is designed to generate the subsequent trajectories, and the Contrastive Module is designed to pull the states in the generated trajectories toward the high-return states and push them away from the low-return states during the training phase.
  • Figure 3: Results of the ablation experiments on different variants.
  • Figure 4: The distribution of state and reward. It is better to view in color mode. CDiffuser achieves higher rewards in out-of-distribution areas (circled with red).
  • Figure 5: Distribution of rewards on (a) Walker2d-Med-Replay and (b) Halfcheetah-Med-Replay.
  • ...and 5 more figures