Linear Mixture Distributionally Robust Markov Decision Processes

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

Q1: Practical motivation for linear mixture DRMDPs. When should we choose linear mixture DRMDPs over other frameworks?

A1: Thanks for the interesting question. To answer this question, we would like to lay out several practical motivations for our framework. Mixture models, in particular Gaussian Mixture Models (GMMs), are widely studied in fields of optimal control and RL. For example:

Model Predictive Control (MPC): [1] studies the density steering for discrete-time linear systems, a variant of optimal mass transport problem widely used in applications such as controlling a swarm of robots/drones to achieve a desired spatial distribution. They model the stochastic dynamical system using GMMs and derive its optimal control policy. [2] studies a stochastic MPC algorithm for linear systems subject to additive Gaussian mixture disturbances. They consider a vehicle control case study in which the vehicle must maintain its position on an ill-maintained road.

Meta RL: Meta learning is one way to increase the data efficiency of learning algorithms by generalizing learned concepts from a set of training tasks to unseen, but related, tasks. Recent works on Meta RL model the dynamics of a new system as a mixture of expert systems, thus realizing knowledge transfer. In particular, [3,4] model the dynamics of the expert systems using Gaussian process with latent embedding. And the dynamics of the new system is constructed by a mixture of Gaussian processes with the distribution over the latent embedding being the mixture weights, which should be learned from data. In experiments, their meta-learning models effectively generalizes to novel tasks under various environments such as Cart-pole & Double-pendulum swing-up, HalfCheetah and Highway-Intersection. Our framework could potentially enhance robustness in unseen tasks.

World Models for model-based RL: Model-based Deep Reinforcement Learning (RL) assumes the availability of a model of an environment’s underlying transition dynamics, which are stochastic in nature and oftentimes multimodal. [5,6] use Mixture-density Networks to construct World Models to capture the multimodality in dynamics. Then model-based RL methods are implemented based on the constructed World Models to solve various tasks such as car racing and Inverse Sine Wave.

The basis modes in our work can be the Gaussian components in GMMs, dynamics of expert systems, or mixture components in the Mixture-density Networks. One potential application of our framework would be the meta-RL. For example, assume we have constructed expert system dynamics (basis modes) from huge amounts of historical data, and would like to solve a new task whose dynamics is a mixture distribution of expert system dynamics. With limited data collected from the new task, there is high uncertainty in weighting parameter estimation. In such cases, the linear mixture DRMDP framework leverages the latent structure in dynamics of the new task, and thus is conceptually more favorable than other DRMDP framework. Other applications involve robust control of drones or vehicles to hedge against environmental uncertainties and distributional perturbations, and robust decision making in model-based RL by modeling the uncertainty in the World Model (specifically, in the weights). Moreover, it serves as an interesting future research direction to also incorporate the uncertainty in basis modes into consideration for robust decision making.

Q2: Linear mixture structure for multimodal distributions.

A2: Thank you for the interesting question. We highlight that the linear mixture structure can naturally model multimodal transitions. In particular, each basis mode can represent a component distribution with corresponding local maximum in the probability density function or probability mass function, such as gaussian distributions with different means and variances. And the mixture of basis modes capture the multimodality in the dynamics, say, for a given state-action pair, the next state follows multiple distinct patterns. For example, [5] discuss the multimodality in dynamics of many real-world RL applications. They introduce world models using mixture distribution networks to approximate the multimodal dynamics, which could potentially be adapted to our framework to construct the basis modes.

Q3: Comparison between the proposed distributionally robust method using the linear mixture ambiguity set and a method based on an (s,a)-rectangular ambiguity set.

A3: We can justify this point. In this work we mainly focus on establishing the framework, which is shown in Section 2 to have advantages in modelling the uncertainty in cases where existing frameworks fall short, from scratch and designing algorithms that achieve provable efficiency. The simple simulation environment is just to validate our theoretical results, i.e., to showcase the robustness of our proposed algorithms. We also would like to note that the linear mixture DRMDP assumes prior information, i.e. the access to basic modes, for algorithm design. It’s not fair to directly compare algorithms designed for different DRMDP settings because they either require other prior information or do not require prior information at all. Thus, we instead compare those two algorithms designed for the linear mixture DRMDP.

As an interesting future work, we would like to explore when it is appropriate to model the change in the dynamics as a linear mixture uncertainty set and how to construct the latent modes in complex environments like MoJuCo. In those real-world applications, we can have a fair comparison between different DRMDP frameworks. However, this requires considerable effort, which is beyond the scope of this work.

References:

