Representation Convergence: Mutual Distillation is Secretly a Form of Regularization

## Abstract

Abstract In this paper, we argue that mutual distillation between reinforcement learning policies serves as an implicit regularization , preventing them from overfitting to irrelevant features. We highlight two separate contributions: (i) Theoretically, for the first time, we prove that enhancing the policy robustness to irrelevant features leads to improved generalization performance. (ii) Empirically, we demonstrate that mutual distillation between policies contributes to such robustness, enabling the spontaneous emergence of invariant representations over pixel inputs. Ultimately, we do not claim to achieve state-of-the-art performance but rather focus on uncovering the underlying principles of generalization and deepening our understanding of its mechanisms. Our website: https://dml-rl.github.io/ .

## 1 Introduction

Humans exhibit a remarkable ability to learn robustly and generalize across diverse environments. Once a skill is acquired, it often transfers seamlessly to new contexts that share the same underlying semantics, even when their visual appearance differs substantially. For example, consider a person who becomes proficient at a video game, even if the background graphics or character textures are altered, the player retains their ability to perform well, effortlessly adapting to the new setting. This suggests that human learning is not overly dependent on low-level visual details, but rather grounded in abstract representations that capture the essential structure of a task. Neuroscientific studies support this view, linking abstract reasoning to the human prefrontal cortex (Bengtsson et al., 2009; Dumontheil, 2014), and highlighting the role of inhibitory neurons in enhancing cognitive processing efficiency (Pi et al., 2013).

In stark contrast, visual reinforcement learning (VRL) agents often struggle with generalization. While they can be trained to solve complex tasks in specific environments, even minor changes, such as shifts in color schemes or background textures, can significantly degrade their performance. This sensitivity indicates that VRL agents tend to overfit to superficial visual features, failing to capture the underlying structure of the task (Cobbe et al., 2019; 2020). These limitations give rise to a fundamental question:

What hinders reinforcement learning agents from generalizing like humans? How can we enable them to learn robust representations that drive human-like generalization behavior?

The core reason behind the limited generalization ability of VRL agents lies in their reliance on convolutional neural networks (CNNs) as visual encoders. While CNNs are the de facto choice for processing high-dimensional visual inputs, they are notoriously sensitive to even small perturbations (Goodfellow et al., 2014). This brittleness significantly hampers the robustness of learned policies and limits their ability to generalize across visually diverse environments. To address this issue, one common strategy is to apply data augmentation (Shorten & Khoshgoftaar, 2019), which improves robustness by diversifying the training distribution and reducing dataset-induced biases. Alternatively, invariant representation learning has emerged as a principled approach to tackle generalization problem from a feature-learning perspective. It aims to extract representations that remain stable under a wide range of input transformations, thereby promoting robustness and transferability (Nguyen et al., 2021).

While data augmentation is an effective bias mitigation technique, its reliance on task-specific strategies that are manually crafted by human experts, poses a challenge for designing task-independent solutions. In contrast, our method enables agents to generalize without any handcrafted augmentations or external priors, relying purely on training experience. Invariant representation learning is a promising approach to enhance model’s cross-domain generalization. However, it relies on transformation correspondences, which are fundamentally inaccessible in the generalization scenarios of reinforcement learning due to the dynamic nature of environments. In addition, invariant representation framework inherently separates the encoder from the model, unnecessarily complicating the theoretical analysis. Instead, our framework is theoretically and empirically end-to-end.

In this paper, we first propose a novel theoretical framework to analyze the generalization problem in reinforcement learning and show that the policy robustness to irrelevant features enhances its generalization performance. Building upon this principled insight, we then provide empirical evidence that deep mutual learning (DML) (Zhang et al., 2018b) can implicitly prevent them from overfitting to such irrelevant features, leading to robust learning process and significant generalization improvements.

In summary, the main contributions of this paper are as follows:

- •
We theoretically prove that improving the policy robustness to irrelevant features enhances its generalization performance. To the best of our knowledge, we are the first to provide a rigorous proof of this intuition.
- •
We propose a hypothesis that deep mutual learning (DML) enhances the generalization performance of the policy by implicitly regularizing irrelevant features. We also provide intuitive insights to support this hypothesis.
- •
Strong empirical results support our theory and hypothesis, showing that DML technique leads to consistent improvements in generalization performance.

## 2 Preliminaries

In this section, we introduce reinforcement learning under the generalization setting in Section [2.1](https://arxiv.org/html/2501.02481v5#S2.SS1), as well as the DML technique in Section [2.2](https://arxiv.org/html/2501.02481v5#S2.SS2).

### 2.1 Markov Decision Process and Generalization

Markov decision process (MDP) is a mathematical framework for sequential decision-making, which is defined by a tuple $\mathcal{M}=\left(\mathcal{S},\mathcal{A},r,\mathcal{P},\rho,\gamma\right)$, where $\mathcal{S}$ and $\mathcal{A}$ represent the state space and action space, $r:\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$ is the reward function, $\mathcal{P}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\mapsto[0,1]$ is the dynamics, $\rho:\mathcal{S}\mapsto[0,1]$ is the initial state distribution, and $\gamma\in(0,1)$ is the discount factor.

Define a policy $\mu:\mathcal{S}\times\mathcal{A}\mapsto[0,1]$, the action-value function and value function are defined as

$$ $Q^{\mu}(s_{t},a_{t})=\mathbb{E}_{\mu}\left[\sum_{k=0}^{\infty}\gamma^{k}r(s_{t+k},a_{t+k})\right],\kern 5.0ptV^{\mu}(s_{t})=\mathbb{E}_{a_{t}\sim\mu(\cdot|s_{t})}\left[Q^{\mu}(s_{t},a_{t})\right].$ (1) $$

Given $Q^{\mu}$ and $V^{\mu}$, the advantage function can be expressed as $A^{\mu}(s_{t},a_{t})=Q^{\mu}(s_{t},a_{t})-V^{\mu}(s_{t})$.

In our generalization setting, we introduce a rendering function $f:\mathcal{S}\mapsto\mathcal{O}_{f}\subset\mathcal{O}$ to obfuscate the agent’s actual observations, which is a bijection(^1^11We define $\mathcal{O}_{f}:=\left\{f(s)|s\in\mathcal{S}\right\}$, which means for any $s_{1}\neq s_{2}$, we have $f(s_{1})\neq f(s_{2})$.) from $\mathcal{S}$ to $\mathcal{O}_{f}$. We now define the MDP induced by the underlying MDP $\mathcal{M}$ and the rendering function $f$, denote it as $\mathcal{M}_{f}=\left(\mathcal{O}_{f},\mathcal{A},r_{f},\mathcal{P}_{f},\rho_{f},\gamma\right)$, where $\mathcal{O}_{f}$ represents the observation space, $r_{f}:\mathcal{O}_{f}\times\mathcal{A}\mapsto\mathbb{R}$ is the reward function, $\mathcal{P}_{f}:\mathcal{O}_{f}\times\mathcal{A}\times\mathcal{O}_{f}\mapsto[0,1]$ is the dynamics, and $\rho_{f}:\mathcal{O}_{f}\mapsto[0,1]$ is the initial observation distribution. We present the following assumptions:

###### Assumption 2.1 .

Assume that $f$ can be sampled from a distribution $p:\mathcal{F}\mapsto[0,1]$, where $f\in\mathcal{F}$.

###### Assumption 2.2 .

Given any $f\in\mathcal{F}$, $o_{0}^{f},o_{t}^{f},o_{t+1}^{f}\in\mathcal{O}_{f}$ and $a_{t}\in\mathcal{A}$, assume that $r_{f}(o_{t}^{f},a_{t})=r(f^{-1}(o_{t}^{f}),a_{t}),\mathcal{P}_{f}(o_{t+1}^{f}|o_{t}^{f},a_{t})=\mathcal{P}(f^{-1}(o_{t+1}^{f})|f^{-1}(o_{t}^{f}),a_{t}),\rho_{f}(o_{0}^{f})=\rho(f^{-1}(o_{0}^{f}))$.

Explanation. Assumption [2.2](https://arxiv.org/html/2501.02481v5#S2.Thmtheorem2) states that all $\mathcal{M}_{f}$ share a common underlying MDP $\mathcal{M}$, in which the agent’s observations are perturbed by different rendering functions while all other components remain unchanged, much like different painters depicting the same scene in their own styles.

Next, consider an agent interacting with $\mathcal{M}_{f}$ following the policy $\pi:\mathcal{O}\times\mathcal{A}\mapsto[0,1]$ to obtain a trajectory

$$ $\tau_{f}=(o_{0}^{f},a_{0},r_{0}^{f},o_{1}^{f},a_{1},r_{1}^{f},\dots,o_{t}^{f},a_{t},r_{t}^{f},\dots),$ (2) $$

where $o_{0}^{f}\sim\rho_{f}(\cdot)$, $a_{t}\sim\pi(\cdot|o_{t}^{f})$, $r_{t}^{f}=r_{f}(o_{t}^{f},a_{t})$ and $o_{t+1}^{f}\sim\mathcal{P}_{f}(\cdot|o_{t}^{f},a_{t})$, we simplify the notation to $\tau_{f}\sim\pi$. During training, the agent is only allowed to access a subset of all MDPs, which is $\left\{\mathcal{M}_{f}|f\in\mathcal{F}_{\mathrm{train}}\subset\mathcal{F}\right\}$, and then tests its generalization performance across all MDPs. Thus, denote $p_{\mathrm{train}}:\mathcal{F}_{\mathrm{train}}\mapsto[0,1]$ as the distribution over $\mathcal{F}_{\mathrm{train}}$, the agent’s training performance $\eta(\pi)$ and generalization performance $\zeta(\pi)$ can be expressed as

$$ $\eta(\pi)=\mathbb{E}_{f\sim p_{\mathrm{train}}(\cdot),\tau_{f}\sim\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{f}(o_{t}^{f},a_{t})\right],\kern 5.0pt\zeta(\pi)=\mathbb{E}_{f\sim p(\cdot),\tau_{f}\sim\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{f}(o_{t}^{f},a_{t})\right].$ (3) $$

The goal of the agent is to learn a policy $\pi$ that maximizes the generalization performance $\zeta(\pi)$.

###### Assumption 2.1 .

###### Assumption 2.2 .

### 2.2 Deep Mutual Learning

Deep mutual learning (DML)
(Zhang et al., 2018b) is a mutual distillation technique in supervised learning. Unlike the traditional teacher-student distillation strategy, DML aligns the probability distributions of multiple student networks by minimizing the KL divergence loss during training, allowing them to learn from each other. Specifically,

$$ $\mathcal{L}_{\mathrm{DML}}=\mathcal{L}_{\mathrm{SL}}+\alpha\mathcal{L}_{\mathrm{KL}},$ (4) $$

where $\mathcal{L}_{\mathrm{SL}}$ and $\mathcal{L}_{\mathrm{KL}}$ represent the supervised learning loss and the KL divergence loss, respectively, $\alpha$ is the weight. Using DML, the student cohort effectively pools their collective estimate of the next most likely classes. Finding out and matching the other most likely classes for each training instance according to their peers increases each student’s posterior entropy, which helps them converge to a more robust representation, leading to better generalization.

## 3 Theoretical Results

In this section, we present the main results of this paper, demonstrating that enhancing the agent’s robustness to irrelevant features will improve its generalization performance.

A key issue is that we do not exactly know the probability distribution $p_{\mathrm{train}}$. Note that $\mathcal{F}_{\mathrm{train}}$ is a subset of $\mathcal{F}$, we naturally assume that the probability distribution $p_{\mathrm{train}}$ can be derived from the normalized probability distribution $p$.

###### Assumption 3.1 .

For any $f\in\mathcal{F}$, assume that

$$ $p_{\mathrm{train}}(f)=\frac{p(f)\cdot\mathbb{I}(f\in\mathcal{F}_{\mathrm{train}})}{Z},\kern 5.0ptp_{\mathrm{eval}}(f)=\frac{p(f)\cdot\mathbb{I}(f\in\mathcal{F}_{\mathrm{eval}})}{1-Z},$ (5) $$

where $Z=\int_{\mathcal{F}_{\mathrm{train}}}p(f)\mathrm{d}f$ and $1-Z$ is the normalization term, $\mathcal{F}_{\mathrm{eval}}=\mathcal{F}-\mathcal{F}_{\mathrm{train}}$, $\mathbb{I}(\cdot)$ denotes the indicator function.

An interesting fact is that, for a specific policy $\pi$, if we only consider its interaction with $\mathcal{M}_{f}$, we can establish a bijection between this policy and a certain underlying policy that directly interacts with $\mathcal{M}$. We now denote it as $\mu_{f}(\cdot|s_{t})=\pi(\cdot|f(s_{t}))$. By further defining the normalized discounted visitation distribution $d^{\mu}(s)=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}(s_{t}=s|\mu)$, we can use this underlying policy $\mu_{f}$ to replace the training and generalization performance of the policy $\pi$. Specifically, we have the following connection:

###### Lemma 3.2 .

For any given policy $\pi$, define its underlying policy as $\mu_{f}(\cdot|s_{t})=\pi(\cdot|f(s_{t}))$, then

$$ $\eta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right],\kern 5.0pt\zeta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right].$ (6) $$

