Balanced Domain Randomization for Safe Reinforcement Learning

Abstract

Reinforcement Learning (RL) has enabled autonomous agents to achieve superhuman performance in diverse domains, including games, navigation, and robotic control. Despite these successes, RL agents often struggle with overfitting to specific training environments, which can hinder their adaptability to new or rare scenarios. Domain randomization is a widely used technique designed to enhance policy robustness by training models in a variety of randomized contexts. In traditional domain randomization, however, models tend to prioritize learning from more common domains, often neglecting rare ones. To address this imbalance, we propose Balanced Domain Randomization (BDR) that balances the learning focus based on the rarity of contexts. We assess context rarity in the embedding space using statistical analysis of context vectors and adjust the loss weights for each transition based on this rarity. This ensures that the agent dedicates adequate attention to rare domains. Our experimental results show that BDR efficiently enhances worst-case performance, significantly improving the robustness of RL-based robotic controllers across diverse conditions. This study provides a robust framework for RL agents, reducing risks in uncommon scenarios and ensuring reliable performance in varied environments.

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.

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 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.

3. Problem Definition

3.1. Contextual Reinforcement Learning Framework

We introduce a contextual Reinforcement Learning (RL) framework for training agents under varying contextual conditions, where the context encapsulates the physical properties of the environment. This framework is grounded in the contextual Markov Decision Process (MDP) [18], defined as $M=(S,A,T,R,\gamma ;c)$. In this formulation, S represents the set of states, A denotes the set of actions, and $T\left(s^{\prime }\right|s,a;c)$ is the transition function, indicating the probability of moving from state s to $s^{\prime }$, given action a in the context c. The reward function $R(s,a,s^{\prime };c)$ provides the immediate reward given the transition, and the discount factor $\gamma$ models the trade-off between current and future rewards. The context c defines system parameters that determine the physical properties of the environment. In domain randomization, environments are considered to belong to the same domain when they share the same context or system parameters.

In practical scenarios, the context c may not be directly observable, making the contextual MDP a special case of Partially Observable MDPs (POMDPs) [42,43], where contexts act as hidden variables. Under this partial observability, contextual RL has two main objectives: (1) encoding a history of recent transitions into a context embedding, and (2) learning a contextual policy $\pi (s;c):S\times C\to A$ that determines the optimal action to maximize the expected return $\mathbb{E}\left[{\sum }_{t=0}^{\infty }{\gamma }^{t}R(s_t,a_t,s_{t+1};c_t)\right]$. Here, the context encoder is utilized to produce a context embedding using neural networks [14]. Specifically, it maps the recent state–action history into a context embedding, which implicitly estimates the underlying environmental contexts. Then, the context embedding is used as an additional input of the policy for more informed decision-making.

3.2. Learning Imbalance in Domain Randomization

Domain randomization enhances the adaptability of RL agents by training them across various environmental contexts. However, our analysis reveals a significant limitation: RL agents often prioritize learning in specific domains, which can lead to suboptimal performance in others.

To explore this issue, we conduct experiments on the Pendulum task, where mass and length are randomized as system parameters. We compare the performance of a baseline model trained with Uniform Domain randomization (UDR) to near-optimal policies trained on individual domains. Figure 2 presents the average episodic return of the models across different domains. The baseline model shows inconsistent performance across various contexts, with notable deficits when the pendulum length is 0.1. In contrast, the results of the near-optimal policies, which represent the potential upper bounds of performance, demonstrate that a policy can be sufficiently trained, even when the pendulum length is 0.1.

These results suggest that RL agents, when trained on uniformly sampled domains, tend to prioritize specific domains over others. We hypothesize that this imbalance stems from the relative rarity of domains in the context embedding space. Although this imbalance might enhance the average performance in common situations, it undermines the agent’s adaptability in rare but potentially critical scenarios, presenting significant challenges in real-world applications.

4. Balanced Domain Randomization

To address the imbalance inherent in domain randomization, we propose a method that balances the learning proportions of contexts according to their rarity in the embedding space. We measure the rarity of a context based on the distance between its context vector and the center of context embeddings. This measure of rarity is then used to evaluate the importance of each context during training, ensuring that the agent pays sufficient attention to rare domains. Our approach aims to improve the agent’s performance in worst-case scenarios, enhancing overall robustness and reliability. The overall procedure is illustrated in Figure 1 and detailed in Algorithm 1.

Algorithm 1: Balanced Domain Randomization

