Beyond The Rainbow: High Performance Deep Reinforcement Learning on a Desktop PC — part 1/4

Cleaned-text code points [0, 15319). Bibliography omitted here; source.md includes it.

# Beyond The Rainbow: High Performance Deep Reinforcement Learning on a Desktop PC

###### Abstract

Rainbow Deep Q-Network (DQN) demonstrated combining multiple independent enhancements could significantly boost a reinforcement learning (RL) agent’s performance. In this paper, we present “Beyond The Rainbow” (BTR), a novel algorithm that integrates six improvements from across the RL literature to Rainbow DQN, establishing a new state-of-the-art for RL using a desktop PC, with a human-normalized interquartile mean (IQM) of 7.4 on Atari-60. Beyond Atari, we demonstrate BTR’s capability to handle complex 3D games, successfully training agents to play Super Mario Galaxy, Mario Kart, and Mortal Kombat with minimal algorithmic changes. Designing BTR with computational efficiency in mind, agents can be trained using a high-end desktop PC on 200 million Atari frames within 12 hours. Additionally, we conduct detailed ablation studies of each component, analyzing the performance and impact using numerous measures. Code is available at https://github.com/VIPTankz/BTR .

###### Keywords:

## 1 Introduction

Deep Reinforcement Learning (RL) has achieved numerous successes in complex sequential decision-making tasks, most rapidly since Mnih et al. (2015) proposed Deep Q-Learning (DQN). With this success, RL has become increasingly popular among smaller research labs, the hobbyist community, and even the general public. However, recent state-of-the-art approaches ( Schrittwieser et al., 2020 ; Badia et al., 2020a ; Hessel et al., 2021 ; Kapturowski et al., 2023 ) are increasingly out of reach for those with more limited compute resources, either in terms of the required hardware or the walltime necessary to train a single agent. This is a unique issue in RL compared to natural language processing or image recognition which have foundation models that can be efficiently fine-tuned for a new task or problem ( Lv et al., 2024 ) . Meanwhile, RL agents must be trained afresh for each environment. Therefore, the development of powerful RL algorithms that can be trained quickly on inexpensive hardware is crucial for smaller research labs and the hobbyist community.

These concerns are not new. Ceron & Castro (2021) highlighted that Rainbow DQN ( Hessel et al., 2018 ) required 34,200 GPU hours (equivalent to 1435 days) of training, making the research impossible for anyone except a few research labs, with more recent algorithms exacerbating this problem. Recurrent network architectures ( Horgan et al., 2018 ) , high update to sample ratio ( D’Oro et al., 2022 ) , and the use of world-models and search-based techniques ( Schrittwieser et al., 2020 ) , all increase the computational resources necessary to train agents. Many of these use distributed approaches requiring multiple CPUs and GPUs (or TPUs), or requiring numerous days and weeks to train a single agent, dramatically decreasing RL’s accessibility.

For this purpose, we develop “Beyond the Rainbow” (BTR), taking the same principle as Rainbow DQN ( Hessel et al., 2018 ) , selecting 6 previously independently evaluated improvements and combining them into a singular algorithm (Section 3 ). These components were chosen for their performance qualities or to reduce the computational requirements for training an agent. As a result, BTR sets a new state-of-the-art score for Atari-60 ( Bellemare et al., 2013 ) (excluding recurrent approaches) with an Interquartile Mean (IQM) of 7.4 1 1 1 All reported IQM scores use the best single evaluation for each environment throughout training as is standard, rather than the agent’s score at 200 million, hence the discrepancy between the overall score and Figure 1 . using a single desktop machine in less than 12 hours, and outperforms Rainbow DQN on Procgen ( Cobbe et al., 2020 ) in less than a fifth of the walltime (Section 4.1 ). Further, we demonstrate BTR’s potential by training agents to solve three modern 3D games for the first time, Mario Kart Wii, Super Mario Galaxy and Mortal Kombat, that each contain complex mechanics and graphics (Section 4.2 ). To verify the effectiveness and effect of the six improvements to BTR, in Section 5.1 , we conduct a thorough ablation of each component, plotting their impact on the Atari-5 environments and in Section 5.2 , we utilize seven different measures to analyse the component’s impact on the agent’s policy and network weights. This allows us to more precisely understand how the components impact BTR beyond performance or walltime.