###### Proof.

See Appendix [E.1](https://arxiv.org/html/2501.02481v5#A5.SS1).
∎

We can thus analyze the generalization problem using the underlying policy $\mu_{f}$. Then, define $L_{\pi}$ as the first-order approximation of $\eta$ (Schulman et al., 2015), we can derive the following lower bounds:

###### Theorem 3.3 (Training performance lower bound) .

Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\eta(\tilde{\pi})\geq L_{\pi}(\tilde{\pi})-\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],$ (7) $$

where $\epsilon_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{\max_{s}\left|\mathbb{E}_{a\sim\tilde{\mu}_{f}(\cdot|s)}\left[A^{\mu_{f}}(s,a)\right]\right|\right\}$.

###### Proof.

See Appendix [E.3](https://arxiv.org/html/2501.02481v5#A5.SS3).
∎

###### Theorem 3.4 (Generalization performance lower bound) .

Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\begin{split}\zeta(\tilde{\pi})&\geq L_{\pi}(\tilde{\pi})-\frac{2r_{\max}(1-Z)}{1-\gamma}-\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\\ &-\frac{2\delta_{\mathrm{train}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],\\ \end{split}$ (8) $$

where $r_{\max}=\max_{s,a}\left|r(s,a)\right|$, $\delta_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{\max_{s,a}\left|A^{\mu_{f}}(s,a)\right|\right\}$, and $\delta_{\mathrm{eval}}=\max_{f\in\mathcal{F}_{\mathrm{eval}}}\left\{\max_{s,a}\left|A^{\mu_{f}}(s,a)\right|\right\}$.

###### Proof.

See Appendix [E.2](https://arxiv.org/html/2501.02481v5#A5.SS2).
∎

Explanation.
Building upon Theorems [3.3](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem3) and [3.4](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem4), we observe that, in contrast to the lower bound on training performance, the lower bound on generalization performance incorporates three additional terms, scaled by the common coefficient $(1-Z)$. This implies that increasing $Z$ contributes to improved generalization performance, with the special case of $Z=1$ resulting in alignment between generalization and training performance. Notably, this theoretical insight was also validated in Figure 2 of Cobbe et al. (2020).

However, once the training level is fixed (i.e., $\mathcal{F}_{\mathrm{train}}$), $Z$ is a constant, improving generalization performance requires constraining the following three terms:

$$ $\underbrace{\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]}_{\text{denote it as }\mathfrak{D}_{1}},\kern 5.0pt\underbrace{\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]}_{\text{denote it as }\mathfrak{D}_{2}},\kern 5.0pt\underbrace{\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]}_{\text{denote it as }\mathfrak{D}_{\mathrm{train}}}.$ (9) $$

During the training process, we can only empirically bound $\mathfrak{D}_{\mathrm{train}}$. Next, we establish the upper bounds of $\mathfrak{D}_{1}$ and $\mathfrak{D}_{2}$. Specifically, we propose the following theorem:

###### Theorem 3.5 .

Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\mathfrak{D}_{1}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)\mathfrak{D}_{\mathrm{train}},\kern 5.0pt\mathfrak{D}_{2}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\right)\underbrace{\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]}_{\text{denote it as }\mathfrak{D}_{\mathrm{eval}}},$ (10) $$

where $\sigma_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{D_{\mathrm{TV}}^{\max}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}$ and $\sigma_{\mathrm{eval}}=\max_{f\in\mathcal{F}_{\mathrm{eval}}}\left\{D_{\mathrm{TV}}^{\max}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}$, $D_{\mathrm{TV}}^{\max}(\tilde{\mu}_{f}\|\mu_{f})[s]$ is defined as $\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]$.

###### Proof.

See Appendix [E.4](https://arxiv.org/html/2501.02481v5#A5.SS4).
∎

The only problem now is finding the relationship between $\mathfrak{D}_{\mathrm{eval}}$ and $\mathfrak{D}_{\mathrm{train}}$. To achieve this, we would like to first introduce the following definition, which represents the policy robustness to irrelevant features.

###### Definition 3.6 ( $\mathcal{R}$ -robust) .

We say that the policy $\pi$ is $\mathcal{R}$-robust if it satisfies

$$ $\sup_{s\in\mathcal{S},\tilde{f},f\in\mathcal{F}}D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]=\mathcal{R}.$ (11) $$

Explanation. This definition demonstrates how the policy $\pi$ is influenced by two different rendering functions, $\tilde{f}$ and $f$, for any given underlying state $s$. If $\mathcal{R}=0$, it indicates that $D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]\equiv 0$, which means that the policy is no longer affected by any irrelevant features.

Our intention in this definition is not to derive the tightest possible bound but rather to demonstrate how policy robustness to irrelevant features can contribute to improved generalization. Subsequently, leveraging Definition [3.6](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem6), we establish an upper bound for $\mathfrak{D}_{\mathrm{eval}}$.

###### Theorem 3.7 .

Given any two policies, $\tilde{\pi}$ and $\pi$, assume that $\tilde{\pi}$ is $\mathcal{R}_{\tilde{\pi}}$-robust, and $\pi$ is $\mathcal{R}_{\pi}$-robust, then the following bound holds:

$$ $\mathfrak{D}_{\mathrm{eval}}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)\mathcal{R}_{\pi}+\mathcal{R}_{\tilde{\pi}}+\mathfrak{D}_{\mathrm{train}}.$ (12) $$

###### Proof.

See Appendix [E.5](https://arxiv.org/html/2501.02481v5#A5.SS5).
∎

Altogether, by combining Theorems [3.4](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem4), [3.5](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem5), and [3.7](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem7), we can derive the following corollary:

###### Corollary 3.8 .

Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\zeta(\tilde{\pi})\geq L_{\pi}(\tilde{\pi})-C_{\mathrm{train}}\mathfrak{D}_{\mathrm{train}}-C_{\pi}\mathcal{R}_{\pi}-C_{\tilde{\pi}}\mathcal{R}_{\tilde{\pi}}-C,$ (13) $$

where

$$ $\begin{split}C_{\mathrm{train}}&=\frac{2\delta_{\mathrm{train}}(1-Z)}{1-\gamma}\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)+\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\left(1+\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\right)+\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}},\\ C_{\pi}&=\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\left(1+\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\right)\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right),\\ C_{\tilde{\pi}}&=\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\left(1+\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\right),\kern 5.0ptC=\frac{2r_{\max}(1-Z)}{1-\gamma}.\\ \end{split}$ (14) $$

Figure: Figure 1: This is a toy environment where the gray agent’s goal is to pick up coins.
Refer to caption: x1.png

