Flow to Better: Offline Preference-based Reinforcement Learning via Preferred Trajectory Generation
Authors: Zhilong Zhang
Yihao Sun
Junyin Ye
Tian-Shuo Liu
Jiaji Zhang
Yang Yu
Venue: ICLR 2024 poster (2024)
Keywords: Preference-based Reinforcement Learning
Offline Reinforcement Learning
Conditional Generative Modeling
Diffusion Models
Primary area: reinforcement learning

Abstract: Offline preference-based reinforcement learning (PbRL) offers an effective solution to overcome the challenges associated with designing rewards and the high costs of online interactions. In offline PbRL, agents are provided with a fixed dataset containing human preferences between pairs of trajectories. Previous studies mainly focus on recovering the rewards from the preferences, followed by policy optimization with an off-the-shelf offline RL algorithm. However, given that preference label in PbRL is inherently trajectory-based, accurately learning transition-wise rewards from such label can be challenging, potentially leading to misguidance during subsequent offline RL training. To address this issue, we introduce our method named $\textit{Flow-to-Better (FTB)}$, which leverages the pairwise preference relationship to guide a generative model in producing preferred trajectories, avoiding Temporal Difference (TD) learning with inaccurate rewards. Conditioning on a low-preference trajectory, $\textit{FTB}$ uses a diffusion model to generate a better one with a higher preference, achieving high-fidelity full-horizon trajectory improvement. During diffusion training, we propose a technique called $\textit{Preference Augmentation}$ to alleviate the problem of insufficient preference data. As a result, we surprisingly find that the model-generated trajectories not only exhibit increased preference and consistency with the real transition but also introduce elements of $\textit{novelty}$ and $\textit{diversity}$, from which we can derive a desirable policy through imitation learning. Experimental results on D4RL benchmarks demonstrate that FTB achieves a remarkable improvement compared to state-of-the-art offline PbRL methods. Furthermore, we show that FTB can also serve as an effective data augmentation method for offline RL.

