{'fraction_non_alphanumeric': 0.040064849151184245, 'fraction_numerical': 0.01026778227083838, 'mean_word_length': 4.599332220367279, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Group Distributionally Robust Reinforcement Learning with Hierarchical Latent Variables

One key challenge for multi-task Reinforcement learning (RL) in practice is the absence of task indicators. Robust RL has been applied to deal with task ambiguity, but may result in over-conservative policies. To balance the worst-case (robustness) and average performance, we propose Group Distributionally Robust Markov Decision Process (GDR-MDP), a flexible hierarchical MDP formulation that encodes task groups via a latent mixture model. GDR-MDP identifies the optimal policy that maximizes the expected return under the worst-possible qualified belief over task groups within an ambiguity set. We rigorously show that GDR-MDP's hierarchical structure improves distributional robustness by adding regularization to the worst possible outcomes. We then develop deep RL algorithms for GDR-MDP for both value-based and policy-based RL methods. Extensive experiments on Box2D control tasks, MuJoCo benchmarks, and Google football platforms show that our algorithms outperform classic robust training algorithms across diverse environments in terms of robustness under belief uncertainties. Demos are available on our project page (\url{https://sites.google.com/view/gdr-rl/home}).


Introduction
Reinforcement learning (RL) has demonstrated extraordinary capabilities in sequential decision-making, even for handling multiple tasks [1, 2, 3, 4]. With policies conditioned on accurate task-specific contexts, RL  perform better than ones without access to context information [5,6]. However, one key challenge for contextual decision-making is that, in real deployments, RL agents may only have incomplete information about the task to solve. In principle, agents could adaptively infer the latent context with data collected across an episode, and prior knowledge about tasks [7,8,9]. However, the context estimates may be inaccurate [10,11] due to limited interactions, poorly constructed inference models, or intentionally injected adversarial perturbations. Blindly trusting the inferred context and performing context-dependent decisionmaking may lead to significant performance drops or catastrophic failures in safety-critical situations. Therefore, in this work, we are motivated to study the problem of robust decision-making under the task estimate uncertainty. Prior works about robust RL involve optimizing over the worst-case qualified elements within one uncertainty set [12,13]. Such robust criterion assuming the worst possible outcome may lead to overly conservative policies, or even training instabilities [14,15,16]. For instance, an autonomous agent trained with robust methods may always assume the human driver is aggressive regardless of recent interactions and wait until the road is clear, consequently blocking the traffic. Therefore, balancing the robustness against task estimate uncertainties and the performance when conditioned on the task estimates is still an open problem. We provide one solution to address the above problem by modeling the commonly existing similarities between tasks under distributionally robust Markov Decision Process (MDP) formulations. Each task is typically represented by a unique combination of parameters or a multi-dimensional context in multitask RL. We argue that some parameters are more important than others in terms of affecting the environment dynamics model and thus tasks can be properly clustered into mixtures according to the more crucial parameters as in Figure 1 (a) and (b). However, existing robust MDP formulations [12] lack the capacity to model task groups, or equivalently, task subpopulations. Thus the effect of task subpopulations on the policy's robustness is unexplored. In this paper, we show that the task subpopulations help balance the worst-case performance (robustness) and average Figure 1: Illustration examples when modeling tasks with a flat latent structure that uses one distribution for all tasks as in (a), and a hierarchical latent structure that clusters tasks to different mixtures as in (b). The graphical model with a hierarchical latent structure for both GDR-MDP and HLMDP is shown in (c). At episode n, a mixture z n is first sampled from a prior distribution w. An MDP m is then sampled according to µ zn (m) and controls the dynamics of the n'th episode.
performance under conditions (Section 5.2). In contrast to prior work [10] that leverages point estimates of latent contexts, we take a probabilistic point of view and represent the task subpopulation estimate with a belief distribution. Holding a belief of the task subpopulation, which is the high-level latent variable, helps leverage the prior distributional information of task similarities. It also naturally copes with distributionally robust optimization by optimizing w.r.t. the worst-possible belief distribution within an ambiguity set. We consider an adaptive setting in line with system identification methods [17], where the belief is initialized as a uniform distribution and then updated during one episode. Our problem formation is related to the ambiguity modeling [18] inspired by human's bounded rationality to approximate and handle distributions, which has been studied in behavioral economics [19,20] yet has not been widely acknowledged in RL.
We highlight our main contributions as follows: 1. We formulate Hierarchical-Latent MDP (HLMDP) (Section 4), which utilizes a mixture model over MDPs to encode task subpopulations. HLMDP has a high-level latent variable z as the mixture, and a lowlevel m to represent tasks (Figure 1 (c)).
2. We introduce the Group Distributionally Robust MDP (GDR-MDP) in Section 5 to handle the overconservative problem, which formulates the robustness w.r.t. the ambiguity of the adaptive belief b(z) over mixtures. GDR-MDP builds on distributionally robust optimization [21,22] and HLMDP to leverage rich distributional information.
3. We show the convergence property of GDR-MDP in the infinite-horizon case. We find that the hierarchical latent structure helps restrict the worst-possible outcome within the ambiguity set and thus helps generate less conservative policies with higher optimal values. 4. We design robust deep RL training algorithms based on GDR-MDP by injecting perturbations to beliefs stored in the data buffer. We empirically evaluate in three environments, including robotic control tasks and google research football tasks. Our results demonstrate that our proposed algorithms outperform baselines in terms of robustness to belief noise.

Related Work
Robust RL and Distributionally Robust RL. RL's vulnerability to uncertainties has attracted large efforts to design proper robust MDP formulations accounting for uncertainties in MDP components [12,13,23,24,25,26].
Existing robust deep RL algorithms [27,28,29,30,31,24] are shown to generate robust policies with promising results in practice. However, it is also known that robust RL that optimizes over the worst-possible elements in the uncertainty set may generate over-conservative policies by trading average performance for robustness and may even lead to training instabilities [16]. In contrast, distributionally robust RL [32,33,34,35,36,37,38,39] assumes that the distribution of uncertain components (such as transition models) is partially/indirectly observable. It builds on distributionally robust optimization [21,22] which optimizes over the worst possible distribution within the ambiguity set. Compared with common robust methods, distributionally robust RL embeds prior probabilistic information and generates less conservative policies with carefully calibrated ambiguity sets [32]. We aim to propose distributionally robust RL formulations and training algorithms to handle task estimate uncertainties while maintaining a trade-off between robustness and performance.
One relevant work is the recently proposed distributionally robust POMDP [37] which maintains a belief over states and finds the worst possible transition model distribution within an ambiguity set. We instead hold a belief over task mixtures and find the worst possible belief distribution.
[38] also maintains a belief distribution over tasks but models tasks with a flat latent structure. , which optimizes a conditional value-at-risk objective and maintains an uncertainty set centered on a context point estimate. Instead, we maintain an ambiguity set over beliefs and further consider the presence of task subpopulations.
[11] also considers the uncertainties in belief estimates but with a flat latent task structure.

Preliminary
This section introduces Latent MDP and the adaptive belief setting, both serving as building blocks for our proposed HLMDP (Section 4) and GDR-MDP (Section 5).
is a tuple (T, S, A, P m , R m , ν m ) where P m is the transition probability, R m is the reward function and ν m is the initial state distribution.
Latent MDP assumes that at the beginning of each episode, one MDP from set M is sampled based on µ(m). It aims to find a policy π that maximizes the accumulated expected return solving max π The Adaptive Belief Setting In general, a belief distribution contains the probability of each possible MDP that the current environment belongs to. The adaptive belief setting [5] holds a belief distribution that is dynamically updated with streamingly observed interactions and prior knowledge about the MDPs. In practice, prior knowledge may be acquired by rule-based policies or data-driven learning methods. For example, it is possible to pre-train in simulated complete information scenarios or exploit unsupervised learning methods based on online collected data [51]. There also exist multiple choices for updating the belief, such as applying the Bayesian rule as in POMDPs [49] and representing beliefs with deep recurrent neural nets [52].

Hierarchical Latent MDP
In realistic settings, tasks share similarities, and task subpopulations are common. Although different MDP formulations are proposed to solve multi-task RL, the task relationships are in general overlooked. To fill in the gap, we first propose Hierarchical Latent MDP (HLMDP), which utilizes a hierarchical mixture model to represent distributions over MDPs. Moreover, we consider the adaptive belief setting to leverage prior information about tasks. Definition 1 (Hierarchical Latent MDPs). An episodic HLMDP is defined by a tuple (Z, M, T, S, A, w). Z denotes a set of Latent MDPs and |Z| = Z. M is a set of MDPs with cardinality |M| = M shared by different Latent MDPs. T , S, and A are the shared episode length (planning horizon), state, and action space, respectively. Each Latent MDP Z z ∈ Z, z ∈ [Z] consists of a set of joint MDPs {M m } M m=1 and their weights µ z satisfying M m=1 µ z (m) = 1. w is the categorical distribution over Latent MDPs and Z z=1 w(z) = 1.
We provide a graphical model of HLMDP in Figure 1 (c). HLMDP assumes that at the beginning of each episode, the environment first samples a Latent MDP z ∼ w(z) and then samples an MDP m ∼ µ z (m). HLMDP encodes task similarity information via the mixture model, and thus contains richer task information than Latent MDP proposed in [45]. For instance, we could always find one Latent MDP for each HLMDP. However, there may exist infinitely many corresponding HLMDPs given one Latent MDP.
HLMDP in Adaptive Belief Setting. When solving multi-task RL problems, the adaptive setting is shown to help generate a policy with a higher performance [5] than the non-adaptive one since it leverages prior knowledge about the transition model as well as the online collected data tailored to the unseen environment. Hence we are motivated to formulate HLMDP in the adaptive belief setting.
HLMDP maintains a belief distribution b(z) over task groups to model the probability that the current environment belongs to each group z. At the beginning of each episode, we initialize the belief distribution with a uniform distribution b 0 . We use the Bayesian rule to update beliefs based on interactions and a prior knowledge base. Note that the knowledge base are not accurate enough and may lead to inaccurate belief updates. At timestep t, we get the next belief estimate b t+1 with the state estimation function SE: wher Under the adaptive belief setting, HLMDP aims to find an optimal policyπ within a history-dependent policy class Π, under which the discounted expected cumulative reward is maximized as in Equation 2. Following general notations in POMDPs, we denote the history at time t as h t = (s 0 , a 1 , s 1 , . . . , s t−1 , a t−1 , s t ) ∈ H t containing state-action pairs (s, a). At timestep t, we use both the observed state s t and the inferred belief distribution b t (z) as the sufficient statistics for history h t .
where r t denotes the reward received at step t. b 0 (z) is the initial belief at timestep 0.

Group Distributionally Robust MDP
The belief update function in Equation 1 may not be accurate, which motivates robust decision-making under belief estimate errors. In this section, we introduce Group Distributionally Robust MDP (GDR-MDP) which models task groups and considers robustness against the belief ambiguity. We then study the convergence property of GDR-MDP in the infinite-horizon case in Section 5.1. We find that GDR-MDP's hierarchical structure helps restrict the worst-possible value within the ambiguity set and provide the robustness guarantee in Section 5.2.
Definition 2 (General Ambiguity Sets). Let ∆ k be a ksimplex. Considering a categorical belief distribution b ∈ ∆ k , a general ambiguity set without special structures is defined as C ∆ k containing all possible distributions for b.
Definition 3 (Group Distributionally Robust MDP). An episodic GDR-MDP is defined by a 8-tuple (C, Z, M, T, S, A, w, SE). C is a general belief ambiguity set. T, S, A, M, Z, w are elements of an episodic HLMDP as in Definition 1. SE : ∆ Z−1 × S → ∆ Z−1 is the belief updating rule. GDR-MDP aims to find a policy π ∈ Π that obtains the following optimal value: where C ∆ Z−1 is a general ambiguity set tailored to beliefs over Latent MDPs in set Z.
GDR-MDP naturally balances robustness and performance by leveraging distributionally robust formulation and rich distributional information. In contrast to HLMDP, which maximizes expected return over nominal adaptive belief distribution (Equation 2), GDR-MDP aims to maximize the expected return under the worst-possible beliefs within an ambiguity set C ∆ Z−1 . Moreover, GDR-MDP optimizes over fewer optimization variables than when directly perturbing MDP model parameters or states. It resembles the group distributionally robust optimization problem in supervised learning [53, 54] but focuses on sequential decision-making in dynamic environments.

Convergence in Infinite-horizon Case
With general ambiguity sets (as in Definition 2), calculating the optimal policy is intractable [33,39]. We propose a belief-wise ambiguity set that follows the b-rectangularity to facilitate solving the proposed GDR-MDP.
Assumption 1 (b-rectangularity). We assume a belief-wise ambiguity set,C := b∈∆ Z−1 C b , where represents Cartesian product. b serves as the nominal distribution of the ambiguity set.
More concretely, the b-rectangularity assumption uncouples the ambiguity set related to different beliefs. When conditioned on beliefs at each timestep, the minimization loop selects the worst-case realization unrelated to other timesteps. The b-rectangularity assumption is motivated by the s-rectangularity first introduced in [23], which helps reduce a robust MDP formulation to an MDP formulation and get rid of the time-inconsistency problem [55]. Ambiguity sets beyond rectangularities are recently explored in [56, 57], which we leave for future works.
With b-rectangular ambiguity sets, we derive Bellman equations to solve Equation 3 with dynamic programming. Detailed proofs are in Appendix Section B.1.

Proposition 1 (Group Distributionally Robust Bellman Equation)
. Define the distributionally robust value of an arbitrary policy π as follows where b t+1 = SE(b t , s t ).
The Group Distirbutionally Robust Bellman expectation equation is Lemma 1 (Contraction Mapping). Let V be a set of realvalued bounded functions on ∆ Z−1 ×S. LV (b, s) : V → V refers to the Bellman operator defined as Theorem 1 (Convergence in Infinite-horizon Case). Define V ∞ (b, s) as the infinite horizon value function. For all b ∈ B and s ∈ S, we have By repeatedly applying the contraction operator in Lemma 1, the value function will converge to a unique fixed point, which corresponds to the optimal value based on Banach fixed point theorem [58].

Robustness Guarantee of GDR-MDP
This section shows how GDR-MDP's hierarchical task structure and the distributionally robust formulation help balance performance and robustness. We compare the optimal value of GDR-MDP denoted as V GDR (π GDR ), with three different robust formulations. Group Robust MDP is a robust version of GDR-MDP with its optimal value denoted as V GR (π GR ). Distributionally Robust MDP holds a belief over MDPs without the hierarchical task structure whose optimal value denoted as V DR (π DR ). Robust MDP is a robust version of Distributionally Robust MDP, denoted as V R (π R ). π · denote optimal policies under different formulations. We achieve the comparison by studying how maintaining beliefs over mixtures affects the worstpossible outcome of the inner minimization problem and the resulting RL policy. We study the worst-possible value via the relationships between ambiguity sets projected to the space of beliefs over MDPs. We first define a discrepancy-based ambiguity set that is widely used in existing DRO formulations [59, 60, 61]. Definition 4 (Ambiguity set with total variance distance). Consider a discrepancy-based ambiguity set defined based on total variance distance. Formally, the ambiguity set is where X ∈ X is the support, ν X is the nominal distribution over X and ξ is the ambiguity set's size.
To achieve a reasonable comparison, we control the adversary's budget ξ the same when perturbing the belief over task groups z and tasks m, which correspond to different model misspecification forms when there is a hierarchical latent structure about tasks. Theorem 2 (Values of different robust formulations). Let U m (π) = E π m T t=1 γ t r t . Let C b(m),d T V ,ξ (m) and C b(z),d T V ,ξ (z) denote the ambiguity sets for beliefs over tasks m and groups z, respectively. b(m) and b(z) satisfy b(m) = Z µ z (m)b(z) and are the nominal distributions. For any history-dependent policy π ∈ Π, its value function under different robust formulations are: We have the following inequalities hold: Theorem 2 shows that with a nontrivial ambiguity set, the distributionally robust formulation in GDR-MDP helps regularize the worst-possible value when compared with robust ones, including the group robust (GR) and task robust (R) formulations. It also shows that GDR-MDP's hierarchical structure further helps restrict the effect of the (c) show the graphical model, the relationship between ambiguity sets, and different robust formulations' optimal values for an example with two groups and two unique tasks.
(d) shows the relationship between ambiguity sets for an example with two groups and three unique tasks.
adversary, resulting in higher values than the distributionally robust formulation with a flat latent structure (DR). To get Theorem 2, we first find that when projecting the ξambiguity set for b(z) to the space of b(m), the resulting ambiguity set is a subset of the ξ-ambiguity set for b(m).
Proofs are detailed in Appendix Section B.2. Our setting is different from [62] which states that DRO is a generalization of point-wise attacks. The key difference is that when the adversary perturbs b(m), we omit the expectation over the mixtures under b(z). Theorem 3 (Optimal values of different robust formulations). Let π · denote the converged optimal policy for different robust formulations, we have Based on Theorem 2, we can compare the optimal values for different robust formulations. Theorem 3 shows that imposing ambiguity set on beliefs over mixtures helps generate less conservative policies with higher optimal values at convergence compared with other robust formulations. Illustration Examples in Figure 2. We provide two hierarchical latent bandit examples in Figure 2. The first example shown in Figure 2 (a) has two latent groups with different weights over two unique MDPs. (b) shows the ambiguity sets of the example in (a). The orange sets denote the ξ-ambiguity sets for the beliefs over mixtures and MDPs. The green set denotes the ambiguity set projected from the ξ-ambiguity set for belief distributions over mixtures. We show that the mapped set is a subset of the original ξ-ambiguity set for the MDP belief distributions. (c) shows the optimal policy and value of different robust formulations for the example in (a). Our proposed GDR has the potential to get a less conservative policy with higher returns than other robust baselines. (d) follows the same notations in (b) but corresponds to an example with three possible MDPs. (b) and (d) together shows that the hierarchical structure helps regularize the adversary's strength.
The detailed procedure for getting the optimal policies is shown in Appendix A.

Algorithms
To solve the proposed GDR-MDP, we propose novel robust deep RL algorithms (summarized in Algorithm 2 and Algo-Algorithm 1: GDR-MDP Trajectory Rollout Input: Mixing weights w(z) and µ z (m), episode index n, episode length T , belief update function SE, rollout policy π θ (b(z), s), exploration Sample action a t with the -greedy method and GDR-DQN and GDR-SAC. We update the Q-net in GDR-DQN and the critic net in GDR-SAC toward TD targets with perturbed beliefs. We follow Definition 4 to construct the ambiguity set C b (z),d T V ,ξ which centers at the originally inferred b (z) and satisfies the brectangularity assumption stated in Assumption 1. At each training step, we sample a batch data {d = (b(z), s, a, r, b (z), s , a , r )} N from the replay buffer D to estimate the perturbed TD target.
We update Q-functions with gradient descents. For both GDR-DQN and GDR-SAC, we have loss as GDR-PPO. GDR-PPO conducts robust training by decreasing the advantages of trajectories that are vulnerable to belief noises. More concretely, given a trajectory d, its advantage for (b t , s t ) is calculated as follows.
We measure the performance drop R drop under worstpossible beliefs within the ambiguity set.
Algorithm 2: Group Distributionally Robust Training for GDR-DQN and GDR-SAC Input: Q-net Q θ (b(z), s, a), ambiguity set C ·,d T V ,ξ , training episodes N , Initialize data buffer D ; for n = 0 to N do Rollout one episode with Algorithm 1 and append data pairs to D ; if Update Q-net parameters then Sample batch data from D ; Worst-possible Beliefs. To obtain the worst case distribution b adv ∈ C b (z),d T V ,ξ , we iteratively apply a stochastic variant of fast gradient sign method (FGSM) [67] to make sure that the perturbed discrete distribution satisfies z p(z) = 1. For each attack to the belief distribution, we randomly sample an index i ∈ Z, and apply the attack to each element in p(z) as follows To stabilize robust training, we pretrain for a small amount of episodes with exact one-hot beliefs to ensure that the value function could approximate the actual state value to some extent. To achieve a certain level of robustness over noisy inferred belief b(z), we fix the ambiguity set size along with robust training, which is analogous to the adversary budget and the robustness level [36].

Experiments
We conduct experiments to empirically study (a) the effect of GDR-MDP's hierarchical structure on the robust training stability and (b) policy's robustness to belief estimate error.

Environments
We evaluate GDR-DQN in Lunarlander [68], GDR-SAC in Halfcheetah [69], and GDR-PPO in Google Research Football [70]. Table 1 shows a summary of environment setups. More details are in Appendix Section C. To initialize each episode, we first sample a group z ∼ w(z), and then a task m ∼ µ z (m) for the episode. Note that both z and m are unknown to the agent. Google Research Football (GRF). This domain presents additional challenges due to its AI randomness, large state-action spaces, and sparse rewards. The RL agent will control one active player on the attacking team at each step and can pass to switch control. The non-active players will be controlled by built-in AI. The dynamics of our designed 3 vs. 2 tasks are determined by the player types including central midfield (CM) and centre back Figure 3: The training performance of GDR and baselines. Each curve is averaged over 5 runs and shaded areas represent standard errors. Our results show that GDR has better training stability than DR by implicitly regularizing the adversary's strength with the hierarchical structure. Table 1: Environment setups. Both parameters affect the environment dynamics. In GRF, the strongest player has a capability level of 1.0. Our tasks are more challenging than the original 3 vs. 2 task in GRF (1.0 vs. 0.6) in terms of the capability level. For notation simplicity, let 1 k be a k-dimensional vectored filled with 1, and E(k) = 1 k 1 k . E i (k) shows the i-th matrix block on the diagonal.
  (CB), and player capability levels. The built-in CM player tends to go into the penalty area when attacking and guard the player on the wing (physically left or right) when defending, while the CB player tends to guard the player in the middle when defending, and not directly go into the penalty area when attacking. Different patterns of policies are required to solve the tasks from different groups.
Box2D Control Task: LunarLander. The Lunarlander's dynamics are controlled by the engine mode and engine power. In the flipped mode, the action turning on the left (or right) engine in normal mode will turn on the right (or left) engine instead.
Mujoco Control Task: HalfCheetah. In HalfCheetah, each task's dynamics are controlled by both the torso mass and the failure joint, to which we cannot apply action. Our setting is similar to the implementation in [10] but with a fixed failure joint within each episode.

Baselines
We compare our Group Distributionlly Robust training methods (GDR) with five baselines. In G-Exact, the RL agent is trained with the exact mixture information encoded in a one-hot vector. The agent in DR maintains a belief dis-tribution b(m) and utilizes distributionally robust training over b(m). It uses the same belief updating rule as in GDR to update b(z) at each timestep but projects b(z) to b(m) with µ z (m). DR utilizes no mixture information and helps ablate the effect of the hierarchical latent structure. The agent in No-Belief has no access to the context information and generates action only based on state s. The No-Belief baseline helps show the importance of the adaptive belief setting. In G-Belief, the agent maintains belief b(z) and is trained towards a nominal TD target. Compared with GDR, G-Belief helps reveal the effect of distributionally robust training. The State-R agent takes both the inferred belief b(z) and state s as input. It updates towards a TD target with perturbed states along with training. For baselines with belief modules, we utilize the Bayesian update rule in Equation 1 and leave the detailed likelihood calculation in Appendix Section D. We study the effect of the hierarchical structure on the adversary's strength based on training performances in Figure 3. We show the importance of mixture information since the No-Belief baseline consistently underperforms G-Exact during training in all three environments. Lunarlander and HalfCheetah have a return much lower than G-Exact since the kinematic observation fed into the neural net does not reveal any mixture information. In GRF, the No-Belief baseline underperforms G-Exact since it could not effectively learn distinct strategies with regard to different types of players as teammates and opponents, while G-Exact could learn group-specific policies. When compared with other robust training baselines including DR and State-R, GDR achieves a higher average return at convergence in all environments as in Figure 3. In LunarLander and HalfCheetah, DR which maintains a belief b(m) over MDPs induces significant training instability, instead of learning a meaningful conservative policy. In GRF, DR has a worse asymptotic performance than GDR. Those observations empirically validate our theoretical result (Section 5.2) in the regime of deep RL, which is that, Figure 4: Robustness evaluations when facing belief inference errors. Each plot is averaged over 5 runs and shaded areas represent standard errors. GDR preserves higher robustness to belief inference errors compared with baselines.
with the same ambiguity set size, perturbing b(m) omitting mixture information will lead to larger value perturbations than perturbing b(z) over mixtures. The State-R baseline leads to more considerable training instability than DR and fails to learn in all three environments. We compare GDR with non-robust training baselines, including G-Exact and G-belief to study the importance of robust training. In LunarLander, GDR has comparable training performance with G-Exact and G-Belief. In GRF, GDR has slightly worse asymptotic performance than G-Exact and better performance than G-Belief. These observations show that GDR successfully extracts task-specific information stored in the noisy beliefs and conditions on the beliefs for action generation. In HalfCheetah, GDR performs better than G-Exact. Although GDR leads to an immediate performance drop after pretraining (100000 steps), the robust training in GDR converges to higher performance. We conjecture that this is due to the perturbed belief helping the algorithm get out of local optima.

Robustness to Belief Inference Errors
We test the robustness against belief noise of the best policies obtained with GDR and baselines along with training. The results are shown in Figure 4. We define the belief noise level as the inaccuracy of the likelihood when updating belief with the Bayesian rule. During robustness evaluation, G-Exact generate actions conditioned on the same noisy beliefs as GDR and G-Belief. In GRF and HalfCheetah, GDR is consistently more robust to belief noise than robust and nominal training baselines. In LunarLander, the mean reward of GDR is better than G-Exact when there is a high belief noise level and is better than DR when a low belief noise level. The large variances in LunarLander are due to the large penalty when crashes which are further exaggerated by the fixed episode length. Although GDR has its performance decreasing along with the increase of the belief noise level, its performance is still an upper bound of DR and G-Exact's performances. These observations show that GDR successfully balances the information between belief distributions and states, and is more robust to belief inference errors. G-Exact is prone to injected belief noise since it heavily relies on accurate mixture information to achieve high performance. G-Belief does not show significant robustness improvement over G-Exact. It shows that the group distributionally robust training procedure instead of the belief randomness along training helps improve the robustness.

Ablation Study
We perform empirical sensitivity analysis to reveal the effect of uncertainty set size on GDR's policy robustness in HalfCheetah. Figure 5 (a) shows that gradually increasing the ambiguity set size up to 0.2 helps improve the robustness. The ambiguity set whose size is greater or equal to 0.25, easily leads to training instability and thus decreases the robustness. In contrast, even with an ambiguity set of size 0.05 and pretraining for 300000 steps, DR without the mixture information still causes unstable training (see Appendix Section E). Figure 5 (b) provides the average belief errors at each time step corresponding to different belief noise levels. Figure 5 (b) and Figure 4 show that GDR only shows significant performance drops when the belief error is nonzero for a large portion of steps.

Conclusion
This paper considers robustness against task estimate uncertainties. We propose the GDR-MDP formulation that can leverage rich distribution information, including adaptive beliefs and prior knowledge about task groups. To the best of our knowledge, GDR-MDP is the first distributionally robust MDP formulation that models ambiguity over belief estimates in an adaptive setting. We theoretically show that GDR-MDP's hierarchical latent structure helps enhance its distributional robustness compared with a flat task structure. We also empirically show that our proposed group distributionally robust training methods generate more robust policies than baselines when facing belief inference errors in realistic scenarios. We hope this work will inspire future research on how diverse domain knowl-edge affects robustness and generalization. One exciting future direction is to scale the group distributionally robust training to high-dimensional and continuous latent task distributions for diverse decision-making applications.

A Toy Example: Hierarchical Latent Bandit
In this section, we show the process of getting the optimal policies for different robust formulations in the Hierarchical Latent Bandit problem as illustrated in Figure 2 (a).
The agent has two possible actions, a 0 and a 1 . There are two possible mixtures/groups denoted as z, and two possible MDPs denoted as m. Given the mixture, we have the conditional probability for each MDP as µ(m = 0|z = 0) = 0.8, µ(m = 1|z = 0) = 0.2, µ(m = 0|z = 1) = 0.0, µ(m = 1|z = 1) = 1.0. We assume the same type of ambiguity set measured by the total variance distance as in the analysis. Let the current belief over groups be b(z) = [0.5, 0.5] and the ambiguity set size be ξ = 0.2.
We compare the optimal policies of four robust formulations, including • our proposed GDR-MDP (shorthanded as GDR) that utilizes both the hierarchical structure and distributionally robust formulation, and optimizes over the worst-possible beliefs over groups, • group robust MDP (GR), which optimizes over the worst-possible groups, • distributionally robust MDP (DR), which holds a belief over MDPs without the hierarchical task structure and optimizes over the worst-possible belief distribution, • robust MDP (R), which is a robust version of distributionally robust MDP and optimizes over the worst-possible MDP.
Optimal policy for R. R desires robustness over the worst possible MDPs. We can see that the worst possible MDP is m 1 since the reward when choosing a 0 or a 1 in m 1 is consistently smaller than the rewards when in m 0 . Since the optimal policy for m 1 is selecting a 1 , the optimal policy for R is a 1 .
Similar results can be derived with the value function. Formally, given ∈ [−ξ, ξ] = [−0.2, 0.2], π(a 0 ) = p, π(a 1 ) = 1−p, we want to solve the following optimization problem Therefore when p = 0, the value is maximized. It shows that the optimal policy is a 1 .
Optimal policy for GDR. GDR instead desires robustness over the worst possible belief distribution over contexts. Similar to the analysis for DR, the value estimate of a 1 ,V (a 1 ), is always equal to 5 regardless of the perturbedb(z). Now we need to investigate the value when selecting deterministic policy as a 0 . The weight on z 0 in the perturbed belief lies in rangeb(z 0 ) ∈ [0.5 − ξ, 0.5 + ξ] = [0.3, 0.7]. The value estimate for a 0 is thusV (a 0 ) =b(z 0 ) * 0.8 * 22 = 17.6b(z 0 ) ∈ [5.28, 12.32]. Since the lower bound is larger than the value ofV (a 1 ) = 5, the optimal policy for GDR is a 0 .
Similarly, we can also write out the value function and the optimization problem.
Therefore when p = 1, the value is maximized. It shows that the optimal policy is a 0 .
To sum up, the Hierarchical Latent Bandit example shows that our proposed GDR-MDP has the potential to find a less conservative policy compared with other robust formulations.

B.1.1 Proofs for Proposition 1
We provide the proof for the Bellman expectation equation as follows. Starting from the definition of V π t (b t , s t ), we first separate the elements at time step t from future timesteps. We then find that the elements related to future timesteps starting from step t + 1 could be aggregated to the group distributionally robust value at step t + 1.
Therefore, the Group Distributionally Robust Bellman expectation equation is Proposition 2. The Group Distributionally Robust Bellman optimality equation is Following a similar process, we could also prove the Bellman optimality equation as follows.
Therefore, the Group Distributionally Robust Bellman optimality equation is s t+1 ) .

B.1.2 Proof for Lemma 1
Let V refer to a set of real-valued bounded functions on ∆ Z−1 × S and LV (b, s) : V → V refer to the Bellman operator defined as Now we start the proof to show that the Bellman operator above is a contraction operator. For notation simplicity, let Given arbitrary U, V ∈ B and based on the definition of the operator LV above, LU, LV are real-valued and bounded.
Let (b U , π U ) and (b V , π V ) be the saddle points for LU and LV , respectively.

B.1.3 Proof for Theorem 1
Since LV (b, s) is a contraction operator based on Lemma 1, we directly follow the Banach's Fixed-Point Theorem [71] to show that (a) there exist a unique solution for LV ∞ (b, s) = V ∞ (b, s), and (b) the value function initiating from any value converge uniformly by iterative applying the Bellman update built in finite horizon case.

B.2 Proofs for Section 5.2: Robustness Guarantee for GDR-MDP
In this section, we prove the robustness guarantee of our proposed GDR-MDP. We compare the GDR-MDP's optimal value with three different robust formulations. We achieve the comparison by studying how maintaining beliefs over mixtures affects the worst-possible outcome of the inner minimization problem and the resulting RL policy. We study the worstpossible value via the relationships between ambiguity sets projected to the space of beliefs over MDPs.

B.2.1 Ambiguity Set Projection and Set Relationships
Recall that we consider a discrepancy-based ambiguity set defined based on total variance distance in Definition 4. Formally, the ambiguity set is where x ∈ X is the support, ν X is the nominal distribution over X , and ξ is the ambiguity set's size. We now define the operator that maps an ambiguity set over distribution for mixtures to an ambiguity set over distributions for MDPs. Definition 5 (Ambiguity Set Projection). The operator T A projects an ambiguity set for distributions over Z to an ambiguity set for distributions over M, and C b(m),d,ξ (m) is the ambiguity set for admissible distributions over supports M, where b(m) is the nominal distribution. d is the distance metric. ξ is the set size and also the adversary's perturbation budget around the nominal distribution. Similarly, C b(z),d,ξ (z) is the ambiguity set for admissible distributions over supports Z.
With the set projection operator T A , we can derive the relationships between the projected ambiguity set T A (C b(z),d,ξ (z)) and the ξ-ambiguity set C b(m),d,ξ (m) which directly represents the model misspecifications over different MDPs. We state the results in Proposision 3.
Proposition 3 (Ambiguity Set Regularization with the Hierarchical Latent Structure). Consider two adversaries with the same attack budget ξ. One adversary perturbs the z-level distribution by selecting the worst possible distribution within C b(z),d,ξ (z) and the other perturbs the m-level distribution by selecting the worst possible distribution within C b(m),d,ξ (m). Given the nominal distribution for Z as b(z), we have the following statements hold: d,ξ (m). The m-level ambiguity set projected from a z-level ξ-ambiguity set is a subset of the m-level ξ-ambiguity set when directly perturbing m-level distributions. It means the hierarchical structure imposes extra regularization/constraints on the adversary. We now provide the proof for Proposition 3 as follows.
Proof for Proposition 3. Item (1) directly follows the definition of operator T A in Definition 5.
Define the ambiguity sets based on Definition 4, where the cost function is the cost total variance distance.
. Denote the i-th element of g as g i , i ∈ [n]. Let a i ∈ [0, 1] 1×m denote the i-th row of A.
Considering that elements in a i are non-negative and lie in interval [0, 1], we have ) + ((·) + is an operator that replaces negative elements with 0) Similarly, we can prove −g i ≤ ξ, ∀i ∈ [n].

