Neural Stochastic Differential Equations for Uncertainty-Aware Offline RL

W3: Computational overhead and scalability to high dimensional domains

In all our experiments, sampling a single point around each training datapoint is enough for the distance-aware term $\eta_\phi$ to cluster the training dataset and capture distance-aware uncertainty accurately. However, we agree with the reviewer that the extra distance-aware training loss increases the computational complexity of our approach as the dimensionality increases. In the MuJoCo example in the paper, our method is computationally more efficient than the Gaussian ensemble as we train a single model to fit the data in contrast to an ensemble of $n$ models.

Q1: How is it guaranteed for the denominator in Eq. 3 to be non-zero?

In the experiments, we define $\mu_\phi$ as $\mu_{\phi}(s,a) = e^{NN_\phi(s,a)}$ ensuring that the output is always positve. Here, $NN_\phi$ is a neural network parametrized by $\phi$. We have clarified this design choice in the appendix of the revised paper.

Q2: What is the intuition for $\log(\det(\Sigma))$ in the context of $\mathcal{L}_{\text{data}}$ in Eq. 6?

The term $\log(\det(\Sigma))$ encapsulates the trade-off between accurately fitting the data and maintaining a model of appropriate complexity by acting as a regularizer that penalizes models with excessive variance. It helps prevent overfitting and promotes generalization, ensuring that the model remains both faithful to the data and as simple as possible.

Q3: Could you provide examples of "unmodeled dynamics" as mentioned in Lines 342–343?

For example, if the assumption of rigid body dynamics is incorrect, the residual term compensates for it. Besides, if the dynamics are not control-affine as assumed in (8), such residual term is required to learn the unmodeled part of the dynamics.

Q4: If the unmodeled dynamics are significant, would it still be beneficial to penalize the residual term?

Great question. If the unmodeled dynamics are significant, the training loss will indicate a failure to fit the data under such model assumption accurately. Then, one can tune the penalization terms to reduce the overreliance on the assumed model and improve the training loss. Such a hyperparameter tuning problem is also present when training an ensemble of Gaussian, and one needs to pick the number of models in the ensemble and the model architecture and size of each element in the ensemble to fit the data accurately.

We would like to thank the reviewer for their thought-provoking comments and questions that helped us improve our work in the rebuttal phase. Below, we first respond to their comments regarding the weaknesses and then their questions.

W1: Relying of state decomposition and having strong assumptions on the environment.

We emphasize that the position/velocity decomposition and rigid body dynamics are specific and minimal examples of prior knowledge that can be incorporated for the MuJoCo-based environments in the paper. These are typical environment choices for comparisons of offline model-based reinforcement algorithms. However, our physics-informed neural SDE model goes beyond rigid body dynamics and is even more relevant for various real-world applications, such as those involving aircraft, drones, or vehicles. In these scenarios, our framework can often leverage a range of a priori physical knowledge such as friction, gravity, and other constraints.

Specifically, a priori physics knowledge can encompass (1) state parametrization, (2) structural knowledge of the vector field (right-hand side of the differential equations), and (3) a priori constraints that are satisfied by some terms in the structural knowledge of the vector field. In its most general form, structural knowledge defines the drift term of the neural SDE as $f_{\theta}(s, a) = F(s, a, g_{\theta_1}, ..., g_{\theta_N})$ where $F$ is a known function that composes a set of unknown terms or functions $g_{\theta_i}(s, a)$. This structural knowledge typically comes from applying first principles techniques, Kinematics, Kinetics, or rigid body dynamics. Furthermore, we may have known constraints that must be satisfied by $g_{\theta_i}$ in the form of $G(s, a,g_{\theta_1}, ..., g_{\theta_N}) \leq 0$, where $G$ is known. In this case, we can train the neural SDE to encode all the prior physics knowledge while minimizing constraint violations via an augmented Lagrangian method formulation.

For instance, in the context of a vehicle operating at the limits of stability (racing or drifting), Kinematics and Kinetics can provide structural knowledge as $f_{\theta}(s, a) = M_{\phi}(s, a) F_{\theta}(s, a)$, where $M_{\phi}$ involves possibly unknown physical parameters like gravity, mass, inertia, and the vector $F_{\theta}$ represents the unknown and challenging-to-model forces resulting from the interaction between the ground and the vehicle. In addition, the friction law constrains even more $F_\theta$ through $F_x^2 + F_y^2 \leq \mu Fz$. This extra information is an example of physics knowledge we can apply to model learning.

