# Balanced Domain Randomization for Safe Reinforcement Learning
## 4 Balanced Domain Randomization
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 \(z\) and covariance matrix \(\Sigma\) is given by

\[d(f(h_{t}),z\ ;\ \Sigma)=\sqrt{(z_{t}-z)^{\top}\Sigma^{-1}(z_{t}-z)}. \tag{2}\]

This distance normalizes the difference between \(z_{t}\) and the mean \(\bar{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. ### 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(h_{t})=\frac{\exp(\text{Rarity}(h_{t}))}{\sum_{h_{t}\in H_{t}}\exp(\text{ Rarity}(h_{t}))} \tag{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_{\text{weighted}}(B_{j})=\sum_{\tau_{t}\in B_{j}}\beta\ w(h_{t})l_{\text{RL}}( \tau_{t})+(1-\beta)\ l_{\text{RL}}(\tau_{t})=\sum_{\tau_{t}\in B_{j}}\{\beta w(h _{t})+(1-\beta)\ \}\ l_{\text{RL}}(\tau_{t}) \tag{4}\]

where \(l_{RL}(\tau_{t})\) is the original RL loss for a single transition \(\tau_{t}\), \(L_{\text{weighted}}(B_{j})\) 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. ### 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, 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:

\[\bar{z}_{j}=\alpha\ \bar{f}(H_{j})+(1-\alpha)\ \bar{z}_{j-1}=\alpha\ \frac{1}{|H_{j}|}\sum_{h_{t}\in H_{j}}f(h_{t})+(1- \alpha)\ \bar{z}_{j-1} \tag{5}\]

where \(\bar{z}_{j}\) is the estimate of the mean context vector after the \(j\)-th update step, \(\bar{f}(H_{j})\) is the average context vector of the \(j\)-th mini-batch, and \(\alpha\) is the smoothing factor for the EMA, with \(0<\alpha\leq 1\). In this work, we set \(\alpha=0.05\). The covariance matrix \(\Sigma_{j}\) is estimated similarly by computing the sample covariance matrix of the context vectors and updating it with EMA. \[\Sigma_{j}=\alpha\ \frac{1}{|H_{j}|-1}\sum_{h_{t}\in H_{j}}(f(h_{t})-\bar{f}( H_{j}))(f(h_{t})-\bar{f}(H_{j}))^{\top}+(1-\alpha)\ \Sigma_{j-1} \tag{6}\]

## 5 Experiments

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