Remark.
A is not a stochastic row matrix, which makes the proof different from the contraction mapping proof in tabular RL settings where the transition matrix is a stochastic row matrix.

B.2.2 Proof for Theorem 2
With the ambiguity set relationships in Proposition 3, we are now ready to prove Theorem 2.
Recall that for notation simplicity, let U m (π) = E π m T t=1 γ t r t . Let C b(m),d T V ,ξ (m) and C b(z),d T V ,ξ (z) denote the ambiguity sets for beliefs over MDPs m and mixtures z, respectively. b(m) and b(z) satisfy b(m) = Z µ z (m)b(z) and are the nominal distributions. For any history-dependent policy π ∈ Π, its value function under different robust formulations are: Proof for Theorem 2.
Given an arbitrary policy π ∈ Π, we have It means that with a nontrivial ambiguity set C b(z),d T V ,ξ (z), the distributionally robust value is more optimistic than the group robust formulation.
Remark The belief robust method with V GR is compatible with a non-adaptive robust problem, where the policy of the decision maker is a Markov policy that only depends on the current state. In contrast, the belief distributionally robust method with V GDR corresponds to an adaptive robust problem, where the decision maker utilizes a history-dependent policy. In other words, it considers both the current state and the information gathered along with interactions. A similar argument but in a non-robust version is presented as Proposition 1. in [5].
Now prove the inequality relationship in item (2) which is V GDR (π) ≥ V DR (π): Based on the projection operator in Definition 5, we change the minimization over belief distribution on mixtures to an equivalent expression that has minimization over belief distribution on MDPs instead.
(based on Definition 5) Then with Proposition 3, which shows the set relationships, we have, It shows that, in general, distributionally robust over high-level latent variable z is more optimistic than that over low-level latent variable m. The hierarchical mixture model structure help regularize the strength of the adversary and generate less conservative policies than the flat model structure.