In environments where images define the observation space, we can encode the images into a latent space whose time evolution follows a neural SDE with prior physics knowledge, if available. However, we emphasize that most offline model-based RL approaches do not consider environments with images as observation space. Besides, it is not straightforward whether Gaussian ensembles are the right choice for modeling under such high-dimensional and non-continuous space.

In environments without prior knowledge, the proposed neural SDE formulation is equivalent to a black box modeling under a time-varying assumption of the evolution of the observations.

W2: Experiments in simplified settings to interpret the design choices for uncertainty estimators

We have revised the paper and its appendix to provide additional details about the distance-aware uncertainty estimator, as well as extensive experiments to illustrate the properties of the estimator. Below is the summary of our changes:

• We reference prior work in which distance-aware uncertainty estimators are theoretically and empirically demonstrated (on toy examples) to provide a better uncertainty estimate than model ensembling and Monte Carlo dropout, state-of-the-art models used in offline model-based reinforcement learning. See Appendix B.2.
• We show on different dataset geometry of a 2-D system that our uncertainty estimator $\eta_{\phi}$ accurately clusters the training dataset and efficiently captures the notion of distance to the training dataset. These results provide further insights into the training mechanism of $\eta_{\phi}$ and how the design choices and desired properties encoded in (3) and (4) are empirically satisfied. See Appendix B.2 and Figure 6.
• We provide new experiments in MuJoCo to demonstrate how the distance-aware uncertainty estimator of the learned models efficiently captures epistemic uncertainty and enforces pessimism during policy learning. In contrast, when using the aleatoric uncertainty term to enforce pessimism, we show that the resulting policy performs poorly in the ground truth environment. See Appendix D and Figure 7.

Q5: Typos in Eq. (8) and Figure 1

About Equation 8: The first variable is $s_{vel}$ because $f_{\theta}^{\text{pos}}$ computes the derivative of position. Notice that Eq. (8) is the derivative of the state.

About Figure 1: We thank the reviewer for their attention to detail. Indeed, there are three typos in this figure: 1) $**s**'$ in uncertainty penalization should be $**s**$. 2) $\nu_{\phi}$ should be replaced with $\eta_{\phi}$, as in Eq. (8). 3) Epistemic uncertainty in uncertainty penalization should be $\eta_{\phi}$. We made corrections in our new submission accordingly.

• References
  1. Djeumou, F., Neary, C., Goubault E., Putot S., Topcu U. (2022). Neural Networks with Physics-Informed Architectures and Constraints for Dynamical Systems Modeling. Learning for Dynamics and Control Conference (L4DC).
  2. Liu, J., Lin, Z., Padhy, S., Tran, D., Bedrax Weiss, T., Lakshminarayanan, B. (2020). Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. Advances in neural information processing systems (NeurIPS).
  3. Van Amersfoort, J., Smith, L., Teh, Y. W., Gal, Y. (2020). Uncertainty estimation using a single deep deterministic neural network. International conference on machine learning (ICML).
  4. Zhang, H., Shao, J., He, S., Jiang, Y., Ji, X. (2023). DARL: distance-aware uncertainty estimation for offline reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence.
  5. Lu, C., Ball, P. J., Parker-Holder, J., Osborne, M. A., & Roberts, S. J. (2021). Revisiting design choices in offline model-based reinforcement learning. arXiv preprint arXiv:2110.04135.

Q2: What are the implications of using dataset distance as an uncertainty estimator?

The reviewer raises a great question about high-distance points (or clusters) that may be either "noise" or actual "stable" data. In both cases, we argue that our distance-aware uncertainty term would either classify them as low-uncertainty areas or high-uncertainty areas in a manner consistent with state-of-the-art approaches based on model ensembling. Suppose a point is the result of noise and does not belong to any cluster in the training dataset. In that case, the point has a low density in the dataset, and model ensemble techniques will provide high epistemic uncertainty due to the disagreement in the predictions. Similarly, since our distance-aware estimator can efficiently cluster the training dataset (see Appendix B.2 and Figure 6 for an illustration), those far-away points will have a high uncertainty estimate. Now, if the point is stable and not the result of noise, then by continuity, other "stable" points in the neighborhood of that point are in the training dataset and will form a new and separate cluster. Therefore, by the construction of the distance-aware uncertainty term, points in the cluster have low uncertainty estimates. If, for some unknown reason, we only have an isolated "stable" point in the training dataset, then due to that point's low mass in the training dataset, both the model ensemble and the distance-aware uncertainty estimators will provide high-epistemic uncertainty.

