Exploration and Anti-Exploration with Distributional Random Network Distillation

The RND method faces challenges with bonus inconsistencies, which can be categorized into initial and final bonus inconsistencies. The initial bonus inconsistency relates to the uneven distribution of bonuses among states at the beginning of training. Addressing this issue is crucial to preventing significant bonus value disparities among states. Conversely, the final bonus inconsistency arises when the final bonuses do not align with the dataset distribution, making it hard for the agent to effectively distinguish between frequently visited states and those encountered relatively fewer times. This issue becomes particularly pronounced after substantial updates to the predictor network, which hinders the agent’s ability to engage in deep exploration. This issue is depicted in Figure 1.

To tackle this, we introduce a method that distills a random distribution, enhancing performance with minimal computational overhead and addressing the bonus inconsistency challenges.

Unlike RND, which only has one target network $\bar{f}$, the DRND algorithm has $N$ target networks $\bar{f}_{1},\bar{f}_{2},...,\bar{f}_{N}$, which are from a random distribution with randomly initialized parameters and do not participate in training. In DRND, we use $s$ as input in the online setting and $(s,a)$ pair as input in the offline setting. For simplicity, we define $x=(s,a)$ (offline setting) or $x=(s)$ (online setting). For each state-action pair $x$, we construct a variable $c(x)$ that satisfies the distribution:

For simplicity, we use some symbols to record the moments of the distribution:

Each time $x$ occurs, $c(x)$ is sampled from this distribution. We use a predictive network $f_{\theta}(x)$ to learn the variable $c(x)$, although using a fixed network to learn a random variable is impossible. We use the MSE loss function to force $f_{\theta}(x)$ to align with $c(x)$ and the loss is

By minimizing the loss, the optimal $f_{*}(x)$ when the state-action pair $x$ appears $n$ times is

where $c_{i}(x)$ is defined as $c(x)$ at the $i$-$th$ occurrence of state $x$. For RND, the more times the predictor network is trained in the same state, the closer the output of the target network is. Therefore, directly using loss as a bonus encourages agents to explore new states. The bonus of our method is not equal to the loss since the loss of random variable fitting is unstable. The expected value of the prediction net is given by

where $n$ is the count of occurrences of $x$. After multiple training iterations, this value approaches the mean of the target network distribution. Hence, to measure the deviation from this mean, the first bonus of DRND is defined as

Compared to predicting the output of one target net, predicting the mean of multiple networks is equivalent to passing through a high-pass filter on the output of multiple networks, which can avoid the problem of initial bonus inconsistency due to extreme values in one network. Especially if the network is linear, this bonus inconsistency can be quantitatively calculated.

The complete proof can be seen in Appendix A.1 . Lemma 4.1 shows that if the predictor parameters and target parameters are sampled from the same distribution, the expectation of the first bonus is a function of input $x$.

Lemma 4.2 suggests that when the input $x$ is confined to a bounded interval, and when Equation 6 is utilized to calculate the initial bonus, the expected maximal difference is modulated by the number of target networks. Importantly, this anticipated discrepancy tends to decrease as $N$ increases. This observation substantiates that our DRND method, equipped with $N$ target networks, exhibits lower bonus inconsistency under a uniform distribution than the RND method, which uses only a single target network.

However, it is essential to note that the network fitting loss determines this bonus. Consequently, it cannot distinguish between states visited multiple times, which stands in contrast to count-based and pseudo-count methods, which do not address the issue of final bonus inconsistency.

It is essential to track data occurrence frequencies to address inconsistent final bonuses. Traditional count-based methods use large tables to record state visitations, while pseudo-count strategies use neural networks for estimation, providing a scalable insight into state visits. However, these methods introduce computational and storage complexities, particularly when dealing with high-dimensional inputs. We constructed a statistic that indirectly estimates state occurrences without extra auxiliary functions.

The complete proof can be found in Appendix A.2. Lemma 4.3 shows that when $n$ is large, this statistic can effectively recover the number of occurrences of the state $n$, thus implicitly recording the number of occurrences of the state like the pseudo-count method. By minimizing Equation 3 can make $f_{\theta}(x)$ and $f_{*}(x)$ infinitely close, so we replace $f_{*}(x)$ in $y(x)$ with $f_{\theta}(x)$ and approximately assume that they are equal. The DRND predictor potentially stores in its weights how much of each state vector is present in the dataset. To correspond to $\sqrt{1/n}$ of the count-based method, the second bonus of the DRND agent is

which is the estimation of $\sqrt{1/n}$.

The total bonus that combines Equations 6 and 9 gives

