SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 [70]	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
Maze2D 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 on Uncertainty-aware Planning
Trajectory optimization for a discrete-time dynamics system defined by  $s_{t+1} = f(\pmb{s}_t, \pmb{a}_t)$  involves identifying a sequence of actions  $\pmb{a}_{0:T}$  that maximizes an objective function  $\mathcal{I}$ , which is the sum of per time-step rewards or costs  $c(\pmb{s}_t, \pmb{a}_t)$ . The optimization problem can be represented as follows:
$$
\boldsymbol {a} _ {0: T} = \underset {\boldsymbol {a} _ {0: T}} {\arg \min } \mathcal {J} (\boldsymbol {s} _ {0}, \boldsymbol {a} _ {0: T}) = \underset {\boldsymbol {a} _ {0: T}} {\arg \min } \sum_ {t = 0} ^ {T} \gamma^ {t} c (\boldsymbol {s} _ {t}, \boldsymbol {a} _ {t}). \tag {10}
$$
Here,  $T$  represents the planning horizon.
While training, we utilize the uncertainty calculated from  $\{\tilde{r}_i\}$  as a cost function to optimize the parameters of the planner model. After computing the reconstruction nonconformity scores on our dataset and setting  $\hat{C} = \mathrm{Quantile}\left(\tilde{r}_1,\dots ,\tilde{r}_k;\frac{\lceil(k + 1)(1 - \alpha)\rceil}{k}\right)$ , we can form the uncertainty loss:
$$
\mathcal {L} (\theta) _ {\text {u n c e r t a i n t y}} = \hat {C}. \tag {11}
$$
After adopting the differentiable ranking and sorting techniques [69], the soft quantile function is differentiable. On the other hand, a diffusion probabilistic model is used to parameterize the learned gradient  $\epsilon_{\theta}(\tau^{n}, n)$  of the trajectory denoising process. The objective for training the  $\epsilon$  model can be expressed as Eq. 12:
$$
\mathcal {L} (\theta) _ {\text {r e c o n}} = \mathbb {E} _ {n, \epsilon , \tau^ {0}} [ \| \epsilon - \epsilon_ {\theta} (\tau^ {n}, n) \| _ {2} ], \tag {12}
$$
where  $n\sim \mathcal{U}\{1,2,\ldots ,N\}$  represents the diffusion timestep.  $\epsilon \sim \mathcal{N}(0,\mathbf{I})$  represents the noise target. Finally,  $\pmb{\tau}^{n}$  denotes the trajectory  $\pmb{\tau}^{0}$  that is corrupted with noise  $\epsilon$
In the context of learning, we hope to use conformal prediction to construct prediction intervals 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, we can make more informed decisions and avoid overly optimistic or pessimistic estimations. Finally, the learning objective for planner training is
$$
\mathcal {L} _ {\text {p l a n}} = \mathcal {L} (\theta) _ {\text {r e c o n}} + \lambda_ {\text {u n c e r t a i n t y}} \mathcal {L} (\theta) _ {\text {u n c e r t a i n t y}}. \tag {13}
$$
Given predicted reward  $\hat{r}$  and ground-truth reward  $r$ , the learning objective for reward model training is
$$
\mathcal {L} _ {\text {r e w a r d}} = \| \hat {r} - r \| _ {2}. \tag {14}
$$
The reward model  $R(\cdot)$  is used as the condition for the diffusion model as Eq. 4 shows. We can obtain reward  $r$  returned by the environment or learned reward model. On the calibration dataset  $D_{cal}$ , we compute the reward conformity scores and set  $Q_i^{\alpha}, Q_i^{1 - \alpha}$  and prediction interval  $\hat{q}$  using conformal prediction as Eq. 15.
$$
Q ^ {\alpha} \left(r _ {k + 1}\right) = \text {Q u a n t i l e} \left(r _ {1}, \dots , r _ {k}; \frac {\lceil (k + 1) (1 - \alpha) \rceil}{k}\right),
$$
$$
Q ^ {1 - \alpha} \left(r _ {k + 1}\right) = \text {Q u a n t i l e} \left(r _ {1}, \dots , r _ {k}; \frac {\lceil (k + 1) \alpha \rceil}{k}\right), \tag {15}
$$
$$
\hat {q} (r _ {k + 1}) = \left[ Q ^ {1 - \alpha} (r _ {k + 1}), Q ^ {\alpha} (r _ {k + 1}) \right].
$$
Once the prediction interval is constructed, we assess its performance on the test set  $D_{test}$ .
Table 2: 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.
Dataset	Diffuser [9]	PlanCP
Rewards ↑	Uncertainty ↓	Interval q ↓	Rewards ↑	Uncertainty ↓	Interval q ↓
Maze2D U-Maze	113.9 ± 3.1	[121.83, 147.19]	25.36	116.4 ± 3.2	[121.11, 142.85]	21.74
Maze2D Medium	121.5 ± 2.7	[132.77, 141.85]	9.08	128.5 ± 2.5	[132.77, 142.99]	10.22
Maze2D Large	123.0 ± 6.4	[92.11, 163.23]	71.12	130.9 ± 5.6	[97.75, 168.46]	70.71
Average	119.5	125.3
3.5 Our Proposed Procedures
We first iterate over the training set  $D_{train}$  and for each expert trajectory  $\pmb{\tau}_i^E$ , obtain a predicted trajectory  $\pmb{\tau}_i$  using the learned dynamics model  $\mathcal{P}_{\theta}$ . Then we compute the reconstruction nonconformity score  $\tilde{r}$  for each  $\pmb{\tau}_i^E$  and construct the uncertainty  $C_i$ . The loss  $\mathcal{L}_{plan}$  and  $\mathcal{L}_{reward}$  are then computed using Eq.13 and Eq.14, respectively. Finally, the dynamics model  $\mathcal{P}_{\theta}$  and reward model  $R(\cdot)$  are updated using the computed loss.
Next, we iterate over the calibration set  $D_{cal}$ , and for each predicted trajectory  $\tau_{i}$ , obtain the reward conformity score  $r_{i}$  and construct the predicted interval  $\hat{q}_{i}$ .
Finally, we verify the predicted interval using the test set  $D_{test}$  to evaluate its performance. The whole procedure is shown in Alg. 1.
In our experiments, we aim to answer the following questions: (1) How to construct uncertainty intervals that can capture the dynamics model's prediction error and provide a 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, we evaluate 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 a wide range of diffusion dynamics models. For illustrative purposes, we demonstrate our framework using a typical Diffuser [9]. Our primary objective is to assess and minimize the model's uncertainty, while performance improvement is a 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  $\lambda_{uncertainty} = 5$  and the failure probability to  $\alpha = 0.1$ . To optimize the model, we use the Adam [73, 74] optimizer with a learning rate of  $2 \times 10^{-4}$ . We train the diffusion dynamics model on the training set  $D_{train}$  for  $2 \times 10^{5}$  iterations. As we lack ground-truth for evaluation, we generate the evaluation datasets  $D_{cal}$  and  $D_{test}$  during the evaluation process. We partition the evaluation data into  $20\%$  for  $D_{cal}$  and  $80\%$  for  $D_{test}$ . The uncertainty measure  $\hat{q}$  is derived from the  $20\%$  calibration set through conformal prediction.
4.2 Long Horizon Task Planning
We use Maze2D to evaluate long-horizon planning.  $\hat{q}$  is a measure of uncertainty for the rewards obtained by a planned trajectory. For most variants of the Maze2D environment, we have a smaller uncertainty interval, especially for Maze2D U-Maze, as Tab. 1 and Tab. 2 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 a feasible path.
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.
Dataset	Environment	Rewards ↑	Diffuser [9] Uncertainty ↓	Interval q ↓	Rewards ↑	PlanCP Uncertainty ↓	Interval q ↓
Medium-Expert	HalfCheetah	85.82 ± 6.91	[82.20, 91.55]	9.34	87.44 ± 6.51	[83.12, 91.81]	8.68
Medium-Expert	Walker2d	108.22 ± 8.52	[108.48, 109.58]	1.10	108.31 ± 8.07	[107.86, 110.43]	2.56
Medium	HalfCheetah	46.10 ± 0.96	[45,05, 47.34]	2.29	46.26 ± 1.35	[44.64, 47.97]	3.33
Medium	Walker2d	75.73 ± 19.59	[40.76, 87.72]	46.96	77.08 ± 18.21	[44.51, 88.46]	43.95
Medium-Replay	HalfCheetah	37.34 ± 6.57	[27.92, 41.58]	13.66	37.65 ± 6.29	[29.25, 42.23]	12.98
Medium-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 a 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 of the returns.
Overall, experimental results show that our framework can construct uncertainty intervals that can capture the dynamics model's prediction error and provide a 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 a diffusion-based 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 a meticulous analysis of how conformal prediction compares to the extensive body of research in planning under uncertainty, including non-parametric models. Our work focuses primarily on proposing an 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 on real-world applications, such as autonomous driving or robotics, where reliable uncertainty estimates are crucial for safety-critical systems, would be an interesting avenue for future research.
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. DGE-1656518. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.
[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] Ilya 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 on 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.07219, 2020.