Explanation.
This represents our central theoretical result, demonstrating that enhancing generalization performance requires not only minimizing $\mathfrak{D}_{\mathrm{train}}$ during training but also improving policy robustness to irrelevant features, specifically by reducing $\mathcal{R}_{\pi}$ and $\mathcal{R}_{\tilde{\pi}}$. Furthermore, we emphasize that these results rely solely on the mild Assumptions [2.1](https://arxiv.org/html/2501.02481v5#S2.Thmtheorem1), [2.2](https://arxiv.org/html/2501.02481v5#S2.Thmtheorem2), and [3.1](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem1). Consequently, this constitutes a novel contribution that is broadly applicable to a wide range of algorithms.

###### Assumption 3.1 .

###### Lemma 3.2 .

###### Proof.

###### Theorem 3.3 (Training performance lower bound) .

###### Proof.

###### Theorem 3.4 (Generalization performance lower bound) .

###### Proof.

###### Theorem 3.5 .

###### Proof.

###### Definition 3.6 ( ℛ \mathcal{R} -robust) .

###### Theorem 3.7 .

###### Proof.

###### Corollary 3.8 .

## 4 Distillation as Regularization

Despite the theoretical advancements, in typical generalization settings, both the underlying MDP and the rendering function remain unknown. Next, we begin by introducing a minimal toy example in Section [4.1](https://arxiv.org/html/2501.02481v5#S4.SS1), which we then provide an in-depth analyze in Section [4.2](https://arxiv.org/html/2501.02481v5#S4.SS2) to motivate our hypothesis.

### 4.1 Toy Example

Let’s consider a simple environment where the agent attempts to pick up coins to earn rewards (see Figure [1](https://arxiv.org/html/2501.02481v5#S3.F1)). The agent’s observations are the current pixels. It is clear that the agent’s true objective is to pick up the coins, and the background color is a spurious feature. However, upon observing the training data for policy A, we can see that in the red background, the coins are always on the right side of the agent, while in the cyan background, the coins are always on the left side. As a result, when training policy A using reinforcement learning algorithms, it is likely to exhibit overfitting behavior, such as moving to the right in a red background and to the left in a cyan background.

However, the overfitting of policy A to the background color will fail in the training data of policy B, because in policy B’s training data, regardless of whether the background color is red or cyan, the coin can appear either on the left or right side of the agent. Therefore, through DML, policy A is regularized by the behavior of policy B, effectively preventing policy A from overfitting to the background color. In other words, any irrelevant features learned by policy A could lead to suboptimal performance of policy B, and vice versa. Thus, we hypothesize that this process will force both policies to learn the true underlying semantics, ultimately improving generalization performance.

Figure: Figure 2: (Left) Independently trained reinforcement learning policies may overfit to irrelevant features. (Right) Through mutual distillation via DML, two policies regularize each other to converge toward a more robust hypothesis space, ultimately improving generalization performance.
Refer to caption: x2.png

### 4.2 Hypothesis

Motivated by Section [4.1](https://arxiv.org/html/2501.02481v5#S4.SS1), DML can be viewed as a form of implicit regularization against irrelevant features, as demonstrated in Figure [2](https://arxiv.org/html/2501.02481v5#S4.F2), which illustrates two randomly initialized policies independently trained using reinforcement learning algorithms. In this case, since the training samples only include a portion of all possible MDPs, the policies are likely to overfit to irrelevant features and fail to converge to a robust hypothesis space.

Applying DML to the training process of both policies facilitates mutual learning, which can mitigate overfitting to irrelevant features. Due to the randomness of parameter initialization and the interaction process, they generate different training samples, DML encourages both policies to make consistent decisions based on the same observations. As discussed in Section [4.1](https://arxiv.org/html/2501.02481v5#S4.SS1), any irrelevant features learned by policy A are likely to degrade the performance of policy B, and vice versa. As training progresses, DML will drive both policies to learn more meaningful and useful representations, gradually reducing the divergence between them. Ideally, we hypothesize that both policies will capture the essential aspects of high-dimensional observations as time grows.

## 5 Experiments

This section presents our main empirical results. Section [5.1](https://arxiv.org/html/2501.02481v5#S5.SS1) introduces the implementation details, Section [5.2](https://arxiv.org/html/2501.02481v5#S5.SS2) validates the effectiveness of DML technique for improving generalization performance, Section [5.3](https://arxiv.org/html/2501.02481v5#S5.SS3) verifies our central hypothesis, and Section [5.4](https://arxiv.org/html/2501.02481v5#S5.SS4) confirms our theoretical results.

### 5.1 Implementation Details

We use Procgen (Cobbe et al., 2019; 2020) as the experimental benchmark for testing generalization performance. Procgen is a suite of 16 procedurally generated game-like environments designed to benchmark both sample efficiency and generalization in reinforcement learning, and it has been widely used to test the generalization performance of various reinforcement learning algorithms (Wang et al., 2020; Raileanu & Fergus, 2021; Raileanu et al., 2021; Lyle et al., 2022; Rahman & Xue, 2023; Jesson & Jiang, 2024).

We employ the Proximal Policy Optimization (PPO) (Schulman et al., 2017; Cobbe et al., 2020) as our baseline. Specifically, given a parameterized policy $\pi_{\theta}$ ($\theta$ represents the parameters), the objective of $\pi_{\theta}$ is to maximize

$$ $J(\theta)=\mathop{\mathbb{E}}_{(o_{t},a_{t})\sim\pi_{\theta_{\mathrm{old}}}}\left\{\min\left[r_{t}(\theta)\cdot\hat{A}(o_{t},a_{t}),\mathrm{clip}\left(r_{t}(\theta),1-\epsilon,1+\epsilon\right)\cdot\hat{A}(o_{t},a_{t})\right]\right\},$ (15) $$

where $\hat{A}$ is the advantage estimate, and $r_{t}(\theta)=\pi_{\theta}(a_{t}|o_{t})/\pi_{\theta_{\mathrm{old}}}(a_{t}|o_{t})$ is the probability ratio, where $\pi_{\theta_{\mathrm{old}}}$ and $\pi_{\theta}$ denote the old and current policies, respectively.

We randomly initialize two agents to interact with the environment and collect data separately. Similar to the DML loss ([4](https://arxiv.org/html/2501.02481v5#S2.E4)) used in supervised learning, we also introduce an additional KL divergence loss term, which leads to

$$ $\mathcal{L}_{\mathrm{DML}}=\mathcal{L}_{\mathrm{RL}}+\alpha\mathcal{L}_{\mathrm{KL}},$ (16) $$

where $\mathcal{L}_{\mathrm{RL}}$ is the reinforcement learning loss and $\mathcal{L}_{\mathrm{KL}}$ is the KL divergence loss, $\alpha$ is the weight. And then we optimize the total loss of both agents, which is the average of their DML losses, as shown in Algorithm [1](https://arxiv.org/html/2501.02481v5#alg1), which we name Mutual Distillation Policy Optimization (MDPO).

Figure: Algorithm 1 Mutual Distillation Policy Optimization (MDPO)

Ultimately, we do not claim to achieve state-of-the-art (SOTA) performance, but rather provide empirical evidence for the non-trivial insight that DML serves as an implicit regularization against irrelevant features, leading to consistent improvements in generalization performance. We also acknowledge the methodological similarities with prior work such as Zhao & Hospedales (2021); despite that, we introduce representation convergence (Section [4.2](https://arxiv.org/html/2501.02481v5#S4.SS2)), a novel insight with further supported by strong theoretical analysis (Section [3](https://arxiv.org/html/2501.02481v5#S3)), constituting our additional contributions.

### 5.2 Empirical Results

We compare the generalization performance of our MDPO against the PPO baseline on the Procgen benchmark, under the hard-level settings (Cobbe et al., 2020), the results are illustrated in Figure [3](https://arxiv.org/html/2501.02481v5#S5.F3). It can be observed that DML technique indeed leads to consistent improvements in generalization performance across all environments. Notably, for the bigfish, dodgeball, and fruitbot environments, we have observed significant improvements. Moreover, the full experimental results for all environments, including training and generalization performance, are provided in Appendix [D](https://arxiv.org/html/2501.02481v5#A4).

A natural concern arises: how can we determine whether DML improves generalization performance by enhancing the policy robustness against irrelevant features, or simply due to the additional information sharing between these two agents during training (each agent receives additional information than it would from training alone)? To answer this question, we conducted robustness testing in Section [5.3](https://arxiv.org/html/2501.02481v5#S5.SS3) and added an ablation study in Section [5.4](https://arxiv.org/html/2501.02481v5#S5.SS4) to support our theory and hypothesis.

Figure: Figure 3: Generalization performance from 500 levels in Procgen benchmark with different methods. Our MDPO gains significant performance improvement compared with the baseline algorithms.
Refer to caption: x3.png

Figure: Figure 4: To test the robustness of the trained policy, we obfuscate the agent’s observations using convolutional layers randomly initialized with a standard Gaussian distribution.
Refer to caption: x4.png

Figure: Table 1: A simple practical measure of $\mathcal{R}$-robustness defined in Definition [3.6](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem6).

### 5.3 Robustness Testing

We design a novel approach to test policy robustness against irrelevant features. For a given frame, we generate adversarial samples using random CNNs initialized with a standard Gaussian distribution, as shown in Figure [4](https://arxiv.org/html/2501.02481v5#S5.F4). Notably, the feature extraction of MDPO encoder is highly stable and focused (red points), whereas the features extracted by the original PPO encoder are significantly dispersed (blue points).

Moreover, we design a practical measure of $\mathcal{R}$-robustness defined in Definition [3.6](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem6). Specifically, for each environment, we run the trained policy (PPO and MDPO) in the environment for 100 steps and obtain observations $o_{1},o_{2},\dots,o_{100}$. Then, for each $o_{i}$ we use 100 random CNNs to simulate rendering function samples $f_{1}^{(i)},f_{2}^{(i)},\dots,f_{100}^{(i)}$ and compute the TV divergence of the policy between the adversarial samples and the original observations, i.e., $D_{\mathrm{TV}}(\pi_{\theta}(\cdot|o_{i})\|\pi_{\theta}(\cdot|f_{j}^{(i)}(o_{i})))$, where $i,j=1,2,\dots,100$. We then take the maximum of these values as a simple practical measure of $\mathcal{R}$-robustness:

$$ $\hat{\mathcal{R}}:=\max_{i,j}D_{\mathrm{TV}}(\pi_{\theta}(\cdot|o_{i})\|\pi_{\theta}(\cdot|f_{j}^{(i)}(o_{i}))),$ (17) $$

the results are shown in Table [1](https://arxiv.org/html/2501.02481v5#S5.T1). We can see that MDPO achieves a significantly lower $\hat{\mathcal{R}}$ than PPO, showing that DML effectively improves the policy robustness to irrelevant features, which serves as further strong evidence for our hypothesis.

### 5.4 Ablation Study

Figure: Table 2: Generalization performance of PPO linear probe on top of the frozen encoders.

We design additional ablation experiments. Specifically, we double the model size, batch size, and total number of interactions for the PPO baseline, as shown in Figure [5](https://arxiv.org/html/2501.02481v5#S5.F5). It can be seen that PPO baseline still fails to match the performance of MDPO, demonstrating that naively scaling up the PPO baseline does not lead to stable improvements in generalization performance.

Furthermore, we retrain a PPO linear probe on top of the frozen encoders of the trained PPO and MDPO policies, training for only 1M steps (2% of the original training steps), the final generalization performance during the last 10% steps is shown in Table [2](https://arxiv.org/html/2501.02481v5#S5.T2). It can be seen that the PPO linear probe trained on the MDPO encoder achieves significantly better generalization performance, indicating that DML helps the policy learn better (more robust) representations. Moreover, we add a sensitivity analysis of the KL divergence weight $\alpha$, and the results are presented in Table [3](https://arxiv.org/html/2501.02481v5#S5.T3).

Figure: Figure 5: Generalization performance of PPO baseline with double model size, batch size, and total number of interactions, compared to original PPO and MDPO (for training results, see Figure [9](https://arxiv.org/html/2501.02481v5#A4.F9)).
Refer to caption: x5.png

## 6 Conclusion

In this paper, we provide a novel theoretical framework to explain the generalization problem in deep reinforcement learning. We further hypothesize that DML, as a form of implicit regularization, effectively prevents the policy from overfitting to irrelevant features. Strong empirical results support our central theory and hypothesis, demonstrating that our approach can improve the generalization performance of reinforcement learning systems by enhancing robustness against irrelevant features. Our work provides valuable insights and elegant solutions into the development of more adaptable and robust policies capable of generalizing across diverse environments.

## References

- Achiam et al. (2017)
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel.
Constrained policy optimization.
In *International conference on machine learning*, pp. 22–31.
PMLR, 2017.
- Bengtsson et al. (2009)
Sara L Bengtsson, John-Dylan Haynes, Katsuyuki Sakai, Mark J Buckley, and
Richard E Passingham.
The representation of abstract task rules in the human prefrontal
cortex.
*Cerebral Cortex*, 19(8):1929–1936, 2009.
- Chalmers (2022)
David J Chalmers.
*Reality+: Virtual worlds and the problems of philosophy*.
Penguin UK, 2022.
- Cobbe et al. (2019)
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman.
Quantifying generalization in reinforcement learning.
In *International conference on machine learning*, pp. 1282–1289. PMLR, 2019.
- Cobbe et al. (2020)
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman.
Leveraging procedural generation to benchmark reinforcement learning.
In *International conference on machine learning*, pp. 2048–2056. PMLR, 2020.
- Cobbe et al. (2021)
Karl W Cobbe, Jacob Hilton, Oleg Klimov, and John Schulman.
Phasic policy gradient.
In *International Conference on Machine Learning*, pp. 2020–2027. PMLR, 2021.
- Cohen (2006)
S. Marc Cohen.
The allegory of the cave, 2006.
URL [https://faculty.washington.edu/smcohen/320/cave.htm](https://faculty.washington.edu/smcohen/320/cave.htm).
- Deckers et al. (2024)
Lucas Deckers, Benjamin Vandersmissen, Ing Jyh Tsang, Werner Van Leekwijck, and
Steven Latré.
Twin network augmentation: A novel training strategy for improved
spiking neural networks and efficient weight quantization.
*arXiv preprint arXiv:2409.15849*, 2024.
- Dumontheil (2014)
Iroise Dumontheil.
Development of abstract thinking during childhood and adolescence:
The role of rostrolateral prefrontal cortex.
*Developmental cognitive neuroscience*, 10:57–76,
2014.
- Falandays & Smaldino (2022)
J Benjamin Falandays and Paul E Smaldino.
The emergence of cultural attractors: How dynamic populations of
learners achieve collective cognitive alignment.
*Cognitive science*, 46(8):e13183, 2022.
- Farebrother et al. (2018)
Jesse Farebrother, Marlos C Machado, and Michael Bowling.
Generalization and regularization in dqn.
*arXiv preprint arXiv:1810.00123*, 2018.
- Gan et al. (2024)
Yaozhong Gan, Renye Yan, Xiaoyang Tan, Zhe Wu, and Junliang Xing.
Transductive off-policy proximal policy optimization.
*arXiv preprint arXiv:2406.03894*, 2024.
- Goodfellow et al. (2014)
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.
Explaining and harnessing adversarial examples.
*arXiv preprint arXiv:1412.6572*, 2014.
- Grigsby & Qi (2020)
Jake Grigsby and Yanjun Qi.
Measuring visual generalization in continuous control from pixels.
*arXiv preprint arXiv:2010.06740*, 2020.
- Huang et al. (2022)
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam
Chakraborty, Kinal Mehta, and JoÃĢo GM AraÃšjo.
Cleanrl: High-quality single-file implementations of deep
reinforcement learning algorithms.
*Journal of Machine Learning Research*, 23(274):1–18, 2022.
- Huh et al. (2024)
Minyoung Huh, Brian Cheung, Tongzhou Wang, and Phillip Isola.
The platonic representation hypothesis.
*arXiv preprint arXiv:2405.07987*, 2024.
- Igl et al. (2019)
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang,
Sam Devlin, and Katja Hofmann.
Generalization in reinforcement learning with selective noise
injection and information bottleneck.
*Advances in neural information processing systems*, 32, 2019.
- Jesson & Jiang (2024)
Andrew Jesson and Yiding Jiang.
Improving generalization on the procgen benchmark with simple
architectural changes and scale.
*arXiv preprint arXiv:2410.10905*, 2024.
- Justesen et al. (2018)
Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa,
Julian Togelius, and Sebastian Risi.
Illuminating generalization in deep reinforcement learning through
procedural level generation.
*arXiv preprint arXiv:1806.10729*, 2018.
- Kakade & Langford (2002)
Sham Kakade and John Langford.
Approximately optimal approximate reinforcement learning.
In *Proceedings of the nineteenth international conference on
machine learning*, pp. 267–274, 2002.
- Kelly & Glymour (1989)
Kevin T Kelly and Clark Glymour.
Convergence to the truth and nothing but the truth.
*Philosophy of Science*, 56(2):185–220,
1989.
- Laskin et al. (2020)
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind
Srinivas.
Reinforcement learning with augmented data.
*Advances in neural information processing systems*,
33:19884–19895, 2020.
- Laudan (1981)
Larry Laudan.
A confutation of convergent realism.
*Philosophy of science*, 48(1):19–49, 1981.
- Lee et al. (2019)
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee.
Network randomization: A simple technique for generalization in deep
reinforcement learning.
*arXiv preprint arXiv:1910.05396*, 2019.
- Lyle et al. (2022)
Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal.
Learning dynamics and generalization in deep reinforcement learning.
In *International Conference on Machine Learning*, pp. 14560–14581. PMLR, 2022.
- Ma et al. (2022)
Guozheng Ma, Zhen Wang, Zhecheng Yuan, Xueqian Wang, Bo Yuan, and Dacheng Tao.
A comprehensive survey of data augmentation in visual reinforcement
learning.
*arXiv preprint arXiv:2210.04561*, 2022.
- Mehta et al. (2020)
Bhairav Mehta, Manfred Diaz, Florian Golemo, Christopher J Pal, and Liam Paull.
Active domain randomization.
In *Conference on Robot Learning*, pp. 1162–1176. PMLR, 2020.
- Moon et al. (2022)
Seungyong Moon, JunYeong Lee, and Hyun Oh Song.
Rethinking value function learning for generalization in
reinforcement learning.
*Advances in Neural Information Processing Systems*,
35:34846–34858, 2022.
- Nguyen et al. (2021)
A Tuan Nguyen, Toan Tran, Yarin Gal, and Atilim Gunes Baydin.
Domain invariant representation learning with domain density
transformations.
*Advances in Neural Information Processing Systems*,
34:5264–5275, 2021.
- Packer et al. (2018)
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen
Koltun, and Dawn Song.
Assessing generalization in deep reinforcement learning.
*arXiv preprint arXiv:1810.12282*, 2018.
- Pi et al. (2013)
Hyun-Jae Pi, Balázs Hangya, Duda Kvitsiani, Joshua I Sanders, Z Josh Huang,
and Adam Kepecs.
Cortical interneurons that specialize in disinhibitory control.
*Nature*, 503(7477):521–524, 2013.
- Rahman & Xue (2023)
Md Masudur Rahman and Yexiang Xue.
Adversarial style transfer for robust policy optimization in deep
reinforcement learning.
*arXiv preprint arXiv:2308.15550*, 2023.
- Raileanu & Fergus (2021)
Roberta Raileanu and Rob Fergus.
Decoupling value and policy for generalization in reinforcement
learning.
In *International Conference on Machine Learning*, pp. 8787–8798. PMLR, 2021.
- Raileanu et al. (2021)
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob
Fergus.
Automatic data augmentation for generalization in reinforcement
learning.
*Advances in Neural Information Processing Systems*,
34:5402–5415, 2021.
- Rajeswaran et al. (2017)
Aravind Rajeswaran, Kendall Lowrey, Emanuel V Todorov, and Sham M Kakade.
Towards generalization and simplicity in continuous control.
*Advances in neural information processing systems*, 30, 2017.
- Schulman et al. (2015)
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp
Moritz.
Trust region policy optimization.
In *International conference on machine learning*, pp. 1889–1897. PMLR, 2015.
- Schulman et al. (2017)
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
Proximal policy optimization algorithms.
*arXiv preprint arXiv:1707.06347*, 2017.
- Scott (2019)
Dan Scott.
The bull by pablo picasso – a lesson in abstraction, 2019.
URL [https://drawpaintacademy.com/the-bull/](https://drawpaintacademy.com/the-bull/).
- Shorten & Khoshgoftaar (2019)
Connor Shorten and Taghi M Khoshgoftaar.
A survey on image data augmentation for deep learning.
*Journal of big data*, 6(1):1–48, 2019.
- Slaoui et al. (2019)
Reda Bahi Slaoui, William R Clements, Jakob N Foerster, and Sébastien Toth.
Robust visual domain randomization for reinforcement learning.
*arXiv preprint arXiv:1910.10537*, 2019.
- Song et al. (2019)
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur.
Observational overfitting in reinforcement learning.
*arXiv preprint arXiv:1912.02975*, 2019.
- Srivastava et al. (2014)
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan
Salakhutdinov.
Dropout: a simple way to prevent neural networks from overfitting.
*The journal of machine learning research*, 15(1):1929–1958, 2014.
- Tobin et al. (2017)
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and
Pieter Abbeel.
Domain randomization for transferring deep neural networks from
simulation to the real world.
In *2017 IEEE/RSJ international conference on intelligent robots
and systems (IROS)*, pp. 23–30. IEEE, 2017.
- (44)
Benjamin Vandersmissen, Lucas Deckers, and Jose Oramas.
Improving neural network accuracy by concurrently training with a
twin network.
In *The Thirteenth International Conference on Learning
Representations*.
- Wang et al. (2020)
Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng.
Improving generalization in reinforcement learning with mixture
regularization.
*Advances in Neural Information Processing Systems*,
33:7968–7978, 2020.
- Xie et al. (2024)
Zhengpeng Xie, Qiang Zhang, Fan Yang, Marco Hutter, and Renjing Xu.
Simple policy optimization.
*arXiv preprint arXiv:2401.16025*, 2024.
- Xu et al. (2024)
Xiaohan Xu, Ming Li, Chongyang Tao, Tao Shen, Reynold Cheng, Jinyang Li, Can
Xu, Dacheng Tao, and Tianyi Zhou.
A survey on knowledge distillation of large language models.
*arXiv preprint arXiv:2402.13116*, 2024.
- Yarats et al. (2021)
Denis Yarats, Ilya Kostrikov, and Rob Fergus.
Image augmentation is all you need: Regularizing deep reinforcement
learning from pixels.
In *International conference on learning representations*, 2021.
- Yuan et al. (2023)
Zhecheng Yuan, Sizhe Yang, Pu Hua, Can Chang, Kaizhe Hu, and Huazhe Xu.
Rl-vigen: A reinforcement learning benchmark for visual
generalization.
*Advances in Neural Information Processing Systems*,
36:6720–6747, 2023.
- Yue et al. (2019)
Xiangyu Yue, Yang Zhang, Sicheng Zhao, Alberto Sangiovanni-Vincentelli, Kurt
Keutzer, and Boqing Gong.
Domain randomization and pyramid consistency: Simulation-to-real
generalization without accessing target domain data.
In *Proceedings of the IEEE/CVF international conference on
computer vision*, pp. 2100–2110, 2019.
- Zhang et al. (2018a)
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio.
A study on overfitting in deep reinforcement learning.
*arXiv preprint arXiv:1804.06893*, 2018a.
- Zhang & Guo (2021)
Hanping Zhang and Yuhong Guo.
Generalization of reinforcement learning with policy-aware
adversarial data augmentation.
*arXiv preprint arXiv:2106.15587*, 2021.
- Zhang et al. (2018b)
Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu.
Deep mutual learning.
In *Proceedings of the IEEE conference on computer vision and
pattern recognition*, pp. 4320–4328, 2018b.
- Zhao & Hospedales (2021)
Chenyang Zhao and Timothy Hospedales.
Robust domain randomised reinforcement learning through peer-to-peer
distillation.
In *Asian Conference on Machine Learning*, pp. 1237–1252.
PMLR, 2021.
- Zhuang et al. (2023)
Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo.
Behavior proximal policy optimization.
*arXiv preprint arXiv:2302.11312*, 2023.

## Appendix A Related Work

The generalization of deep reinforcement learning has been widely studied, and previous work has pointed out the overfitting problem in deep reinforcement learning (Rajeswaran et al., 2017; Zhang et al., 2018a; Justesen et al., 2018; Packer et al., 2018; Song et al., 2019; Cobbe et al., 2019; Grigsby & Qi, 2020; Cobbe et al., 2020; Yuan et al., 2023). A natural approach to avoid the overfitting problem in deep reinforcement learning is to apply regularization techniques originally developed for supervised learning such as dropout (Srivastava et al., 2014; Farebrother et al., 2018; Igl et al., 2019), data augmentation (Laskin et al., 2020; Yarats et al., 2021; Zhang & Guo, 2021; Raileanu et al., 2021; Ma et al., 2022), domain randomization (Tobin et al., 2017; Yue et al., 2019; Slaoui et al., 2019; Mehta et al., 2020), or network randomization technique (Lee et al., 2019). On the other hand, in order to improve sample efficiency, previous studies encouraged the policy network and value network to share parameters (Schulman et al., 2017; Huang et al., 2022). However, recent works have explored the idea of decoupling the two and proposed additional distillation strategies (Cobbe et al., 2021; Raileanu & Fergus, 2021; Moon et al., 2022). In particular, Raileanu & Fergus (2021) demonstrated that more information is needed to accurately estimate the value function, which can lead to overfitting.

Knowledge distillation is a learning paradigm that aims to align the student network with the teacher network to achieve knowledge transfer. A commonly used practice is to distill the knowledge learned by a large model into a smaller model to reduce inference costs after deployment (Xu et al., 2024). On the other hand, distillation technique can also be used to distill a model with privileged information into a model with access to only partial information to improve its generalization ability. However, research has shown that knowledge distillation can also be applied to multiple student networks during training to encourage them to learn from each other, called deep mutual learning (DML) (Zhang et al., 2018b). Building upon this observation, Zhao & Hospedales (2021) further demonstrate that DML can improve the generalization performance of reinforcement learning agents, yet no in-depth analysis of why this happens. In addition, recent studies suggest that aligning the student networks at the output layer may be suboptimal, and recommend alignment at the logits layer instead (Deckers et al., 2024; Vandersmissen et al.,).

## Appendix B Hyperparameters

Table [4](https://arxiv.org/html/2501.02481v5#A2.T4) shows the detailed hyperparameter settings in our code, with the main hyperparameters consistent with the hard-level settings in Cobbe et al. (2020), except that we train for 50M steps instead of 200M. We train the policy on the initial 500 levels and then test its generalization performance across the full distribution of levels.

**Table 4: Detailed hyperparameters in Procgen.**
| Hyperparameter\Algorithm | PPO (Schulman et al., 2017) | MDPO (ours) |
| --- | --- | --- |
| Number of workers | 64 | 64 |
| Horizon | 256 | 256 |
| Learning rate | 0.0005 | 0.0005 |
| Learning rate decay | No | No |
| Optimizer | Adam | Adam |
| Total interaction steps | 50M | 50M |
| Update epochs | 3 | 3 |
| Mini-batches | 8 | 8 |
| Batch size | 16384 | 16384 |
| Mini-batch size | 2048 | 2048 |
| Discount factor $\gamma$ | 0.999 | 0.999 |
| GAE parameter $\lambda$ | 0.95 | 0.95 |
| Value loss coefficient $c_{1}$ | 0.5 | 0.5 |
| Entropy loss coefficient $c_{2}$ | 0.01 | 0.01 |
| Clipping parameter $\epsilon$ | 0.2 | 0.2 |
| KL divergence weight $\alpha$ | - | 1.0 |

## Appendix C Philosophical Insight

An intriguing analogy for our hypothesis is the process of truth emergence. Typically, each scholar offers their unique perspective, but for it to be widely accepted, it must garner consensus from peers within the field, or even from the broader academic community. We can draw a parallel between DML and the peer review process: when a particular viewpoint is accepted by the majority, it is more likely to reflect an objective truth. Going deeper, our hypothesis aligns with the philosophical concept of convergent realism (Laudan, 1981; Kelly & Glymour, 1989; Huh et al., 2024), which posits that science progresses towards an objective truth.

In the Allegory of the Cave (Cohen, 2006), Plato describes a group of people who have been chained in a cave from birth, facing the wall. They can only see the shadows cast on the wall by objects behind them, illuminated by a fire. These shadows are all they know and perceive as reality. One day, a prisoner is freed and exits the cave. Upon seeing the world outside, he discovers that the true reality is entirely different from the shadows, and he realizes that what they thought was “real” was only an illusion.

The Allegory of the Cave explores the concept of different levels of knowledge. From a metaphysical standpoint, Plato is distinguishing between the “appearance” and the “reality”. The shadows in the cave represent the sensory world, which is just an illusion, while the outside world symbolizes the realm of true reality, which can only be apprehended through rational thought and philosophical inquiry. Plato argues that what we perceive with our senses is not the true essence of reality, what he called the ideal reality, but rather a mere shadow of the higher, eternal truths.

However, even though the freed prisoner has seen the world outside the cave, how can he be certain that what he now perceives is not a shadow of something even more fundamental? This is exactly skepticism, the philosophical stance that questions the possibility of certain knowledge. It suggests that every layer of perceived reality might itself be an illusion, prompting thinkers from Descartes onward to doubt not only the senses but even the existence of the external world, demanding ever more rigorous foundations for truth. To counter skepticism, some philosophers have found it necessary to expand the definition of “reality” and propose the virtual realism (Chalmers, 2022).

Figure: Figure 6: Picasso’s The Bull (Scott, 2019). By focusing on and exaggerating specific details, rather than trying to capture every detail realistically, artists can convey the core meaning or essence of the subject, more powerfully.
Refer to caption: figures/abstraction.jpg

This paper offers another powerful perspective against skepticism: if we regard the underlying state $s_{t}$ as the essence of things, and the rendering function $f$ as the projection (possibly a composition of multiple layers of projections), then our perception of the world becomes an observation $o_{t}=f(s_{t})$. By encouraging agents to make consistent decisions based on the same observations, our method fosters a process akin to cognitive alignment (Falandays & Smaldino, 2022), which has been fundamental in human societal development. For instance, in voting, the majority rule is employed because decisions supported by the majority are perceived as more reliable. Similarly, our method facilitates cognitive alignment between agents, enabling them to converge on objective truths despite noisy or irrelevant features in their observations. Over time, the cognitive alignment between agents encourages the convergence of their individual representations toward a more accurate understanding of the environment.

It is worth noting that metaphysical ideas can also be observed in certain works of art, such as Picasso’s The Bull (Figure [6](https://arxiv.org/html/2501.02481v5#A3.F6)). Most people think that art is all about seeing more detail, but it is really about seeing less. Seeing basic patterns amongst the “noise”; seeing basic forms amongst the complex; seeing the few important details which convey the majority of meaning.

## Appendix D More Results

### D.1 Full Results

Figure: Figure 7: Generalization performance of PPO and MDPO from 500 levels in each environment.
Refer to caption: x6.png

Figure: Figure 8: Training performance of PPO and MDPO from 500 levels in each environment.
Refer to caption: x7.png

### D.2 Other Baseline

We select Simple Policy Optimization (SPO) (Xie et al., 2024) as another baseline and compare the generalization performance of SPO with DML-based SPO, the results are shown in Table [5](https://arxiv.org/html/2501.02481v5#A4.T5).

**Table 5: generalization performance of SPO and SPO with DML.**
| Algo\Env | bigfish | dodgeball | fruitbot | starpilot |
| --- | --- | --- | --- | --- |
| SPO | 1.04${}^{\pm\text{0.96}}$ | 1.59${}^{\pm\text{0.87}}$ | 2.28${}^{\pm\text{1.64}}$ | 6.06${}^{\pm\text{1.54}}$ |
| SPO with DML | 4.62${}^{\pm\text{2.60}}$ | 4.66${}^{\pm\text{1.44}}$ | 1.13${}^{\pm\text{1.15}}$ | 8.65${}^{\pm\text{2.47}}$ |

We can see that DML can also significantly improve the generalization performance of SPO, which fully demonstrates that DML-based policy optimization is a general learning framework for enhancing the generalization ability of diverse reinforcement learning algorithms.

### D.3 More Ablation Results

Here, we additionally present the training curves from the Ablation Study (Section [5.4](https://arxiv.org/html/2501.02481v5#S5.SS4)), as shown in Figure [9](https://arxiv.org/html/2501.02481v5#A4.F9).

Figure: Figure 9: Training performance of PPO baseline with double model size, batch size, and total number of interactions, compared to original PPO and MDPO.
Refer to caption: x8.png

Interestingly, although the scaled-up PPO nearly matches MDPO in training performance during the final stage of training in the fruitbot environment, there remains a substantial gap in their generalization performance (as shown in Figure [5](https://arxiv.org/html/2501.02481v5#S5.F5)). This provides further strong evidence that DML effectively enhances the policy robustness to irrelevant features, as MDPO achieves significantly better generalization performance despite comparable training performance.

### D.4 Additional Visualizations

We also generate adversarial samples by adjusting the brightness, contrast, saturation, and hue of the images, and test the robustness of the PPO encoder and our MDPO encoder, as shown in Figure [10](https://arxiv.org/html/2501.02481v5#A4.F10).

Figure: Figure 10: The robustness of PPO and MDPO to brightness, contrast, saturation, and hue.
Refer to caption: x9.png

We can see that the MDPO policy has also learned robustness representations to these irrelevant factors, while the PPO policy remains sensitive to them. Additionally, we present adversarial samples generated by random CNNs, as shown in Figure [11](https://arxiv.org/html/2501.02481v5#A4.F11), as well as those generated by randomly adjusting brightness, contrast, saturation, and hue, as can be seen from Figure [12](https://arxiv.org/html/2501.02481v5#A4.F12).

Figure: Figure 11: Adversarial samples generated by random CNNs.
Refer to caption: figures/89.png

Figure: Figure 12: Adversarial samples generated by different brightness, contrast, saturation, and hue.
Refer to caption: figures/1.png

## Appendix E Proofs

Let’s start with some useful lemmas.

###### Lemma E.1 (Performance difference) .

Let $\mu_{f}(\cdot|s_{t})=\pi(\cdot|f(s_{t}))$ and $\tilde{\mu}_{f}(\cdot|s_{t})=\tilde{\pi}(\cdot|f(s_{t}))$, define training and generalization performance as

$$ $\eta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right],\kern 5.0pt\zeta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right].$ (18) $$

Then the differences in training and generalization performance can be expressed as

$$ $\eta(\tilde{\pi})-\eta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right],\kern 5.0pt\zeta(\tilde{\pi})-\zeta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right].$ (19) $$

###### Proof.

This result can be directly derived from Kakade & Langford (2002).
∎

###### Lemma E.2 .

The divergence between two normalized discounted visitation distribution, $\|d^{\tilde{\mu}}-d^{\mu}\|_{1}$, is bounded by an average divergence of $\tilde{\mu}$ and $\mu$:

$$ $\|d^{\tilde{\mu}}-d^{\mu}\|_{1}\leq\frac{\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu}(\cdot)}\left[\|\tilde{\mu}-\mu\|_{1}\right]=\frac{2\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}\|\mu)[s]\right],$ (20) $$

where $D_{\mathrm{TV}}(\tilde{\mu}\|\mu)[s]=\frac{1}{2}\sum_{a\in\mathcal{A}}|\tilde{\mu}(a|s)-\mu(a|s)|$ represents the Total Variation (TV) distance.

###### Proof.

See Achiam et al. (2017).
∎

###### Lemma E.3 .

Given any state $s\in\mathcal{S}$, any two policies $\tilde{\mu}$ and $\mu$, the average advantage, $\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]$, is bounded by

$$ $\left|\mathbb{E}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right|\leq 2D_{\mathrm{TV}}(\tilde{\mu}\|\mu)[s]\cdot\max_{a}\left|A^{\mu}(s,a)\right|.$ (21) $$

###### Proof.

Note that

$$ $\begin{split}\mathbb{E}_{a\sim\mu(\cdot|s)}\left[A^{\mu}(s,a)\right]=&\mathbb{E}_{a\sim\mu(\cdot|s)}\left[Q^{\mu}(s,a)-V^{\mu}(s)\right]\\ =&\mathbb{E}_{a\sim\mu(\cdot|s)}\left[Q^{\mu}(s,a)\right]-V^{\mu}(s)\\ =&V^{\mu}(s)-V^{\mu}(s)\\ =&0,\\ \end{split}$ (22) $$

thus,

$$ $\begin{split}\left|\mathbb{E}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right|&=\left|\mathbb{E}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]-\mathbb{E}_{a\sim\mu(\cdot|s)}\left[A^{\mu}(s,a)\right]\right|\\ &\leq\left\|\tilde{\mu}-\mu\right\|_{1}\cdot\left\|A^{\mu}(s,a)\right\|_{\infty}\\ &=2D_{\mathrm{TV}}(\tilde{\mu}\|\mu)[s]\cdot\max_{a}\left|A^{\mu}(s,a)\right|.\\ \end{split}$ (23) $$

This is a widely used trick (Schulman et al., 2015; Zhuang et al., 2023; Gan et al., 2024).
∎

In addition, using the above lemmas, the following corollary can be obtained, which will be repeatedly used in our proof.

###### Corollary E.4 .

Given any two policies, $\tilde{\mu}$ and $\mu$, the following bound holds:

$$ $\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\tilde{\mu}}(\cdot)\\ a\sim\tilde{\mu}(\cdot|s)\end{subarray}}\left[A^{\mu}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu}(\cdot)\\ a\sim\tilde{\mu}(\cdot|s)\end{subarray}}\left[A^{\mu}(s,a)\right]\right|\leq\frac{2\epsilon\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}\|\mu)[s]\right],$ (24) $$

where $\epsilon=\max_{s}\left|\mathbb{E}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right|$.

###### Proof.

We rewrite the expectation as

$$ $\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\tilde{\mu}}(\cdot)\\ a\sim\tilde{\mu}(\cdot|s)\end{subarray}}\left[A^{\mu}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu}(\cdot)\\ a\sim\tilde{\mu}(\cdot|s)\end{subarray}}\left[A^{\mu}(s,a)\right]\right|=\left|\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}}(\cdot)}\left\{\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right\}-\mathop{\mathbb{E}}_{s\sim d^{\mu}(\cdot)}\left\{\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right\}\right|,$ (25) $$

where the expectation $\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]$ is a function of $s$, then

$$ $\left|\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}}(\cdot)}\left\{\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right\}-\mathop{\mathbb{E}}_{s\sim d^{\mu}(\cdot)}\left\{\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right\}\right|\leq\left\|d^{\tilde{\mu}}-d^{\mu}\right\|_{1}\cdot\left\|\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right\|_{\infty}.$ (26) $$

Next, according to Lemma [E.2](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem2), we have

$$ $\left\|d^{\tilde{\mu}}-d^{\mu}\right\|_{1}\cdot\left\|\mathop{\mathbb{E}}_{a\sim\tilde{\mu}(\cdot|s)}\left[A^{\mu}(s,a)\right]\right\|_{\infty}=\left\|d^{\tilde{\mu}}-d^{\mu}\right\|_{1}\cdot\epsilon\leq\frac{2\epsilon\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}\|\mu)[s]\right],$ (27) $$

concluding the proof.
∎

###### Lemma E.1 (Performance difference) .

###### Proof.

###### Lemma E.2 .

###### Proof.

###### Lemma E.3 .

###### Proof.

###### Corollary E.4 .

###### Proof.

### E.1 Proof of Lemma 3.2

Lemma [3.2](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem2).
For any given policy $\pi$, define its underlying policy as $\mu_{f}(\cdot|s_{t})=\pi(\cdot|f(s_{t}))$, then

$$ $\eta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right],\kern 5.0pt\zeta(\pi)=\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right].$ (28) $$

###### Proof.

According to the definition of training and generalization performance in ([3](https://arxiv.org/html/2501.02481v5#S2.E3)), we have

$$ $\eta(\pi)=\mathbb{E}_{f\sim p_{\mathrm{train}}(\cdot),\tau_{f}\sim\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{f}(o_{t}^{f},a_{t})\right],\kern 5.0pt\zeta(\pi)=\mathbb{E}_{f\sim p(\cdot),\tau_{f}\sim\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{f}(o_{t}^{f},a_{t})\right].$ (29) $$

To prove Lemma [3.2](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem2), we only need to show that for any given $f\in\mathcal{F}$, the following equation holds:

$$ $\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]=\mathbb{E}_{\tau_{f}\sim\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{f}(o_{t}^{f},a_{t})\right].$ (30) $$

According to the definition of the normalized discounted visitation distribution $d^{\mu}(s)=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}(s_{t}=s|\mu)$, we have

$$ $\begin{split}\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]&=\frac{1}{1-\gamma}\sum_{s\in\mathcal{S}}(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}(s_{t}=s|\mu_{f})\sum_{a\in\mathcal{A}}\mu_{f}(a|s)\cdot r(s,a)\\ &=\sum_{t=0}^{\infty}\sum_{s\in\mathcal{S}}\mathbb{P}(s_{t}=s|\mu_{f})\sum_{a\in\mathcal{A}}\mu_{f}(a|s)\cdot\gamma^{t}r(s,a)\\ \end{split}$ (31) $$

Next, according to Assumption [2.2](https://arxiv.org/html/2501.02481v5#S2.Thmtheorem2), we have

$$ $\begin{split}\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]&=\sum_{t=0}^{\infty}\sum_{s\in\mathcal{S}}\mathbb{P}(s_{t}=s|\mu_{f})\sum_{a\in\mathcal{A}}\mu_{f}(a|s)\cdot\gamma^{t}r(s,a)\\ &=\sum_{t=0}^{\infty}\sum_{s\in\mathcal{S}}\mathbb{P}(f(s_{t})=f(s)|\mu_{f})\sum_{a\in\mathcal{A}}\pi(a|f(s))\cdot\gamma^{t}r_{f}(f(s),a)\\ &\overset{f(s)=o^{f},f(s_{t})=o_{t}^{f}}{=}\sum_{t=0}^{\infty}\sum_{o^{f}\in\mathcal{O}_{f}}\mathbb{P}(o_{t}^{f}=o^{f}|\pi)\sum_{a\in\mathcal{A}}\pi(a|o^{f})\cdot\gamma^{t}r_{f}(o^{f},a)\\ &=\mathbb{E}_{\tau_{f}\sim\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{f}(o_{t}^{f},a_{t})\right],\\ \end{split}$ (32) $$

concluding the proof.
∎

###### Proof.

### E.2 Proof of Theorem 3.4

Theorem [3.4](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem4).
Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\begin{split}\zeta(\tilde{\pi})&\geq L_{\pi}(\tilde{\pi})-\frac{2r_{\max}(1-Z)}{1-\gamma}-\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\\ &-\frac{2\delta_{\mathrm{train}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right].\\ \end{split}$ (33) $$

###### Proof.

Let’s start with the first-order approximation of the training performance (Schulman et al., 2015), denote it as

$$ $L_{\pi}(\tilde{\pi})=\eta(\pi)+\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right].$ (34) $$

Then, we are trying to bound the difference between $\zeta(\tilde{\pi})$ and $L_{\pi}(\tilde{\pi})$, according to Lemma [E.1](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem1), that is,

$$ $\begin{split}&\left|\zeta(\tilde{\pi})-L_{\pi}(\tilde{\pi})\right|\\ =&\left|\zeta(\pi)-\eta(\pi)+\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\frac{1}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|\\ =&\frac{1}{1-\gamma}\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|\\ \leq&\frac{1}{1-\gamma}\left\{\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]\right|+\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|\right\}.\\ \end{split}$ (35) $$

We can bound these two terms separately. Simplifying the notation, denote $g(f)=\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot),a\sim\mu_{f}(\cdot|s)}\left[r(s,a)\right]$, we can thus rewrite the first term as