4.1. Context Embedding

In our training framework, the agent interacts with the environment to collect data stored in a replay buffer, which is used for learning steps. At the j-th learning step, we sample a mini-batch $B_j$ of transitions from the replay buffer. ${\tau }_t\in B_j$ is a transition of time t stored in the buffer. ${\tau }_t$ includes the current state $s_t$, the action $a_t$, the reward $r_t$, the next state $s_{t+1}$, and the history of state–action pairs $h_t$. The context encoder f maps $h_t$ into a context vector $z_t$, which is then fed into the policy as a condition. The context vector captures contextual information from the agent’s past interactions with the environment. This structure allows the agent to take adaptive actions based on the current environment dynamics, even when it gains the same state as its input.

4.2. Accessing the Rarity of Contexts

As explained in Section 3.2, RL agents tend to prioritize specific domains over others in randomized environments. We hypothesize that this imbalance problem in domain randomization arises from the rarity of contexts encountered during training. It is caused by the typical learning framework that the policy is trained to increase the expected reward of its rollouts, which may lead to overfitting to domains with common dynamics. In this regard, it is necessary to evaluate the rarity of the contexts to detect domains in which the policy may potentially underfit.

Given a history $h_t$, the corresponding context vector is computed as $z_t=f\left(h_t\right)$. We define the distance between $z_t$ and the center of context embeddings as its rarity,

$$\mathrm{Rarity}\left(h_t\right)=d(f\left(h_t\right),\overline{z})$$

(1)

where $d(·,·)$ is a distance metric and $\overline{z}$ represents the expected context vector. Any distance metric can be employed, such as Euclidean or cosine distance.

Without any specific restrictions to context vectors, variables in the context embedding usually have different variances and exhibit correlations with each other. It is not considered in traditional distance measures like Euclidean distance, since it assumes all features are independent and equally scaled. Therefore, we utilize the Mahalanobis distance [44,45], which accounts for both the variance of individual variables and the correlations between them using the covariance matrix. This makes it particularly well-suited for high-dimensional spaces, such as our context embedding space, where features often exhibit complex relationships. Mathematically, the Mahalanobis distance between the embedding vector $z_t$ and the distribution with mean $\overline{z}$ and covariance matrix $\mathsf{\Sigma }$ is given by

$$d(f\left(h_t\right),\overline{z};\mathsf{\Sigma })=\sqrt{{(z_t-\overline{z})}^{\top }{\mathsf{\Sigma }}^{-1}(z_t-\overline{z})}.$$

(2)

This distance normalizes the difference between $z_t$ and the mean $\overline{z}$ by the variance and correlations present in the embedding space, ensuring that deviations in directions with low variance are weighted more heavily. Consequently, Mahalanobis distance allows us to effectively capture rare contexts, which may be overlooked in traditional distance metrics.

4.3. Reweighting Training Loss

To ensure that the policy is adequately trained on rare domains during training, we adjust the training loss weights based on the computed rarity of each context. We introduce Balanced Domain Randomization (BDR), a method that encourages the agent to focus on transitions sampled from rare domains during updates, enhancing generalization performance. This approach involves the process of assigning weights to training samples by evaluating the rarity scores based on their context vectors. Specifically, we normalize the rarity scores to determine weights using the softmax function,

$$w\left(h_t\right)={\displaystyle \frac{exp\left(\mathrm{Rarity}\left(h_t\right)\right)}{{\sum }_{h_t\in H_j}exp\left(\mathrm{Rarity}\left(h_t\right)\right)}}$$

(3)

where $H_j$ is the set of histories in the j-th mini-batch. The weighted loss for the RL model is then computed as

$$L_{\mathrm{weighted}}\left(B_j\right)=\sum _{{\tau }_t\in B_j}\beta w\left(h_t\right)l_{\mathrm{RL}}\left({\tau }_t\right)+(1-\beta )l_{\mathrm{RL}}\left({\tau }_t\right)=\sum _{{\tau }_t\in B_j}\{\beta w\left(h_t\right)+(1-\beta )\}{l}_{\mathrm{RL}}\left({\tau }_t\right)$$

(4)

where $l_{RL}\left({\tau }_t\right)$ is the original RL loss for a single transition ${\tau }_t$, $L_{\mathrm{weighted}}\left(B_j\right)$ is the total weighted loss for a batch $B_j$, and $\beta$ is a scalar value that balances the weighted loss and the original loss. In this work, $\beta$ is set to 0.5.

