Conformal Prediction for Uncertainty-Aware Planning with Diffusion Dynamics Model
Authors: Jiankai Sun
Yiqi Jiang
Jianing Qiu
Parth Talpur Nobel
Mykel Kochenderfer
Mac Schwager
Venue: NeurIPS 2023 poster (2023)
Keywords: Uncertainty
Conformal Prediction
Dynamics Model

Abstract: Robotic applications often involve working in environments that are uncertain, dynamic, and partially observable. Recently, diffusion models have been proposed for learning trajectory prediction models trained from expert demonstrations, which can be used for planning in robot tasks. Such models have demonstrated a strong ability to overcome challenges such as multi-modal action distributions, high-dimensional output spaces, and training instability. It is crucial to quantify the uncertainty of these dynamics models when using them for planning. In this paper, we quantify the uncertainty of diffusion dynamics models using Conformal Prediction (CP). Given a finite number of exchangeable expert trajectory examples (called the “calibration set”), we use CP to obtain a set in the trajectory space (called the “coverage region”) that is guaranteed to contain the output of the diffusion model with a user-defined probability (called the “coverage level”). In PlanCP, inspired by concepts from conformal prediction, we modify the loss function for training the diffusion model to include a quantile term to encourage more robust performance across the variety of training examples. At test time, we then calibrate PlanCP with a conformal prediction process to obtain coverage sets for the trajectory prediction with guaranteed coverage level. We evaluate our algorithm on various planning tasks and model-based offline reinforcement learning tasks and show that it reduces the uncertainty of the learned trajectory prediction model. As a by-product, our algorithm PlanCP outperforms prior algorithms on existing offline RL benchmarks and challenging continuous planning tasks. Our method can be combined with most model-based planning approaches to produce uncertainty estimates of the closed-loop system.

Conformal Prediction for Uncertainty-Aware Planning with Diffusion Dynamics Model
Jiankai Sun Stanford University jksun@stanford edu
Yiqi Jiang Stanford University
Jianing Qiu
Imperial College London
yqj i angistannotarrrroo ord. edu
jianing. giu17@imperial. ac.
Parth Talpur Nobel
Mykel Kochenderfer
Mac Schwager Stanford University
Stanford University
Stanford University mykel@stanf ord. edu
ptnobel@stanfore edu
schwager@stanf oord edu
Abstract
Robotic applications often involve working in environments that are uncertain, dynamic, and partially observable. Recently, diffusion models have been proposed for learning trajectory prediction models trained from expert demonstrations, which can be used for planning in robot tasks. Such models have demonstrated strong ability to overcome challenges such as multi-modal action distributions, high- dimensional output spaces, and training instability. It is crucial to quantify the uncertainty of these dynamics models when using them for planning. this paper, quantify the uncertainty of diffusion dynamics models using Conformal Prediction (CP). Given finite number of exchangeable expert trajectory examples (called the "calibration set"), use CP obtain set in the trajectory space (called the "coverage region") ") that is guaranteed to contain the output of the diffusion model with user-defined probability (called the "coverage level"). In PlanCP, inspired by concepts from conformal prediction, we modify the loss function for training the diffusion model to include a quantile robust performance across the variety of training examples. At test time, we then for
to encourage
calibrate PlanCP with conformal prediction process obtain coverage the trajectory prediction with guaranteed coverage level. We evaluate our algorithm various planning tasks and model-based offline reinforcement learning tasks and show that it reduces the uncertainty of the learned trajectory prediction model. As a by-pu sy-product, our algorithm PlanCP outperforms prior algorithms on existing offline RL benchmarks and challenging continuous planning tasks. Our method can be combined with most model-bases planning approaches to produce uncertainty estimates of the closed-loop system.
Introduction
Uncertainty is major challenge in robotics, arising from variety of sources, including errors in modeling, sensing, planning, and control [ B. Understanding the effect of uncertainty is important for developing successful motion strategies, it allows robots anticipate and plan for unexpected events. Therefore, planning with uncertainty is a crucial challenge in robotics and accurately modeling this uncertainty is important for ensuring the safety and reliability of robotic systems, particularly in safety-critical applications. A range of techniques, such as probabilistic modeling [4] and robust optimization (5,6), can be used to account for the uncertainty of dynamics models in robotics.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

The dynamics model is especially critical for model- planning 7 and reinforcement learning (MBRL) 8 Due the advantages such as expressing multi-modal action distributions, high- dimensional output space, and stable training, there has been a growing trend in learning-based planning approaches to utilizing diffusion models as dynamics models [9] [Q]. Diffusion models can also address the challenge of multi-modal distribution of actions in planning, which arises when multiple actions may be appropriate for a given state. This is achieved by enabling policies to express distribution over multiple actions rather than deterministic choice of a single action ([]]. Furthermore, diffusion models hold the potential to be applied to Offline Reinforcement Learning, which aims to learn effective strategies using offline datasets without exploration [1].. However, existing model-bard based algorithms struggle when encountering cutt-ofddistribute
(OOD) actions or states [9] A crucial element absent from these methods is the adequate management of uncertainty in
the offline scenario [33 14 Uncertainty-awan dynamics models that convey algorithmic confidence valuable if they possess accurate calibration and sharpness [1]]. Calibration entails aligning predicted probabilities with actual outcomes, whereas sharpness refers to the concentration of predicted probabilities around the true probability [16]. Nevertheless, many machine learning algorithms generate poorly calibrated probability estimates, which can result in overconfident or underconfident predictions and lead suboptimal decision-making [17 [8]. Therefore, ensure accuracy and optimal decision-making, it is crucial to utilize uncertainty-awart models that are well-calibrated and sharp.
Conformal Prediction (CP) 19 is a potent approach to assessing the uncertainty associated with predictions generated by black-box models. CP offers several advantages traditional methods for measuring uncertainty. For example, CP is model-agnnstic (20), meaning it can be applied to any type of model without requiring knowledge of the model's inner mechanisms. This is particularly advantageous in cases where the underlying model complex and not well understood. Additionally, CP allows for customization of the desired confidence level, providing the ability control the trade-off between predictive accuracy and the level of uncertainty (21). This feature is especially important in safety-critical applications where a higher level of confidence is required. Furthermore, CP can be used for online learning (22), meaning that it can adapt to new data points as they become available. This makes a particularly useful approach for planning dynamic environments. Lastly, CP is distribution-free free [23], providing statistical guarantees (24] and reliable uncertainty estimates from finite samples (25] Given these advantages, we choose to use CP as our method of uncertainty quantification for diffusion dynamics models in planning. The flexibility and reliability of CP provide a rigorous and mathematically proven framework for uncertainty quantification, which is critical for ensuring the safety and reliability of robotic systems in a variety of applications.
In this paper, we connect conformal prediction planning and propose PlanCP, adopting a practical and effective conformal prediction uncertainty estimation method for planning algorithms. Empiri- cally, we observe that PlanCP substantially improves model stability. In addition, PlanCP reduces uncertainty without sacrificing performance on datasets with demonstrations collected from experts. Plan PlanCP can be applied to reinforcement learning or imitation learning from demonstrations [26] 27) as well. However, since we build on the Diffuser [9), which calls itself planning method, we will generally stick with the planning terminology for consistency with the literature.
In summary, our main contributions are as follows:
We introduce PlanCP, a framework that enables the use of a conformal prediction to estimate uncertainty in planning tasks, without sacrificing the accuracy of dynamics models.
Our approach is compatible with diffusion dynamics models. We demonstrate that by employing a diffusion model as the foundational dynamics model, we can construct conformal predictor that guarantees trustworthy and accurate uncertainty estimates.
We extensively evaluate our approach, and the experimental results demonstrate that exhibits lower uncertainty without compromising performance across multiple benchmarks, including Offline Reinforcement Learning.
2
Related Work
2.1
Conformal Prediction
Conformal Prediction (CP) 19 28 29
effective for measuring the un- certainty associated with predictions produced by black-box models. As variant of conformal
32 33

---

prediction, Transductive Conformal Predictor (TCP) can be applied on top of existing machine learning methods to establish valid prediction regions, 19. TCP takes into account the structure of the data being predicted and uses the entire dataset during the prediction phase to construct prediction regions around each data point [34]. To improve the computational efficiency of the initial transduc- tive approach, a more computationally efficient framework Inductive Conformal Predictor (ICP) was proposed. which involves dividing the training data into a proper training set and a calibration set (35) However, ICP can be informationally inefficient as some examples are used solely for modeling while others are used only for calibration. To address this, ensembles of conformal predictors have been proposed 36 37] and (3). Examples include Cross Conformal Predictor (CCP) [39], Bootstrap Conformal Predictor (BCP) [99 and the generalized Aggregated Conformal Predictor (ACP) [40] In recent study [1] conformal prediction
used to construct prediction regions for recurrnt neural networks (RNNs). Strawn et al. [42) propose conformal predictive safety filters that use statistical techniques to provide uncertainty intervals around predictions and learn an additional safety filter that closely follows the RI. controller but avoids uncertainty intervals. Tonkens et al. [45] expand upon prior work in planning with probabilistic safety guarantees using conformal prediction. Conceptually
closest to our work are [44] 45 42 43) However, in these works, only quantification of uncertainty
provided, and reduction of the uncertainty during training is absent. In comparison, we are able to reduce uncertainty in training and extend conformal prediction to diffusion models.
2.2 Diffusion Models
In recent years, diffusion models have gained significant attention as powerful class of generative
models that formulates the data- gate-eeneat process as sequence of denoising steps [46] 47 48 [.].
Inspired by score matching [49 and energy-based models (EBMs) 50 51 52 53]. diffusion models
rely on a denoising procedure, which can be viewed as a way of parameterizing the gradients of the data distribution [54]. From the model-based trajectory-
perspective, Diffuser [9] applies a diffusion model as a trajectory generator. From the model-free volicy-optimization perspective, Wang et al. [1]] utilize diffusion model to represent the policy, applying it to the action space and conditioning it on the states to form a conditional diffusion model. Diffusion Policy [[O] utilizes a conditional diffusion-based process to model policies for real robots, exhibiting robust handling of multi-modal action distributions. However, the aforementioned approaches lack the capability to quantify the associated uncertainty, which is precisely the focus of our paper.
2.3 Planning, Predictive Dynamies Model, and Modcl-Basee Reinforcement Learning
Model-based reinforcement learning (MBRL) has emerged as promising solution for real-world sequential decision making problems with distinct advantages such as analytic gradient computa-
tion 155 A significant body of literature focuses on learning dynamics model and utilizing the
learned model for policy learning through model-based planning [56 57 58 [9]. PIL.CO 60) utilizes
Gaussian process to model system dynamies, while World Models 61 and Dreamer 62 learn
latent dynamics through two-stage and online planming, respectively. MGAIL [63] introduces a forward model in GAIL, but its discriminators output uniform representation that is less transfer- able. PETS (64] combines probabilistic neural networks and deterministic neural networks to create uncertainty- models. To capture long- long-term dependencies, FNNs, RNNs, and Transformers have been used to develop internal models [65] 66, 67]. Our study distinguishes itself from prior research by investigating techniques to measure and minimize uncertainty using conformal prediction for diffusion-based dynamics models.
3 Method
In this section, we describe our framework PlanCP (as Fig. D shows) solve the problem of uncertainty-aware planning. We begin by outlining our overall methodology, which involves three key steps. First, we introduce the diffusion- based dynamics model to predict future trajectories. Second, we apply the conformal prediction technique on the dynamics model to quantify and reduce the uncertainty. Finally, we summarize the whole procedure of our algorithm. In the following subsections, we provide more details on each of these steps and explain the specific techniques we used to implement them. As for notation, we use superscript n to denote diffusion timestep, or i to denote different trajectories, and to denote trajectory timestep.
3