$$ $\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]\right|=\left|\mathop{\mathbb{E}}_{f\sim p(\cdot)}\left[g(f)\right]-\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left[g(f)\right]\right|,$ (36) $$

then

$$ $\left|\mathop{\mathbb{E}}_{f\sim p(\cdot)}\left[g(f)\right]-\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left[g(f)\right]\right|=\left|\int_{\mathcal{F}}p(f)\cdot g(f)\mathrm{d}f-\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\cdot g(f)\mathrm{d}f\right|.$ (37) $$

Next, according to Assumption [3.1](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem1),

$$ $\begin{split}&\left|\int_{\mathcal{F}}p(f)\cdot g(f)\mathrm{d}f-\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\cdot g(f)\mathrm{d}f\right|=\left|\int_{\mathcal{F}}p(f)\cdot g(f)\mathrm{d}f-\int_{\mathcal{F}_{\mathrm{train}}}\frac{p(f)}{Z}\cdot g(f)\mathrm{d}f\right|\\ =&\left|\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f-\int_{\mathcal{F}_{\mathrm{train}}}\frac{p(f)}{Z}\cdot g(f)\mathrm{d}f+\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|\\ =&\left|\int_{\mathcal{F}_{\mathrm{train}}}\frac{Z-1}{Z}p(f)\cdot g(f)\mathrm{d}f+\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|,\\ \end{split}$ (38) $$

