Spatially Structured Recurrent Modules
Authors: Nasim Rahaman
Anirudh Goyal
Muhammad Waleed Gondal
Manuel Wuthrich
Stefan Bauer
Yash Sharma
Yoshua Bengio
Bernhard Schölkopf
Venue: ICLR 2021 Poster (2021)
Keywords: spatio-temporal modelling
modular architectures
recurrent neural networks
partially observed environments

Abstract: Capturing the structure of a data-generating process by means of appropriate inductive biases can help in learning models that generalise well and are robust to changes in the input distribution. While methods that harness spatial and temporal structures find broad application, recent work has demonstrated the potential of models that leverage sparse and modular structure using an ensemble of sparingly interacting modules. In this work, we take a step towards dynamic models that are capable of simultaneously exploiting both modular and spatiotemporal structures. To this end, we model the dynamical system as a collection of autonomous but sparsely interacting sub-systems that interact according to a learned topology which is informed by the spatial structure of the underlying system. This gives rise to a class of models that are well suited for capturing the dynamics of systems that only offer local views into their state, along with corresponding spatial locations of those views. On the tasks of video prediction from cropped frames and multi-agent world modelling from partial observations in the challenging Starcraft2 domain, we find our models to be more robust to the number of available views and better capable of generalisation to novel tasks without additional training than strong baselines that perform equally well or better on the training distribution. 