4.4. Context Statistics with Exponential Moving Average

To compute the rarity of each context vector, we estimate the statistics of the context vectors. However, since the context encoder is continually updated during training, re-encoding all past histories in the replay buffer would be necessary to accurately compute these statistics. To reduce this computational burden, we employ Exponential Moving Average (EMA) to gradually update the approximation with mini-batch samples. The mean context vector is updated as follows:

$${\overline{z}}_j=\alpha \overline{f}\left(H_j\right)+(1-\alpha ){\overline{z}}_{j-1}=\alpha {\displaystyle \frac{1}{|H_j|}}\sum _{h_t\in H_j}f\left(h_t\right)+(1-\alpha ){\overline{z}}_{j-1}$$

(5)

where ${\overline{z}}_j$ is the estimate of the mean context vector after the j-th update step, $\overline{f}\left(H_j\right)$ is the average context vector of the j-th mini-batch, and $\alpha$ is the smoothing factor for the EMA, with $0<\alpha \le 1$. In this work, we set $\alpha =0.05$. The covariance matrix ${\mathsf{\Sigma }}_j$ is estimated similarly by computing the sample covariance matrix of the context vectors and updating it with EMA.

$${\mathsf{\Sigma }}_j=\alpha {\displaystyle \frac{1}{|H_j|-1}}\sum _{h_t\in H_j}(f\left(h_t\right)-\overline{f}\left(H_j\right)){(f\left(h_t\right)-\overline{f}\left(H_j\right))}^{\top }+(1-\alpha ){\mathsf{\Sigma }}_{j-1}$$

(6)

5. Experiments

5.1. Setup

To evaluate the effectiveness of the proposed Balanced Domain Randomization (BDR) method, we conduct experiments on two control tasks: Pendulum and Walker. We utilize the OpenAI Gymnasium [46] for the Pendulum task and the DeepMind Control (DMC) suite [47] for the Walker task. To test the adaptability and robustness of the Reinforcement Learning (RL) agents across varied dynamic conditions, we apply domain randomization facilitated by the CARL framework [19].

In the Pendulum task, the goal is to swing a pendulum from a hanging position to an upright position. The agent controls the torque applied to the pendulum, based on its position and angular velocity. The agent is penalized for deviations from the upright position as well as control costs. We introduce variability in the environment by altering the mass and length of the pendulum. For both training and evaluation, we use combinations of mass and length values from the set {0.1, 0.5, 1.0, 1.5, 2.0}, creating a broader range of scenarios to assess the agent’s adaptability.

In the Walker task, the goal of a bipedal walker is to move forward efficiently while keeping balance. The agent controls the torques at its joints, based on its height, orientation, and velocity. The agent is rewarded for reaching the target velocity and maintaining balance. We alter the friction coefficients and the gravitational forces, using combinations of gravity values {1, 5, 10, 15, 20} and friction coefficients {0.1, 0.5, 1.0, 1.5, 2.0}.

In this framework, the context encoder maps the history of the agent’s state–action pairs into a context vector using a Multi-Layer Perceptron (MLP) network. Here, we use five recent state–action pairs as we empirically find it sufficient to infer the environmental dynamics. This MLP consists of four layers, with 256, 128, 64, and 10 neurons, respectively. The context encoder outputs a context vector $z_t$, which represents a compressed embedding of the environmental dynamics. This vector is then fed into the policy network to condition the policy on the current environmental context. The context encoder is trained jointly with the contextual policy in an end-to-end manner to maximize the reward. For RL optimization, we employ the Twin Delayed Deep Deterministic policy gradient (TD3) algorithm [48].

5.2. Evaluation Metric

We evaluate the performance of the agents based on the episodic return, defined as the total sum of rewards accumulated over an episode. For a comprehensive evaluation, we run N episodes and aggregate the episodic returns by calculating the average for each domain. In this work, we run 100 episodes for each domain. Our evaluation focuses on both worst-case and average-case performance, using the following metrics:

$$\mathrm{Minimum}\mathrm{Episodic}\mathrm{Return}=\underset{d\in D}{min}\left(R_d\right)$$

The minimum episodic return metric captures the agent’s performance in the most challenging domain, reflecting its robustness in difficult scenarios.

$$\mathrm{Mean}\mathrm{Episodic}\mathrm{Return}\left(\mathrm{Worst-}k\right)={\displaystyle \frac{1}{k}}\sum _{i=1}^k\left(R_{\mathrm{sorted},i}\right)$$

