# Balanced Domain Randomization for Safe Reinforcement Learning
## 1 Introduction  
Reinforcement Learning (RL) has significantly advanced the capabilities of autonomous agents to perform complex tasks across various domains [1; 2; 3], including board games [4; 5], real-time strategy games [6; 7], robotic manipulation [8; 9], and navigation [10; 11]. The core strength of RL lies in its ability to autonomously learn policies from environmental interactions through trial and error. However, RL models often overfit to the specific contexts encountered during training, severely limiting their adaptability and robustness in real-world applications where scenarios may differ significantly from their training data.  
Domain randomization is a widely adopted technique to enhance the robustness of RL agents by exposing them to various simulation environments during training [9; 12; 13; 14; 15; 16; 17]. Despite its success, models often prioritize learning from frequent domains with similar dynamics, resulting in insufficient learning from rare but potentially critical scenarios. This imbalance can result in suboptimal policies that perform well in common situations, but fail in less frequent, challenging ones. To develop truly robust RL agents, it is essential to ensure adequate training on these rare domains.  
In this paper, we propose a novel approach that balances the learning focus based on the rarity of contexts (as illustrated in Figure 1). We measure the rarity of a context in the embedding space using the statistics of context vectors, under an assumption that policies may not be sufficiently trained in contexts where their embeddings significantly deviate from the distribution. We adjust the training loss weights according to the rarity of these contexts, encouraging the agent to focus more on rare ones. This method ensures that the agent is thoroughly exposed to a wide range of domains. Our main contributions are as follows:Identification of rare domains: we identify rare domains where policies may be undertrained by analyzing context embedding distances.
* Balancing domain randomization: we propose a balancing mechanism that assigns greater weight to rare domains during training to achieve balanced domain randomization.
* Empirical validation: our experiments demonstrate that the proposed method efficiently improves worst-case performance, enhancing the robustness of RL agents, especially in novel situations.  
In summary, this study addresses the critical need for robustness in RL by introducing a balanced domain randomization approach that ensures a harmonious representation of both frequent and rare domains, improving the reliability of RL agents for deployment in diverse and unpredictable environments.  
The remainder of this paper is organized as follows: Section 2 reviews related work on contextual RL and domain randomization. Section 3 explains the contextual reinforcement learning framework, and discusses the imbalance problem in standard domain randomization. In Section 4, we propose the Balanced Domain Randomization (BDR) method, detailing our training framework and context rarity evaluation based on the statistics of context embedding. Section 5 shows the experimental results on benchmark tasks. We conclude with a discussion of our findings and future directions in Section 6.