Accelerating Diffusion Planners in Offline RL via Reward-Aware Consistency Trajectory Distillation
Authors: Xintong Duan
Yutong He
Fahim Tajwar
Ruslan Salakhutdinov
J Zico Kolter
Jeff Schneider
Venue: ICLR 2026 Poster (2026)
Keywords: offline RL
consistency models
diffusion models
Primary area: reinforcement learning
TL;DR: A method that makes diffusion planners dramatically faster in offline RL by distilling them into single-step consistency trajectory models that directly optimizes for rewards, achieving both better performance and significant speedups

Abstract: Although diffusion models have achieved strong results in decision-making tasks, their slow inference speed remains a key limitation. While consistency models offer a potential solution, existing applications to decision-making either struggle with suboptimal demonstrations under behavior cloning or rely on complex concurrent training of multiple networks under the actor-critic framework. In this work, we propose a novel approach to consistency distillation for offline reinforcement learning that directly incorporates reward optimization into the distillation process. Our method achieves single-step sampling while generating higher-reward action trajectories through decoupled training and noise-free reward signals. Empirical evaluations on the Gym MuJoCo, FrankaKitchen, and long horizon planning benchmarks demonstrate that our approach can achieve a $9.7$% improvement over previous state-of-the-art while offering up to $142\times$ speedup over diffusion counterparts in inference time.

Published as a conference paper at ICLR 2026
ACCELERATING DIFFUSION PLANNERS IN OFFLINE RL VIA REWARD-AWARD CONSISTENCY TRAJECTORY DIS- TILLATION
Xintong Duan"
Yutong He Jeff Schneider
Fahim Tajwar
Ruslan Salakhutdinov
J. Zico Kolter
Cameegie Mellon University
(xintongd, yutonghe, fta atajwaz, rsalakhu, zkottr, jefff) cmu, edu
ABSTRACT
Although diffusion models have achieved strong results in decision decision-making tasks, their slow inference speed remains a key limitation. While consistency models offer a potential solution, existing applications to decision-making either struggle with suboptimal demonstrations under behavior cloning or rely on complex concurrent training of multiple networks under the actor-critic framework. In this work, we propose novel approach to consistency distillation for offline reinforcement learning that directly incorporates reward optimization into the distillation process. Our method achieves single-step sampling while generating higher- action trajectories through decoupled training and noise-free reward signals. Empirical evaluations on the Gym MuJoCo, FrankaKitchen, and long horizon planning benchmarks demonstrate that approach can achieve 9.7% improvement over previous state-of-the-ar while offering to 142x speedup over diffusion counterparts in inference time.
INTRODUCTION
Recent advances in diffusion models have demonstrated their remarkable capabilities across various
domains (Song et al. 2020a Karras et al. 2022 Liu et al.
2023 Chi et al.
2023 Janner et al. 20222.
including decision-makee making tasks in reinforcement learning (RI,) ). These models excel particularly in
capturing multi-modal behavior patterns (Janner et al.) 2022;
0222:/C Chi et al. 2023 Ajay et al. 2022) and
achieving strong out-t--diitributiin generalization (Duan al. 2025 Block al. 2023), making
them powerful tools for complex docision-making scenarios. However, their practical deployment faces significant challenge: the computational overhead of the iterative sampling procedures, which requires numerous denoising steps to generate high-quality outputs.
To address this limitation, various diffusion acceleration techniques have been proposed, including ordinary or stochastic differential equations (ODE or SDF) solvers with flexible step sizes (Song
al.
2020a Lu al, 2022 Karras et al., (002), sampling step distillation (Song et al. 2023 Kim
al. 2023) and improved noise schedules and parametrizations (Salimans & Ho 2022, Song &
Dhariwal 2024). particular, consistency distillation (Song et al.) 2023) and consistency trajectory models (Kim et al. 2023) have emerged
of the most promising solutions for image generation, in which a many-step diffusion model serves as a teacher to train a student consistency model that achieves comparable performance while enabling faster sampling through single-step few-step
generation process.
This breakthrough has sparked considerable interest in applying consistency-base distillation decision-making making tasks. However, current applications either adopt a behavior cloning approach (Lu al. 2024 Prasad et al. 2024 Wang et 2024 integrate few-step diffusion based samplers in actor-critic frameworks (Chen et al. 2023 Ding & Jin 2023 Li et al. 2024). While promising, these approaches face inherent challenges: behavior cloning performs well only with expert demonstrations but struggles with suboptimal data (e.g., median-qualin replay buffers), while actor-critic methods
"Equal contribution

---

Published as a conference paper at ICLR 2026
Reward Gradient
Maximize Reward Function
Distillation Gradient
P(a|s)
Reward
M: c
Reward
Loss
Higher Reward
Predicted
State
Student Prediction
action
Noise
Match Teacher ODE
Student
Teacher
Teacher
Dg
Prediction
DSM CTM Loss
Backward PFODE
Xy-1
Figure 1: Overview of Reward Aware Consistency Trajectory Distillation (RACTD). We incorporate reward guidance with consistency trajectory distillation to train student model that generate actions with high rewards with only one denoising step.
require concurrent optimization of multiple networks from scratch with sensitive hyperparameters, leading to training complexity, instability, and computational overhead.
This raises an important question: can we develop more effective approach to consistency-hased distillation that tackles the two fundamental challenges introduced previous framework: suboptimal data and the concurrent training of actor critic networks for offline RL? We address this challenge by introducing a novel method that directly incorporates reward optimization into the consistency trajectory distillation process. Our approach builds existing infrastructure, leveraging pre-traan mr-itrined diffusion planner as teacher and a standalone reward model to augment the standard consistency trajectory distillation (Kim et al., 2023) with an explicit reward objective for single-step sampling student. The vanilla consistency trajectory distillation helps the student planner distill the diverse behavior modes learned by the teacher from mixed-quality offline data into single-step sampling model. The additional reward objective acts as mode selection mechanism: effectively steering the student planner toward selecting high reward modes from the multi-modal distributions captured by the teacher diffusion planner.
Our decoupled training framework offers two additional key advantages: training simplicity and flexibility. By generating clean action trajectories through single-step diffusion sampling, our approach removes the need for noise-aware reward models required for classifier-guided sampling from multi-step diffusion models (Janner et al.) 2022), which also suffer from reduced accuracy at high noise levels. Furthermore, the reward model is trained independently from the teacher diffusion planner and the distillation process, avoiding the tight coupling that arises from concurrent multi- network training present actor-critic methods. By integrating reward optimization directly into the single-step distillation process, method simultaneously enables efficient generation and achieves superior performance, while maintaining straightforward training procedures.
We demonstrate the performance and sampling efficiency of our RACTD the suboptimal hetero- geneous D4RI. Gym-MuloCo and D4RL. FrankaKitchen benchmark and challenging long-horizon planning task Maze2d (Fu et al., 2020). Our method demonstrates both superior performance and substantial sampling efficiency compared to existing approaches, achieving an 9.7% improvement compared existing state-of-the-arr (SOTA) while maintaining a 142-fold reduction in sampling time inherited from the consistency trajectory distillation.
Our contributions include: (1) We propose novel reward reward-aware aware consistency trajectory distillation method for offline RL that enables single-step action trajectory generation while achieving superior performance, (2) We demonstrate that our approach enables decoupled training without the com- plexity of concurrent multi-network optimization or noise aware reward model training, (3) Through comprehensive experiments on multi-modal suboptimal dataset and long-horizon planning tasks, we show that our method achieves 9. 9.7% improvement over prior SOTA while achieving up to 142x speedup by leveraging consistency trajectory distillation. We will publicly release our code upon acceptance of this paper.
2

---