Published as a conference paper at ICLR 2024
FLOW TO BETTER: OFFLINE PREFERENCE-BASED RE- INFORCEMENT LEARNING VIA PREFERRED TRAJEC- TORY GENERATION
Zhilong Zhang-.2; Yihao Sun'; Junyin Yc!.2, Tian-Shuo Liu Jiaji Zhang', Yang Yu
1 National Key Laboratory for Novel Software Technology, Nanjing University, China &School of Artificial Intelligence, Nanjing University, China
Pollxxrr Technologies
{zhangzl, sunyh, yejy, liuts, zhangii} ( lamda nju edu.cn yuyênju. edu. on
ABSTRACT
Offline preference-based reinforcement learning (PbRL.) offers an effective solu- tion overcome the challenges associated with designing rewards and the high costs of online interactions. Previous studies mainly focus on recovering rewards from preferences, followed by policy optimization with an RL algorithm. However, given that preference labels in PbRL are inherently trajectory-based accurately learning transition-wise rewards from such labels can be challenging, potentially leading to misguidance during subsequent offline RL training. To address this issue, we introduce our method named Flow-to-Beite (FTB), which leverages the pairwise preference relationship to guide a generative model in producing preferred trajectories, avoiding Temporal Difference (TD) learning with inaccurate rewards. Conditioning FTB uses a diffusion model to generate a better one, achieving high-fidelity full- horizon trajectory improvement. During diffusion training, nique called Preference Augmentation to alleviate the problem of insufficient pref- erence data. As result, we surprisingly find that the model geneeraed trajectories not only exhibit increased preference and consistency with the real transition but also introduce elements of novelty and diversity, from which we can derive a de- sirable policy through imitation learning. Experimental results on several bench- marks demonstrate that FTB achieves remarkable improvement compared state-of-the offline PbRL methods. Furthermore, we show that FTB can also serve as an effective data augmentation method for offline RL. Our project's ss web- site can be found at https://iihubb.oo//Z/5///// tottetttt.............
offline
low- prw-prefeeence trajectory,
propose tech-
INTRODUCTION
In reinforcement learning (RL), agents interact with an environment and receive feedback to learn a policy that maximizes cumulative rewards. This paradigm has demonstrated efficacy in many domains, including games (Silver et al., 2017; Mnih et al., 2015), robotics (Levine et al., 201 8), and large language model (Ouyang et al., 2022). However, crafting well-designed rewards that align with the task's objectives or human intentions presents a formidable challenge. First, we aim for the reward function to densely informative, like transition-we wise, to facilitate task learning. Secondly, the reward function must exhibit robustness to prevent policy exploitation loopholes that could result in unreasonably high rewards. Therefore, reward engineering necessitates a substantial foundation of prior knowledge and rigorous testing by human experts, making it extremely complex and even unfeasible in some cases.
Preference-based reinforcement learning (PbRL) tackles the challenge of designing reward func- tions by utilizing trajectory preferences (Akrour et al., 2011; 2012; Christiano et al., 2017). In this
"Equal contribution.
Corresponding author.

---

Published as a conference paper at ICLR 2024
Diffusion Model
X low-preference trijectory
high-preference trajectory
flow better
Figure I:
Illustration of the key idea of our method. Given low-preference trajectory (left), the
FTB model generates a higher-preference trajectory (right).
framework, agents are provided with preferences between pairs of trajectory segments, which are relative judgments labeled by humans. However, existing approaches still necessitate a substantial number of queries to human experts, making it difficult to scale online preference- baseee RL to var- ious applications. For instance, even in simple video games or control tasks, existing methods still need millions of requests (Lee et al., 2021). Therefore, it becomes imperative explore PbRL in an offline setting. where agents have access to a fixed offline dataset with preference labels, without concerns about safety and sample efficiency.
In offline PbRL, existing approaches (Shin & Brown, 2021; Kim et al., 2023) typically involve two steps: reward learning and offline RI. Specifically, agents employ the Bradley-Tery model (Bradley & Terry. 1952) supervised manner to learn a transition-wise reward function, followed by an off-the-shelf offline RL algorithm for policy optimization. However, preference labels in PbRL trajectory-wise posing challenges in recovering precise transition-wise rewards, which will signifi- cantly impact offline RL training and lead undesirable policy performance.
To address this challenge, we aim to optimize policy behaviors directly at the trajectory level, as opposed decomposing preferences into rewards and performing TD learning with these inaccu- rate rewards. Our innovative method, Flow-to-Better (FTB), leverages the preference relationship to learn generative model capable of improving a trajectory, thereby generating more high- preferreee improved full-horizon trajectory condi- tioned on the less preferred one, as illustrated in Figure 1. Additionally, we introduce Preference Augmentation to generate sufficient preference pairs for training the diffusion model. After diffusion model training. we apply conditional generation process in an autoregressive manner to the original trajectories iteratively, effectively "flowing" Now-preference trajectories to high-preferenee trajecto- ries. With substantial high-pree gegh-rrfference trajectories generated by our model, we can extract desirable policy through simple imitation learning. Overall, this approach introduces fline PbRL to circumvent the pitfalls associated with inaccurate reward models.
trajectories. We employ a diffusion model to generate
paradigm for of-
We highlight the main contributions of work below:
We present novel framework for offline PbRL, i.e., Flow-to-Beeterr, which uses trajectory diffuser to achieve trajectory optimization, avoiding TD learning with inaccurate rewards.
We introduce Preference Augmentution, an innovative technique designed alleviate issue of insufficient preference labels in our approach.
We demonstrate that FTB consistently outperforms previous offline PbRL methods across various complex locomotion and manipulation tasks.
Our results show that the proposed trajectory diffuser our method can also be used as an effective data augmentation method for existing offline RI. approaches.
2 BACKGROUND
2.1 REINFORCEMENT LEARNING
We consider environments represented as a finite Markov decision process (MDP) (Sutton & Barto, 2018; Puterman, 2014) 4), which is described by tuple AS, P, p. Hf), where H is the length of an episode, S and A are the state and action space respectively, P = [PA(- a) (s, E A and P represent the dynamics and the initial state distribution, function. policy Fa S + A a mapping from states to actions. The decision process runs as follows: initial state 81 is drawn from p the beginning, for each time step ht the agent observes
(r1, ... TH) specifics the reward
2

---

Published as a conference paper at ICLR 2024
state Sh, selects an action On from the distribution [=)))))),, and subsequently, the environment provides a reward ri an) to the agent and transits to new state $1+11 according to Pr as). The goal of the agent is to maximize the expected cumulative reward E E Th (sh 0z)].
Offiine preference-based reinforcement learning (Offline PbRL). In contrast to online RL, offline PbRL (Shin et al., 2022; Kim et al., 2023) prohibits the agent from interacting with an environment, and there is no access to the reward function. Instead, the agent provided with two sets of offline datasets. The first dataset, denoted as DL = [[Tm> Tm Mm } m== comprises a limited number of trajectory pairs, each associated with a preference label ym E {0, 13,, where 0 implies that To > T 1 implies TH > Th and 0.5 implies T T The second dataset, referred to as De {Tn N 1e=1, comprises large number of unlabeled trajectories. Most previous methods try to
leamn reward function from DL and label transitions in D for offline RL training.
2.2 DIFFUSION MODELS
Diffusion Models (DMs) are sort of generative models that can model complex distributions and have achieved significant success in text to image (Dhariwal & Nichol, 2021; Rombach et al., 2022). The fundamental concept behind DMs is to progressively refine noise-perturbed input to generate samples that closely resemble the target distribution. The forward process of DMs g(xx,(2;-1) N(z;; V (1 a1)a gradually diffuses from the target distribution to Gaus- sian distribution in T timesteps, while the reverse process pe(x;-1/o = (( Me (x t), Σx) starts from Gaussian distribution and iteratively denoises samples using trained model, ulti- mately recovering the target distribution. The training objective of DMs is to minimize the negative log-likelihood Eazgl log po (z))] which can be simplified using the variational lower bound in Equation 1.
LVLB [1.T (0.1) |||e c 1)
Guidance Diffusion Models. Guidance DMs model the conditional distribution q(xly). This pro- cess can be understood through score function V log ((c:/y) = V log p(xx)) V log p(y (xx) (Song et al., 2020). Intuitively, we need to simultaneously model both the unconditional score log p(x,) and classifier guidance score V log p(y a.2). Currently, methods for guidance DMs primarily fall into two categories: classifier-guidance (Dhariwal & Nichol, 2021) and classifier-free (Nichol et al., 2022). The former entails training an additional classifier guidance estimator, denoted as fy(lzz)) alongside the DMs. During generation, it samples perturbed noise Ez,,, t) == c, t) wv1 - ā V log ( xz)) On the other hand, the latter approach involves training both conditional models (x,,, y. and unconditional models cœ t) throughout the DM training process. During generation, it samples perturbed noise ep (I,, t) := 6g (x,, t) + (e) y, t) - Eg (Z,,,)) t)). Here, w serves hyperparameter to adjust the degree of conditional guidance.
METHOD
This section introduces our Flow-to-Better (FTB) method, which optimizes policies at the trajec- tory level without TD learning under inaccurate learned rewards. We treat offline preference-bes bbased reinforcement learning as conditional generation task, where a diffusion model generates higher- preference trajectories based on inferior trajectories (Section 3.1). Given that a significant portion of the offline dataset lacks labels, we introduce a method called Preference Augmentation, which can provide more preference pairs for diffusion training. (Section 3.2). Furthermore, we explain how to derive a deployable policy (Section 3.3) and provide a full procedure of our method (Section 3.4).
3.1 OFFLINE PREFERENCE-BASED RL AS A CONDITIONAL GENERATION TASK
In offline PbRL, have pairs of trajectories with preference labels, enabling agents to learn from relative talgmentss-pogreedd
from an inferior trajectory (low preference) to superior one (high preference). To accomplish this, formulate the problem depicted in Equation 2. In this task, our goal is to maximize the log-likelihood of the conditional distribution of generating improved trajectories. Here, the condition is the less preferred trajectory in human judgment.
conditional generation task, as
min B(rigr")-D |- log po (o (r)))
(2)
3

---

Published as a conference paper at ICLR 2024
Atsentionllcc (tea(r))
Mulli-lead
Eg
Gaussian Noise
Conold
Conold
"Traffacttoyyy Diffuer U-Nett
Temporittlock
Conall
tv
-Atterice Boot
steps
Mah
Grouphorm
Timestep
Figure 2: The architecture of trajectory diffuser, which is classifier- free diffusion model.
Note that this generation task is highly challenging, as full-horizon trajectories typically have high dimensions. Furthermore, there exist multiple diverse and distinct trajectories that outperform the conditioned trajectory, leading to multi-modal target distribution. Therefore, we employ classifier-free diffusion model, Trajectory Diffuser, as the central implementation of our generative model. The training loss is defined as Equation 3:
C CTajDDifluser (0) (r. ,0 Bera(p) [lle co(T) (1
(1- B2, t)),
(3)
where is the diffusion timestep sampled from t ~ Uf1, .i......,, Th, To is the low- low-preference conditional trajectory, and of is the eigh-preference target trajectory the t diffusion timestep. Note that with probability p. we ignore the conditional information, and Ø is dummy value taking the place of condition 7
In Figure 2, we illustrate the architecture of the trajectory diffuser in FTB, which includes several encoders. Notably, introduce trajectory encoder denoted as ep designed embed the refer- ence trajectory. This encoder employs one-dimensional convolutional network map trajectory E R (ds to representation Ç E "d,,,, where H represents the trajectory horizon. This encoder compresses the dimensions of the state (ds) and action (d.a) into d, channels. Subse- quently, Ç is passed through an attention block with a U-Net architecture, undergoing multi-head cross-attention. This produces intermediate representations (p(z) E xd, in preceding layers. Following similar approach to (Ho et al., 2020), we utilize timestep encoder denoted as lo to generate representations Rd=, These derived representations & are integrated into the temporal block, serving as conditions alongside the intermediate representation (2(())
3.2 PREFERENCE AUGMENTATION
Although the method in Section 3.1 trains a generator that improves trajectorics, there are still the following issues in practice: 1) The aforementioned method requires preference labels for each pair of trajectories to conduct training, meaning that it can only use the dataset DL which each pair of trajectories has preference label while wasting the unlabeled dataset D 2) The performance differences between trajectory pairs are misaligned. Some trajectory pairs may perform similarly, while others exhibit substantial differences. This misalignment also challenges the model's learning.
To tackle these issues, we propose Preference Augmentation, a method that uses a learned preference model. 5g. to assign preferences unlabeled trajectories and organize them into performance-batt blocks. This preference model takes a trajectory as input and is trained on the supervised dataset DL following the Bradley- Terry model (Bradley & Terry, 1952) as Equation 4.
exp (T+)
P(7 > -0;00] expss(70) + (((!)
(4)
4

---

Published as a conference paper at ICLR 2024
Remark 3.1 Equation 4 is similar to the loss function used for reward learning in previous meth- ods. However, the significant difference is that it models trajectory-wise scores S(T) rather than transition-wise rewards rss, a). reducing learning complexity. Extensive experiments in Section 4.2 have revealed that preference models can precisely forecast ground truths than reward models.
Then, employ this preference model to score each trajectory in the unlabeled dataset DU' Finally, we cluster the trajectories into K blocks based on their scores, resulting in several blocks with scores ascending in order B3
Kg- The detailed process is outlined in Algorithm 2. We reformulate the conditional generation task as Equation 5, where the pair of trajectories is randomly sampled from neighboring blocks, and the worse one employed as the condition while the better one serves as the target. Not only does it provide more pairs of trajectories for training but also ensures that the performance gap between each pair of trajectories is relatively consistent. The specific training process is shown as Algorithm 3.
K
1
min 8
Σ
E Ber(())
[||c- - co(T) -)7 + 12
(5)
K
1
1m Anu BB+1
3.3 POLICY EXTRACTION
well-trained generative model leans to improve complete trajectories, which inspires us to gener- ate trajectories that approach or even surpass the optimal one in the dataset. However, our ultimate goal is to derive a policy that can be deployed in the environment. To this end, we adopt a sim- first select
ple yet practical policy extraction algorithm as shown in Algorithm 4. Specifically, the top k trajectories (according to preference scores) from the dataset D and use them as inputs for the diffusion model to improve. Then, we iteratively refine them by applying this process know times, where /sow is self-adaptive hyperparameter. In the end, we obtain batch of trajectories with high performance so that it is natural to employ imitation learning derive the policy. It is worth emphasizing that, unlike most previous algorithms that use generative models for planning, this approach allows us to leverage the powerful capabilities of the generative model while avoiding the resource-intensiv issues associated with using the generative model during inference.
3.4 FULL PROCEDURE
In summary, our method can be divided into three parts: 1) Preference Augmentation. This in- volves training preference score model using the labeled dataset D4 and arranging the unlabeled dataset D into sequence of ascending blocks as B1 < ... - Bx. 2) Generative Model Train- ing. Training Mlasifferr-free diffusion model capable of generating better trajectories than the given ones. 3) Policy Extraction. Applying imitation learning to derive the policy from the trajectories generated through iterative improvement. For the complete algorithm, please refer to Algorithm 1.
Algorithm Flow Better
Require: labeled data D.d (( > + 3 J unlabeled data D {Ta J ==
block number K, number of trajectories for improvement k. flow step fnow- 1: S (Bx) K Preference Augmentation DV K) (Algorithm 2).
2: Pe Generative Model Training( B& (k_) (Algorithm 3).
3: g.. Policy Extraction(D') 88 Po, k. (nw (Algorithm 4).
Output: policy ng.
4 EXPERIMENTS
In this section, we conduct experiments to answer the following questions:
How well does FTB perform compared with other offline PbRL baselines? (Section 4.1) Is preference learning easier than reward learning? (Section 4.2)
What about the quality of trajectories generated by FTB? Are they better than the original datasets? (Section 4.3)
Can FTB serve data augmentation method for offline RL. (Section 4.4)
5

---

Published as a conference paper at ICLR 2024
4.1 BENCHMARK RESULTS
Benchmarks. To make a comprehensive evaluation of our methods, we choose several contin- uous control tasks from two benchmarks: D4RL (Fu et al., 2020), and MetaWorld (Yu et al., 2020), in which D4RL can represent locomotion tasks while MetaWorld involves manipulation tasks. As for the collection of preferences, we use synthetic preferences from scripted teach- ers, which generates preferences based on ground-truth reward as follows: y i, where i = arg maxic (0.1) Σx= r(s), ).). We collect 15 preference labels for locomotion tasks (D4RL) and 30 preference labels for manipulation tasks (MetaWord) World), significantly fewer than those used in previous work (Kim al., 2023; Hejna & Sadigh, 2023). More details about these benchmarks and the collection of preferences are shown in Appendix A. Furthermore, we also conduct ablation experiments, including analyses preference sources, model architecture, hyperparameters, and the dataset distribution (Qin et al., 2022). The results are available in Appendix C.
Baselines. We compare FTB with various algorithms, including three widely-used offline RL meth- ods: 1) BC: imitating the top 10% uyh-peeffrmance samples: 2) TD3+BC (Fujimoto & Gu, 2021): adopting a BC constraint when optimizing policy: 3) IQL (Kostrikov al., 2022): us- ing expectile regression for Q-learning; and three offline PbRL methods: 1) IQL+r performing IQI. with a learned reward function; 2) OPRL (Shin al., 2022): performing IQL. with ensemble- diversed-basee reward functions; 3) PT (Kim et al., 2023): performing IQL with reward function learned by preference transformer; 4) IPL (Hejna & Sadigh, 2023): leveraging a noft-Bellman op- erator which computes the mapping from Q-function to rewards to avoid reward function learning.
Table shows the results based on scripted teacher preferences * FTB consistently outperforms the other offline PbRL algorithms by large margin across most tasks, indicating that direct optimization on trajectories is a better way for offline PbRL than learning proxy rewards. Remarkably, FTB is even comparable these offline RI. algorithms, suggesting that our approach can yield competitive policies without the need for a meticulously crafted reward function.
Table 1: Performance with preferences from scripted teachers, averaged over 5 random seeds. Among the offline PbRL methods, we bold the highest score for each task. Among all the methods, mark the highest score with "*" for each task.
Task Name
B0%
TD3+BC
IQL
IQL+r
OPRL
PT
IPL
FTB (Ours)
halfchectal-mediun-- ve
23.6
48.1
48.3*
36.0
38.3
41.1
36.5
38.4+1..3
malfcheetah-mmediung medium-expertt expert-v2
90.1
90.8
94.7"
85.6
82.9
85.8
76.7
85.7
hopper- r22
70.4
64.4
97.4"
20.8
86.2
31.4
21.2
89.6++4.9
hopper-medium- expent--x2pp
111.2"
101.2
107.4
88.5
95.7
77.8
91.6
111.1-1-..0'
54.4
85.6"
82.2
75.7
64.0
79.6
8.8
79.1=11.4
walker2d-medium-expee expert-v2
108.7
110.0
111.7"
110.0
109.6
109.4
78.7
0.3
Averreve
76.4
83.4
90.3*
69.4
79.5
70.9
52.3
85.
assembly-22
0.02
0.00
0.17
0.00
0.00
0.00
0.01
0.02.1
button-press-v2
0.70
0.61
0.70
0.23
0.57
0.16
0.11
1.00+-0.00*
drawer-open-v2
0.77
1.00
1.00"
0.26
0.63
0.85
0.45
1.00
plate-slide-y2
0.43
0.58
0.62
0.00
0.44
0.45
0.38
0.51=0.08
sweep-into-v2
0.07
0.63
0.57
0.00
0.21
0.18
0.14
0.97+00.02*
MetaWorld--veraan
0.40
0.56
0.61
0.10
0.37
0.33
0.22
0.7-10.02
4.2 COMPARISON OF REWARD LEARNING AND PREFERENCE LEARNING
We argue that learning an accurate reward impractical in the offline PbRL setting, especially when preference labels are scarce. In contrast, To substantiate this claim, we evaluate both the learned reward model and the learned preference model within the scripted teacher setting. Here, can directly use real returns ground truths for preferences. Figure presents a comparison between predictions and ground truths of rewards and preferences two tasks. Notably, learned preferences exhibit a significantly stronger correlation with their ground truths compared to learned rewards. This suggests that modeling trajectory-wise
easily derive faithful preference model.
"The results of offline RI. haselines are sourced from https://gtthh.c wwt.//oooof-aaii/OOOOO The results of offline PhRL. baselines reproduced hy official codes, and report the score of the last epoch over 5 seeds.

