Year: 2021

Abstract:
Neural networks leverage robust internal representations in order to generalise. Learning them is difficult, and often requires a large training set that covers the data distribution densely. We study a common setting where our task is not purely opaque. Indeed, very often we may have access to information about the underlying system (e.g. that observations must obey certain laws of physics) that any “tabula rasa” neural network would need to re-learn from scratch, penalising data efficiency. We incorporate this information into a pre-trained reasoning module, and investigate its role in shaping the discovered representations in diverse selfsupervised learning settings from pixels. Our approach paves the way for a new class of data-efficient representation learning.

heading: 1 INTRODUCTION
text: Neural networks are able to learn policies in environments without access to their specifics (Schrittwieser et al., 2020), generate large quantities of text (Brown et al., 2020), or automatically fold proteins to high accuracy (Senior et al., 2020). However, such “tabula rasa” approaches hinge on having access to substantial quantities of data, from which robust representations can be learned. Without a large training set that spans the data distribution, representation learning is difficult (Belkin et al., 2019; LeCun, 2018; Sun et al., 2017).
Here, we study ways to construct neural networks with representations that are robust, while retaining a data-driven approach. We rely on a simple observation: very often, we have some (partial) knowledge of the underlying dynamics of the data, which could help make stronger predictions from fewer observations. This knowledge, however, usually requires us to be mindful of abstract properties of the data—and such properties cannot always be robustly extracted from natural observations.
Motivation Consider the task of predicting the future state of a system of n bouncing balls, from a pixel input x (Figure 1). Reliably estimating future pixel observations, y, is a challenging reconstruction task. However, the generative properties of this system are simple. Assuming knowledge of simple abstract inputs (radius, rc, position, xc, and velocity, vc) for every ball, x̄c, the future movements in this abstract space are the result of applying the laws of physics to these low-dimensional quantities. Hence, future abstract states, ȳ, can be computed via a simple algorithm that aggregates pair-wise forces between objects.
While this gives us a potentially simpler path from pixel inputs to pixel outputs, via abstract inputs to abstract outputs (x → x̄ ȳ → y), it still places potentially unrealistic demands on our task setup, every step of the way:
x→ x̄: Necessitates either upfront knowledge of how to abstract away x̄ from x, or a massive dataset of paired (x, x̄) to learn such a mapping from;
x̄ ȳ: Implies that the algorithm perfectly simulates all aspects of the output. In reality, an algorithm may often only give partial context about y. Further, algorithms often assume that x is provided without error, exposing an algorithmic bottleneck (Deac et al., 2021): if x̄ is incorrectly predicted, this will negatively compound in ȳ, hence y;
ȳ→ y: Necessitates a renderer that generates y from ȳ, or a dataset of paired (ȳ,y) to learn it.
We will assume a general setting where none of the above constraints hold: we know that the mapping x̄ ȳ is likely of use to our predictor, but we do not assume a trivial mapping or a paired dataset which would allow us to convert directly from x to x̄ or from ȳ to y. Our only remaining assumption is that the algorithm x̄ ȳ can be efficiently computed, allowing us to generate massive quantities of paired abstract input-output pairs, (x̄, ȳ).
Present work In this setting, we propose Reasoning-Modulated Representations (RMR), an approach that first learns a latent-space processor of abstract data; i.e. a mapping x̄
f−→ z P−→ z′ g−→ ȳ, where z ∈ Rk are high-dimensional latent vectors. f and g are an encoder and decoder, designed to take abstract representations to and from this latent space, and P is a processor network which simulates the algorithm x̄ ȳ in the latent space.
We then observe, in the spirit of neural algorithmic reasoning (Veličković & Blundell, 2021), that such a processor network can be used as a drop-in differentiable component for any task where the
x̄ ȳ kind of reasoning may be applicable. Hence, we then learn a pipeline x f̃−→ z P−→ z′ g̃−→ y, which modulates the representations z obtained from x, forcing them to pass through the pretrained processor network. By doing so, we have ameliorated the original requirement for a massive natural dataset of (x,y) pairs. Instead, we inject knowledge from a massive abstract dataset of (x̄, ȳ) pairs, directly through the pre-trained parameters of P . This has the potential to relieve the pressure on encoders and decoders f̃ and g̃, which we experimentally validate on several challenging representation learning domains.
Our contributions can be summarised as follows:
• Verifying and extending prior work, we show that meaningful latent-space models can be learned from massive abstract datasets, on physics simulations and Atari 2600 games;
• We then show that these latent-space models can be used as differentiable components within neural pipelines that process raw observations. In doing so, we recover a dataefficient neural network pipeline that relies solely on the existence of massive abstract datasets (which can often be automatically generated).
• Finally, we demonstrate early signs of processor reusability: latent-space abstract models can be used in tasks which do not even directly align with their environment, so long as these tasks can benefit from their underlying reasoning procedure.