B.2.3 Proof for Theorem 3
Based on Theorem 2, we can derive the relationships between the optimal values for different formulations.
Since π GDR is the optimal policy for GDR-MDP, we have Since V GDR (π) ≥ V DR (π), ∀π, base on Theorem 2, we have Therefore we have Following similar procedures, we prove that V GDR (π GDR ) ≥ V GR (π GR ) ≥ V R (π R ).
V GDR (π GDR ) ≥ V GDR (π GR ) (since π GDR is the optimal policy for GDR-MDP) (since π GR is the optimal policy for group robust MDP) ≥ V R (π R ).
(since V GR (π) ≥ V R (π), ∀π in Theorem 2) Therefore, we have shown the following inequalities hold: C Environment Details

C.1 Google Research Football
Google Research Football (GRF) is a physics-based 3D soccer simulator for reinforcement learning. This domain presents additional challenges due to its AI randomness, large state-action spaces, and sparse rewards. The RL agent will control one active player on the attacking team at each step and can pass to switch control. The non-active players will be controlled by built-in AI. In our designed 3 vs. 2 tasks, there are three attacking players of a certain type and two defending players, including one player of a chosen type and a goalkeeper.
The dynamics of the 3 vs. 2 tasks are determined by the player types, including central midfield (CM) and centre back (CB), and player capability levels. The mixture index set Z has a cardinality of two, z = 0 and z = 1, corresponding to CM vs. CB (with the goalkeeper) and CB vs. CM (with the goalkeeper), respectively. The built-in CM player tends to go into the penalty area when attacking and guard the player on the wing (physically left or right) when defending, while the CB player tends to guard the player in the middle when defending, and not directly go into the penalty area when attacking. Thus, different patterns of policies are required to solve the tasks from different groups. As shown in Figure 6, in a CM-attacking-CB-defending task, a good solution is to first pass the ball to the player on the wing and then shoot. In a CB-attacking-CM-defending task, a good policy is to directly run into the penalty area and shoot. To further encourage task diversity, we add some noisy actions to a run-into-penalty policy in a CM-attacking-CB-defending task, and to a pass-and-shoot policy in a CB-attacking-CM-defending task, when the controlled player faces high-intensity defense.
For the player capability level, we have two types of settings, players with 1.0 capability attacking while players with 0.7 capability defending (1.0 vs. 0.7), and players with 0.9 capability attacking while players with 0.6 capability defending (0.9 vs. 0.6). The strongest player has a capability level of 1.0. It is worth noting that these settings are more challenging than the original 3 vs. 2 task in GRF (1.0 vs. 0.6) in terms of capability level. Detailed descriptions of the state and action space are shown in Table 2.