[1] Balci et. al., 2024, Density steering of gaussian mixture models for discrete-time linear systems

[2] Engelaar et. al., 2024, Stochastic MPC for Finite Gaussian Mixture Disturbances with Guarantees

[3] Sæmundsson et. al., 2018, Meta reinforcement learning with latent variable gaussian processes

[4] Xu et. al., 2020, Task-agnostic online reinforcement learning with an infinite mixture of gaussian processes

[5] Sedlmeier et. al., 2021, Quantifying multimodality in world models

[6] Ha et. al., 2018, Recurrent world models facilitate policy evolution

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

Q1: Dependence of suboptimality on the size of the uncertainty set.

A1: Thank you for your question. For the case of TV divergence, the suboptimality bound does not depend on the radius $\rho$, as shown in Theorem 4.6 of our submission. For the case of KL divergence, the suboptimality bound is inversely proportional to the radius $\rho$ as shown in Theorem 4.11 of our submission. This dependence aligns with prior works on DRMDP with KL divergence defined uncertainty set (see Theorem 3 in [1] and Corollary 4.1 in [2]), arising from the geometry of KL divergence in analysis. For $\chi^2$ divergence, the suboptimality bound is proportional to $(\sqrt{\rho}H^3+H^2)$ as shown in Theorem 4.13 of our submission. For $\rho=O(1/H^2)$, this term is in order of $O(H^2)$ and our bound matches that for non-robust linear mixture MDP; for moderate $\rho$, our bound suggests that policy learning would be harder due to the complex geometry of the $\chi^2$ divergence uncertainty set. A similar pattern has also been identified by a prior work [3] on tabular DRMDPs with $\chi^2$ divergence defined uncertainty sets.

Q2: Dependence of your models on the number of data trajectories available for training.

A2: Our suboptimality bound can be translated into sample complexity bounds. To learn an $\epsilon$-optimal policy, we need sample sizes in order of $\tilde{O}(\frac{d^2H^4(C^{\pi^{\star}})^2}{\epsilon^2})$ for TV divergence, $\tilde{O}(\frac{d^2H^4(C^{\pi^{\star}})^2e^{2H/{\underline{\lambda}}}}{\rho^2\epsilon^2})$ for KL divergence, and $\tilde{O}(\frac{d^2(\sqrt{\rho}H^3+H^2)^2(C^{\pi^{\star}})^2}{\epsilon^2})$ for $\chi^2$ divergence.

Classic DRO approaches are expected to be effective in scenarios with limited data because they are explicitly designed to guard against uncertainty from insufficient data samples. In contrast, our current framework mostly focuses on uncertainty from dynamics shift, which is also common in real world applications. Moreover, our DRMDP framework hedging against dynamics shifts can also be adapted to achieve robustness w.r.t. parameter estimation uncertainty due to limited data. Actually this idea has been explored in the prior work [4], where they show DRMDP can be used to solve offline RL. Experiments therein demonstrate DRO approaches achieve significantly better performance than non-robust approaches with limited samples. Thus, we can expect the same is true for our framework. We leave a comprehensive examination for future study.

Q3: The impact of the robustness approach on scalability, especially compared to the standard (s, a) and d-uncertainty sets?

A3: From a theoretical perspective, our suboptimality bound does not depend on the size of state-action space. Thus even when the state-action space is of high dimensional and infinitely large, our approach scales well thanks to the linear mixture structure. d-rectangular DRMDPs also can scale, leveraging a different modeling on latent structure of dynamics and uncertainty set. The linear mixture MDP or linear MDP assumption on the nominal transition kernel naturally admits linear representation in terms of the dynamics or Q-function, thus enabling linear function approximation for scalability. While for $(s,a)$-rectangular DRMDPs, most existing model-based and model-free approaches [1,3,6] cannot scale, they either rely on frequency-based transition estimation or rely on a Q-table for robust Bellman update. Existing literature has also studied general function approximation for scalability in the context of $(s,a)$-rectangular DRMDPs. For example, [5] study $(s,a)$-rectangular DRMDPs with an infinitely large state-action space, they leverage a double layer function approximation routine for dual optimization approximation and robust Q-function estimation. Their proposed algorithm, Robust Fitted Q-Iteration, can scale in applications with continuous states and actions, like CartPole and Hopper. [2] study function approximation in $(s,a)$-rectangular DRMDPs from a theoretical perspective. They present suboptimality bounds depending on the complexity measure of the function class and independent on the size of the state-action spaces.

Q4: how difficult to obtain good estimates of the latent modes.

