# Balanced Domain Randomization for Safe Reinforcement Learning
## 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. ```
1 Initialize context encoder \(f\) and RL network \(\theta\);
2 Initialize average context vector \(\mathbf{z}\leftarrow\mathbf{0}\);
3 Initialize covariance matrix \(\mathbf{\Sigma}\leftarrow\mathbf{\mathrm{I}}\); /* Learning Phase */
4 forach update step \(j\)do
5 Sample a random mini-batch of transitions \(B_{j}=(S_{j},A_{j},R_{j},S_{j+1},H_{j})\) from replay buffer \(\mathcal{D}\)
6 Compute the mean of context embeddings in the mini-batch: \(f(H_{j})=\sum_{h_{t}\in H_{j}}f(h_{t})/|H_{j}|\)
7 Update the statistics \(\mathbf{z}_{j},\mathbf{\Sigma}_{j}\)
8 Compute rarity scores: \(\mathrm{Rarity}(h_{t})\gets d(f(h_{t}),\mathbf{z}_{j}\ ;\ \mathbf{\Sigma}_{j})\) for each \(h_{t}\) in \(H_{j}\)
9 Assign weights using softmax: \(w(h_{t})\leftarrow\exp(\mathrm{Rarity}(h_{t}))/\sum_{h_{t}\in H_{j}}\exp( \mathrm{Rarity}(h_{t}))\)
10 Compute weighted RL loss: \(L_{\mathrm{weighted}}(B_{j})=\sum_{\tau_{i}\in B_{j}}\{\beta w(h_{t})+(1- \beta)\}\ l_{\mathrm{RL}}(\tau_{i})\)
11 Update RL network using weighted RL loss
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
310
32
334
34
356
367
378
389
4001
411
4202
4391
444
4503
461
471
482
492
493
494
5045
5106
5107
5108
5109
6010
6111

51120
52012
53013
5314
5415
55216
5530
5417
55418
5



### 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. ### 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(h_{t})\). We define the distance between \(z_{t}\) and the center of context embeddings as its rarity,

\[\text{Rarity}(h_{t})=d(f(h_{t}),z) \tag{1}\]

where \(d(\cdot,\cdot)\) is a distance metric and \(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.