Published as a conference paper at ICLR 2021
SPATIALLY STRUCTURED RECURRENT MODULES
Nasim Rahaman 1.2
Anirudh Goyal?
Muhammad Waleed Gondal Manuel Wuthrich
Stefan Bauer*
Yash Sharma
Yoshua Bengio
Bernhard Schölkopf
ABSTRACT
Capturing the structure of a data-generating process by means of appropriate inductive biases can help in learning models that generalize well and are robust to changes in the input distribution. While methods that harness spatial and temporal structures find broad application, recent work has demonstrated the potential models that leverage sparse and modular structure using an ensemble of sparingly interacting modules. In this work, we take a step towards dynamic models that are capable of simultaneously exploiting both modular and spatiotemporal structures. To this end, we model the dynamical system as a collection of autonomous but sparsely interacting sub-systems that interact according to learned topology which is informed by the spatial structure of the underlying system. This gives rise to a class of models that are well suited for capturing the dynamics of systems that only offer local views into their state, along with corresponding spatial locations of those views. On the tasks of video prediction from cropped frames and multi multi-agent world modeling from partial observations in the challenging Starcraft2 domain, we find our models to be more robust to the number of available views and capable better generalization to novel tasks without additional training than strong baselines that perform equally well or better on the training distribution.
INTRODUCTION
Many spatiotemporal complex systems can be abstracted collection of autonomous but sparsely interacting sub-systems, where sub-systems tend to interact if they are in each others' vicinity. As an illustrative example, consider grid of traffic intersections. Traffic flows from a given intersection to the adjacent ones, and the actions taken by some "agent", say an autonomous vehicle, may at first only affect its immediate surroundings.
Now suppose we want to forecast the future state of the traffic grid (say for the purpose of avoiding traffie jams). There is spectrum of possible strategies for modeling the system at hand. On one extreme lies the most general strategy which considers the entirety of all intersections simultaneously to predict the next state of the grid (Figure Tc.. The resulting model class can in principle account for interactions between any two intersections, irrespective of their spatial distance. However, the number of interactions such models must consider does not scale well with the size of the grid, and this strategy might be rendered infeasible for large grids with hundreds of intersections. On the other end of the spectrum a strategy which abstracts the dynamics of each intersection as an autonomous sub-system, with each sub-system interacting only with its immediate neighbors (Figure Ta). The interactions may manifest as messages that one sub-system passes to another and possibly contain information about how many vehicles are headed towards which direction, resulting in a collection of message passing entities (i.e. sub-systems) that collectively model the entire grid. By adopting this strategy, assumes that the immediate future of any given intersection is affected only by the present states of the neighboring intersections, and not some intersection at the opposite end of the grid. The resulting class of models scales well with the size of the grid, but is possibly unable to model certain long-range interactions that could be leveraged to efficiently distribute traffic flow.
The spectrum above parameterizes the extent to which the spatial structure of the underlying system informs the design of the model. The former extreme ignores spatial structure altogether, resulting
Max-Planck Institute for Intelligent Systems Tübingen, Mila,, Quétee, Bethgelab, Eberhard Karls Univer- sität Tübingen, "Université de Montreal. Correspondence <nasim. rahaman@tuelt inginnnnnn.. mpg.de>.

---

Published as a conference paper at ICLR 2021
(a) Fully localized sub-systems.
(b) Middle ground.
(c) Single, monolithic system.
Figure I: A schematic representation of the spectrum of modeling strategies. Solid arrows with speech bubbles denote (dynamic) messages being passed between sub-systems (dotted arrows denote the lack thereof). Gist: on end of the spectrum, (Figure Iap. we have the strategy of abstracting cach intersection interact with neighboring sub-systems. On the other end of the spectrum (Figure Ic] we have the strategy of modeling the entire grid with one monolithic system. The middle ground (Figure [by we explore involves letting the model develop a notion of locality by (say) abstracting entire avenues with single sub-system.
sub-system that
in a class of models that can be expressive but whose sample and computational complexity do not scale well with the size of the system. The latter extreme results in class of models that can scale well, but its adequacy (in terms of expressivity) is contingent on predefined notion of locality (in the example above: the immediate four-neighborhood of an intersection). In this work, we aim to explore a middle -rrouud between the two extremes: namely, by proposing class of models that learns notion of locality instead of relying on a predefined one (Figure [bb]. Reconsidering the traffic grid example: the proposed strategy results in model that may learn to abstract (say) entire avenues with a single sub-system, if it is useful towards solving the prediction task. This yields a scheme where single sub-system might account for events that are spatially distant (such those in the opposite ends of an avenue), while events that are spatially closer together (like those on two adjacent avenues of the same street, where streets run perpendicular to avenues) might be assigned to different sub-systems.
To implement this scheme, we build framework wherein the sub-systems are modelled as independent recurrent neural network (RNN) modules that interact sparsely via a bottleneck of attention (a variant of which is explored in Goyal et al. (2019)) while extending it along two salient dimensions. First, we learn an interaction topology between the sub-systems, instead of assuming that all sub-systems interact with all others in an all-to-all topology. We achieve this by learning to embed each sub-system in space endowed with metric, and attenuate the interaction between two given sub-systems according to their distance in this space (i.e., sub-systems too far away from each other in this space are not allowed to interact). Second, we relax a common assumption that the entire system perceived simultaneously; instead, we only assume access to local (partial) observations alongside with the associated spatial locations, resulting in setting that partially resembles that of Eslami et al. (2018). Expressed in the language of the example above: we do not expect a bird's eye view of the traffic grid, but only (say) LIDAR observations from autonomous vehicles at known GPS coordinates, or video streams from traffic cameras at known locations. The spatial location associated with an observation plays a crucial role in the proposed architecture in that map it to the embedding space of sub-systems and address the corresponding observation only to sub-systems whose embeddings lie in close vicinity. Likewise, to predict future observations at a queried spatial location, we again map said location to the embedding space and poll the states of sub-systems situated nearby. The result is a model that can learn which spatial locations are to be associated with each other and be accounted for by the same sub-system. As an added plus, the parameterization we obtain is not only agnostic to the number of available observations and query locations, but also the number of sub-systems.
To evaluate the proposed model, we choose a problem setting where (a) the task is composed of different sub-systems or processes that locally interact both spatially and temporally, and (b) the environment offers local views into its state paired with their corresponding spatial locations. The challenge here lies in building and maintaining a consistent representation of the global state of the system given only a set of partial observations. To succeed, a model must learn to efficiently capture the available observations and place them in an appropriate spatial context. The first problem consider is that of video prediction from crops, analogous to that faced by visual systems of many animals: given a set of small crops of the video frames centered around stochastically sampled pixels (corresponding where the fovea is focused), the task is to predict the content of crop
2

---

Published as a conference paper at ICLR 2021
around any queried pixel position at future time. The second problem is that of multi-agent world modeling from partial observations in spatial domains, such as the challenging Starcraft2 domain (Samvelyan et al. 2019 Vinyals al. 2017). The task here is to model the dynamics of the global state of the environment given local observations made by cooperating agents and their corresponding actions. Finally, also include visualizations on multi-agent grid- world environment designed for simulating railroad traffic (Eichenberger et al. 2019). Importantly and unlike prior work (Sun et al., 2019), our parameterization is agnostic to the number of agents in the environment, which can be flexibly adjusted on the fly as new agents become available or existing agents retire. This is beneficial for generalization in settings where the number of agents during training and testing are different.
Contributions. (a) We propose new class of models, which we call Spatially Structured Recurrent Modules or S2RMs, which perform attention-driven modular computations according to a learned spatial topology. (b) We evaluate S2RMs (along with several strong baselines) on selection of challenging problems and find that S2RMs are robust the number of available observations and can generalize to novel tasks.
PROBLEM STATEMENT
In this section, we build on the intuition from the previous section to formally specify the problem we aim to approach with the methods described in the later sections. To that end, let X' a metric space, O some set of possible observations, and Ox a set of mappings X' - O. Now, consider the evolution function of a discrete-time dynamical system:
Φ Z x Dx - Dx satisfying:
(1)
(00, o) = where o E xx and 6(22, (11, Φ(11 + (2) for 1,,2 E Z
Informally, can be interpreted as the world state of the system: together with spatial location x E X, it gives the local observation O == o(x) E O. Given an initial world state o, the mapping ((1, o) yields the world state o, at some (future) time I, thereby characterizing the dynamics of the system (which might be stochastic). The problem we consider is the following:
Problem: At every time step O,... T, we given set of positions {x]] and the corre- sponding observations [ where O == 0, (x))) The task is to infer the world state or at some future time-step t' T in order to predict O ou (x*) at some arbitrary query position x²,
In the traffic grid example of Section one could imagine a as indexing traffic cameras or autonomous vehicles (i.e., observers), xi the GPS coordinates of observer a, and Of the corresponding sensor feed (e.g. LIDAR observations or video streams from vehicles or traffic cameras).
3 MODELLING ASSUMPTIONS
Given the problem in Section we now make certain modelling assumptions. These assumptions will ultimately inform the inductive biases we select for the model (proposed in Section 4; nevertheless, we remark beforehand that as with any inductive bias, their applicability is subject to the properties of the system being modeled and the objectives being optimized (OOD generalization, etc).
Recurrent Dynamics Modeling, While there exist multiple ways of modeling dynamical systems, we shall focus on recurrent neural networks (RNNs). Typically, RNN-based dynamics models are expressed as functions of the form:
h ((O,,h.)
O, -(hb)
where O is the observation at time E Z, and h,41 is the hidden state of the model. F can be thought of as the parameterized forward-evolution function the hidden state h conditioned on the observation O, whereas D is decoder that maps the hidden state to observations.
Decomposition into Locally Interacting Sub-systems. We make the assumption that the dynamical system Ø can be decomposed to constituent sub-systems (ø1, .2.,.. @M) that dynamically and sparsely interact with each other while respecting some interaction topology. By interaction topology, mean that each module 6; can be identified with embedding p, in topological space S equipped with a similarity kernel Z and that the sub-system &; may preferentially interact with another sub-system oj if their respective embeddings are close in S with respect to Z, i.e. if
3

---

Published as a conference paper at ICLR 2021
Set of Interacting RNNs
Encoder
Decoder
Observation
P
Predicted Observation Querled Location
Location
Input Attention
Output Attention
Observation
Query Location
Location
F4
Positional
Positional
Embedding
Embedding
Figure 2: Schematic illustration of the proposed architecture. An observation is addressed to modules with embeddings situated in vicinity of its embedded location. Likewise, modules with embeddings in the vicinity of an embedded query location are polled to produce prediction.
(p., Py) is large. Intuitively, one may think of Z as inducing a notion of locality between sub- systems, according to which Øj lies in the local vicinity of Φi-
Locality of Observations. The notion of locality between sub-systens induced by Z is distinct from that induced by the metric of space X' of locations in the environment (cf. Section 2). and one important modelling decision is how these two should interact. We propose to embed the position X' associated with observation O to the metric space of sub-systems S via continuous and injective mapping P: - S. This allows us to match the observation O to all sub-systems om that are in the vicinity of P(x) E S, i... where (x((),, Pm) sufficiently large. Each subsystem &n therefore accounts for observations made at a set of positions Xm. C X, which we call its enclave.
4 SPATIALLY STRUCTURED RECURRENT MODULES (S2RMs)
Informed 2 by the model assumptions detailed in the previous section, we now describe the proposed model (Figure 2 comprising the following components:
Model Inputs. Recall from Section 2 that we have for every time step t = 0,..., T a set of tuples of positions and observations {xx] ) -==1 where x;² X and Of E O for all and a. To simplify, we assume that X C R", and denote by Ię the i-th component of the vector x E X.
The encoder E is a parameterized function mapping observations O to a corresponding vector representation e = (O)... Here, E processes all observations in parallel across t. and a to yield representations ef.
Encoder.
Positional Embedding. The positional embedding P is a fixed mapping from to S. We choose S to be the unit sphere in d-dimensions, d being multiple of 2n, and the positional encoder the following function:
P(x) = s/11=11 E S where (8i+m, 8i+14m) (sin (5m/10000'). cos (zm/110000))
with 0, ..., n 1 and i 0, ..., 4/1 1. The above function finds common use Vaswwnn
eet al. 2017 Mildenhall et al. 2020 Zhong et al. 2020) and can be motivated from the perspective of
Reproducing Kernel Hilbert Spaces (Rahimi & Recht 2007) (see Appendix C.3 for a discussion).
We henceforth refer to P(x) as s and P(x;) as sç.
Set of Interacting RNNs. To model the dynamics of the world state, we use set of M independent RNN modules, which we denote as J To each Fm, we associate an embedding vector p" E S, where all (pm) 71M Jm=1 are learnable parameters. The RNNs F,, interact with each other via an inter-cell attention, and with the input representations cf, via input attention. Precisely, at a given time step I, each Fm expects an input u,", an aggregated hidden state hy and optionally, a memory state ct" to yield the hidden and memory states at the next time step:
(h 1 ) = Fm (u," m c.
(4)
1 'The continuity of P ties the notions of locality by requiring that an infinitesimal change corresponds to one S. Injectivity ensures that no two points X are mapped to the same point S. In doing so, we the assumptions merely as guiding principles; we do not claim that we infer the true decomposition of the ground-truth system, even if all assumptions satisfied.
4

---

Published as a conference paper at ICLR 2021
where the input ut" results from the input attention and h." from the inter-cell attention (see below).
Kernel Modulated Dot-Product Attention. A central component of the proposed architecture is the kemel modulated dot-product attention (KMDPA), which now define. First, we let Z S x S [0, 1| be the following kernel:
[x -p -p.s)] if 2 0,
(p..) s)
(5)
otherwise
where (0, 0o) is the kernel bandwidth, and T E 1)), is the truncation parameter (additional details in Appendix C.I.p. Now, KMDPA maps two sets A and B to a third set A, where:
A = {(aaa, DiE B = [[bb, z') z') ==((c {(a", ((a""""")) "
KMDPA B)
(6)
Here, a', b' E S,, and y", z are vectors of not necessarily the same dimension. In order to evaluate A, we first compute the interaction weights Wi between any two pairs of entities y') and
(b', z²), which depends local term W.f.) and non-local term Waz. We have:
Wij softmax, (Q(Quer)(y')) . (((e)) (2))): wh( = Z(a b)); W5 = Wo, W
(7)
where (Querry) and (Key) are learnable linear mappings that project yf and x¹ to the same space.
The penultimate step computes the following two quantities:
Σ, W if yd;
Ei
Wo Value) (z')
where (Value) is another learnable linear function mapping from the vector space of z to that of Finally, we have:
y = Gly y") y' + (1-
where G is a gating layer with sigmoid non-linearity implementing a soft selection mechanism between the linear combination of values (y') and the inputs weighted by the local weights (ý²). In what follows, we will refer to the set A as query set, B as key set and Å as output set.
Input Attention. The input attention mechanism is KMDPA, mapping between sets of observation tuples [s] et
14 ) Ja=1 (key set) and the current RNN-states {(DT ) (query set) to that of RNN inputs {PD", u, ) (output set). Now the one hand, we observe that the input u," to RNN Fm can contain information about an observations of only if the embedded location of the said observation SÊ is close enough to the embedding of the RNN p" in S (i.e. if Z(s", p't) > 0), thereby implementing the assumption of locality of observations. On the other hand, the non-local term allows module Fm to reject (or accept) an observation based on its content, which can be beneficial if two modules attend to overlapping regions in the environment but specialize to different aspects of the dynamics. Please refer to Appendix C.1 for precise description of the mechanism, in particular the (optional) use of multiple dot-product attention heads.
Inter-cell Attention. The intercell attention mechanism is another KMDPA, mapping two copies of the current RNN-states {p"". .h." J (one as query and another as key set) to the set of aggregated hidden states {P"" h |n= (output set). This enables local interaction between the RNNs Fm. in that the local term ensures that RNN F.. interacts with RNN F, only if their respective embeddings Td and p" are close enough in S (i.e. if (p" Th p²) > 0), thereby implementing the assumption of local interactions between sub-systems. The non-local term allows two modules to interact with each other based on their hidden states, i.c. it provides the mechanism for a module to (not) interact with another other based on their respective states, even if their embeddings are similar enough in S. Appendix C.2 contains a precise description of the attention mechanism.
Output Attention. The output attention mechanism together with the decoder (described below) serve as an apparatus to evaluate the world state modeled implicitly by the set of RNNs ((Fm)) M at time t + 1 (for one-step forward models). Given query location x X and its corresponding embedding sº, the output attention mechanism polls the RNNs F,,, whose embeddings nd are similar enough to s", as measured by the kernel Z. Denoting by hmj the j-th component of h and by dy the j-th component of the vector d associated with the query location x", we have:
d3 = Σ Z(s 2(s".p")) Pm') hims
(10)
m
5

---

Published as a conference paper at ICLR 2021
Decoder. The decoder D is parameterized function that predicts the observation 02+, O at x given the representation d 1 from the output attention.
This concludes the description of the generic architecture, which allows for flexibility in the choice of the RNN architecture (i.e., the internal architecture of Fm). In practice, we find Gated Recurrent Units (GRUs) (Cho et al.) 2014) to work well, and call the resulting model Spatially Structured GRU or S2GRU. Moreover, Relational Memory Cores (RMCss) (Santoro al. 2018) also profit from our architecture (with a modification detailed in Appendix E.3) and we call the resulting model S2RMC.
5 RELATED WORK
Recall that the problem setting we consider is one where the environment offers local (partial) views into its global state paired with the corresponding spatial locations. With Generative Query Networks (GQNs), Eslami et al. (2018) investigate a similar setting where the 2D images of 3D scenes are paired with the corresponding viewpoint (camera position, yaw, pitch and roll). Given that GQNs are feedforward models, they do not consider the dynamics of the underyling scene and such cannot be expected to be consistent over time (Kumar al.) 2018). (Singh et al.) (2019) and Kumar et al. 2018) propose variants that are temporally consistent, but unlike us, they do not focus on the problem of predicting the future state of the system.
Problem Setting.
Modularity. Modularity has been a recurring topic in the context of meta- meta-learning (Alet et al. Bengio al. 2019; Ke et al. (2019), sequence modeling (Ghahramani & Jordan [996] Henaff et al. 2016 Li et al. [2018] Goyal el al.) 2019; Mei et al. 2020
2018
Mittal et al. (2200) and beyond (Jacobs
ct al. 1991] Shazeer al. [017]] Parascandolo et 2017 In the context of RNNs, ei al. (2018)
explore setting where the recurrent units operate entirely independently of each other. Closer to our work, Goyal al. (2019) explores the setting where autonomous RNN modules interact with each other via the bottleneck of sparse attention. However, instead of leveraging the spatial structure of the environment, they induce sparsity using a scheme inspired by the k-winners-take-i principle (Majani et al. 1988) where only the k modules that attend the most to the input are activated and propagate their state forward, whereas the remaining modules that do not receive an input follow default dynamics in that their hidden states are not updated. This can be contrasted with S2RMs, where the modules that do not receive inputs may still evolve their states forward in time, reflecting that the environment may evolve even when observations are available.
Attention Mechanisms and Information Flow, Attention mechanisms have been used to attenuate the flow of information between components of the network, e.g. (Graves et al.) 2014 2016 [Santoro iet al. 2018 Ke et al. 2018 Veličković et al.) 2017 [Battaglia et 2018)) There is growing interest in efficient attention mechanisms for use in transformers (Vaswani et al. 2017). and like KMDPA, some recently proposed methods rely on learned sparsity (Kitaev et al. 2020 Tay et al. 2020).
However, these induce sparsity by dynamically clustering or sorting based on content, while we make explicit use of the spatial information accompanying observations to learn a ppatially-grounde
sparsity pattern. Moreover, mechanisms for spatial attention have also been studied (Jaderberg et al., 2015. Wang et al. [2017] Zhang et al. 2018 Parmar et al. (018)), but they typically operate on image pixels. Our setting is different in that we do not assume that the world-state (from which we sample local observations) can be represented as an image.
EXPERIMENTS
In this section, we present selection of experiments to quantitatively evaluate S2RMs and gauge their performance against strong baselines on two data domains, namely video prediction from crops on the well-known bouncing-balls domain and multi-agent world modelling from partial observations in the challenging Stareraft2 domain. We also include qualitative visualizations on grid-worlh world task in Appendix A Additional tables, results and supporting plots can be found in Appendix F
Baselines. To draw fair comparisons between various RNN architectures, we require an architec- tural scaffolding that is agnostic to the number of observations A, is invariant to the ordering of {(x], D with respect to a and features a querying mechanism to extract a predicted observation O4 at a given query location x" in future time-step t' > Fortunately. it is possible to obtain performant class of models fulfilling our requirements by extending prior work on Generative Query Networks GQNs (Eslami et al. 2018). The resulting model has three components: encoder,

---

Published as a conference paper at ICLR 2021
.0,
Figure 3: Rollouts (OOD) with bouncing balls, from top to bottom: ground-truth, S2GRU, RIMs, RMC, LSTM. Note that all models were trained on sequences with 3 bouncing balls, and the global state is reconstructed by stitching together 16 patches of size 11 x 11 produced by the models (queried on 4 x grid). Gist: S2GRU succeeds at keeping track of all bouncing balls over long rollout horizons (25 frames).
1.00
1.00
Binery F1-Score
Balanced Accurscy
0.95
0.90
0.90
Training Distr.
0.80
0.85
TTO
TTO
RMC
0.80
RMC
0.70
0.75
Number of Bouncing Balls
Number of Bouncing Balls
Figure Performance metrics on OOD one-step forward prediction task. Gist: S2GRU outperforms all RNN baselines OOD.
RNN, and decoder, which deseribe in detail in Appendix D. In experiments, fix the encoder and decoder to be essentially identical to those in S2RMs, but vary the architecture of the RNN, where experiment with LSTMs (Hochreiter & Schmidhuber 1997), RMCs (Santoro et al. 2018) and RIMs (Goyal et al. 2019). As a sanity check, we also show results with a Time Travelling Oracle (TTO), which time-step has access to the (partially observed) state at + Its purpose to verify that the architectural scaffolding around the baseline RNNs (defined in Appendix D) does not constrain their performance and that the comparison S2RMs is indeed fair.
Video Prediction from Crops.
We consider the problem of predicting the future frames of simulated videos of balls bouncing in a closed box, given only crops from the past video frames which are centered at known pixel positions. Using the notation introduced in Section at every time step I,, we sample
pixel positions (x( Jaz1 from the 1-th full video frame o, of size 48 48. Around the sampled central pixel positions x", we extract 11 x 11 crops, which we use the local observations Og The task now to predict 11 11 crops O% corresponding to query central-ixell-ositiie at future time-step E' > t. Observe that at any given time-step the model has access to at most 52% the global video frame assuming that the crops never overlap (which rather unlikely).
x".
Having trained on the training dataset with 3 bouncing balls, evaluate the forward- preerddprediction performance on all test datasets with to bouncing balls. Given that we treat the predic- tion problem pixel-wise binary classification problem, report the balanced accuracy (i.e. arithmetic mean recall and specificity) or Fi-scores (i.e. harmonic mean of preci- sion and recall) account for class-imbalance. In Figure see that S2GRUs out-perform
Figure 5: Visualization the spatial locations each module is responsible for modeling (i.e. the enclaves Xm. defined in Section B). The central ball does not bounce, i.e. it is stationary in all sequences. Gist: the modules focus attention on challenging regions, e.g. the corners of the arena and the surface of the fixed ball.
7

---

Published as a conference paper at ICLR 2021
Tember
Oreinal
without Local Torn at rputton
Non-Lccll form wonn
53.4444 CACA 2214
Fracken CniArodaae Fraciee
Slow Eidlo
Avoilaada
Figure 6: Ablation over the local and the non-local terms in the input and inter-cell attention mechanisms (KMDPAs). For set number of bouncing balls, each sub-plot shows how the balanced accuracy changes with the fraction of views (crops) available to the model. Gist: Both local and non-local terms in KMDPA contribute to the overall performance. The non-local term is more important for the input attention, whereas the local term is more important for the inter-cell attention.
STORP
KIORP
T1
ETD
-0.3
LTT
-120
-2.1
-025
-335
50
2.4
3.55
9.5 CA
-0415
244
02
Agent
Figure 8: Performance metrics (larger the better) as function of the probability that an agent will not supply information to the world model but still query it. Gist: while all models lose performance as fewer agents share observations, we find S2RMs to be most robust.
all non-oracle baselines on the one-step forward prediction task and strike a good balance with regard to in-distribution and OOD performance. In Figure 3 we qualitatively show re- constructions from 25 step rollouts on the out of distribution dataset with 5 balls to demon- strate that S2GRUs can perform OOD rollouts over long horizons without losing track of balls. Figure shows the result of an ablation study where we disable the local and non-local terms in each of the two KMDPAs while keeping everything else same. We see that both local and non-
local terms contribute the overall performance; moreover, the input attention relies on the non-local term and the performance severely affected by its absence, whereas the inter-cell attention is dependent on the local-term to yield good performance. This suggests that the modules indeed rely on the content of the input observations as they select their inputs, and learning an interaction topology between modules is strong contributor to the final performance. In Figure 5 we show for each module its corresponding enclave, which is the spatial region that is responsible for modelling. i.e. for pixels at position x, we plot (Z(((x) p (cf. Section[4. 4 We find that the modules learn to share the responsibility of modelling the entire spatial domain. Finally, in Figure (we we see the effect of removing (ran- domly sampled) modules at test time, i.c. without additional retraining. The performance degrades gracefully as fewer mod- ules are available, suggesting that the individual modules can function while other modules are missing. We include details and additional results in Appendix F.I
Batanced Accurscy
Batanced
0.7
0.6
0.5
Figure 7: The effect of removing ran- dom modules al test time. Gist: Perfor- mance degrades gracefully modules are removed, suggesting that modules can function even when their counter- parts are removed, and that there is lim- ited co -daapttaionn between them.
Multi-Agent World Modeling on Starcraf12. In Section 2 we formulated the problem of modeling what we called the world state of dynamical system @ given local observations {(x,, O₂) d where O = o(t, (x?). Under certain restrictions, this problem can be mapped to that of multi- agent world modeling from partial and local observations, allowing us to evaluate the proposed model in rich and challenging setting. In particular, we consider environments that are (a) spatial, i.e. all agents a have a well-defined and known location xf (at time t), (b) the agents' actions uf are local, in that their effects propagate away (from the agent) only at finite speed, (c) the observations local and centered around agents, in the sense that the agent only observes the events in its local vicinity, i.e., O Observe that we do not fix the number of agents in the environment, and allow
8

---

Published as a conference paper at ICLR 2021
for agents to dynamically enter or exit the environment. Now, the task is: given observations O from a team of (cooperating) agents at position x and their corresponding actions uy, predict the observation Og that would be made by agent at time t' t if were at position x".
Starcraft2 unit- micromanagement (Samvelyan et al. 2019) is a multi-agent reinforcement learning benchmark, wherein teams of hetero- geneously typed units must defeat a team of opponents in melee and ranged combat.
UT-FI
FM-FI
NMSE
LL
(1s2z)
LSTM
0.6267
0.8464
-0.0040
-0.0382
RMC
0.6839
0.8597
-0.0033
-0.0334
S2GRU
0.7488
0.8627
-0.0023
-0.0233
The observations Of and actions už are both multi-channel images represented in polar co- ordinates centered around the agent position x²_ The field of view (FOV) of each agent is therefore a circle of fixed radius centered around it. The channels of the image cor- respond to (a) a binary indicator marking whether position in FOV is occupied by a living friendly agent (friendly marker), (b) a categorical indicator marking the type of liv- ing units at a given position in FOV (unit-rype marker), and (c) four channels marking the health, energy, weapon-cooldown and shields (HECS markers) of all agents FOV. With a heuristic, we gather a total of 9K trajectories ([xf, Of, u, Ja=1 1 spread over three train- ing scenarios, corresponding to 1 1c3s52 3s52 and 2s52 in Samvelyan et al. 2019)). We also sample 1K trajectories (each) from two OOD scenarios 1s2z and 5s3z. Details in Appendix B.
S2RMC
0.7317
0.8563
-0.0026
-0.0261
(TTO)
0.7518
0.8883
-0.0025
-0.0259
(5s32)
LSTM
0.4975
0.7123
-0.0134
-0.1251
RMC
0.5414
0.7486
-0.0132
-0.1167
S2GRU
0.5310
0.7058
-0.0119
-0.1108
S2RMC
0.5114
0.6945
-0.0124
-0.1205
(TTO)
0.6115
0.7872
-0.0107
-0.0940
Table 1: Performance metrics on OOD scenarios 1s22 and 5s32 (larger numbers are better): unit- type macro Fl score (UT-FI), friendly-marker F1 score (FM-F1), HECS Negative Mean Squared Error (NMSE) and Log Likelihood (LL).
Having trained all models on scenarios 133552, 3s5z and 2s52, we test their robustness to dropped agents (Figure 8 and their performance on OOD scenarios (Table We only include baselines that achieve similar or better validation scores than S2RMs. Figure shsuwss shows that S2RMCs remain robust when fewer agents supply their observations to the world model, whereas Table I shows that S2GRUs outperforms the baselines in the OOD scenario 1 ss 22 but is matched by RMCs in 5s32 (see Appendix F2 for details). The strong performance of RMCs suggests that the task benefits from the inductive bias of relational memory. One hypothesis to why is that the pace of the considered environments requires fast communication between agents, which can be achieved by shared memory where all agents may read from and write to. Further, we observe that while the oracle (TTO) can generalize well out of distribution, siowrs that it is less robust to the number of available observations. This is explained by the fact that unlike recurrent models, TTO does not leverage the temporal dynamics to fill in the missing information due to fewer available observations. This pattern also holds for the bouncing balls task, cf, Figures 2Te and 20e in Appendix F.
CONCLUSIONS, LIMITATIONS AND FUTURE WORK
We proposed Spatially Structured Recurrent Modules, leverage both spatial and modular structure in data, and explored its potential in the challenging problem setting of predicting the forward dynamics from partial observations known spatial locations. In the tasks of video prediction from crops and multi-agent world modeling in the Starcraft2 domain, found that it compares favorably against strong baselines in terms of out-of- distribution generalization and robustness to the number of available observations. Future work may attempt to extend the idea to parallel-in inniell--i-timee methods like universal transformers (Dehghani et al.) 2018) and thereby address the computational bottleneck of recurrent processing, which is a current limitation. Another interesting avenue of research could be to explore how latent random variables can be used in tandem with the spatial structure to obtain a variational version of S2RMs. Finally, efficient implementations using block-sparse methods (Gray et al., (2017) might hold the key to unlock applications to significantly larger scale spatiotemporal forecasting problems encountered in domains like climate change research (Rolnick et al.) 2019).
class of models constructed to jointly
Here, the code cz refers to scenario where each team comprises colossus (1 c), 3 stalkers (3s),
and 5 zealots (5z).

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
The authors would like to thank Georgios Arvanitidis, Luigi Gresele, Michael Cobos for their feedback the paper, and Murray Shanahan for the discussions. The authors also acknowledge the important role played by their colleagues at the Empirical Inference Department of MPI-IS Tübingen and Mila throughout the duration of this work.
REFERENCES
Ferran Alet, Tomás Pérez, and Leslie Kaelbling. Modular meta-learning, arXiv preprint arXiv:1806. 100.:6660000166. 2018.
Peter Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanzaez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 1806.01261. 2018.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. arXiv preprint arXiv:1901.1091: 2019.
Alberto Cenzato, Alberto Testolin, and Marco Zorzi On the difficulty of learning and predicting the long-term dynamics of bouncing objects. arXiv preprint arXiv:1907 2019.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint urXiv: 1078, 2014.
Mostafa Dehghani. Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprini arXiv:1807.03817 2018.
Christian Eichenberger, Adrian Egli, Mattias Ljungström, Sharada Mohanty, Guillaume Mollard, Erik Nygren, Giacomo Spigler, and Jeremy Watson. Flatland, 2019. URL
http:://iatlland-11- readme htmll
s3- wtt-website.uu-cceee eu- aww.zaoaa.aaaaaaaa....aaaa amazonaws com/
SM Ali Eslami, Danilo Jimenez Rezende, Frederic Besse, Fabio Viola, Ari S Morcos, Marta Gar- nelo, Avraham Ruderman, Andrei A Rusu, Ivo Danihelka, Karol Gregor, et al. Neural scene representation and rendering. Science, 104--1110, 2018.
Gregory E Fasshauer. Positive definite kernels: past, present and future. 2011.
Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Advances Neural Information Processing Systems, pp. 3601-3610, 2017.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint 187.00180. 2018.
Zoubin Ghahramani and Michael Jordan. Factorial hidden markov models. In Advances in Neural Information Processing Systems, pp. 472-478, 1996.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909, 1909.1999,, 2019. 10893,
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401, 2014.
arXiv preprint
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiñska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626): 471-476, 2016.
10

---

Published as a conference paper at ICLR 2021
Scott Gray, Alec Radford, and Diederik P Kingma. Gpu kernels for block-sparse weights. arXiv preprint arXiv:1711.09221 3, 2017.
Mikael Henaff, Jason Weston, Arthur Szlam, Antoine Bordes, and Yann LeCun. Tracking the world state with recurrent entity networks, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1785 1780, 1997.
Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Max Jaderberg, Karen Simonyan, Andrew Zisserman, and Koray Kavukcuoglu. Spatial transformer networks, 2015.
Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. In Advances in neural information processing systems, pp- 7640-7651, 2018.
Nan Rosemary Ke, Olexa Bilaniuk, Anirudh Goyal, Stefan Bauer, Hugo Larochelle, Chris Pal, and Yoshua Bengio. Learning neural causal models from unknown interventions. arXiv preprint arXiv:1910.01075. 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint vrXiv:2001.0445/1 2020.
Jannik Kossen, Karl Stelzner, Marcel Hussing, Claas Voelcker, and Kristian Kersting, Structured object-aware physies prediction for video modeling and planning. arXiv preprint 150022911 2019.
Ananya Kumar, SM Eslami, Danilo J Rezende, Marta Garnelo, Fabio Viola, Edward Lockhart, and Murray Shanahan. Consistent generative query networks. arXiv preprint 10X7vv1807.00000 2018.
Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently recurrent neural network (indrnn): Building longer and deeper In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5457-5466, 2018.
E Majani, Ruth Erlanson, and Yaser Abu-Mostafa. On the k-winners-take-all network. Advances in neural information processing systems, 6634-642, 1988.
Hongyuan Mei, Guanghui Qin, Minjie Xu, and Jason Eisner. Informed temporal modeling via
logical specification of factorial {Istm}s, 2020. URL
forum??
id=S1 ghz
Djordje Miladinovié, Muhammad Waleed Gondal, Bernhard Schölkopf, Joachim M Buhmann, and Stefan Bauer. Disentangled state space representations. arXiv preprint 1rXiv:1906.03255, 2019.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes neural radiance fields for view synthesis. arXiv preprint arXiv:2003.08934, 2020.
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning combine top-down and bottom-up signals in recurrent neural networks with attention over modules. arXiv preprint 20X0::2006.1681.. 2020.
Giambatrista Parascandolo, Niki Kilbertus, Mateo Rojas-Carulla, and Bernhard Schölkopf. Learning independent causal mechanisms. arXiv preprint
2017.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer, 2018
11

---

Published as a conference paper at ICLR 2021
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, hirg-peefformacc deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Bie Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20:1177-1184, 2007.
David Rolnick, Priya Donti, Lynn H Kaack, Kelly Kochanski, Alexandre Lacoste, Kris Sankaran, Andrew Slavin Ross, Nikola Hilojevic-Dupont Natasha Jaques, Anna Waldman Waldman-Brown, Brown, et al. Tackling climate change with machine learning. arXiv preprint arXiv: /906/ 050000 2019.
Mikayel Samvelyan, Tabish Rashid, Christian Schroeder de Witt, Gregory Farquhar, Nantas Nardelli, Tim G. J. Rudner, Chia-Man Hung. Philip H. S. Torr, Jakob Foerster, and Shimon Whiteson. The starcraft multi-agent challenge, 2019.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. In Advances in neural information processing systems, pp. 4967-4976, 2017.
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap. Relational recurrent neural networks. In Advances in Neural Information Processing Systems, pp. 7299-7310, 2018.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- gated mixture-of -x-eere layer. arXiv preprint arXiv:1701.00 06538, 2017.
Gautam Singh, Jaesik Yoon, Youngsung Son, and Sungjin Ahn. Sequential neural processes, 2019.
Chen Sun, Per Karlsson, Jiajun Wu, Joshua B Tenenbaum, and Kevin Murphy. Predicting the present and future states of multi-agent systems from partially-observed visual data. In International Conference on Learning Representations, 2019. bttpstt:: //tp://es///peeeeee.... www.:eee/foooee ld Fidrrxxe 1xcK3CcKX
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. arXiv preprint arXiv:2002.1 /296, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10903, 2017.
Oriol Vinyals, Timo Ewalds, Sergey Bartunov, Petko Georgiev, Alexander Sasha Vezhnevets, Michelle Yeo, Alireza Makhzani, Heinrich Küttler, John Agapiou, Julian Schrittwieser, al. Starcraft ii: A new challenge for reinforcement learning. arXiv preprint arXiv: 2017.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks, 2017.
Nicholas Watters, Daniel Zoran, Theophane Weher, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. In Advances in neural information processing systems, pp. 4539-4547, 2017.
Han Zhang, Ian Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. arXiv preprint arXiy: 1805.08315 2018.
Ellen D. Zhong, Tristan Bepler, Joseph H. Davis, and Bonnie Berger. Reconstructing continuous distributions of 3d protein structure from cryo-em images, 2020.
12

---

Published as a conference paper at ICLR 2021
A QUALITATIVE VISUALIZATIONS A GRID- WORLD NAVIGATION TASK
In this section, we show qualitative results on grid-world task defined in Eichenberger et al. 2019 which formulates the problem of navigation on a railway network in multi-agent reinforcement learning framework. The environment comprises network of railroads, on which agents (trains) may move in order to reach their destination. our experiments, the entire railway network is defined on a 60 x 60 grid-world world and we let each agent only observe a partial and local view of the environment, which is 5 x crop centered around itself.
We gather 10000 multi-agent trajectories with 10 agents and maximum length 128, from which we use 8000 for training and reserve 2000 for validation. We train S2GRU with 10 modules for 100 epochs and early stop when the validation loss is at its minimum. With the trained model, we visualize the following two things.
First: for each module F.. we visualize the spatial locations it may attend to. To this end, we consider all 60 x 60 = 3600 pixel locations in the grid-world, say Xij where j E {1,..., 60). For each such Xij, we evaluate the quantity:
X = (p"",..,)
(11)
where Stj = P(xi)) (see Eqn 3). pm is the embedding of module Fm. and Amm is 60 x 60 image indexed by i and j. which we call the enclave of module Fm. Note that this is identical to what we visualize in Figure
we identify each module with its enclave, and visualize the graph of interactions between them. In Figure 9 we plot as nodes the enclaves X²m Further, we draw an edge between enclaves x and X" iff Z (p"'.'"' >> O.
Next:
We make the following two observations. First, the images in Figure 9 show that each module learns to account for spatial region in the environment, as we imagined in Figure [b] Section || Second, we find that the modules interact sparsely with each other while some modules learn interact with up-to five other modules, other modules learn to operate independently.
B DETAILED TASK DESCRIPTIONS
B. STARCRAFT2
The Starcraft2 Environment we use is a modified version of the SMAC-Eny proposed in Samvelyan eet al. (2019) and built PySC2 wrapper around Blizzard SC2 API (Vinyals et al. 2017) Starcraft2 is a real-time-strateg armies of units (airborne or land-based) defeat the opponent's army (where the opponent can be an AI or another human). The players must choose their alien race before starting the game; available options are Protoss, Terran and Zerg. All unit types (of all races) have their strengths and weaknesses against other unit types, be it in terms of maximum health, shields (Protoss), energy (Terran), DPS (damage per second, related to weapon cooldown), splash damage, or manufacturing costs (measured in minerals and vexpene gas, which must be mined).
(RTS) game where players are tasked with manufacturing and controlling
The key engineering contribution of Samvelyan et al, 2019) to repurpose the RTS game as a multi-agent environment, where the individual units in the army become individual agent:]] F1 The result a rich and challenging environment where heterogeneous teams of agents must defeat each other in melee and ranged combat. The composition of teams vary between scenarios, of which Samvelyan et al. (2019) provide a selection. Further, new scenarios can be easily created with the SC2MapEditor, which allows for practically endlessly many possibilities.
We build on Samvelyan et al. (2019) by modifying their environment to better expose the transfer and dutrof-distribution aspects of the domain by (a) standardizing the state and action space across large class of scenarios and (b) standardizing the unit stats to better reflect the game-defined notion of hit-points.
Pleasee note that this a game-specific notion.
"Note that this is rather unconventional, since each player usually controls entire armies and must switch between macro- and micro-management of units unit-groups.
13

---

Published as a conference paper at ICLR 2021
Figure 9: Joint visualization of spatial enclaves and the interaction graph between modules in the grid-world environment of Eichenberger et al. (2019), as detailed in Appendix A The images show which spatial locations module attends to via the local attention (spatial enclaves), whereas the presence of an edge indicates that the corresponding modules may interact via inter-cell attention. Gist: The modules indeed learn notion of spatial locality, while interacting sparsely with each other.
14

---

Published as a conference paper at ICLR 2021
(a) 1s2z (1 Stalker and Zealots per team).
(b) 5s3z (5 Stalkers and Zealots per team).
(c) 2s3z Stalkers and 3 Zealots per team).
(d) 3s5z (3 Stalkers and Zealots per team).
(c) 1c3852 (1 Colossus, 3 Stalkers 5 Zealots per team).
Figure 10: Human readable illustrations of the Stareraft2 (SMAC) scenarios we consider in this work. Figures 10a and 10b show the OOD scenarios, whereas Figures 10c 10d and [0e show the training scenarios (provided by Samvelyan et al. (2019).
B.1.1 STANDARDIZED STATE SPACE FOR ALL SCENARIOS
In the environment provided by Samvelyan al. 2019). the dimensionality of the vector state space varies with the number of friendly and enemy agents, which in turn varies with the scenario. While this is an issue the typical use case of training MARL agents in a fixed scenario, it is convenient for designing models that seamlessly handle multiple scenarios. In the following, we propose an alternate state representation that preserves the spatial structure and consistent across multiple scenarios.
Instead of representing the state an agent a with vector of variable dimension, we represent with multi-channed polar image Jo of shape C' x J, where C' is the number of channels and J) is the image size. Given the radial and angular resolutions and y (respectively), the pixel coordinate i = 0, O..,, 1,j O,... J corresponds to coordinates (i p.j (p) with respect to polar coordinate system centered on the agent a, where the positive r-axis (j = 0) points towards the
east. Further, the field of view (FOV) of an agent is characterized by circle of radius I p centered on the agent at 2D game-coordinates x (af. xg) to which the Starcraft2 API (Vinyals et al. 2017) provides raw access.
The polar image Ja therefore provides an agent-centric view of the environment, where pixel coordinates i.j in J be mapped to global game coordinates x (x1, in FOV via:
i pcos I vp + x
(12)
22 sin j ip] + x.
(13)
In what follows, we denote this transformation with T,, in (i,))) x2).
Now, the channels in the polar image can encode various aspects of the observation; in our case: friendly markers (one channel), unit-type markers (nine channels, one-hot), health-energy-colldovr.
15

---

Published as a conference paper at ICLR 2021
shields (HECS, four channels) and terrain height (one channel). As an example, let consider the friendly markers, which is a binary indicator marking units that are friendly. If we have an agent at game position (21, 222 that is friendly to agent a, then we would expect the pixel coordinate (i,j) = T (x1, =2) of the corresponding channel in the polar image 3 to be 1, but 0 otherwise. Likewise, the value of 3 at the channels corresponding HECS pixel position i, j gives the HECS of the corresponding uni at Ta (t.j). This representation has the following advantages: (a) it does not depend on the number of units in the field of view, (b) exposes the spatial structure in the arrangement of units which can naturally processed by convolutional neural networks (e.g. with circular convolutions).
Nevertheless, it has the disadvantage that the positions are quantized to pixels. but the euclidean distance between the locations represented by pixels (i.j) and (i,j + increases with increasing i. Consequently, this representation may not remain suitable for larger FOVs.
Further, this representation is also appropriate for the action space. Given an agent, we represent the one-hot categorical actions of all friendly agents in FOV as multi-channel polar image. In this representation, the pixel position i gives the action taken by an agent at at position Ta (i.j). Unfriendly agents get assigned an "unknown action", whereas positions not occupied by a living agent are assigned a "no-op" action.
B.1.2 STANDARDIZED UNIT STATS
At any given point in time, an active unit in Starcraft2 has certain stats, c.g. its health, energy (Terran), shields (Protoss) and wcapon-cooldown (for armed units). A large and expensive unit-type like the Colossus has more max-health (hit-points) than smaller units like Stalkers and Marine Likewise, unit-types differ in the rate which they deal damage (measured in caanage-per-seconee excluding splash damage), which in turn depends on the cooldown duration of the active weapon.
DPS,
Now, the environment provided by (Samvelyan et al. (0119) normalizes the stats by their respective maximum value, resulting in values between O and 1. However, given that different units may have different normalization, the stats are rendered incomparable between unit types (without additionally accounting the unit-type). We address this by standardizing stats (instead of normalizing) by dividing them by fixed value. In this scheme, the stats are scaled uniformly across all unit-types, enabling models to directly rely on them instead of having to account for the respective unit-types.
B.2 VIDEO PREDICTION FROM CROPS ON THE BOUNCING BALLS TASK
The bouncing balls task is a well-known test-bed for evaluating the performance of video prediction models (Fraccaro al. 2017 Watters et al. 2017: Miladinović et al. 2019 Kossen et al. 2019 Cenzato et al. 2019). We modify the problem by introducing partial observability concretely, instead of providing the model with the full image frames, we only provide with crops randomly sampled locations.
As mentioned in Section 6 at every time step we sample 10 pixel positions (xf 1.00 from the 1-th full video frame o, of size 48 x 48. Around the sampled central pixel positions x² we extract 11 11 crops, which we use as the local observations Og. The task now to predict 11 11 crops 0% corresponding to query central pixel positions x", at future time-step t' > t. Observe that at any given time-step t, the model has access to at most 52% of the global video frame assuming that the crops never overlap (which is rather unlikely).
We train all models on a training dataset of 20K vidco sequences with 100 frames of balls bouncing in arena of size 48 48. We also include additional fixed ball in the center make the task more challenging. We use another 1K video sequences of the same length and the same number of balls held-out validation set. In addition, we also have out-of-distributio (OOD) test sets with various number of bouncing balls (ranging from to 6) and each containing 1K sequences of length 100.
'If health drops to zero, the unit is considered dead and the representation does not differentiate between dead and absent units.
These stats may change with game-versions, and are catalogued here: http:: //1i:: pediipeiia net./ (tareraft2/um tss (Sta (CCaff))
16

---

Published as a conference paper at ICLR 2021
In Figure 4 for each number of balls (i.e. point on the x-axis), we plot the respective metries which are aggregated over 10 randomly selected 11 x 11 crops of a total of 100000 frames spread over 1000 trajectories with 100 frames each.
C PRECISE DESCRIPTION OF ATTENTION MECHANISMS
C.I INPUT ATTENTION
Recall from Section that the input attention mechanism a mapping between sets: namely, from that of observation encodings { to that of RNN inputs {u," |m In what follows, we use the einsum notation to succintly describe the exact mechanism. But before that, we repeat the definition of the truncated spherical Gaussian kernel (Fasshauer) 2011) quantify the similarity between two points p.s E S:
exp [-2c - p.p.... s)], if otherwise
((p,s)
(14)
0,
where E RT and 7 E [-1, 1) are hyper-parameters (kernel bandwidth and truncation parameter, respectively). and 0 < < since p and s are unit vectors. We observe that both and controls the sparsity of the kernel: determines the size of the neighborhood of p. i.e. the size of the set B(p) S of all E (p)) such that p.s 2 T and accordingly ((p.s) > O, whereas the bandwidth e controls how the attention decays inside (p)) Intuitively, r determines lower bound to the amount of sparsity that the kernel induces (irrespective of the bandwidth e), whereas for fixed т, sparsity can be increased by increasing e. We find 7 E [-1, 1, 0.6] and e E [0.9, 2] to work well; setting and e to much larger values destabilizes the training due to excessive sparsity, whereas setting to much smaller values results in being flat inside (p)) and therefore poor propagation of gradients.
Now, we use k to index the attention heads, d to index the dimension of the key and query vectors, and denote with cor the i-th component of ef and with hmj the j-th component of h," Given learnable parameters (K), ()Q), O we obtain:
Qakd = Cas θ
(Q) ikd
Kinkd = hym 1
(15)
Vaku = Cas (V)
Wman Qakd Kmid
(16)
ske
W Wmakk = sm, (maa)
ma = Z(p", s²)
(17)
W mak W W mak
(mn(ke) W Wmak Vako
(18)
where: sm, denotes softmax along the a-dimension, (L4) is what we will call the local weights, we omit the time subseript in s" for notational clarity, and (m((kw) is the (kv))tt component of a vector u Finally, we obtain the components Umi of RNN inputs u," via gating operation:
Uyni
Glimp) .mi + (1 G(HP)) u.xi
(19)
where the gating weight Ghiup) E 10 is obtained by passing ümi and bmi Whild cas through two-layer MLP with sigmoidal output (in parallel across m)). Now, observe that by weighting the MHDPA attention outputs ( in Equation 18 by the kernel Z (via WF(L),, we construct a scheme where the interaction between input Où and RNN Fm is allowed only if the embedding st of the corresponding position x² has large enough cosine similarity (2 r) to the embedding Pm F...
This partially implements the assumption of Locality of Observation detailed in Section 3
C.2 INTER-CELL ATTENTION
Recall from Section 4 that the inter-cell attention maps the hidden states of each RNN (h.") M to the set of aggregated hidden states (hr" M thereby enabling interaction between the RNNs F... While its mechanism is identical to that of the input attention, formulate it below for completeness.
Indices not appearing both sides of an equation summed over; this is implemented as e ninummi in most DI. frameworks.
17

---

Published as a conference paper at ICLR 2021
To proceed, we denote with his the i-th component of hj (in addition to the notation introduced before Equation 15), and take 41Q), 4(K) and a(V) to be learnable parameters. We have:
amke = hmg 1
jkd
(Q)
Kidd = hys ( (K)
ikd
(20)
Vike = h4g 1
(V)
W milk = Kidd
(21)
sku
w Z(p",, .p')
Wmit mllk sm/ (W (Wiik) mllk)
(22)
W Wmik W ml
Amn(ke) = Wmik View
(23)
where hmc(kw the (kw))-h component of vector hT Finally, the j-th component kay of the aggregated hidden state h in Equation is given by gating operation:
haj - Cmf + (1 - hmj
(24)
where the gating weight G E (0, is obtained by passing Riy and Cmjj = ml huz through
two-layer MLP with sigmoid output (in parallel across m.) The weighting by Z (in Equation 23
left) ensures that the interaction is constrained to be only between RNNs whose embeddings in S similar enough, thereby implementing the assumption of Local Interactions between Sub-systems in Section
C.3 POSITIONAL. ENCODING
In Section recall that we used the following positional embedding P:
P(x) = */||=]| E S where (si+m Si+14m) (sin (#m/10000'), cos (sm/10000-))
(25)
In this section, explore how the choice of positional embedding function P determines function space of spatial functions (defined on X') that the local-attention can represent. To this end, consider the distance in of module with embedding p an observation made location x as a function of
x. given by
w (0) p- P(x)
(26)
Here, the local weight of interaction between the module at p and an observation made at x is given by:
exp [-2c w/(1) (x))), if (x( 2 otherwise
((p, P(x))
(27)
0,
In particular, observe that in order for two locations x and y to be connected by the module, we require from w(() that be flexible enough such that wx()(x) 2 t and wx()(y) Z for a chosen т. This flexibility stems from the fact that we implicitly express wL) as a linear combination of sinusoidal basis functions with learned weights:
w (L) (x)
Σ [22) cos (wy x) + P2j+1 sin (wy x)]
(28)
Here, P2j and P2 1 are learnable parameters (as components of learnable vector p of dimension 2.)), and wy are frequency vectors. Now, if the dimension of the embedding vector were to tend to infinity, we may have a growing number of frequencies Wj to gradually recover the full Fourier basis of (2') (assuming is Euclidean for simplicity). In the limit, (x)) can be an arbitrary function lying on unit sphere in (22) (i.e. f' [2 = 1; recall that p is normed to unity). In other words, in large dimensional embedding space, the system is afforded large amount flexibility to learn any spatial structure or topology on X by connecting pairs of points x and y in X via a module.
For computational tractability, however, require p to be finite-dimensional, implying that w, must be sampled. By using P as defined in Equation 25 we essentially sample wy as coordinate (one-hot) vectors with log ||wyy] |I sampled on uniform grid. This sampling step is not only computationally favorable, but also justified in the theory of RKHS Raahim Recht (2007) use Bochner's 's theorem to show that any proper distribution p(w) (from which
be sampled) leads to feature map,
18

---

Published as a conference paper at ICLR 2021
the inner product of which in expectation over p corresponds to positive-defintte kernel. The convergence to such a kernel is exponential in the number of samples (equivalent to the dimension of the embedding).
Further, we note that while the sampling constrains the function space in which w( can lie in, we find (empirically) that this can fact have a regularizing effect. Nevertheless, this raises the question whether other choices of basis function are viable. We speculate that polynomial basis (e.g. feature maps of a degree d polynomial kernel) might also be viable, but leave extensive exploration to future work.
D BASELINE ARCHITECTURE
As mentioned in 6 in order to ensure fair comparison between the baselines and method, we describe a baseline architecture constructed to satisfy a few critically important desiderata that naturally satisfied by S2RMs. Namely, (a) it must be parameterically agnostic to the number of available observations and (b) it must be invariant to the permutation of the observations. For this, we extend the framework Generative Query Networks (Eslami et al.. (018) by predicting the forward dynamics of an aggregated representation. While we invest effort in ensuring that the resulting class of models can perform at least well as S2RMs on in-distribution (validation) data, we do consider it novel contribution of this work.
Encoder. At given timestep t. the encoder E jointly maps the embedding s S of the position x E X' and the corresponding observations Of to encodings ef, which are then summed over a to obtain an aggregated representation:
F Σ E(O st)
(29)
The additive aggregation scheme use is well known from prior work (Santoro et al. 2017 Eslami ct al. 2018 Garnelo et al. 2018) and makes the model agnostic to A and to permutations of (x; O a. The encoder E is seven-layer CNN with residual layers, and the positional embedding st injected after the second convolutional layer via concatenation with the feature tensor. The exact architectures can be found in Appendices E.] and E.2
RNN. The aggregated representation r, is used as an input to a RNN model F as following:
h C++1 (1,,,h,,
(30)
where h, and are hidden and memory states of the RNN F' respectively. We experiment with various RNN models, including LSTMs (Hochreiter & Schmidhuber, 1997). RMCs (Snntoro et al. 2018) and Recurrent Independent Mechanisms (RIMs) (Goyal et al. (011)).
As a samity check, we also show results with a Time Travelling Oracle (TTO), which has access to T1+1 (but at time step t), and produces h Frro with two layer MLP Frro- TTO therefore does not model the dynamics, but merely verifies that the additive aggregation scheme (Equation 29) and the querying mechanism (Equation 31) are sufficient for the task at hand.
Decoder. Given the embedding s" of the query position x", the decoder D predicts the corresponding observation Oz
Ö D(h++1,s))
(31)
We parameterize D with deconvolutional network with residual layers, and inject the positional embedding of the query 8 after single convolutional layer by concatenating with the layer features (see Appendices E.I and E.2
E
HYPERPARAMETERS AND ARCHITECTURES
E. ENCODER AND DECODER FOR BOUNCING BALLS
The architectures of image encoder and decoder was fixed for all models after initial experimentation. We converged the following architectures.
19

---

Published as a conference paper at ICLR 2021
Representation
Observation
(IC)
Convolution
(128 128)
(omml padding stride
Convolution (kernel size padding 1, strid
IC)
Skip-Add
Deconvolution
Convolution
(128 128)
(128 128)
(kernel size packing 0 sbride
(xernll size padding stride
Skip-Add
Convolution
(128 128)
(xernel size padding 0, stride 1)
Convolution
(128 128)
(kemel padding stride
Skip-Add
Deconvolution
Convolution
(128 128)
(128 - 128) padding stride
(kemel size padding stride 1)
Skip-Add
Convolution
(128 CC 128) padding stride
Convolution
(128 128)
(kemal padding stride
Concatenate
Deconvolution
Skip-Add
(128 CC 128) padding stride
Convolution
Concatenate
(128 128) padding
Deconvolution
Convolution
(RC 128)
(IC 128) (kernel size padding stride 1)
(kemel size 5, padding stride 1)
11x11
Representation (RC)
Observation (IC)
Positional Embedding
Positional Embedding
(CC)
(CC)
(a) Encoder.
(b) Decoder.
Figure 11: Baseline encoder and decoder architectures for the Bouncing Ball task.
20

---

Published as a conference paper at ICLR 2021
Representation
Observation
(IC)
Convolution
(128 128)
(eemel size padding stride
Convolution
IC) (kernel size 3, padding 1, sbride
Skip-Add
Deconvolution (128 128)
Convolution
(128 128)
(kernel size padding 0 sbride 1)
(xernll size padding stride
Skip-Add
Convolution
(128 128)
(eernel size padding 0, stride
Convolution
(128 128)
padding
Skip-Add
Deconvolution
Convolution
(128 128)
(128 128) (kamnl padding stride
(kemel size padding stride 1)
Skip-Add
Convolution
(128 - 128)
Convolution
(128 128)
(kemall size 3, padding stride 1)
Skip-Add
Deconvolution
Convolution
(128 128)
(kemal stride
(128 128)
(kermal size 3, padding stride 1)
Deconvolution
Convolution
(RC 128)
(IC 128) (kemal padding stride
padding 1)
11x11
Representation (RC)
Observation (IC)
(a) Encoder.
(b) Decoder.
Figure 12: S2RM encoder and decoder architectures for the Bouncing Ball task.
21

---

Published as a conference paper at ICLR 2021
E.1.1 S2RMs
The encoder (decoder) a (de)convolutional network with residual connections (Figure 12
E.1.2 BASELINES
Like in the case of S2RMs, the encoder (decoder) is a (de)convolutional network with residual connections (Figure [T]. but with the positional embeddings injected after the second convolutional layer. This is loosely inspired by the encoders used in Eslami et al. (2018).
E.2 ENCODER AND DECODER FOR STARCRAFT2
E.2.1
S2RMs
Recall from Appendix [B. that the states polar images. We therefore polar convolutions, which entails zero-padding the input image along the first (radial) dimension but circular padding along the second (angular) dimension. The encoder and decoder architectures be found in
Figure 14
E.2.2 BASELINES
Like for S2RMs, we use polar convolutions while injecting the positional embeddings further downstream in the network. The corresponding encoder and decoder architectures are illustrated in Figure 13
E.3 SPATIALLY STRUCTURED RELATIONAL MEMORY CORES (S2RMCs)
Embedding Relational Memory Cores Santoro et al. 2018) naïvely in the S2RM architecture did not result in working model. We therefore had to adapt it by first projecting the memory matrix (M in Santoro et al. (2018) of the mn-th RMC to message h," This message is then processed by the intercell attention to obtain hn"", which finally concatenated with the memory matrix and current input before applying the attention mechanism (i.e. in Equation 2 of Santoro et al. (2018), replace [M with [M; 2, h
E.4 HYPERPARAMETERS
E.4.1 BOUNCING BALL MODELS
The hyperparameters used can be found in Table Further, note that in Equation 5 pass the gradients through the constant region the kernel as if the kernel had not been truncated.
E.4.2
STARCRAFT2 MODELS
The hyperparameters we used can be found Table Note that we only report models that attained validation loss similar to better than S2RMs.
E.4.3 TRAINING
All models trained using Adam Kingma & Ba 2014 with initial learning rate 0.000: We use Pytorch's 's (Paskke et al. 2019) ReduceLROnP Paateu learning rate scheduler to decay the learning rate by factor if the validation loss does not improve by least 0.01% over the span of 5 epochs. We initially train all models for 100 epochs. select the best of three successful runs, fine-tune for another 100 epochs, and finally select the checkpoint with the lowest validation loss (i.e. we early stop). We train all models with batch-size (Starcraft2) or 32 (Bouncing Balls) on single V100-32GB GPU (each).
22

---

Published as a conference paper at ICLR 2021
Representation
Observation
(IC)
Convolution
(128 128)
Polar Convolution
(128 IC) Oamel sice Se5, padking 1x2, stnid Txl1)
Skip-Add
Deconvolution
Polar Convolution
(128 + 128) (emee paddligg 0x0, stride
(128 128)
Demel sico 1x3, padding 0x1, strdd 1x11
Skip-Add
Polar Convolution (128 128)
Demee size 363, padding Ox1, soide 1x21
Polar Convolution
(128 128)
Jonmel padding etridn 1x1)
Skip-Add
Deconvolution (128 128)
Polar Convolution
Homel size 9a5, pudding Beb, strido 2x00
(128 128)
Oermel iiin pedding 1x2, stride 1x]]
Skip-Add
Polar Convolution
(128 128)
0omel size 3x5, paddlng 0x2, sondo
Polar Convolution
(128 128)
Concatenate
Deconvolution
Skip-Add
(128 CC 128) padsling
Polar Convolution (128
Concatenate
Jarmal padding
Fully Connected (RC 128)
Polar Convolution [IC 128)
Bammd size 2x5. pudding 1x2, atrico 2x:0
x4
Representation (RC)
Observation (IC)
Positional Embedding
Positional Embedding
(CC)
(CC)
(a) Encoder.
(b) Decoder.
Figure 13: Baseline encoder and decoder architectures for the Starcraft2 task.
23

---

Published as a conference paper at ICLR 2021
Representation
Observation
(IC)
Convolution
(128 - 128)
Demet paddling 0x0, strie
Polar Convolution
(128 IC) (emeel dite 3x5, padding 1x2, condo 1x11
Skip-Add
Deconvolution (128 - 128)
Polar Convolution
(temel dite 345, padsing Co0, condo 2x31
(128 128)
Bemel size 1x3, padding 0x1, soido 1651
Skip-Add
Polar Convolution (128 128)
Oemel size 3x3, padding Oxt, sorde 1x21
Polar Convolution
(128 128)
Jomel padding
Skip-Add
Deconvolution
Polar Convolution
(128 128)
(128 - 128)
Home size 3x5, padding Dmo, strido 2x50
Skip-Add
Polar Convolution
(128 128) size padding
Polar Convolution
(128 + 128)
Skip-Add
Deconvolution (128 128)
Polar Convolution
(128 + 128)
Jamel padding atrid 1x1)
Fully Connected (RC + 128)
Polar Convolution |IC 128)
x4 36
Representation (RC)
Observation (IC)
(a) Encoder.
(b) Decoder.
Figure 14: S2RM encoder and decoder architectures for the Starcraft2 task.
24

---

Published as a conference paper at ICLR 2021
Model
Hyperparameter
Value
S2GRU
Number of modules (M)
10
GRU: hidden size per module
128
Module embedding size (d)
16
Kernel bandwidth (e)
Kernel truncation (7)
0.6
shape ((4/K)
(128, 2, 016)
shape +
(128, 2, 128)
shape k(Q/K)
(128, 4, 016)
shape Φ(V)
(128, 4, 128)
RMC (Santoro et al.
2018
Number of attentron heads
4
Size of attention head
128
Number of memory slots
Key size
128
LSTM (Hochreiter & Schmidhuber 1997
Hidden size
512
RIMs Goyal et al.
2019
Number of RIMs (NT)
Update Top-k (kA)
Hidden size (heise)
510
Input key size Input value size
32
400
TTO
MLP hidden size
512
Table Hyperparameters used for various models on the Bouncing Ball task. Hyperparameters not listed here were left at their respective default values.
25

---

Published as a conference paper at ICLR 2021
Model
Hyperparameter
Value
S2GRUs
Number of modules (M)
10
GRU: hidden size per module
128
Module embedding size (d)
Kernel bandwidth (c)
Kernel truncation (7)
0.5
shape ((Q/K)
(128, 2, 016)
shape (
(128, 2, 128)
shape ((Q/K)
(128, 4, 016)
shape ((V)
(128, 4, 128)
S2RMC
Number of modules (M)
10
RMC: number attention heads
RMC: size of attention head
64
RMC: number memory slots
RMC: key size
64
Module embedding size (d)
8
Kernel bandwidth (e)
Kernel truncation (T)
0.5
shape (Q//K)
(128, 2, 016)
shape ((V)
(128, 2, 128)
shape Φ(Q/K)
(128, 4, 016)
shape Φ(V)
(128, 4, 128)
RMC (Santoro et al. 2018
Number of attention heads
4
Size of attention head Number of memory slots
128
Key size
16
LSTM (Hochreiter & Schmidhuber 1997
Hidden size
2048
TTO
MLPP hidden size
512
Table 3: Hyperparameters used for various models on the Starcraft2 task. Hyperparameters not listed here were left at their respective default values.
26

---

Published as a conference paper at ICLR 2021
O
0
8
6
O
O
O
0
Figure 15: Rollouts (OOD) with bouncing ball, from top to bottom: ground-truth S2GRU, RIMs, RMC, LSTM. Note that all models were trained on sequences with 3 bouncing balls, and the global state was reconstructed by stitching together 11 11 patches from the models (queried on a 4 x 4 grid).
O
O
O
9
-
5
6
Φ
+
S
Figure 16: Rollouts (OOD) with bouncing balls, from top to bottom: ground truth, S2GRU, RIMs. RMC, LSTM. Note that all models were trained on sequences with 3 bouncing balls, and the global state was reconstructed by stitching together 11 11 patches from the models (queried on a 4 x grid).
E.4.4 OBJECTIVE FUNCTIONS
In the Starcraft2 task, predicting the next state entails predicting images of binary friendly markers, categorical unit type markers and real valued HECS markers. Accordingly, the loss function is sum of binary cross-entropy term (on friendly markers), a categorical cross-entropy term (on unit-type markers) and a mean squared error term (on HECS markers).
In the Bouncing Balls task, the model output is a binary image. Accordingly, pixel-wise binary cross-entropy loss.
F ADDITIONAL RESULTS
F. BOUNCING BALLS
F.1.1 ROLLOUTS
To obtain the rollouts in Figure 3. we adopt the following strategy. For the first 20 prompt-steps, we present all models with exactly the same 11 x 11 crops around randomly sampled pixel positions for 20 time-steps. For the next 25 steps, all models are queried at random pixel positions [T]] and the resulting predictions (on crops) thresholded at 0.5 and fed back in to the model for the next step (at known pixel positions from the previous step).
Also at every time-step, the models are queried for their predictions 16 pixel locations placed 4 x 4 grid. The resulting predictions are stitched together and shown in Figures [5]
[6
[8]
and
19
27

---

Published as a conference paper at ICLR 2021
26
+
4P
P
O
O
+
2
9
e
1
Figure 17: Rollouts (ID) with 3 bouncing balls, from top to bottom: ground-truth, S2GRU, RIMs, RMC, LSTM. Note that all models were trained on sequences with 3 bouncing balls, and the global state was reconstructed by stitching together 11 11 patches from the models (queried on a 4 x grid).
"0"
G
Figure 18: Rollouts (OOD) with 4 bouncing balls, from top to bottom: ground-truth, S2GRU, RIMs, RMC, LSTM. Note that all models were trained on sequences with 3 bouncing balls, and the global state was reconstructed by stitching together 11 11 patches from the models (queried 4 x 4
grid).
-
.
Figure 19: Rollouts (OOD) with 6 bouncing balls, from top to bottom: ground ground-truth, S2GRU, RIMs, RMC, LSTM. Note that all models were trained sequences with 3 bouncing balls, and the global state was reconstructed by stitching together 11 x 11 patches from the models (queried on a 4 x grid).
28

---

Published as a conference paper at ICLR 2021
0.5
0.0
6.049
500
0.087
0.71
8.314
0.015
5977
0.11
8:000
0.71
0.75
0.59
0.90
3.9
2039
6.40
8.66
5.54
8856
0.66
0.86
5.00
03:33
5.27
8:42
0:45
3017
0:45
0.82
5,08
2684
8.92
6.00
5.23
0.K43
0.41
2:8
0867
E.27
0.80
2388
2.07
0.01
0.32
C357
K623
Fracdion
Fraction baliool
S2GRU
(b) RIMs
5.004
0144
0993
8.90
5914
8.000
914
0844
0.5%
5914
0.
0335
6,05
8:57
6371
6.74
6576
577
0,90
030
6:211
14
0557
6952
6567
3.99
0.91
6025
6855
C475
0.30
DEEL
0.57
6674
0.71
-02
602
(c) RMC
(d) LSTM
0.99
6.227
145
0.145
0.77
0.90
5.503
3397
1.00
0.99
C.00
511
0.39
0.90
8.71
5.57
995
0.955
0550
6078
C.71
D656
6.09
6:44
0.90
6072
65
6.8
1
(e) TTO
Figure 20: Balanced accuracy (arithmetic mean of recall and specificity) achieved by all evaluated models for one-step forward prediction task with various number of balls and fractions of available views. All models were trained on video sequences with balls and constant number of crops views (10 views, corresponding to the right-most columns labelled 1.0). The color map is consistent across all plots.
29

---

Published as a conference paper at ICLR 2021
0.88
0.55
6.001
0.913
0.003
0.00
K:00
0051
0:489
3370
0.913
CAB
3.8
2024
6924
8.46
5.25
8922
0.943
5.46
5915
8902
0326
3920
0.923
3.14
5.004
2658
8.72
0.00
0.26
03:44
6355
C.75
0.00
C150
Fracdinn
Fraction
S2GRU
(b) RIMs
0.54
0.54
2:30
1992
0.91
8:107
0.40
PM
1:0
0341
5111
0.35
000
6.65
8337
6.57
6.544
8555
0.907
0.92
6.500
8G4
0537
6.44
8:46
6552
3553
0,90
6:70
0412
6823
6.27
0.MG
C750
0417
6435
0.32
(c) RMC
(d) LSTM
0.866
5.323
0.559
6.722
8.90
5.555
3397
1,00
0.99
C.05
0.13
0.51
8.50
8.544
0.955
6072
0637
6859
6522
0.004
65
01
6.
10
(e) TTO
Figure 21: F1-Score (harmonic mean of precision and recall) achieved by all evaluated models for one-step forward prediction task with various number of balls and fractions of available views. All models were trained on video sequences with balls and constant number of crops views (10 views, corresponding to the right-most columns labelled 1.0). The color map is consistent across all plots.
30

---

Published as a conference paper at ICLR 2021
Model % of Active Agents
LSTM
RMC
S2GRU
S2RMC
TTO
20%
0.570565
0.586541
0.642292
0.637618
0.550806
30%
0.599391
0.606114
0.660127
0.653950
0.578965
40%
0.630606
0.640435
0.678752
0.671476
0.605867
50%
0,638374
0.657472
0.688528
0.685988
0.627444
60%
0.681040
0.704552
0.713851
0.708786
0.671961
70%
0.709861
0.737436
0.734256
0.727980
0.723238
80%
0.721041
0.748138
0.738611
0.732114
0.740936
90%
0.750449
0.778647
0.755476
0.747613
0.786931
100%
0.765592
0.795049
0.763126
0.754637
0.813504
Table 4: Friendly marker F1 scores the validation set of the training distribution. Larger numbers are better.
Model % of Active Agents
LSTM
RMC
S2GRU
S2RMC
TTO
20%
0.323482
0.326685
0.435318
0.377538
0.297192
30%
0.345108
0.350621
0.491934
0.433945
0.323736
40%
0.373612
0.387733
0.540163
0.485278
0.350733
50%
0,385550
0.406048
0.552589
0.510371
0,371088
60%
0.430793
0,481986
0.599470
0.566149
0.435724
70%
0.497964
0.590214
0.635928
0.606039
0.539652
80%
0.579952
0.649277
0.650682
0.623040
0.617973
90%
0.657643
0.694158
0.675294
0.655581
0.699008
100%
0.677952
0.715929
0.689669
0.672186
0.737745
Table 5: Unit-type marker (macro averaged) Fl scores on the validation set of the training distribution. Larger numbers better.
F.1.2 ROBUSTNESS TO DROPPED VIEWS
In this section, evaluate the robustness of all models dropped crops on in-distribution and OOD data. We measure the performance metrics on one-step forward prediction task on all datasets (with 1-6 balls), albeit by dropping given fraction of the available input observations.
Figure 220 and 21 visualize the performance of all evaluated models. We find that S2GRU maintains performance on OOD data even with fewer views (or crops) than it was trained on. Interestingly, find that the time trav-rravelling oracle (TTO), while robust OOD, is adversely affected by the number of available views. This could be because unlike the other models, it cannot leverage the temporal information to compensate for the missing observations.
F.2 STARCRAFT2
F.2.1 TABULAR RESULTS
The results used to plot Figure 8 can be found tabulated in Tables 4
and
Thehee random pixel positions the same for all models, hut change between time-steps
31

---

Published as a conference paper at ICLR 2021
Model % of Active Agents
LSTM
RMC
S2GRU
S2RMC
ТTO
20%
-0.014035
00.033599
-0.011491
-0.011921
-0.014144
30%
-0,013355
-0.012747
-0.010631
-0,011101
-0,013539
40%
-0.012567
10.0
-0.009906
-0.010367
-0,012916
50%
-0,012220
-0,011305
-0.009637
-0,009887
-0,012481
60%
-0,010888
-0.009799
-0.008751
-0.009034
-0.010929
70%
-0.009738
-0.008469
-0.008068
-0.008359
-0,009184
80%
-0.009081
-0.008027
-0.007873
-0.008162
-0.008466
90%
-0.007970
-0.007180
-0.007347
-0.007615
-0.007038
100%
-0.007638
-0.006823
-0.007103
-0.007362
-0.006401
Table HECS Negative MSE on the validation set of the training distribution. Larger numbers better.
Model
LSTM
RMC
S2GRU
S2RMC
TTO
% of Active Agents
20%
-0.303051
0.300892
-0.141989
0.116553
00.33499
30%
-0.258878
-0.256878
-0.126037
-0.137025
-0.347899
40%
-0.216924
-0,211048
-0.113317
-0.126882
-0.276596
50%
-0.206582
-0.191644
-0.108643
-0.113293
-0.245019
60%
-0.158170
-0.142643
-0.094380
-0.099989
-0.175233
70%
-0.126446
-0.109634
-0.084129
-0.095527
-0.120694
80%
-0.111735
-0.099229
-0.081624
-0.086723
-0.104133
90%
-0.082463
-0.074518
-0.073439
-0.078197
-0.071243
100%
-0.070488
-0.063183
-0.069856
-0.074041
-0.057276
Table 7: Log Likelihood (negative loss) on the validation of the training distribution. Larger numbers are better.
32


Reviews:
[{"review_id": "uR5-Os7-4Eb", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a recurrent state space model for partially observed data, where at any given time, only a partial subset of observations are accessible to the model to make future predictions. This is an important and challenging problem in ML as we often only have access to a partial view of temporally evolving data.\n\nThe authors propose uses a recurrent model, which models a dynamically evolving partially observed process by introducing spatio-temporal interactions across multiple recurrent neural nets. The key contribution of their work is the use of a shared embedding space, that allows them to employ a mechanism to modulate spatial and temporal interactions between the RNNs.\n\nThe proposed method learns a mapping function for both, the embeddings corresponding to the RNNs, and the embedding corresponding to the observations and locations, to a shared metric space. Since the functions used to map the observations and locations, and the embeddings vectors associated with each RNN, are learned, the joint learning objecting corresponds to learning a topological structure on the points in this shared metric space. Additionally, the authors propose a truncated kernel to attenuate effects of observations far apart in time. \n\nThe final learning objective then encourages the shared metric space to be topologically organized in a way such that the embedding space is partitioned into regions, where each region places a varying “responsibility” over each of the different RNNs. Since the observations are also embedded into this space, any given observation can then be mapped into the embedding space, in the embedding space, the observation embedding be close to one or more embedding points corresponding different RNNs (which in turn models the dynamics in that regime).\n\nThe idea itself is quite interesting and allows for a “soft” representation of a state space model which through the embedding respace resembles an explicit duration + factorial Markov model. The shared embedding space is an interesting idea and can be extended to other time series models as well. \n\nThe experiments along comparisons to strong baselines demonstrates the validity of their approach. In addition to the interesting core idea, the authors employ a fairly involved attention mechanism. Given the space limitation it is understandable, however, it would have been informative to have some experiments with some ablation studies i.e. what is the simplest model structure they could have used while still incorporating the shared metric space idea to make the model more parameter efficient.\n\n\nSome minor notational issues:\nD_x seems to be defined as a function space in the model, however, it is not a function space. One could think of the forward evolutions of an RNN with a memory cell having an equal representation in function space, however, I feel the function space characterization of dynamical system section 2 obfuscates the model structure. Additionally, s_t^a can be inferred to be P(x_t^a), however, s_t^a is not defined in the text.\n\n\nA few clarifications from the textual descriptions:\nCould the authors elaborate on how exactly does E(.) process all observations in parallel across t and a? \n\n“Explained by the fact that unlike recurrent models, it does not leverage the temporal dynamics to fill in the missing information due to fewer available observations.” Could you elaborate on why you don’t expect the individual RNNs to leverage the temporal dynamics?\n\nIn the experiments, does the reported metric F-1 score account for label switching?\n\nAs currently modeled, only the forward rollout of the RNNs are used, could this be extended to having a Bi-directional RNN – one might expect better global estimates as we switch from “filtering” to a “smoothing” estimate\n\nOverall, I believe this an interesting research direction and the approach proposed by the authors can be extended to other useful models. The experiments are well thought out. The paper organization could be clearer but as it stands it is easy to understand after a thorough read.\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "nfSWIaebFY", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThis paper is concerned with making predictions about the (global) state of a dynamical system in a partially observable setting, where only local observations are available. Concretely, this paper studies video prediction from glimpses and world-modeling in a multi-agent setting.\n\nTo solve this task the paper proposes Spatially Structured Recurrent Modules (S2RMs), an RNN-based dynamics model based on interacting sub-systems. S2RMs are an extension of RIMs (Goyal et al., 2019) with the main difference being that each module additionally contains an embedding vector that is meant to reflect its location in some learned metric space. This positional information can then be used to limit the interactions between modules, to distribute inputs (that come with an associated location) to the modules, and to query a particular subpart of the global state at a future point in time.  \n\nIt is shown how S2RMs outperforms a number of baselines, which are obtained by combining the query mechanism of a GQN with either an LSTM, RIMs, or a Relational RNN to model the dynamics.\n\n### Pro’s / Con’s / Justification\n\nOverall I find that the paper is reasonably well written, although the clarity can be improved significantly. The first _four_ pages essentially only focus on motivation and problem statement, which is excessive, especially since the main contribution is an improved RNN architecture. In contrast, Section 4 that describes the method leaves important details about the adapted input attention mechanisms and inter-cell attention mechanisms to the appendix. Similarly, the experiments on the grid-world are entirely discussed in the Appendix, which I have therefore not considered as part of the contribution. More generally, the problem formulation seems unnecessarily broad given the relatively narrow contribution that is made.\n\nThe considered problem setting is interesting and I believe relevant to several real-world settings. The contribution itself is rather incremental since it essentially only involves associating a location with each module in RIMs. The adaptations to the existing attention mechanisms in RIMs that this then necessitates are straightforward, although additionally having the RIM input attention over the local observations is interesting and appears novel (I couldn't figure out if the original RIMs already attend to a set-representation of the input -- i.e. where each element is a different spatial location of the learned CNN representation). However, this paper does not compare alternative ways of implementing these changes or provide an ablation, which leaves it open what the effect is of adapting the attention mechanisms in this way. I also note that requiring the local observations to include an associated spatial location is arguably a more narrow setting than that explored in RIMs (if the input attention in RIMs is additionally directed to attend over the encoded local observations) and it would have been interesting to understand the importance of this. For example, how do S2RMs compare to RIMs when a similar encoder is used for RIMs (except for using the spatial coordinate)?\n\nThe experimental evaluation indicates that S2RMs outperforms the considered baselines and most notably RIMs, although there are some concerns. Firstly, for the baselines, the aggregated representations are computed as a sum of the output of the encoder applied to each local observation, which seems like a major bottleneck. While S2RMs can use attention to attend to each patch separately, information is almost certainly lost for the baselines in this way. I can understand why this may be necessary for the LSTM, but why not let RIM attend to the encoded local observations separately? Currently, this makes it difficult to understand in what way S2RMs improve over RIMs, especially since the reported margin already is quite small. Secondly, and related to this, I believe that it is important to include an ablation of the changes proposed to the attention mechanism. For example, what happens if positional information is only used for the recurrence, but not for the input attention (i.e. only using the global term)? Or what if it is only used for input attention, but not for the recurrence? Understanding the effect of these design choices would help strengthen the contribution and make it more significant compared to RIMs.\n\nGiven these issues I can not quite yet recommend an accept at this point in time, but I would be willing to increase my score depending on the outcome of some of the experiments I have suggested. More generally, I encourage the authors to revise Section 2-4 to put more emphasis on the actual contribution and discuss the specific design in detail.\n\n### Detailed comments\n\n* Why is it necessary to choose between the product of local and global terms and only the local term for the input attention? It seems to me that the local terms should suffice, and I wonder what the effect is of having this additional term.\n\n* I would appreciate a more detailed discussion of prior work in the related work section. Currently, the paper only makes ‘sweeping claims’ about how the considered setting is different from a bunch of prior approaches. Rather, it would be more interesting to point out specific parallels, or discuss how ideas from prior work may be included in the considered set-up, or in what way those ideas have already been re-used.\n\n* I am also missing a discussion of object-centric approaches to performing physical prediction tasks, like RNEM (Steenkiste et al., 2018), OP3 (Veerapaneni et al., 2020), etc. especially in relation to the considered bouncing balls task. Arguably, on this task these methods provide the best trade-off between locally interacting sub-systems and more global modeling since each RNN specializes to a specific object, which is evidently not achieved by S2RMs. Since each RNN learns to specialize on a single object, it can easily be viewed as modeling the global state of the system through locally interacting subsystems acting on partial observation. On the other hand, it is clear that S2RMs (and RIMs by extension) offer other advantages. For example, an advantage of S2RMs is that they consider modules having their own weights, which can thereby specialize on specific interactions between objects (or break down interactions between objects further). RNEM and the likes are limited to modeling the same set of interactions between interacting subsystems given by objects, although evidently, the notion of an object can also be flexible in this case (i.e. when an object-centric representation corresponds to two physical objects in pixel space).\n\n* The visualization in Figure 5, suggests that individual modules specialize, but there also seems to be quite some redundancy. Would it not be possible to quantify the achieved modularity somehow, eg. by measuring IoU? In that case it would be interesting to see how the system behaves when the number of provided modules is insufficient to model each individual actor in the considered system, and when the number of modules is equal to the number of actors (eg objects on bouncing balls). Additionally, do you have any intuition for what happens if modules are removed at test-time? \n\n* Why are RIMs not included in Table 1? I would be surprised if it performs that much worse to S2RM, since both models perform similar on the bouncing balls task. If this is the case then please at the very least provide some intuition or insight to explain this difference. \n\n* Please include a discussion of limitations in the conclusion. \n\n* Please take a moment to go through the references and correctly cite papers that have been published.\n\n%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n\nI have improved my score following the improvements made by the authors. See my reply below for details.", "extra_scores": {}, "extra_text": {}}, {"review_id": "BcJV6wqVLgP", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n### Summary of the paper\n\nThe authors propose a novel architecture for synthesizing information from multiple local observations and making predictions into the future.\nGiven a set of observations, each with an associated location, the model maps these into an embedding space.\nA set of RNNs, each with a learned embedding that corresponds to a vector in the embedding space, is used to process observations that are close in the embedding space.\nThe activations that the RNNs can receive, exchange between each other and output for a given query location are modulated by the distance in the embedding space.\nThis is achieved through three separate attention mechanisms with a similar design: A fully non-local (all-to-all) attention layer is followed by a multiplication with a kernel that falls off (and is eventually cut off to 0) with distance.\nThis type of attention is used once for computing the inputs to the RNNs, once for the hidden state of the LSTM (exchanging information between RNN modules that are close by) and once when computing an output response given a query location.\n\n### Relation to prior work\n\nThe paper references previous work that deals with synthesizing information from multiple localized observation and positions itself reasonably with respect to prior work.\n\n### Experimental evaluation\n\nThe authors demonstrate results on two domains: A bouncing balls video prediction task which is well-known and useful but could be considered a toy task and predicting sequences from StarCraft 2 battles. In the case of StarCraft 2 the target for the prediction are raw unit stats and actions rather than images.\nThe model generally outperforms baselines on the bouncing balls task and generalizes better on the StarCraft 2 domain (but does not outperform baselines on the training task itself).\nIt's difficult for me to judge whether the model will also work on other domains based on these experiments.\n\nI'm curious whether some of the components of the model are actually needed (e.g. attention between RNN modules). In my opinion it would improve the paper significantly if some results on ablation experiments would be provided. Another aspect of the model that could be investigated is the dependence of the performance on the kernel that modulates information exchange by distance in the embedding space. What happens if the modulation is disabled?\n\n### Presentation and clarity\n\nThe paper is generally well written. But I did have some difficulties understanding the goal and approach because the paper relies on a lot of wordy exposition (some of which is highly speculative in my opinion). I believe that the paper would be clearer if the introduction was shortened. This could also leave more space to discuss the experiments in more detail or add an ablation study.\n\n### Conclusions\n\nI found the architecture introduced in the paper genuinely interesting and would like to see more work in this direction. I would give the paper a higher rating if it included ablation experiments, or if it was adjusted to simplify and shorten some of the discussion in the introduction.\n\n\n==========\n\nEdit after author comments:\nI've read the author comments and the updated version of the paper.\nAlthough the authors claim that they have shortened the introduction of the paper by 1 page, this doesn't actually seem to be the case in the last version that was uploaded, where the section titled \"Introduction\" is almost unchanged compared to the original upload (I've used the diff tool between the latest and original version).\nMaybe there was a misunderstanding and the authors have shortened a different part of the paper?\nAlthough it would have been nice to shorten and streamline the introduction to make the paper easier to read, it's not critical to my rating.\nThe added ablation experiments demonstrate that each of the different attention modules proposed in the paper improve results, which I think really improves the paper. I was originally not sure whether the complexity of the model was justified, but the new experiments demonstrate that each of the components seems to be needed.\nI've also carefully read the rest of the paper and the author comments explaining details of the tasks under study and now feel that I have a much better understanding of what was done and how the model could be used for other tasks.\nI agree with R4 that the novelty of the paper might not be groundbreaking, but I believe the paper could be relevant and interesting for other researchers who want to incorporate attention mechanisms into their architectures, so I recommend accepting the paper.\nI've increased my rating from 6 to 7.", "extra_scores": {}, "extra_text": {}}, {"review_id": "0AjQK9NmlC", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper models noisy observations from complex dynamical systems, consisting of multiple interacting subsystems, by a set of sparsely interacting recurrent networks. The interactions between the recurrent networks (or modules) are constrained to be spatially localized, this is done by embedding the position of each module in a metric space and scaling the strength of the interactions between modules using this metric.\n\nThe overall effect of this constraint is to induce spatial structure in the global dynamics of the interacting modules. The paper argues that this constraint is a good inductive bias for capturing the dynamics of systems that consist of sparsely interacting agents, like cars and traffic lights in a traffic simulation, or characters and actions in a video game.\n\nThe paper uses recurrent neural networks for modeling the individual sub-systems, as these are flexible nonlinear dynamical systems that can model the (local) dynamics of individual modules. The interactions between these systems is captured through an attention mechanism. These attention weights are modulated (scaled) by a similarity between embedded positions of pairs of modules.\n\nI found the paper interesting, but I thought the presentation could be clearer, and the paper could better demonstrate the importance of the spatial structure.\n\nFirst, as far as I can tell, the main difference between this work and previous work (specifically, the recurrent interacting modules (RIM) paper) is the addition of the positional embedding and corresponding weighted interaction terms.\n\nGiven that the key novelty is the positional embedding, I think the paper should spend more time introducing and motivating how that is implemented, and experiments testing that implementation.\n\nThe positional embedding is introduced in a few lines at the top of section 4, with the motivation that the chosen function is \"commonly used\" (and cites the original transformer paper). However, the motivation for positional encodings in NLP tasks for the transformer seems to me to be completely different from the motivation for positional or spatial embeddings in this work. This work is (largely) aimed at modeling interacting systems in the physical world, whereas the positional encoding in Vaswani et al was motivated as a simple way to introduce signals to the network that represented position. Given that these are different domains, I think the positional embedding should be better motivate here. In addition, the paper suggests that \"other choices might also be viable\". What are these other choices? Did the authors experiment with them?\n\nSecond, how are the hyperparameters for the similarity metric (epsilon and tau) chosen? These govern how close two subsystems need to be to interact, but as far as I can tell they are simply presented as constants in a table in the appendix. How does performance vary as you change these? At the very least, the main text should state how these were chosen.\n\nThird, from what I can tell, the S2GRU has less capacity than a system that models all of the interactions, such as a global LSTM. If so, then how come the LSTM underperforms on simple tasks such as the bouncing balls? Presumably, the LSTM could just learn the same interactions that are in the S2GRU. Is this underperformance due to issues with trainability (the global LSTM is harder to train?) or perhaps it still has limited capacity (e.g. it does not have enough units or layers to model the dynamics?). If the latter, it would be nice to see experiments with different numbers of units and/or more training data, showing that for large systems the LSTM is as good as (or better) than the S2GRU (since we expect the LSTM to be able to learn any nonlinear dynamics). Then, as you reduce either the number of units or training data, perhaps the S2GRU starts to outperform (as the inductive bias of the S2GRU starts to win). Basically, if the spatial structure is really an inductive bias, I would expect the benefit to go away with a higher capacity model or more training data (where inductive biases are not needed, as we can just learn directly from data).\n\nFourth, are there cases where the spatial structure is *not* a good inductive bias? For example, situations where modeling the interactions between all pairs of agents or subsystems is necessary to capture the dynamics. If so, I would like to see experiments on these problems that show that the S2GRU does *not* outperform a global model such as an LSTM. This would be a nice control to show that the limits of the inductive biases of the imposed spatial structure.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a model for dynamical systems with multiple interacting components. Each component is modeled as an RNN, and the interactions between components are functions of their distance in a learned embedding space. It's an interesting idea and well motivated inductive bias. The results were made more compelling with the addition of "ablation" studies during the discussion phase, which showed how various aspects of the model combined to yield the best performance.  Overall, this paper should be of interest to many in the ICLR community working on complex, multi-agent systems.
Decision: Accept (Poster)