In summary, we make the following contributions to state-of-the-art RL.

• High Performance (Section 4.1 ) - BTR outperforms the state-of-the-art for non-recurrent RL on the Atari-60 benchmark, with an IQM of 7.4 (compared to Rainbow DQN’s 1.9), outperforming humans on 52/60 games. Furthermore, BTR outperforms Rainbow DQN with Impala on the Procgen benchmark despite using a smaller model and 80% less walltime.

• Modern Environments (Section 4.2 ) - Testing beyond Atari, we demonstrate BTR can train agents for three modern games: Super Mario Galaxy (final stage), Mario Kart Wii (Rainbow Road), and Mortal Kombat (Endurance mode). These environments contain 3D graphics and complex physics and have never been solved using RL.

• Computationally Accessible (Figure 6 ) - Using a high-end desktop PC, BTR trains Atari agents for 200 million frames in under 12 hours, significantly faster than Rainbow DQN’s 35 hours. This increases RL research’s accessibility for smaller research labs and hobbyists without the need for GPU clusters or excessive walltime.

• Component Impact Analysis (Section 5 ) - We conduct thorough ablations of BTR without each component, investigating performance and other measures. We discover that BTR widens action gaps (reducing the effects of approximation errors), is robust to observation noise, and reduces neuron dormancy and weight matrix norm (shown to improve plasticity throughout training).

## 2 Background

Before describing BTR’s extensions, we outline standard RL mathematics, how DQN is implemented, and Rainbow DQN’s extensions.

### 2.1 RL Problem Formulation

We adopt the standard formulation of RL ( Sutton & Barto, 2018 ) , described as a Markov Decision Process (MDP) defined by the tuple ( 𝒮 , 𝒜 , 𝒫 , ℛ ) \mathcal{(S,A,P,R)} , where 𝒮 \mathcal{S} is the set of states, 𝒜 \mathcal{A} is the set of actions, 𝒫 : 𝒮 × 𝒜 → Δ ⁡ ( 𝒮 ) \mathcal{P:S\times A\rightarrow}\Delta\mathcal{(S)} is the stochastic transition function, and ℛ : 𝒮 × 𝒜 → ℝ \mathcal{R:S\times A\rightarrow}\mathbb{R} is the reward function. The agent’s objective is to learn a policy π : S → Δ ⁡ ( 𝒜 ) \pi:S\rightarrow\Delta\mathcal{(A)} that maximizes the expected sum of discounted rewards 𝔼 π ​ [ ∑ t = 0 ∞ γ t ​ r ​ ( s t , a t ) ] \mathbb{E}_{\pi}[\sum^{\infty}_{t=0}\gamma^{t}r(s_{t},a_{t})] , where γ ∈ [ 0 , 1 ) \gamma\in[0,1) is the discount rate.

### 2.2 Deep Q-Learning (DQN)

One popular method for solving MDPs is Q-Learning ( Watkins & Dayan, 1992 ) where an agent learns to predict the expected sum of discounted future rewards for a given state-action pair. To allow agents to generalize over states and thus be applied to problems with larger state spaces, Mnih et al. (2013) successfully combined Q-Learning with neural networks. To do this, training minimizes the error between the predictions from a parameterized network Q θ \mathrm{Q}_{\theta} and a target defined by r t + γ max a ∈ A Q θ ′ ( s t + 1 , a ) , r_{t}+\gamma\max_{a\in A}\mathrm{Q_{\theta^{{}^{\prime}}}}(s_{t+1},a)\;, (1) where Q θ ′ Q_{\theta^{{}^{\prime}}} is an earlier version of the network referred to as the target network, which is periodically updated from the online network Q θ \mathrm{Q}_{\theta} . The data used to perform updates is gathered by sampling from an Experience Replay Buffer ( Lin, 1992 ) , which stores states, actions, rewards, and next states experienced by the agent while interacting with the environment. To effectively explore the environment, ϵ \epsilon -greedy exploration is used, where each observation has a ϵ 1 \frac{\epsilon}{1} probability of choosing a random action.