We also want to emphasize that existing work [2-4] has shown that MC dropouts or model ensembles, popular modeling techniques in offline model-based RL, are unaware of the distance between unseen samples and training datasets, even in simple toy examples. Theoretical and empirical results have been provided to illustrate the benefits of using a distance-aware uncertainty estimate as an epistemic uncertainty estimator.

Q3: Why isn't aleatoric uncertainty used in penalization?

We consider this a design choice. Our intuition is not to penalize an agent with respect to aleatoric uncertainty because aleatoric uncertainty is a property of the dynamical system, not of data. When the offline dataset lacks coverage in certain areas, penalizing with respect to epistemic uncertainty is a way to signal the agent that the model may not have sufficient knowledge of said areas.

As the reviewer suggests, there are many options for uncertainty estimators, especially in existing works that use Gaussian ensembles to learn dynamics models [5].

Q4: How critical is this particular distance-aware uncertainty estimator? What about other uncertainty surrogates?

In our global response, we provide an ablation study that evaluates the impact of the choice of uncertainty estimator to penalize/truncate rollouts during policy learning. More specifically, we look into two choices:

1. Distance-aware uncertainty estimate, which measures epistemic uncertainty, $\eta_{\phi}$, and
2. Aleatoric uncertainty $\sigma_{\phi}$.

Our results showcase that in random datasets, the distance-aware uncertainty estimates $\eta_{\phi}$ performs similarly to aleatoric uncertainty, except in walker2d, where the epistemic uncertainty results in significantly better policies. However, in medium-expert datasets, the epistemic uncertainty estimator is clearly the best, as the aleatoric uncertainty estimator can result in trajectories with very low returns. One possible explanation behind these results is related to data coverage. Random datasets have better data coverage, as the data-logging policies behave randomly, resulting in diverse behaviors. In comparison, medium-expert datasets have better quality data yet low coverage hence the need for estimating epistemic uncertainty is more evident. For more details, please see the document shared in our global response.

An uncertainty surrogate that estimates frequency of visitation may be unsuitable for continuous state space domains such as MuJoCo continuous control benchmarks. In addition, as [5] showcases, in offline model-based RL, surrogates based on learned dynamics models offer wide-ranging options. Our work follows a similar approach and learns uncertainty estimators as part of the dynamics model for more accurate next-state prediction. In addition, our work utilizes them to enforce conservatism in policy learning.

We would like to thank the reviewer for their positive comments, valuable criticism, and thought-provoking questions that helped us improve the presentation of our work.

First, we respond to the reviewer's comments and then their questions.

W1: Some of the main contributions listed by the authors are not representative of the actual content(...)

Could the reviewer clarify what is incomplete or misleading in the contributions we listed, please? We respond to their questions, as well, but are unsure which questions are related to this comment.

W2 and W3: The concept of uncertainty is underdeveloped in the manuscript(...) and (...) there is a lack of exploration into the inner workings of the method.

We thank the reviewer for raising these points. We have revised the paper and its appendix to provide additional details about the distance-aware uncertainty estimator and extensive experiments to illustrate the properties of the estimator. Below is the summary of our changes, which are provided in detail in the responses to Q1-Q4:

• We reference prior work in which distance-aware uncertainty estimators are theoretically and empirically demonstrated (on toy examples) to provide a better uncertainty estimate than model ensembling and Monte Carlo dropout, state-of-the-art models used in offline model-based reinforcement learning. See Appendix B.2.
• We show on different dataset geometry of a 2-D system that our uncertainty estimator $\eta_{\phi}$ accurately clusters the training dataset and efficiently captures the notion of distance to the training dataset. These results provide further insights into the training mechanism of $\eta_{\phi}$ and how the desired properties encoded in (3) and (4) are satisfied. See Appendix B.2 and Figure 6.
• Extensive experiments in MuJoCo demonstrate how the distance-aware uncertainty estimator efficiently captures epistemic uncertainty. We also show that such an estimator enforces pessimism during policy learning in contrast to using aleatoric uncertainty, which we show cannot be used to enforce pessimism. See Appendix D and Figure 7.

Q1: What are the differences between priori physics knowledge and conveniently parameterizing the state? What about applications other than MuJoCo? What about other types of physics knowledge?

Incorporating a priori knowledge of physics goes beyond simply parametrizing the state. Section 4.2 presents a general formulation for integrating prior physics knowledge. We then emphasized that equation (8) reflects a specific and minimal understanding of rigid body dynamics, which we apply specifically for benchmark MuJoCo environments commonly used in offline model-based RL. However, our physics-informed neural SDE model is even more relevant for various real-world applications, such as those involving aircraft, drones, or vehicles. In these scenarios, our framework can often leverage a range of a priori physical knowledge such as friction, gravity, and other constraints.