The mean episodic return of the worst-k samples is calculated by first sorting all episodic returns in ascending order, from worst to best. The worst-performing k episodes are then selected, and their mean is computed. This metric provides insight into its resilience under unfavorable conditions, as it focuses on the episodes where the agent performs the worst.

$$\mathrm{Mean}\mathrm{Episodic}\mathrm{Return}={\displaystyle \frac{1}{\left|D\right|}}\sum _{d\in D}\left(R_d\right)$$

The mean episodic return indicates the average-case performance, highlighting how well the agent adapts to the variety of environments on average.

5.3. Results

We evaluate the performance of the Uniform Domain Randomization (UDR) baseline and Balanced Domain Randomization (BDR) by keeping their configurations identical, including model architecture and hyperparameters, with the only distinction being the weighted RL loss component. Additionally, we compare the performance with Automatic Domain Randomization (ADR) [49], which gradually expands the randomization ranges during training. ADR realizes a training curriculum that gradually expands a distribution over training environments, from the initial range concentrated on a single environment.

Table 1. Mean and minimum episodic returns in Pendulum and Walker tasks. The best performance is marked in bold.

Task	Method	Minimum Episodic Return	Mean Episodic Return (Worst 10%)	Mean Episodic Return
Pendulum	Baseline	−1263.97	−1221.57	−392.70 ± 319.62
	ADR	−1312.10	−1332.14	−569.46 ± 428.23
	BDR	−1199.37	−1132.30	−379.14 ± 295.16
Walker	Baseline	539.62	516.80	878.21 ± 106.04
	ADR	224.65	220.40	560.88 ± 198.98
	BDR	673.49	638.59	885.18 ± 88.89

summarizes the quantitative results. In the Pendulum task, BDR improves the minimum episodic return by 5.11% (from −1263.97 to −1199.37), the mean episodic return (worst 10%) by 7.31% (from −1221.57 to −1132.30), and the mean episodic return by 3.45% (from −392.70 to −379.14), over the baseline. In the Walker task, BDR enhances the minimum episodic return by 24.81% (from 539.62 to 673.49), the mean episodic return (worst 10%) by 23.57% (from 516.80 to 638.59), and the mean episodic return by 0.79% (from 878.21 to 885.18). These consistent improvements across both tasks suggest that BDR reduces overfitting to common domains and enhances adaptability in rare scenarios.

Figure 3 and Figure 4 show the average episodic return for each domain in the Pendulum and Walker tasks, respectively. We observe that the improvements are especially evident in domains with extreme conditions, such as the low length in the Pendulum task and the low gravity in the Walker task. Specifically, when the length is 0.1 in the Pendulum task, the BDR achieves an average episodic return of −670.6, outperforming the baseline, which exhibits an average return of −772. This represents a 13.13% improvement. When the gravity is 1 in the Walker task, the BDR improves an average episodic return by 11.38%, achieving an average episodic return of 763.4 while the baseline only achieves an average return of 685.4.

We also find that ADR is not effective in learning a balanced behavior over diverse environmental contexts, as it does not improve the performance over the baseline both in the Pendulum and Walker tasks. Although there is still room for further improvement in worst-case scenarios, BDR effectively improves performance in challenging domains while maintaining its average-case performance.

5.4. Analysis

Weight assignment. We also analyze how BDR adjusts the weights for training samples from different domains. Figure 5 visualizes the average sample weights assigned to each domain. The analysis indicates that BDR gives higher weights to domains with smaller lengths in the Pendulum task (Figure 5a) and lower gravity in the Walker task (Figure 5b). It demonstrates that BDR effectively targets and emphasizes rare domains where the baseline performs poorly by measuring the rarity in the embedding space.

Finally, we assess the efficiency of BDR by observing performance changes throughout the training process. Figure 6 illustrates the differences in learning performance between the two methods at each training checkpoint. The results demonstrate that BDR efficiently optimizes worst-case performance while maintaining average-case performance, highlighting its robustness and reliability in varied conditions during training.

Balancing parameter. To further investigate the impact of the balancing parameter $\beta$, we conducted an ablation study by varying its value from 0 (which denotes uniform randomization) to 1 (which places full emphasis on rare domains).

Table 2. Mean and minimum episodic returns of the BDR with different values of $\beta$. The best performance is marked in bold.