### 2.3 Rainbow DQN and Improvements to DQN

In collecting 6 different improvements to DQN, Rainbow DQN ( Hessel et al., 2018 ) proved cumulatively that these improvements could achieve a greater performance than any individually. We briefly explain the individual improvements, ordered by performance impact, most of which are preserved within BTR (see Table 1 ). For more detail, we refer readers to the extension’s respective papers:

1. Prioritized Experience Replay - To select training examples, DQN sampled uniformly from an Experience Replay Buffer, assuming that all examples are equally important to train with. Schaul et al. (2015) proposed sampling training examples proportionally to their last seen absolute temporal difference error, increasing training on samples for which the network most inaccurately predicts their future rewards.

2. N-Step - Q-learning utilizes bootstrapping to minimize the difference between the predicted value and the resultant reward plus the maximum value of the next state (Eq. 1 ). N-step ( Sutton et al., 1998 ) reduces the reliance on this bootstrapped next value by considering the next n n rewards and the observation in n n timesteps (Rainbow DQN used n = 3 n=3 ).

3. Distributional RL - Due to the stochastic nature of RL environments and agent policies, Bellemare et al. (2017) proposed learning the return distribution rather than scalar expectation. This was done through modeling the return distributions using probability masses and the Kullbeck-Leibler divergence loss function.

4. Noisy Networks - Agents can often insufficiently explore their environment resulting in sub-optimal policies. Fortunato et al. (2018) added parametric noise to the network weights, causing the model’s outputs to be randomly perturbed, increasing exploration during training, particularly for states where the agent has less confidence.

5. Dueling DQN - The agent’s Q-value can be rewritten as the sum of state-value and advantage ( Q ⁡ ( s , a ) = V ⁡ ( s ) + A ⁡ ( s , a ) Q(s,a)=V(s)+A(s,a) ). Looking to improve action generalization, Wang et al. (2016) split the hidden layers into two separate streams for the value and advantage, recombining them with Q ⁡ ( s , a ) = V ⁡ ( s ) + ( A ⁡ ( s , a ) − 1 | 𝒜 | ​ ∑ a ′ A ⁡ ( s , a ′ ) ) Q(s,a)=V(s)+(A(s,a)-\frac{1}{|\mathcal{A}|}\sum_{a^{\prime}}A(s,a^{\prime})) .

6. Double DQN - Selecting a target Q-value with the maximum Q-value from the next observation (Eq. 1 ) can frequently cause overestimation, negatively affecting agent performance. To reduce this overestimation, Van Hasselt et al. (2016) propose utilizing the online network rather than the target network to select the next action when forming targets, defined as: r t + γ Q θ ′ ( s t + 1 , arg ​ max a ∈ A Q θ ( s t + 1 , a ) ) . r_{t}+\gamma\mathrm{Q_{\theta^{{}^{\prime}}}}(s_{t+1},\argmax_{a\in A}\mathrm{Q_{\theta}}(s_{t+1},a))\;. (2)

## 3 Beyond the Rainbow - Extensions and Improvements

Building on Rainbow DQN ( Hessel et al., 2018 ) , BTR includes 6 more improvements undiscovered in 2018. 2 2 2 After the completion of our work, we additionally found Layer Normalization applied after the stem of each residual block and between dense layers to be beneficial (see Appendix H for a discussion) Additionally, as hyperparameters are critical to agent performance, Section 3.2 discusses key hyperparameters and our choices. In the appendices, we include a table of hyperparameters, a figure of the network architecture and the agent’s loss function (Appendices D.2 , E and E.2 ). Finally, the source code using Gymnasium ( Towers et al., 2024 ) is included within the supplementary material to help future work build upon or utilize BTR.

