SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Due to discretization, for each quantile  $\tau_{j}$ , the minimizer of the 1-Wasserstein distance within that quantile is at  $\theta$  that satisfies  $F_{Z}(\theta) = \frac{1}{2} (\tau_{j - 1} + \tau_{j}) = \hat{\tau}_{j}$  (see Lemma 2 of Dabney et al. [23]). Consequently, to approximate  $Z$ , we can simultaneously optimize for all of the quantiles,  $\{\theta_j\}_{j = 1}^N$ , by minimizing the following loss:
$$
\sum_ {j = 1} ^ {N} \mathcal {L} _ {\mathrm {Q R}} ^ {\hat {\tau} _ {j}} (\theta_ {j}) = \sum_ {j = 1} ^ {N} \mathbb {E} _ {z \sim Z} \left[ \rho_ {\hat {\tau} _ {j}} (z - \theta_ {j}) \right].
$$
To stabilize optimization, QR-DQN uses an modified version of quantile loss called quantile Huber loss which utilizes the Huber loss with hyperparameter  $\kappa$ :
$$
\mathcal {L} _ {\kappa} (u) = \frac {1}{2} u ^ {2} \mathbb {1} (| u | \leq \kappa) + \kappa \left(| u | - \frac {1}{2} \kappa\right) \mathbb {1} (| u | > \kappa).
$$
The quantile Huber loss is defined as:
$$
\rho_ {\tau} ^ {\kappa} (u) = | \tau - \mathbb {1} (u <   0) | \mathcal {L} _ {\kappa} (u).
$$
In QR-DQN, both the approximating distribution and the target distribution are discretized with quantiles. The target distribution is computed with bootstrapping through the distributional Bellman operator. Concretely, given a target network,  $\theta^{\mathrm{target}}$ , and transition tuple  $(s,a,r,s')$ , the target distribution is approximated as:
$$
Q \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right) = \frac {1}{N} \sum_ {j = 1} ^ {N} \boldsymbol {\theta} _ {j} ^ {\text {t a r g e t}} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\prime}\right)
$$
$$
\boldsymbol{a}^{\star} = \operatorname *{arg  max}_{\boldsymbol{a}^{\prime}}Q(\boldsymbol{s}^{\prime},\boldsymbol{a}^{\prime})
$$
$$
\mathcal {T} \boldsymbol {\theta} _ {j} = r + \gamma \boldsymbol {\theta} _ {j} ^ {\text {t a r g e t}} \left(\boldsymbol {s} ^ {\prime}, \boldsymbol {a} ^ {\star}\right). \tag {15}
$$
The TD target Q-value distribution's quantiles and distribution are:
$$
\{\mathcal {T} \boldsymbol {\theta} _ {1}, \mathcal {T} \boldsymbol {\theta} _ {2}, \dots , \mathcal {T} \boldsymbol {\theta} _ {N} \}, \quad \mathcal {T} Z \stackrel {{d}} {{=}} \frac {1}{N} \sum_ {j = 1} ^ {N} \delta (\mathcal {T} \boldsymbol {\theta} _ {j}).
$$
Treating the target distribution as an oracle (i.e., not differentiable), we update the current quantile estimates to minimize the following quantile regression loss:
$$
\sum_ {j = 1} ^ {N} \mathbb {E} _ {z \sim \mathcal {T} Z} \left[ \rho_ {\hat {\tau} _ {j}} ^ {\kappa} (z - \boldsymbol {\theta} _ {j} (\boldsymbol {s}, \boldsymbol {a})) \right] = \frac {1}{N} \sum_ {j = 1} ^ {N} \sum_ {i = 1} ^ {N} \rho_ {\hat {\tau} _ {j}} ^ {\kappa} (\operatorname {s g} [ \mathcal {T} \boldsymbol {\theta} _ {i} ] - \boldsymbol {\theta} _ {j} (\boldsymbol {s}, \boldsymbol {a})). \tag {16}
$$
sg denotes the stop gradient operation as we do not optimize the target network (standard DQN practice).
E Algorithmic Details
In this section, we describe EDE in detail and highlight its differences from previous approaches. In Algorithm 1, we describe the main training loop of EDE. For simplicity, the pseudocode omits the
initial data collection steps during which the agent uses a uniformly random policy to collect data and does not update the parameters. We also write the parallel experience collection as a for loop for easier illustration, but in the actual implementation, all loops over  $K$  are parallelized. The main algorithmic loop largely resembles the standard DQN training loop. Most algorithmic differences happen in how exploration is being performed i.e., how the action is chosen and how the ensemble members are updated.
In Algorithm 2, we describe how EDE chooses exploratory actions during training. For each state, we first extract the features with the shared feature extractor  $f$  and then compute the state-action value distributions for each ensemble head. With the state-action value distributions, we compute the estimate for the epistemic uncertainty as well as the expected Q-values. Using these quantities, we choose the action based on either Thompson sampling or UCB.
In Algorithm 3, we describe how EDE updates the parameters of the feature extractor and ensemble heads. This is where our algorithm deviates significantly from prior approaches. For each ensemble head, we sample an independent minibatch from the replay buffer and compute the loss like deep ensemble. Inside the loop, the loss is computed as if each ensemble member is a QR-DQN. However, doing so naively means that the feature extractor  $f$  will be updated at a much faster rate. To ensure that all parameters are updated at the same rate, for each update, we randomly choose a single ensemble member that will be responsible for updating the feature extractor and the gradient of other ensemble heads is not propagated to the feature extractor. This also saves some compute. In practice, this turns out to be crucial for performing good uncertainty estimation and generalization. We show the effect of doing so in Figure 17c.
Attentive readers would notice that EDE is more expensive than other methods since it needs to process  $M$  minibatches for each update. In our experiments, the speed of EDE with 5 ensemble members is approximately 2.4 times slower than the base QR-DQN (see Appendix H). On the other hand, we empirically observed that the deep ensemble is crucial for performing accurate uncertainty estimation in CMDPs. UA-DQN [20], which uses MAP sampling, performs much worse than EDE even when both use the same number of ensemble members e.g., 3 (see Figure 5). In this work, we do not investigate other alternatives of uncertainty estimation but we believe future works in improving uncertainty estimation in deep learning could significantly reduce the computation overhead of EDE
Algorithm 1 EDE
1: Initialize feature extractor  $f$  and ensemble heads  $\{g_i\}_{i=1}^M$   
2: Initialize target feature extractor  $f^{target}$  and ensemble heads  $\{g_i^{target}\}_{i=1}^M$   
3: Initialize replay buffer Buffer  
4:  $\{s^{(k)}, \text{done}^{(k)}\}_{k=1}^K \leftarrow$  reset all environment  
5: for  $t$  from 1 to  $T$  do ▷ Standard DQN training loop  
6: for  $k$  from 1 to  $K$  do  
7:  $a \rightarrow \text{CHOOSEACTION}(s^{(k)}, f, \{g_i\}_{i=1}^M, k)$   
8:  $s', r, \text{done}^{(k)} \leftarrow$  Take  $a$  in the environment  
9: add  $\{s^{(k)}, a, r, s', \text{done}^{(k)}\}$  to Buffer  
10:  $s^{(k)} \leftarrow s'$   
11: if done(k) then  
12:  $s^{(k)}$ , done(k)  $\leftarrow$  reset the  $k^{\text{th}}$  environment  
13: end if  
14: end for  
15: UPDATE(Buffer, f, \{g_i\}_{i=1}^M, f^{target}, \{g_i^{target}\}_{i=1}^M)  
16: if  $t$  mod update_freqency = 0 then  
17:  $f^{target}, \{g_i^{target}\}_{i=1}^M \leftarrow f, \{g_i\}_{i=1}^M$   
18: end if  
19: end for
Algorithm 2 CHOOSEACTION  $\left(s,f,\{g_i\}_{i = 1}^M,k\right)$
1: feature  $\leftarrow f(s)$  Only one forward pass on the feature extractor  
2: for  $i$  from 1 to  $M$  do  
3:  $\{\theta_{ij}(s, a)\}_{j=1}^{N} \leftarrow g_i(\text{feature})$  Compute individual ensemble prediction  
4: end for  
5: Compute  $\hat{\sigma}_{\mathrm{epi}}^2(s, a)$  using Equation 2 Estimate epistemic uncertainty with ensemble  
6: Compute  $a^\star$  using Equation 5 or Equation 3 Explore with estimated uncertainty  
7: return  $a^\star$
Algorithm 3 UPDATE  $\left(\text{Buffer}, f, \{g_i\}_{i=1}^M, f^{\text{target}}, \{g_i^{\text{target}}\}_{i=1}^M\right)$
1:  $\mathcal{L}\gets 0$    
2: grad_index  $\sim \mathcal{U}\{1,M\}$  ▷Randomly select an ensemble member   
3: for  $i$  from 1 to  $M$  do Train each ensemble member independently   
4:  $\{s,a,r,s^{\prime},$  done}  $\leftarrow$  Sample from Buffer   
5: feature(s')  $\leftarrow$  ftarget(s')   
6:  $\{\theta_{ij}^{\mathrm{target}}(s',a)\}_{j = 1}^{N}\leftarrow g_{i}^{\mathrm{target}}(\mathrm{feature}(s'))$    
7: Compute  $\{\mathcal{T}\theta_{i1},\mathcal{T}\theta_{i2},\dots,\mathcal{T}\theta_{iN}\}$  using Equation 15 for the  $i^{\mathrm{th}}$  ensemble member   
8: feature(s)  $\leftarrow$  f(s)   
9: if  $i\neq$  grad_index then Prevent over-training the feature extractor   
10: feature(s)  $\leftarrow$  STOPGRADIENT (feature(s))   
11: end if   
12:  $\{\theta_{ij}(s,a)\}_{j = 1}^{N}\leftarrow g_{i}(\mathrm{feature}(s))$    
13:  $\mathcal{L}_i\gets$  compute QR loss using Equation 16 for the  $i^{\mathrm{th}}$  ensemble member   
14:  $\mathcal{L}\gets \mathcal{L} + \mathcal{L}_i$    
15: end for   
16: Compute gradient of L w.r.t the parameters of  $f,\{g_i\}_{i = 1}^M$    
17: Update the parameters with Adam
F Procgen Results
In Table 3 and Table 2, we show the mean and standard deviation of the final unnormalized train and test scores (at the end of 25M environment steps) of EDE long side with other methods in the literature. We adapt the tables from Raileanu and Fergus [90]. In Figure 13, we show the curves of min-max normalized test score that we reproduced using the code from Raileanu and Fergus [90]. We see that for a subset of games, EDE s significantly more sample-efficient. For other games, EDE s either on par with the policy optimization methods or fails to train, indicating there are other issues behind its poor performance on these games. Note that IDAAC results tune the hyperparameters for individual games whereas we only tune the hyperparameters on bigfish.
Game	PPO	MixReg	PLR	UCB-DRAC	PPG	DAAC	IDAAC	EDE
bigfish	3.7 ± 1.3	7.1 ± 1.6	10.9 ± 2.8	9.2 ± 2.0	11.2 ± 1.4	17.8 ± 1.4	18.5 ± 1.2	22.1 ± 2.0
StarPilot	24.9 ± 1.0	32.4 ± 1.5	27.9 ± 4.4	30.0 ± 1.3	47.2 ± 1.6	36.4 ± 2.8	37.0 ± 2.3	49.6 ± 2.4
FruitBot	26.2 ± 1.2	27.3 ± 0.8	28.0 ± 1.4	27.6 ± 0.4	27.8 ± 0.6	28.6 ± 0.6	27.9 ± 0.5	25.7 ± 1.4
BossFight	7.4 ± 0.4	8.2 ± 0.7	8.9 ± 0.4	7.8 ± 0.6	10.3 ± 0.2	9.6 ± 0.5	9.8 ± 0.6	10.0 ± 0.5
Ninja	5.9 ± 0.2	6.8 ± 0.5	7.2 ± 0.4	6.6 ± 0.4	6.6 ± 0.1	6.8 ± 0.4	6.8 ± 0.4	6.1 ± 0.6
Plunder	5.2 ± 0.6	5.9 ± 0.5	8.7 ± 2.2	8.3 ± 1.1	14.3 ± 2.0	20.7 ± 3.3	23.3 ± 1.4	4.9 ± 0.5
CaveFlyer	5.1 ± 0.4	6.1 ± 0.6	6.3 ± 0.5	5.0 ± 0.8	7.0 ± 0.4	4.6 ± 0.2	5.0 ± 0.6	7.9 ± 0.4
CoinRun	8.6 ± 0.2	8.6 ± 0.3	8.8 ± 0.5	8.6 ± 0.2	8.9 ± 0.1	9.2 ± 0.2	9.4 ± 0.1	6.7 ± 0.5
Jumper	5.9 ± 0.2	6.0 ± 0.3	5.8 ± 0.5	6.2 ± 0.3	5.9 ± 0.1	6.5 ± 0.4	6.3 ± 0.2	5.7 ± 0.3
Chaser	3.5 ± 0.9	5.8 ± 1.1	6.9 ± 1.2	6.3 ± 0.6	9.8 ± 0.5	6.6 ± 1.2	6.8 ± 1.0	1.6 ± 0.1
Climber	5.6 ± 0.5	6.9 ± 0.7	6.3 ± 0.8	6.3 ± 0.6	2.8 ± 0.4	7.8 ± 0.2	8.3 ± 0.4	5.7 ± 1.1
Dodgeball	1.6 ± 0.1	1.7 ± 0.4	1.8 ± 0.5	4.2 ± 0.9	2.3 ± 0.3	3.3 ± 0.5	3.2 ± 0.3	13.3 ± 0.1
Heist	2.5 ± 0.6	2.6 ± 0.4	2.9 ± 0.5	3.5 ± 0.4	2.8 ± 0.4	3.3 ± 0.2	3.5 ± 0.2	1.5 ± 0.3
Leaper	4.9 ± 2.2	5.3 ± 1.1	6.8 ± 1.2	4.8 ± 0.9	8.5 ± 1.0	7.3 ± 1.1	7.7 ± 1.0	6.4 ± 0.3
Maze	5.5 ± 0.3	5.2 ± 0.5	5.5 ± 0.8	6.3 ± 0.1	5.1 ± 0.3	5.5 ± 0.2	5.6 ± 0.3	3.4 ± 0.5
Miner	8.4 ± 0.7	9.4 ± 0.4	9.6 ± 0.6	9.2 ± 0.6	7.4 ± 0.2	8.6 ± 0.9	9.5 ± 0.4	0.7 ± 0.1
Table 2: Procgen scores on test levels after training on 25M environment steps. The mean and standard deviation are computed using 5 runs with different seeds.
Game	PPO	MixReg	PLR	UCB-DRAC	PPG	DAAC	IDAAC	EDE
bigfish	9.2 ± 2.7	15.0 ± 1.3	7.8 ± 1.0	12.8 ± 1.8	19.9 ± 1.7	20.1 ± 1.6	21.8 ± 1.8	27.5 ± 2.0
StarPilot	29.0 ± 1.1	28.7 ± 1.1	2.6 ± 0.3	33.1 ± 1.3	49.6 ± 2.1	38.0 ± 2.6	38.6 ± 2.2	46.9 ± 0.7
FruitBot	28.8 ± 0.6	29.9 ± 0.5	15.9 ± 1.3	29.3 ± 0.5	31.1 ± 0.5	29.7 ± 0.4	29.1 ± 0.7	27.7 ± 1.0
BossFight	8.0 ± 0.4	7.9 ± 0.8	8.7 ± 0.7	8.1 ± 0.4	11.1 ± 0.1	10.0 ± 0.4	10.4 ± 0.4	10.6 ± 0.4
Ninja	7.3 ± 0.3	8.2 ± 0.4	5.4 ± 0.5	8.0 ± 0.4	8.9 ± 0.2	8.8 ± 0.2	8.9 ± 0.3	9.1 ± 0.4
Plunder	6.1 ± 0.8	6.2 ± 0.3	4.1 ± 1.3	10.2 ± 1.76	16.4 ± 1.9	22.5 ± 2.8	24.6 ± 1.6	8.2 ± 0.8
CaveFlyer	6.7 ± 0.6	6.2 ± 0.7	6.4 ± 0.1	5.8 ± 0.9	9.5 ± 0.2	5.8 ± 0.4	6.2 ± 0.6	10.6 ± 0.1
CoinRun	9.4 ± 0.3	9.5 ± 0.2	5.4 ± 0.4	9.4 ± 0.2	9.9 ± 0.0	9.8 ± 0.0	9.8 ± 0.1	6.6 ± 0.4
Jumper	8.3 ± 0.2	8.5 ± 0.4	3.6 ± 0.5	8.2 ± 0.1	8.7 ± 0.1	8.6 ± 0.3	8.7 ± 0.2	9.0 ± 0.4
Chaser	4.1 ± 0.3	3.4 ± 0.9	6.3 ± 0.7	7.0 ± 0.6	10.7 ± 0.4	6.9 ± 1.2	7.5 ± 0.8	2.2 ± 0.1
Climber	6.9 ± 1.0	7.5 ± 0.8	6.2 ± 0.8	8.6 ± 0.6	10.2 ± 0.2	10.0 ± 0.3	10.2 ± 0.7	10.0 ± 0.3
Dodgeball	5.3 ± 2.3	9.1 ± 0.5	2.0 ± 1.1	7.3 ± 0.8	5.5 ± 0.5	5.2 ± 0.4	4.9 ± 0.3	15.9 ± 0.3
Heist	7.1 ± 0.5	4.4 ± 0.3	1.2 ± 0.4	6.2 ± 0.6	7.4 ± 0.4	5.2 ± 0.7	4.5 ± 0.3	7.2 ± 0.1
Leaper	5.5 ± 0.4	3.2 ± 1.2	6.4 ± 0.4	5.0 ± 0.9	9.3 ± 1.1	8.0 ± 1.1	8.3 ± 0.7	9.0 ± 0.3
Maze	9.1 ± 0.2	8.7 ± 0.7	4.1 ± 0.5	8.5 ± 0.3	9.0 ± 0.2	6.6 ± 0.4	6.4 ± 0.5	5.7 ± 0.9
Miner	11.7 ± 0.5	8.9 ± 0.9	9.7 ± 0.4	12.0 ± 0.3	11.3 ± 1.0	11.3 ± 0.9	11.5 ± 0.5	1.1 ± 0.2
Table 3: Procgen scores on train levels after training on 25M environment steps. The mean and standard deviation are computed using 5 runs with different seeds.
Figure 13: Min-max normalized test performance of a few representative methods on individual Progen games. Mean and standard deviation are computed over 5 seeds. We can see that for games on which value-based methods can make meaningful progress, EDE almost improve upon QR-DQN. In many cases, EDE is significantly more sample efficient compared to the other methods. There are notably 4 games (chaser, maze, heist, miner) where value-based methods (including EDE) still perform much worse than policy optimization. These games all require generalization with long-horizon planning which appears to remain challenging for value-based methods. Getting value-based methods to generalize on these games would be an interesting future direction.
F.1 Additional Ablation
Figure 14: The results on Procgen for DQN+UCB+TEE. The performance is much worse than the other ablations. This is not surprising. Since DQN+UCB performed unfavorably, it is unlikely that TEE would improve its performance.
(a) Meta Analysis
(b) Singleton Experiments
(c) DQN with Task ID
Figure 15: Investigations of other potential sources of poor generalization in RL.
G Additional Analysis
While our method improves train and test performance on two different CMDP benchmarks, it still does not fully close the generalization gap. In this section, we study other potential reasons for the poor generalization of RL algorithms, particularly value-based ones.
G.1 Alternative hypothesis
The second hypothesis for the superior performance of UCB relative to Greedy is that better exploration improves sample efficiency resulting in better training performance, which can, in turn, lead to better test performance. Indeed, even in the simple tabular MDPs from Figures 2 and 11, UCB converges faster on the training MDP. Here, we study this hypothesis in more detail. Note that this hypothesis and the one proposed in Section 3 can be simultaneously true, but the hypothesis here assumes the model readily learns generalizable representations, whereas the first hypothesis applies even if the agent has perfect state representation (e.g., tabular). In both cases, adapting uncertainty-driven exploration to deep RL on complex CMDPs has additional challenges such as representation learning and unobserved contexts.
Optimization vs. Generalization. In supervised learning, deep neural networks have good generalization, meaning that they perform similarly well on both train and test data. However, in RL, the main bottleneck remains optimization rather than generalization (which also boils down to sample complexity). For example in Procgen, the state-of-the-art training performance is far below the theoretical optimum [22, 90], suggesting that current RL algorithms do not yet "overfit" the training distribution. To verify this hypothesis, we first conduct a meta-analysis of existing algorithms. We choose 7 representative algorithms (Appendix F) and measure whether better performance on training MDPs translates into better performance on test MDPs. More specifically, we compute the Spearman's rank correlation,  $r_s$ , between the training and test performance of the algorithms, for each Procgen environment (Figure 15a). For many Procgen environments, training performance is strongly correlated with test performance, suggesting that further improvements in optimization could also result in better generalization.
In Figure 16a and Figure 16b, we show the scatter plot of the training performance and test performance of each algorithm on games with positive correlation and games with negative correlation (each color represents a different game). For games with positive correlation, we see that the correlation is almost linear with a slope of 1 (i.e., almost no generalization gap), which suggests that for these games, we can expect good generalization if the training performance is good. On the games with negative correlation, it is less clear that improving training performance will necessarily lead to better generalization or more severe overfitting. Most of the games with negative correlations are either saturated in terms of performance or too challenging for current methods to make much progress. This suggests that despite their apparent similarity, there are actually two "categories" of games based on the algorithm's tendency to overfit. Note that this dichotomy likely depends on the algorithm; here, all algorithms are PPO-based and we hypothesize the picture would be different for value-based methods. It would be interesting to understand what structural properties of these CMDPs are responsible for this difference in overfitting, but this is out of the scope of the current work.
Another interesting thing to notice is that most of the outliers in the two plots are PLR [48] which is an adversarial method that actively tries to reduce training performance. This may explain why the trend does not hold for PLR as all other methods use a uniform distribution over the training environments.
G.2 Other Sources of Poor Generalization
Sample efficiency on a single environment. One hypothesis for the poor performance of value-based methods relative to policy optimization ones in CMDPs is that the former is less sample efficient than the latter in each individual MDP  $\mu \in \mathcal{M}_{\mathrm{train}}$ . If that is the case, the suboptimality can accumulate across all the training MDPs, resulting in a large gap between these types of methods when trained in all the environments in  $\mathcal{M}_{\mathrm{train}}$ . To test this hypothesis, we run both DQN and PPO on a single environment from each of the Progen games. As shown in Figure 15b, we see that there is not a significant gap between the average performances of DQN and PPO on single environments (across 3 seeds for 10 environments on which both methods reach scores above 0.5). In fact, DQN is usually slightly more sample-efficient. Thus, we can rule this out as a major factor behind the observed discrepancy.
Partial observability. Another sensible hypothesis for why DQN underperforms PPO in CMDPs is the partial observability due to not knowing which environment the agent is interacting with at any given time [35]. Policy optimization methods that use trajectory-wise Monte Carlo update are less susceptible to partial observability whereas value-based methods that use temporal difference updates can suffer more since they rely heavily on the Markovian assumption. To test this hypothesis, we provide the task ID to the Q-network in addition to the observation similar to Liu et al. [65]. The access to the task ID means the agent knows exactly which environment it is in (even though the environment itself may still be partially observable like Atari). In Figure 15b, we see that both methods do well on the singleton environments, so with task IDs, the algorithms should in theory be able to do as well as singleton environments even if there is no transfer between different MDPs because the model can learn them separately. In practice, we embed the discrete task IDs into  $\mathbb{R}^{64}$  and add them as input to the final layers of the Q-network. Since there is no semantic relationship between discrete task IDs, we do not expect this to improve generalization performance, but, surprisingly, we found that it does not improve training performance either (Figure 15c). This suggests that partial observability may not be the main problem in such environments as far as training is concerned. Note that this issue is related to but not the same as the aforementioned value-overfitting issue. Having task ID means that the agent can have an accurate point estimate for each MDP (as far as representation is concerned), but optimization remains challenging without proper exploration.
Value overfitting. Most deep RL algorithms model value functions as point estimates of the expected return at each state. As discussed in [90], this is problematic in CMDPs because the same state can have different values in different environments. Hence, the only way for the values to be accurate is to rely on spurious features which are irrelevant for finding the optimal action. For policy optimization algorithms, the policy can be protected from this type of overfitting by using separate networks to train the policy and value, as proposed in [90]. However, this approach cannot be applied to value-based methods since the policy is directly defined by the Q-function. An alternative solution is distributional RL [10] which learns a distribution (rather than a point estimate) over all possible Q-values for a given state. This models the aleatoric uncertainty resulting from the unobserved contexts in CMDPs (i.e., not knowing which MDP the agent is interacting with), thus being able to account for one state having different potential values depending on the MDP. Our method uses distributional RL to mitigate the problem of value overfitting. As seen in Figure 5, while learning a value distribution leads to better results than learning only a point estimate (i.e., QR-DQN > DQN), the largest gains are due to using a more sophisticated exploration method (i.e., QR-DQN + UCB + TNN  $\mathcal{G}(\alpha)$  DQN).
Our analysis indicates that sample efficiency and partial observability are not the main reasons behind the poor generalization of value-based methods. In contrast, value overfitting is indeed a problem, but it can be alleviated with distributional RL. Nevertheless, effective exploration proves to be a key factor in improving the train and test performance of value-based approaches in CMDPs. Our results also suggest that better optimization on the training environment can lead to better generalization to new environments, making it a promising research direction.
(a) Games with positive train-test correlation.
(b) Games with negative train-test correlation.
Figure 16: Additional empirical results.
(c) Comparing different uncertainties (mean test min-max normalized scores).
H Uncertainty