Large-Scale Retrieval for Reinforcement Learning
Abstract
Effective decision making involves ﬂexibly relating past experiences and relevant contextual information to a novel situation.
In deep reinforcement learning (RL), the dominant paradigm is for an agent to amortise information that helps decisionmaking into its network weights via gradient descent on training losses.
Here, we pursue an alternative approach in which agents can utilise large-scale contextsensitive database lookups to support their parametric computations.
This allows agents to directly learn in an end-to-end manner to utilise relevant information to inform their outputs.
In addition, new information can be attended to by the agent, without retraining, by simply augmenting the retrieval dataset.
We study this approach for ofﬂine RL in 9x9 Go, a challenging game for which the vast combinatorial state space privileges generalisation over direct matching to past experiences.
We leverage fast, approximate nearest neighbor techniques in order to retrieve relevant data from a set of tens of millions of expert demonstration states.
Attending to this information provides a signiﬁcant boost to prediction accuracy and game-play performance over simply using these demonstrations as training trajectories, providing a compelling demonstration of the value of large-scale retrieval in ofﬂine RL agents.
Introduction
How can reinforcement learning (RL) agents leverage relevant information to inform their decisions?
Deep RL agents have typically been represented as a monolithic parametric function, trained to gradually amortise useful information from experience by gradient descent.
This has been effective [25, 38] but is a slow means of integrating experience, with no straightforward way for an agent to incorporate new information without many additional gradient updates.
In addition, this requires increasingly massive models (see e.g. [31]) as environments become more complex – scaling is driven by the dual role of the parametric function, which must support both computation and memorisation.
Finally, this approach has a further drawback of particular relevance in RL – the only way in which previously encountered information (that is not contained in working memory) can aid decision making in a novel situation is indirectly through weight changes mediated by network losses.
There is no end-to-end means for an agent to attend to information outside of working memory to directly inform its actions.
While there has been a signiﬁcant amount of work focused on increasing the information available from previous experiences within an episode (e.g., recurrent networks, slot-based memory [19, 28]), more extensive direct use of more general forms of experience or data has been limited, although some recent works have begun to explore utilising inter-episodic information from the same agent [6, 10, 29, 34, 41].
We seek to drastically expand the scale of information that is accessible to an agent, allowing it to attend to tens of millions of pieces of information, while learning in an end-to-end manner how to use this information for decision making.
We view this as a ﬁrst step towards a vision in which an agent can ﬂexibly draw on diverse and large-scale information sources, including its own (inter-episodic) experiences along with experiences from humans and other agents.
In addition, given that retrieved information need not match the agent’s observation format, retrieval could enable agents to integrate information sources that have not typically been utilised such as videos [2], text, and third-person demonstrations.
We investigate a semi-parametric agent architecture for large-scale retrieval, in which fast and efﬁcient approximate nearest neighbor matching is used to dynamically retrieve relevant information from a dataset of experience.
We evaluate this approach in an ofﬂine RL setting for an environment with a combinatorial state space – the game of 9x9 Go with ≈ 10 38 possible games, where generalisation from past data to novel situations is challenging.
We equip our agent with a large-scale dataset of ∼ 50M Go board-state observations, ﬁnding that a retrieval-based approach utilising this data is able to consistently and signiﬁcantly outperform a strong non-retrieval baseline.
Several key advantages of this retrieval approach are worth highlighting: Instead of having to amortise all relevant information into its network weights, a retrieval-augmented network can utilise more of its capacity for computation.
In addition, this semi-parametric form allows us to update the information available to the agent at evaluation time without having to retrain it.
Strikingly, we ﬁnd that this enables improvements to agent performance without further training when games played against the evaluation opponent are added to its knowledge base.
Methods
[Table]
Before introducing the details of our method, let us ﬁrst present its high-level ingredients.
We train a model-based agent [36] that predicts future policies and values conditioned on future actions in a given state.
This semiparametric model incorporates a retrieval mechanism, which allows it to utilise information from a large-scale dataset to inform its predictions.
We train the agent in a supervised ofﬂine RL setting, which allows us to directly evaluate the degree to which retrieval improves the quality of the model predictions.
We subsequently evaluate the resulting model, augmented by MonteCarlo tree search, against a reference opponent.
This allows us to determine how these improvements translate to an effective acting policy for out-of-training-distribution situations.
To effectively improve model predictions using auxiliary information, we need 1) a scalable way to select and retrieve relevant data and 2) a robust way of leveraging that data in our model.
As with many successful attention mechanisms, we establish relevance for 1) through an inner product in an appropriate key-query space and select the top N relevant results.
Choosing the relevant key-query space is critical, since performing nearest-neighbor on the raw observations will only deliver desirable results for the simplest of domains.
For example, in the game of Go, a single stone difference on the board can dramatically change the outcome and reading of the board, while seemingly different board positions might still share high-level characteristics (e.g. inﬂuence, partial local sequences, and life-and-death of groups).
At the scale we want to operate, learning the key & query embeddings end-to-end in order to optimize the ﬁnal model predictions is challenging.
Instead, as in the language modelling work by Borgeaud et al.
[7], we learn an embedding function through a surrogate procedure and use the resulting frozen function to describe domain-relevant similarity.
Moreover, scale also constrains the nearest-neighbors lookup to be approximate, since getting the true nearest-neighbors is too time-consuming at inference time and good approximations can be obtained.
To address 2) and make the best use of the relevant data, we provide the data associated with the nearest-neighbor as additional input features to the parametric part of our model, so that the network can learn to interpret it.
We provide some light inductive biases in the architecture to ensure permutation invariance in the neighbors and robustness to outliers and distribution shifts (see Secs.
Letting the network interpret the nearest-neighbor data is essential in large and complex environments, such as Go, because the typically imperfect match between the current situation and retrieved neighbors will require different aspects of this retrieved data to be ignored or emphasized in a highly context speciﬁc manner.
This is in contrast to episodic RL approaches like [6, 29] which prescribe how to use the retrieved data.
Subsequent sections describe the model, retrieval process, and the full algorithm in more detail.
2.1 Retrieval-augmented agents
We consider a setting in which we wish to train a neural network model m θ on a set of environment trajectories τ ∈ D (for example, the dataset D π of trajectories produced by a policy π ).
For each timestep t of τ , the model takes an observation o t 2 and must produce predictions ˆ y of targets y t .
In our retrieval paradigm, in addition to having access to o t , the model also has access to an auxiliary dataset of knowledge or experience D r .
The auxiliary set D r could be the same as D , but it can also contain more information, or in the most general case, information from a completely different distribution and in a different format.
If there is overlap between D r and D , it is important for robustness to ensure that the model cannot retrieve the same trajectory as it is being trained on (otherwise the network will tend to overly trust information retrieved from D r ).
Common information sources in RL are trajectories from other agents or experts (ofﬂine RL) or the agent’s own previous experiences (episodic memory, replay buffer).
Note that, in contrast to ofﬂine RL, we do not assume that we should directly use this auxiliary data as trajectories to train on.
We wish to use D r to inform the model predictions ˆ y t , such that ˆ y t
= m θ ( o t , D r ) .
This is challenging, as D r is typically far too large to directly be consumed as a model input.
One solution, shown in
Fig. 1, is to adopt a retrieval-based model, wherein the parametric and differentiable portion of the model m θ is provided with an informative subset of data { x 1 t , x 2 t , . . .
, x Nt } retrieved from D r , conditioned on o t :
[Equation]
This approach requires a number of design choices relating to the retrieval mechanism, which we explore further below.
2.1.1 Scalable nearest-neighbor retrieval using SCaNN
Inspired by previous work in language modelling, we chose to leverage the SCaNN architecture [11] for fast approximate nearest-neighbor retrieval.
This requires each entry
o i in the dataset D r to be associated with a key vector k i ∈ R d , and a given observation o t to be mapped to a query vector q t ∈ R d .
During retrieval, the squared Euclidean distances between q t and the dataset keys k i are used to determine which neighbors to retrieve – reminiscent of neural attention mechanisms.
This retrieval process is very efﬁcient and can be scaled to datasets with billions of items.
We will typically want to retrieve further associated meta-data, or context, for each neighbor.
For example, if a neighbor is part of a trajectory, we would like to retrieve information about action choices and their consequences.
The neighbor observation o i , together with its meta-data, forms the auxiliary input x it to the model.
The nearest-neighbor retrieval process is non-differentiable, which means that the query and key mappings cannot be trained end-to-end directly.
Instead, in this ﬁrst study, we pre-train a non-retrieval prediction network m eφ on our experience dataset D (details in Sec.
A.4).
We then use this network to generate an embedding corresponding to a given observation o t by retrieving the network activations from a speciﬁed layer of m eφ .
We use principal component analysis to compress these activations to a d dimensional vector representing this observation state.
The embedding and projection step together form our key (& query 4 ) network k i
= g φ ( o i ) . [Table]
[Table] [Figure] [Table] We preprocess all of the observations in D r using g φ to produce corresponding keys.
The resulting dataset of keys and observations is then used by ScaNN to retrieve neighbors given a query vector.
For ofﬂine RL training with a ﬁxed training dataset, we can also preprocess the nearest neighbor lookups for all training dataset observations, i.e., q t =
g φ ( o ) ∀ o t ∈ D .
However, to act during evaluation, we must do this nearest-neighbor lookup online for each new observation.
In future experiments, we intend to incorporate end-to-end learning of the query vector [12] to improve agent performance.
A future online RL pipeline will also require us to dynamically retrieve neighbors during training.
2.1.2 Policy optimization with retrieval
The objective in our ofﬂine RL setting is to leverage the available data in order to optimize a policy π for acting.
Our proposed semi-parametric model (Eqn. 1) is compatible with several methods for policy optimization in this ofﬂine setting.
Indeed, it can represent the actor π θ ( o t , D r ) and/or the critic Q θ ( o t , a, D r ) in an ofﬂine actor-critic method [21], or the Q-value in a value iteration approach [32].
Here we focus on a model-based approach, inspired by MuZero [36], where the learned model is employed in a search procedure based on Monte-Carlo Tree Search (MCTS) [8].
This is motivated by the efﬁcacy of MuZero in the ofﬂine RL setting [37, 24].
Model-based search with retrieval Following MuZero, our prediction model m θ is conditioned on the current observation o t but also on a sequence of future actions (cid:126)a t = a t +1 , a t +2 , . .
. , a t + K .
The model m θ is therefore redeﬁned as ˆ y t = m θ ( o t ,(cid:126)a t , x 1 t , x 2 t , . . .
, x Nt ) .
The architecture for this retrieval prediction model is illustrated in Fig.
The model m θ can be decomposed into an encoding step s t = f θ ( o t , x 1 t , x 2 t , . . .
, x nt ) , which incorporates the observation and neighbor information, followed by iterative action-conditioned inference s k +1 t
= h θ ( s kt , a t + k ) to produce embeddings s kt corresponding to subsequent time steps (where s 0 t =
s t ).
The model output ˆ y t is composed of K + 1 value and policy distributions for the current and next K time-steps: { ˆ v kt , ˆ π kt } k =0 ...
K .[-
-]
These outputs are obtained from their respective embeddings s kt .
The target for value predictions is the game outcome and the targets for the policy predictions are the actions taken by the expert a t , a t +1 , . . .
, a t + K .
The sample loss L ( θ, ˆ y t , y t ) is obtained by summing the K + 1 individual loss terms for value and policy predictions (see detail in App. A.2).
The training procedure for this retrieval prediction model is summarized in Alg 1.
In order to act using the trained model, online search is used to generate an improved policy π s = Search ( m θ , D r ) .
We use MCTS with a pUCT rule for internal action selection [35, 40], carrying out n sims simulations per time step.
Model-based search is implemented as follows: after retrieving the observation’s neighbors, the encoded state s t is computed.
Search is carried out from s t by varying the input action sequence (cid:126)a for each simulation and collecting the model outputs to update search statistics.
Details of this process can be found in [36].
At the end of the search, the resulting policy π s ( a | o t ) can be sampled to select the next action.
Due to the semi-parametric nature of the model supporting the search, introducing changes to D r will have an immediate effect on the acting policy π s , even if the model parameters θ remain ﬁxed.
As we explore in Sec. 3.4, this provide a mechanism for fast adaption to new information or experiences.
2.1.3 Using neighbor information
Several choices are possible for the network f θ that processes the observation and neighbors.
Since this is not the main focus of this work, we chose a straightforward approach.
We ﬁrst compute an embedding o et of the observation o t .
We then compute an embedding for each neighbor e ti = p
θ ( o e t , x t i ) , using the same network p θ for all neighbors.
These streams are combined in a permutation invariant way through a sum, and then concatenated with o e t to produce the ﬁnal embedding s t , :
[Equation]
2.2 Evaluating retrieval in the domain of Go
In order to test the utility of retrieval in RL, we wish to evaluate whether agents can effectively generalise between related but distinct experiences, as opposed to simply retrieving the outcome of a previous instance of an identical situation.
This motivates our choice of Go as a domain.
We focus on 9x9 Go, instead of full-scale 19x19 Go, as the less demanding computational costs of 9x experiments enable a more thorough analysis.
Even for 9x9 Go, the state space of 10 38 possible games is vastly larger than the number of positions that we could hope to query.
We collected a dataset of ∼ 3.5M expert 9x9 Go self-play games from an AlphaZero-style agent [40].
We randomly subsampled 15% of the positions from these games, leaving us with ∼ 50M board state observations.
These, along with metadata on the game outcome, ﬁnal game board, and future actions (all encoded as input planes), form our retrieval dataset D r .
We chose to subsample as we hypothesised that this would reduce the chance of retrieving multiple neighbors from the same trajectory, and therefore boost the retrieved neighbor diversity.
However, we did not perform ablations on this choice.
A future solution would be to use a ﬁltering mechanism to reject neighbors from the same trajectory.
In this initial work, training and retrieval datasets are the same – at least during training.
During training, we split D r in two halves such that each game’s observations are only in one of the datasets.
We retrieve neighbors for an observation o t from the half it is not contained in.
This is simply to avoid retrieving the same position as the query.
Other ways to obtain the same effect could be devised.
2.3 Regularisation
We wish to make our network robust to poor quality neighbors, in order to ensure that the network can perform well in settings for which there is lower overlap with the retrieval dataset than encountered in training.
We therefore explore several techniques to improve network robustness to irrelevant neighbors.
We randomly zero-out a subset of retrieved neighbors during training ("neighbor dropout"), and/or more adversarially, randomly replace a subset of retrieved neighbors with the neighbors of a different observation ("neighbor randomisation").
Inspired by [10], we also explore using a loss to regularise the embedding produced by the neighbor retrieval towards the embedding produced with the observation alone ("neighbor regularisation").
Further details are given in Sec.
A.5.
We carried out ablations of these techniques (Appendix Fig. 7), which show that neighbor randomisation is important in some contexts for maintaining performance, but that the others do not seem to have a signiﬁcant effect in our ﬁnal conﬁguration.
The results reported in this study utilise all of these augmentations, as during the development process they had been found to slightly beneﬁt performance.
Interestingly, as we explore in later sections, MCTS with enough simulations compensates for the harmful effect of low-quality neighbors, and can perform effectively without these training augmentations.
Results 3.1 Qualitative examination of retrieved neighbors
[Figure]
Changing a single stone in a Go board can dramatically affect the game, and the number of Go positions is enormous.
We ﬁrst wanted to assess whether, despite the combinatorial aspect of the domain, meaningful nearest-neighbors could be retrieved from the dataset using our learned keys/queries.
Examples of retrieved positions in Go are shown in Fig.
3, where we observed relevant matches in terms of both local and global structure.
In some cases, especially in the early game, the retrieved position is an exact match even though the rest of the game (and therefore the associated nearest-neighbor meta-data) differs.
Row 1 in Fig.
3 is an example of this – in this case, the retrieved data effectively provides sample rollouts from the query position akin to those derived from MCTS.
[Figure]
3.2 Impact of retrieval on supervised training
We next evaluated the extent to which the model can learn to exploit retrieved information for better model predictions and better decision-making.
First, we evaluated the impact of retrieval on supervised learning losses.
As a baseline comparison, we trained the same network architectures but setting x i = 0 for all i — this has the effect of maintaining the number of parameters, ﬂops, and useful capacity in the network, while removing access to the retrieval data.
We evaluated losses for retrieval networks trained with different N (Fig. 4b) and model sizes (Fig.
4c, see Sec.
A.1 for details).
Across conditions, we consistently observed a signiﬁcant boost in test-set accuracy for all metrics and over the course of training.
This improvement to predictions is further observed across game trajectories, and is not limited, for example, to opening play positions.
One advantage of the semi-parametric approach we outlined is that we can modify the retrieval dataset and potentially see immediate effects on the prediction, without changing the parameters θ .
We ﬁrst veriﬁed this by evaluating our model when allowed access to varying fractions of the full dataset D r (Fig. 4d).
We observed clear gains in prediction accuracy from increasing the size of the retrieval set (only half is used at train time).
A further important observation is that large-scale datasets are clearly important - the evaluation metrics drop below the baseline level for a dataset D r that is 1% the size of the full dataset.
3.3 Evaluation against a reference opponent
While the results observed on the ofﬂine dataset suggest a strong positive effect from learning to leverage the retrieval data, this is in the context of a ﬁxed test data distribution that matches the retrieval data distribution.
When deployed, games played against different opponents will likely diverge from that distribution.
We evaluated the performance of our search policy π s ( n sims = 200 ) by playing against a ﬁxed reference opponent — the Pachi program [5], which can perform beyond strong amateur level of play in 9x
Go given sufﬁcient simulation budget (we evaluate against 400k simulations).
We observed a signiﬁcant boost to performance for retrieval-based networks over the equivalent baseline network of the same capacity (Fig. 5).
Interestingly, as shown in Appendix Fig. 8, playing using only the base policy prior ˆ π 0 shows a much smaller boost over the baseline network.
As we explore further below, the quality of retrieved neighbors available during play is much lower than for training - we hypothesise that the search policy is more robust to this distributional shift than ˆ π 0 .
[Figure]
3.4 Augmenting the retrieval dataset
[Figure]
We observed a signiﬁcant distribution shift between game positions observed in play against the Pachi opponent versus positions in the retrieval dataset (Fig. 6a), as measured by the empirical distance distribution of game observations to their approximate nearest neighbors.
Changing or augmenting the retrieval dataset D r modiﬁes this distributional shift.
A particularly interesting modiﬁcation in this setting is to augment the dataset with play recorded between our agents and the Pachi opponent.
As can be seen in Fig.
6a, augmenting the dataset with a set of ∼ 600k agent-Pachi game states increases the similarity between positions observed in play and their retrieved neighbors.
Strikingly, we found that integrating these games into our retrieval dataset leads to improvements in performance without further training for subsequent play against the Pachi opponent (play with the MCTS policy π s is shown in Fig.
6b, play with the policy prior ˆ π 0 is shown in Appendix Fig. 9).
This highlights the potential of a semi-parametric model to rapidly integrate recent experience without further training.
As a related intervention, we instead inputted randomly retrieved neighbors into the model at evaluation time (see Appendix Fig. 6).
This signiﬁcantly impaired performance for play using the policy prior ˆ π 0 , but somewhat unexpectedly, for play with the MCTS policy π s , performance did not fully regress to the level of the baseline network.
This suggests that our retrieval network is able to amortise some information from neighbors during training, leading to better performance even when no relevant neighbors are available.
Related work
The idea of supporting decision making by directly attending to a cache of related events has been visited many times in different contexts, under the names of case-based reasoning [30], nonparametric RL [16, 4, 1, 27], or episodic memory [20, 23].
The aim of some methods is to better support the working memory of an agent during a given episode [26, 42] or a series of successive related episodes [33].
Other methods, including ours, aim to leverage a broader class of relevant experience or persistent knowledge (including across episodes, and from other agents) to better support reasoning and planning.
One differentiating factor in our work to these past approaches is that we do not prescribe how to process the information from the available data (e.g. through specifying the agent’s action-value directly in terms of previously generated value estimates [6, 13, 15, 29], or a model from observed transitions [39]) but rather learn end-to-end how the data can support better predictions within the parametric model.
A recent approach by Goyal et al.
[10] has considered an attention mechanism to select where and what to use from available trajectories, but over a small retrieval batch of data rather than the full available experience data.
Another class of method to leverage a transition dataset is to replay the data at training time in order to perform more gradient steps per experience, this is a widespread technique in modern RL algorithms [21, 22, 25, 36] but it does not beneﬁt the agent at test time, requires additional learning steps to adapt to new data, and does not allow end-to-end learning of how to relate past experience to new situations.
Discussion
Our approach and empirical results highlight how reinforcement learning agents can beneﬁt from direct access to a large collection of raw interaction data at inference time, through a retrieval mechanism, in addition to their already effective parametric representation.
We showed this was the case even 1) when the domain is large enough to require generalisation in how to interpret past data, 2) when there is signiﬁcant distribution shift when acting, and 3) at a scale where only approximate nearest-neighbors can be retrieved.
We believe this already demonstrates the potential of this approach for many possible scenarios and applications.
We show that retrieval can be effectively combined with model-based search.
We ﬁnd that the beneﬁts from retrieval and search are synergistic, with increasing numbers of retrieved neighbors and increasing simulations both leading to performance increases in almost all contexts we investigated.
Furthermore, empirical evidence suggests that search signiﬁcantly improves agent robustness to distributional shift as compared to playing with the policy prior.
In Appendix Fig.
10, we compare the boost in performance as a function of parametric-model compute cost for increasing MCTS simulations versus increasing the number of retrieved neighbors processed by the model.
This provides a tentative indication that retrieval is also a compute efﬁcient means of improving performance.
A key future direction is to investigate the online learning scenario in which recent experience of the agent is rapidly made available for retrieval, hence progressively growing the retrieval dataset over the course of training.
While there are additional challenges associated with the online paradigm (e.g., it may be desirable to update the queries and/or keys during training [12]), the fast adaptation effect we highlighted in this work may have even more impact there.
There are many potentially relevant sources of information beyond an agent’s own experience, or that of other humans or agents.
For example, it has been shown that YouTube videos are useful for learning to play the Atari game Montezuma’s revenge [2].
Training an embedding network on sufﬁciently diverse data may enable retrieval of information from a wide range of contexts [43], including third-person demonstrations, videos and perhaps even books. Disclosure of Funding The authors received no speciﬁc funding for this work
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] No theoretical results.
(b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental
material or as a URL)?
[No] Instructions and links to relevant libraries were provided.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A]
No people were involved in the data gathering process.
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
The appendix details the architecture and training mechanisms for our results, it also lists additional ablations and results. A.1 Architecture Observations consist of 9x9 feature planes representing the current state of the board (indicating stones for the last two moves, and the player color).
The networks we train are convolutional, composed of residual blocks, and maintain the 9x9 topology throughout — except for the ﬁnal output layers. As illustrated in Figure 1, the main inference network s t = f θ ( o t , x 1 t , x 2 t , . . . , x nt ) is composed of an encoding of the observation o t using m enc =2 residual 3 × 3 convolutional blocks with 256 channels (each block consists of two convolutional layers with a skip connection) [14], preceded by a ﬁrst convolution layer with 256 ﬁlters and followed by a convolution layer that goes back to 128 channels. Each neighbor information x i is concatenated separately with the encoded observation o e to be processed by another residual network p θ (blue layers in Fig. 1), but with m nn =3 blocks andchannels.
The neighbor information, in addition to the feature planes corresponding to the retrieved position, also include 1-hot 9x9 input planes for the 10 next actions, a tiled input plane for the game outcome, and the feature planes corresponding to the ﬁnal game board.
The parameters for each neighbor processing stream are shared.
The output from each of the neighbors e i are then summed (and normalized) and concatenated back with the encoded observation along the channel dimension, see Eq.
The concatenated tensor is processed by m root =
6 additional residual blocks with the same hyperparameters to produce the main output embedding s t = s 0 t (top pink layers in Fig. 1).
Each internal transition in the recurrent action-conditional model (in purple Fig. 1) takes the previous embedding s kt , concatenates a one-hot representation of the action a t + k , and outputs the next embedding s k +1 t .
These internal models are composed of m tran =3 residual blocks preceded by a convolution layer to bring the activations back to 256 layers.
All internal transitions, and output networks for k > 1 , share parameters.
Non-linear transformations after each layer are rectiﬁed linear units (ReLU), preceded by a layer norm layer applied spatially within the residual blocks.
The model outputs are obtained at each step k from the embedding s kt with a value and policy head.
The value head ﬁrst applies a single 1 × 1 convolution ﬁlter, followed by a ReLU, a ﬂattening operation, and an MLP with a single hidden layer of 256 units to a ﬁnal scalar output passed through a tanh non-linearity to obtain ˆ v kt .
{+The policy head follows the same pattern, but applies two initial convolution ﬁlters, and the ﬁnal output has dimension 
+}
= 9 × 9 + 1 (all possible actions including pass in 9x9 Go) and corresponds to the logits for ˆ π kt .
Together, this describes the parametric network m θ .
The search policy π s = Search ( m θ , D r ) employs the trained network to plan and act.
The main embedding s t = f θ ( o t , x 1 t , x 2 t , . .
, x nt ) , and the retrieval process to obtain the neighbors, is computed once per search.
Multiple rollouts can be computed from s t with different input action sequences for a given root observation.
The details of the MCTS search mechanism using the trained model follows [36] closely, we refer the reader to that paper for details.
For experiments that vary the model size, model size of 1 is as described above.
For model sizes n = 2 & 4 , we multiply m root , m tran , m enc , and m nn by a factor of 2 , with the number of channels at all layers doubled for n = 4 .
The model size labeled 0 .
5 corresponds to m root =3 , m tran =2 , m enc =1 , and m nn =2 .
A.2 Losses and optimization
The total loss, for a single datapoint, is:
[Equation]
with w k = 1 K for k > 0 and 1 otherwise ( K = 5 in our experiments), g the empirical return (the − 1 , 1 relative game outcome for Go), and α = 1 e- 4 controls the weight decay.
The individual loss terms are l v ( v, g )
= 12 ( v − g ) 2 for the value outputs (MSE when averaged over the batch elements), and l p ( π, a )
= − log( π ( a )) for the policy outputs.
In addition, an optional regularisation loss is described below in Sec.
A.5.
We use the Adam optimizer [18] with a learning rate schedule, training for 600k steps with a batch size of 1024.
The learning rate is the initial learning rate divided by { 2 , 8 , 64 , 256 } after respectively {180k, 360k, 480k, 570k} steps, with the initial learning rate of 1 e- 3 .
Training was performed on
TPU [17] using JAXline and Mctx within the DeepMind JAX Ecosystem [3].
A.3 Go game conﬁguration
The komi throughout the experiments is 5.5.
The Pachi program is conﬁgured with 16 threads, a max_tree_size of 2000, a resign_threshold of 0.1, no pondering, and the Chinese ruleset.
The search policy π s when testing is deﬁned as the action with the maximum visit counts.
To ensure diversity in the evaluation games against Pachi, we seed each game by letting Pachi play the ﬁrstmoves in the opening.
After an evaluation game goes past the opening moves and 50 subsequent moves, we stop the game if the value function is above 0.995, or below 1-0.995, to avoid playing long end-games when the game is already clearly settled.
A.4 Embedding network
To train the embedding g φ to map observations o t to keys and queries, we ﬁrst train a MuZero model m eφ ( o t ,(cid:126)a t ) following the same method as in A.1-A.2.
We train this on the same data D that we train the retrieval network on, but without any retrieval input or neighbor processing network.
We also use 8 residual blocks for the main encoding stage to compute s t = f φ ( o t ) , and 3 residual blocks per internal model transition.
At the end of training, we choose the output of one of the f φ network layers as a (pre)-embedding ¯ g φ ( o ) .
To reduce the dimensionality of ¯ g φ , we project it using PCA into the ﬁrst d = 512 principal components to obtain g φ ( o ) =
(¯ g φ ( o ) − µ ) V ∈ R d , where V contains the ﬁrst d eigenvectors of the covariance matrix of the pre-embedding vectors, estimated with a subset of the data after removing their mean µ .
In our experiments, we used the output of the 6th residual block of f φ as ¯ g φ ( o ) .
We have not extensively searched for the best such embedding due to the expense of testing each – further investigation of this design choice will be the subject of future research.
A.5 Regularisation ablations
Here we provide more detail on techniques we investigated to improve network robustness to low quality retrieved neighbors (as introduced in Sec. 2.3).
Neighbor dropout
We randomly zero-out a subset of retrieved neighbors during training.
For N retrieved neighbors, the number of neighbors to zero out M is uniformly randomly chosen from [0 , N ] .
We then randomly choose M out of N neighbors, and mask out their contribution to the ﬁnal invariant neighbor embedding (Sec. 2.1.3).
Neighbor randomisation
We randomly replace a subset of retrieved neighbors with the neighbors of a different observation within the mini-batch.
This is implemented by randomly choosing M neighbors as described for dropout above, and then replacing each by a randomly chosen member of the mini-batch (which may with a small chance be the correct neighbor).
Neighbor regularisation
Inspired by [10], we also explore using a loss to regularise the embedding produced by the neighbor retrieval towards the embedding produced with the observation alone.
In this case, we apply a convolutional layer to the embedding produced from the game state observation, and use this to predict the output of the neighbor processing tower.
We apply a mean-squared-error loss to encourage the base embedding to be predictive of the neighbor output, and similarly apply an equivalent loss to the neighbor output.
[Figure]
Without this
, the performs worse
A.6 Additional results
[Figure]
[Figure]
[Figure]