A4: Thanks for the interesting question. To answer this question, we would like to lay out several practical motivations for our framework. Mixture models, in particular Gaussian Mixture Models (GMMs), are widely studied in fields of optimal control and RL. For example:

Model Predictive Control (MPC): [1] studies the density steering for discrete-time linear systems, a variant of optimal mass transport problem widely used in applications such as controlling a swarm of robots/drones to achieve a desired spatial distribution. They model the stochastic dynamical system using GMMs and derive its optimal control policy. [2] studies a stochastic MPC algorithm for linear systems subject to additive Gaussian mixture disturbances. They consider a vehicle control case study in which the vehicle must maintain its position on an ill-maintained road.

Meta RL: Meta learning is one way to increase the data efficiency of learning algorithms by generalizing learned concepts from a set of training tasks to unseen, but related, tasks. Recent works on Meta RL model the dynamics of a new system as a mixture of expert systems, thus realizing knowledge transfer. In particular, [3,4] model the dynamics of the expert systems using Gaussian process with latent embedding. And the dynamics of the new system is constructed by a mixture of Gaussian processes with the distribution over the latent embedding being the mixture weights, which should be learned from data. In experiments, their meta-learning models effectively generalizes to novel tasks under various environments such as Cart-pole & Double-pendulum swing-up, HalfCheetah and Highway-Intersection.

World Models for model-based RL: Model-based Deep Reinforcement Learning (RL) assumes the availability of a model of an environment’s underlying transition dynamics, which are stochastic in nature and oftentimes multimodal. [5,6] use Mixture-density Networks to construct World Models to capture the multimodality in dynamics. Then model-based RL methods are implemented based on the constructed World Models to solve various tasks such as car racing and Inverse Sine Wave.

These applications naturally give rise to distributionally robust MDP problems, for example 1) robust control of drones or vehicles is needed to hedge against environmental uncertainties and distributional perturbations; 2) given established expert systems, usually the mixture weights are estimated from limited data collected from the unseen new task, and thus of high uncertainty. Our framework could potentially enhance robustness of decision making in unseen tasks by modeling the weighting parameter uncertainty; 3) by modeling the uncertainty in the World Model (specifically, in the weights), our framework enables robust decision making in model-based RL. The basis modes in our work can be the Gaussian components in GMMs, dynamics of expert systems, or mixture components in the Mixture-density Networks. As for the basis modes construction, we can adopt the estimation procedure in prior works. The difficulty should depend on specific tasks. It serves an interesting future research direction to also incorporate the uncertainty in basis modes into consideration for robust decision making.

Q5: Potential limitations of the choice to use mixture models.

A5: Our model assumes the dynamics is a mixture distribution with finite modes/component distributions. This restricts its application in scenarios where mixture distributions cannot approximate the dynamics well. Furthermore, the current manuscript assumes perfect knowledge of the basis modes and focuses solely on uncertainty in the weighting parameters. However, in practice, these basis modes often need to be estimated from data, introducing an additional layer of uncertainty that can significantly affect decision-making. Addressing this source of uncertainty is nontrivial and falls outside the scope of the current work; we leave it for future investigation.

References:

[1] Shi et. al., 2024, Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity

[2] Blanchet et. al., 2023, Double Pessimism is Provably Efficient for Distributionally Robust Offline Reinforcement Learning: Generic Algorithm and Robust Partial Coverage

[3] Shi et. al., 2023, The Curious Price of Distributional Robustness in Reinforcement Learning with a Generative Model

[4] Wang et. al., 2024, A Unified Principle of Pessimism for Offline Reinforcement Learning under Model Mismatch

[5] Panaganti et. al., 2022, Robust reinforcement learning using offline data

[6] Liu et. al., 2022, Distributionally Robust Q-Learning

[7] Balci et. al., 2024, Density steering of gaussian mixture models for discrete-time linear systems

[8] Engelaar et. al., 2024, Stochastic MPC for Finite Gaussian Mixture Disturbances with Guarantees

[9] Sæmundsson et. al., 2018, Meta reinforcement learning with latent variable gaussian processes

[10] Xu et. al., 2020, Task-agnostic online reinforcement learning with an infinite mixture of gaussian processes

[11] Sedlmeier et. al., 2021, Quantifying multimodality in world models

[12] Ha et. al., 2018, Recurrent world models facilitate policy evolution

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

Q1: Practical motivation for the linear mixture distributionally robust setting.

A1: Mixture models, in particular Gaussian Mixture Models (GMMs), are widely studied in fields of optimal control and RL. For example:

Model Predictive Control (MPC): [1] studies the density steering for discrete-time linear systems, a variant of optimal mass transport problem widely used in applications such as controlling a swarm of robots/drones to achieve a desired spatial distribution. They model the stochastic dynamical system using GMMs and derive its optimal control policy. [2] studies a stochastic MPC algorithm for linear systems subject to additive Gaussian mixture disturbances. They consider a vehicle control case study in which the vehicle must maintain its position on an ill-maintained road.

Meta RL: Meta learning is one way to increase the data efficiency of learning algorithms by generalizing learned concepts from a set of training tasks to unseen, but related, tasks. Recent works on Meta RL model the dynamics of a new system as a mixture of expert systems, thus realizing knowledge transfer. In particular, [3,4] model the dynamics of the expert systems using Gaussian process with latent embedding. And the dynamics of the new system is constructed by a mixture of Gaussian processes with the distribution over the latent embedding being the mixture weights, which should be learned from data. In experiments, their meta-learning models effectively generalizes to novel tasks under various environments such as Cart-pole & Double-pendulum swing-up, HalfCheetah and Highway-Intersection.

World Models for model-based RL: Model-based Deep Reinforcement Learning (RL) assumes the availability of a model of an environment’s underlying transition dynamics, which are stochastic in nature and oftentimes multimodal. [5,6] use Mixture-density Networks to construct World Models to capture the multimodality in dynamics. Then model-based RL methods are implemented based on the constructed World Models to solve various tasks such as car racing and Inverse Sine Wave.

The basis modes in our work can be the Gaussian components in GMMs, dynamics of expert systems, or mixture components in the Mixture-density Networks. These applications naturally give rise to distributionally robust MDP problems, for example 1) robust control of drones or vehicles is needed to hedge against environmental uncertainties and distributional perturbations; 2) given established expert systems, usually the mixture weights are estimated from limited data collected from the unseen new task, and thus of high uncertainty. Our framework could potentially enhance robustness of decision making in unseen tasks by modeling the weighting parameter uncertainty; 3) by modeling the uncertainty in the World Model (specifically, in the weights), our framework enables robust decision making in model-based RL.

Q2: In many real-world applications where model dynamics are uncertain and robustness is desired, it often seems more appropriate to model uncertainty in the features $\phi$, rather than in the mixture weights $\theta$.

A2: Thank you for raising this insightful point. Take the Meta-RL application in our response to Q1 for an example, the expert systems are constructed from huge amounts of historical data and thus of low uncertainty. Once constructed, they can be deemed as fixed. To estimate the dynamics of the new tasks, we have to estimate the mixture weights from data collected from the new tasks, which is usually of limit size and cause large uncertainty in the mixture weights estimation. This motivates us to model the uncertainty in the mixture weights $\theta$ to enable robust decision making in the new task, which turns out to be natural and of great importance.

Moreover, we agree that accounting for the uncertainty in the feature $\phi$ is also important for achieving robust decision-making. However, addressing it requires careful and non-trivial treatment in both algorithm design and theoretical analysis—particularly due to its interaction with uncertainty in the weighting parameters. A thorough investigation of this issue is beyond the scope of our current work, and we leave it as an important direction for future research.

Q3: Algorithm design: Decoupling of the learning and planning phase

A3: Thank you for this interesting question. We would like to provide more insights on our choice of algorithm design. We highlight that the decoupling of the learning and planning phase is the key to make theoretical guarantees possible. In particular, the analysis of the meta-algorithm only needs to control one-source of estimation error, which arises from the transition-targeted estimation of $\theta_h$. Moreover, the transition-targeted estimation simply uses $\phi$ as the feature vector. These two attributes are essential to allow Assumptions 3.4, 3.10 and 3.14 to neatly control the coverage of the offline dataset. And our analysis also leverages recently established concentration lemma for dependent error structure as we discussed in Remark 4.2.

For algorithms that embed learning with learning step-wisely, like the value-targeted regression algorithm widely studied in standard linear mixture MDP literature [7], they involve two layers of estimation error, i.e., the estimation error in $\theta$ and the estimation error in $\hat{V}$. During the preparation of this work, we found that the consideration of dynamics shift in DRMDP induces essential hardness to control these two sources of estimation error, unlike in standard MDPs. It is still unclear how to quantify/control the data coverage of the offline dataset with this type of algorithm. We would like to leave this problem for future study.

References:

[1] Balci et. al., 2024, Density steering of gaussian mixture models for discrete-time linear systems

[2] Engelaar et. al., 2024, Stochastic MPC for Finite Gaussian Mixture Disturbances with Guarantees