where $\alpha$ represents the scaling factor for the two bonus terms. Figure 2 shows the diagram of our method and the RND method. For smaller values of $n$, the variance of the second bonus estimate is substantial, rendering the first bonus $b_{1}$ a more dependable measure for states with infrequent occurrences. Conversely, as $n$ increases, the variance of the second bonus $b_{2}$ approaches zero, enhancing its reliability. From the experiments, we found that during the initial training phase, the magnitudes of the two bonuses, denoted as $b_{1}$ and $b_{2}$, are nearly the same. However, as training progresses, the rate at which the first bonus $b_{1}$ decreases is much faster than that of the second bonus $b_{2}$. Eventually, the magnitude of $b_{1}$ becomes approximately two orders of magnitude lower than $b_{2}$. Therefore, if we want $b_{1}$ to dominate during the early stages of training and have $b_{2}$ dominate during the later stages, we do not need to set a dynamic coefficient to achieve this functionality. Instead, we can simply set it as a fixed constant $\alpha$. The distribution of various bonuses as well as the total bonus of the DRND method before and after training can be seen in Figure 3.

The aforementioned heatmap serves as a qualitative illustration. We now provide quantitative experimental results to demonstrate that DRND offers superior bonuses compared to RND. In count-based exploration methods, it is common to assess the rationality of the bonus distribution by comparing it with the $1/\sqrt{n}$ distribution, where $n$ represents the state visitation count. We evaluate the discrepancy between the intrinsic reward distribution and the uniform distribution $U$ before training, and the distribution of $1/\sqrt{n}$ after training using KL divergence. We randomly select 100 initial dataset distributions and report the mean and variance of the KL divergence. This assessment is denoted as $D_{KL}(P\|U)$ and $D_{KL}(P\|1/\sqrt{n})$, where $P$ represents the distribution of RND bonus or DRND bonus. The results obtained are presented in Table 1. It is evident that the KL divergence between the bonuses of DRND and the uniform distribution is smaller during the initial training stages, effectively promoting uniform exploration by the agent. In the later stages of training, DRND better aligns with the distribution of $1/\sqrt{n}$, thereby facilitating deeper exploration by the agent.

To validate that $b_{1}$ gives better and more uniform bonuses during the early stages of training, while $b_{2}$ better captures the state visitation count $n$ later on, we similarly measure the KL divergence between each bonus term $b_{1}$, $b_{2}$, the overall bonus $b$ and (a) the uniform distribution $U$ before training; (b) the distribution of $1/\sqrt{n}$ after training. The results are presented in Table 2.

From the table, it can be observed that in the early stages of training, $b_{1}$ performs better, providing a better solution to the initial bonus inconsistency. Conversely, in the later stages of training, $b_{2}$ performs better, addressing the final bonus inconsistency more effectively. Moreover, from the KL divergence of the total bonus, it is evident that setting a fixed value for $\alpha$ achieves the desired effect of having $b_{1}$ dominate during the early stages and $b_{2}$ dominate during the later stages, without the need for dynamic adjustments.

The combination of $b_{1}$ and $b_{2}$ depends on the environment and task at hand. In some environments, $b_{1}$ may be more crucial, while in others, the contribution of $b_{2}$ might be greater. However, regardless of the specifics, properly integrating both can lead to improved performance.

The target network of the DRND method remains static throughout the training process. Moreover, its loss and intrinsic reward computations do not entail additional backpropagation steps, thereby maintaining computational efficiency comparable to that of the RND algorithm. Specifically, when the hyperparameters $\alpha$ and $N$ are set to 1, the expressions for the loss and intrinsic reward become simplified to $\left\|f_{\theta}(x)-\bar{f}(x)\right\|^{2}$, which coincides with the formulation utilized in the RND framework. In contrast to count-based and pseudo-count methods, we forgo the utilization of an additional network or table to track state occurrences. Instead, we estimate these occurrences using the statistical information derived from the prediction network itself. When $\alpha$ is set to 0 and $c(x)$ follows the distribution $c(x)\sim\begin{array}[]{c|c c }\hline\cr X&-1&1\\ \hline\cr P&0.5&0.5\\ \hline\cr\end{array}$, the expressions for the loss and intrinsic reward are simplified to $\|f_{\theta}(x)\|^{2}$ and $\|f_{\theta}(x)\|$, which align with the pseudo-count approach CFN.

