O FFLINE Q-L EARNING ON D IVERSE M ULTI -T ASK D ATA B OTH S CALES A ND G ENERALIZES
Anonymous authors
Paper under double-blind review
A BSTRACT
The potential of offline reinforcement learning (RL) is that high-capacity models trained on large, heterogeneous datasets can lead to agents that generalize broadly, analogously to similar advances in vision and NLP.
However, recent works argue that offline RL methods encounter unique challenges to scaling up model capacity.
Drawing on the learnings from these works, we re-examine previous design choices and find that with appropriate choices: ResNets, cross-entropy based distributional backups, and feature normalization , offline Q-learning algorithms exhibit strong performance that scales with model capacity.
Using multi-task Atari as a testbed for scaling and generalization, we train a single policy on 40 games with near-human performance using up-to 80 million parameter networks, finding that model performance scales favorably with capacity.
In contrast to prior work, we extrapolate beyond dataset performance even when trained entirely on a large (400M transitions) but highly suboptimal dataset (51% human-level performance).
Compared to return-conditioned supervised approaches, offline Q-learning scales similarly with model capacity and has better performance, especially when the dataset is suboptimal.
Finally, we show that offline Q-learning with a diverse dataset is sufficient to learn powerful representations that facilitate rapid transfer to novel games and fast online learning on new variations of a training game, improving over existing state-of-the-art representation learning approaches.
I NTRODUCTION
High-capacity neural networks trained on large, diverse datasets have led to remarkable models that can solve numerous tasks, rapidly adapt to new tasks, and produce general-purpose representations in
NLP and vision (Brown et al., 2020; He et al., 2021).
The promise of offline RL is to leverage these advances to produce polices with broad generalization, emergent capabilities, and performance that exceeds the capabilities demonstrated in the training dataset.
Thus far, the only offline RL approaches that demonstrate broadly generalizing policies and transferable representations are heavily-based on supervised learning (Reed et al., 2022; Lee et al., 2022).
However, these approaches are likely to perform poorly when the dataset does not contain expert trajectories (Kumar et al., 2021b).
Offline Q-learning performs well across dataset compositions in a variety of simulated (Gulcehre et al., 2020;
Fu et al., 2020) and real-world domains (Chebotar et al., 2021;
Soares et al., 2021), however, these are largely centered around small-scale, single-task problems where broad generalization and learning general-purpose representations is not expected.
Scaling these methods up to high-capcity models on large, diverse datasets is the critical challenge.
Prior works hint at the difficulties: on small-scale, single-task deep RL benchmarks, scaling model capacity can lead to instabilities or degrade performance (Van Hasselt et al., 2018; Sinha et al., 2020; Ota et al., 2021) explaining why decade-old tiny 3-layer CNN architectures (Mnih et al., 2013) are still prevalent.
Moreover, works that have scaled architectures to millions of parameters (Espeholt et al., 2018;
Teh et al., 2017;
Vinyals et al., 2019; Schrittwieser et al., 2021) typically focus on online learning and employ many sophisticated techniques to stabilize learning, such as supervised auxiliary losses, distillation, and pre-training.
Thus, it is unclear whether offline Q-learning can be scaled to high-capacity models trained on a large, diverse dataset.
In this paper, we demonstrate that with careful design decisions, offline Q-learning can scale to highcapacity models trained on large, diverse datasets from many tasks, leading to policies that not only generalize broadly, but also learn representations that effectively transfer to new downstream tasks and
[Figure]
exceed the performance in the training dataset.
Crucially, we make three modifications motivated by prior work in deep learning and offline RL.
First, we find that a modified ResNet architecture (He et al., 2016) substantially outperforms typical deep RL architectures and follows a power-law relationship between model capacity and performance, unlike common alternatives.
Second, a discretized representation of the return distribution with a distributional cross-entropy loss (Bellemare et al., 2017) substantially improves performance compared to standard Q-learning, that utilizes mean squared error.
Finally, feature normalization on the intermediate feature representations stabilizes training and prevents feature co-adaptation (Kumar et al., 2021a).
To systematically evaluate the impact of these changes on scaling and generalization, we train a single policy to play 40 Atari games (Bellemare et al., 2013; Agarwal et al., 2020), similarly to
Lee et al.
and evaluate performance when the training dataset contains expert trajectories and when the data is sub-optimal.
This problem is especially challenging because of the diversity of games with their own unique dynamics, reward, visuals, and agent embodiments.
Furthermore, the sub-optimal data setting requires the learning algorithm to “stitch together” useful segments of sub-optimal trajectories to perform well.
To investigate generalization of learned representations, we evaluate offline fine-tuning to never-before-seen games and fast online adaptation on new variants of training games (Section 5.2).
With our modifications,
• Offline Q-learning learns policies that attain more than 100% human-level performance on most of these games, about 2x better than prior supervised learning (SL) approaches for learning from sub-optimal offline data (51% human-level performance).
• Akin to scaling laws in SL (Kaplan et al., 2020), offline Q-learning performance scales favorably with model capacity (Figure 6).
• Representations learned by offline Q-learning give rise to more than 80% better performance when fine-tuning on new games compared to representations learned by state-of-the-art return-conditioned supervised (
Lee et
al.
) and self-supervised methods
(He et
al.
Oord
et al.
By scaling Q-learning, we realize the promise of offline RL: learning policies that broadly generalize and exceed the capabilities demonstrated in the training dataset.
We hope that this work encourages large-scale offline RL applications, especially in domains with large sub-optimal datasets.
R ELATED W ORK
Prior works have sought to train a single generalist policy to play multiple Atari games simultaneously from environment interactions, either using off-policy RL with online data collection (Espeholt et al., 2018; Hessel et al., 2019a; Song et al., 2019), or policy distillation (Teh et al., 2017; Rusu et al., 2015) from single-task policies.
While our work also focuses on learning such a generalist multi-task policy, it investigates whether we can do so by scaling offline Q-learning on suboptimal offline data, analogous to how supervised learning can be scaled to large, diverse datasets.
Furthermore, prior attempts to apply transfer learning using RL-learned policies in ALE (Rusu et al., 2015; Parisotto [Figure] n
[Figure]
M
[Figure] i a [Figure] et al., 2015; Mittel & Sowmya Munukutla, 2019) are restricted to a dozen games that tend to be similar and generally require an “expert”, instead of learning how to play all games concurrently.
Closely related to our work, recent work train Transformers (Vaswani et al., 2017) on purely offline data for learning such a generalist policy using supervised learning (SL) approaches, namely, behavioral cloning (Reed et al., 2022) or return-conditioned behavioral cloning (Lee et al., 2022).
While these works focus on large datasets containing expert or near-human performance trajectories, our work focuses on the regime when we only have access to highly diverse but sub-optimal datasets.
We find that these SL approaches perform poorly with such datasets, while offline Q-learning is able to substantially extrapolate beyond dataset performance (Figure 2).
Even with near-optimal data, we observe that scaling up offline Q-learning outperforms SL approaches with 200 million parameters using as few as half the number of network parameters (Figure 6).
There has been a recent surge of offline RL algorithms that focus on mitigating distribution shift in single task settings (Fujimoto et al., 2018; Kumar et al., 2019; Liu et al., 2020; Wu et al., 2019;
Fujimoto & Gu, 2021; Siegel et al., 2020; Peng et al., 2019; Nair et al., 2020; Liu et al., 2019;
Swaminathan & Joachims, 2015; Nachum et al., 2019; Kumar et al., 2020; Kostrikov et al., 2021;
Kidambi et al., 2020; Yu et al., 2020b; 2021).
Complementary to such work, our work investigates scaling offline RL on the more diverse and challenging multi-task Atari setting with data fromdifferent games (Agarwal et al., 2020; Lee et al., 2022).
To do so, we use CQL (Kumar et al., 2020), due to its simplicity as well as its efficacy on offline RL datasets with high-dimensional observations.
P RELIMINARIES AND P ROBLEM S ETUP
We consider sequential-decision making problems (Sutton & Barto, 1998) where on each timestep, an agent observes a state s , produces an action a , and receives a reward r .
The goal of a learning algorithm is to maximize the sum of discounted rewards.
Our approach is based on conservative
Q-learning (CQL) (Kumar et al., 2020), an offline Q-learning algorithm.
CQL uses a sum of two loss functions to combat value overestimation on unseen actions: (i) standard TD-error that enforces
Bellman consistency, and (ii) a regularizer that minimizes the Q-values for unseen actions at a given state, while maximizing the Q-value at the dataset action to counteract excessive underestimation.
Denoting Q θ ( s , a ) as the learned Q-function, the training objective for CQL is given by:
[Equation]
where α is the regularizer weight, which we fix to α = 0 .
based on preliminary experiments unless noted otherwise.
Kumar et al.
(2020) utilized a distributional TDError ( θ ; D ) from C51 (Bellemare et al., 2017), whereas (Kumar et al., 2021a) showed that similar results could be attained with the standard mean-squared TD-error.
Lee et al.
(2022) use the distributional formulation of CQL and found that it underperforms alternatives and performance does not improve with model capacity.
In general, there is no consensus on which formulation of TD-error must be utilized in Equation 1, and we will study this choice in our scaling experiments.
Problem setup.
Our goal is to learn a single policy that is effective at multiple Atari games and can be fine-tuned to new games.
For training, we utilize the set of 40 Atari games used by Lee et al.
(2022), and for each game, we utilize the experience collected in the DQN-Replay dataset (Agarwal et al., 2020) as our offline dataset.
We consider two different dataset compositions:
Sub-optimal dataset consisting of the initial 20% of the trajectories (10M transitions) from DQN-Replay for each game, containing 400 million transitions overall with average human-normalized interquartile-mean (IQM)
(Agarwal et al., 2021) score of 51%.
Since this dataset does not contain optimal trajectories, we do not expect methods that simply copy behaviors in this dataset to perform well.
On the other hand, we would expect methods that can combine useful segments of sub-optimal trajectories to perform well.
Near-optimal dataset, used by Lee et al.
(2022), consisting of all the experience (50M transitions) encountered during training of a DQN agent including human-level trajectories, containing 2 billion transitions with average human-normalized IQM score of 93.5%.
Evaluation .
We evaluate our method in a variety of settings as we discuss in our experiments in
Section 5.
Due to excessive computational requirements of running huge models, we are only able to run our main experiments with one seed.
Prior work (Lee et al., 2022) that also studied offline multi-game Atari evaluated models with only one seed.
That said, to ensure that our evaluations are reliable, for reporting performance, we follow the recommendations by Agarwal et al.
Specifically, we report interquartile mean (IQM) normalized scores, which is the average scores across middle 50% of the games, as well as performance profiles for qualitative summarization.
O UR A PPROACH FOR S CALING O FFLINE RL
In this section, we describe the critical modifications required to make CQL effective in learning highly-expressive policies from large, heterogeneous datasets.
Parameterization of Q-values and TD error.
In the single game setting, both mean-squared TD error and distributional TD error perform comparably online (Agarwal et al., 2021) and offline (Kumar et al., 2020; 2021a).
In contrast, we observed, perhaps surprisingly, that mean-squared TD error does not scale well, and performs much worse than using a categorical distributional representation of return values (Bellemare et al., 2017) when we train on many Atari games.
We hypothesize that this is because even with reward clipping, Q-values for different games often span different ranges, and training a single network with shared parameters to accurately predict all of them presents challenges pertaining to gradient interference along different games (Hessel et al., 2019b; Yu et al., 2020a).
While prior works have proposed to use adaptive normalization schemes (Hessel et al., 2019b; Kurin et al., 2022), preliminary experiments with these approaches were not effective to close the gap
Q-function architecture.
Since large neural networks has been crucial for scaling to large, diverse datasets in NLP and vision (e.g., Tan & Le, 2019; Brown et al., 2020; Kaplan et al., 2020)), we explore [Table] [Figure] [Table]
[Figure]
[Table]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
using bigger architectures for scaling offline Q-learning.
We use standard feature extractor backbones from vision, namely, the Impala-CNN architectures (Espeholt et al., 2018) that are fairly standard in deep RL and ResNet 34 , 50 and 101 models from the ResNet family (He et al., 2016).
We make modifications to these networks following recommendations from prior work (Anonymous, 2022): we utilize group normalization instead of batch normalization in ResNets, and utilize point-wise multiplication with a learned spatial embedding when converting the output feature map of the vision backbone into a flattened vector which is to be fed into the feed-forward part of the Q-function.
To handle the multi-task setting, we use a multi-headed architecture where the Q-network outputs values for each game separately.
The architecture uses a shared encoder and feedforward layers with separate linear projection layers for each game (Figure 3).
The training objective (Eq. 1) is computed using the Q-values for the game that the transition originates from.
In principle, explicitly injecting the task-identifier may be unnecessary and its impact could be investigated in future work.
Feature Normalization via DR3 (Kumar et al., 2021a).
While the previous modifications lead to significant improvements over naïve CQL, our preliminary experiments on a subset of games did not attain good performance.
In the single-task setting, Kumar et al.
(2021a) proposes a regularizer that stabilizes training and allows the network to better use capacity, however, it introduces an additional hyperparameter to tune.
Motivated by this approach, we regularize the magnitude of the learned features of the observation by introducing a “normalization” layer in the Q-network.
This layer forces the learned features to have an ℓ 2 norm of 1 by construction, and we found that this this speeds up learning, resulting in better performance.
We present an ablation study analyzing this choice in Table 2.
We found this sufficient to achieve strong performance, however, we leave exploring alternative feature normalization schemes to future work.
To summarize, the primary modifications that enable us to scale CQL are: (1) use of large
ResNets with learned spatial embeddings and group normalization, (2) use of a distributional representation of return values and cross-entropy loss for training (i.e., C51 (Bellemare et al., 2017)), and (3) feature normalization at intermediate layers to prevent feature co-adaptation, motivated by Kumar et al.
(2021a).
For brevity, we call our approach Scaled Q-learning .
E XPERIMENTAL E VALUATION
In our experiments, we study how our approach, scaled Q-learning, can simultaneously learn from sub-optimal and optimal data collected from 40 different Atari games.
We compare the resulting multi-task policies to behavior cloning (BC) with same architecture as scaled QL, and the prior state-of-the-art method based on decision transformers (DT) (Chen et al., 2021), which utilize returnconditioned supervised learning with large transformers (Lee et al., 2022), and have been previously proposed for addressing this task.
We also study the efficacy of the multi-task initialization produced by scaled Q-learning in facilitating rapid transfer to new games via both offline and online fine-tuning, in comparison to state-of-the-art self-supervised representation learning methods and other prior approaches.
Our goal is to answer the following questions: (1) How do our proposed design decisions impact performance scaling with high-capacity models?
, (2) Can scaled QL more effectively leverage higher model capacity compared to naïve instantiations of Q-learning?
, (3) Do the representations learned by scaled QL transfer to new games?
We will answer these questions in detail through multiple experiments in the coming sections, but we will first summarize our main results below.
Main empirical findings.
Our main results are summarized in Figures 2 and 5.
These figures show the performance of scaled QL, multi-game decision transformers (Lee et al., 2022) (marked as “DT”), a prior method based on supervised learning via return conditioning, and standard behavioral cloning baselines (marked as “BC”) in the two settings discussed previously, where we must learn from: (i) near optimal data, and (ii) sub-optimal data obtained from the initial 20% segment of the replay buffer (see Section 3 for problem setup).
See Figure 4 for a direct comparison between DT and BC.
[Figure]
In the more challenging sub-optimal data setting, scaled QL attains a performance of 77.8% IQM human-normalized score, although trajectories in the sub-optimal training dataset only attain 51% IQM human-normalized score.
Scaled QL also outperforms the prior DT approach by 2.
times on this dataset, even though the DT model has more than twice as many parameters and uses data augmentation, compared to scaled QL.
In the 2 nd setting with near-optimal data, where the training dataset already contains expert trajectories, scaled QL with 80M parameters still outperforms the DT approach with 200M parameters, although the gap in performance is small (3% in IQM performance, and 20% on median performance).
Overall, these results show that scaled QL is an effective approach for learning from large multi-task datasets, for a variety of data compositions including sub-optimal datasets, where we must stitch useful segments of suboptimal trajectories to perform well, and near-optimal datasets, where we should attempt to mimic the best behavior in the offline dataset.
To the best of our knowledge
[Figure]
these results represent the largest performance improvement over the average performance in the offline dataset on such a challenging problem.
We will now present experiments that show that offline Q-learning scales and generalizes
[Figure]
5. D OES O FFLINE Q-L EARNING S CALE F AVORABLY ?
One of the primary goals of this paper was to understand if scaled Q-learning is able to leverage the benefit of higher capacity architectures.
Recently, Lee et al.
(2022) found that the performance of
CQL with the IMPALA architecture does not improve with larger model sizes and may even degrade with larger model sizes.
To verify if scaled Q-learning can address this limitation, we compare our value-based offline RL approach with a variety of model families: (a) IMPALA family (Espeholt et al., 2018): three IMPALA models with varying widths ( 4 , 8 , 16 ) whose performance numbers are taken directly from Lee et al.
(2022) (and was consistent with our preliminary experiments), (b)
[Figure]
ResNet 34, 50, 101 and 152 from the ResNet family, modified to include group normalization and learned spatial embeddings.
These architectures include both small and large networks, spanning a wide range from 1M to 100M parameters.
As a point of reference, we use the scaling trends of the multi-game decision transformer and BC transformer approaches from Lee et al.
Observe in Figure 6 that the performance of scaled Q-learning improves as the underlying Q-function model size grows.
Even though the standard mean-squared error formulation of TD error results in worse absolute performance than C51 (blue vs orange), for both of these versions, the performance of scaled Q-learning increases as the models become larger.
This result indicates that value-based offline RL methods can scale favorably, and give rise to better results, but this requires carefully picking a model family.
This also explains the findings from Lee et al.
(2022): while this prior work observed that CQL with IMPALA scaled poorly as model size increases, they also observed that the performance of return-conditioned RL instantiated with IMPALA architectures also degraded with higher model sizes.
Combined with the results in Figure 6 above, this suggests that poor scaling properties of offline RL can largely be attributed to the choice of IMPALA architectures, which may not work well in general even for supervised learning methods (like return-conditioned BC).
5. C AN O FFLINE RL L EARN U SEFUL I NITIALIZATIONS THAT E NABLE F INE -T UNING ?
Next, we study how multi-task training on multiple games via scaled QL can learn general-purpose representations that can enable rapid fine-tuning to new games.
We study this question in two scenarios: fine-tuning to a new game via offline RL with a small amount of held-out data (1% uniformly subsampled datasets from DQN-Replay (Agarwal et al., 2020)), and finetuning to a new game mode via sample-efficient online RL initialized from our multi-game offline Q-function.
For finetuning, we transfer the weights from the visual encoder and reinitialize the downstream feedforward component (Figure 1).
For both of these scenarios, we utilize a ResNet101 Q-function trained via the methodology in Section 4, using C51 and feature normalization.
Scenario 1 (Offline fine-tuning) :
First, we present the results for fine-tuning in an offline setting: following the protocol from Lee et al.
(2022), we use the pre-trained representations to rapidly learn a policy for a novel game using limited offline data (1% of the experience of an online DQN run).
In
Figure 7, we present our results for offline fine-tuning on 5 games from Lee et al.
A LIEN ,
M S P ACMAN , S PACE I NVADERS , S TAR G UNNER and P ONG , alongside the prior approach based on decision transformers (“DT (pre-trained)”), and fine-tuning using pre-trained representations learned from state-of-the-art self-supervised representation learning methods such as contrastive predictive coding (CPC) (Oord et al., 2018) and masked autoencoders (MAE)
(He et al., 2021).
For CPC performance, we use the baseline reported in Lee et al.
MAE is a more recent self-supervised approach that we find generally outperformed CPC in this comparison.
For MAE, we first pretrained a vision transformer (ViT-Base)
(Dosovitskiy et al., 2020) encoder with 80M parameters trained via a reconstruction loss on observations from multi-game Atari dataset and freeze the encoder weights as done in prior work (Xiao et al., 2022).
Then, with this frozen visual encoder, we used the same feed forward architecture, Q-function parameterization, and training objective (CQL with C51) as scaled QL to finetune the MAE network.
We also compare to baseline methods that do not utilize any multi-game pre-training (DT (scratch) and Scaled QL (scratch)).
[Figure]
Results.
Observe in Figure 7 that multi-game pre-training via scaled QL leads to the best fine-tuning performance and improves over prior methods, including decision transformers trained from scratch.
Importantly, we observe positive transfer to new games via scaled QL.
Prior works (Badia et al., 2020) running multi-game Atari (primarily in the online setting) have generally observed negative transfer across Atari games.
We show for the first time that pre-trained representations from Q-learning enable positive transfer to novel games that significantly outperforms return-conditioned supervised learning methods and dedicated representation learning approaches.
Scenario 2 (Online fine-tuning) :
Next, we study the efficacy of the learned representations in enabling online fine-tuning.
While deep RL agents on ALE are typically trained on default game modes (referred to as m 0 d 0 ), we utilize new variants of the ALE games designed to be challenging for humans (Machado et al., 2018; Farebrother et al., 2018) for online-finetuning.
We investigate whether multi-task training on the 40 default game variants can enable fast online adaptation to these never-before-seen variants.
In contrast to offline fine-tuning (Scenario 1), this setting tests whether scaled QL can also provide a good initialization for online data collection and learning, for closely related but different tasks.
Following Farebrother et al.
(2018), we use the same variants investigated in this prior work: B REAKOUT , H ERO , and F REEWAY , which we visualize in Figure 8 (left).
To disentangle the performance gains from multi-game pre-training and the choice of Q-function architecture, we compare to a baseline approach (“scaled QL (scratch)”) that utilizes an identical
Q-function architecture as pre-trained scaled QL, but starts from a random initialization.
As before, we also evaluate fine-tuning performance using the representations obtained via masked auto-encoder pre-training (He et al., 2021; Xiao et al., 2022).
We also compare to a single-game DQN performance attained after training for 50M steps, 16 × more transitions than what is allowed for scaled QL, as reported by Farebrother et al.
Results .
Observe in Figure 8 that fine-tuning from the multi-task initialization learned by scaled
QL significantly outperforms training from scratch as well as the single-game DQN run trained with 16x more data.
Fine-tuning with the frozen representations learned by MAE performs poorly, which we hypothesize is due to differences in game dynamics and subtle changes in observations, which must be accurately accounted for in order to learn optimal behavior (Dean et al., 2022).
Our results confirm that offline Q-learning can both effectively benefit from higher-capacity models and learn multi-task initializations that enable sample-efficient transfer to new games.
5. A BLATION S TUDIES
Finally, in this section we perform controlled ablation studies to understand how crucial the design decisions introduced in Section 4 are for the success of scaled Q-learning.
In particular, we will attempt to understand the benefits of using C51 and feature normalization.
MSE vs
C
: We ran scaled Q-learning with identical network architectures (ResNet 50 and ResNet 101), with both the conventional squared error formulation of TD error, and compare it to C51, which our main results utilize.
Observe in Table 1 that C51 leads to much better performance for both
[Table] [Table]
ResNet 50 and ResNet 101 models.
The boost in performance is the largest for ResNet 101, where
C51 improves by over 39% as measured by median human-normalized score.
This observation is surprising since prior work (Agarwal et al., 2021) has shown that C51 performs on par with standard
DQN with an Adam optimizer
[Table]
which all of our results use.
One hypothesis is that this could be the case as TD gradients would depend on the scale of the reward function
[Table]
and hence some games would likely exhibit a stronger contribution in the gradient.
This is despite the fact that our implementation of MSE TD-error already attempts to correct for this issue by applying the unitary scaling technique from (Kurin et al., 2022) to standardize reward scales across games.
That said, we still observe that
C51 performs significantly better.
Importance of feature normalization: We ran small-scale experiments with and without feature normalization (Section 4).
In these experiments, we consider a multi-game setting with only 6 games:
A STERIX , B REAKOUT , P ONG , S PACE I NVADERS , S EAQUEST , and we train with the initial 20% data for each game.
We report aggregated median human-normalized score across the 6 games in Tablefor three different network architectures (ResNet 34, ResNet 50 and ResNet 101).
Observe that the addition of feature normalization significantly improves performance for all the models.
Motivated by this initial empirical finding, we used feature normalization in all of our main experiments.
To summarize, these ablation studies validate the efficacy of the two key design decisions introduced in this paper.
However, there are several avenues for future investigation: 1) it is unclear if C51 works better because of the distributional formulation or the categorical representation and experiments with other distributional formulations could answer this question, 2) we did not extensively try alternate feature normalization schemes which may improve results.
D ISCUSSION
This work shows, for the first time (to the best of our knowledge), that offline Q-learning can scale to high-capacity models trained on large, diverse datasets.
As we hoped, by scaling up model capacity, we unlocked analogous trends to those observed in vision and NLP.
We found that scaled Q-learning trains policies that exceed the average dataset performance and prior methods, especially when the dataset does not already contain expert trajectories.
Furthermore, by training a large-capacity model on a diverse set of tasks, we show that Q-learning alone is sufficient to recover general-purpose representations that enable rapid learning of novel tasks.
Although we detailed an approach that is sufficient to scale Q-learning, this is by no means optimal.
The scale of the experiments limited the number of alternatives we could explore, and we expect that future work will greatly improve performance.
Given the strong performance of transformers, we suspect that offline Q-learning with a transformer architecture is a promising future direction.
For example, contrary to DT (Lee et al., 2022), we did not use data augmentation in our experiments, which we believe can provide significant benefits.
While we did a preliminary attempt to perform online fine-tuning on an entirely new game (S PACE I NVADERS ), we found that this did not work well for any of the pretrained representations (see Figure A.1). Addressing this is an important direction for future work.
We speculate that this challenge is related to designing methods for learning better exploration from offline data, which is not required for offline fine-tuning.
Another important avenue for future work is to scale offline Q-learning on other RL domains such as robotic navigation, manipulation, locomotion, education, etc. This would require building large-scale tasks, and we believe that scaled QL would provide for a good starting point for scaling in these domains.
Appendices A A DDITIONAL R ESULTS A. A DDITIONAL R ESULTS FROM THE
P APER
[Figure] [Figure]
A. R ESULTS FOR S CALING D ISCRETE -BCQ
To implement discrete BCQ
[Figure]
we followed the official implementation from Fujimoto et al. We first trained a model of the behavior policy, (cid:98) π β ( a | s ) , using an architecture identical to that of the Q-function, using negative log-likelihood. Then, following Fujimoto et al.
(2019), we updated the
Bellman backup to only perform the maximization over actions that attain a high likelihood under the probabilities learned by the behavior policy, as shown below:
[Equation]
where τ is a hyperparameter.
To tune the value of τ , we ran a preliminary initial sweep over τ = { 0 . 05 , 0 . 1 , 0 . 3 } .
When using C51 in our setup, we had to use a smaller CQL α of 0.05 (instead of 0.1 for the MSE setting from Kumar et al.
(2021a)), possibly because a discrete representation of Q-values used by C51 is less prone to overestimation.
Therefore, in the case of discrete-BCQ, we chose to perform an initial sweep over τ values that were smaller than or equal to (i.e., less conservative) the value of τ = 0 .
3 used in Fujimoto et al.
Since BCQ requires an additional policy network, it imposes a substantial memory overhead and as such, we performed a sweep for initial 20 iterations to pick the best τ .
We found that in these initial
[Figure]
experiments, τ = 0 .
05 performed significantly worse, but τ = 0 .
1 and τ = 0 .
3 performed similarly.
So, we utilized τ = 0 .
3 for reporting these results.
We ran these scaling experiments with ResNet 34 and ResNet 50 in the six-game setting and report human-normalized IQM performance after 100 epochs = 6.25M gradient steps in Figure A.3.
We also present the results for CQL on the side for comparisons.
Observe that we find favorable scaling trends for BCQ: average performance over all games increases as the network size increases, indicating that other offline RL algorithms such as BCQ can scale as we increase network capacity.
A. A BLATION FOR B ACKBONE A RCHITECTURE
In this section, we present some results ablating the choice of the backbone architecture.
For this ablation, we ablate the choice of the spatial embedding while keeping group normalization fixed in both cases.
We perform this study for the 40-game setting, but due to time and compute constraints, we are only able to present preliminary results, 60 epochs into training, which is at least 3 times fewer steps than our results in the main paper.
Observe that using the learned spatial embedding results in better performance, and improves in 33 out of 40 games compared to not using the learned embeddings.
[Table]
Num games where backbone is better
[Table]
Regarding the choice of group normalization vs batch normalization, note that we have been operating in a setting where the size of the batch per device / core is only 4.
Particularly, we use Cloud TPU v accelerators with 64 / 128 cores, and bigger batch sizes than 4 do not fit in memory, especially for larger-capacity ResNets.
This means that if we utilized batch normalization, we would be computing batch statistics over only 4 elements, which is known to be unstable even for standard computer vision tasks, for example, see Figure 1 in Wu & He (2018).
A. R ESULTS FOR S CALED QL W ITHOUT P ESSIMISM
In Table A.2, we present the results of running scaled Q-learning with no conservatism, i.e., by setting the value of α in Equation 1 to 0.0 in the six game setting. We utilize the entire DQN-replay dataset (Agarwal et al., 2020) for each of these six games that would be present in the full 40-game dataset, to preserve the per-game dataset diversity. Observe that while utilizing no conservatism does still learn, the performance of scaled QL without conservatism is notably worse than standard scaled QL. Interestingly, on A STERIX , the performance without pessimism is better than performance with pessimism, whereas, the use of pessimism in
S PACE I NVADERS and S EAQUEST leads to at least 2x improvement in performance.
[Table]
We have launched this study in the full 40-game setting, but have so far only observed quite poor performance without pessimism.
These observations are preliminary, and we will run it for longer and update the final paper with the results.
B I MPLEMENTATION D ETAILS AND H YPER - PARAMETERS
In this section, we will describe some of the implementation details behind our method and will provide implementation details for our approach, including the details of the network architectures, the details of feature normalization and the details of our training and evaluation protocol.
B. N ETWORK A RCHITECTURE
In our primary experiments, we consider variants of ResNet architectures for scaled Q-Learning.
The vision backbone in these architectures mimic the corresponding ResNet architectures from He et al.
(2016), however, we utilize group normalization (Wu & He, 2018)
(with a group size of 4) instead of batch normalization, and instead of applying global mean pooling to aggregate the outputs of the ResNet, we utilize learned spatial embeddings (Anonymous, 2022), that learn a matrix that point-wise multiplies the output feature map of the ResNet.
The output volume is then flattened to be passed as input to the feed-forward part of the network.
The feed-forward layer part of the network begins with a layer of size 2048, and then applies layer norm on the network.
After this we apply 3 feed-forward layers with hidden dimension 1024 with
ReLU activations, to obtain the representation of the image observation.
Then, we apply feature normalization to the representation, by applying a normalization layer which divides the representation of a given observation by its ℓ 2 norm.
Note that we do pass gradients through this normalization term.
Now, this representation is passed into different heads that are supposed to predict the Q-values.
The total number of heads is equal to the number of games we train on.
Each head consists of a linear layer that maps the 1024-dimensional normalized representation to a vector of K elements, where K = |A| (i.e., the size of the action space) for the standard real-valued parameterization of Q-values, and K = |A| × 51 for C51.
The network does not apply any output activation in either case, and the Q-values are treated as logits for C51. [Table]
B. D ETAILS OF C
For the main results in the paper [Table] we utilize C51.
The main hyperparameter in C51 is the size of the support set of Q-values.
Unlike the paper from Bellemare et al.
(2017) which utilizes a support set of [ − 10 , 10] [Table] we utilize a support set of [ − 20 , 20] to allow for flexibility of CQL: Applying the
CQL regularizer can underestimate or overestimate Q-values [Table] and this additional flexibility aids such scenarios.
Though, we still utilize only 51 atoms in our support set [Table] and the average dataset Q-value in our training runs is generally always smaller, around ∼ 8 − [Table]
B. T RAINING AND E VALUATION P ROTOCOLS AND H YPERPARAMETERS
We utilize the initial 20% (sub-optimal) and 100% (near-optimal) datasets from Agarwal et al.
(2020) for our experiments.
These datasets are generated from runs of standard online DQN on stochastic dynamics Atari environments that utilize sticky actions, i.e. , there is a 25% chance at every time step that the environment will execute the agents previous action again, instead of the new action commanded.
The majority of the training details are identical to a typical run of offline RL on single-game Atari.
We discuss the key differences below.
We trained our ResNet 101 network for 10 M gradient steps with a batch size of 512.
The agent hasn’t converged yet, and the performance is still improving gradually.
When training on multiple games, we utilize a stratified batch sampling scheme with a total batch size of 512 .
To obtain the batch at any given training iteration, we first sample 128 game indices from the set all games ( games in our experiments) with replacement, and then sample 4 transitions from each game.
This scheme does not necessarily produce an equal number of transitions from each game in a training batch, but it does make sure that all games are seen in expectation throughout training.
Since we utilize a larger batch size, that is 16 times larger than the standard batch size of 32 on Atari, we scale up the learning rate from 5 e − 05 to 0 . 0002 , but keep the target network update period fixed to the same value of 1 target update per 2000 gradient steps as with single-task Atari.
We also utilize n -step returns with n = 3 by default, with both our MSE and C51 runs.
Evaluation Protocol.
Even though we train on Atari datasets with sticky actions, we evaluate on
Atari environments that do not enable sticky actions following the protocol from Lee et al.
This allows us to be comparable to this prior work in all of our comparisons, without needing to re-train their model, which would have been too computationally expensive.
Following standard protocols on Atari, we evaluate a noised version of the policy with an epsilon-greedy scheme, with ε eval = 0 . 001 .
Following the protocol in Castro et al.
(2018), we compute average return over 125K training steps.
B. F INE -T UNING P ROTOCOL
For offline fine-tuning we fine-tuned the parameters of the pre-trained policy on the new domain using a batch size of 32, and identical hyperparameters as those used during pre-training.
We utilized α = 0 .
for fine-tuning, but with the default learning rate of 5 e −
(since the batch size was the default 32).
We attempted to use other CQL α values { 0 . 07 , 0 .
1 } for fine-tuning
but found that retaining the value of α = 0 .
05 for pre-training worked the best.
For reporting results, we reported the performance of the algorithm at the end of 300k gradient steps.
For online fine-tuning , we use the C51 algorithm (Bellemare et al., 2017), with n -step = 3 and all other hyperparameters from the C51 implementation in the Dopamine library (Castro et al., 2018).
We swept over two learning rates, { 1 e − 05 , 5 e − 05 } for all the methods and picked the best learning rate per-game for all the methods.
For the MAE implementation, we used the Scenic library (Dehghani et al., 2022) with the typical configuration used for ImageNet pretraining, except using 84 × 84 × 4 sized Atari observations, instead of images of size 224 × 224 × 3 .
We train the
MAE for 2 epochs on the entire multi-task offline Atari dataset and we observe that the reconstruction loss plateaus to a low value.
The “MT Impala DQN” comparison in Figures 2 & 5 is a multi-task implementation of online DQN, evaluated at 5x many gradient steps as the size of the sub-optimal dataset. This comparison is taken directly from Lee et al. To explain this baseline briefly, this baseline runs C51 in conjunction with n-step returns with n = 4 , with an IMPALA architecture that uses three blocks
C R AW T RAINING S CORES FOR DIFFERENT M ODELS
[Table]
[Table]