Task	Method	Minimum Episodic Return	Mean Episodic Return (Worst 10%)	Mean Episodic Return
Pendulum	Baseline ($\beta =0$)	−1263.97	−1221.57	−392.70 ± 319.62
	$\beta =0.25$	−1299.66	−1244.98	−450.21 ± 317.04
	$\beta =0.5$	−1199.37	−1132.30	−379.14 ± 295.16
	$\beta =0.75$	−1227.09	−1134.78	−411.02 ± 319.18
	$\beta =1$	−1386.99	−1186.47	−464.74 ± 323.14
Walker	Baseline ($\beta =0$)	539.62	516.80	878.21 ± 106.04
	$\beta =0.25$	668.35	633.97	900.73 ± 85.59
	$\beta =0.5$	673.49	638.59	885.18 ± 88.89
	$\beta =0.75$	454.86	455.92	767.93 ± 109.67
	$\beta =1$	497.80	567.64	851.82 ± 106.13

shows that $\beta =0.5$ yields the best performance in most cases. However, performance degrades for higher values of $\beta$, since overemphasizing rare domains may hinder the agent learning its general strategies in more common scenarios. Although it may be necessary to appropriately tune the value of $\beta$, depending on tasks, the empirical results show that $\beta =0.5$ provides the best trade-off between learning from rare and common scenarios in both tasks.

5.5. Evaluation in Safety-Critical Navigation Tasks

To evaluate the efficacy of the BDR framework in more complex and safety-critical scenarios, we adopt a 2D safe navigation task using the SafetyGymnasium environment [50]. This environment introduces hazards as well as goal-reaching objectives, representing safety-critical real-world challenges such as autonomous navigation and robotic deployment. The objective is to reach a randomly generated goal while avoiding randomly placed unsafe regions (hazards). The agent controls actuators for linear and rotational velocity, receiving rewards for reaching the goal and penalties for entering unsafe regions. To introduce variability in dynamics, we alter both the mass and the moment of inertia (MOI). For training and evaluation, we use combinations of mass and MOI from the set {0.1, 0.3, 0.5, 0.7, 1.0}. As well as episodic return, we measure episodic cost as a safety-specific metric. The episodic cost is measured as accumulated penalties for entering unsafe regions.

Table 3. Mean and minimum episodic returns and costs in the safe navigation task. The best performance is marked in bold.

Task	Method	Minimum Episodic Return	Mean Episodic Return (Worst 10%)	Mean Episodic Return
Navigation	Baseline	0.12	−1.82	2.81 ± 1.37
	BDR	0.72	−0.35	9.88 ± 4.84
		Maximum Episodic Cost	Mean Episodic Cost (Worst 10%)	Mean Episodic Cost
	Baseline	109.97	484.53	61.62 ± 16.96
	BDR	82.09	317.56	55.69 ± 11.09

reports the quantitative results. BDR improves the minimum episodic return by 500% (from 0.12 to 0.72), the mean episodic return (worst 10%) by 80.77% (from −1.82 to −0.35), and the mean episodic return by 251.60% (from 2.81 to 9.88), over the baseline. For safety-specific evaluation, BDR reduces the minimum episodic cost by 25.35% (from 109.97 to 82.09), the mean episodic cost (worst 10%) by 34.46% (from 484.53 to 317.56), and the mean episodic cost by 9.62% (from 61.62 to 55.69), over the baseline. These results demonstrate the efficacy of BDR in reducing safety risks while maintaining task performance.

6. Conclusions

This paper presents Balanced Domain Randomization (BDR) to address the imbalance problem in contextual Reinforcement Learning (RL) by prioritizing the agent to learn its policy in rare domains during training. Our key idea is to automatically identify insufficiently learned domains by evaluating context rarity based on the statistics of context vectors. BDR leverages Mahalanobis distance in the context embedding space to assess rarity and adaptively adjust the training weights, ensuring sufficient exposure to rare scenarios. Our experimental results on the Pendulum and Walker tasks demonstrate that BDR improves worst-case performance, with a 5.11% gain over the baseline in Pendulum and 24.81% in Walker, while maintaining average performance. BDR has the potential to enhance the robustness of RL agents for real-world applications, such as robotics and autonomous driving, by improving sim-to-real transfer and reducing the risk of failure in unpredictable scenarios. While some hyperparameter tuning is required, this could be further explored in future work. Moreover, we expect that BDR’s generalization capabilities can be extended to entirely unseen domains.