In this sub-section, we introduce our inconsistency experiments to compare bonus predictions for both RND and DRND. We created a mini-dataset resembling those used in offline RL or online RL replay buffers in the experiments. This small dataset contains $M$ data categories labeled from $1$ to $m$, with each data type occurring $i$ times proportional to its label. Each data point is represented as a one-hot vector with $M$ dimensions, where $M$ is set to 100. We train both the RND and DRND networks on the dataset and record both the initial intrinsic reward and the final intrinsic reward.

The left panel in Figure 4 illustrates the difference in initial intrinsic rewards between RND and DRND, with the x-axis representing the number of target networks. As $N$ increases, the y-axis, representing the range of intrinsic rewards, becomes narrower, resulting in a more uniform distribution of rewards. In the right panel of Figure 4, we display the intrinsic reward distribution trained on the mini-dataset, showing that DRND’s rewards have a stronger correlation with sample count than RND, as indicated by the regression lines.

Like many other exploration methods, we conduct our DRND approach in Atari games, Adroit environments (Rajeswaran et al., 2017), and fetch manipulation tasks (Plappert et al., 2018), which need deep exploration to get a high score. We integrate our method with the PPO (Schulman et al., 2017) algorithm. We compare our approach with the RND method, the pseudo-count method CFN (Lobel et al., 2023), the curiosity-driven method ICM (Pathak et al., 2017), and the baseline PPO method. The solid lines in the figures represent the mean of multiple experiments, and the shading represents the standard deviation interval.

Atari Games. We chose three Atari games — Montezuma’s Revenge, Gravitar, and Venture — to evaluate our algorithms. These games require deep exploration to achieve high scores, making them ideal for assessing algorithmic exploratory capabilities. We benchmarked our method against the RND and PPO algorithms, with results presented in Figure 5. Our DRND method converges faster and attains the highest final scores in these environments.

Adroit Experiments. We further delve into the Adroit continuous control tasks. In these challenges, a robot must skillfully manipulate a hand to perform various actions, such as adjusting a pen’s orientation or unlocking a door. Considering the complexity of the tasks and the robot’s high-dimensional state space, it becomes imperative to explore methods that can facilitate the robot’s learning. Figure 6 illustrates that our DRND method outperforms all the other methods in exploration, especially in the challenging ‘Hammer’ and ‘Relocate’ environments. However, in the ‘Pen’ environment, our method does not exhibit a significant improvement compared to other exploration algorithms. This could be attributed to the relatively simpler nature of this environment, which does not demand deep exploration.

Fetch Manipulation Tasks. The Fetch manipulation tasks involve various gym-robotics environments, challenging the Fetch robot arm with complex tasks like reaching, pushing, sliding, and pick-and-place actions. Due to their complexity, these tasks demand advanced exploration strategies. Our evaluation of exploration algorithms in this context highlights their effectiveness in handling intricate robotic manipulations. As shown in Figure 7, our DRND approach excels in assisting the robot in these tasks. Our DRND method effectively combines the strengths of these approaches, outperforming results achievable with either pseudo-count or curiosity-driven methods alone. Consequently, our DRND algorithm performs significantly better than the RND method and other exploration algorithms.

We assessed our method using the D4RL (Fu et al., 2020) offline datasets, integrating the DRND approach with the SAC algorithm (Haarnoja et al., 2018). Considering all available datasets in each domain, we tested SAC-DRND on Gym-MuJoCo and the more intricate AntMaze D4RL tasks. Our analysis compares against notable algorithms as detailed in (Rezaeifar et al., 2022), including IQL (Kostrikov et al., 2021), CQL (Kumar et al., 2020), and TD3+BC (Fujimoto & Gu, 2021). It is worth noting that although our method also has $N$ target networks, they are fixed and not trained, making it ensemble-free. Our SAC-DRND is ensemble-free and only involves training double critics networks. We compare our methods against recent strong model-free offline RL algorithms in Table 3. Additionally, we compare SAC-DRND against strong ensemble-based algorithms like SAC-N in Appendix E.1. Only the results of the ensemble-free methods are shown in the main text. The results are evaluated at the final gradient step over 10 different seeds.

It can be seen that SAC-DRND excels in the majority of MuJoCo tasks, attaining the best results among all ensemble-free methods. On Antmaze tasks, DRND also reached a level similar to SOTA. Compared to SAC-RND, which has comparable computational and storage requirements as our approach, SAC-DRND more effectively captures the dataset distribution, as reflected in its superior average scores and decreased variance. We also conducted experiments on Adroit tasks (Appendix E.2), hyperparameters sensitivity experiments (Appendix E.3) using Expected Online Performance (EOP, (Kurenkov & Kolesnikov, 2022)) and offline-to-online experiments (Appendix E.8).