Group Distributionally Robust RL 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.

To handle belief uncertainty, we formulate the group distributionally robust Markov Decision Process (GDR-MDP). GDR-MDP naturally balances the worst-case (robustness) and average performance by leveraging the adaptive belief and the distributionally robust formulation. To build GDR-MDP, we first formulate Hierarchical-Latent MDP (HLMDP), which utilizes a mixture model over MDPs to encode task subpopulations. HLMDP has a high-level latent variable z as the mixture, and a low-level m to represent tasks. GDR-MDP also encodes task subpopulations and additionally formulates the robustness w.r.t. the ambiguity of the adaptive belief b(z) over mixtures.

Figure 4. Robustness to belief noise. GDR is more robust to belief noises compared with all the other baselines.

We visualize the 4 tasks from 2 groups in the Google Research Football environment. In this section, the videos are generated with the policies using ground-truth group indexes.

We show the rollouts of different methods under various levels of noisy beliefs. The reported success rates here are averaged across the visualized 20 episodes. We term our proposed method as GDR and denote the best success rates in red.