C.2 LunarLander
We modify the LunarLander environment [68] by changing the engine mode and engine power. The mixture index set Z has a cardinality of two, z = 0 and z = 1, corresponding to two different engine operation modes, normal mode and left-right-flip mode, respectively. When in left-right-flip mode, the action turning on the left engine in normal mode will turn on the right engine instead, and the action turning on the right engine in normal mode will turn on the left instead. We visualize the tasks in Figure 7. The engine power has two choices which are 3.0 and 6.0. The MDP set M has carnality four corresponding to four combinations of engine mode and engine power. Detailed descriptions of the state and action space are shown in Table 5. Figure 7: LunarLander task visualization. Task 0 and task 1 belong to group 0 (normal mode). Tasks 2 and task 3 belong to group 1 (flip mode).

C.3 HalfCheetah
We modify the joint failure and torso mass of HalfCheetah and build 18 tasks with different dynamics. The joint failure has six choices which correspond to the 6 joints of HalfCheetah. For instance, when the joint failure index is 0, we cannot apply control torque (action) to joint 0. The torso mass has three choices, which are 0.9, 1.0, and 1.1 times the original torso mass. We visualize the joint indexes in Figure 8. Detailed descriptions of the state and action space are shown in Table 6.

D Implementation Details
Trajectory rollout. In both training and testing, we initialize the environment by sampling first a mixture and then an MDP realization. The sampled mixture and MDP are fixed throughout one episode. In our environments with discrete mixtures and MDPs, we can represent the ground truth mixture indexẑ with a one-hot vector eẑ, which is used in the pretraining phase of all baselines and in the whole training phase of baseline G-Exact. For baselines with belief module including GDR, G-Belief, DR, State-R, the actual mixtureẑ and MDP weights µ(m|ẑ) are unknown to the RL agent. Instead, the RL agent is given the number of possible mixtures Z and is able to infer a belief over mixtures b(z) based on a belief update function SE. A detailed algorithm for trajectory rollout is Algorithm 1. For baseline No-Belief, we mask out the beliefs in the input by replacing them with zeros.
Belief update mechanism. In our implementation (Section 7), we use the Bayesian update rule to update beliefs based on the interaction at each timestep. At the beginning of each episode, we initialize a uniform belief distribution b 0 (i) = 1/(|Z|), ∀i ∈ [|Z|]. At timestep t, we update the belief as follows where L represents the likelihood. Letẑ denote the true mixture index for the episode. We let the likelihood L vector be a soft version of the actual one-hot mixture encoding eẑ.
More concretely, at each time step, we first sample a noisy index j where j =ẑ with probability l and j is uniformly sampled from [Z] otherwise. The likelihood L is a vector with dimension |Z|, and ∀i ∈ [|Z|], the i-th element L(i) is There are lots of literature on accurate belief updates [72]. In this work, we utilize a simple but controllable belief update mechanism above, which is more suitable for robustness evaluations since we could explicitly vary the hyperparameters. We leave a more sophisticated design of belief update mechanism for future work.
Belief noise level During robustness evaluation in Section 8, we control the belief noise level ẑ which affects the likelihood L. More concretely, we add another layer of randomness on the estimate ofẑ. Define the noisy mixture index at test-time as z text , we have z test = ẑ with probability ẑ a random index uniforms samples from [|Z|], otherwise During the robust evaluation, the likelihood L test is calculated based on z test . More concretely, at each time step, we first sample a noisy index j where j = z test with probability l and j is uniformly sampled from [Z] otherwise. The likelihood and belief updates are as follows: Distributionally robust training with belief distribution over MDPs (DR) DR has an agent that takes the belief distribution b(m) and state s as inputs. DR uses the same belief updating rule as in GDR to update b(z) at each timestep and then project b(z) to b(m) with µ z (m).
This is a variant of our proposed Group Distributionally Robust DQN, which has a perturbed target taking m-level belief distribution as part of its input. Note that in DR, we still update z-level belief b(z) based on the same belief updating function SE as in GDR. However, in DR, for data pair d, the ambiguity set C b (m),d T V ,ξ is centered at b (m) = T A (b (z)) which is mapped from b (z). We also modify the fast gradient sign attack over b(m) accordingly. We first sample i ∈ [M ] and apply attacks as p(m) j = p(m) j + α b · sign(∇ p(m),j V (p(m), s )), ∀j = i and p(m) i = p(m) i − j =i p(m) j . We iteratively apply the gradient sign attack to find b adv (m) = arg min p(m)∈C b (m),d T V ,ξ a∈A Q θ (p(m), s , a).