heading: 2 RELATED WORK
text: Neural algorithmic reasoning RMR relies on being able to construct robust latent-space models that imitate abstract reasoning procedures. This makes it well aligned with neural algorithmic reasoning (Cappart et al., 2021), which is concerned with constructing neuralised versions of classical algorithms (typically by learning to execute them in a manner that extrapolates). Leveraging the ideas of algorithmic alignment (Xu et al., 2019), several known algorithmic primitives have already been successfully neuralised. This includes iterative computation (Veličković et al., 2019; Tang et al., 2020), linearithmic algorithms (Freivalds et al., 2019), and data structures (Strathmann et al., 2021; Veličković et al., 2020). Further, the XLVIN model (Deac et al., 2021) demonstrates how such primitives can be re-used for data-efficient planning, paving the way for a blueprint (Veličković & Blundell, 2021) that we leverage in RMR as well.
Physical simulation with neural networks Our work also has contact points with prior art in using (graph) neural networks for physics simulations. In fact, there is a tight coupling between algorithmic computation and simulations, as the latter are typically realised using the former. Within this space, abstract GNN models of physics have been proposed by works such as interaction networks (Battaglia et al., 2016) and NPE (Chang et al., 2016), and extended to pixel-based inputs by visual interaction networks (Watters et al., 2017). The generalisation power of these models has increased drastically in recent years, with effective models of systems of particles Sanchez-Gonzalez et al. (2020) as well as meshes (Pfaff et al., 2020) being proposed. Excitingly, it has also been demonstrated that rudimentary laws of physics can occasionally be recovered from the update rules of these GNNs (Cranmer et al., 2020), and that they can be used to uncover new physical knowledge (Bapst et al., 2020).
Recent work has also explored placing additional constraints on learning-based physical simulators, for example by using Hamiltonian ODE integrators in conjunction with GNN models (SanchezGonzalez et al., 2019), or by coupling a (non-neural) differentiable physics engine directly to visual inputs and optimizing its parameters via backprop (Jaques et al., 2019; Jatavallabhula et al., 2021).
Object-centric and modular models for dynamic environments RMR with factored latents can be viewed as a form of object-centric neural network, in which visual objects in an image or video are represented as separate latent variables in the model and their temporal dynamics and pairwise interactions are modeled via GNNs or self-attention mechanisms. There is a rich literature on discovering objects and learning their dynamics from raw visual data without supervision, with objectcentric models such as R-NEM (Van Steenkiste et al., 2018), SQAIR (Kosiorek et al., 2018), OP3 (Veerapaneni et al., 2020), SCALOR (Jiang et al., 2019), G-SWM (Lin et al., 2020). Recent work has explored using contrastive losses (Kipf et al., 2019) in this context or other losses directly in latent space (François-Lavet et al., 2019). Related approaches discover and use keypoints (Kulkarni et al., 2019) to describe objects and even discover causal relations from visual input (Li et al., 2020) using neural relational inference (Kipf et al., 2018) in conjunction with a keypoint discovery method. A related line of works integrate attention-mechanisms in modular and object-centric models to interface latent variables with visual input, including models such as RMC (Santoro et al., 2018), RIM (Goyal et al., 2019), Slot Attention (Locatello et al., 2020), SCOFF (Goyal et al., 2021b), and NPS (Goyal et al., 2021a).
3 RMR ARCHITECTURE
Having provided a high-level overview of RMR and surveyed the relevant related work, we proceed to carefully detail the blueprint of RMR’s various components. This will allow us to ground any subsequent RMR experiments on diverse domains directly in our blueprint. Throughout this section, it will be useful to refer to Figure 2 which presents a visual overview of this section.
Preliminaries We assume a set of natural inputs, X , and a set of natural outputs, Y . These sets represent the possible inputs and outputs of a target function, Φ : X → Y , which we would like to learn based on a (potentially small) dataset of input-output pairs, (x,y), where y = Φ(x).
We further assume that the inner workings of Φ can be related to an algorithm, A : X̄ → Ȳ . The algorithm operates over a set of abstract inputs, X̄ , and produces outputs from an abstract output set Ȳ . Typically, it will be the case that dim X̄ dimX ; that is, abstract inputs are assumed substantially lower-dimensional than natural inputs. We do not assume existence of any aligned input pairs (x, x̄), and we do not assume that A perfectly explains the computations of Φ. What we do assume is that A is either known or can be trivially computed, giving rise to a massive dataset of abstract input-output pairs, (x̄, ȳ), where ȳ = A(x̄).
Lastly, we assume a latent space, Z , and that we can construct neural network components to both encode and decode from it. Typically, Z will be a real-valued vector space (Z = Rk) which is high-dimensional; that is, k > dim X̄ . This ensures that any neural networks operating over Z are not vulnerable to bottleneck effects.
Note that either the natural or abstract input set may be factorised, e.g., into objects; in this case, we can accordingly factorise the latent space, enforcing Z = Rn×k, where n is the assumed maximal number of objects (typically a hyperparameter of the models if not known upfront).
Abstract pipeline RMR training proceeds by first learning a model of the algorithm A, which is bound to pass through a latent-space representation. That is, we learn a neural network approximator g(P (f(x̄))) ≈ A(x̄), which follows the encode-process-decode paradigm (Hamrick et al., 2018). It consists of the following three building blocks: Encoder, f : X̄ → Z , tasked with projecting the abstract inputs into the latent space; Processor, P : Z → Z , simulating individual steps of the algorithm in the latent space; Decoder, g : Z → Ȳ , tasked with projecting latents back into the abstract output space. Such a pipeline is now widely used both in neural algorithmic reasoning (Veličković et al., 2019) and learning physical simulations (Sanchez-Gonzalez et al., 2020), and can be trained end-to-end with gradient descent.
For reasons that will become apparent, it is favourable for most of the computational effort to be performed by P . Accordingly, encoders and decoders are often designed to be simple learnable linear projections, while processors tend to be either deep MLPs or graph neural networks—depending on whether the latent space is factorised into nodes.
Natural pipeline Once an appropriate processor, P , has been learned, it may be observed that it corresponds to a highly favourable component in our setting. Namely, we can relate its operations to the algorithm A, and since it stays high-dimensional, it is a differentiable component we can easily plug into other neural networks without incurring any bottleneck effects. This insight was originally recovered in XLVIN (Deac et al., 2021), where it yielded a generic implicit planner. As an ablation, we have also rediscovered the bottleneck effect in our settings; see Appendix C. We now leverage similar insights for general representation learning tasks.
On a high level, what we need to do is simple and elegant: swap out f and g for natural encoders and decoders, f̃ : X → Z and g̃ : Z → Y , respectively. We are then able to learn a function g̃(P (f̃(x))) ≈ Φ(x), which is once again to be optimised through gradient descent. We would like P to retain its semantics during training, and therefore it is typically kept frozen in the natural pipeline. Note that P might not perfectly represent A which in turn might not perfectly represent Φ. While we rely on a skip connection in our implementation of P , it has no learnable parameters and does not offer the system the ability to learn a correction of P in the natural setting. Our choice is motivated by the desire to both maintain the semantics and interpretability of P and to make this processor a bottleneck forcing the model to rely on it. We show empirically that our pipeline is surprisingly robust to imperfect P models even with weak (linear) encoders/decoders.
It is worth noting several potential challenges that may arise while training the natural pipeline, especially if the training data for it is sparsely available. We also suggest remedies for each:
• If x and/or y exhibit any nontrivial geometry, simple linear projections will rarely suffice for f̃ and g̃. For example, our natural inputs will often be pixel-based, necessitating a convolutional neural network for f̃ .
• Further, since the parameters of P are kept frozen, f̃ is left with a challenging task of mapping natural inputs into an appropriate manifold that P can meaningfully operate over. While we demonstrate clear empirical evidence that such meaningful mappings definitely occur, we remark that its success may hinge on carefully tuning the hyperparameters of f̃ .
• A very common setting assumes that the abstract inputs and latents are factorised into objects, but the natural inputs are not. In this case, f̃ is tasked with predicting appropriate object representations from the natural inputs. This is known to be a challenging feat (Greff et al., 2020), but can be successfully performed. Sometimes arbitrarily factorising the feature maps of a CNN (Kipf et al., 2019) is sufficient, while at other times, models such as slot attention (Locatello et al., 2020) may be required.
• One corollary of using automated object extractors for f̃ is that it’s very difficult to enforce their slot representations to line up in the same way as in the abstract inputs. This implies that P should be permutation equivariant (and hence motivates using a GNN for it).