where $Z=\int_{\mathcal{F}_{\mathrm{train}}}p(f)\mathrm{d}f\leq 1$, thus,

$$ $\begin{split}&\left|\int_{\mathcal{F}_{\mathrm{train}}}\frac{Z-1}{Z}p(f)\cdot g(f)\mathrm{d}f+\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|\\ \leq&\left|\int_{\mathcal{F}_{\mathrm{train}}}\frac{Z-1}{Z}p(f)\cdot g(f)\mathrm{d}f\right|+\left|\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|\\ \leq&\frac{1-Z}{Z}\left|\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|+\left|\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|.\\ \end{split}$ (39) $$

Meanwhile,

$$ $\begin{split}\left|g(f)\right|=\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]\right|=&\left|\sum_{s\in\mathcal{S}}(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}(s_{t}=s|\mu_{f})\sum_{a\in\mathcal{A}}\mu_{f}(a|s)\cdot r(s,a)\right|\\ \leq&(1-\gamma)\sum_{t=0}^{\infty}\sum_{s\in\mathcal{S}}\mathbb{P}(s_{t}=s|\mu_{f})\sum_{a\in\mathcal{A}}\mu_{f}(a|s)\cdot\gamma^{t}\left|r(s,a)\right|\\ \leq&(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}r_{\max}=r_{\max},\\ \end{split}$ (40) $$