Specifically, a priori physics knowledge can encompass (1) state parametrization, (2) structural knowledge of the vector field (right-hand side of the differential equations), and (3) a priori constraints that are satisfied by some terms in the structural knowledge of the vector field. In its most general form, structural knowledge defines the drift term of the neural SDE as $f_{\theta}(s, a) = F(s, a, g_{\theta_1}, ..., g_{\theta_N})$ where $F$ is a known function that composes a set of unknown terms or functions $g_{\theta_i}(s, a)$. This structural knowledge typically comes from applying first principles techniques, Kinematics, Kinetics, or rigid body dynamics. Furthermore, we may have known constraints that must be satisfied by $g_{\theta_i}$ in the form of $G(s, a,g_{\theta_1}, ..., g_{\theta_N}) \leq 0$, where $G$ is known. In this case, we can train the neural SDE to encode all the prior physics knowledge while minimizing constraint violations via an augmented Lagrangian method formulation.

For instance, in the context of a vehicle operating at the limits of stability (racing or drifting), Kinematics and Kinetics can provide structural knowledge as $f_{\theta}(s, a) = M_{\phi}(s, a) F_{\theta}(s, a)$, where $M_{\phi}$ involves possibly unknown physical parameters like gravity, mass, inertia, and the vector $F_{\theta}$ represents the unknown and challenging-to-model forces resulting from the interaction between the ground and the vehicle. In addition, the friction law constrains even more $F_\theta$ through $F_x^2 + F_y^2 \leq \mu Fz$. This extra information is an example of physics knowledge we can apply to model learning.

We refer to the method outlined in [1] as an example of integrating general structural knowledge and constraints within the specific context of neural ordinary differential equations for modeling dynamical systems.

Thank you to the authors for their thorough answers to my concerns. The authors’ comments helped clarify some details and improve my understanding of the narrative of the paper. Overall, I think the paper is strong. It presents a well-studied method, and the authors have effectively addressed key details and the inner workings of their approach. This work opens many interesting avenues, which is a positive feature rather than a burden for the authors to develop every possible aspect of their method. That said, I believe there are two important points to clarify:

(W1): Contributions and Narrative Alignment

Apologies, I realized I was not clear in my previous comment. When I stated that “the contributions listed by the authors are not representative of the actual content presented in the paper,” my intention was to highlight a different focus. I believe this paper makes a valuable contribution; however, I disagree with the authors on what that contribution is.

In the contribution list, point (2) emphasizes the physics-governed dynamics while relegating the distance-aware uncertainty estimator to a secondary role within point (1). I would argue the opposite. The distance-aware uncertainty estimator is the standout idea in the revised manuscript, as it could potentially be applied to other problems. Meanwhile, the physics-aware feature seems more specifically engineered for the rigid body dynamics case, although it remains a meaningful contribution.

Evidence for this change of focus includes the theoretical guarantees, the central role of the uncertainty measure in the model, and the additional results and descriptions provided for the estimator. In contrast, the functional form of prior physics knowledge is less explored and appears tailored to this specific scenario. I understand that the prior physics component of the proposed method can be adapted to other cases. However, the level of detail and relevance of this aspect seems less significant compared to the distance-aware uncertainty estimator. I would like to know your thoughts on this, considering the development of the distance-aware uncertainty estimator in the paper and the incorporation of prior physics knowledge.

Epistemic and Aleatoric Uncertainty

Thanks to the author's responses, I now understand what initially confused me. However, I believe the use of uncertainty terms in the text requires greater care. I hope my argument here will be helpful.

There are two types of uncertainty discussed in this work: epistemic uncertainty (which can be reduced through learning) and aleatoric uncertainty (intrinsic noise in the data that cannot be learned). In this work, they appear at two levels, each with its own epistemic and aleatoric components.

• Distance-Aware Uncertainty Estimator Level: At this level, the estimator treats epistemic and aleatoric uncertainty from the data together as a single quantity. This is why I raised questions about aleatoric uncertainty (Q2 and Q3).

• Model Level ($h_{\phi}$ and $\sigma_{\phi}$): Here, the stochastic ODE model separates epistemic (provided by the distance estimator) and aleatoric uncertainty. However, the text entangles these two pairs of uncertainties.