Published as a conference paper at ICLR 2026
2 BACKGROUND
2.1 PROBLEM SETTING
In this paper consider the classic setting of offline reinforcement learning, where the goal is to learn a policy и to generate actions that maximize the expected cumulative discounted reward in Markov decision process (MDP). A MDP is defined by the tuple (8,A,7 P, R, where S is the set of possible states 8 E S, A the set of actions a E A. P(s 8, a) is the transition dynamics, R(s, a) reward function, and 7 E [0, 1] discount factor. In offline RL, we further assume that the agent can no longer interact with the environment and is restricted to learning from a size M static dataset D {T, J where (su, ag. ro. $1, 01 r....., SH. aH (H,TH) represents rollout of episode horizon H collected by following behavior policy TB.
.......,.............
Mathematically, we want to find policy कt" that
arg max E
Σ n=U
y" R( an)
(I)
subject to the constraint that all policy evaluation and improvement must rely on D alone.
2.2 DIFFUSION MODELS
Diffusion models generate data by learning to reverse a gradual noise corruption process applied to training examples. Given a clean data sample zy. we define 2, for E [0, T] as increasingly noisy versions of øg. The forward (or noising) process is commonly formulated an Itô stochastic differential equation (SDE):
dæ [ax,, t)dt + g(t)du:
where w is a standard Wiener process. As t approaches the final timestep T, the distribution of ær converges known prior distribution, typically Gaussian. At inference time, the model reverses this corruption process by following the corresponding reverse-time SDE, which depends on the score function V. log (ar). In practice, this score function approximated by a denoiser network Da, enabling iterative denoising from ØT back to gg. An alternative, deterministic interpretation of the reverse process is given by the probability flow ODE (PFODE):
dx f(x, t) (t) Va log (a:) dt
which preserves the same marginal distribution (æ) as the reverse SDE at each timestep t. This ODE formulation often enables more efficient sampling through larger or adaptive step sizes without significantly compromising the sample quality.
EDM (Krrras al.) 2022) refine both the forward and reverse processes through improved noise parameterization and training objectives. Concretely, they reparametrize the denoising score matching (DSM) loss so that the denoiser network learns to predict scaled version the clean data:
CEDM
E F, a ( [d (oo (x,,, t))]
where d is distance metric in the clean data space. In this paper, we train an EDM model as the teacher using the pseudo huber loss as d following Prasad et al. (2024). At inference time, EDM solves the associated PFODE with 2nd-order Heun solver.
2.3 CONSISTENCY TRAJECTORY DISTILLATION
The iterative nature of the diffusion sampling process introduces significant computational overhead. Among various acceleration techniques proposed, consistency distillation (Song al. (0023) has emerged as particularly effective approach. The core idea is to train a student model that can emulate the many-step denoising process of teacher diffusion model a single step.
Building upon this framework, Kim et al. (2023) introduced Consistency Trajectory Models (CTM). Instead of learning only the end-to-end mapping from noise to clean samples, CTM learns to predict across arbitrary time intervals in the diffusion process. Specifically, given three arbitrary timesteps 0 s ls < u < t < T, CTM aims to align two different paths to predict σ4; (1) direct prediction from time to & using the student model, and (2) a two-stage prediction that first uses numerical solver (e.g.. Heun) with the teacher model to predict from time t to u, and then uses the student model to predict from time to k.
3

---

Published as a conference paper at ICLR 2026
Student
Stopgrad Student
Teacher
L Reward
GolX, (,0)
210
(xgaa (xl00, k.0)
Solven I
(n
Reward
DSM
X
(las)
No
Backward PFODE
LeTM
(g (oooo&RRkkk) k,0)
G.(X,.111 (.
Figure 2: Visualization of CTM loss, DSM loss and reward loss.
Since the distance metric is defined on the clean data space and may not be well-defined in the noisy data space, in practice we further map all predictions to time using the student model. Formally, denote the student model as Gg, the CTM loss is defined as:
LсTM E odl ( k, 0), (x/ k 4,0)]]
(5)
where Gø(æ,, t, u) represents the student prediction from time t to given noisy sample æ, at time t, sg) represents stop-grade gradient student parameters and
$ Go t, k),
Fk
(t,u) (g) (Solver( t, us (),, u, k))
(6)
Here Solver(xt, t, u; σ) is the numerical solver result from time t to using the teacher model D given noisy sample x2 at time t.
In addition to the CTM loss. Kim et al. 2023) also incorporates the DSM loss to enforce the generated samples to be close to the training data. The DSM loss for the student model is the same as the one
for EDM in Equation
LDSM = E,,,0,.. æe [d( Go (a(,, (x,, (,
CTM loss together with DSM loss constitutes consistency trajectory distillation (CTD), resembling the model in Kim et al. (2023) without the optimal GAN loss. In Figure 2 we provide visualization of these objectives on a PFODE trajectory. After the distillation, the student model can perform "anytime-to-annting jumps along the PFODE trajectory. One-step sampling
then be achieved by calculating atT) (T T, 0). This one-step sampling ability unlocks significant speedups and many potential design choices that are beneficial in RI,, which we will elaborate next section.
METHOD
3.1 MOTIVATION AND INTUITION
Diffusion models and their consistency-based counterparts have demonstrated promising results in decision-making tasks, particularly in capturing multimodal behavior patterns (Janner al. 2022
Chi et al. 2023 Ajay et al. 2022). Common recipes for using these models in decision-making generally fall one of the three paradigms: (1) training a diffusion
consistency planner on expert demonstrations via behavior cloning and deploying it directly as policy (Chi et al. 2023
Prasad et 2024); (2) integrating diffusion or consistency models into actor-critic frameworks for
RL tasks (Wang et al. 2022 Hansen-Estruch et al.) 2023 Ding & Jin 2023) ; (3) using reward
agnostic diffusion planner with guided diffusion sampling for RI. tasks (Janner et al. 2022 Ajay
ct al. 2022).
While behavior cloning pipelines perform well when trained on expert demonstrations, they often struggle with suboptimal datasets (e.g., medium-replay buffers) collected from diverse behavior policies. Such datasets typically exhibit complex multimodal behavior patterns, where only some modes lead to high rewards. Although could potentially rejection sampling filter out low-reward training data, this approach becomes prohibitively sample inefficient, particularly as the quality of the training data deteriorates. On the other hand, generate high reward actions, actor-critic
4

---

Published as a conference paper at ICLR 2026
approaches require concurrent training of multiple neural networks with sensitive hyperparameters. Finally, guided diffusion sampling necessitates training awire reward models and multi-step sampling, which could be detrimental for time-sensitive decision decision-making making tasks like self-driving.
So how can we better leverage potentially suboptimal datasets to design a ceatusioo--aaseee single- step sampling model with simple training procedure? Our key idea is to utilize the multimodal information captured by the teacher diffusion planner and encourage the student diffusion planner favor the high reward modes. We achieve this by incorporating reward objective directly in the consistency distillation process. Since our student model can achieve single-step denoising, we can incorporate reward model trained in the clean sample space and avoid the multi-step reward optimization for diffusion models.
3.2 MODELING ACTION SEQUENCES
When applying diffusion models RL, several modeling choices are available: modeling actions (as a policy). modeling rollouts (as a planner), or modeling state transitions (as world model). Following Chi (202)). we adopt a planner approach and model a fixed-length sequence of future actions conditioned on a fixed-length sequence of observed states. This formulation ensures consecutive actions form coherent sequences, and reduces the chances of generating idle actions.
Formally, let s, = {sn-h, Sn-h- 1............................... s,} denote a length-h sequence of past states at rollout time n, and an = {an, an+1, (n+a] represent length-c sequence of future actions. Both the teacher and the student learn to model the conditional distribution p(@, ).). In the context of diffusion notations, x = a, Sn,
During exccution, we can either execute only the first predicted action a, in the environment before replanning at the next step, or follow the entire predicted sequence of actions at once.
3.3 REWARD- CONSISTENCY TRAJECTORY DISTIILLATION
Based on action trajectory formulation, optimization into the consistency trajectory distillation process.
present approach to integrating reward
Let Ro be pre-trained differentiable return-to-go network (i.e. reward model) that takes the state Sn and action an at rollout time n as inputs and predicts the future discounted cumulative reward fn = =j= r We implement this reward model with four ConvBlocks used by previous heward-guided sampling method (Janner et al., 2022), followed by a Linear layer to map output to the correct dimension (Appendix B When the student model generates prediction
an 8, = 2 = Go (wT, T, 0), we extract the action at time n, denoted as ån, from the predicted sequence and pass it along with 8, to the frozen reward model RyG. estimate A.. The goal of our reward-aware training is to maximize the estimated discounted cumulative reward. Mathematically, the reward objective is defined as
Crwward = -Ry (8n, ān)
The final loss for reward-award consistency trajectory distillation (RACTD) combines all three terms:
= == aLerM 3LDSM + LRewardd
where a, B, and a are hyperparameters balance different loss terms. This objective bears resem- blance to established offline RI. algorithms like off-policy deterministic policy gradient (Silver et al., 2014) (Appendix G). suggesting our approach builds upon well recognized theoretical foundations.
3.4 DECOUPLED TRAINING
A key advantage of our method of combining reinforcement learning, diffusion models, and con- sistency distillation the ability to support fully decoupled training all components. Traditional actor-critic frameworks, which commonly used incorporate diffusion models into reinforce- ment learning, require simultancous training of both the actor and critic networks from scratch. This concurrent optimization presents considerable challenges, often demanding extensive hyperparameter tuning and careful balancing of different learning objectives.
Guided diffusion sampling, as proposed in Janner al. (2022), offers alternative approach by taking inspiration from classifier guided diffusion (Dhariwal & Nichol| 2021] [Song et al. (020b). However, since these classifiers (i.e. reward models) must evaluate noisy states at every denoising step
5

---

Published as a conference paper at ICLR 2026
to provide gradient guidance, they require noise-aware training and thus cannot be separately trained from the diffusion model. Also, predicting the correct reward from highly corrupted input could be very challenging, which can lead to inaccurate guidance that accumulates during its multi-step sampling process.
Our method, in contrast, fully leverages the advantages of single-step denoising models by operating entirely in the noise-free state-action space. This design choice enables the reward model to provide stable and effective signals without requiring noise-aware training. Importantly, the reward model can be pre-trained completely decoupled from the teacher model and distillation process. This separation not only simplifies the training process but also allows for flexible integration of different reward models using the same teacher model.
3.5 REWARD OBJECTIVE AS MODE SELECTION
In offline RL, models often have to learn from datasets containing behaviors of varying quality, While diffusion models excel at capturing these diverse behavioral modes, they inherently lack the ability to differentiate between actions that lead to high versus low rewards. Our RACTD addresses this limitation by transforming the award-agnostie teacher diffusion sampling distribution into one that preferentially samples from high-reward modes. We empirically verify this through comparative analysis using the D4RL peopeer-ediumm-cee mixture of expert demonstrations and suboptimal rollouts from a partially trained policy.
dataset Fu et al. (2020), which contains an equal
Figure 3 illustrates the reward distributions of roll- outs sampled from three models: the unconditioned teacher, unconditioned student, and RACTD. The dataset (grey) exhibits two distinct modes correspond- ing medium-quality and expert rollouts. The uncon- ditioned teacher model (blue) accurately captures this bimodal distribution, and the unconditioned student model (orange) faithfully replicates it. contrast, RACTD (green) concentrates its samples the higher-reward mode, demonstrating that our reward guidance effectively identifies and selects optimal be- haviors from the teacher's multi-modal distribution. We also include the discussion between sample diver- sity and mode selection in Appendix
uncoed
RACTO Student
Datastt
Figure 3: The reward distribution of the D4RL hopper-mediummexxer dataset and 100 rollouts from unconditioned teacher, an unconditioned student, and RACTD.
4 EXPERIMENT
In this section, we conduct experiments to demonstrate: (1) the effectness of our RACTD in identify- ing high reward behavior patterns among multimodal mix quality data, (2) the expressiveness of our single step model compared to its multistep counterpart complex high dimensional long horizon planning tasks, and (3) the speed-up achieved over the teacher model and existing policy- diffusion models inherited from appropriate application of CTD.
4.1 OFFLINE RL
Baselines We compare our approach against comprehensive set of baselines, including vanilla
behavior cloning (BC) and Consistency Policy (Consistency BC) (Ding & Jin (002)); model-free
RL algorithms CQL (Kumar et al. 2020) and IQL (Kostrikov al. 202 1); model-based algorithms
Trajectory Transformer (TT) (Janner et al. 2021), MOPO Yuu et al. 2020). MOReL (Kidambi et al.
2020), MBOP (Argenson & Dulac-Arold 2020); autoregressive model Decision Transformer (DT)
Chen et al. (202)); diffusion based planner Diffuser
DiffuserDanner Janner et al. (2022); and diffusion/flow: blasedewwwwdddddddd
actor-critic methods Diffusion QL (Wang al. 2022 Consistency AC (Ding Jin 2023). and Flow
Q-learning (Park et al. 2025).
Setup We evaluate our method and the baselines on D4RL Gym-MuJoCo and FrankaKitchen benchmark (Fu et al. 2020). Gym-MuloCo is popular offline RL benchmark that contains three environments (hopper, walker2d, halfcheetah) mixtures of varying quality data (medium-replas medium, medium-expent expert). The FrankaKitchen dataset features 9-DoF robotic arm performing multi- task manipulation across four kitchen objects (microwave, kettle, light, cabinets) with mixed- -exxity demonstrations (partial, mixed).

---

Published as a conference paper at ICLR 2026
Following the setups in the prior works, report the performance of both online and offline model selection if available in their original paper. Online model selection reports the best rest-performing checkpoint observed during training, while offline model selection reports the performance of the last training epoch. Results for non-diffusion-basee models and Diffuser are taken from Janner et al. (2022), and results for Diffusion QI. and Consistency AC/BC are sourced from their respective papers. The results for RACTD are reported as the mean and standard error over 100 planning seeds. We use past observation length h and prediction horizon с 16 and closed-loop planning in all experiments following the same setup in previous diffusion planners (Chi et al. 2023 Prasad et al.
2024) and include the ablations on observation and planning horizons Appendix E.6 The best score is emphasized in bold and the second-best is underlined.
Table I: (Offline RL: MymmM Performance and sampling efficiency (NFE: Number of Function Evaluations) of RACTD and a variety of baselines on the D4RL Gym-MuJoCo benchmark.
Medium Expert
Medium
Medium Replay
Avg NEE
HallCheetah Hopper Walker2d HalfCheelah Hopper Walker2d HalfCheetah Hopper Walker2d
Offline model selection
BC
55.2
52.5
107.5
42.6
52.9
75.3
36.6
18.1
26.0
51.9
CQL
91.6
105.1
105.8
41.0
58.5
72.5
15.5
95.0
17.2
77.6
JQL
86.7
91.5
109.6
47.1
66.3
78.3
44.2
94.7
73.9
77.0
86.8
107.6
108.1
42.6
67.6
74.0
36.6
82.7
66.6
71.7
TT
95.0
110.0
101.9
46.9
61.1
79.0
41.9
91.5
82.6
78.9
MOPO
63.3
23.7
11.6
42.3
28.0
17.8
53.1
67.5
39.0
12.1
MOReL
53.3
108.7
95.6
42.1
95.4
77.8
40.2
93.6
49.8
72.9
MBOP
105.9
55.1
70.2
41.6
18.8
11.0
12.3
12.4
9.7
47.8
Diffusion QL
96.8 #0.3
111.1
+1.3
110.1
+0.3
51.1 #0.5
90.5
87.0
#0.9
47.8 #0.a
1013 3
05.5
#1.5
87.9
Consistency AC
84.3 1:.1
100.4 13.5
110.4
10.7
69.1 10.7
80.7
410.5
83.1 20.3
58.7 aan
90.7 20.5
70.5
#3.6
85.1
Consistency
32.7 #1.2
90.6 =0. 110.4
31.0 +0.4
71.7
83.1
34.4
99.7
13.3
Diffiser
$8.9 10.3
103.3
106.9 20-2
42.8 10.33
74.3
21-4 79.6 =0.55
37.7 20.5
93.8 40.4
70.6
#16
77.5
RACTD(OUrS)
68.5 #2.1
120.2
6
122.3 #0.3
56.6 +0.6
102.4 #2.4 112.8 =1.
51.4
108.3 #1.1
105.2
#1.
96.4
Online model selection
Diffusion QL
97.2 20.4
112.3 s 111.2 20.
51.5 t:0.3
96.6 23.4 87.8 20.
48.3 20.2
102 20.4
98.0
19.5
89.3
Consistency AC
69.2
#3.5
106.0
+1.
111.6 #0.7
71.9 #0.8
99.7 +2.3
81.1 #0.3
62.7 #0.0
100,1=0,0
83,0
#1.5
89.8
Consistency BC
39.6
$3.4
96.8
111.6
#0.7
46.2 +0.4
78.3
22.8
84.1
45.4 d0.7
100.4
20.8
80.8
#34
75.9
RACTD(ÓUrs)
95.9 #15
129.0
#1.3
122.3
+0.3
59.3 #0.2
115.5 #1.7
118.8
=0.3
57.9 #1.0
109.5
#0.3
105.2
#1.
101.5
Results As shown in Table 1 RACTD achieves the highest average score by substantial margin and best or second-best performance on 8/9 tasks in Gym-MuJoCo, with the only exception being a medium-expeett dataset where reward guidance is less beneficial. RACTD is also the only planning- based method that achieves single-step sampling compared to consistency model based actor-critic methods (which require double sampling steps) and other diffusion-based planners (which require 20x more sampling steps).
As shown in Table 2 RACTD consistently out- performs its counterpart that incorporates an actor-critic framework for consistency models (Consistency AC) or flow-based models (Flow q-learning) across all dataset qualities. On sub- optimal partial and mixed datasets where reward guidance crucial, our method achieves near- equivalent performance to multi-step diffusion approaches with 5x fewer sampling steps. These findings highlight the effectiveness of our reward- aware distillation framework, which combines decoupled training, noise-free reward modeling, and single-step sampling efficiency.
Table 2: (Oftline RL: FrankaKitchen) Performance and sampling efficiency (NFE) of RACTD and diffu- sion based baselines D4RL FrankaKitchen bench- mark. Each cell has values: the first is offline model selection and the second (in brackets) is online model selection.
Methed kitchen-partial
kitchen-mixed
Asg
NFEL
Diffusion Qr. 68.5 a (63.7 2t.0 62.6 2ha (66.6 =h.h 61.6/65.15
Consistency AC 38 a1s (39.8 424; 45.8 a15 (16.7 ( usa
42.0 (2.5)
Consistency BC: 22.6 4& (23.8 as a5 23.5 alla 2243 61
22.1/24.)
Flow Q-learning 67.5 424 (53.3 43., 46.0 sea (53.5
16. 86.8(53.4)
RACTI(OOuS) 18.0 415 (63.
(66
(0.061255
4.2 LONG HORIZON PLANNING
Next, we showcase the expressiveness of the single-step sampling model RACTD in complex high- dimensional long-horizon planning tasks. Previously, Diffuser has shown great potential in open-loop long-horizon planning, but requires significantly larger number of denoising steps compared to closed-loop planning like MuJoCo. We demonstrate that our model can achieve superior performance with single-step denoising process under the same problem formulation.
7

---

Published as a conference paper at ICLR 2026
Setup We test this ability on D4RL Maze2d (Fu et al. 2020), which is sparse reward Long-horizon planning task where an agent may take hundreds of steps to reach the goal in static environments.
Maze2d medium and large are also shared environments in other benchmarks like OGBench (Park) et al. 2024), where it is demonstrated to be even more challenging than higher-dimensional envi- ronments like AntMaze. Following the setup in Janner et al. (2022), use planning horizon 128, 256, 384 for U-Maze, Medium and Large respectively, all conditioned on two observations: the start and goal position. We perform open-loop planning by generating the entire state sequence followed by reverse dynamics model to infer all the actions from the predicted state sequence. The reward model returns if the current state reaches the goal and otherwise. The baseline results are reported from Janner et al. (2022) and RACTD results are reported as the mean and standard error of 100 planning seeds.
Results As shown in Table 3. both Diffuser and RACTD outperform previous model-free RL algorithms CQL, IQL, Flow Q-learning, and MPPI. Our approach surpasses the Diffuser baseline in all settings, highlighting its ability to effectively capture complex behavioral patterns and high- dimensional information in the training dataset. Notably, the planning dimension for this task (384 for the Large Maze) is substantially higher than that of MuJoCo tasks (16). As a result, Diffuser requires significantly more denoising steps (256 for the I Large Maze) compared to MuJoCo (20 steps). On the contrary, despite the increased task complexity, RACTD still only requires a single denoising step to achieve 11.6x performance boost, showcasing its expressiveness. Furthermore, maintaining the speedup inherited from CTD, our RACTD able to achieve significant performance improvement over CTD, especially on the more challenging long-horiza tasks.
Table 3: (Long-horizon planning) The performance of RACTD, Diffuser, Flow Q-learning and prior model-free algorithms in the Maze2D environment. Flow Q-learning close loop planning and Diffuser, RACTD are open loop planning.
Method
U-Maze
Medium
Large
Avg Score
Score
NFE
Time
Score
NFE
Time
Score
NFE
Time
MPPI CQL
33.2
10.2
5.1
16.2
5.7
12.5
7.7
IQL
47.4
31.9
58.6
17.0
Flow Q-learning
106.7 + 6.4
89.7 + 1.9
1
107.3 + 7.1
1
101.2
Diffuser
113.9 #9.1
61
1.661
121.5 =2.7
256
1.312
123.0 #:.44
256
5.508
119.5
123.4
0.029
119.8
0.047
127.1
1
123.4
RACTD (Ours)
125.7 19.8
1
0.029
130.8 =18
1
0.047
143.8 1:000
1
0.049
133.4
4.3 INFERENCE TIME COMPARISON
Beyond performance improvements, thanks to the single-step sampling ability from consistency trajectory distillation, our RACTD significantly accelerates diffusion-based models while maintaining their superior performance for decision-making diffusion models arises from the multiple function evaluations (NFEs) required by the iterative denoising process. By reducing the number of denoising steps to a single NFE, our approach achieves up roughly proportional to the number of denoising steps originally required.
tasks. The primary computational bottleneck in
Setup To evaluate sampling efficiency, we compare RACTD with different samplers, including DDPM (Ho et al., 2020), DDIM (Song et al.) 2020a), and EDM (Karras et al. (202)) using the same network architecture of Diffuser. Note that since Diffuser employs a different model architecture and generates future state-action pairs, its sampling time may also be influenced by these factors. Table 4 and Table present the wall clock sampling time and NFE for MuJoCo thopper-mediun-replay and Maze2d. All experiments conducted on one NVIDIA Tesla V100-SXM2-32GBB
teacher and student model. Additionally, we report the efficiency
Results In hopper-meddum-replly, RACTD achieved 20 reduction in NFEs and 43x speed-up compared to Diffuser. Additionally, our student model requires 80x fewer NFEs and samples 142x faster than the teacher model. In Maze2d, RACTD significantly accelerates computation compared to Diffuser, achieving approximately 57 x, 92x, and 114x speed-ups on Umaze, Medium and Large mazes, by reducing NFEs by a factor of 256 for the Medium and Large mazes.
8

---

Published as a conference paper at ICLR 2026
Table 4: Wall clock time and NFEs per action for different samplers and Diffuser on MuloCo
Table 5: We compare incorporating the reward
model in different stages of training on MuJoCo
hopper-medium medium medium-replay. Results are presented as the mean and standard error across 100 seeds.
Method
Time (s)
NFE
Score
Diffuser
0.644
20
93.6
Hopper
Unconditioned teacher
Reward- Aware teacher
DDPM
0.236
15
24.2
Meplam-Replay
DDIM
0,208
15
60.6
Unconditioned student
50.8 =0.3
96.2 #0,2
EDM (Teacher)
2.134
80
114.2
Reward-Aware student
109,5 10.3
96.0 10.3
RACTD (Ours, Student)
0.015
1
109.5
ABLATION STUDY
We evaluate the necessity of incorporating reward objectives during student versus teacher training and include extended analyses reward loss weights and multi-step sampling in Appendix E.3
5.1 IMPACT REWARD OBJECTIVE
To understand the unique advantages of incorporating reward objectives during distillation, we conduct systematic comparison across four model configurations: the baseline combination of an unconditioned teacher and student, a reward-aware teacher paired with an unconditioned student, fully reward-aware teacher-student pair, and our proposed RACTD, which combines an unconditioned teacher with a reward-awar student. The results for hopper-medium-reply and walker-medium is shown in Table and Table I2
Our analysis reveals that while incorporating reward objectives at any stage yields substantial improvements, optimal performance achieved through our RACTD framework, which combines unconditioned teacher with reward-aware student distillation. Our method allows the teacher to capture a comprehensive range of behavioral patterns, while enabling the student to selectively distill the most effective strategies. Although incorporating reward objectives in the teacher model also enhances performance, this approach risks discarding suboptimal behaviors that may be valuable in novel testing scenarios, potentially limiting the model's generalization capabilities.
6 RELATED WORK
Diffusion Models in Reinforcement Learning Diffusion models have emerged as a powerful approach for decision-makingg tasks in reinforcement learning Janner et al. (2022) y Ajay et al. (2022);
Wang et al. (2022); Hansen-Estruck et al. (2023); Chi et al. (2023). The integration of diffusion models into RL frameworks typically follows two main paradigms: actor-critic approaches, where diffusion models serve policy or value networks (Wang al. 2022 Mansen-Estrus et al. (2023)
and policy-based approaches, where diffuusion models directly generate action trajectories (astoriessJaneerr Janner et al.) (0022) ; Ajay al. (2022): ; Chi et al. (2023). While these methods have demonstrated impressive performance on standard RL benchmarks, their practical deployment is hindered by the slow sampling time inherent to vanilla diffusion policies based on DDPM Ho et al. (2020). This limitation poses particular challenges for speed-sensitvee real-world applications such robotics.
Accelerating Difftusion Model Sampling Various approaches have been proposed to accelerate the sampling process in diffusion models. One prominent direction leverages advanced ODE solvers
to reduce the number of required denoising steps (Song et al.) (2020a] Karras et al.) 2022 Lu et al.
2022). Another line of work explores knowledge distillation techniques Luhman & Luhman (2021): Kiim Salimans & Ho 2022); Berthelot et al. (2023): Kim et al. (2023); [Song et al. (2023). where student models learn take larger steps along the ODE trajectory. Particularly, consistency trajectory models
(Kim et al. (023) enable one-step sampling by learning anytime-to-ane to-ynytime jumps along the PFODE
trajectory, achieving performance that surpasses teacher models.
Consistency Models in Decision Making Consistency models have emerged as promising policy class for behavior cloning from expert demonstrations in robotics (Lu et al.) 2024 Prasad et al.. 2024: ( et al. 2024). In RL, several works have enhanced actor-critic methods by replacing traditional diffusion based value/policy networks with consistency models, showing faster inference
and training speed (Chen et al. 2023 Ding & Jin 2023
Li al. 2024). These approaches directly

---

Published as a conference paper at ICLR 2026
incorporate consistency loss (Song et al. 2023) into the value/policy network training, rather than distilling separate student model. Park et al. (2025) distilled a flow-based model into a single-step student but still requires concurrent network optimization under the actor critic framework. While Wang et al. made some initial attempts to apply consistency distillation in policy-based RL through classifier-free guidance and reverse dynamics, their approach requires two NFEs and under-performs the state-of-the-art. In contrast, RACTD is straightforwar approach of using a separate reward model and incorporating reward objective during student distillation, achieving superior performance with only one NFE.
7 CONCLUSION
In this work, we address the challenge of accelerating diffusion planners at sampling by introduc- ing reward-aware consistency trajectory distillation (RACTD), which predicts high-reward action trajectories in a single denoising step. RACTD uses pre-trained diffusion teacher planner and separately trained reward model, leveraging the teacher's ability to capture multi-modal distributions while prioritizing higher-reward modes to generate high-quality action trajectories from suboptimal training data. Its decoupled training approach avoids the complex concurrent optimization of multiple networks and enables the use of standalone, noise-free reward model. RACTD outperforms previous state-of- the-art by 9.7% while leveraging consistency trajectory distillation to accelerate its diffusion counterparts up to a factor of 142.
IMPACT STATEMENT
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
ACKNOWLEDGMENT
This work was partly supported by DSTA, ONR N000142312368 and ONR MURI N00014-25- 1-2116. FT was partially supported by the U.S. Army Futures Command under Contract No. W519TC-23-C-0033 during the project.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is con- ditional generative modeling you need for decision-making makiing? arXiv preprint arXiv:2211.15657 2022.
Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. arXiv preprint arXiv:2008.0555 05556, 2020.
David Berthelot, Amaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv:2303.042443 2023.
Adam Block, Ali Jadbabaie, Daniel Pfrommer, Max Simchowitz, and Russ Tedrake. Provable guarantees for generative behavior cloning: Bridging low-level stability and high-level behavior. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023. URL https /
net http://ww.nee//rrrun11-------/.- Ogwid
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in newral information processing systems, 10:1154841109, 2021.
Yuhui Chen, Haoran Li, and Dongbin Zhao. Boosting continuous control with consistency policy. arXiv preprint 10000:2310.06343 2023.
10

---

Published as a conference paper at ICLR 2026
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel. Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal Robotics Research, pp. 02783649241273668, 2023. 2023.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Zihan Ding and Chi Jin. Consistency models as a rich and efficient policy class for reinforcement learning. arXiv preprint arXiv:2309. 16984, 2023.
Xintong Duan, Yutong He, Fahim Tajwar, Wen-Tse Chen, Ruslan Salakhutdinov, and Jeff Schneider. State combinatorial generalization in decision making with conditional diffusion models. arXiv preprint arXiv:2501.13241 2025.
Justin Fu. Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 2020.
Philippe Hansen- -sttruh,, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as actor-critic method with diffusion policies. arXiv preprint 1eXiv:2304.1057. 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-685 6851, 2020.
Audrey Huang, Adam Block, Dylan Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz, Jor- dan T Ash, and Akshay Krishnamurthy, Self-improvement in language models: The sharpening mechanism. arXiv preprint arXiv:2412.0133 2024.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34: 14:11273-18 1286, 2021.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.099% 099991999999,. 2022.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 15:265655-65777 2022.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Ue- saka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279, 2023.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint urXiv:2110.06111 2021.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 33:1 179-1191,, 2020.
Haoran Li, Zhennan Jiang, Yuhui Chen, and Dongbin Zhao. Generalizing consistency policy to visual rl with prioritized proximal experience regularization. arXiv preprint 4rXiv:2410.00055, 2024.
Haohe Liu, Zehua Chen, Yuan, Xinhao Mei, Xubo Liu, Danilo Mandic, Wenwu Wang, and Mark D Plumbley. Audioldm: Text-to-audio generation with latent diffusion models. arXiv preprint arXiv:2301.12500 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.

---

Published as a conference paper at ICLR 2026
Guanxing Lu, Zifeng Gao, Tianxing Chen, Wenxun Dai, Ziwei Wang, and Yansong Tang. Maniem: Real-time 3d diffusion policy via consistency model for robotic manipulation. arXiv preprint NrXiv:2406.01586 2024.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02332 2021.
Diganta Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint
arXiv:/908. 08681, 2019.
Seohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine. Ogbench: Benchmarking offline goal-conditioned rl. arXiv preprint arXiv:2410.20092, 2024.
Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. arXiv preprint arXiv:250..02538, 2025.
Aaditya Prasad, Kevin Lin, Jimmy Wu, Lingi Zhou, and Jeannette Bohg. Consistency policy: Accelerated visuomotor policies via consistency distillation. arXiv preprint arX1v:2405.005003 2024.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffiusion models. In International Conference on Learning Representations, 2022. URL [tLL/attpp:///peeeevies net/ net//Eorum.Ice-TTlXll0ass
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pp. 387-395. Pmlr, 2014.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020a.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In The Twelfth International Conference on Learning Representations, 2024. URL https
7/openreview.o net.: forum? 1iw.ummww-wwNNeywwwddc
Yang Song, Jascha
Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1.Xiv:2011.13455: 2020b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01449, 2023.
Guanquan Wang, Takuya Hiraoka, and Yoshimasa Tsuruoka. Planning with consistency models for model-based offline reinforcement learning. In Deployable RL: From Research to Practice Reinforcement Learning Conference 2024.
Zhendong Wang. Jonathan Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint 12XXi:22008.61993 2022.
Zhendong Wang, Zhaoshuo Li, Ajay Mandlekar, Zhenjia Xu, Jiaojiao Fan, Yashraj Narang, Linxi Fan, Yuke Zhu, Yogesh Balaji, Mingyuan Zhou, et al. One-step diffusson policy: Fast visuomotor policies via diffusion distillation. arXiv preprint arXiv:2410.21257, 2024, 2024.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Newral Information Processing Systems, 13:14129-14142 2020.
12

---

Published as a conference paper at ICLR 2026
A MODEL ARCHITECTURE
We follow the model architecture used in Chi et al. (2023) and Prasad et al. (2024) and continue to use the ID temporal CNN layer for our Unet and FiLM layers to process the conditioning information.
A.1 MODEL SIZES FOR MAZE2D
Model parameters for teacher models of Umaze, Medium, and Large Maze are shown below in Table The student model has the architecture the teacher model except it also takes extra variable of denoising timestep as conditioning.
Parameter
Umaze
Medium
Large
diffusion ceciiionntte____ee step_embed_dine dim
256
256 [512, 1024, 2048]
256
down dims
[256, 512, 1024]
[256, 512, 1024, 2048]
horizon kernel size
128
256
384
5
Table 6: Model parameters for Unet in Maze2d.
A.2 MODEL SIZES FOR GYM-MUJOCO
Unet parameters for teacher and student models in the MuJoCo task are shown below in Table 7 Model sizes are fixed through 9 different environments.
Parameter
MuJoCo
diffusion_ cetooonstttpeememeeeeeee tie__embed_dee dim
128 [512, 1024, 2048]
down dims
horizon kernel size
16
5
Table 7: Model parameters for Unet in MuJoCo.
B REWARD MODEL PARAMETERS
We follow the setup in Janner et al. (2022), where we use Linear layers and Mish layers Misra (0199) for the reward model. We stack Blocklocks for the reward model, each with dimension shown in Table [4] Each ConvBlock consists of Conv ID layer, groupnorm, and a Mish layer. Reward model architecture remains the same across all MuloCo benchmarks.
C TRAINING DETAILS
Our models are trained D4RL dataset (Fu et al. 2020), which released under Apache-2.0 license.
C.I NOISE SCHEDULER
We follow the setup in Prasad et al. (2024) and use EDM noise scheduler (Karras et al. 2022) for the teacher model. Particularly, discretization bins are chosen to be 80.
Student model used CTM scheduler (Kim et al. 2023) also with discretization bins of 80.
C.2 WEIGHT OF DIFFERENT LOSSES
The weights for CTM, DSM, and Reward loss we used in the experiment are shown below in Table 9 Generally, if the training dataset includes more expert samples, the weight for reward guidance
13

---

Published as a conference paper at ICLR 2026
Parameter
MuJoCo
layer dimensions
[32, 64, 128, 256]
Table 8: Reward model parameters in MuJoCo.
Parameter
CTM
DSM
Reward
hopper-medium-rellay
1.0
1.0
0.8
hopper-medium
1.0
1.0
3.0
hopper-medium-exper
1.0
1.0
0.0
1.0
1.0
1.0
malker2d-mediim
1.0
1.0
0.4
walker2d-mediue expert
medium-experr
1.0
1.0
0.2
1.0
halfcheetah-medium
1.0
1.0
0.5
hallchectah-mmdiim-eexe medium-experrt
1.0
1.0
0.0
Table 9: Weights for CTM, DSM, and Reward loss used in MuJoCo benchmark.
smaller. A reward weight of 0.0 resembles behavior cloning with consistency trajectory distillation. We found that as long as the loss weights are chosen to keep the individual loss terms within the same order of magnitude, the model will achieve reasonable performance.
D MORE RESULTS
D.1 EXPERT AND RANDOM DATASET PERFORMANCE FOR GYM-MUJOCo
We include the performance on expert and random datasets along with a comparison between dataset quality, ranging from random expert, with another diffusion-based planning method Diffuser (Janner et al. 2022). Under the trajectory modeling setup, the random dataset provides little valid trajectory and the expert dataset will degrade the problem to behavior cloning. Nevertheless, our method outperforms Diffuser in all settings tested.
hopper
random
medium-replay
medium
medium-expert
expert
NFE
RACTD
31.8+0.0
104.942.1 70.6+1.6
121.0+0.5 79.6+0.6
129.0+1.3 106.9+0.2
137.
Diffuser
6.7+0.1
110.3+0.1
20
Table 10: Performance of Diffuser and RACTD on Gym-MuJoCo random and expert dataset.
D.2 EXPERT DATASET PERFORMANCE FOR FRANKA KITCHEN
Since kitchen-complete consists of pure expert demonstrations, the problem degrades to behavior cloning and reward guidance becomes less effective.
Table 11: (Offline RL: FrankaKitchen) Performance and sampling efficiency (NFE) of RACTD and diffusion based baselines on FrankaKitchen complete. Each cell has two values: the first is offline model selection and the second (in brackets) online model selection.
Method
kitchen-complete
NFE.
Diffusion QL
84.0 =7.4 (84.5 1
5
Consistency AC
51.9 46.0 (67.6 +2.7)
2
Consistency BC
45.2 15.0 (50.9 13.6)
2
RACTD(OUrs)
56.3 (58.1 18.3)
14

---

Published as a conference paper at ICLR 2026
E MORE ABLATIONS
E. REWARD OBJECTIVE FOR
Table 12: We compare incorporating the reward model in different stages of training MuloCo walker-medium. Results are presented as the mean and standard error across 100 seeds.
Walker Medium
Unconditioned teacher
Reward-Aware teacher
Unconditioned student
93.3 + 1.8
97.0 1100
Reward-Aware student
118.8 #0.3
94.5 #2.6
E.2 COMPARING WITH FAST SAMPLING ALGORITIMS FOR MAZE2D
Table 13: We compare fast sampling algorithms DDIM and CTD, along with our method RACTD on Mazeld environment. DDIM performs fast sampling based on a DDPM model, while CTD and RACTD (ours) distill an EDM teacher. The number of function evaluations (NFE) reflects the sampling speed of each algorithm. Results are reported as the mean and standard error over 100 random seeds.
U-Maze
Medium
Large
Method
NFE
Average Score
Score
Score
Score
DDPM
100
126.310.7
126.8 =3.0
144.8 14.9
132.6
DDIM
10
121.2 11.1
126.2
#2.8
143.1 #4.9
130.2
DDIM
1
3.5 #4.7
226
112.5
-
-1.5 #0.7
-0.2
EDM
80
125.4
10.6
120.1
14.2
149.0
10.5
131.5
CTD
1
123.4 11.0
119.8
14.1
127.1
16.4
123.4
RACTD (ours)
1
125.7
10.6
130.8
11.8
143.8
10.0
133.4
E.3 EFFECT OF REWARD OBJECTIVE WEIGHTS
The reward loss weight is hyperparameter our pipeline that impacts both training stability and performance. We plot the reward curve achieved over 200 epochs of student training with reward loss weights ranging from [0.3, 0.8, 1.5] on MuJoCo hopper-medium-repphe in Figure 4 The mean and standard error are reported across 20 rollouts from intermediate checkpoints.
80-
80-
50
75
320
125
150
125
703
spoch
Figure 4: Ablation reward objective weight.
With lower weights, increasing the weight leads to higher performance and relatively stable training. However, when the weight too high (e.g. 1.5 in this plot), evaluation initially increases but fluctuates
15

---

Published as a conference paper at ICLR 2026
training progresses. This occurs when the reward loss dominates DSM and CTM losses, resulting in unstable training.
E.4 DSM AND REWARD LOSS WEIGIITS
We include the ablation for different DSM and reward loss weights. The other two losses fixed at optimal when altering the target loss.
DSM weight
0.0
0.3
0.6
0.8
1.0
1.2
1.5
2.0
RACTD
0.0
23.8 + 1.0
35.6 6 2.2
105.2
1.8
108.3 # 1.1
100.0 + 2.3
104.9 = 1.5
93.6 1.66::1
3.4
Table 14: DSM loss weight ablation for RACTD in MuJoCo.
Reward weight
0.0
0.3
0.5
0.7
0.8
0.9
1.0
1.5
2,0
50.8 0.3
00.8 2.5
108.4
1.1
100.9
Table 15: Reward loss weight ablation for RACTD in MuloCo.
E.5 NUMBER OF SAMPLING STEPS
Since our student model is trained for anytime-to-anytime jumps, it naturally extends to multi- step denoising without additional training. Following the approach in [Song et al. (2023), given intermediate denoising timesteps 0 < ti < f2 < T, we first denoise from T to 0 as usual. We then add noise again to & and denoise it back to 0, and repeat this process for f2. This iterative refinement can enhance generation quality. We evaluate the student using 2, 3, and 4 denoising steps reported Table [6 Results show that multi-step sampling barely improves model performance.
Table 16: Inference time, NFE, and score comparison for student model multi-step sampling MuloCo
NFE
Time(s)
Score
1
0.0147
109.5 +0.3
2
0.0241
109,8
3
0.0377
108.7 40.1
4
0.0517
107.9 41.6
E.6 OBSERVATION AND PLANNING HORIZON
We compare the performance of incorporating different steps of past observations and predicting different numbers of future actions for MuloCo walker medium-replay here in Table 17 We can see that the performance be further improved by conditioning larger number of observations and a longer planning horizon.
Table 17: We compare the performance of using different numbers of past observations (h) and future actions (c) on MuloCo walker medium-replay. Results are presented as the mean and standard error across 100 seeds.
walker medium-reelay
Future actions c=16
Future actions c=8
Past observation h=1
105.2 +1.8
104.3 #1.4
Past observation h=4
108.0 #1.4
103.3 #2.2
F THE TRADE-OFF BETWEEN MODE SELECTION AND SAMPLE DIVERSITY
In this section, we include a discussion about the trade-off between mode selection induced by our reward-aware training and the sample diversity of the student. Naturally, favoring selected modes
16

---

Published as a conference paper at ICLR 2026
led generation with limited sample diversity summarized in Table 18 This trade-off between sample diversity and sample optimality observed in RACTD is similar to what has been seen in other generative domains (e.g., language model RLHF (Huang al. 2024). classifier-guidee
diffusion, conditional image generation), where preference alignment also often reduces sample diversity. In case, the reward model acts similarly to classifier model toward desirable behaviors and sacrificing some of the sample diversity by design.
alignment reward model, guiding the
Importantly, our decoupled framework allows the use of single, unconditioned teacher with strong generalization capabilities across tasks. For multi-task or unseen-task settings, different reward models can be trained per task, and corresponding student models can be distilled from the same teacher using different reward models.
Table 18: A summarization of the trade off between sample diversity and model performance.
Sample diversity
Performance
Sample time
Reward agnostic diffusion
High
Low
Slow
Reward aware diffusion
Low
High
Slow
Reward agnostic consistency distillation
High
Low
Fast
Reward aware consistency distillation
Low
High
Fast
G
THEORETICAL INSIGHTS
Here provide proof sketch of how reward-aware consistency trajectory distillation training resembles off-policy deterministic policy gradient (Silver et al. 2014), thus sharing the well-studied theoretical analysis of these methods for offline RL.
According to CTM (Kim et al. 2023) (Section 3), our student model generates a single-step prediction by
Gø T,0) = (x T.0) 2T + / Tu [6xx) [.. du
From Taylor expansion, we have Go(æT, T, T) = E[ [TT] + O(T)
Let μs Ea ar)) Then the gradient of reward model is
VgRiGo(27,,,.00 T,)) Vel(up) E.ooaa]] Voll(u))oor
For off- off-policy deterministic policy gradient[cite), we have
VaJa(uo) R E arpe
where p is the behavior policy.
VoJa(uo)
R
E sapp
Our method first samples state from the offline dataset collected by the behavior policy pa and asks the student model to predict actions a given s. Then, the reward model takes and a and predicts the corresponding value. Therefore R(Ha) is equivalent to Q"(s, alewp" a=))(())) Thus, our reward-aware student model training resembles off-policy deterministic policy gradient.
H LIMITATIONS AND FUTURE WORK
One limitation of our approach is the need to train three separate networks: the teacher, student, and reward model. Training the teacher can be time-consuming, achieving strong performance often
17

---

Published as a conference paper at ICLR 2026
requires higher number of denoising steps. Additionally, consistency trajectory distillation is prone to loss fluctuations, and incorporating a reward model into the distillation process may further amplify this instability. Future work will focus developing more stable and efficient training procedure, as well as exploring methods to integrate non-differentiable reward models into the framework.
18


Reviews:
[{"review_id": "u7Du4Ozk6y", "reviewer": "Reviewer_LYwS", "summary": "This work introduces Reward-Aware Consistency Trajectory Distillation (RACTD), which presents a novel application of consistency distillation to diffusion-based planning in reinforcement learning. The method trains a single-step student model to emulate a multi-step diffusion teacher, effectively accelerating inference while maintaining high performance. The authors incorporate a reward model alongside the consistency trajectory loss and denoising score matching loss within a decoupled training framework. Experimental results demonstrate that the proposed RACTD substantially outperforms state-of-the-art diffusion planning methods in both performance and inference speed.", "questions": "1. The authors claim that their proposed decoupled training outperforms the simultaneous training of actor and critic networks from scratch. However, there does not appear to be any experimental evidence supporting this statement. Could the authors provide an ablation study or empirical comparison demonstrating the benefits of decoupled training?\n\n2. In Table 2 and Table 11, Diffusion-QL achieves better performance than RACTD, although RACTD uses 5x fewer sampling steps as stated by the authors. In this case, why did the authors not compare all methods under the same number of sampling steps? Such a comparison would be more meaningful than simply showing that the baseline achieves the best performance, particularly in Table 11.\n\n3. The results in Table 1 on the D4RL benchmarks show that the proposed RACTD method significantly outperforms prior state-of-the-art approaches on average. What do the authors believe is the main factor behind this improvement? Is the performance gain primarily attributable to the application of consistency distillation?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper introduces a novel application of consistency distillation, resulting in a single-step model that effectively emulates multi-step diffusion processes, thereby significantly improving inference efficiency.\n\n2. The proposed RACTD demonstrates strong empirical performance, substantially outperforming prior state-of-the-art diffusion-based reinforcement learning methods, particularly on the D4RL benchmark.\n\n3. The paper presents comprehensive experiments and detailed implementation descriptions, accompanied by code availability, which enhances the credibility and reproducibility of the reported results.", "weaknesses": "1. The idea of the paper is interesting, and the contribution of accelerating the sampling stage through consistency distillation is clear. However, the work appears to rely heavily on existing consistency distillation techniques, with limited novelty beyond their direct application to diffusion-based planning.\n\n2. The contribution of the proposed reward-aware consistency trajectory distillation is somewhat unclear. The method appears to employ a standard reward model as an auxiliary loss applied to the student’s output, rather than introducing a fundamentally new formulation or architecture.", "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": "LpnKV0TReV", "reviewer": "Reviewer_7SAP", "summary": "The manuscript proposes a method termed RACTD aimed at accelerating diffusion-based components within offline reinforcement learning by reducing the number of function evaluations (NFE) while maintaining policy performance.\nThis work presents a novel approach to consistency distillation that directly incorporates reward optimization into the distillation process. The proposed method achieves single-step diffusion sampling while generating\nhigher-reward action trajectories through decoupled training and noise-free reward guidance.", "questions": "The theoretical contribution appears weak; please strengthen the theoretical analysis and add comprehensive comparisons with state-of-the-art methods (e.g., one-step diffusion/consistency baselines) to substantiate the claims.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Clear experimental scope and benchmarks: Evaluation on widely used D4RL Gym-MuJoCo and FrankaKitchen datasets enhances relevance and comparability, with both offline and online model selection reported where applicable.\n- Focus on sampling efficiency: Explicit reporting of NFE alongside performance indicates attention to practical efficiency, which is crucial for diffusion-based methods.", "weaknesses": "- Limited novelty and contribution: The core idea is to augment the distillation process with a cumulative reward maximization objective. This training pipeline has appeared in prior work (e.g., Flow Q-Learning), and the paper does not clearly isolate what is fundamentally new beyond this template.\n- Central claim lacks rigorous empirical validation: The paper emphasizes incorporating a reward objective directly into consistency distillation rather than optimizing via a critic (e.g., Q-values or advantages). This design choice requires thorough ablations and head-to-head comparisons against critic-based alternatives to substantiate the claim.\n- Motivation section is superficial: The insights in Section 3.1 are relatively basic and widely known; they would be more appropriate in the introduction. The main text should instead focus on deeper methodological details, theoretical justification, or non-trivial empirical findings.\n- Incomplete experimental comparisons: The evaluation is missing strong, direct baselines from state-of-the-art one-step diffusion/consistency methods.", "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": "xLpcrSLwGr", "reviewer": "Reviewer_ixyB", "summary": "This paper introduces Reward-Aware Consistency Trajectory Distillation (RACTD), a method that extends consistency trajectory models to the offline reinforcement learning setting by incorporating a reward objective into the distillation process. The approach allows a student model to generate high-reward action trajectories in a single denoising step, combining efficiency from consistency models with reward-guided behavior selection. Experiments across standard offline RL benchmarks demonstrate that RACTD improves policy quality and sampling efficiency compared to diffusion-based and actor-critic baselines by a large margin in many cases.", "questions": "- In Section 3.3 the reward model is frozen during training. How sensitive is the overall planner to the diversity and coverage of the dataset used to train this reward model? The paper emphasizes robustness to suboptimal data, but what happens when planned trajectories drift outside the dataset’s support, potentially leading to compounding errors in both the reward model and the planner?\n- How sensitive is your method to alpha, beta, and sigma in the RACTD loss? Do you perform ablations to determine this sensitivity?\n- In Figure 3, the diffusion-based methods appear to have one of their high-probability density modes lying outside the support of the high-reward mode in the training dataset. Am I misinterpreting this figure? Intuitively, the highest-probability regions for the unconditional diffusion-based methods should align with the high-reward mode of the training distribution. The rightward shift of the RACTD distribution is understandable due to the added reward objective, but why do the unconditional diffusion-based models also shift toward higher rewards than the dataset’s high-reward mode?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well written and was genuinely an enjoyable read. Most questions that arise while reading are quickly addressed by the text, and all of the necessary context to follow the technical discussion is included.\n- The motivation is strong, with clear explanations of the limitations of prior work such as the slow inference speed of diffusion models and the sensitivity of actor-critic frameworks to hyperparameters.\n- The introduction of decoupled training in Section 3.4 is a genuine strength, making the method simpler and more stable in practice. The idea of operating directly in the clean reward space rather than training noise-aware reward models is both elegant and effective.\n- The proposed method demonstrates substantial empirical gains over prior work, with strong improvements in performance and efficiency. \n- The inference-time acceleration that this method inherits from prior work makes the approach especially appealing for real-world decision-making tasks.", "weaknesses": "- Not a substantial weakness, but there is a typo between lines 202 and 203 (struggle vs. struggles)\n- In Table 3, the paper introduces CTD as a comparison point without tying this acronym to a particular method and without comparing CTD's results to anything else in the main text (unless I am mistaken). Presumably this is Consistency Trajectory Distillation. Reading back through the paper, it appears that CTD corresponds to the authors’ approach without the reward aware component, that is, Consistency Trajectory Model (CTM) [1] augmented with a denoising score matching (DSM) loss. However, this relationship is never made explicit. Clarifying whether CTD is simply CTM+DSM loss applied to offline reinforcement learning or a modified variant beyond the addition of the DSM loss is important for readers to properly interpret the results. \n- Based on the prior point, the novelty of the proposed method is therefore somewhat unclear when compared to the backbone it builds off of. The paper introduces Reward Aware Consistency Trajectory Distillation (RACTD), which extends CTD by incorporating a reward term. This alongside the incorporation of the DSM loss is novel and shows good results compared to prior methods in the literature. So the motivation for these additions is clear, as shown in Figure 3, but the relative benefit of the reward aware RACTD component remains mostly untested against the unconditional variants CTD and CTM. A broader ablation comparing CTM, CTD, and RACTD across all benchmarks (for example, Tables 1, 2, 4, and 5), alongside their relative inference times and sample quality, would better isolate the benefit of the reward awareness and DSM loss integration. \n- If CTD is effectively a reimplementation of CTM adapted to the offline RL setting with minimal modifications, then the faster inference times should be credited to prior work rather than presented as new contributions. This is backed by Table 3's comparison of CTD to RACTD where both methods have identical inference times. Applying existing techniques in new settings is a contribution that is noteworthy and part of this work, but the main contribution here would then lie in the reward aware adaptation, DSM loss, and application in a new domain with great score improvement results, rather than claiming the inference speedups already established by CTM. The acceleration results therefore represent a successful domain adaptation of a pre-existing backbone you built on top of, not a new methodological innovation, and should be framed as such in the paper (e.g. the leading sentence in Section 4.3 \"Beyond performance improvements, another major contribution of our work is significantly accelerating diffusion-based models for decision-making tasks.\" or the last sentence of the conclusion: \"RACTD outperforms previous state-of-the-art by 9.7% while accelerating its diffusion counterparts up to a factor of 142.\").\n- Given that the primary acceleration mechanism is inherited from Consistency Trajectory Models [1], the title “Accelerating Diffusion Planners in Offline RL via Reward-Aware Consistency Trajectory Distillation” may overstate the novelty of the acceleration component. A more accurate framing would emphasize the reward-aware extension rather than the inference-speed improvement itself, which seemingly follows directly from prior work.\n\n[1] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" arXiv preprint arXiv:2310.02279 (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": {}}]

Author rebuttal:

Decision: Accept (Poster)