where $r_{\max}=\max_{s,a}\left|r(s,a)\right|$, then we can bound the first term as

$$ $\begin{split}\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\mu_{f}(\cdot|s)\end{subarray}}\left[r(s,a)\right]\right|\leq&\frac{1-Z}{Z}\left|\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|+\left|\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot g(f)\mathrm{d}f\right|\\ \leq&\frac{1-Z}{Z}\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot\left|g(f)\right|\mathrm{d}f+\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot\left|g(f)\right|\mathrm{d}f\\ \leq&\frac{(1-Z)r_{\max}}{Z}\int_{\mathcal{F}_{\mathrm{train}}}p(f)\mathrm{d}f+r_{\max}\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\mathrm{d}f\\ =&\frac{(1-Z)r_{\max}}{Z}\cdot Z+r_{\max}\cdot(1-Z)=2r_{\max}(1-Z).\\ \end{split}$ (41) $$

Now we are trying to bound the second term, which can be expressed as

$$ $\begin{split}&\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|\\ =&\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|\\ \leq&\underbrace{\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|}_{\text{denote as }\Phi}+\underbrace{\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|}_{\text{denote as }\Psi}.\\ \end{split}$ (42) $$

Using Corollary [E.4](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem4), $\Psi$ can be bounded by

$$ $\begin{split}\Psi&=\left|\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right\}\right|\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|\right\}\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\frac{2\epsilon\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right\},\\ \end{split}$ (43) $$

where $\epsilon=\max_{s}\left|\mathbb{E}_{a\sim\tilde{\mu}_{f}(\cdot|s)}\left[A^{\mu_{f}}(s,a)\right]\right|$, denote $\epsilon_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{\epsilon\right\}$, we obtain

$$ $\Psi\leq\frac{2\gamma\epsilon_{\mathrm{train}}}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right].$ (44) $$

Next, with a little abuse of notation $g(f)$, denote

$$ $g(f)=\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right],$ (45) $$

we can rewrite $\Phi$ as

$$ $\Phi=\left|\mathop{\mathbb{E}}_{f\sim p(\cdot)}\left[g(f)\right]-\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left[g(f)\right]\right|,$ (46) $$

then, similar to ([37](https://arxiv.org/html/2501.02481v5#A5.E37)), ([38](https://arxiv.org/html/2501.02481v5#A5.E38)), ([39](https://arxiv.org/html/2501.02481v5#A5.E39)) and ([41](https://arxiv.org/html/2501.02481v5#A5.E41)),

$$ $\Phi\leq\frac{1-Z}{Z}\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot\left|g(f)\right|\mathrm{d}f+\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot\left|g(f)\right|\mathrm{d}f.$ (47) $$

According to Lemma [E.3](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem3), we can bound $g(f)$, which can be expressed as

$$ $g(f)=\mathop{\mathbb{E}}_{\begin{subarray}{c}s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]=\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left\{\mathop{\mathbb{E}}_{a\sim\tilde{\mu}_{f}(\cdot|s)}\left[A^{\mu_{f}}(s,a)\right]\right\},$ (48) $$

thus,

$$ $\left|g(f)\right|\leq\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left\{\left|\mathop{\mathbb{E}}_{a\sim\tilde{\mu}_{f}(\cdot|s)}\left[A^{\mu_{f}}(s,a)\right]\right|\right\}\leq\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left\{2D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\cdot\max_{a}\left|A^{\mu_{f}}(s,a)\right|\right\}.$ (49) $$

Denote $\delta=\max_{s,a}\left|A^{\mu_{f}}(s,a)\right|$, then we have

$$ $\left|g(f)\right|\leq 2\delta\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],$ (50) $$

which means that

$$ $\begin{split}\Phi\leq&\frac{1-Z}{Z}\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot\left|g(f)\right|\mathrm{d}f+\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot\left|g(f)\right|\mathrm{d}f\\ \leq&\frac{2\delta_{\mathrm{train}}(1-Z)}{Z}\int_{\mathcal{F}_{\mathrm{train}}}p(f)\cdot\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\mathrm{d}f\\ &+2\delta_{\mathrm{eval}}\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}p(f)\cdot\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\mathrm{d}f\\ =&2\delta_{\mathrm{train}}(1-Z)\int_{\mathcal{F}_{\mathrm{train}}}\frac{p(f)}{Z}\cdot\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\mathrm{d}f\\ &+2\delta_{\mathrm{eval}}(1-Z)\int_{\mathcal{F}-\mathcal{F}_{\mathrm{train}}}\frac{p(f)}{1-Z}\cdot\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\mathrm{d}f\\ =&2\delta_{\mathrm{train}}(1-Z)\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]+2\delta_{\mathrm{eval}}(1-Z)\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],\\ \end{split}$ (51) $$

where $\delta_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{\max_{s,a}\left|A^{\mu_{f}}(s,a)\right|\right\}$ and $\delta_{\mathrm{eval}}=\max_{f\in\mathcal{F}_{\mathrm{eval}}}\left\{\max_{s,a}\left|A^{\mu_{f}}(s,a)\right|\right\}$.

Finally, combining ([35](https://arxiv.org/html/2501.02481v5#A5.E35)), ([41](https://arxiv.org/html/2501.02481v5#A5.E41)), ([42](https://arxiv.org/html/2501.02481v5#A5.E42)), ([44](https://arxiv.org/html/2501.02481v5#A5.E44)), and ([51](https://arxiv.org/html/2501.02481v5#A5.E51)), we have

$$ $\begin{split}\left|\zeta(\tilde{\pi})-L_{\pi}(\tilde{\pi})\right|&\leq\frac{2r_{\max}(1-Z)}{1-\gamma}+\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\\ &+\frac{2\delta_{\mathrm{train}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]+\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],\\ \end{split}$ (52) $$

thus, the generalization performance lower bound is

$$ $\begin{split}\zeta(\tilde{\pi})&\geq L_{\pi}(\tilde{\pi})-\frac{2r_{\max}(1-Z)}{1-\gamma}-\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\\ &-\frac{2\delta_{\mathrm{train}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\frac{2\delta_{\mathrm{eval}}(1-Z)}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],\\ \end{split}$ (53) $$

concluding the proof.
∎

###### Proof.

### E.3 Proof of Theorem 3.3

Theorem [3.3](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem3).
Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\eta(\tilde{\pi})\geq L_{\pi}(\tilde{\pi})-\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right].$ (54) $$

###### Proof.

Since

$$ $\begin{split}\left|\eta(\tilde{\pi})-L_{\pi}(\tilde{\pi})\right|&=\frac{1}{1-\gamma}\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\\ a\sim\tilde{\mu}_{f}(\cdot|s)\end{subarray}}\left[A^{\mu_{f}}(s,a)\right]\right|=\frac{\Psi}{1-\gamma}\\ &\leq\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],\end{split}$ (55) $$

thus,

$$ $\eta(\tilde{\pi})\geq L_{\pi}(\tilde{\pi})-\frac{2\gamma\epsilon_{\mathrm{train}}}{(1-\gamma)^{2}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],$ (56) $$

concluding the proof.
∎

###### Proof.

### E.4 Proof of Theorem 3.5

Theorem [3.5](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem5).
Given any two policies, $\tilde{\pi}$ and $\pi$, the following bound holds:

$$ $\mathfrak{D}_{1}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)\mathfrak{D}_{\mathrm{train}},\kern 5.0pt\mathfrak{D}_{2}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\right)\underbrace{\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]}_{\text{denote it as }\mathfrak{D}_{\mathrm{eval}}},$ (57) $$

where $\sigma_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{D_{\mathrm{TV}}^{\max}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}$ and $\sigma_{\mathrm{eval}}=\max_{f\in\mathcal{F}_{\mathrm{eval}}}\left\{D_{\mathrm{TV}}^{\max}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}$, $D_{\mathrm{TV}}^{\max}(\tilde{\mu}_{f}\|\mu_{f})[s]$ represents $\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]$.

###### Proof.

According to Lemma [E.2](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem2), we have

$$ $\begin{split}\left|\mathfrak{D}_{1}-\mathfrak{D}_{\mathrm{train}}\right|&=\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\\ &=\left|\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right\}\right|\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\left|\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\right\}\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\left\|d^{\tilde{\mu}_{f}}-d^{\mu_{f}}\right\|_{1}\cdot\left\|D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\|_{\infty}\right\}\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left\{\frac{2\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}\\ &\leq\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]=\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\cdot\mathfrak{D}_{\mathrm{train}},\\ \end{split}$ (58) $$

as a result,

$$ $\mathfrak{D}_{1}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)\mathfrak{D}_{\mathrm{train}}.$ (59) $$

Similarly, using Lemma [E.2](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem2) again, we have