[3] Sæmundsson et. al., 2018, Meta reinforcement learning with latent variable gaussian processes

[4] Xu et. al., 2020, Task-agnostic online reinforcement learning with an infinite mixture of gaussian processes

[5] Sedlmeier et. al., 2021, Quantifying multimodality in world models

[6] Ha et. al., 2018, Recurrent world models facilitate policy evolution

[7] Jia et. al., 2020, Model-based reinforcement learning with value-targeted regression.

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

Q1: The theoretical section is dense and could benefit from clearer intuition or visual aids.

A1: Thank you for the suggestion. The current layout of section 4 is due to the space limit. Given additional space in the camera ready, we will separate section 4 into two sections, Algorithm Design and Theoretical Guarantees, and expand our current discussion and interpretation on our analysis and assumptions.

Q2: How sensitive is the framework to mis-specified basis modes in practice?

A2: From a theoretical point of view, mis-specification in basis modes can cause error and result in suboptimal performance, then there would be no guarantee for the algorithm to learn the optimal robust policy. In practice, the basis modes are typically constructed from data, such as the Gaussian components in Gaussian Mixture Model [1,2], or dynamics of expert systems in Meta-RL [3,4], thus approximation error is inevitable. Nevertheless, these prior works have shown certain empirical success in using mixture models to approximate the complex true dynamics and performances actually vary from tasks to tasks. Another empirical evidence is from the recent work [5] on d-rectangular linear DRMDPs. They use linear MDP to model nonlinear dynamics. Even with large mis-specification errors in the feature mapping, their algorithm still achieves certain robustness performance. Based on these evidences, we conjecture that sensitivity of our framework with respect to the misspecification in basis modes should depend on specific tasks. But even with certain approximation errors, we can still expect certain robustness performances in practice. It serves an interesting future research direction to include the uncertainty in basis modes into consideration.

Q3: Extend the framework to linear MDPs and low-rank MDPs.

A3: We would like to note that distributionally robust MDPs have been separately studied in prior works under the framework of linear MDPs [5] and low rank MDPs [6]. They study different types of uncertainty sets constructed based on the posed structures of transition models, propose corresponding algorithms and provide theoretical guarantees. In our work, we particularly focus on a mixture structure of the dynamics and modeling uncertainty in the mixture weights. Thus, our work is orthogonal to previous ones on linear MDPs and low-rank MDPs.

Q4: Would it be possible to extend the framework to online settings with exploration guarantees?

A4: Thanks for the great question. Yes, our framework can potentially be extended to online settings where the agent has to actively interact with the nominal environment to collect data. Challenges lie in how to design algorithms capable of reasoning about the exploration and address the exploration-exploitation trade off in a data efficient manner. A potential solution would be leveraging the principle of optimism in the face of uncertainty [7]. For example, at each episode $k$, we can maintain a confidence set $\hat{\mathcal{P}}^k$ based on the weighting parameter $\theta$ estimation and update it when new trajectories are collected. Then we define the value estimator $J^k(\pi)$ as the largest robust value function with nominal kernel included in the confidence set: $J^k(\pi) = \sup_{P\in \hat{\mathcal{P}}^k}\inf_{\tilde{P}\in \mathcal{U}_h^{\rho}(P)}V_{\tilde{P}}^{\pi}(s_1)$, instead of the double pessimism estimator in our work. Then we use the optimal policy wrt $J^k(\pi)$ as the execution policy at the current episode to collect data. Another potential challenge would be quantifying the difficulty of exploration in online DRMDPs. There is a line of work studying this problem and several interesting phenomenons have been identified that are unique to the online setting. We leave a rigorous theoretical development for future research.

References:

[1] Balci et. al., 2024, Density steering of gaussian mixture models for discrete-time linear systems

[2] Engelaar et. al., 2024, Stochastic MPC for Finite Gaussian Mixture Disturbances with Guarantees

[3] Sæmundsson et. al., 2018, Meta reinforcement learning with latent variable gaussian processes

[4] Xu et. al., 2020, Task-agnostic online reinforcement learning with an infinite mixture of gaussian processes

[5] Liu et. al., 2024, Distributionally robust off-dynamics reinforcement learning: Provable efficiency with linear function approximation

[6] Hu et. al., 2024, Efficient Duple Perturbation Robustness in Low-rank MDPs

[7] Azar et. al., 2017, Minimax Regret Bounds for Reinforcement Learning

[8] He et. al., 2025, Sample Complexity of Distributionally Robust Off-Dynamics Reinforcement Learning with Online Interaction

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?