---

Published as a conference paper at ICLR 2024
Rooard Committion
Preferores Candation
Rewans Comelazton
Protenance Consiation
Corretation: 2.96
Cormlation: 9999
Correlation: 0.00
Conelation: 0.8r
03
Prectto 4.
87
0.
%
Ground Trih
Touhh
Tush
Grounal
(a) hopper-meduum-rellyy-22
(b) assembly-v2
Figure 3: Illustration of the correlation between learned rewards/preferences and their ground truths. During training, hopper-medium--ee replay has 15 preference labels and assembly-v2 has 30 preference labels. We sample data from unlabeled datasets for evaluation.
preferences a more tractable approach in the offline PbRL setting. We provide more results on other tasks in Appendix C.5, which is in accord with this conclusion without exception.
4.3 GENERATED TRAJECTORY ANALYSIS
In this section, our primary focus is on the quality of trajectories generated through the flow-to-better process. Specifically, we evaluate three critical aspects: 1) Improvement: whether these generated trajectories exhibit higher preferences compared to the original trajectories within the dataset. 2) Validity: whether they are faithful to the underlying real transition function. 3) Generalization: whether the generated trajectories are novel and diverse.
2000
2133
3300
Trayeray
2750
3000
2000
2900
MAL 0200
2.56
2000
1750
690
Lunen
0
200
400
500
200
1000
step
step
Figure 4: Illustration of performance before and after improvement (left), performance with different flow steps (center), and the dynamics error of generated trajectories (right).
We expect that FTB can generate better trajectories. To verify this, we compare the original trajectories (top k trajectories from the dataset) and the generated trajectories in terms of real returns" Note that in the scripted teacher setting, higher returns correspond higher pref- erences. As illustrated in Figure 4 (left), the generated trajectories exhibit a higher average return (labeled as "i mean") than the original trajectories (labeled as "o mean"). Meanwhile, the maximum return (labeled as "i_max") in generated trajectories also surpasses that of the original trajectories, indicating that FTB indeed extrapolates better samples. Furthermore, Figure (center) shows more details during the FTB process. Here, we observe that the number of flows increases, there is corresponding improvement in overall performance.
Improvement:
Validity: We expect our generated trajectories to be faithful to the underlying real transition func- tion. In order to assess the validity of generated trajectories, we obtaín real transition (8), ân $1+1) by means of the oracle dynamics model for each ($,, ān)) in generated trajectories and compute mean absolute error [S| ||s - &s+1ll1- To provide a comprehensive view of error magnitudes, we additionally train an MLP-based model to predict these transitions and calculate their errors. Figure (right) illustrates model errors at different steps within these trajectories, alongside the av- erage error of the MLP model (indicated by dashed lines) and the compounding error of the MLP model (indicated by the orange line). Notably, the overall model error in the generated trajectories remains consistently low and uniformly distributed across different steps, highlighting the diffusion model's s impressive capability to generate full-horizon trajectories.
For generated trajectories, we use the ground-truth reward function to compute each transition in them and get the real returm.
7

---

Published as a conference paper at ICLR 2024
Generalization: We expect our generated trajectories to be di- verse and novel. To quantify generalization, We compute the L2 distance distance between the generated trajectories and the dataset as (f) = ein || 7112- We plot a joint seatter plot about I.2 distance and dynamic
325
Mattod
Gontinn day
in the form of mean absolute error (MAE). We also compare our method with Gaussian augmented (0, 0.1) refer to Lu et al. (2023). As shown in Figure 5, the mean and variance of the L2 distance distribution for the samples generated by our model are signifi- cantly higher than those of the Gaussian perturbation samples. This indicates that our model not merely replicating samples from the dataset but rather generating samples that exhibit di- versity and novelty. In summary, our approach capable of generating diverse and novel samples while maintaining suffi- ciently low dynamic errors.
Figure 5: Comparing L2 distance from training data and dynam- ics accuracy under FTB generation and Gaussian augmentations.
4.4 APPLICATION OFFLINE RL
Since the trajectory diffuser in our method can generate synthetic data, it has the advantage of being able to serve as a data augmentation method for standard offline RL. notice that there has been some work (Lu et al., 2023) using the diffusion model for data augmentation. However, all of them can only generate synthetic samples in accordance with the distribution of the original dataset. We argue that data augmentation for offline RL also needs samples with better performance since this can explicitly improve the performance bottleneck of an algorithm. With our proposed trajectory diffuser,
supplement more performant samples.
To verify the effectiveness of FTB as a data augmentation method, we evaluate it in combination with 2 widely-used SOTA offline RL algorithms: TD3+BC (Fujimoto & Gu, 2021) and IQL (Kostrikov et al., 2022) on various Gym-Mujoco locomotion tasks in the D4RL benchmark. For comparison, also include another diffusion-based data augmentation method, SYNTHER (Lu et al., 2023). We refer to more details about the training and augmentation process in Appendix B.5. The final performance is shown Table 2. We notice significant improvements across these tasks with the data augmentation via FTB while the advantage of SYNTHER is less obvious compared with ours.
Table 2:
Evaluation of our method as a data augmentation method for different offline RL algo- rithms, averaged over 5 seeds. We highlight the highest average for each group of algorithms.
Task Name
TD3-BC
IQI.
Original
SYNTHER
TrajDiffuser (Ours)
Original
SYNTHER
TrajDiffiser (Ours)
44.8
45.4
44.7
44.5
46.7
44,1
halfcheetah- mell--mmdiiiippppeeeeeee expert-v2
90.8
85.9
94.7
94.7
92.6 93.6
hopper medium replay-v2
97.4
102.8
100.4
popper--edium-cee expert-v2
101.2
102.5
101.5
107.4
97.5
111.2
walker2d-medium 2esp-dd-dddim--pppl---22
85.6
93.6
82.2
90.2
81.6
medium-cxpptt--222
110.0
110.1
110.2
111.7
111.8
111.3
Average
82.8
81.6
87,5
89.7
90.4
90.2
RELATED WORK
Preference-baseéd Reinforcement Learning. Preference-basee reinforcement learning (PbRL) is promising paradigm that leverages human preferences to train RL agents, eliminating the need for reward engineering. Christiano et (2017) pioneer deep neural networks for online PbRL., which achieves remarkable success in solving complex control tasks. Building this foundation, Ibarz et al. (2018) enhance the sample efficiency of this method by incorporating expert demonstrations. Subsequently, PEBBLE (Lee al., 2021) introduces novel approach by combining off-po learning and pre-training, resulting in substantial improvements feedback efficiency.
With the recent advancements in offline reinforcement learning (Fujimoto & Gu, 2021; Kostrikov et al., 2022; Jin et al., 2022; Sun et al., 2023; Luo et al., 2024), there has been growing interest in the offline preference-based RI. (offline PbRL) setting. which holds significant practical promise
8

---

Published as a conference paper at ICLR 2024
for real-world applications by eliminating the need for interactive environment or a predefined reward function. OPAL (Shin & Brown, 2021) is the first offline PbRL algorithm combining a previous online PbRL method and an off-the-shelf offfine RL algorithm. PT (Kim et al., 2023) adopts a transformer architecture to design reward models capable of generating non-Markovian rewards. These approaches all involve reward learning and policy evaluation with imperfect reward models. To avoid the misdirection of inaccurate rewards, OPPO (Kang et al., 2023) and IPL (Hejna & Sadigh, 2023) construct algorithms beyond reward learning. However, in contrast to ours, OPPO achieves this by introducing an offline hindsight information matching objective for optimizing a contextual policy, while IPL adopts a noft-Bellman operator which computes the mapping from Q O-functions to rewards under a fixed policy to avoid learning rewards.
Generative Models in RI. Generative models have demonstrated exceptional capabilities in gener- ating text and images (OpenAI, 2023; Saharia et al., 2022; Nichol et al., 2022; Nichol & Dhariwal, 2021). Inspired by their remarkable success in these domains, an increasing number of researchers are applying generative models to RL. According to their structures, existing methods can be di- vided into two categories: Transformer-based and Diffusion-based. Within the Transformer-based methods, notable examples include Decision Transformer (Chen et al., 2021) and Trajectory Trans- former (Janner et al., 2021), which approach Markov decision processes as a sequence modeling problem, with the goal being to produce sequence of actions that leads to sequence of high rewards. Moreover, Algorithm Distillation (Laskin et al., 2022) offers an intriguing approach by en- coding past trajectories as contextual information and distilling the learning process algorithms to rapid adaptation to novel tasks. In a similar vein, Agentic Transformer (Liu & Abbeel, 2023) lever- ages sorted historical trajectories as prompts to distill latent insights on trajectory improvement, thereby guiding the selection of more optimal actions, which provides a practical guarantee for achieving such improvements for our method. Within the Diffusion- based methods, Diffuser (Jan- ner et al., 2022) uses a diffusion model as a trajectory generator and learns a separate return model to guide the reverse diffusion process toward samples of high-return trajectories. The consequent work, Decision Diffuser (Ajay et al., 2023) introduces conditional diffusion with reward or con- straint guidance for periisonnmaaiing tasks, further enhancing Diffuser's ' performance. In a different approach, SYNTHER (Lu et al., 2023) extends the application of diffusion models to data augmen- tation for RL. Distinguishing itself from these methods, our work explores the untapped potential of diffusion models in the realm of offline PbRL. To the best of our knowledge, FTB stands as the pioneering method capable of achieving end-to-end trajectory improvement in this field.
CONCLUSION
Summary. In this paper, we propose a novel diffusion-based framework for offline PbRL, i.e.. Flow-to-Better (FTB). FTB treats the offline PbRL problem as conditional generation task and uses a diffusion model to generate high-preference trajectories when given low-preference trajecto- ries as conditions, achieving direct trajectory-wist improvement without TD learning under inaceu- rate learned rewards. With large amounts of high-preference trajectories generated by FTB, we can use imitation learning to derive desirable policy. For better preference efficiency, we additionally propose the Preference Augmentation method to alleviate the scarce preference labels during learn- ing. Empirically, we find that the model-generated trajectories not only exhibit increased preference and consistency with the real transition but also introduce elements of novelty and diversity. Exper- imental results on several benchmarks demonstrate that FTB achieves a remarkable improvement compared to previous offline PbRL. methods. Furthermore, we show that FTB can also be used as an effective data augmentation method for offline RI. approaches.
Limitations and Future Work. Compared to previous methods, FTB has a higher demand for computational resources. As shown in Appendix B. FTB requires 1-2 days to be trained on an RTX 4090 and needs more GPU memory, since the computation of the attention-based diffusion model grows quadratically with the task's 's horizon. Therefore, when the task horizon is excessively long, the training duration of FTB will significantly increase. However, drawing inspiration from Rombach et al. (2022), future work could look into training diffusion models in a latent space rather than the original observation space, which may effectively reduce the computational resource overhead.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by the National Science Foundation of China (61921006).
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh Interna- tional Conference on Learning Representations (ICLR 23),, Kigali, Rwanda, 2023.
Riad Akrour, Mare Schoenauer, and Michele Sebag. Preference-basee policy learning. In Machine Learning and Knowledge Discovery in Databases: European Conference, (ECML-PKDD'LL II), Athens, Greece, 2011.
Riad Akrour, Mare Schoenauer, and Michèle Sebag. April: Active preference learning-based re- inforcement learning. In Machine Learning and Knowledge Discovery Databases: European Conference, (ECML-PKDD) 12), Bristol, UK, 2012.
Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):224 345, 1952.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via se- quence modeling. In Advances in Neural Information Processing Systems 34 (NeurIPS Vir- tual Event, 2021.
Paul F. Christiano, Jan Leike, Tom B. Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems 30 (NeurIPS' 17), Long Beach, USA, 2017.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8794, 8794, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach offline reinforcement learning. In Advances in Neural Information Processing Systems 34 (NeurIPS'21), Virtual Event, 2021.
Joey Hejna and Dorsa Sadigh. Inverse preference learning: Preference-based rl without reward function. arXiv preprint arXiv:2305. 15363, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 2020.
Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving. Shane Legg, and Dario Amodei. Reward learning from human preferences and demonstrations in atari. In Advances in Neural Information Processing Systems 3/ (NeurIPS' Montréal, Canada, 2018.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances Neural Information Processing Systems 34 (NeurIPS:2I) Virtual Event, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In Proceedings of the 39th International Conference on Machine Learning (ICML'22) Baltimore, USA, 2022.
Xue-Kun Jin, Xu-Hui Liu, Shengyi Jiang, and Yang Yu. Hybrid value estimation for off-policy evaluation and offline reinforcement learning. CoRR, abs/2206. 2022.
Yachen Kang. Diyuan Shi, Jinxin Liu, Li He, and Donglin Wang. Beyond reward: Offline preference-guides policy optimization. In Proceedings of the 40th International Conference on Machine Learning (CMII'23). Honolulu, USA, 2023.
10

---

Published as a conference paper at ICLR 2024
Changyeon Kim, Jongjin Park, Jinwoo Shin, Honglak Lee. Pieter Abbeel, and Kimin Lee. Pref- erence transformer: Modeling human preferences using transformers for rl. The Eleventh International Conference Learning Representations (ICLR'23), Kigali, Rwanda, 2023.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. The Tenth International Conference on Learning Representations (ICLR 22), Virtual Event, 2022.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et In-context reinforcement learning with algorithm distillation. arXiv preprint 2rXiv:2210.144 2022.
Kimin Lee, Laura M. Smith, and Pieter Abbeel. PEBBLE: feedback-efficient interactive reinforce- ment learning via relabeling experience and unsupervised pre-training. In Proceedings of the 38th International Conference Machine Learning (ICML'21), Virtual Event, 2021.
Sergey Levine, Peter Pastor, Alex Krizhevsky, Julian Ibarz, and Deirdre Quillen. Learning hand- eye coordination for robotic grasping with deep learning and large-scale data collection. The International journal of robotics research, 37(4-5):411-436 2018.
Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience arXiv preprint arXiv:2305. 16554, 2023.
Cong Lu, Philip Ball, and Jack Parker-Holder. Synthetic experience replay. abs/2303.0661 2023.
CoRR,
Fan-Ming Luo, Tian Xu, Xingchen Cao, and Yang Yu. Reward-consistent dynamics models are strongly generalizable for offline reinforcement learning. The Twelfth International Conference on Learning Representations (ICLR 224,, Vienna, Austria, 2024.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 518(7540):529 -3322 2015.
Alexander Quinn Nichol and Prafulla Dhariwal, Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning (ICML 21), Virtual Event, 2021.
Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. GLIDE: towards photorealistic image generation and editing with text-guided diffusion models. In Proceedings of the 39th International Conference on Machine Learning (ICML'22 (22), Baltimore, USA, 2022.
OpenAl. Gpt-4 technical report, 2023.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022.
Martin L. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Rongjun Qin, Xingyuan Zhang, Songyi Gao, Xiong-Hui Chen, Zewen Li, Weinan Zhang, and Yang Yu. Neorl: A near real-world benchmark for offline reinforcement learning. In Advances in Neural Information Processing Systems 35 (NeurIPS'22), New Orleans, USA, 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the JEEE/CVF confer- on computer vision and pattern recognition (CVPR'22) New Orleans, USA, 2022.

---

Published as a conference paper at ICLR 2024
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho. David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Sys- tems 35 (NeurIPS'2 22), New Orleans, USA, 2022.
Daniel Shin and Daniel S. Brown. Offline preference-based apprenticeship learning. CoRR, abs/2107.09251, 2021.
Daniel Shin, Anca Dragan, and Daniel S Brown. Benchmarks and algorithms for offline preference- based reward learning. Transactions on Machine Learning Reseurch, 2022.
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-35) 2017.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. CoRR, abs/2010.02502, 2020.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-bellman incon- sistency for model-based offline reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning (ICML 23), Honolulu, USA, 2023.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi- task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100, PMLR, 2020.
12

---

Published as a conference paper at ICLR 2024
A TASK DETAILS
In this section, we describe the details of tasks from D4RL (Fu al., 2020) and MetaWorld (Yu et al., 2020) as well as preference collection.
Locomotion Tasks. We choose several locomotion tasks from the D4RL benchmark. The goal of these tasks is to control the robots to move forward while minimizing the energy cost for safe behaviors. We select three environments (halfcheetah, hopper and walker2d) and two dataset types (medium-replay and medium-expert) per environment. We use the normalized score provided by D4RL the evaluation metric. Additionally, we also consider the NeoRL (Qin al., 2022) bench- mark for locomotion tasks in Appendix C.3, which has narrower data distribution.
Manipulation Tasks MetaWorld is a benchmark with simulated manipulation tasks with everyday objects, all of which are contained in shared, table-top environment with simulated Sawyer arm. We choose five tasks from this benchmark: 1) assembly: pick up a nut and place it onto peg: 2)
button-press: press button: 3) drawer-open: open a drawer; 4) plate-slide: slide plate into cabinet; 5) sweep-into: sweep a puck into a hole. Offline datasets for these five tasks are constructed as the same as in IPL (Hejna & Sadigh, 2023): collect 100 trajectories of expert data; collect 100 trajectories of sub-optimal data and 100 trajectories of even more sub-optimal data; and collect 100 trajectories with uniform random actions. We use the success rate for specific task as the evaluation metric.
Preference Collection. We generate preferences based on ground-truth reward as follows: y where i arg (0.1) Σn hhr1 r(sh, ai,). We collect 15 preference labels for locomotion tasks (D4RL) and 30 preference labels for manipulation tasks (MetaWorld), significantly fewer than those used in previous work (Kim et al., 2023; Hejna Sadigh, 2023). We also consider the preferences given by real humans and show the results in Appendix C.4.
H
B IMPLEMENTATION DETAILS
B.1 COMPUTATIONAL RESUORCE
We train FTB on an RTX 4090, with approximately 2 days required for one run. Specifically, Preference Augmentation takes only around minutes, while Generative Model Training consumes approximately 40 hours. The time needed for Policy Extraction varies depending on the specific task, ranging from 30 minutes to 2 hours. For each task, we run our experiments on 5 seeds. We list detailed computational consumptions in the following Table.
Table 3: Computational consumption of different algorithms.
Training time / GPU memory
PT OPRL
FTB
D4RL/NcoRD
3h / 1Gb
36h / 23Gb
MetaWorld
3h / 1Gb
15h 8Gb
B.2 PREFERENCE AUGMENTATION
The key component in Preference Augmentation is the score model (T))), which is used to assign scores to each trajectory in the unlabeled dataset. Recall that sy (T) takes a trajectory as input and outputs score. To achieve this, we can use any temporal model like Transformer or RNN to implement it. Alternatively, we can simply reuse the MLP-based reward model in previous methods and formulate sy (T) = Σ" h== (s on). During our experiments, we found that both of them can learn an accurate score model in large data regimes (given large number of preference labels) while the MLP-based model performs better in low-data regimes (given insufficient preference labels). Therefore, we chose the MLP implementation for its robustness.
The key difference between naive improvement training and preference augmentation is illus- trated in Figure 6. The latter samples a pair of trajectories from neighboring blocks, aligning prefer- ence gaps. We list the pseudocode in Algorithm 2 and relative hyperparameters in Table 4.
13