$$ $\begin{split}\left|\mathfrak{D}_{2}-\mathfrak{D}_{\mathrm{eval}}\right|&=\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\tilde{\mu}_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\\ &=\left|\mathop{\mathbb{E}}_{f\sim p_{\mathrm{eval}}(\cdot)}\left\{\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right\}\right|\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{eval}}(\cdot)}\left\{\left|\mathop{\mathbb{E}}_{s\sim d^{\tilde{\mu}_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\right\}\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{eval}}(\cdot)}\left\{\left\|d^{\tilde{\mu}_{f}}-d^{\mu_{f}}\right\|_{1}\cdot\left\|D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\|_{\infty}\right\}\\ &\leq\mathop{\mathbb{E}}_{f\sim p_{\mathrm{eval}}(\cdot)}\left\{\frac{2\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}\\ &\leq\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]=\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\cdot\mathfrak{D}_{\mathrm{eval}},\\ \end{split}$ (60) $$

as a result,

$$ $\mathfrak{D}_{2}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{eval}}}{1-\gamma}\right)\mathfrak{D}_{\mathrm{eval}},$ (61) $$

concluding the proof.
∎

###### Proof.

### E.5 Proof of Theorem 3.7

Theorem [3.7](https://arxiv.org/html/2501.02481v5#S3.Thmtheorem7).
Given any two policies, $\tilde{\pi}$ and $\pi$, assume that $\tilde{\pi}$ is $\mathcal{R}_{\tilde{\pi}}$-robust, and $\pi$ is $\mathcal{R}_{\pi}$-robust, then the following bound holds:

$$ $\mathfrak{D}_{\mathrm{eval}}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)\mathcal{R}_{\pi}+\mathcal{R}_{\tilde{\pi}}+\mathfrak{D}_{\mathrm{train}}.$ (62) $$

###### Proof.

Let’s first rewrite $\mathfrak{D}_{\mathrm{eval}}$ as

$$ $\mathfrak{D}_{\mathrm{eval}}=\mathop{\mathbb{E}}_{\begin{subarray}{c}\tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\mu_{\tilde{f}})[s]\right].$ (63) $$

For another $f\in\mathcal{F}_{\mathrm{train}}$, by repeatedly using the triangle inequality of the TV distance, we have

$$ $\begin{split}\mathfrak{D}_{\mathrm{eval}}&=\mathop{\mathbb{E}}_{\begin{subarray}{c}\tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\mu_{\tilde{f}})[s]\right]\\ &\leq\mathop{\mathbb{E}}_{\begin{subarray}{c}\tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\tilde{\mu}_{f})[s]+D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]+D_{\mathrm{TV}}(\mu_{f}\|\mu_{\tilde{f}})[s]\right]\\ &=\mathop{\mathbb{E}}_{\begin{subarray}{c}\tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\tilde{\mu}_{f})[s]\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}\tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}\tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\mu_{f}\|\mu_{\tilde{f}})[s]\right],\\ \end{split}$ (64) $$

taking the expectation of both sides of the inequality with respect to $f\sim p_{\mathrm{train}}(\cdot)$, we obtain

$$ $\mathop{\mathbb{E}}_{f\sim p_{\mathrm{train}}(\cdot)}\left[\mathfrak{D}_{\mathrm{eval}}\right]\leq\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\tilde{\mu}_{f})[s]\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\mu_{f}\|\mu_{\tilde{f}})[s]\right].$ (65) $$

Since $\mathfrak{D}_{\mathrm{eval}}$ is independent of $f$, it becomes a constant after taking the expectation, which is

$$ $\mathfrak{D}_{\mathrm{eval}}\leq\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\tilde{\mu}_{f})[s]\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\mu_{f}\|\mu_{\tilde{f}})[s]\right].$ (66) $$

Note that $\tilde{\pi}$ is $\mathcal{R}_{\tilde{\pi}}$-robust, and $\pi$ is $\mathcal{R}_{\pi}$-robust, we can thus bound the first term:

$$ $\begin{split}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\tilde{\mu}_{f})[s]\right]=&\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\end{subarray}}\left[\sum_{s\in\mathcal{S}}d^{\mu_{\tilde{f}}}(s)\cdot D_{\mathrm{TV}}(\tilde{\mu}_{\tilde{f}}\|\tilde{\mu}_{f})[s]\right]\\ \leq&\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\end{subarray}}\left[\sum_{s\in\mathcal{S}}d^{\mu_{\tilde{f}}}(s)\cdot\mathcal{R}_{\tilde{\pi}}\right]=\mathcal{R}_{\tilde{\pi}}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\end{subarray}}\left[\sum_{s\in\mathcal{S}}d^{\mu_{\tilde{f}}}(s)\right]=\mathcal{R}_{\tilde{\pi}}.\\ \end{split}$ (67) $$

Similarly, we can bound the third term:

$$ $\begin{split}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]\right]=&\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\end{subarray}}\left[\sum_{s\in\mathcal{S}}d^{\mu_{\tilde{f}}}(s)\cdot D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]\right]\\ \leq&\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\end{subarray}}\left[\sum_{s\in\mathcal{S}}d^{\mu_{\tilde{f}}}(s)\cdot\mathcal{R}_{\pi}\right]=\mathcal{R}_{\pi}\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\end{subarray}}\left[\sum_{s\in\mathcal{S}}d^{\mu_{\tilde{f}}}(s)\right]=\mathcal{R}_{\pi}.\\ \end{split}$ (68) $$

Next, we are trying to bound the second term, which is similar to $\mathfrak{D}_{\mathrm{train}}$. Note that $\mathfrak{D}_{\mathrm{train}}$ is independent of $\tilde{f}$, we can thus rewrite it as

$$ $\mathfrak{D}_{\mathrm{train}}=\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]=\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right],$ (69) $$

then

