SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiankai Sun
Stanford University jksun@stanford.edu
Yiqi Jiang
Stanford University yqjiang@stanford.edu
Jianing Qiu
Imperial College London jianing.qiu17@imperial.ac.uk
Parth Talpur Nobel
Stanford University ptnobel@stanford.edu
Mykel Kochenderfer
Stanford University mykel@stanford.edu
Mac Schwager
Stanford University schwager@stanford.edu
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.
Uncertainty is a major challenge in robotics, arising from a variety of sources, including errors in modeling, sensing, planning, and control [1, 2, 3]. Understanding the effect of uncertainty is important for developing successful motion strategies, as it allows robots to 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.
The dynamics model is especially critical for model-based planning [7] and reinforcement learning (MBRL) [8]. Due to 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, 10]. 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 a distribution over multiple actions rather than a deterministic choice of a single action [11]. Furthermore, diffusion models hold the potential to be applied to Offline Reinforcement Learning, which aims to learn effective strategies using offline datasets without exploration [12]. However, existing model-based algorithms struggle when encountering out-of-distribution (OOD) actions or states [9]. A crucial element absent from these methods is the adequate management of uncertainty in the offline scenario [13, 14]. Uncertainty-aware dynamics models that convey algorithmic confidence are valuable if they possess accurate calibration and sharpness [15]. 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 to suboptimal decision-making [17, 18]. Therefore, to ensure accuracy and optimal decision-making, it is crucial to utilize uncertainty-aware 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 over traditional methods for measuring uncertainty. For example, CP is model-agnostic [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 is complex and not well understood. Additionally, CP allows for customization of the desired confidence level, providing the ability to 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 it a particularly useful approach for planning in dynamic environments. Lastly, CP is distribution-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 to planning and propose P1anCP, adopting a practical and effective conformal prediction uncertainty estimation method for planning algorithms. Empirically, we observe that P1anCP substantially improves model stability. In addition, P1anCP reduces uncertainty without sacrificing performance on datasets with demonstrations collected from experts. P1anCP 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 a planning method, we will generally stick with the planning terminology for consistency with the literature.
- 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 a conformal predictor that guarantees trustworthy and accurate uncertainty estimates.
- We extensively evaluate our approach, and the experimental results demonstrate that it exhibits lower uncertainty without compromising performance across multiple benchmarks including Offline Reinforcement Learning.
2.1 Conformal Prediction
Conformal Prediction (CP) [19, 28, 29, 30, 21, 24, 31, 32, 33] is effective for measuring the uncertainty associated with predictions produced by black-box models. As a variant of conformal
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 transductive 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 [38]. Examples include Cross Conformal Predictor (CCP) [39], Bootstrap Conformal Predictor (BCP) [19] and the generalized Aggregated Conformal Predictor (ACP) [40]. In a recent study [41], conformal prediction was used to construct prediction regions for recurrent 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 RL controller but avoids uncertainty intervals. Tonkens et al. [43] 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 is 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.
In recent years, diffusion models have gained significant attention as a powerful class of generative models that formulates the data-generating process as a sequence of denoising steps [46, 47, 48, 11]. 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-planning perspective, Diffuser [9] applies a diffusion model as a trajectory generator. From the model-free policy-optimization perspective, Wang et al. [11] utilize a 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 [10] 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 Dynamics Model, and Model-Based Reinforcement Learning
Model-based reinforcement learning (MBRL) has emerged as a promising solution for real-world sequential decision-making problems with distinct advantages such as analytic gradient computation [55]. A significant body of literature focuses on learning a dynamics model and utilizing the learned model for policy learning through model-based planning [56, 57, 58, 59]. PILCO [60] utilizes a Gaussian process to model system dynamics, while World Models [61] and Dreamer [62] learn latent dynamics through two-stage and online planning, respectively. MGAIL [63] introduces a forward model in GAIL, but its discriminators output a uniform representation that is less transferable. PETS [64] combines probabilistic neural networks and deterministic neural networks to create uncertainty-aware models. To capture 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.
In this section, we describe our framework PlanCP (as Fig. 1 shows) to 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,  $i$  or  $k$  to denote different trajectories, and  $t$  to denote trajectory timestep.
Figure 1: PlanCP Framework: To update the dynamics model during training, we perform uncertainty-aware optimization using the training set  $D_{train}$ . During calibration, we construct the uncertainty interval  $\hat{q}$  based on  $D_{cal}$ , and then we verify the uncertainty interval using the test set  $D_{test}$ .
We consider a Markov Decision Process (MDP) represented by the tuple  $(\mathcal{S},\mathcal{A},\mathcal{T},r,\rho_0,\gamma)$  with state space  $\mathcal{S}$ , action-space  $\mathcal{A}$ , dynamics  $\mathcal{T}(s'|a,s)$ , reward function  $r(s,a)$ , initial state distribution  $\rho_0(s)$ , and discount factor  $\gamma \in (0,1)$ . We assume a stochastic policy  $\pi$  that maps a state to a distribution over actions. We define  $\pmb{\tau}$  as trajectories, which consist of sequences of state-action pairs  $(s_0,a_0,\dots ,s_T,a_T)$  generated by a policy  $\pi$ .  $\pmb{\tau}^{E}$  are expert trajectories generated by an expert policy  $\pi^E$ .  $T$  represents the time horizon. We use a two-dimensional array to represent the trajectory  $\pmb{\tau}$ :
$$
\boldsymbol {\tau} = \left[ \begin{array}{l l l l} \boldsymbol {s} _ {0} & \boldsymbol {s} _ {1} & \dots & \boldsymbol {s} _ {T} \\ \boldsymbol {a} _ {0} & \boldsymbol {a} _ {1} & \dots & \boldsymbol {a} _ {T} \end{array} \right], \tag {1}
$$
with each column corresponding to a 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 to the high dimensionality and the complexity of the model. One type of latent variable model approximates posterior by the forward or diffusion process  $q(\pmb{\tau}^n|\pmb{\tau}^{n-1})$ , which is modeled as a fixed Markov chain that progressively adds Gaussian noise to the data. Diffusion Probabilistic Models, as described in [46, 47], characterize the process of data generation as an iterative denoising procedure, formalized as the probability distribution  $p_{\theta}(\pmb{\tau}^{n-1}|\pmb{\tau}^n)$ . The denoising operation (as Eq. 2 shows) is performed in reverse of a forward process  $q(\pmb{\tau}^n|\pmb{\tau}^{n-1})$ , which progressively introduces noise and alters the underlying structure of the data.
$$
p _ {\theta} \left(\boldsymbol {\tau} ^ {0}\right) = \int p \left(\boldsymbol {\tau} ^ {N}\right) \prod_ {n = 1} ^ {N} p _ {\theta} \left(\boldsymbol {\tau} ^ {n - 1} \mid \boldsymbol {\tau} ^ {n}\right) d \boldsymbol {\tau} _ {1: N}, \tag {2}
$$
where  $p(\boldsymbol{\tau}^N) = \mathcal{N}(\boldsymbol{\tau}^N; \mathbf{0}, \mathbf{I})$  is a standard Gaussian prior and  $\boldsymbol{\tau}^0 \sim q(\boldsymbol{\tau}^0)$  is noiseless problem data. The reverse process is parameterized as Gaussian (Eq. 3):
$$
p _ {\theta} \left(\boldsymbol {\tau} ^ {n - 1} \mid \boldsymbol {\tau} ^ {n}\right) = \mathcal {N} \left(\boldsymbol {\tau} ^ {n - 1}; \mu_ {\theta} \left(\boldsymbol {\tau} ^ {n}, n\right), \boldsymbol {\Sigma} _ {n}\right), \tag {3}
$$
in which  $\mu_{\theta}$  and  $\Sigma$  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-designed condition function  $w(\cdot)$  is added to the diffusion model, which yields a planner  $\mathcal{P}_{\theta}$  described by Eq. 4:
$$
\mathcal {P} _ {\theta} (\boldsymbol {\tau}) \propto p _ {\theta} (\boldsymbol {\tau}) w (\cdot). \tag {4}
$$
where  $p_{\theta}(\tau)$  is the diffusion model's probability distribution, and  $w(\cdot)$  is a 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(\cdot)$  at all diffusion timesteps. We can use a learned reward function as a condition, as introduced in Sec. 3.2.
Algorithm 1: PlanCP: Conformal Prediction for Planning with Diffusion Dynamics Models
Result: predicted trajectory  $\pmb{\tau}$ , uncertainty  $C$ .
Data: Training set  $D_{train}$  with  $K_{train}$  samples, calibration set  $D_{cal}$  with  $K_{cal}$  samples, test set  $D_{test}$  with  $K_{test}$  samples, planner  $\mathcal{P}_{\theta}$ , expert trajectory  $\tau^{E}$ , and failure probability  $\alpha$ .
1 while not converged do  
2 for  $i = 0,1,2,\dots ,K_{train} - 1$  do  
3 Get predicted trajectory  $\pmb{\tau}_{i}$  corresponding to expert trajectory  $\pmb{\tau}_{i}^{E}\in D_{train}$   
4 Compute reconstruction nonconformity score as Eq. 6 for each  $\pmb{\tau}_{i}^{E}\in D_{train}$   
5 Construct uncertainty  $C_i$  as Eq. 9;  
6 Compute total loss  $\mathcal{L}_{plan}$  as Eq. 13 and  $\mathcal{L}_{reward}$  as Eq. 14;  
7 Update the dynamics model  $\mathcal{P}_{\theta}$  using loss  $\mathcal{L}_{plan}$  and the reward model  $R(\cdot)$  using loss  $\mathcal{L}_{reward}$ ;  
8 end  
9 end  
10 for  $i = 0,1,2,\dots ,K_{cal} - 1$  do  
11 Get predicted trajectory  $\pmb{\tau}_{i}$  as  $D_{cal}$ ;  
12 Compute reward conformity score  $r_i$  for each  $\pmb{\tau}_{i}\in D_{cal}$ ;  
13 Construct predicted interval  $\hat{q}$  as Eq. 15;  
14 end  
15 Verify the predicted interval using the test set  $D_{test}$ ;
3.2 Learned Reward Function
Learning a reward model is important because it allows an intelligent agent to estimate the expected rewards of a given action and state, which is useful for the agent to make informed decisions in an environment. Additionally, using a learned reward model can enable us to adapt to changes in the environment and adjust our predictions accordingly. However, learning an accurate reward model from data is challenging. Thus, we introduce a transformer-based reward model as Eq. 5, which can perform accurate long-term temporal predictions.
$$
\hat {r} _ {0: T} = R (\text {M u l t i H e a d} (Q = s _ {0: T}, K = s _ {0: T}, V = \boldsymbol {a} _ {0: T})),
$$
$\operatorname{MultiHead}(Q, K, V) = \operatorname{Concat}(\operatorname{head}_1, \operatorname{head}_2, \dots, \operatorname{head}_h) W^O$ ,
$$
\operatorname {h e a d} _ {j} = \operatorname {A t t e n t i o n} \left(Q W _ {j} ^ {Q}, K W _ {j} ^ {K}, V W _ {j} ^ {V}\right), \tag {5}
$$
Attention  $(Q,K,V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_f}}\right)V,$
where  $R(\cdot)$  is an MLP layer to map the multi-head attention feature to a predicted reward.  $d_{f}$  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  $\mathcal{P}_{\theta}$  can lead to 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 in machine learning [24]. This allows us to quantify the uncertainty associated with the predictions and provides more robust estimates of the model's performance.
We have split the dataset  $D$  into three parts —  $D_{train}$ ,  $D_{cal}$ , and  $D_{test}$  — for training, calibration, and testing, respectively. The planner  $\mathcal{P}_{\theta}$  and reward function  $R(\cdot)$  are learned from the  $D_{train}$  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
PlanCP
(a) Maze2D U-Maze
Diffuser (Failed)
PlanCP
(b) Maze2D Medium
Figure 2: Learned long-horizon planning: By learning the planning procedure, PlanCP is capable of conducting long-horizon planning with a specified start  $\Theta$  and goal  $\star$  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.
PlanCP
(c) Maze2D Large
the mean squared error of the learned trajectories from the expert trajectories. Reward conformity scores are exactly the reward returned by the environment for a given trajectory.
Training To be able to compute and reduce the uncertainty during training, we need to perform conformal prediction on the training set  $D_{train}$ . First, we compute the reconstruction nonconformity score as Eq. 6 for each trajectory  $\tau_i$  of the training set  $D_{train}$ . A lower reconstruction nonconformity score corresponds to better predictions.
$$
\tilde {r} _ {i} = \left| \left| \boldsymbol {\tau} _ {i} - \boldsymbol {\tau} _ {i} ^ {E} \right| \right| _ {2}, \tag {6}
$$
where  $\pmb{\tau}_i$  is the predicted  $i$ th trajectory while  $\pmb{\tau}_i^E$  is the ground-truth  $i$ th trajectory from the expert demonstration of  $D_{train}$ .
Calibration and Testing During calibration and testing, we use reward conformity score to quantify the uncertainty. Let  $r_0, r_1, \dots, r_k$  be  $k + 1$  rewards corresponding to  $\tau_0, \tau_1, \dots, \tau_k$  trajectories predicted by our planner  $\mathcal{P}_{\theta}$ . 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 a metric that reflects the planner's ability to predict the quality of plans. A higher reward conformity score corresponds to better predictions, indicating that the planner is more accurate in its estimations. Our objective is to quantify the uncertainty associated with the reward  $r_i$  based on a set of reference rewards  $\{r_0, \dots, r_k\} \setminus r_i$ . Specifically, we aim to construct a prediction interval  $C_i$  that contains  $r_i$  with probability at least  $1 - \alpha \in (0, 1)$ . Formally, we seek to find a valid prediction interval  $C_i$  that satisfies the following inequality:
$$
\mathbb {P} \left(r _ {i} \in C _ {i} \left(\left\{r _ {0}, \dots , r _ {k} \right\} \backslash r _ {i}\right)\right) \geq 1 - \alpha , \tag {7}
$$
where  $C_i(\{r_0, \dots, r_k\} \setminus r_i)$  denotes the prediction interval constructed using the reference rewards  $\{r_0, \dots, r_k\} \setminus r_i$ . The prediction interval  $C_i$  should provide a measure of the uncertainty associated with the reward  $r_i$ , and it should be constructed in such a way as to guarantee that with a probability of at least  $1 - \alpha$ , it contains  $r_i$ .
Lemma 3.1. Given  $k + 1$  exchangable random trajectories  $(\tau_0^E, \tau_1^E, \dots, \tau_k^E) \sim D_{train}$ , predicted trajectory  $(\tau_0, \tau_1, \dots, \tau_k)$ , corresponding (non)conformity scores  $d_0, d_1, \dots, d_k$ , and a failure probability  $\alpha \in (0, 1)$ , then
$$
\mathbb {P} \left[ d _ {i} \leq C _ {i} \right] \geq 1 - \alpha , \tag {8}
$$
where the prediction interval  $C_i$  are defined by
$$
C _ {i} = \text {Q u a n t i l e} \left(d _ {0}, \dots , d _ {i - 1}, d _ {i + 1}, \dots , d _ {k}; \frac {\lceil (k + 1) (1 - \alpha) \rceil}{k}\right). \tag {9}
$$
During training, we use reconstruction nonconformity score as  $d_{i}$ , while during calibration and testing, we use the reward conformity score as  $d_{i}$ . Note that if  $\alpha$  is less than  $1 / k$ , our confidence intervals become infinitely wide as we have insufficient data to make predictions with that accuracy.