### 3.1 Extensions

Impala Architecture + Adaptive Maxpooling - Espeholt et al. (2018) proposed a convolutional residual neural network architecture based on He et al. (2016) , featuring three residual blocks 3 3 3 The network architecture is referred to as Impala due to the accompanying training algorithm IMPALA proposed in Espeholt et al. (2018) , substantially increasing performance over DQN’s three-layer convolutional network. Following Cobbe et al. (2020) , we scale the width of the convolutional layers by 2 to improve performance. We include an additional 6x6 adaptive max pooling layer after the convolutional layers ( Schmidt & Schmied, 2021 ) , which was found to speed up learning and support different input resolutions. The adaptive maxpooling is identical to a standard 2D maxpooling layer, but can be used with any input resolution as it automatically adjusts the stride and kernel size to fit a specified output size.

Spectral Normalization (SN) - To help stabilize the training of discriminators in Generative Adversarial Networks (GANs), Miyato et al. (2018) proposed Spectral Normalization to help control the Lipschitz constant of convolutional layers. SN works to normalize the weight matrices of each layer in the network by their largest singular value, ensuring that the transformation applied by the weights does not distort the input data excessively, which can lead to instability during training. Bjorck et al. (2021) and Gogianu et al. (2021) found that SN could improve performance in RL, especially for larger networks and Schmidt & Schmied (2021) found SN reduced the number of updates required before initial progress is made.

Implicit Quantile Networks (IQN) - Dabney et al. (2018) improved upon Bellemare et al. (2017) , used in Rainbow DQN, learning the return distribution over the probability space rather than probability distribution over return values. This removes the limit on the range of Q-values that can be expressed, and enables learning the expected return at every probability.

Munchausen RL - Bootstrapping is a core aspect of RL; used to calculate target values (Eq. 1 ) with most algorithms using the reward, r t r_{t} , and the optimal Q-value of the next state, Q ∗ Q^{*} . However, since in practice the optimal policy is not known, the current policy π \pi is used. Munchausen RL ( Vieillard et al., 2020 ) leverages an additional estimate in the bootstrapping process by adding the scaled-log policy to the loss function (Eq. 3 where α ∈ [ 0 , 1 ] \alpha\in[0,1] is a scaling factor, σ \sigma is the softmax \operatorname{softmax} function, and τ \tau is the softmax temperature). This assumes a stochastic policy, therefore DQN is converted to Soft-DQN with ​ π θ ′ = σ ⁡ ( Q θ ′ τ ) \text{with }{\color[rgb]{0,0,1}\pi_{\theta^{\prime}}=\sigma(\frac{Q_{\theta^{\prime}}}{\tau})} . As Munchausen does not use argmax \operatorname{argmax} over the next state, Double DQN is obsolete. Munchausen RL’s update rule is

Q θ ​ ( s t , a t ) = r t + α ​ τ ​ ln ⁡ π θ ′ ​ ( a t | s t ) + γ ∑ a ′ ∈ A π θ ′ ( a ′ | s t + 1 ) ( Q θ ′ ( s t + 1 , a ′ ) − τ ln ( π θ ′ ( a ′ | s t + 1 ) ) . \begin{split}Q_{\theta}(s_{t},a_{t})=r_{t}+{\color[rgb]{1,0,0}\alpha\tau\ln\pi_{\theta^{\prime}}(a_{t}|s_{t})}\,+\\ \gamma\sum_{a^{\prime}\in A}{\pi_{\theta^{\prime}}(a^{\prime}|s_{t+1})(Q_{\theta^{\prime}}(s_{t+1},a^{\prime})-{\color[rgb]{0,0,1}\tau\ln(\pi_{\theta^{\prime}}(a^{\prime}|s_{t+1}))}}\;.\end{split} (3)