$$ $\begin{split}&\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathfrak{D}_{\mathrm{train}}\right|\\ =&\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{f}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\\ =&\left|\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\int_{\mathcal{F}_{\mathrm{eval}}}p_{\mathrm{eval}}(\tilde{f})\left\{\mathop{\mathbb{E}}_{s\sim d^{\mu_{\tilde{f}}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right\}\mathrm{d}\tilde{f}\mathrm{d}f\right|\\ \leq&\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\int_{\mathcal{F}_{\mathrm{eval}}}p_{\mathrm{eval}}(\tilde{f})\left\{\left|\mathop{\mathbb{E}}_{s\sim d^{\mu_{\tilde{f}}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\right\}\mathrm{d}\tilde{f}\mathrm{d}f.\\ \end{split}$ (70) $$

Note that,

$$ $\left|\mathop{\mathbb{E}}_{s\sim d^{\mu_{\tilde{f}}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\leq\left\|d^{\mu_{\tilde{f}}}-d^{\mu_{f}}\right\|_{1}\cdot\left\|D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\|_{\infty}.$ (71) $$

According to Lemma [E.2](https://arxiv.org/html/2501.02481v5#A5.Thmtheorem2),

$$ $\left\|d^{\mu_{\tilde{f}}}-d^{\mu_{f}}\right\|_{1}\leq\frac{2\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]\right],$ (72) $$

$\pi$ is $\mathcal{R}_{\pi}$-robust, so,

$$ $\left\|d^{\mu_{\tilde{f}}}-d^{\mu_{f}}\right\|_{1}\leq\frac{2\gamma}{1-\gamma}\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]\right]=\frac{2\gamma}{1-\gamma}\sum_{s\in\mathcal{S}}d^{\mu_{f}}(s)\cdot D_{\mathrm{TV}}(\mu_{\tilde{f}}\|\mu_{f})[s]\leq\frac{2\gamma}{1-\gamma}\mathcal{R}_{\pi}.$ (73) $$

As a result,

$$ $\begin{split}&\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathfrak{D}_{\mathrm{train}}\right|\\ \leq&\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\int_{\mathcal{F}_{\mathrm{eval}}}p_{\mathrm{eval}}(\tilde{f})\cdot\left\{\left|\mathop{\mathbb{E}}_{s\sim d^{\mu_{\tilde{f}}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathop{\mathbb{E}}_{s\sim d^{\mu_{f}}(\cdot)}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\right|\right\}\mathrm{d}\tilde{f}\mathrm{d}f\\ \leq&\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\int_{\mathcal{F}_{\mathrm{eval}}}p_{\mathrm{eval}}(\tilde{f})\cdot\left\{\frac{2\gamma}{1-\gamma}\mathcal{R}_{\pi}\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}\mathrm{d}\tilde{f}\mathrm{d}f\\ =&\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\cdot\left\{\frac{2\gamma}{1-\gamma}\mathcal{R}_{\pi}\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}\cdot\int_{\mathcal{F}_{\mathrm{eval}}}p_{\mathrm{eval}}(\tilde{f})\mathrm{d}\tilde{f}\mathrm{d}f\\ =&\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\cdot\left\{\frac{2\gamma}{1-\gamma}\mathcal{R}_{\pi}\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}\mathrm{d}f=\frac{2\gamma}{1-\gamma}\mathcal{R}_{\pi}\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\mathrm{d}f.\\ \end{split}$ (74) $$

We previously defined $\sigma_{\mathrm{train}}=\max_{f\in\mathcal{F}_{\mathrm{train}}}\left\{\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right\}$, so that

$$ $\begin{split}\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]-\mathfrak{D}_{\mathrm{train}}\right|\leq&\frac{2\gamma}{1-\gamma}\mathcal{R}_{\pi}\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\cdot\max_{s}D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\mathrm{d}f\\ \leq&\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\mathcal{R}_{\pi}\int_{\mathcal{F}_{\mathrm{train}}}p_{\mathrm{train}}(f)\mathrm{d}f=\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\mathcal{R}_{\pi},\\ \end{split}$ (75) $$

thus, the second term is bounded by

$$ $\mathop{\mathbb{E}}_{\begin{subarray}{c}f\sim p_{\mathrm{train}}(\cdot)\\ \tilde{f}\sim p_{\mathrm{eval}}(\cdot)\\ s\sim d^{\mu_{\tilde{f}}}(\cdot)\end{subarray}}\left[D_{\mathrm{TV}}(\tilde{\mu}_{f}\|\mu_{f})[s]\right]\leq\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\mathcal{R}_{\pi}+\mathfrak{D}_{\mathrm{train}}.$ (76) $$

Finally, combining ([67](https://arxiv.org/html/2501.02481v5#A5.E67)), ([68](https://arxiv.org/html/2501.02481v5#A5.E68)) and ([76](https://arxiv.org/html/2501.02481v5#A5.E76)), we have

$$ $\mathfrak{D}_{\mathrm{eval}}\leq\left(1+\frac{2\gamma\sigma_{\mathrm{train}}}{1-\gamma}\right)\mathcal{R}_{\pi}+\mathcal{R}_{\tilde{\pi}}+\mathfrak{D}_{\mathrm{train}},$ (77) $$

concluding the proof.
∎

###### Proof.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1) Have the authors analysed or discussed the additional cost introduced by mutual distillation? 

2) Has the method been tested across tasks with different sources of irrelevant variation, or is it specific to the evaluated benchmarks?

3) Can the representational changes be visualised or quantified to support the theoretical claims?

4) Could the authors comment on how their approach differs from or relates to the works mentioned above?
soundness: 2
strengths: The paper presents a new theoretical framework to investigate generalisation issues in deep RL. Generalisation in RL is a major and actively researched topic. The insights provided by the paper will have far reaching impact.
summary: This paper presents a theoretical framework to demonstrate that improving the policy robustness to irrelevant features enhances its generalisation performance. The paper further shows that deep mutual learning forms an implicit regularisation and prevents policy from overfitting to irrelevant features. Empirical results are given on the ProcGen benchmark, designed to test generalisability under controlled environments, as well as for toy examples. The proposed method, mutual distillation policy optimisation, demonstrates benefits as compared to the selected baseline approaches.
weaknesses: Although impactful, the experimental evaluation is limited. In the sense that, it doesn't demonstrate the phenomenon exists, beyond testing on the ProcGen benchmark and presenting performance. Also, apart from the toy example. 

There are other methods focusing on distillation (mutual or peer). However, these papers seem not to be mentioned in the paper. It would be good to see a comparison, for example, 

* Periodic Intra-Ensemble Knowledge Distillation for Reinforcement Learning, https://arxiv.org/pdf/2002.00149

* Robust Domain Randomised Reinforcement Learning through Peer-to-Peer Distillation, https://arxiv.org/pdf/2012.04839 (this has been cited in the paper) 

* Online Policy Distillation with Decision-Attention, https://arxiv.org/pdf/2406.05488

While the theoretical framework seems to be strong, the paper lacks interpretability. More experiments would be helpful to understand the impact of mutual distillation on the representation space. 

Running multi distillation may increase the computational overhead, which has not beed discussed in the paper. 

Mutual distillation offers regularisation by reducing the reliance on irrelevant features. However, the generalisability of the approach hasn't been investigated.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: I am wondering how MDPO will scale to large number of policies beyond only two policies as discussed. Can you share some insights on that?
soundness: 2
strengths: ### Strengths:

1. This work presents a formal proof of a long standing assumption that robustness of policy against irrelevant features improves generalization. In particular, they derive a lower bound for generalization performance that includes minimization of a robustness term, which is defined how a policy is influenced by two different rendering (perturbation) functions.

2. The paper is presented in a clear and well-organized way. Especially, Fig. 1 and 2 helps the readers to better understand the intuition and impact of DML in the discussed setting. 

3. Legitimate ablations are conducted and the results validate the claims.
summary: This paper attempts to prove the conjecture that the policy which is robust to irrelevant features would lead to improved generalization performance. In addition, they propose mutual distillation of policies to achieve such robustness and presents the intuition behind that. While they do not show state-of-the-arts performance, they present proof-of-concept for their approach with basic regularization baseline on all environments of Procgen benchmark.
weaknesses: ### Weaknesses:

1. The proposed method has been validated only on the ProcGen benchmark. Experiments on more diverse set up is needed to show the applicability of such methods. 

2. While I understand that the target is not to outperform the state-of-the arts, but how DML stands against other data augmentation based approaches such as [1] are not evident. While the authors present result with SPO, it seems SPO performance itself is not upto the current standard. 

3. The proposed method relies on multiple policies for distillation. However, the computational overhead compared to single policy methods is not discussed.

[1] Raileanu, Roberta, et al. "Automatic data augmentation for generalization in reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 5402-5415.

---

confidence: 3
contribution: 3
initial_rating: 2
presentation: 1
questions: - Does it matter how $\pi$ and $\tilde{\pi}$ are related? Can we view them through the lens of the two policies used in mutual distillation somehow? 
- The policy $\pi$, $\tilde{\pi}$ and the linear approximation $L_\pi(\tilde{\pi})$ (the only positive term in the lower bound in Corollary 3.8), seem to originate from [9], which assumes $\pi$ and $\tilde{\pi}$ are close to each other. Does this mean the lower bound will only be non-vacuous if $\pi$ and $\tilde{\pi}$ are similar? How does this impact the significance of the derived bounds?
- Line 200: "During the training process, we can only empirically bound $\mathfrak{D}_{train}$" Why? 
- Figure 2: "(Right) Through mutual distillation via DML, two policies regularize each other to converge toward a more robust hypothesis space,..." I agree that they are regularized towards each other, but why would that be toward a more robust hypothesis space, and not just any other part of the non-robust hypothesis space?
- Why is an algorithm that regularizes two policies towards each other, the solution for the problem identified in Figure 1 that requires the two policies to collect substantially different data?
- For the experiments, how many seeds are used and what do the shaded regions indicate? Also, how is significance of the results determined?


**Things to improve that did not impact decision:**
- There many missing words, making sentences incomplete (for example, line 58 and 65 of the introduction). 
- Section 2.1: Stating that a function $f: X \to [0,1]$ does not sufficiently define it as a probability distribution (it is missing the constraint that all the probabilities sum up to 1). 
- Section 2.1: The authors seem to introduce a new MDP framework with the rendering function in the background section. Either this is new, which means it should not be in the background section, or it is an existing framework from somewhere else, which means it is missing a citation. 
	- The framework also seems reminiscent of a specific type of contextual MDP [10], perhaps it could be useful to frame this work within the CMDP framework. 
- Table 2: It would be interesting to also include the training performance for this experiment.
- Figure 3: Does the x-axis for MDPO include the timesteps of both agents (since they collect data independently)? In other words, at timestep 50 million in the figure, the individual MDPO policies will only have trained on 25 million steps each? If not, I feel this figure is slightly misleading. 

**References:**\
[1] Learning Invariant Representations for Reinforcement Learning Without Reconstruction. Zhang et al. 2021\
[2] Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL. Mazoure et al. 2022\
[3] Learning Dynamics and Generalization in Reinforcement Learning. Lyle et al. 2022\
[4] How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning. Weltevrede et al. 2025\
[5] Dual Policy Distillation. Lai et al. 2020\
[6] On the Importance of Exploration for Generalization in Reinforcement Learning. Jiang et al. 2023\
[7] Policy Confounding and Out-of-Trajectory Generalization in RL. Suau et al. 2024\
[8] Exploration Implies Data Augmentation: Reachability and Generalisation in Contextual MDPs. Weltevrede et al. 2025\
[9] Trust Region Policy Optimization. Schulman et al. 2015\
[10] A survey of zero-shot generalisation in deep reinforcement learning. Kirk et al. 2023
soundness: 2
strengths: - The central idea of using mutual distillation to induce robustness to spurious correlations in the training data is interesting. 
- The empirical results for MDPO look promising.
- The analytical experiments of the robustness of the MDPO policy to visual disturbances and the quality of the learned representations in Sections 5.3 and 5.4 are very insightful.
summary: The paper presents a theoretical and empirical study arguing that mutual distillation between reinforcement learning policies acts as an implicit regularization to prevent overfitting to irrelevant features, thereby improving generalization. It claims to be the first to theoretically prove that enhancing policy robustness to irrelevant features leads to better generalization. The paper introduces Mutual Distillation Policy Optimization (MDPO), which uses Deep Mutual Learning (DML) via a KL divergence loss between to (otherwise independently) trained PPO agents. Empirically, MDPO shows improved generalization performance over a PPO baseline on the hard level configuration of the Procgen benchmark and demonstrates enhanced robustness to visual disturbances like random convolutions or changes in brightness, contrast, saturation or hue.  

**Recommendation:**\
I recommend to reject the paper because it is lacking in certain fundamental areas: insufficient positioning with respect to related work,  issues with clarity regarding the significance of theoretical results, empirical results, and the novelty of the ideas/approach.
weaknesses: - The positioning of this work is severely lacking, especially concerning previous literature.
	-  The paper seems to miss several related works on topics such as:
		- Representation learning in RL (for example, [1,2]) 
		- Policy distillation for generalization (for example, [3,4]) 
		- Mutual distillation (for example, [5])
		- Overfitting to training data in RL (for example, [6,7,8])
		- The above are not necessarily exhaustive. 
	- The claims regarding theoretical novelty are strong, but difficult for me to verify due to insufficient discussion of existing work, and my own lack of knowledge on this specific topic.  
	- The novelty of the motivation in Section 4 is difficult to judge, since it lacks comparison with related ideas in the existing literature (for example, see [7,8]). 
- The theoretical contribution has some issues:
	- The core theoretical claim is formulated too strongly. Corollary 3.8 only shows an increase in the _lower bound_ on generalization performance as robustness increases. However, increasing a lower bound does not guarantee an increase in actualized performance, which is what is promised in the abstract and introduction. 
	- Clarity of the theoretical framework is lacking. For example, the proofs depend on two policies $\pi$ and $\tilde{\pi}$, but there is no mention of what the significance of these two policies are or how they are related to the overall narrative of the paper. 
	- A discussion on the significance of the theoretical results is missing. Whether increasing a lower bound has any bearing on the realized performance depends on whether the bound is vacuous or not. This means the bounds derived in Section 3 would benefit from analysis or discussion on how tight they are. For example, the bound's dependence on the relative performance of two arbitrary policies $\pi$ and $\tilde{\pi}$ complicates the interpretation of its significance. 
- The motivation for mutual distillation seems conceptually incomplete. The logic presented in Sections 4.1 and 4.2 hinges on the policies encountering different spurious correlations due to collecting different trajectories as a result of different initializations. However, the DML mechanism's goal is to regularize the two policies to converge to the same policy. This seems to introduce a paradox: in order to benefit from the two policies collecting different training data, they need to be regularized to be the same. 
- Experimental details are missing, making it difficult to judge the validity and significance of the results. For example, the paper claims significant improvements, but there is no mention of number of seeds used, what the shaded areas in the figures denote, or how significance was tested.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: - can you provide a theoretical statement (e.g., in linearized form) showing that mutual distillation dynamics reduce the $R$-robustness term? 
- how do you define the “representation convergence” concept?
- do the mutual distillation learners together use the same amount of samples as PPO baselines? If not, how is data shared or counted? 
- could you describe the random CNNs used to generate alternate renderings and the motivation behind it?
soundness: 3
strengths: - The paper is easy to follow for the most part and has a clear motivated given a widespread interest in generalization and representation learning in RL.
- exploring generalization in RL through distillation is well-motivated and interesting. As far as I know, this technique is not exhaustively researched / understood in the context of RL.
- The formalization of rendering families and the decomposition of generalization error into robustness and train-test divergence components is intuitive and provides a neat mathematical framing. I believe this framing ends being equivalent to classical POMDP definitions, so it may be worth to align language (e.g., rendering function - emission function / observation function). 
- MDPO performs consistently better than PPO and other baselines on Procgen tasks, even when controlling for model size or training budget (at least partly).
summary: This paper studies generalization in reinforcement learning by examining the robustness of learned features. It proposes (1) a theoretical framework showing how policies sensitive to “rendering” features that do not generalizable suffer from worse performance in expected returns and (2) the conjecture that mutual distillation between two agents regularizes representation learning toward such robust, invariant features and thus improves generalization.

The paper derives bounds relating expected test returns to terms measuring representation robustness and train/test differences, and empirically investigates whether mutual distillation indeed yields more robust policies on Procgen environments. The empirical results exhibit strong performance gains of the mutual distillation variant over PPO and several controls.
weaknesses: **Theoretical novelty and completeness**:

I found that the novelty of the theoretical exposition is overstated in serveral places. One of the main theorems (Theorem 3.3) mirrors the first-order performance difference bound from Schulman et al. (2015, TRPO) so closely that I believe it should be cited as a known existing bound. 
The paper furthermore claims to be the first to “ first to provide a rigorous proof of {the} intuition {that robustness to irrelevant features enhances generalization performance}", which is an overstatement in this phrasing in my view. Several prior works in the literature provide generalization (some as part of sample-complexity) bounds explicitly dependent on representation adequacy. A non-exhaustive list: 

- kernel RL and kernel complexity as a measure of representation capacity: Yeh et al. (2023)
- from a causality perspective: Kallus et al. (2020), Suau et al. (2023)
- from a symmetry and invariance perspective: Weltevrede et al. (2025)

The latter of the above even considers distillation for deep RL policies and should most certaintly be discussed in this work. 

**Disconnect between theory and practice**
The biggest weakness of this paper in my view is the strong disparity between the derived theoretical claims and the conjectured effect of mutual distillation. The argument or intuition for how mutual distillation learns more robust features is very handwavy. For example it is not obvious at all that policies from different intialization or different batch orders or other factors will learn significantly different features. The work above indeed does show a related effect theoretically using neural tangent kernel theory (Weltevrede et al.). Knowing that feature learning is difficult to treat theoretically in deep learning I would suggest an alternative path would be to support the argument more empirically. For example with experiments showing indeed that different initialization / different exploration seeds / different batch orders etc. lead to feature learning of distinct spurious features which can then be eliminated through mutual distillation. 

**Presentation**
The paper has several issues with presentation and clariity. Apart from several writing / language issues, the are also a number of ambiguities: 
- Inconsistency in robustness metric. The robustness term is defined as a maximum total variation distance max TV distance, yet Table 1 lists values larger than 1.
- Algorithm 1 implies two agents collecting data separately; it is unclear whether the 50M total steps refer to per agent or shared. The appendix does not explicitly resolve this. Without this clarification, comparisons to PPO may reflect more total experience rather than representational effects.
- Adversarial/robustness experiments underspecified. The procedure for generating “adversarial renderings” (random CNNs) lacks some detail, and, importantly, motivation. I find this a very perculiar way of generating adversarial examples, could the authors elaborate this choice?
- some quantities (e.g., $L_\pi$) is not defined in such a way that the paper is self-sufficient. One should not need to refer to TRPO to recall its definition.
- Despite being in the title, the term "representation convergence" is never formally defined. I'm not following what this term aims to imply, what converges to what? 


- Yeh, Sing-Yuan, et al. "Sample complexity of kernel-based q-learning." International Conference on Artificial Intelligence and Statistics. PMLR, 2023.
- Kallus, Nathan, and Angela Zhou. "Confounding-robust policy evaluation in infinite-horizon reinforcement learning." Advances in neural information processing systems 33 (2020): 22293-22304.
- Suau, Miguel, Matthijs TJ Spaan, and Frans A. Oliehoek. "Bad habits: Policy confounding and out-of-trajectory generalization in RL." arXiv preprint arXiv:2306.02419 (2023).
- Weltevrede, Max, et al. "How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning." arXiv preprint arXiv:2505.16581 (2025).


Decision: Reject