The key distinction is that the agent is not directly exposed to the data uncertainty. Instead, the agent deals with model uncertainty, which separates epistemic and aleatoric components ($h_{\phi}$ and $\sigma_{\phi}$). In contrast, the distance-aware uncertainty estimator, which turns out to be the epistemic uncertainty of the stochastic ODE, does not have access to this separation and combines data epistemic and aleatoric uncertainty into one measure (as in the example of a point generated by noise, or being "stable"). While the separation of uncertainties in the distance estimator could be refined in future work, I think clarifying it in the current manuscript could help readers familiar with the literature on epistemic and aleatoric uncertainty.

Thank you for increasing your score! We appreciate your thought-provoking questions and comments, which led to discussions that eventually improved our paper in several parts, from the contributions to the limitations.

We agree with your comment that there are cases where the distance-aware uncertainty estimator $\eta_{\phi}$ may not separate aleatoric and epistemic uncertainty. Your explanation resonates with our perspective. Hence, our newly added discussion covers these cases.

Indeed, you are also correct that $\sigma_{\phi}$ is trained to capture the aleatoric uncertainty. Both $\eta_{\phi}$ and $\sigma_{\phi}$ are part of the neural SDE; hence, the trajectories generated by the model allow the agent to perceive both uncertainty sources, namely epistemic and aleatoric, respectively. NUNO utilizes the distance-aware uncertainty estimator to further enforce conservatism in policy training via uncertainty penalization and truncation. As our ablation study evidences, using the distance-aware uncertainty estimator $\eta_{\phi}$ to impose pessimism is more effective than the aleatoric uncertainty estimator $\sigma_{\phi}$. This is likely because the distance-aware uncertainty estimator represents the model uncertainty more accurately than the aleatoric uncertainty estimator, corresponding to the stochasticity of dynamics or noise in measurements.

Thank you again!

We thank the reviewer for their positive comments and fair criticism, which guided us to run an ablation study on our design choices and visualize their impact on toy domains. Below, we respond to the reviewer's comments.

W1: Lack of ablation for design decisions.

In our global response, we provide an ablation study that evaluates the impact of the choice of uncertainty estimator on policy learning. More specifically, we look into two choices:

1. Distance-aware uncertainty estimate, which measures epistemic uncertainty, $\eta_{\phi}$, and
2. Aleatoric uncertainty $\sigma_{\phi}$.

Our results showcase that in random datasets, the distance-aware uncertainty estimates $\eta_{\phi}$ perform similarly to aleatoric uncertainty, except in walker2d, where the epistemic uncertainty results in significantly better policies. However, in medium-expert datasets, the epistemic uncertainty estimator is the best, as the aleatoric uncertainty estimator can result in trajectories with very low returns. One possible explanation behind these results is related to data coverage. Random datasets have better data coverage, as the data-logging policies behave randomly, resulting in diverse behaviors. In comparison, medium-expert datasets have better quality data yet low coverage; hence, the need for estimating epistemic uncertainty is more evident. For more details, please see the document shared in our global response.

In addition, our global response provides experiments on different dataset geometry of a 2-D system to demonstrate that our distance-aware uncertainty estimator $\eta_{\phi}$ accurately clusters the training dataset and efficiently captures the notion of distance to the training dataset.

W2: Focus on low-dimensional domains

HalfCheetah, Hopper, and Walker2d environments in MuJoCo domains are standard evaluation benchmarks for offline model-based RL [1-4]. Therefore, we focus on these domains. Nevertheless, we acknowledge that datasets in the D4RL benchmarks may not reflect real-world characteristics from conservative data-logging policies. To address this, we also evaluate NUNO in NeoRL datasets designed accordingly. Our results show that NUNO outperforms all other approaches being assessed in terms of average among all datasets in both benchmarks.

W3: Limited scope on partial observability and image-based domains.

As the reviewer indicates, our discussion of this work's limitations already mentions that we do not address partially observable domains or domains with image-based observations. Extending our work to partially observable domains that do not have image-based observations is relatively more straightforward. We plan to extend our uncertainty-aware approach to image-based observations by learning dynamics on the latent space via a neural SDE, which is likely more challenging.

• References
  1. Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., ... & Ma, T. (2020). Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33, 14129-14142.
  2. Yang, Y., Jiang, J., Zhou, T., Ma, J., & Shi, Y. (2021). Pareto policy pool for model-based offline reinforcement learning. In International Conference on Learning Representations.
  3. Zhang, H., Shao, J., He, S., Jiang, Y., Ji, X. (2023). DARL: distance-aware uncertainty estimation for offline reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence.
  4. Sun, Y., Zhang, J., Jia, C., Lin, H., Ye, J. & Yu, Y.. (2023). Model-Bellman Inconsistency for Model-based Offline Reinforcement Learning. Proceedings of the 40th International Conference on Machine Learning. 202:33177-33194