heading: 4 RMR FOR BOUNCING BALLS
text: To evaluate the capability of the RMR pipeline for transfer from the abstract space to the pixel space, we apply it on the “bouncing balls” problem. The bouncing balls problem is an instance of a physics simulation problem, where the task is to predict the next state of an environment in which multiple balls are bouncing between each other and a bounding box. Though this problem had been studied in the the context of physics simulation from (abstract) trajectories (Battaglia et al., 2016) and from (natural) videos (Watters et al., 2017; Van Steenkiste et al., 2018; Löwe et al., 2020), here we focus on the aptitude of RMR to transfer learned representations from trajectories to videos.
Our results affirm that strong abstract models can be trained on such tasks, and that including them in a video pipeline induces more robust representations. See Appendix A for more details on hyperparameters and experimental setup.
Preliminaries Here, trajectories are represented by 2D coordinates of 10 balls through time, defining our abstract inputs and outputs X̄ = Ȳ = R10×2. We slice these trajectories into a series of moving windows containing the input, x̄∗, spanning a history of three previous states, and the target, ȳ, representing the next state. We obtain these trajectories from a 3D simulator (MuJoCo (Todorov et al., 2012)), together with their short-video renderings, which represent our natural input and output space X = Y = R64×64×3. Our goal is to train an RMR abstract model on trajectories and transfer learned representations to improve a dynamics model trained on these videos.
Abstract pipeline So as to model the dynamics of trajectories, we closely follow the RMR desiderata for the abstract model. We set f to a linear projection over the input concatenation, P to a Message Passing Neural Network (MPNN), following previous work (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2020), and g to a linear projection.
Our model learns a transition function g(P (f(x̄∗)) ≈ ȳ, supervised using Mean Squared Error (MSE) over ball positions in the next step. It achieves an MSE of 4.59 × 10−4, which, evaluated qualitatively, demonstrates the ability of the model to predict physically realistic behavior when unrolled for 10 steps (the model is trained on 1-step dynamics only). Next, we take the processor P from the abstract pipeline and re-use it in the natural pipeline.
Natural pipeline Here we evaluate whether the pre-trained RMR processor can be reused for learning the dynamics of the bouncing balls from videos. The pixel-based encoder f̄ is concatenation of per-input-image Slot Attention model (Locatello et al., 2020), passed through a linear layer, and a Broadcast Decoder (Watters et al., 2019) for the pixel-based decoder ḡ.
The full model is a transition function ḡ(P (f̄(x∗)) ≈ y, supervised by pixel reconstruction loss over the next step image. We compare the performance of the RMR model with a pre-trained processor P against a baseline in which P is trained fully end-to-end. The RMR model achieves an MSE of 7.94±0.41 (×10−4), whereas the baseline achieves 9.47±0.24 (×10−4). We take a qualitative look at the reconstruction rollout of the RMR model in Figure 3, and expose the algorithmic bottleneck properties for this task in Appendix C.

heading: 5 CONTRASTIVE RMR FOR ATARI
text: We evaluate the potential of our RMR pipeline for state representation learning on the Atari 2600 (Bellemare et al., 2013). We find the RMR applicable here because there is a potential wealth of information that can be obtained about the Atari’s operation—namely, by inspecting its RAM traces.
Preliminaries Accordingly, we will define our set of abstract inputs and outputs as Atari RAM matrices. Given that the Atari has 128 bytes of memory, X̄ = Ȳ = B128×8 (where B = {0, 1} is the set of bits). We collect data about how the console modifies the RAM by acting in the environment and recording the trace of RAM arrays we observe. These traces will be of the form (x̄, a, ȳ) which signify that the agent’s initial RAM state was x̄, and that after performing action a ∈ A, its RAM state was updated to ȳ. We assume that a is encoded as an 18-way one-hot vector.
We would like to leverage any reasoning module obtained over RAM states to support representation learning from raw pixels. Accordingly, our natural inputs, X , are pixel arrays representing the Atari’s framebuffer.
Mirroring prior work, we perform contrastive learning directly in the latent space, and set Y = Z; that is, our natural outputs correspond to an estimate of the “updated” latents after taking an action. All our models use latent representations of 64 dimensions per slot, meaning Z = R128×64. We note that it is important to generate a diverse dataset of experiences in order to train a robust RAM model. To simulate a dataset which might be gathered by human players of varying skill, we sample our data using the 32 policy heads of a pre-trained Agent57 (Badia et al., 2020). Each policy head collects data over three episodes in the studied games. Note that this implies a substantially more challenging dataset than the one reported by Anand et al. (2019), wherein data was collected by a purely random policy, which may well fail to explore many relevant regions of the games.
Abstract pipeline Firstly, we set out to verify that it is possible to train nontrivial Atari RAM transition models. The construction of this abstract experiment follows almost exactly the abstract RMR setup: f and g are appropriately sized linear projections, while P needs to take into account which action was taken when updating the latents. To simplify the implementation and allow further model re-use, we consider the action a part of the P ’s inputs. See Appendix F for detailed equations.
This implies that our transition model learns a function g(P (f(x̄), a)) ≈ ȳ. We supervise this model using binary cross-entropy to predict each bit of the resulting RAM state. Since RAM transitions are assumed deterministic, we assume a fully Markovian setup and learn 1-step dynamics.
For brevity purposes, we detail our exact hyperparameters and results per each Atari game considered in Appendix B. Our results ascertain the message passing neural network (MPNN) (Gilmer
et al., 2017) as a highly potent processor network in Atari: it ranked most potent in 17 out of 24 games considered, compared to MLPs and Deep Sets (Zaheer et al., 2017). Accordingly, we will focus on leveraging pre-trained MPNN processors for the next phase of the RMR pipeline.
Natural pipeline We now set out to evaluate whether our pre-trained RMR processors can be meaningfully re-used by an encoder in a pixel-based contrastive learning pipeline.
For our pixel-based encoder f̃ , we use the same CNN trunk as Anand et al. (2019)—however, as we require slot-level rather than flat embeddings, the final layers of our encoder are different. Namely, we apply a 1 × 1 convolution computing 128m feature maps (where m is the number of feature maps per-slot). We then flatten the spatial axes, giving every slot m × h × w features, which we finally linearly project to 64-dimensional features per-slot, aligning with our pre-trained P . Note that settingm = 1 recovers exactly the style of object detection employed by C-SWM (Kipf et al., 2019). Since our desired outputs are themselves latents, g̃ is a single linear projection to 64 dimensions.
Overall, our pixel-based transition model learns a function g̃(P (f̃(x), a)) ≈ f̃(y), where y is the next state observed after applying action a in state x. To optimise it, we re-use exactly the same TransE-inspired (Bordes et al., 2013) contrastive loss that C-SWM (Kipf et al., 2019) used.
Once state representation learning concludes, all components are typically thrown away except for the encoder, f̃ , which is used for downstream tasks. As a proxy for evaluating the quality of the encoder, we train linear classifiers on the concatenation of all slot embeddings obtained from f̃ to predict individual RAM bits, exactly as in the abstract model case. Note that we have not violated our assumption that paired (x, x̄) samples will not be provided while training the natural model—in this phase, the encoder f̃ is frozen, and gradients can only flow into the linear probe.
Our comparisons for this experiment, evaluating our RMR pipeline against an identical architecture with an unfrozen P (equivalent to C-SWM (Kipf et al., 2019)) is provided in Table 1. For each of 20 random seeds, we feed identical batches to both models, hence we can perform a paired Wilcoxon test to assess the statistical significance of any observed differences on validation episodes. The results are in line with our hypothesis: representations learnt by RMR are significantly better (p < 0.05) on 15 out of 24 games, significantly worse only on Pitfall!—and indistinguishable to C-SWM’s on others. This is despite the fact their architectures are identical, indicating that the pretrained abstract model induces stronger representations for predicting the underlying data factors. As was the case for bouncing balls, we expose the algorithmic bottleneck here too; see Appendix C.
As a relevant initial baseline—and to emphasise the difficulty of our task—we also include in Table 1 the performance of the latent embeddings extracted from a pre-trained Agent57 (Badia et al., 2020). These embeddings are, perhaps unsurprisingly, substantially worse than both the RMR and C-SWM models. As Agent57 was trained on a reward-maximising objective, its embeddings are likely to capture the controllable aspects of the input, while filtering out the environment’s background.
Abstract model transfer While the results above indicate that endowing a self-supervised Atari representation learner with knowledge of the underlying game’s RAM transitions yields stronger representations, one may still argue that this constitutes a form of “privileged information”. This is due to the fact we knew upfront which game we were learning the representations for, and hence could leverage the specific RAM dynamics of this game.
In the final experiment, we study a more general setting: we are given access to RAM traces showing us how the Atari console manipulates its memory in response to player input, but we cannot guarantee these traces came from the same game that we are performing representation learning over. Can we still effectively leverage this knowledge?
Specifically, we test abstract model transfer in the following way: first, we take a pre-trained abstract processor network P from one game (the “train game”) and freeze it. Then, using this processor, we perform the aforementioned natural pipeline training and testing over frames from another game (the “test game”). We then evaluate whether the recovered performance improves over the C-SWM model with unfrozen weights—once again using a paired one-sided Wilcoxon test over 20 seeds to ascertain statistical significance of any differences observed. The final outcome of our experiment is hence a 24 × 24 matrix, indicating the quality of abstract model transfer from every game to every other game. Table 1 corresponds to the “diagonal” entries of this matrix.
The results, presented in Figure 4, testify to the performance of RMR. Representations learned by RMR transfer better in 64.6% of the train/test game pairs, are indistinguishable from C-SWM in 33.7% of the game pairs and perform worse than C-SWM in only 1.7% of game pairs. Therefore, in plentiful circumstances, the answer to our original question is positive: discovering a trace of Atari RAM transitions of unknown origin can often be of high significance for representation learning from Atari pixels, regardless of whether the underlying games match.
Qualitative analysis of transfer While we find this result interesting in and of itself, it also raises interesting follow-up questions. Is representation learning on certain games more prone to being improved just by knowing anything about the Atari console? Figure 4 certainly implies so: several games (such as Seaquest or Space Invaders) have “fully-green” columns, making them “universal recipients”. Similarly, we may be interested about the “donor” properties of each game – to what extent are their RAM models useful across a broad range of test games?
We study both of the above questions by performing a hierarchical (complete-link) clustering of the rows and columns of the 24 × 24 matrix of transfer performances, to identify clusters of related donor and recipient games. Both clusterings are marked in Figure 4 (on the sides of the matrix).
The analysis reveals several well-formed clusters, from which we are able to make some preliminary observations, based on the properties of the various games. To name a few examples:
• The strongest recipients (e.g. Yars’ Revenge, H.E.R.O., Seaquest, Space Invaders and Asteroids) tend to include elements of “shooting” in their gameplay.
• Conversely, the weakest recipients (River Raid, Berzerk, Private Eye, Ms. Pac-Man and Pitfall!) are all games in which movement is generally unrestricted across most of the screen, indicating a larger range of possible coordinate values to model.
• Pong, Breakout, Battlezone and Skiing cluster closely in terms of donor properties—and they are all games in which movement is restricted to one axis only.
• Lastly, strong donors (Montezuma’s Revenge, Yars’ Revenge, Q*Bert and Venture) all feature massive, abrupt, changes to the game state (as they feature multiple rooms, for example). This implies that they might generally have seen a more diverse set of transitions. Further, on Yars’ Revenge, a massive laser features, which, conveniently, prints an RGB projection of the RAM itself on the frame.

heading: 6 CONCLUSIONS
text: We presented Reasoning-Modulated Representations (RMR), a novel approach for leveraging background algorithmic knowledge within a representation learning pipeline. By encoding the underlying algorithmic priors as weights of a processor neural network, we alleviate requirements on alignment between abstract and natural inputs, protect our model against bottlenecks, and make it more data-efficient in the process. We believe that RMR paves the way to a novel class of algorithmic reasoning-inspired representations, with a high potential for transfer across tasks—a feat that has largely eluded deep reinforcement learning research.
REPRODUCIBILITY STATEMENT
Our work proposes a general pipeline for representation learning under algorithmic priors. We have made a careful effort to expose all of the key elements of our RMR pipeline, both within figures and in text. We also report the hyperparameters of our experiments in the Appendices.
We carefully described the manner in which our datasets have been collected. Further, we will explore the options for open-sourcing the specific datasets we employed in this paper post-publication. We note that Atari trajectory datasets similar to the one we explore here (albeit with a significantly weaker actor) are already publicly available e.g. through the work of Anand et al. (2019).

heading: A BOUNCING BALLS MODELLING SETUP
text: Abstract pipeline The f is a linear projection over the concatenation of inputs, outputting a 128- long representation for each of the 10 balls in the input. The P is a MPNN over the fully connected graph of the ball representations. It is a 2-pass MPNN with a 3-layered ReLU-activated MLP as a message function, without the final layer activation, projecting to the same 128-dimensional space, per object. Finally, g is a linear projection applied on each object representation of the output of P .
The model is MSE-supervised with ball position on the next step. It is trained on a 8-core TPU for 10000 epochs, with a batch size of 512, and the Adam optimizer with the initial learning rate of 0.0001.
Natural pipeline f̄ is a Slot Attention model (Locatello et al., 2020) on each input image, concatenating the images and passing them through a linear layer, outputting a 128-dimensional vector for each of the objects. ḡ is a Broadcast Decoder (Watters et al., 2019) containing a sequence of 5 transposed convolutions and a linear layer mapping before calculating the reconstructions and their masks.
The model is MSE-supervised by pixel reconstruction (per-pixel MSE) over the next step image. Both the RMR and the baseline are trained on a 8-core TPU for 1000 epochs, with a batch size of 512, with the Adam optimiser and the initial learning rate of 0.0001, all over 3 random seeds.

heading: B ATARI ABSTRACT MODELLING SETUP AND RESULTS
text: Our best processor network is a MPNN (Gilmer et al., 2017) over a fully connected graph (Santoro et al., 2017) of RAM slots, which concatenates the action embedding to every node (as done in Kipf et al. (2019)). It uses three-layer MLPs as message functions, with the ReLU activation applied after each hidden layer. The entire model is trained for every game in isolation, over 48 distinct episodes of Agent57 experience. We use the Adam SGD optimiser (Kingma & Ba, 2014) with a batch size of 50 and a learning rate of 0.001 across all Atari experiments. To evaluate the benefits of message passing, we also compare our model to Deep Sets (Zaheer et al., 2017), which is equivalent to our MPNN model—only it passes messages over the identity adjacency matrix. Lastly, we evaluate the benefits of factorised latents by comparing our methods against a three-layer MLP applied on the flattened RAM state.
One immediate observation is that RAM updates in Atari are extremely sparse, with a copy baseline already being very strong for many games. To prevent the model from having to repeatedly re-learn identity functions, we also make it predict masks of the shapeM = B128, specifying which cells are to be overwritten by the model at this step. This strategy, coupled with teacher forcing (as done by Veličković et al. (2020); Strathmann et al. (2021)) yielded substantially stronger predictors. We also use this observation to prevent over-inflating our prediction scores: we only display prediction accuracy over RAM slots with label entropy larger than 0.6 (as done by Anand et al. (2019)).
The full results of training Atari RAM transition models, for the games studied in Anand et al. (2019), are provided in Table 2. We evaluate both bit-level F1 scores, as well as slot-level accuracy (for which all 8 bits need to be predicted correctly in order to count), over the remaining 48 Agent57 episodes as validation. To the best of our knowledge, this is the first comprehensive feasibility study for learning Atari RAM transition models.

heading: C REDISCOVERING THE ALGORITHMIC BOTTLENECK
text: As mentioned in the main text body, one of the key reasons in favour of a high-dimensional algorithmic component is to avoid the algorithmic bottleneck, as first exposed by Deac et al. (2021).
In short, the performance guarantees of running classical algorithms rely on having the exactly correct inputs for them. If there are any errors in the predictions of these (usually very low-dimensional) inputs, these errors may propagate to the algorithmic computations and yield suboptimal results. Further, there is no room for any kind of fallback if such an event occurs.
In contrast, the high-dimensional neural processors like the ones we study here are not vulnerable to bottleneck effects: if any dimensions of the latent state are poorly predicted, the other components of it could step in and compensate for this. Further, we can easily support skip connections in the case where the algorithm is not fully descriptive of the problem we’re solving.
In this section, we reaffirm the bottleneck effect for both Atari and bouncing ball experiments by providing additional sets of ablations on the processor network’s latent size.
We ablate various RMR processor network architectures, for dim z ∈ {2, 4, 8, 16, 32, 64}, but leaving all other components and operations unchanged.
The results of this ablation for the Atari representation learning setting are provided in Figure 5. It can be clearly observed that, as we reduce the size of the latents, this typically induces a performance regression in the downstream bit F1 scores. This indicates the algorithmic bottleneck effect.
On the bouncing balls task, we observe the algorithmic bottleneck as well, with more pronounced effects. Namely, for all latent sizes 32 and under, the validation MSE shoots up even though the training MSE keeps improving. Ultimately, we also attempted using an Interaction Network (Battaglia et al., 2016) as a bottlenecked RMR processor which requires projecting on to x̄ rather than z, which exhibited exactly the same behaviour, at a larger scale. The IN processor achieved a final validation MSE of 341.40± 5.60 (×10−4). This is roughly 43× worse than the non-bottlenecked RMR.

heading: D ON THE WEAK ACCEPTOR PROPERTIES OF BATTLEZONE
text: Out of all the Atari games studied in our work, Battlezone could been singled out as the least “receptive” to RMR processors—with only four games successfully donating their RAM representations to its pixel based encoder (Figure 4).
We set out to study why this effect took place. Upon inspection of the game’s dynamics1, we determined that Battlezone has certain properties that make representation learning uniquely challenging and somewhat decoupled from the underlying console computation.
Namely, Battlezone is the only game in our dataset that is played in the “first-person”. Therefore, from the point of view of the pixel inputs, it may seem as if the player is always in the same place, and we would expect the mechanics and the way of thinking about the ’avatar’ to be substantially different from all other Atari games considered.

heading: E ON THE DATA DISTRIBUTION AND SETUP USED FOR TRAINING P
text: In many cases, while an algorithmic prior may be known, generative aspects of the relevant abstract inputs for the natural task could be unknown. For example, it may be known that the natural task could benefit from a sorting subroutine, but it may not be known upfront how many objects will need to be sorted.
When such details are unknown, it may still be possible to fall back to abstract inputs sampled from some sensible generic random distribution—so long as the processor network is trained in a way that promotes extrapolation. For a recent theoretical treatment of OOD generalisation in algorithmic reasoners, we refer the reader to Xu et al. (2020). Further, as observed by Deac et al. (2021) in the case of implicit planning, even generic random abstract distributions can promote useful transfer to noisy pixel-based acting settings such as Atari.
Lastly, it is important to ensure that, in the abstract pipeline, the processor network P carries the brunt of the computational burden; otherwise, the reasoning task may be partially captured in either f or g, and not left to P ’s weights. In our work, we generally promote such behaviour by keeping f and g to be only linear layers; but even in settings where this is not appropriate, we recommend taking care to not overpower the abstract encoder and decoder.

heading: F ATARI ABSTRACT MODEL EQUATIONS
text: In this appendix, we provide a “bird’s eye” view of the modelling steps taken by our abstract Atari pipeline, aiming to support future implementations of the RMR blueprint. We will readily re-use the notation from the main text for the various components.
1https://www.youtube.com/watch?v=9X4_xy7rC1A
Firstly, the abstract encoder f is applied on the relevant Atari RAM representations, x̄, augmented by a one-hot representation which is aiming to provide a generic embedding of the semantics of each RAM slot. f is implemented as a linear layer, hence:
z = f(x̄‖o) = Wf x̄ + Ufo + bf (1)
where Wf ,Uf ,bf are learnable weights, and o ∈ B128×128 is a one-hot encoding s.t. o = I128. In a separate pipeline (officially part of the processor network), the performed actions are also encoded using a linear action encoder:
α = fa(a) = W faa + bfa (2)
where a is a one-hot encoded action representation.
Then, the processor network GNN is called to update these latents, and a sum-based skip connection is employed to promote a model-free path:
z′ = P (z +α) + z +α (3)
Here, the processor network P is implemented as a standard message passing neural network (Gilmer et al., 2017) over a fully connected graph. Equations of such a P are commonly exposed in e.g. Veličković et al. (2019).
Lastly, the relevant decoder networks predict two properties: a mask which signifies which RAM slots have been changed as a result of applying a, and the updated states ȳ.
µ = gµ(z ′) = σ (Wµz′ + bµ) (4) ȳ = g(z′) Iµ>0.5 = (Wgz′ + bg) Iµ>0.5 (5)
Here, σ is the logistic sigmoid activation, is the elementwise product, and I is the indicator function, which thresholds the mask. While this thresholding is non-differentiable, we note that the mask values can be directly supervised from known trajectories, and at training time, teacher forcing can be applied—slotting ground-truth masks in place of the indicator function.

Review summaries:
This paper contributes an implementation of the neural algorithmic reasoning blueprint outlined in Velickovic & Blundell (2021). The setting that is considered is one where, for a given complex reasoning task that one is interested in solving, there is access to information about the underlying system that should be incorporated to simplify learning. For example, in the case of making predictions from videos about the future state of a physical system composed of bouncing balls bouncing balls, given access to each ball's radius, position, and velocity it is well known how their future movements follow the laws of physics.   In order to incorporate this prior knowledge, Velickovic & Blundell (2021) proposed the idea of training a neural processor to capture this prior knowledge between abstract inputs (eg. the 3-dimensional state of each ball) under the assumption that sufficient data obeying the known input-output mapping in this space can be generated (eg. pairs of states at time t and time t+1). An encoder then maps abstract inputs to high dimensional space, the processor transforms it to another high-dimensional space, and a decoder maps it back onto the corresponding abstract output pair. In this way, the processor can distill knowledge about the transformation between abstract inputs (eg. governing the laws of physics in its weights) into its weights. By incorporating this processor in a network that is trained to solve the original task (eg. prediction in pixel spaces) it is expected that this extracted prior knowledge can be leveraged to improve learning and transfer to new tasks.  The implementation considered here (Reasoning-Modulated Representations), is evaluated on two environments: the bouncing balls environment, and various Atari games. On bouncing balls it is shown how RMR can transfer learned representations from trajectories to videos. However, the absence of a stronger baseline makes it difficult to interpret these results. On Atari, RMR is trained on transitions of RAM states to improve predicting latent space transitions from pixels. Compared to C-SWM (Kipf et al., 2019), it is shown how RMR leads to better representations in several games, and it is shown how processors may transfer across games.  

This paper proposes a training method based on ideas on neural algorithmic reasoning. As well using a model architecture that is aligned with the underlying algorithmic process this method proposes to pre-train the reasoning module using data from an abstract output stimulator (the work considers problems where this simulator is readily available, and so can generate lots of abstract data cheaply). A key motivation is that the reasoning module $P$ should ultimately be a mapping between two medium dimensional latent spaces (lower dimension than input space, higher dimension than the abstract representation space). Finding this balance allows the module to avoid modeling high dimensional data, whilst also not introducing low-dimensional bottlenecks into the architecture.  This work finds that inserting the pre-trained $P$ into a model that processes raw input data leads to improved performance compared to end-to-end training of a model on raw input data. 

This paper presents a learning paradigm referred to as "Reasoning-Modulated Representation" (RMR). The main argument is that it first learns a latent-space processor for abstract simplification of a problem, then use the processor for any tasks that can be modeled with algorithmic reasoning. The authors then verify RMR in two scenarios, bouncing balls problem (physics simulation) and Atari 2600 games.  For real-world problems, the proposed method alleviates requirements on alignment between abstract and natural inputs, making the learning pipeline more data-efficient, because the pipeline could rely solely on automatically generated, abstract training data. In addition, it is possible that the latent-space processor can be re-used to tasks that do not even directly align with the abstract environment.