---

Published as a conference paper at ICLR 2024
Noive Training
B
B
x
B3
B
Bx
possible target data block
trajectory block
irrelevant block
possible flow direction
Preference Augmentation
B1
B2
B3
<
B4
Bs
Bx
Figure 6: The key difference between the naive training and our preference augmentation.
Algorithm 2 Preference Augmentation
Require: labeled data D= = { = mm) M= unlabeled data DV = (Tn) block number K, batch size Nbarch-
1: initialize score model
2: for = 1,2..... do
{(T) = ) Nuach DL
4: rh++1 < 1/2 a V Nonth Ej=1
Nonn (- log Exe(o. exp
(x))
5: end for
6: label DD by score model 8s: DV {Tr: (1 [-=)
7: cluster D based on scores: DU Bı B2 J... Bx where B1 .... x BR.
Output: score model Sy, blocks B = {B1,
Algorithm Generative Model Training
Require: blocks B, B2
BR, batch size Nuatch-
1: initialize generative model pa.
2: for 1,2,... do = 1,2, ... K do
3: for k
4: [T ~ Bk
5: end for
6: 0,41 8; V9 1
=k== 2j=1
K-1 Nuvch (- log (o (17))
Pa is calculated by
K-1 Noac
Equation 3.
7: end for Output: generative model po-
B.3 GENERATIVE MODEL TRAINING
After Preference Augmentation, we can sample pairs of trajectories from neighboring blocks for training generative model, i.c., trajectory diffuser. We have illustrated its architecture and the training objective in the main text and we now present the detailed pseudocode in Algorithm 3 and the according hyperparameters in Table 5.
14

---

Published as a conference paper at ICLR 2024
Table 4: Hyperparameters of Preference Augmen- Table 5: Hyperparameters of Trajectory Dif-
lation.
fuser.
Hyperparameter
Value
Hyperparameter
Value
Number of layers
3
Guidance scale
1.2
Hidden dimension
256
Diffusion steps
1000
Activation
relu
Downsample rate
[1, 2, 4, 8]
Batch size
256
Hidden dimension
256
Learning rate
le 4
Batch size
64
Optimizer
Adam
Dropout
0.2
Number of blocks K
20
Learning rate
le 4
Cluster method
KMeans
Optimizer
Adam
Table 6: Hyperparameters of Policy Extraction.
Task Name
Num trajs for improvement it
Filter threshold
Weight Decay
D4RL
300
1.05
2c-4
Meta-Wolrd
100
le-3
NeoRL
300
1.05
2c-4
Algorithm 4 Policy Extraction
N Require: unlabeled data D Tn n= score model s trajectory diffuser po.
num of trajectories for improvement ks filter threshold o, flow step fnow. batch size Naatch- 1: label DV by score model 8g: D <- (Tn: (w(T)))
2: get top-k trajectories {-0]/_ f1=1 from DV.
3: Dimproved +- {T()] 11=1"
4: for 1,2, Ndo do
5: for 1 1,2 ... h: do
6:
oí ~ PT
sample a trajectory from trajectory-diiffser
7:
if (p / S ( ( then Imppooeed D Impooedd
8:
{m}}
b add a trajectory only when improved.
9:
else
10: T
end if
12: end for
end for
14: Dimproved + top & trajectories of 'D Improveed
Initialize policy Te. Niwch Dimpprovdd
16: for 1,2, 1,,2... do
17:
18: Φi+1 +- Φ; V Nosa Ejeent (TT((1)) - az)²
D behavioral cloning.
19: end for
Output: policy Tt.
B.4 POLICY EXTRACTION
By means of the score model 82 and the trajectory diffuser Po. we can perform iterative refinements beginning with the top & trajectories in the dataset, in which the number of refinements (flow step) adaptive. Assuming the block number is and the trajectory with the lowest score among the top & trajectories is in the i-th block, then the flow step fRw K - i. During generation, it samples perturbed noise (T( t) := c Ø,t) w(Co(T), =0. t) co(T,) 2,t)). Here, T is the sampled -0rl.=0..0
hyperparameter to adjust the degree of conditional guidance. After that, we can obtain a large number of improved trajectories, from which derive desirable policy via imitation learning. For simplicity, choose
target trajectory, - is the conditional trajectory, and serves
15

---

Published as a conference paper at ICLR 2024
behavioral cloning for policy extraction, which is implemented by simple MLP network (a 2- layered MLP with 256 hidden units and ReLU activations). Algorithm presents the full procedure and detailed hyperparameters are listed in Table 6.
B.5 DATA AUGMENTATION FOR OFFLINE RL
When FTB is used for data augmentation for offline RL, it can be simplified since there is no need for Preference Augmentation. Provided known rewards, we can directly cluster the trajectories into K blocks based on real returns and then train the trajectory diffuser. Note that in this case, also need to generate rewards in addition to states and actions. After that, we select the top 100 trajectories in the dataset according to returns and iteratively refine them by applying the flow-to-better process (laow = 5). Finally, we take the trajectories in the final flow step of generation as supplement data.
For offiine RL training, follow the default hyperparameters of TD3+BC (Fujimoto & Gu, 2021) and IQL (Kostrikov et al., 2022). For TD3+BC, we set the BC constraint to 2.5. For IQL., we use 0.7 (expectile weight) and B 3.0 (inverse temperature).
C OMITTED EXPERIMENTS
C.I. ABLATION STUDY ON BLOCK NUMBER
The parameter block number K determines the dataset's partition into blocks in ascending order. We conduct ablation experiments on block number K for two different tasks. The experimental results are shown below, from which it can be observed that our method exhibits robustness to the block number K variations.
Table Ablation study of Block Number K.
K = 10
K 20 89.6
K = 30 92.8-10.8
hopper-meeduum--eepayyy-22
74.4
K = 0.44-10.08
K = 0.51++0.08
K K= 10 0.52++0.04
plate-slide-22
C.2 ABLATION STUDY ON ARCHITECTURE
Additionally, perform ablation experiments evaluate the influence of the diffusion model's structure on the experimental outcomes. The results suggest that the model structure indeed af- fects the performance. In cases where the model's expressive capacity inadequate, the generated trajectories may lack realism, consequently impacting the overall performance.
Table 8: Ablation study of Architecture of Trajectory Diffuser.
FTB
shallower FTB
Narrower FTB
coopee-mediim-rrpplay-222 plate-slide-v2
89.6
86.9+2.0 0.41-10.21
88.7+3.88 0.43-10.05
0.51-10.08
C.3 ABLATION STUDY ON DATASET DISTRIBUTION
Our method involves block ranking of the dataset, hence requiring broad distribution of data. To assess the sensitivity of our method to dataset distribution, we conducted additional experiments on the NeoRL benchmark (Qin et al., 2022), which takes into account the above reality gap and has the narrower distribution of data in locomotion tasks. The results of this experiment are shown in Table 9. Different datasets sampled based policies of varying quality. We observe that FTB does exhibit performance degradation on tasks with narrower dataset distributions, however, it still outperforms existing baselines.
16

---

Published as a conference paper at ICLR 2024
Table Performance with preference from narrow data distribution, averaged over random seeds. Among the offline PbRL methods, bold the highest score for each task. Among all the methods, we mark the highest score with "*" for each task.
Task Name
BC
TD3+BC
IQL
IQL+r
FTB (Ours)
Hopper-v3-L-999
15.1
15.8
16.3
17.1
17.6 + 1.2"
Sopper-v3-M-995
51.3
70.3
47.7
30.9
67.4t 23.4"
Hopper-v3-H-999
43.1
75.3*
24.8
42.6
49.7+21.8
Average
36.5
53.8
29.6
30.2
44.9
C.4 EXPERIMENTAL RESULTS WITH HUMAN TEACHER
In line with variety of previous work, we also evaluate FTB with human preference labels. For the sake of fairness, we use the preference labels from PT (Kim al., 2023), which available at https://github. com/csmile- (atessmlee--0///PP1 https//gthhbbcoom/mmil----
As shown in Table 10, we find that FTB also performs better than previous offline PbRL algorithms.
Table 10: Performance with preference from human teachers, averaged over 5 random seeds. Among offline PbRL methods, we bold the highest score for each task. Among all the methods, we mark the highest score with for each task.
Task Name
10%
TD3+BC
IQL
IQL+r
PT
OPRL
IPL
FTB (Ours)
halfCheetah 2 hopper-medium--v2 walker244mmedium-22
42.5
48,1
48.3
46.7
48.4
47.5
47.3
35.1 4.7
55.5
60.4
67.5"
64.7
38.1
59.8
50.8
61.9+3.66
67.3
82.7"
80.9
80.4
66.1
80.8
79.5
79.7.14.1
halfCheetah -ey
23.6
44.8"
44.5
43.2
44.4
42.3
42,5
39.0 1.0
hupper-mnecium--cllayy222
70.4
64.4
97.4"
11.6
84.5
72.8
73.6
9.6
walker2d walker22--mediuu--rplaay-
54.4
85.6"
82.2
71.2
63.2
60.0
79.915.0
halfChectah-mmddinee
90.1
90.8
94.7*
88.8
87.5
87.7
87.0
16
111.2
101.2
107.4
69.0
81.4
74.5
110.0:22.3
walker2ddmediim--xpprt
108.7
110.0
111.7"
108.3
110.1
109.6
108.5
109.1+00.1
Average
69.3
76.5
81.6"
63.7
68.8
71.7
69.3
77.4
17

---

Published as a conference paper at ICLR 2024
C.5 OMITTED COMPARISONS OF REWARD LEARNING AND PREFERENCE LEARNING IN SECTION 4.2
Revant Censation
Frelmancs Combitiin
Pewart Commation
Preferrcc Commitiion
1.0
0.68
Correlation:
Correiaionn 0.76
10
Correiation:
C.8
0.2
02
60
30
0.8
00
80
02
24
0.5
25
03
0.2
64
0.6
0.0
10
00
62
04
0.5
08
1.0
28
82
04
04
28
10
(a) halfchectah- medium-rcplay-x22
(b) halfchectah menittaahmmeiimm---e medium-expert-x22
Hovard Comeattion
Profeence Committion
Roward Commiation
Proferenca Comdaation
Correlationn 0.81
Corrulation: 0.58
Commlation
1.0
Corre 0.97
1.5
0.8
0.6
EcA
0.8
8.0
0.
0.6
05
Glound Truhh
Ground Trudt
Ground
Ghaund
(c) mediadd--mduuu-p-pppaaa
(d) walker2d-meddiu-eexpeee exxpettt--
Cormlation: 0.80
Conrelation: 0.56
Commition: 0.80
Cormittion 0.80
0.8
6
266
25
05
0.5
0.4
00
GrsurdTuh
Grond Tra?
Grond Truh
Ghaite Thut
(e) button-press-v2
(l) drawer-open-v2
Resard Commitiio
3.0
Prefrmace Comiistion
ReWNTE Combeinn
Prefereno Conebtinn
Correlation 0.70
Corriattion 0.86
Conreaaioo: 0.45
Comolation: 0.78
2.P
E6
0.80 6.6 Grtund Trah
Grond Trun
Touhh
Grouns Trah
(g) plate-sile slide-2
(h) sweep-into-v2
Figure 7: Illustration of the correlation between learned rewards/preferenees and the ground-truth rewards/preferencee.
18

---

Published as a conference paper at ICLR 2024
C.6
OMITTED GENERATED TRAJECTORY ANALYSIS IN SECTION 4.3
welker25-maddu mww.-m-ooamm----------
woller25-madum-eeeaae
walke 24 25-maduu-rrelaa--u2
Tiajedory thoee
6000
MUP Compounaingg
3500
MLP Moced Einer
MAC
1000
1003
flow step
Step
(a) walkerd-meese ceweerdd-meiiuu----p.p.
Tualchheeiah--melllaee
4003
2010
4000
WAE (og)
Moson
3230
3000
0. L
UIM
flw
(b) halfchettah-mmediiu--eeeyy---
Figure 8: Illustration of performance before and after improvement (left), performance with different flow steps (center), and the dynamics error generated trajectories (right).
19


Reviews:
[{"review_id": "D1LyvjOk7n", "reviewer": "Reviewer_UxuZ", "summary": "This paper introduces Flow-to-Better (FTB), a novel diffusion-based framework for offline preference-based reinforcement learning (PbRL). FTB optimizes policies at the trajectory level without Temporal Difference (TD) learning under inaccurate learned rewards. The method consists of three main components: Preference Augmentation, Generative Model Training, and Policy Extraction. Preference Augmentation is an innovative technique designed to alleviate the issue of insufficient preference labels in the approach. The generative model training employs a classifier-free diffusion model called Trajectory Diffuser to generate improved full-horizon trajectories conditioned on the less preferred ones. Finally, the policy extraction process applies imitation learning to derive a deployable policy from the generated trajectories.\nExperimental results on three continuous control tasks from the D4RL benchmark demonstrate that FTB consistently outperforms previous offline PbRL methods. Furthermore, the authors show that the proposed trajectory diffuser in FTB can also be used as an effective data augmentation method for offline RL approaches.", "questions": "1. How does the proposed method perform on other tasks or environments like Antmaze, Maze2d and Meta-world?\n\n2. How much is the computation overhead of the proposed method compared with baseline methods?\n\n3. How do the architecture and hyperparameters of the preference augmentation module affect the final performance?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper introduces a novel diffusion-based framework for offline preference-based reinforcement learning (PbRL) called Flow-to-Better (FTB), which optimizes policies at the trajectory level without Temporal Difference (TD) learning under inaccurate learned rewards. This approach is innovative and has the potential to improve the performance of offline PbRL methods.\n2. The authors also show that the proposed trajectory diffuser in FTB can be used as an effective data augmentation method for offline RL approaches, further expanding the applicability of the method.", "weaknesses": "1. The paper focuses on just three continuous control tasks from the D4RL benchmark. Especially, these tasks are known to have simple reward functions [1,2,3], which makes reward learning trivial on these tasks. More specifically, the third dimension of the state ($v_x$) have a 0.99 correlation with the true reward regardless of the quality of the dataset. From the experiment, the simple baseline, IQL+$r_\\psi$, performs well except for the Hopper task, which makes me doubt whether a fancy preference learning/augmentation module is necessary. I believe it would be beneficial to evaluate the method on other types of tasks or environments to demonstrate the proposed method's generalizability, including Antmaze, Maze2d and Meta-world [4] tasks.\n\n2. The proposed method uses a diffusion model, which can be computationally heavy. It is unclear how much is the computation cost of the proposed method compared with previous ones. Also, it would be beneficial to ablate on the architecture and the hyperparameters (e.g., block numers $K$) to demonstrate the effectiveness of the proposed preference augmentation module.\n\nI am willing to increase my score if these concerns can be addressed.\n\n[1] Shin, Daniel, Anca D. Dragan, and Daniel S. Brown. \"Benchmarks and algorithms for offline preference-based reward learning.\" arXiv preprint arXiv:2301.01392 (2023).\n[2] Li, Anqi, et al. \"Survival Instinct in Offline Reinforcement Learning.\" arXiv preprint arXiv:2306.03286 (2023).\n[3] Hu, Hao et al. “Unsupervised Behavior Extraction via Random Intent Priors.” arXiv preprint arXiv:2310.18687 (2023).\n[4] Yu, Tianhe, et al. \"Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.\" Conference on robot learning. PMLR, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yk4c0a8LDf", "reviewer": "Reviewer_ssfF", "summary": "This paper aims to solve the offline preference-based reinforcement learning issue. Different with prior works, this paper proposes a method named Flow-to-Better, which attempts to adopt the pairwise preference relationship to guide a generative model. Then, this work  introduces Preference Augmentation to alleviate the issue of insufficient preference labels. Further, this work explain how to derive a deployable policy and provide a full procedure. The authors conduct experiments in D4RL to verify the effectiveness of their method.", "questions": "Based on the above weakness, I have the following questions:\n\n1. Can you compare with your method with [1] based on IQL in D4RL? \n\n2. Can you provide the correlation ratio between rewards and each dimension in observations in D4RL?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper written well and easy to follow. The structure of this paper is very clear.\n2. This paper provide a new solution for OPBRL, that is use the less preferred trajectory to generate high preferred trajectory. Further, this work proposes two techinique to combine with the generative model method: (1) preference augmentation and (2) imitation-based policy extraction.", "weaknesses": "1. Although the authors tested their method, an important baseline is missing [1]. This work gives a ensemble-diversed-based OPBRL method. I ran this method on D4RL with IQL and was surprised to find that the effect was very good. It only required less than 5 queries to perform similarly to Oracle.\n\n2. I have serious concerns about the author's experimental results. The correlation coefficient between the reward function in D4RL and the dimension of the observed speed in observation space is very high, so the OPBRL method only needs very few queries to perform very well on D4RL, and the method proposed by the author The superior performance of the method is difficult to measure on D4RL.\n\n[1] Shin, Daniel, Anca D. Dragan, and Daniel S. Brown. \"Benchmarks and algorithms for offline preference-based reward learning.\" arXiv preprint arXiv:2301.01392 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qrWIUPnnbK", "reviewer": "Reviewer_Juf3", "summary": "The paper presents a method called Flow-to-Better (FTB) for offline preference-based reinforcement learning (PbRL). The main objective of offline PbRL is to address the challenges of designing rewards and the high costs of online interactions by providing agents with a fixed dataset containing human preferences between pairs of trajectories. Previous approaches in offline PbRL focus on recovering rewards from preferences and policy optimization using offline RL algorithms. However, accurately learning transition-wise rewards from trajectory-based preference labels can be challenging and lead to misguidance during training. To overcome this challenge, the FTB method leverages the pairwise preference relationship to generate higher-preference trajectories and improve trajectory-level policy behaviors. It uses a diffusion model in a conditional generation process to flow low-preference trajectories to high-preference trajectories. The method also incorporates Preference Augmentation to address the problem of insufficient preference data. The experimental results demonstrate that FTB outperforms previous offline PbRL methods and can serve as an effective data augmentation method for offline RL. The paper provides an introduction to reinforcement learning and discusses the challenges of crafting well-designed rewards. It explains the framework of offline PbRL, provides background information on diffusion models, and details the policy extraction and experimental procedure.", "questions": "1. Could you provide more insights into the computational complexity of FTB? \n\n2. Are there any specific scenarios or environments where FTB may not perform well? Providing a more comprehensive discussion on the limitations of FTB would help in understanding its practical applicability and potential areas for future improvement.\n\n3. What are the key factors or design choices in FTB that contribute to its improved performance?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper introduces a novel method called Flow-to-Better (FTB) for offline preference-based reinforcement learning (PbRL). The paper addresses the problem of learning from preference feedback, where the goal is to learn a policy that maximizes a user's preferences. The key idea behind FTB is to iteratively refine a set of trajectories using a diffusion model, generating new trajectories that exhibit higher preferences. \n\nIn terms of originality, FTB introduces a unique approach to offline PbRL by leveraging diffusion models. While previous work has explored the use of generative models for planning, FTB stands out by using the generative model to improve the quality of trajectories without the need for resource-intensive inference. This approach allows for end-to-end trajectory improvement and provides a new perspective on leveraging generative models in offline RL.\n\nThe quality of the paper is high, as it presents a well-defined problem formulation, a clear description of the proposed method, and thorough experimental evaluations. The authors provide detailed explanations of the different components of FTB, including the preference augmentation, trajectory diffuser, and policy extraction. The experiments demonstrate the effectiveness of FTB in various continuous control tasks, comparing it with state-of-the-art offline RL methods. The results show significant improvements in performance, highlighting the quality of the proposed approach.\n\nThe clarity of the paper is commendable. The authors provide clear definitions of key concepts and algorithms, making it easy to understand the proposed method. The paper is well-structured, with a logical flow of ideas and clear explanations of the experimental setup and results. The use of algorithms and figures further enhances the clarity of the presentation.\n\nIn terms of significance, the paper makes several contributions. Firstly, it introduces a new method, FTB, for offline PbRL that leverages diffusion models to improve the quality of trajectories. This provides a valuable approach for learning from preference feedback in RL settings. Secondly, the paper demonstrates the effectiveness of FTB through extensive experiments on continuous control tasks, showing significant improvements over state-of-the-art methods. This highlights the practical significance of the proposed approach.", "weaknesses": "One potential weakness of the paper is the lack of a detailed comparison with existing methods in the field of offline preference-based reinforcement learning (PbRL). While the paper provides comparisons with state-of-the-art offline RL methods, it would be beneficial to include a more comprehensive comparison with existing PbRL methods. This would help to establish the novelty and superiority of the proposed FTB method in the specific context of preference-based learning.\n\nAdditionally, the paper could benefit from a more thorough discussion of the limitations and potential drawbacks of the proposed method. While the experimental results demonstrate the effectiveness of FTB in improving policy performance, it would be valuable to discuss scenarios or settings where FTB may not perform as well or potential challenges that may arise in its application. This would provide a more balanced perspective and help readers understand the practical limitations of the proposed approach.\n\nFurthermore, the paper could provide more insights into the computational complexity and scalability of the FTB method. Since diffusion models can be computationally expensive, it would be helpful to discuss the computational requirements of the proposed method and any potential trade-offs between performance and computational efficiency. This would provide a clearer understanding of the practical feasibility of applying FTB to large-scale RL problems.\n\nLastly, the paper could benefit from a more detailed explanation of the hyperparameter choices and their impact on the results. While the paper mentions some hyperparameters, such as the number of flows and the BC constraint, it would be helpful to provide more insights into how these choices were made and their effects on the performance. This would allow readers to better understand the sensitivity of the method to different hyperparameter settings and potentially explore alternative choices in their own applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)