D.1 GDR-PPO
We represent the pseudo algorithm of GDR-PPO in Algorithm 3. We collect rollouts with un-perturbed beliefs and use the online rollouts to update the value network. To enhance the robustness to belief ambiguity, we tend to down-weight the probability of trajectories that may lead to large performance drops under the worst-possible belief within the ambiguity set. Hence we construct a pseudo-advantageÂ by subtracting the performance drop R drop from the actual accumulated return. The worst-possible belief is calculated by FGSM.
Algorithm 3: Group Distributionally Robust Training for GDR-PPO Input: Value-net V θ (b(z), s), ambiguity set C ·,d T V ,ξ , training episodes N Initialize data buffer D ; for n = 0 to N do Rollout several episode with Algorithm 1 and append data pairs to D ; if Update Actor-net parameters then Sample batch data from D ; for Each trajectory in the batch do Get advantage for the data pair at timestep t A(b t , s t ) = ; Update Actor-net with PPO ; Return: Actor-net

D.2 Hyperparameters
We show the hyperparameters for training Google Research Football, Lunarlander and Halfcheetah in Table 8, Table 9 and  Table 10, respectively. We select hyperparameters via grid search.

E Additional Ablation Study
In this section, we show how the ambiguity set size and the pretrain episodes affect the training stability and robustness of DR, which maintains a belief over MDPs. Compared with our proposed GDR, DR omits the hierarchical structure.
E.1 The effect of Ambiguity Set Size Figure 9 shows the effect of the ambiguity set in HalfCheetah. All curves in Figure 9 are pre-trained in the first 100000 episodes. With ambiguity set size 0.01 and 0.05, the DR does not crash and converge to a non-negative value. Comparing Figure 9 (b) for DR with Figure 3 (c) for GDR, we can conclude that GDR is less sensitive to the ambiguity set size along training since it converge to a non-negative value with a larger range of ambiguity set size. Comparing Figure 5 (a) with Figure 4 (c) for our proposed GDR, we can conclude that the hierarchical structure enhances the robustness to belief noise since the robustness performance for GDR consistently outperforms that of DR for different ambiguity set sizes.  Figure 10 shows the effect of the pretrain episodes in HalfCheetah. All curves in Figure 10 has an ambiguity set size 0.2. Figure 10 shows that even pretraining for 900000 episodes, DR still will crash after the pretraining phase. It shows that DR is less sensitive to the pretraining episodes compared with the ambiguity set size. Figure 10: Ablation study about the effect of pretrain episodes on DR's robustness and training stability. We choose the number of pretrain episodes among 100000, 300000, 500000, 700000, and 900000.