---

Calibration (Deat
Dynamics (Diffusion Medell)
Validating
planner Pe reward function R(-))
Pradictoo
Uncertainty-aware Optimization
Uncertainty Quantification
Testing (Drest)
Training (Derain)
VeLreward
Conformal Quantile
Vo Volpptnn
Now
Figure I: PlanCP Framework: To update the dynamics model during training, we perform uncertainty-aware optimization using the training set Dirain. During calibration, we construct the uncertainty interval y based on Deat and then we verify the uncertainty interval using the test set Dreat-
3.1 Background
We consider Markov Decision Process (MDP) represented by the tuple (S, A, T. Po, 7) with state space S, action-space A, dynamics T (a, s), reward function r(s, a), initial state distribution Po(s), and discount factor distribution over actions. We define T as trajectories, which consist of sequences of state-action pairs (so. a...... 8T, aT) generated by policy - are expert trajectories generated by expert policy π represents the time horizon. We use a two-dimensional array to represent the trajectory T:
(0, 1). We assume stochastic policy и that maps state
80
81
8T
(1)
ag
0.
with each column corresponding to timestep in the planning horizon.
3.1.1 Planning with Conditional Diffusion Probabilistic Models
In Bayesian inference, it is intractable to compute the exact posterior distribution due the high dimensionality and the complexity of the model. One type of latent variable model approximates
posterior by the forward or diffusion process q(>" ((7")>"-)) )). which modeled fixed Markov chain that progressively adds Gaussian noise to the data. Diffusion Probabilistic Models, as described in 46 7 characterize the process of data generation as an iterative denoising procedure, formalized as the prohability distribution in reverse of a forward process q(T") ([(""" which progressively introduces noise and alters the underlying structure of the data.
(7"""["") The denoising operation (as Eq. 2 shows) is performed
N
po(+ = / p(T N
II Po(T=~1|p "dTINNN
where P(T N 0, is standard Gaussian prior and - ~ g(70) is noiseless problem data. The reverse process parameterized as Gaussian (Eq. 3)
Po ( >e(7~~1|7 (1") = -NT" #0(T" n). E,,)
in which μi and Σ are the mean and covariance of the Gaussian distribution respectively. Janner et
al. (9] proposed a way to turn diffusion models into planners that can incorporate different conditions. To achieve this, a learned/expert-desiinee condition function ue(-) added the diffusion model,
which yields planner Po described by Eq. 4
Po(T) O Po(T)ur(-)
where PT( is the diffusion model's probability distribution, and w(+) is conditional function that can include information such as observations history, desired goals, or rewards. Practically, this can be implemented by sampling from the unperturbed reverse process and replacing the sampled values with conditioning values w(-)) at all diffusion timesteps. We can use a learned reward function as condition, introduced in Sec. 3.2

---

Algorithm 1: PlanCP: Conformal Prediction for Planning with Diffusion Dynamics Models
Data:
Training set Dtrain with train samples, calibration set Deat with Keat samples, test set Drest with Keett samples, planner Po. expert trajectory = and failure probability a. Result: predicted trajectory т, uncertainty C.
while not converged do
for 0, 1, Ktrain do
Get predicted trajectory т, corresponding to expert trajectory T." E Diroins
Compute reconstruction nonconformity score as Eq. 6 foo each T. Darain's
Construct uncertainty C; as Eq.
Compute total loss Lplan as Eq. 13 and Creward Eq. 14
Update the dynamics model Pa using loss Epian and the reward model R(-) using loss
Lrewordi
end end for i 2, , do
10
11
Get predicted trajectory т; as Deat;
12
Compute reward conformity score for each Ti Deal;
13
Construct predicted interval 0 as Eq. 15
end
15
Verify the predicted interval using the test set Diest'
3.2 Learned Reward Function
Learning a reward model important because allows an intelligent agent to estimate the expected rewards of given action and state, which is useful for the agent to make informed decisions in an environment. Additionally, using learned reward model can enable us to adapt to changes in the environment and adjust our predictions accordingly. However, learning accurate reward model from data is challenging. Thus, we introduce transformer based reward model as Eq. 5 which can perform accurate long-ten temporal predictions.
for R(MultHHead( 80.T, K 80.7, ao.T)); MultiHead K, V) = Concat(heade beadd,,,,ed,,.00 head2, head, W
head, (QWR, K vWY),
(5)
QKT VJ
Attention K, softmax
V
where R(-) is an MLP layer to map the multi-head attention feature to a predicted reward. dy is the dimension of the attention feature. For more details of transformer architecture, please refer to [68]. The transformer based reward model is used as a condition for diffusion models to generate reasonable trajectories.
3.3 Quantifying Uncertainty with Conformal Prediction
The inherent uncertainty in the diffusion-based dynamics model Po lead inaccurate predictions. To address these challenges, we employ the conformal prediction technique to obtain prediction uncertainty. Unlike other approaches [4) conformal prediction does not rely on assumptions about the underlying distribution or the predictive model, making it a flexible and robust method for uncertainty quantification machine learning 24 This allows to quantify the uncertainty associated with the predictions and provides more robust estimates of the model's 's performance.
We have split the dataset D into three parts Dtroins Deats and Dteat for training, calibration, and testing, respectively. The planner Pa and reward function R(-) are learned from the D Drrain dataset.
In our problem setting, we use reconstruction nonconformity scores during training and reward conformity scores during calibration and testing. Reconstruction nonconformity scores is defined as
5

---

RC
Diffuser
PlanCP
Diffuser (Failed)
PlanCP
Diffuser
PlanCP
(a) Maze2D U-Maze
(b) Maze2D Medium
(c) Maze2D Large
Figure 2: Learned long-horizon planning: By learning the planning procedure, PlanCP is capable of conducting long-horizon planning with a specified start and goal ( condition. even under conditions of sparse rewards. Sub-environment visualization for Maze2D: (a) U-Maze, (b) Medium, (c) Large. We have marked the risky region in the baseline with purple dashed boxes.
the mean squared error of the learned trajectories from the expert trajectories. Reward conformity scores exactly the reward returned by the environment for given trajectory.
Training To be able compute and reduce the uncertainty during training, we need perform conformal prediction on the training set Dtrain- First, we compute the reconstruction nonconformity score as [ [6 each trajectory т; of the training set Derain- A lower reconstruction nonconformity score corresponds to better predictions.
i = ||T; - 7. 1.2
(6)
where is the predicted ith trajectory while 7. is the ground-truth ith trajectory from the expert demonstration of Dtrain-
Calibration and Testing During calibration and testing, we use reward conformity score to quantify the uncertainty. Let ro, r-,-.- Tk be + rewards corresponding to To, T.. Th trajectories predicted by our planner Pa. These rewards are defined and returned by the environments, which can be regarded as independent and identically distributed random variables. For planning, the reward conformity score can be seen as metric that reflects the planner's ability to predict the quality of plans. A higher reward conformity score corresponds better predictions, indicating that the planner is more accurate in its estimations. Our objective is to quantify the uncertainty associated with the reward r, based on set of reference rewards (ro.... .rk) r; Specifically, we aim to construct prediction interval C; that contains with probability at least 1 E (0, 1). Formally, we seek to find a valid prediction interval C, that satisfies the following inequality:
P(r: E C,(tro. .-. .. \ (r)))
21
where Ciltro," rk] \ rr)) denotes the prediction interval constructed using the reference rewards (ro. ... r Tt. The prediction interval C2 should provide a measure of the uncertainty associated with the reward rı, and should be constructed in such way to guarantee that with a probability of at least a, it contains rü.
Lemma 3.1. Gäven k + exchangable random trajectories (T T.. T
,Th) Derains predicted
trajectory (To, TI, ... ,T), corresponding (non)conformii scores do, di...., dhs and a failure probability (0, 1), then
P(d. CC)=
(8)
where the prediction interval Os are defined by
[kk + 1)(1 - a)l C = Quantile (do...., di dib de;
(9)
k.
During training, we use reconstruction nonconformity score as d,, while during calibration and testing, we use the reward conformity score as d, Note that if is less than 1/k, our confidence intervals become infinitely wide as we have insufficient data to make predictions with that accuracy.

---

Table 1: Long-horizon planning in Maze2D. Maze2D features a sparse reward structure and requires long-horizon planning. PlanCP has lower uncertainty and better performance. Bold font means better performance.
Environment
BCQ
770
CQL
71
IQL
72
Diffuser 9
PlanCP
Maze2D U-Maze
12.8
5.7
47.4
113.9 3.1
116.4 + 3.2
Maze2D Medium
8.3
5.0
34.9
121.5 + 2.7
128.5 + 2.5
Maze2I Large
6.2
12.5
58.6
123.0 + 6.4
130.9 + 5.6
Average
9.1
7.7
47.0
119.5
125.3
3.4 Trajectory Optimization Uncertainty-aware Planning
Trajectory optimization for a discrete-time dynamics system defined by 88+ = [(c a₂) involves identifying sequence of actions do.r that maximizes objective function J, which is the sum of per time-step rewards or costs c(s;, a.). The optimization problem can be represented as follows:
T
a.:T = arg min J (80, ao:T) = arg min
y c(( az).
(10)
09:7
0g:?
1=0
Here, T represents the planning horizon.
While training, we utilize the uncertainty calculated from { as a cost function to optimize the parameters of the planner model. After computing the reconstruction nonconformity scores
dataset and setting Ć' = Quantile ff.....,Pe: [(k+1)11-a) we can form the uncertainty loss:
(()) uncertaintt
(11)
After adopting the differentiable ranking and sorting techniques [69,, the soft quantile function is differentiable. On the other hand, diffusion probabilistic model is used to parameterize the learned gradient co (T", n) of the trajectory denoising process. The objective for training the e model can be
expressed as Eq. 12
[( reecm = En,,, - c(("",,
(12)
where n (1, ................................ N} represents the diffusion timestep. e ~ N I) represents the noise target. Finally, -" denotes the trajectory 7O that is corrupted with noise
In the context of learning. hope for expected cumulative rewards, which can be used to estimate the uncertainty of the predictions. By incorporating these prediction intervals into the decision-making process, make more informed decisions and avoid overly optimistic or pessimistic estimations. Finally, the learning objective for planner training is
conformal prediction construct prediction intervals
Eytan = L(0) + Auncertainty (()) uncertainty-
(13)
Given predicted reward and ground- truth reward r, the learning objective for reward model training is
Creward ||f -r12--
(14)
The reward model R(-) is used as the condition for the diffusion model Eq. shows. We can obtain reward r returned by the environment or learned reward model. On the calibration dataset Deats we compute the reward conformity scores and set Q Q!-aa and prediction interval q using conformal prediction as Eq. 15
[(k: 1)(1 -
Qa (Tk+1) = Quantile
TI.....,TRI TRi
[Kk + 1)a b
(15)
Q (rx=1) Quantile Ti....,Tki
((T++1) = (Q²-*(rx+1), (rk+1).
Once the prediction interval is constructed, assess its performance the test set Drest-
7

---

Table Uncertainty. The uncertainty quantification of PlanCP and prior algorithms in the Maze2D environment. PlanCP has lower uncertainty than previous methods for most tasks. Bold font means better performance.
Diffuser 9
PlanCP Uncertainty 4
Dataset
Rewards f
Uncertainty
Interval
Rewards f
Interval @
Mazc2D U-Maze
113.9 3.1 [121,83, 147.19] 121.5 [132.77, 141.85]
25.36
116.4 = 3.2 [121.11, 142.85]
21.74
Maze2D Medium
9.08
128.5 2.5 [132.77, 142.99]
10.22
Maze2D Large
123.0 6.4 192.11, 163.23]
71.12
130.9 + 5.6 197.75, 168.46)
70.71
Average
119.5
125.3
3.5 Our Proposed Procedures
We first iterate over the training set Dtrain and for each expert trajectory T," obtain a predicted tra- jectory т, using the learned dynamics model Pa. Then we compute the reconstruction nonconformity score 7' for each T. and construct the uncertainty C. The loss Cpion and Lyeward are then computed using Eq 13 and Eq 14 respectively. Finally, the dynamics model Py and reward model R(-) updated using the computed loss.
Next, we iterate over the calibration set Deat, and for each predicted trajectory Tú, obtain the reward conformity score т; and construct the predicted interval 0:-
Finally, we verify the predicted interval using the test set Diest to evaluate its performance. The whole procedure is shown in Alg.
Experiments
In our experiments, we aim to answer the following questions: (1) How to construct uncertainty intervals that can capture the dynamics model's 's prediction error and provide measure of confidence in the planned trajectories? (2) Can the uncertainty of the dynamics model be reduced during the training process? To answer these questions,
evaluate Pl PlanCP on multiple planning and offline RL tasks. We show planning with our algorithm can update the dynamics model while quantifying uncertainty. In principle, our approach can be applied to wide range of diffusion dynamics models. For illustrative purposes, we demonstrate our framework using a typical Diffuser [9]. Our primary objective is assess and minimize the model's uncertainty, while performance improvement is by-product of our approach and is not necessarily guaranteed. We assess the performance of our models based on several quantitative metrics, including credibility and rewards.
4.1 Training setup
We set the uncertainty weight to Auncertainty = 5 and the failure probability to a 0.1. To optimize the model, we use the Adam 773 74] optimizer with a learning rate of x 10-4, We train the diffusion dynamics model on the training set Dirain for x 10 iterations. As lack ground-truth for evaluation, we generate the evaluation datasets Deat and Dtest during the evaluation process. We partition the evaluation data into 20% for Deal and 80% for Drest. The uncertainty measure is derived from the 20% calibration set through conformal prediction.
4.2 Long Horizon Task Planning
We use Maze2D to evaluate long-horizon planning. q is measure of uncertainty for the rewards obtained by a planned trajectory. For most variants of the Maze2D environment, we have smaller uncertainty interval, especially for Maze2D U-Maze, as Tab. [ and Tab. show. As Fig. 2 shows, Compared to Diffuser, our PlanCP has successfully generated a smoother path with lower uncertainty, which achieves a higher reward. Furthermore, in cases where the previous Diffuser method fails to plan, our PlanCP is able to derive feasible path.
8

---

Table 3: Uncertainty of Offline Reinforcement Learning. The uncertainty quantification of PlanCP and Diffuser (9] in the D4RL environment. PlanCP has lower uncertainty than Diffuser 9 for most tasks. Bold font means better performance.
Diffuser g
PlanCP Uncertainty
Dataset
Enviromment
Rewards
Uncertainly
Interval
Rewards
Interval 0
Seddimm-Expenn
HalsCheetah
85.82 1: 6.91
[82.20. 91.55]
9.34
87.44 1 6.51
[83.12, 91.81]
8.68
Medium-Exppe Expert Walker2d
108.22 8.52
[108.48. 109.58)
1.10
108.31 + 8.07
[107.86. 110.43]
2.56
Medium
HalfChectah
46,10 0.96
[45,05. 47.34]
2.29
46.26 1.35
[44.64, 47.97]
3.33
Mcdium
Walker2d
75.73 = 19.59
[40.76, 87.72]
46.96
77.08 + 18.21
[44,51, 88.46]
43.95
Medium-Reppe Replay
HalfCheetah
37.34 +: 6.57
[27.92. 41.58)
13.66
37.65 + 6.29
[29.25, 42.23]
12.98
Sodium-Replay Walker2d
52.24 24.44
[21.18. 87.83]
66.65
54.53 23.00
[25.41, 87.78]
62.37
Average
67.58
68.55
4.3 Offline Reinforcement Learning
Finally, we evaluate the ability to recover an effective single-task controller from heterogeneous data of varying quality using the D4RL Benchmark (75). The D4RL. benchmark consists of several challenging continuous control tasks, such as the HalfCheetah and Ant environments. For each task, we measured the average return achieved by our approach and the baseline methods over multiple runs. We also computed the standard deviation of the returns to measure the stability of the algorithms. Our experimental results demonstrate that our approach outperforms the baseline methods on most of the D4RL environments, as Tab. 3 shows. For most of the environments, our approach achieved higher return with a smaller uncertainty interval compared to the baseline methods. We also observed that our approach achieved a more stable performance across multiple runs, as evidenced by the smaller standard deviation the returns.
Overall, experimental results show that our framework can construct uncertainty intervals that can capture the dynamics model's prediction error and provide measure of confidence in the planned trajectories.
5 Discussion, Limitations, and Conclusions
In this work, we tackled the problem of obtaining uncertainty estimates for planning and offline RL tasks by proposing a new approach that applies conformal prediction to diffusion-basee dynamics model. Experimental results show that our framework maintains the original model's accuracy while reducing uncertainty. One limitation of our study is that we have not conducted evaluations on real robotic platforms. An intriguing avenue for future research involves gaining a comprehensive understanding and conducting meticulous analysis of how conformal prediction compares to the extensive body research in planning under uncertainty, including non- paraamettei models. Our work focuses primarily proposing algorithmic framework. Future work could also explore the use of more recent and effective diffusion models, as well as other types of dynamics models. Moreover, evaluating the proposed method real-world applications, such autonomous driving or robotics, where reliable uncertainty estimates are crucial for safety-critical systems, would be an interesting avenue for future research.
Acknowledgments and Disclosure of Funding
The NASA University Leadership initiative (grant #80NSSC20M0163) provided funds to assist the authors with their research, but this article solely reflects the opinions and conclusions of its authors and not any NASA entity.
This research was also partly funded by a gift from Meta's Project Aria.
Parth Nobel was supported in part by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGF-16565188 Any opinions, findings, and conclusions or recommenda- tions expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.

---

References
[1] Michael Andreas Erdmann. On motion planning with uncertainty. 1984.
[2] Mykel Kochenderfer, Tim A Wheeler, and Kyle H Wray. Algorithms for decision making. MIT press, 2022.
[3] Jonathan Ko and Dieter Fox. Gp-bayesfilter: Bayesian filtering using gaussian process prediction and observation models. In 2008 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 3471-3476, 2008.
[4] Christian Soize. Uncertainty quantification. Springer, 2017.
[5] Aharon Ben- Tal, Laurent El Ghaoui, and Arkadi Nemirovski. Robust optimization, volume 28. Princeton university press, 2009.
[6] Chancellor Johnstone and Bruce Cox. Conformal uncertainty sets for robust optimization. In Conformal and Probabilistic Prediction and Applications, pages 72-90. PMLR, 2021.
[7] Rinu Boney, Juho Kannala, and Alexander Ilin. Regularizing model-based planning with energy-based models. In Conference on Robot Learning, pages 182-191. PMLR, 2020.
[8] Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT press, 2018.
[9] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
[10] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04437, 2023.
[11] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[12] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 141424 Curran Associates, Inc., 2020.
[13] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, pages 104-114. PMLR, 2020.
[14] Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua M Susskind, Jian Zhang, Ruslan Salakhut- dinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11319-1132 11328. PMLR. 18-24 Jul 2021.
[15] Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):359 378, 2007.
[16] Tilmann Gneiting, Fadoua Balabdaoui, and Adrian E Raftery, Probabilistic forecasts, calibration and sharpness. Journal of the Royal Statistical Sociery: Series (Statistical Methodology). 69(2):243-25 2688 2007.
[17] Shengjia Zhao, Michael Kim, Roshni Sahoo, Tengyu Ma, and Stefano Ermon. Calibrating predictions to decisions: novel approach multi-class calibration. Advances in Neural Information Processing Systems, 34:22313-222 2021.
[18] Mathias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai, Neil Houlsby, Dustin Tran, and Mario Lucie. Revisiting the calibration of modern neural networks. Advances in Neural Information Processing Systems, 34: 14:15682-1566 2021.
10

---

[19] Vladimir Vovk, Alexander Gammerman, and Glenn Shafer. Algorithmic learning in a random world. Springer Science & Business Media, 2005.
[20] Ulf Johansson, Henrik Linusson, Tuve Löfström, and Henrik Boström. Model-agnostic non- conformity functions for conformal classification. In 2017 International Joint Conference Neural Networks (ICNN), pages 2072-2077 IEEE, 2017.
[21] Matteo Fontana, Gianluca Zeni, and Simone Vantini. Conformal prediction: a unified review of theory and new challenges. Bernoulli, 29(1):1-23, 2023.
[22] Isaac Gibbs and Emmanuel Candes. Adaptive conformal inference under distribution shift. Advances in Neural Information Processing Systems, 34:1660-167,, 2021.
[23] Jing Lei, James Robins, and Larry Wasserman. Distribution-free prediction sets, Journal of the American Statistical Association, 108(501):278-8877 2013.
[24] Anastasios N Angelopoulos and Stephen Bates. A gentle introduction to conformal prediction and distribution-frre uncertainty quantification. arXiv preprint arXiv:2107.07511, 2021.
[25] Yachong Yang and Arun Kumar Kuchibhotla. Finite-sample efficient conformal prediction. arXiv preprint 1rXh::2104.13871 2021.
[26] Junning Huang, Sirui Xie, Jiankai Sun, Qiurui Ma, Chunxiao Liu, Dahua Lin, and Bolei Zhou, Learning a decision module by imitating driver's 's control behaviors. In Jens Kober, Fabio Ramos, and Claire Tomlin, editors, Proceedings of the 2020 Conference on Robot Learning, volume 155 of Proceedings of Machine Learning Research, pages 1-10. PMLR, 16-18 Nov 2021.
[27] Jiankai Sun, Hao Sun, Tian Han, and Bolei Zhou. Neuro-symbolic program search for au- tonomous driving decision module design. In Jens Kober, Fabio Ramos, and Claire Tomlin, editors, Proceedings of the 2020 Conference on Robot Leurning, volume 155 of Proceedings of Machine Learning Research, pages 21-30. PMLR, 16-18 Nov 2021.
[28] Harris Papadopoulos, Kostas Proedrou, Volodya Vovk, and Alex Gammerman. Inductive confidence machines for regression. In European Conference on Machine Learning, pages 345-356. Springer, 2002.
[29] Jing Lei and Larry Wasserman, Distribution-fier prediction bands for non parametric regression. Journal of the Royal Statistical Sociery: Series (Statistical Methodology). (6(1):71-96 2014.
[30] Glenn Shafer and Vladimir Vovk. A tutorial on conformal prediction. Journal of Machine Learning Research, 9(3), 2008
[31] Benjamin Kompa. Jasper Snock, and Andrew Beam. Empirical frequentist coverage of deep learning uncertainty quantification procedures. Entropy, 23(12):1608, 2021.
[32] Rina Foygel Barber, Emmanuel J Candes, Aaditya Ramdas, and Ryan J Tibshirani. Conformal prediction beyond exchangeability. arXiv preprint UrXiv:2202./3411 2022.
[33] Yaniv Romano, Evan Patterson, and Emmanuel Candes. Conformalized quantile regression. Advances in neural information processing systems, 32, 2019.
[34] Vladimir Vovk. Transductive conformal predictors. In Artificial Intelligence Applications and Innovations: 9th IFIP WG 12.5 International Conference, AIAI 2013, Paphos, Cyprus, September 30 Octobrr 2, 2013, Proceedings 9, pages 360. Springer, 2013.
[35] Harris Papadopoulos. Inductive conformal prediction: Theory and application to neural net- works. In Tools in artificial intelligence. Citeseer, 2008
[36] Paolo Toccaceli and Alexander Gammerman. Combination of inductive mondrian conformal predictors. Machine Learning, 108:489-510, 2019.
[37] Tuve Löfström, Ulf Johansson, and Henrik Boström. Effective utilization of data in inductive conformal prediction. International Joint Conference on Neural Networks, Dallas, TX, USA, August 4-9, 2013. IEEE, 2013.

---

[38] Vineeth N Balasubramanian, Shayok Chakraborty, and Sethuraman Panchanathan. Conformal predictions for information fusion: A comparative study of p-value combination methods. Annals of Mathematics and Artificial Intelligence, 74(1-2):45 65, 2015.
[39] Vladimir Vovk. Cross-conformal predictors. Annals of Mathematics and Artificial Intelligence, 74:9-28, 2015.
[40] Lars Carlsson, Martin Eklund, and Ulf Norinder. Aggregated conformal prediction. In Artificial Intelligence Applications and Innovations: AIAI 2014 Workshops: CoPA, MHDW, IIV,, and MT4BD, Rhodes, Greece, September 19-21, 2014. Proceedings 10, pages 231-240. Springer, 2014.
[41] Kamile Stankeviciute, Ahmed M. Alaa, and Mihacla van der Schaar. Conformal time-series forecasting. In M. Ranzato. A. Beygelzimer. Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 6216-6228. Curran Associates, Inc., 2021.
[42] Kegan J. Strawn, Nora Ayanian, and Lars Lindemann. Conformal predictive safety filter for controllers in dynamic environments, 2023.
[43] Sander Tonkens, Sophia Sun, Rose Yu, and Sylvia Herbert. Scalable safe long-horizon planning in dynamic environments leveraging conformal prediction and temporal correlations. In ICRA, 2023.
[44] Lars Lindemann, Matthew Cleaveland, Gihyun Shim, and George J. Pappas. Safe planning in dynamic environments using conformal prediction. IEEE Roboties and Automation Letters, 8(8):5116-511 2023.
[45] Yuxiao Chen, Ugo Rosolia, Chuchu Fan, Aaron Ames, and Richard Murray, Reactive motion planning with probabilisticsafety guarantees. In Conference on Robot Learning, pages 1958- 1970. PMLR, 2021.
[46] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265 PMLR, 2015.
[47] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6855 2020.
[48] Jiankai Sun, Rui Liu, and Bolei Zhou. Hiabp: Hierarchical initialized abp for unsupervised representation learning. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 35, pages 9747-9755, 2021.
[49] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[50] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato. and Fujie Huang. A tutorial on energy- based learning. Predicting structured data. 1(0). 2006.
[51] Yilun Du and Igor Mordatch. Implicit generation and generalization in energy- based models. arXiv preprint arXiv:1903. 08689, 2019.
[52] Erik Nijkamp, Mitch Hill, Song-Chun Chun Zhu, and Ying Nian Wu. Learning non-convergent non-persistent short-run meme toward energy- model. Advances in Neural Information Processing Systems, 32, 2019.
[53] Will Grathwohl, Kuan-Chieh Wang, Jörn-Heriik Jacobsen, David Duvenaud, and Richard Zemel. Learning the stein discrepancy for training and evaluating energy-based models without sampling. In International Conference on Machine Learning, pages 3732-3747. PMLR, 2020.
[54] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, Advances in neural information processing systems, 32, 2019.
12

---

[55] Jiankai Sun, Lantao Yu, Pinqian Dong, Bo Lu, and Bolei Zhou. Adversarial inverse reinforce- ment learning with self-attention dynamics model. IEEE Robotics and Automation Letters, ((2):1880-188) 2021.
[56] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-But Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 22519-12535 Curran Associates, Inc., 2019.
[57] 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 Neural Information Processing Systems,
2020.
[58] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 110-21823, Curran Associates, Inc., 2020.
[59] Xianyuan Zhan, Xiangyu Zhu, and Haoran Xu. Model-based offline planning with trajectory pruning. arXiv preprint arXiv:2105.07335 2021.
[60] Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference machine learning (ICML-11), pages 463-472 2011.
[61] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 2450-2462. Curran Associates, Inc., 2018
[62] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representa- tions, 2020.
[63] Nir Baram, Oron Anschel, Itai Caspi, and Shie Mannor, End-to-end differentiable adversarial imitation learning. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 390-399, International Convention Centre, Sydney, Australia, 06-11 Aug 2017.
[64] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, and R. Garnett, editors, Advances in Neural
H. Larochelle, K. Grauman, N. Information Processing Systems, volume 31, pages 3754-4765 Curran Associates, Inc., 2018.
[65] Tingwu Wang. Xuchan Bao, Ignasì Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model- based reinforcement learning. arXiv preprint arXiv:
2019.
[66] Jiankai Sun. Shreyas Kousik, David Fridovich-Keil and Mac Schwager. Connected autonomous vehicle motion planning with video predictions from smart, self-supervised infrastructure, 2023.
[67] Jiankai Sun, De-An Huang, Bo Lu, Yun-Hui Liu, Bolei Zhou, and Animesh Garg. Plate: Visually-groundee planning with transformers in procedural tasks. IEEE Robotics and Automa- tion Letters, 7(2):492444933 2022.
[68] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[69] Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert. Differentiable ranking and sorting using optimal transport. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
13

---

[70] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062. PMLR, 2019.
[71] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179. 1191, 2020.
[72] Iya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning, 2021.
[73] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[74] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations, 2019.
[75] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07119, 2020.
14


Reviews:
[{"review_id": "aqdYTBAjdv", "reviewer": "Reviewer_gfdz", "summary": "This paper describes a method for learning a dynamics model that uses conformal prediction for explicit representation of the uncertainty. The dynamics model is used for sequential decision making such as planning in a maze or learning control in one of the D4RL problems.", "questions": "* The definition of calibration is very unclear. Why is a separate calibration data set is used for the reward conformity loss function, rather than use the training data.\n\n* What is the exact definition of the Quantile function?\n\n* Line 7 of Algorithm 1 says \"update the dynamics model\" but it is not explained in the text how to do this. Figure 1 suggest the update is done through gradient descent ($\\nabla_\\theta \\mathcal{L}_{reward}$, etc.), but it is not obvious to me that you can differentiate through the Quantile function. How is this done? This is a minor point.", "limitations": "The authors provide a limitation section, but I disagree with them that the limitation of the approach is the absence of real robot experiments. Instead, I think the primary limitation is the lack of understanding and careful analysis of what conformal prediction achieves relative to the large body of work in planning under uncertainty, including non-parametric models.", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The primary strength of this paper is the novel use of conformal prediction. The authors correctly describe several possible benefits of conformal prediction as a representation of uncertainty, including not needing a commitment to a particular representation or parameterisation of the uncertainty in the dynamics.\n* The authors give a complete algorithm --- it is (mostly) clear how all the pieces work. \n* The paper is reasonably well written and clear.", "weaknesses": "* The primary weakness of the paper is that it is not clear what exactly the authors have accomplished. Developing a new method for sequential decision making should either show that we can solve problems that we could not solve before ((or find better solutions than was possible), or alternatively, provide some understanding and analysis of the value of a dynamics model that encodes uncertainty using conformal prediction. Unfortunately, this paper does not show that it can solve problems that could not be solved before, or outperform the state of the art. The only comparison is to Diffusion RL which is not the best performing RL algorithms according to the original Diffusion RL paper (Janner et al) or the current leader in the [D4RL benchmarks](https://paperswithcode.com/sota/gym-halfcheetah-expert-on-d4rl).\n\n* It's fine to not have the best performer on D4RL, but then what is conformal prediction buying for the algorithm? A far better comparison would have been to a model-based algorithm that had an explicit parametric representation of uncertainty, and an analysis of what is actually being learned. What is the effect of the loss functions in 12-14 on $P(\\theta)$?\n\n* The experimental results in general are not very compelling, at least in part because there is no systematic evaluation of the uncertainty. The paper needed to remind the reader of the uncertainty models in the D4RL and Maze engines, and even better, show that these uncertainties are difficult to capture with parametric models. Even better would be to show that conformal prediction outperforms non-parametric dynamics models such as Ko and Fox (2009, 2011) in important ways.\n\n* Additionally, while the paper overall is well-written, it is not careful about the distinction between planning and RL. This paper appears to be fundamentally an RL paper, and not a planning paper in the sense that if the model changes or the reward changes, there is no optimisation that can be used to recover from this change. If the loss function was solely equation 13, and did not include equation 14, then at least the model would be robust to changes in the reward function, and a new sequence of actions could be optimised over.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: What did this paper accomplish?\n\nA1: Thanks for the comments! Instead of developing a method for sequential decision-making that achieves the highest performance (reward) on D4RL, we are the first to augment the `diffusion` RL algorithm with uncertainty awareness. The progress in the D4RL benchmark is rapidly changing, but at the time of our paper submission, Diffuser (Janner et al.) was the state-of-the-art open-source Diffusion RL algorithm. The original Diffuser algorithm was evaluated only 100 times, which is insufficient to capture uncertainty. To address this limitation, we re-run 1000 evaluations of the Diffuser algorithm on our setup and present the performance results. Since our PlanCP is a modification of the Diffuser algorithm, we compare PlanCP to Diffuser and experimentally prove that our method has lower uncertainty. We focus on finding better solutions for reducing the uncertainty of sequential decision-making than before.\n>Q2: What is conformal prediction buying? \n\n**Comparison with a model-based algorithm that had an explicit parametric representation of uncertainty.**\nWe compare with a model-based algorithm MOPO that had an explicit parametric representation of uncertainty. We can see that the performance (reward) of MOPO is lower because it does not use the diffusion model as a dynamics model and the uncertainty (std and interval) is still large, as attached Table 1 shows (Summary Rebuttal). Also, MOPO has a limitation: it assumes that the data conforms to a Gaussian distribution.\n\n**What is actually being learned?**\nWhat the model is actually learning in conformal prediction is a measure of how well a data point fits. This is captured by the conformal score, which is used to construct the prediction intervals or regions. \n\n**The effect of the loss functions in 12-14 on $P(\\theta)$?**\n\nEq. (12) represents the reconstruction loss used to improve the accuracy of diffusion models by guiding the diffusion planner $P_\\theta$ in predicting dynamics more effectively.\n\nEq. (13) presents the joint learning objective for training the planner $P_\\theta$, aiming to enhance both accuracy and reduce uncertainty in the learned planners.\n\nEq. (14) measures the $L_2$ loss between the predicted reward and the actual reward, helping to optimize the reward model as the condition $w(\\cdot)$ in the conditional diffusion planner $P_\\theta$, as shown in Eq. (4).\n>Q3: How to interpret the experimental results?\n\nA3: GP-BayesFilters (Ko and Fox, 2009) show how GP prediction and observation models can be combined with particle filters (GP-PF), and extended Kalman filters (GP-EKF). Bayesian filters and conformal prediction, although both involving probability and uncertainty, have distinct goals and applications. Bayesian filters estimate system states based on observations, while conformal prediction provides confidence intervals for predictions. The main advantage of conformal prediction over Bayesian filtering is that it does not rely on prior knowledge of the data distribution, making it more robust in uncertain or unknown probability distributions. Additionally, conformal prediction offers a validity measure for each prediction, which is valuable in applications where the impact of incorrect predictions can be significant.\nWe compare with Ko and Fox, 2009, a non-parametric dynamics method, and MOPO. The result is shown as attached Table 1 (attached to the Summary Rebuttal). The findings show that conformal prediction outperforms non-parametric dynamics models Ko and Fox, 2009, and these uncertainties are challenging to capture with parametric models like MOPO.\n>Q4: Is this an RL paper or a planning paper? \n\nA4: This is a planning paper. Eq. (14) is used to learn a better reward model, which guides the conditional diffusion planner $P_\\theta$. Unlike the RL paradigm, the planner $P_\\theta$ is learned through expert data, allowing for re-planning if the model or reward changes. Planning and RL are two different approaches to sequential decision-making problems. Planning involves constructing a model of the environment and using it to simulate future outcomes, assuming complete knowledge of the environment. In contrast, RL involves learning from feedback in the form of rewards or penalties without complete knowledge of the environment and its dynamics. We build on Diffuser (Janner et al.), which calls itself a planning method.\n\n>Q5: Why is a separate calibration dataset used for the reward conformity loss function?\n\nA5: In conformal prediction, calibrating with the training data invalidates the statistical guarantees of the method.  Specifically, it results in overly confident (too small) prediction intervals that do not accurately reflect performance variation at test time.  Therefore, separating training and calibration data is standard practice in CP.\n\n>Q6: What is the exact definition of the Quantile function?\n\nA6: We follow the same mathematical definition of traditional Quantile function: Formally, for a random variable $X$ with cumulative distribution function $F(x)$, the Quantile function $Q(p)$ is defined as:\n$$\nQ(p) = \\inf\\lbrace x: F(x) \\ge p\\rbrace\n$$\nwhere $p$ is a probability between 0 and 1, and \"inf\" denotes the infimum of the set of values $\\lbrace x: F(x) \\ge p \\rbrace$.\n\n>Q7: How to update the dynamics model and differentiate through the Quantile function?\n\nAfter adopting the differentiable ranking and sorting techniques (Cuturi et al. Differentiable ranking and sorting using\noptimal transport, NeurIPS 2019), the soft quantile function is differentiable. We can use gradient descent to update the dynamics model. \nSpecifically, Let $X$ be a random variable that has a smooth density function $f$. Let $w = w(p)$ be the $p$-th quantile. Then the first derivative of the quantile function is\n$$\n\\frac{dw}{dp}=\\frac{1}{f(w(p))}\n$$\nprovided that the denominator is not zero. The second derivative is\n$$\n\\frac{d^2w}{dp^2}=-\\frac{f'(w)}{f(w)}(\\frac{dw}{dp})^2=\\frac{-f'(w)}{f(w)^3}\n$$\n"}, {"review_id": "qtyHaNGjy4", "reviewer": "Reviewer_zTTG", "summary": "The work proposes uncertainty quantification for learned dynamics model and imitation learning. The authors incorporates an uncertainty statistic as part of the loss to train their dynamics model that uses a diffusion model architecture. They show that doing so brings performance improvement empirically on common RL simulation environments. The method also performs conformal prediction on the learned reward forecast. ", "questions": "See weaknesses 3, 4, and 5.", "limitations": "See weaknesses. I do not think the discussion & limitation section fully address the limitations of this paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Using UQ methods and optimizing to reduce uncertainty in dynamics modeling is a great idea. The authors are able to make the first step of incorporating conformal statistics directly into the loss function, and is able to show results of (slight) improvement in performance. ", "weaknesses": "There are many areas that the paper needs improvement. \n\n1. English. There are grammatical errors, awkward sentence structures, and confusing statements throughout the paper. Although you don't need perfect English for a CS paper, improving on writing will help you get your message across. Using the abstract as an example, I would edit:\n\n- line 3-5: \"to overcome the xx, yy, and zz challenges\" -> to overcome challenges such as xx, yy, and zz.\n- line 14-15: \"Furthermore, during the test, PlanCP can also measure the model uncertainty\" -> Unclear what you are trying to say. Do you mean that model uncertainty is also used for planning during test time? \n- line 20-21: \"Our method is highly flexible and can combine .... and produces ...\" ->  Our model can be combined with ... to produce .... And what do you mean by flexible? is it with regard to tasks, or underlying models? \n- Figure 1: \"Dynamic model\" -> dynamics model\n\n2. There has been various recent works that uses CP for planning (see list below). You can argue that your setting is different from theirs, but it's important to cite them to provide context. \n\n- Lindemann, Lars, et al. \"Safe planning in dynamic environments using conformal prediction.\" IEEE Robotics and Automation Letters (2023).\n- Strawn, Kegan J., Nora Ayanian, and Lars Lindemann. \"Conformal Predictive Safety Filter for RL Controllers in Dynamic Environments.\" arXiv preprint arXiv:2306.02551 (2023).\n- Tonkens, Sander, et al. \"Scalable Safe Long-Horizon Planning in Dynamic Environments Leveraging Conformal Prediction and Temporal Correlations.\" ICRA (2023)\n\n3. More explanation is needed for the experiments section. \n- The term \"Validity %\" should be changed to \"Coverage %\". Validity is a binary metric; a method either satisfies the validity condition, or not. [1]\n- On the topic of validity - you chose $\\alpha = 0.1$ (line 261), does that mean that your coverage should be at least 90%? (or should it be 80% since you are choosing $[Q^{1-\\alpha}, Q^{\\alpha}]$ in Eq. 15?) Does that mean for all of your experiments, the CP uncertainty interval is _invalid_? That's problematic and means that your data violates the assumptions you are making. You can't say your method captures uncertainty in this case. \n- You need to say in your table captions that bold font means better performance, or the meaning is unclear.\n- How exactly is credibility calculated? it's not explained in either the main text or the appendix. What is the theoretical basis for using the $p$-value as metric?\n\n4. Lastly The term \"learning to reduce uncertainty\" is unsound. The \"reconstruction nonconformity scores\" are calculated on the training set, which violates the algorithm of conformal prediction [1]. The $\\mathcal{L}(\\theta)_{uncertainty}$ in equation 13 is hence not a measure of uncertainty, but just a training heuristics. The only valid CP component is the interval for rewards, which is not used anywhere in training. \n\n5. Contribution is not significant. Form my understanding, the authors used a common CP technique on an existing diffusion dynamics model. The uncertainty results are not new (in CP literature), and improvements in reward are rather insignificant. \n\n[1] Angelopoulos, Anastasios N., and Stephen Bates. \"A gentle introduction to conformal prediction and distribution-free uncertainty quantification.\" arXiv preprint arXiv:2107.07511 (2021).\n\nI think the directions the authors are going in is interesting and insightful; with some work this paper can be a great contribution to the community. However in its current state, the paper is not fit for publishing at NeurIPS yet.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: English\n\nA1:Thanks for pointing this out. We have fixed these grammatical errors, awkward sentence structures, and confusing statements throughout the paper. We will update them in the new version of the paper.\n\n**Is model uncertainty used for planning during test time?**\nNo. Currently, PlanCP can measure model uncertainty during testing. We have not used model uncertainty for planning during testing.\n\n>Q2: Cite recent works to provide context.\n\nA2: Thanks for mentioning these papers! In our NeurIPS submission, we already cited and discussed Lindemann, et al. RA-L (2023) in the related work section. Strawn et al. and Tonkens et al. were posted online after our NeurIPS submission, and are not peer-reviewed. We will cite and discuss these papers in the related work section in the updated version. Lindemann, et al. RA-L (2023), Strawn et al., and Tonkens et al. have made significant contributions to safe planning. These works are noteworthy because they effectively measure uncertainty with conformal prediction, even if it cannot be reduced by optimization.\n\n>Q3: More explanation for the experiments.\n\nA3: Thanks for pointing these out. We will fix them in the updated version.\n\n3.1. We will change the term \"Validity \\%\" to \"Coverage \\%\".\n\n3.2. We choose $\\alpha = 0.1$. Theoretically, the coverage will be exactly 80\\% only in expectation.  In CP, the empirical coverage follows a beta distribution with mean at the desired coverage (see, e.g., Angelopoulos and Bates, pp. 14). Empirically, our coverages are closely distributed around 80\\%, consistent with a beta distribution, which satisfies the assumptions.\n\n3.3. We will add \"Bold font means better performance.\" to the table captions.\n\n3.4. As we explained in our paper *Section 4.6 Credibility*: \"Credibility is defined as the p-value in a prediction set\". Previously, we believed this to be a trivial mathematical definition, and due to page limitations, we have not explained it much. Mathematically, we obtain credibility by doing T-test between the calibration trajectory and the predicted trajectory. \n$$\nt = \\frac{\\mu_1-\\mu_2}{\\sqrt{\\frac{\\sigma_1^2}{n_1} + \\frac{\\sigma_2^2}{n_2}}}\n$$\nwhere $\\mu_1$ is the mean of calibration trajectory samples, $\\mu_2$ is the mean predicted trajectory samples, $\\sigma^2_1$ is the variance of calibration trajectory samples, $\\sigma^2_2$ is the variance of predicted trajectory samples, $n_1$ is the sample size of calibration trajectory samples, and $n_2$ is the sample size of predicted trajectory samples.\n\\paragraph{What is the theoretical basis for using the p-value as a metric?} The use of the p-value as a metric is based on the principles of probability theory and statistical inference. The p-value is the probability of obtaining test results at least as extreme as the result actually observed. We follow Giovannotti et al. Transformer-based conformal predictors for paraphrase detection, PMLR 2021, to use $p$-value as metric.\n\n>Q4: Explanation of the term \"learning to reduce uncertainty\"\n\nThanks for the suggestion. We understand that using the training data to calibrate invalidates the guarantees of conformal prediction. That's why we split the data into a separate calibration set and test set. When we mention \"learning to reduce uncertainty,\" we are not referring to reducing $L_{uncertainty}$ on the training set. Instead, it means that our model exhibits lower reward uncertainty on the test set. We observed that our algorithm produces smaller reward intervals, which are calculated without relying on the training set.  Our method carefully follows the underlying requirements of CP (exchangeability of calibration data, etc.), and inherits the guarantees of CP.\n\n>Q5: Clarification of contributions\n\nA5: Thanks for the suggestion.\n\n5.1. The novelty of this work lies in using conformal prediction to (i) measure the uncertainty for the conditional diffusion planner using calibration data and (ii) reduce the uncertainty of the planner by including a quantile loss term during training. Previous CP literature can only measure the uncertainty but cannot reduce it. For experimental results, we emphasize that we accomplished reducing uncertainty without sacrificing performance.\n\n5.2. The application of CP to the diffusion dynamics model is novel and provides insights into the uncertainty associated with the predictions of the model. This can be useful for decision-making in scenarios where accurate predictions are critical. There is no prior work that attempts to introduce conformal prediction on diffusion models.\n\n5.3. The experimental results also show that we can reduce the uncertainty of the model without compromising performance.\n\n"}, {"review_id": "GqDcRYp9mx", "reviewer": "Reviewer_wAeC", "summary": "This work addresses the challenge of uncertainty estimation for planning. The authors propose the use of diffusion models for learning dynamics, which have demonstrated effectiveness in overcoming challenges such as multi-modal action distributions. To quantify the uncertainty of these dynamics models, they employ Conformal Prediction (CP), a technique for constructing prediction sets with valid coverage. They introduce PlanCP, a framework that connects conformal prediction to planning and optimizing the model explicity to minimize uncertainty. The effectiveness of uncertainty sets is evaluated through coverage and optimization performance, and the algorithm is tested in D4RL benchmarks and block stacking problem, showcasing reduced uncertainty and outperforming prior algorithms. The authors highlight the flexibility of their method, as it can be combined with different model-based planning approaches and provides uncertainty estimates of the dynamics model.\n\n", "questions": " I was wondering if the authors had some thoughts on what adapting CP framework to other conditioning approaches will entail, such as temporal condition guidance mechanisms.", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper provides a clear problem statement that is well motivated and clearly outlines the proposed approach with intuitive evaluation metrics.  Overall I found the paper well written and easy to follow.\n    \n- The proposed methodology presents a robust framework for quantifying and mitigating uncertainties within diffusion-based dynamics models, which have garnered significant attention in recent times.\n    \n- The quantitative results on D4RL although not SOTA(in terms of rewards), consistently outperforms the baseline without uncertainty quantification. It is important to note that the limited performance gains could be attributed to the constraints of the D4RL dataset itself, potentially reaching the performance ceiling. Nevertheless, the outlined approach consistently outperforms the baseline methods across the metrics identified.", "weaknesses": "-    The computational complexity of the proposed framework, especially during training and calibration, could hinder its scalability and applicability in real-time or resource-constrained scenarios. It would be beneficial if the authors provided more details regarding the computational aspects of training and inference.\n    \n- The handling of partial observability in the framework is not adequately explained. Naive extensions to partially observed Markov Decision Processes (POMDPs) may lead to unreliable estimations in diffusion models. Further elaboration on how the framework addresses partial observability would be valuable.\n    \n- The paper lacks clarity on how uncertainty estimates will manifest in highly stochastic domains. It would be insightful to include experiments on a benchmark dataset with increasing levels of stochasticity to demonstrate the framework's performance under such conditions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: More details regarding the computational aspects of training and inference.\n\nA1: Thanks for the suggestion! Yes, the framework with conformal prediction is a little slower than the framework without conformal prediction. Regarding the computational aspects of training and inference, as the supplementary material states, we train our model for 200 epochs with a batch 521 size of 256, utilizing a single GTX 1080 Ti GPU for computation. It was measured that the computational framework with the addition of conformal prediction increased the training time by 18\\% over the computational framework without conformal prediction. This is due to the fact that to compute reconstruction nonconformity. When making inferences, we use the same amount of time as the other frameworks and do not add extra overhead.\n\n>Q2: Further elaboration on how the framework addresses partial observability.\n\nThanks for pointing this out. Addressing partial observability in the CP framework applied to diffusion models requires careful modeling of observability, employing appropriate state estimation techniques, and potentially incorporating informative features. By considering these approaches, the framework can better handle partial observability and provide more reliable estimations and uncertainty assessments.\n\nWe model partial observability by incorporating relevant information into the model. This involves integrating additional features and conditions that capture observable aspects of the diffusion process into the diffusion model. By incorporating these observations, the model can better account for the uncertainties arising from partial observability.\n\nBy considering the uncertainty associated with partial observability, the resulting predictions would be more reliable and better reflect the inherent limitations of the diffusion models in such scenarios. This would provide decision-makers with a better understanding of the confidence and uncertainty associated with the predictions, particularly in situations where partial observability can significantly impact the accuracy of predictions.\n\n>Q3: How uncertainty estimates will manifest in highly stochastic domains?\n\nGood suggestions. Actually, the Maze2D environment is a stochastic environment. The data is generated by selecting goal locations at random. We have already shown the results for Maze2D in **Tables 1 and 2 in the main text**.\n\nNow, we modify the D4RL environment to introduce a random noise $\\epsilon_1 \\sim N(0, 1)$ and $\\epsilon_2 \\sim N(0, 0.5)$ add to the observations. **The framework's performance is shown in the attached Table 2 (please refer to the attached .pdf in the Summary Rebuttal).** We can see the stochastic D4RL is more challenging but PlanCP is still effective and can still quantify the uncertainty.\n\n>Q4:Some thoughts on adapting the CP framework to other conditioning approaches, such as temporal condition guidance mechanisms.\n\nA4: Thanks for the suggestions. Adapting the conformal prediction (CP) framework to other conditioning approaches, such as temporal condition guidance mechanisms, would require careful consideration of the specific requirements and characteristics of the conditioning approach. To adapt the CP framework to temporal condition guidance mechanisms, one approach could be to incorporate the temporal information into the conformal score calculation. Another approach could be to use a hybrid approach that combines CP with other methods that are specifically designed for temporal condition guidance, such as Kalman filters or particle filters. These methods are commonly used in control and tracking applications, where the goal is to estimate the state of a system based on noisy measurements and temporal dependencies."}, {"review_id": "otMbooOVpx", "reviewer": "Reviewer_8c3j", "summary": "In this paper the authors study how to measure uncertainty in a planner with a generative diffusion model, and how to reduce uncertainty during prediction. Results are presented on several MDPs.", "questions": "n/a", "limitations": "not an adequate discussion", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "+ coupling planners to diffusion models is an interesting idea worth exploring.\n+ measuring and reducing uncertainty is a good goal.", "weaknesses": "- I found this paper really hard to read and follow. Although the paper appears to address planning, the text refers to expert trajectories the \"planner\" is supposed to reconstruct. It is not clear if the task is planning or imitation learning. Various aspects are not well defined, for example, how is a quantile directly a loss function? How do you optimize a quantile? At the same time there is a lot of repetition of reconstruction nonconformity and reward conformity. Multiple experimental tasks, such as \"test time flexibility\", are introduced without adequate discussion of why this approach should be applied to them. The results are in tables that are almost unreadable and there is not enough discussion of the meaning of the results.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: It is not clear if the task is planning or imitation learning.\n\nThanks for pointing this out. We aim to address planning with uncertainty-aware diffusion models. We involve generating a sequence of actions to achieve a desired goal or optimize an objective. We focus on decision-making and determining the optimal course of action based on a given model or environment. Although we use expert data, we are more like offline reinforcement learning (RL) rather than imitation learning. Planning algorithms often utilize search algorithms, optimization techniques, or reinforcement learning methods to find the best action sequence. For more details about such kind of planner, please refer to prior work Diffuser (Janner et al Planning with Diffusion for Flexible Behavior Synthesis)\n\n>Q2: How is a quantile directly a loss function? How do you optimize a quantile? \n\nWe aim to minimize the size of the quantile, concentrating the probability mass into a smaller area, therefore reducing uncertainty.  After adopting the differentiable ranking and sorting techniques (Cuturi et al. Differentiable ranking and sorting using optimal transport, NeurIPS 2019), the soft quantile function is differentiable. We can use gradient descent to update the dynamics model. \nSpecifically, Let $X$ be a random variable that has a smooth density function $f$. Let $w = w(p)$ be the $p$-th quantile. Then the first derivative of the quantile function is\n$$\n\\frac{dw}{dp}=\\frac{1}{f(w(p))}\n$$\nprovided that the denominator is not zero. The second derivative is\n$$\n\\frac{d^2w}{dp^2}=-\\frac{f'(w)}{f(w)}(\\frac{dw}{dp})^2=\\frac{-f'(w)}{f(w)^3}\n$$\n\n>Q3: Repetition of reconstruction nonconformity and reward conformity.\n\nA3: Thanks for pointing this out. We will make the writing more clear and concise. We'll also organize the paper better to make it easier to read and understood.\n\n>Q4: Discussion of why this approach should be applied to experimental tasks such as \"test time flexibility\".\n\nA4: Thanks for pointing this out. We will update the writing to add the following explanation. In order to evaluate the ability to generalize to new test-time goals, we run a \"test time flexibility\" evaluation. \"Test time flexibility\" is useful to evaluate the capability of the planner as the Diffuser (Janner et al) paper shows. \n\nLong-horizon multi-tasking planning is an important capability of the planner that we need to evaluate experimentally\n\nOffline Reinforcement Learning allows us to evaluate the capacity of our method to recover an effective single-task controller from heterogeneous data of varying quality, which is useful for planning.\n\nHence, we demonstrate our framework has a number of useful properties and is particularly effective in offline control settings that require long-horizon reasoning and test-time flexibility.\n\n>Q5: Interpretation of the results in tables.\n\nThe \"Rewards\" serve as metrics for evaluating the model. The \"mean\" and \"standard deviation (std)\" are statistical measures of the rewards. A higher \"mean\" indicates better performance, while a smaller \"std\" reflects lower uncertainty. The \"Uncertainty\" is represented by the prediction interval in conformal prediction. Interval is the result of the upper bound of uncertainty minus the lower bound. A smaller \"interval\" indicates reduced uncertainty. Additionally, a higher \"overall interval\" is desired, as it signifies a higher reward prediction.\n\nThe concept of \"Coverage“ (”Validity\"), discussed in Section 4.6 of the paper, evaluates the validity of a prediction sample in the test dataset ($D_{test}$) by checking if it falls within the uncertainty interval computed using the calibration dataset ($D_{cal}$). A higher coverage is desirable, indicating better performance. The \"Credibility\", defined as the p-value in a prediction set, is obtained by conducting a T-test between the calibration trajectory and the predicted trajectory. The provided equation calculates the T-value,\n$$\nt = \\frac{\\mu_1-\\mu_2}{\\sqrt{\\frac{\\sigma_1^2}{n_1} + \\frac{\\sigma_2^2}{n_2}}}\n$$\nwhere $\\mu_1$ is the mean of calibration trajectory samples, $\\mu_2$ is the mean predicted trajectory samples, $\\sigma^2_1$ is the variance of calibration trajectory samples, $\\sigma^2_2$ is the variance of predicted trajectory samples, $n_1$ is the sample size of calibration trajectory samples, and $n_2$ is the sample size of predicted trajectory samples. Higher credibility values are preferred, indicating better performance."}, {"review_id": "3oaedPvgUo", "reviewer": "Reviewer_BQNc", "summary": "The paper combines diffusion modeling and conformal prediction to predict state-action trajectories with uncertainty quantification. The proposed approach then performs uncertainty-aware model-based planning with strong results on several established offline RL benchmarks.", "questions": "I am confused about the role that the expert trajectories play in this approach. The approach is evaluated on datasets without labeled expert trajectories, i.e. the datasets include suboptimal demonstration data. But the Algorithm 1 block references expert trajectories, and I'm not sure why? It seems to me that the approach doesn't really assume optimal demonstrations.\n\nRelated to the above, I'm not sure what Algorithm 1 is actually showing. The \"result\" line indicates that the output is a prediction $\\tau$ with associated uncertainty $C$, but L1-9 describe the complete training loop and there aren't really details on inference in this figure.\n\nI think a brief definition of the $Quantile$ operator would be useful when it's first introduced in Equation 9.\n\nI feel like Equation 5 and the accompanying description of transformers is not necessary to include in the main body, and working to include more of the actual experimental results (currently in the appendix) would be a better use of space in my opinion.\n\n===\nI have read the rebuttal and still lean towards acceptance.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The approach has fairly strong results on several established offline RL benchmarks and outperforms several state-of-the-art baselines. The results show that uncertainty-aware planning does perform better than simply rolling out the underlying diffusion model alone.\n\nThe introduction/abstract are well-written and motivate the approach well. Most of the writing is quite good (some feedback on parts I think are unclear below). The figures are also well done.\n\nTo my knowledge, the proposed approach is novel and technically sound. I think working to incorporate more ideas from uncertainty quantification into offline RL methods is worth investigating deeper as well.", "weaknesses": "There are no comparisons to other uncertainty quantification methods. Even within the offline RL literature, several existing methods [1][2][3] rely on uncertainty estimates and could be compared to. I think at the very least a comparison to popular uncertainty estimation methods in the model-based RL literature (e.g. bootstrapping disagreement with ensembles, predicting variances) could be helpful.\n\nI struggle with Sections 3.3 and 3.4 a bit and think they could use some improvements clarity-wise.\n\nWhile the approach is novel, it is a fairly straightforward application of conformal prediction to offline model-based planning. While I don't think the result is groundbreaking, it is certainly interesting enough to merit acceptance in my view.\n\n[1] Zhan, Xianyuan, Xiangyu Zhu, and Haoran Xu. \"Model-based offline planning with trajectory pruning.\" arXiv preprint arXiv:2105.07351 (2021).\n\n[2] Yu, Tianhe, et al. \"Mopo: Model-based offline policy optimization.\" Advances in Neural Information Processing Systems 33 (2020): 14129-14142.\n\n[3] Kidambi, Rahul, et al. \"Morel: Model-based offline reinforcement learning.\" Advances in neural information processing systems 33 (2020): 21810-21823.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Q1: Comparisons to other uncertainty quantification methods Zhan et al [1], Yu et al. [2], Kidambi et al [3].\n\nA1:Thanks for pointing this out. We will make the writing of Sections 3.3 and 3.4 more concise with clearer expression.\n\nZhan et al. [1], Yu et al. [2], and Kidambi et al. [3] have made significant contributions to model-based offline RL. These works are noteworthy because even without utilizing techniques like conformal prediction, they are able to measure uncertainty effectively.\n**We compare our method with Zhan et al [1], Yu et al. [2], and Kidambi et al [3]. The results are shown below as the attached Table 3 (please refer to the attached .pdf in the Summary Rebuttal)**. The results show that our method achieves smaller uncertainties. We can see that the performances (rewards) of these baselines are lower because they do not use the diffusion model as a dynamics model. Despite MOPO (Yu et al. [2]) and Kidambi et al [3] have an explicit parametric representation of uncertainty, the uncertainty (std and interval) is still large. Also, MOPO and Kidambi et al [3] have a limitation: it assumes that the data conforms to a Gaussian distribution.\n\n>Q2: The role that the expert trajectories play in this approach\n\nA2: We do not assume optimal demonstrations and the datasets include suboptimal demonstration data, which is the same setting as offline reinforcement learning. Different from imitation learning, an offline reinforcement learning algorithm applied to a dataset collected by a suboptimal non-learning-based algorithm can still result in a reasonable policy (sometimes even outperforms the behavior agent used to collect the data).\n\n>Q3: Clarification of Algorithm 1\n\nA3: Sorry for the confusion. We get predicted trajectories in L 3 (training) and L 11 (calibration). We get associated uncertainty in L 5 (training) and L 13 (calibration). During inference, we also get trajectories predicted by our planner and compute the associated uncertainty with Eq. 15.\n\n>Q4: A brief definition of the Quantile operator\n\nA4: Thanks for the suggestions. \nWe follow the same mathematical definition of the traditional Quantile function:\nThe quantile function, also known as the inverse cumulative distribution function, is a mathematical function that maps a probability to the corresponding value in a random variable's distribution. Formally, for a random variable $X$ with cumulative distribution function $F(x)$, the quantile function $Q(p)$ is defined as:\n$$\nQ(p) = \\inf\\lbrace x : F(x) \\ge p \\rbrace,\n$$\nwhere $p$ is a probability between 0 and 1, and \"inf\" denotes the infimum, or greatest lower bound, of the set of values $\\lbrace x: F(x) \\ge p \\rbrace$.\n\nIn other words, the quantile function gives the value $x$ such that the probability of observing a value less than or equal to $x$ in the distribution of $X$ is at least $p$. \n\n>Q5: Equation 5 and the accompanying description of transformers\n\nA5: Thanks for the suggestions. We will update this part in the new version. We will include more of the actual experimental results instead of Equation 5 and the accompanying description of transformers."}]

Author rebuttal:
We are grateful to the reviewers for their valuable feedback on our work. Thank you for the many positive comments: (i) acknowledging the novel use of conformal prediction in diffusion models (all reviewers), (ii) noting the reasonably well-written and clear exposition (gfdz, wAeC, BQNc), (iii) recognizing the  importance of a robust framework for quantifying and reducing uncertainties in diffusion dynamics models (all reviewers), and (iv) noting that our approach consistently outperforms the baseline methods across the metrics identified (zTTG, wAeC, BQNc).  We are happy to hear the reviewers find our direction interesting and potentially impactful.  As noted by Reviewer zTTG, we hope, with some work, this paper can be a great contribution to the community. 

Below are some important questions:

>Q1: Clarification of Contributions

A1: 
1. Instead of developing a method for sequential decision-making that achieves the highest performance (reward) on D4RL, we are **the first** to augment the `diffusion` RL algorithm with uncertainty awareness. 

2. The progress in the D4RL benchmark is rapidly changing, but at the time of our paper submission, Diffuser (Janner et al.) was the state-of-the-art open-source Diffusion RL algorithm. The original Diffuser algorithm was evaluated only 100 times, which is insufficient to capture uncertainty. To address this limitation, we re-run 1000 evaluations of the Diffuser algorithm on our setup and present the performance results.

3. We focus on finding better solutions for reducing the uncertainty of sequential decision-making than before.

4. The novelty of this work lies in using conformal prediction to (i) accurately measure the uncertainty for conditional diffusion planners, and **(ii) reduce the uncertainty of diffusion planners by including a quantile loss term in the training**. Previous CP literature can only measure the uncertainty but cannot reduce it. For experimental results, we emphasize that we accomplished **reducing uncertainty without sacrificing performance**.

5. The application of CP to the diffusion dynamics model is novel and gives an accurate quantification of uncertainty for the predictions of the model with statistical guarantees. This can be useful for decision-making in scenarios where accurate predictions are critical. **There is no prior work that attempts to introduce conformal prediction on diffusion models**.

6. The experimental results also show that we can **reduce the uncertainty** of the model without compromising performance.

>Q2: More explanation for the experimental results

A2: *We have added 3 tables* **(attached to the Summary Rebuttal)***. Please refer to the attached .pdf file for more details.*

Table 1: Comparison of our approach with MOPO. PlanCP has lower uncertainty than previous methods for most tasks.

Table 2: Uncertainty of Stochastic D4RL. We modify the D4RL environment to introduce a random noise $\epsilon_1 \sim N(0, 1)$ and $\epsilon_2 \sim N(0, 0.5)$ add to the observations. The framework's performance is shown in the attached Table 2. We can see the stochastic D4RL is more challenging but PlanCP is still effective and can still quantify the uncertainty.

Table 3: Comparison of our approach with other model-based RL. PlanCP has lower uncertainty than previous methods for most tasks

>Q3: The definition of the Quantile operator

A3: We follow the same mathematical definition of the traditional Quantile function:
The quantile function is a mathematical function that maps a probability to the corresponding value in a random variable's distribution. Formally, for a random variable $X$ with cumulative distribution function $F(x)$, the quantile function $Q(p)$ is defined as:
$$
Q(p) = \inf\lbrace x : F(x) \ge p \rbrace,
$$
where $p$ is a probability between 0 and 1, and "inf" denotes the infimum, or greatest lower bound, of the set of values $\lbrace x: F(x) \ge p \rbrace$.

In other words, the quantile function gives the value $x$ such that the probability of observing a value less than or equal to $x$ in the distribution of $X$ is at least $p$.

>Q4: How to update the dynamics model and differentiate through the Quantile function?

A4: After adapting the differentiable ranking and sorting techniques (Cuturi et al. Differentiable ranking and sorting using optimal transport, NeurIPS 2019), the soft quantile function is differentiable. We can use gradient descent to update the dynamics model.

>Q5: Discussion of why this approach should be applied to experimental tasks such as "test time flexibility".

A4: Thanks for pointing this out. We will update the writing to add the following explanation. In order to evaluate the ability to generalize to new test-time goals, we run a "test time flexibility" evaluation. "Test time flexibility" is useful to evaluate the capability of the planner as the Diffuser (Janner et al) paper shows. 

Long-horizon multi-tasking planning is an important capability of the planner that we need to evaluate experimentally

Offline Reinforcement Learning allows us to evaluate the capacity of our method to recover an effective single-task controller from heterogeneous data of varying quality, which is useful for planning.

Hence, we demonstrate our framework has a number of useful properties and is particularly effective in offline control settings that require long-horizon reasoning and test-time flexibility.

Meta-review: This paper considers the use of conformal prediction framework to explicitly represent uncertainty in learned dynamics model for sequential decision-making. Overall, this is a good direction and the proposed approach is reasonable. The paper is mostly well-written and showed good experimental results.

All reviewers' liked the general approach and agree on its merits. They also raised a number of questions and concerns. Some of them were addressed by authors' rebuttal resulting in score update from some reviewers'.

I recommend accepting the paper because it is a good direction and with the hope that this paper can inspire the NeurIPS community to use conformal prediction based uncertainty quantification for solving sequential decision-making tasks.

I strongly encourage the authors to improve the exposition of the paper and to incorporate the rebuttal discussion/results in the revised paper. Specifically
1. Please clarify the distinction between planning and RL very early in the paper. Clearly mention the paper's scope of contribution as per one of the reviewer comments.
2. Please mention that CP requires calibration data and discuss the practical aspects such as selection and amount from a practitioners' point of view.
3. You may be able to use CP methods to handle distribution shift (see https://arxiv.org/abs/1904.06019 and https://www.stat.berkeley.edu/~ryantibs/statlearn-s23/lectures/conformal_ds.pdf) to demonstrate the applicability of your method to changing rewards to address one of the reviewer comments. It would be great if you can add some results in the final paper. This is important to justify the planning claim of the paper. 
4. Please add all the new results and incorporate rebuttal discussion in the final paper.
Decision: Accept (poster)