# Balanced Domain Randomization for Safe Reinforcement Learning
## 2 Related Work  
Contextual Reinforcement Learning (RL) extends the standard RL framework by incorporating environmental contexts, which represent varying physical properties that affect the transition dynamics of environments. This approach enables agents to learn policies that generalize across a range of different situations. Contextual RL has been extensively studied to develop adaptable and robust RL agents, with applications ranging from game playing to robotic control, where agents must adapt to different levels or environmental dynamics [18, 19].  
Domain randomization is a prominent technique employed to improve the robustness and generalization of RL agents by exposing them to a diverse set of simulated environments during training [9, 12, 13, 14, 15, 16, 17, 20, 21, 22, 23]. By varying the environmental contexts, such as physical properties, domain randomization encourages the agent to learn policies that are adaptable across various environments. This approach has been particularly successful in sim-to-real transfer, where agents trained in simulation are deployed in the real world [24, 25, 26]. However, while domain randomization improves generalization, RL agents often prioritize learning in certain domains, which may share similar dynamics. Our analysis shows that uniformly sampling domains does not ensure balanced learning, leaving rare but critical scenarios  
Figure 1: An overall framework of Balanced Domain Randomization (BDR). BDR reweights training domains based on their rarity in the embedding space. The framework uses a context encoder to capture contextual information about the current environment from a recent history of state–action pairs. The rarity evaluation module estimates the rarity of each context, adaptively adjusting the weight assigned to each domain. This ensures that rare scenarios are sufficiently visited during training.  
underrepresented. This imbalance within domain randomization is the core challenge addressed by our work.  
Data imbalance is a well-documented issue across various machine learning domains. To mitigate its effects, several approaches have been proposed. Data rebalancing techniques, such as oversampling rare classes or undersampling common ones, aim to create more balanced datasets [27; 28; 29]. Another strategy is loss reweighting, which adjusts the training weights to emphasize rare classes, thus generating more balanced training signals [30; 31; 32; 33; 34]. Transfer learning approaches have also been used to transfer features learned from common classes to rare ones [35; 36]. Additionally, curriculum learning methods gradually introduce more complex or rarer scenarios as training progresses, encouraging the agent to learn simpler tasks before more challenging ones [37; 38; 39].  
Our proposed method builds on the concept of loss reweighting, but addresses an additional challenge: autonomously detecting rare domains. Existing approaches in visual domains have attempted to measure the rarity of high-dimensional data using embedding distances [40; 41]. Inspired by these approaches, we evaluate the rarity of contexts based on their embedding vectors in contextual RL. Our proposed method effectively mitigates data imbalance by integrating loss reweighting with the rarity measure, ensuring that RL agents receive adequate exposure to rare scenarios.