SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Miltiadis Kofinas
University of Amsterdam m.kofinas@uva.nl
Erik J. Bekkers
University of Amsterdam e.j.bekkers@uva.nl
Naveen Shankar Nagaraja
BMW Group
Naveen-Shankar.Nagaraja@bmw.de
Efstratios Gavves
egavves@uva.nl
Systems of interacting objects often evolve under the influence of field effects that govern their dynamics, yet previous works have abstracted away from such effects, and assume that systems evolve in a vacuum. In this work, we focus on discovering these fields, and infer them from the observed dynamics alone, without directly observing them. We theorize the presence of latent force fields, and propose neural fields to learn them. Since the observed dynamics constitute the net effect of local object interactions and global field effects, recently popularized equivariant networks are inapplicable, as they fail to capture global information. To address this, we propose to disentangle local object interactions –which are  $\mathrm{SE}(n)$  equivariant and depend on relative states– from external global field effects –which depend on absolute states. We model interactions with equivariant graph networks, and combine them with neural fields in a novel graph network that integrates field forces. Our experiments show that we can accurately discover the underlying fields in charged particles settings, traffic scenes, and gravitational n-body problems, and effectively use them to learn the system and forecast future trajectories.
Systems of interacting objects are omnipresent in nature, with examples ranging from the subatomic to the astronomical scale –including colliding particles and n-body systems of celestial objects– as well as settings that involve human activities, governed by social dynamics, like traffic scenes. The majority of these systems does not evolve in a vacuum; instead, systems evolve under the influences of underlying fields. For example, electromagnetic fields may govern the dynamics of charged particles, while galaxies
swirl around supermassive black holes that create gravitational fields. In traffic scenes, the road network and traffic rules govern the actions of traffic scene participants. Despite the ubiquity of fields, previous works on modelling interacting systems have only focused on the in vitro case of systems evolving in a vacuum.
Earlier work on learning interacting systems proposed graph networks [3, 23, 40]. Recently, state-of-the-art methods for interacting systems propose equivariant graph networks [50, 41, 24, 5, 10] to
Figure 1: N-body system with underlying gravitational field. We uncover fields that underlie interacting systems using only the observed trajectories.
model dynamics while respecting the symmetries that often underlie them. These networks exhibit increased robustness and performance, while maintaining parameter efficiency due to weight sharing. They are, however, not compatible with underlying field effects, since they can only capture local states, such as relative positions, while fields depend on absolute states (e.g. positions or orientations). In other words, global fields violate the strict equivariance hypothesis.
Within the context of modelling interacting systems, a function  $f$  that predicts future trajectories is SE(3) equivariant -equivariant to the special Euclidean group of translations and rotations - if  $f(\mathbf{R}\mathbf{x} + \boldsymbol{\tau}) = \mathbf{R}f(\mathbf{x}) + \boldsymbol{\tau}$  for a translation vector  $\boldsymbol{\tau}$  and a rotation matrix  $\mathbf{R}$ . While strict equivariance holds in idealized settings, it does not hold in many real-world settings. That is, even if the symmetries exist in a particular setting, they only manifest themselves in local interactions, yet they are entangled with global effects that stem from absolute states. N-body systems from physics, for example, exhibit E(3) symmetries, since gravitational forces only depend on relative positions. Dynamics, however, may be influenced by external force fields, e.g. black holes, which are either unknown or not subject to transformations. Thus, strict equivariance is violated, since equivariant object interactions are entangled with global field effects.
We make the following contributions. First, we introduce neural fields to discover global latent force fields in interacting dynamical systems, and infer them by observing the dynamics alone. Second, we introduce the notion of entangled equivariance that intertwines global and local effects, and propose a novel architecture that disentangles equivariant local object interactions from global field effects. Third, we propose an approximately equivariant graph network that extends equivariant graph networks by using a mixture of global and local information. Finally, we conduct experiments on a number of field settings, including real-world traffic scenes, and extending state-of-the-art setups from the literature. We observe that explicitly modelling fields is mandatory for effective future forecasting, while their unsupervised discovery opens a window for model explainability.
We term our method Aether, inspired by the postulated medium that permeates all throughout space and allows for the propagation of light.
Interacting dynamical systems An interacting dynamical system comprises trajectories of  $N$  objects in  $d$  dimensions,  $d \in \{2, 3\}$ , recorded for  $T$  timesteps. The snapshot of the  $i$ -th object at timestep  $t$  describes the state  $\mathbf{x}_i^t = [\mathbf{p}_i^t, \mathbf{u}_i^t]$ ,  $i \in \{1, \dots, N\}$ ,  $t \in \{1, \dots, T\}$ , where  $\mathbf{p} \in \mathbb{R}^d$  denotes the position and  $\mathbf{u} \in \mathbb{R}^d$  denotes the velocity, using  $[\cdot, \cdot]$  to denote vector concatenation along the feature dimension. We are interested in forecasting future trajectories, i.e. predict the future states for all objects and for a number of timesteps. Interacting dynamical systems can be naturally formalized as spatio-temporal geometric graphs [3, 23, 14],  $\mathcal{G} = \{\mathcal{G}^t\}_{t=1}^T$ , with graph snapshots  $\mathcal{G}^t = (\mathcal{V}^t, \mathcal{E}^t)$  at different time steps. The set of graph nodes  $\mathcal{V}^t = \{v_1^t, \ldots, v_N^t\}$  describes the objects in the system;  $v_i^t$  corresponds to  $\mathbf{x}_i^t$ . The set of edges  $\mathcal{E}^t \subseteq \left\{\left(v_j^t, v_i^t\right) \mid \left(v_j^t, v_i^t\right) \in \mathcal{V}^t \times \mathcal{V}^t\right\}$  describes pair-wise object interactions;  $\left(v_j^t, v_i^t\right)$  corresponds to an interaction from node  $j$  to node  $i$ . Finally,  $\mathcal{N}(i)$  denotes the neighbors of node  $v_i$ .
Local coordinate frame graph networks Local coordinate frame graph networks have been popularized in recent years [24, 26, 10, 52, 20, 30] as a method to achieve SE(3) – or E(3) – equivariance, due to their low computational overhead and high performance. Kofinas et al. [24] proposed LoCS and introduced local coordinate frames for all node-objects at all timesteps. They define augmented node states  $\mathbf{v}_i^t = [\mathbf{p}_i^t,\boldsymbol {\omega}_i^t,\mathbf{u}_i^t]$ , where  $\boldsymbol{\omega}_{i}^{t}$  denotes the angular position of node  $i$  at timestep  $t$ . Kofinas et al. [24] use velocities as a proxy to angular positions, while Luo et al. [26] use another network that predicts latent orientations. Each local coordinate frame is translated to match the target object's position and rotated to match its orientation. Considering the representation of node  $j$  in the local coordinate frame of node  $i$ , denoted as  $\mathbf{v}_{j|i}^{t}$ , they first compute the relative positions  $\mathbf{r}_{j,i}^{t} = \mathbf{p}_{j}^{t} - \mathbf{p}_{i}^{t}$  and then they rotate the state using the matrix representation of the angular position  $\mathbf{Q}(\boldsymbol{\omega}_i^t)$ :
$$
\mathbf {v} _ {j \mid i} ^ {t} = \tilde {\mathbf {R}} \left(\boldsymbol {\omega} _ {i} ^ {t}\right) ^ {\top} \left[ \mathbf {r} _ {j, i} ^ {t}, \boldsymbol {\omega} _ {j} ^ {t}, \mathbf {u} _ {j} ^ {t} \right], \tag {1}
$$
where  $\tilde{\mathbf{R}} (\boldsymbol{\omega}_i^t) = \mathbf{Q}(\boldsymbol{\omega}_i^t)\oplus \mathbf{Q}(\boldsymbol{\omega}_i^t)\oplus \mathbf{Q}(\boldsymbol{\omega}_i^t)$ , and  $\oplus$  denotes a direct sum. LoCS then proposes a graph neural network [42, 25, 13] that uses local states:
$$
\mathbf {h} _ {j, i} ^ {t} = f _ {e} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \mathbf {v} _ {i | i} ^ {t} \right]\right), \tag {2}
$$
$$
\boldsymbol {\Delta} \mathbf {x} _ {i | i} ^ {t + 1} = f _ {v} \left(g _ {v} \left(\mathbf {v} _ {i | i} ^ {t}\right) + C \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {t}\right), \tag {3}
$$
where  $f_{v}, f_{e}$ , and  $g_{v}$  are MLPs, and  $C = 1 / |\mathcal{N}(i)|$ . The output of this graph network comprises differences in positions and velocities from the previous time step, in the local frame of each object. Since these outputs are invariant, LoCS performs an inverse transformation to convert them back to the global coordinate frame and achieve equivariance,  $\mathbf{x}_i^{t + 1} = \mathbf{x}_i^t +\mathbf{R}(\boldsymbol{\omega}_i^t)\cdot \Delta \mathbf{x}_{i|i}^{t + 1}$ , where  $\mathbf{R}(\boldsymbol{\omega}_i^t) = \mathbf{Q}(\boldsymbol{\omega}_i^t)\oplus \mathbf{Q}(\boldsymbol{\omega}_i^t)$ .
Neural fields Finally, we make a brief introduction to neural fields. Neural fields, or coordinate-based MLPs, are a class of neural networks that parameterize fields using neural networks (see Xie et al. [54] for a survey). They take as input states like spatial coordinates and predict some quantity. Neural fields can learn prior behaviors and generalize to new fields via conditioning on a latent variable  $\mathbf{z}$  that encodes the properties of a field. Perez et al. [33] proposed Feature-wise Linear Modulation (FiLM), a conditioning mechanism that modulates a signal. It comprises two sub-networks  $\alpha, \beta$  that perform multiplicative and additive modulation to the input signal, and can be described by  $\mathrm{FiLM}(\mathbf{h},\mathbf{z}) = \alpha (\mathbf{z})\odot \mathbf{h} + \beta (\mathbf{z})$ , where  $\mathbf{z}$  is the conditioning variable,  $\mathbf{h}$  is the signal to be modulated, and  $\alpha, \beta$  are MLPs that scale the signal, and add a bias term, respectively.
In this section, we present our method, termed Aether. First, we describe the notion of entangled equivariance, and introduce our architecture that disentangles global field effects from local object interactions. Then, we continue with the description of the neural field that infers latent fields by observing the dynamics alone. Finally, we formulate approximately equivariant global-local coordinate frame graph networks. We note that throughout this work, we focus on fields that are unaffected by the observable objects and their interactions thereof.
3.1 Aether
Interacting dynamical systems rarely evolve in a vacuum, rather they evolve under the influence of external field effects. While object interactions depend on local information, the underlying fields depend on global states. On the one hand, locality in object interactions stems from the fact that dynamics obey a number of symmetries. By extension, object interactions are equivariant to a particular group of transformations. On the other hand, field effects are non-local; they depend on absolute object states. Thus, strict equivariance is violated, since equivariant object interactions are entangled with global field effects. We refer to this phenomenon as entangled equivariance.
As an example, in Figure 2 we observe a system of two objects that evolve in a gravitational field. The arrows positioned on the objects represent the forces exerted on them. One constituent of the net force is caused by object interactions, and is thus equivariant, while the other can be attributed to the gravitational pull. However, we can only observe the net force at each particle, i.e. the sum of equivariant pairwise forces and non-equivariant field effects. Hence, in this system, we say that equivariance is entangled.
We now propose our architecture that disentangles local object interactions from global field effects. We model object interactions with local coordinate frame graph networks [24], and field effects with neural fields. During
training, and given a multitude of input systems, neural fields will, in principle, be able to isolate
Figure 2: Two objects in a gravitational field. We only observe the total force exerted at each particle, i.e. the sum of equivariant pairwise particle forces and global field effects.
global from local effects, since only global effects are recurring phenomena. We hypothesize that field effects can be attributed to force fields, and therefore, our neural fields learn to discover latent force fields. The pipeline of our method is shown in Figure 3. Our inputs comprise augmented states  $\{\mathbf{v}_i^t\}$  for trajectories of  $N$  objects for  $T$  timesteps. Since neural fields model global fields, they depend on absolute states. Thus, we feed the states of the trajectories  $\mathbf{v}_i^t$  -or a subset of state variables- as input to a neural field that predicts latent forces  $\mathbf{f}_i^t = \mathbf{f}(\mathbf{v}_i^t)$ .
Figure 3: The pipeline of our method, Aether. In the latent neural field (a), a graph aggregation module summarizes the input trajectories in a latent variable  $\mathbf{z}$ . Query states from input trajectories, alongside  $\mathbf{z}$ , are fed to a neural field that predicts a latent force field. In (b), a graph network integrates predicted forces with input trajectories to predict future trajectories. The graph aggregation module and the FiLM layers exist only in a dynamic field setting.
The predicted field forces can be now considered part of the node states, and further, they can be treated similarly to other state variables like velocities; as vectors, forces are unaffected by the action of translations, while they covariantly transform with rotations. Thus, moving onward, we can treat the problem setup as if we were once again back in the strict equivariance regime. We append the predicted forces  $\mathbf{f}_i^t$  for each node-object  $i$  and each timestep  $t$  to the node states, and transform them to corresponding local coordinate frames, similarly to Equation (1). Namely, the force exerted on node  $j$ , expressed in the local coordinate frame of node  $i$  is computed as:  $\mathbf{f}_{j|i}^{t} = \mathbf{Q}^{\top}(\boldsymbol{\omega}_{i}^{t})\mathbf{f}_{j}^{t}$ . We feed the new local node states to a local coordinate frame graph network as follows:
$$
\mathbf {h} _ {j, i} ^ {t} = f _ {e} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \boxed {\mathbf {f} _ {j | i} ^ {t}}, \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {f} _ {i | i} ^ {t}} \right]\right) \tag {4}
$$
$$
\Delta \mathbf {x} _ {i | i} ^ {t + 1} = f _ {v} \left(g _ {v} \left(\left[ \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {f} _ {i | i} ^ {t}} \right]\right) + C \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {t}\right) \tag {5}
$$
$$
\mathbf {x} _ {i} ^ {t + 1} = \mathbf {x} _ {i} ^ {t} + \mathbf {R} \left(\omega_ {i} ^ {t}\right) \cdot \Delta \mathbf {x} _ {i | i} ^ {t + 1}, \tag {6}
$$
where  $\mathbf{R}(\boldsymbol{\omega}_i^t) = \mathbf{Q}(\boldsymbol{\omega}_i^t) \oplus \mathbf{Q}(\boldsymbol{\omega}_i^t)$ ,  $C = 1 / |\mathcal{N}(i)|$ . The equations above are similar to Equations (2) and (3), with the addition of the highlighted parts that denote the predicted forces expressed at local coordinate frames. In practice, in most experiments, we closely follow [23, 14, 24] and formulate our model as a variational autoencoder [21, 37] with latent edge types. The exact details are presented in Appendix A.1.2.
3.2 Field discovery
Oftentimes, fields might not be directly observable for us to probe them at will and use them for supervision. For example, astronomical observations of solar systems and galaxies might not include black holes, yet we can observe their effects. Moreover, fields are often not even measurable or quantifiable, or they are defined implicitly. For instance, "social fields" that guide traffic, cannot be measured or defined explicitly, but we can safely assume they exist. Motivated by these observations, we design an architecture that performs unsupervised field discovery, while solving the surrogate supervised task of trajectory forecasting.
In this work, we aim to discover two different types of fields, which we term "static" and "dynamic" fields. Static fields refer to settings in which we have a single field shared throughout the whole dataset. On the other hand, dynamic fields refer to settings in which we have a different field for each input system, and consequently, fields also differ between train, validation, and test sets.
We now describe neural fields, used in this work, to model the underlying field effects. Neural fields depend on absolute states and predict latent force fields. When dealing with static fields, we use unconditional neural fields, i.e. neural fields that are functions only of the query states, as the field values are common across data samples. Note that unconditional neural fields are not functions of the input states; they will make the same predictions regardless of the inputs. In contrast, for dynamic fields, we use a conditional neural field, i.e. a neural field that also depends on a latent vector  $\mathbf{z} \in \mathbb{R}^{D_z}$  that represents the underlying field. The latent  $\mathbf{z}$  will be inferred from the input trajectories and can be thought of as representing unusual non-equivariant dynamics. We use  $\mathbf{z}$  to explicitly condition the neural field, and thus, its general form is  $\mathbf{f}: \mathbb{R}^d \times \mathrm{SO}(d) \times \mathbb{R}^d \times \mathbb{R}^{D_z} \to \mathbb{R}^d$ , where  $d \in \{2,3\}$ , depending on the setting.
During training, both for conditional and unconditional neural fields, we only sample the field at query states that coincide with the states of the input objects, since we only have supervision about their future trajectories there.
Static fields We start with the description of unconditional neural fields used in static field settings, since conditional neural fields share the same backbone. First, we encode the query positions using Gaussian random Fourier features [46], as follows:  $\gamma (\mathbf{p}) = [\cos (2\pi \mathbf{B}\mathbf{p}),\sin (2\pi \mathbf{B}\mathbf{p})]^{\top}$ , where  $\mathbf{p} \in \mathbb{R}^d$  are the query coordinates, and  $\mathbf{B} \in \mathbb{R}^{\frac{D_c}{2} \times d}$  is a matrix with entries sampled from a Gaussian distribution,  $\mathbf{B}_{kl} \sim \mathcal{N}(0,\sigma^2)$ . The variance  $\sigma^2$  can be chosen per task with a hyperparameter sweep.
We encode velocities using a simple linear layer  $\zeta(\mathbf{u}) = \mathbf{W}_u\mathbf{u}$ . For orientations, in  $d = 3$  dimensions, we use a unit vector representation for each angle in  $\boldsymbol{\omega} = (\theta, \phi, \psi)^\top$ ,  $\hat{\boldsymbol{\omega}} = [\cos \boldsymbol{\omega}, \sin \boldsymbol{\omega}]^\top$ . In  $d = 2$  dimensions, we use the same encoding, except that we now have a single angle  $\boldsymbol{\omega} = \boldsymbol{\theta}$ . Then, we use a linear layer to encode the orientation vectors,  $\delta(\boldsymbol{\omega}) = \mathbf{W}_{\boldsymbol{\omega}}\hat{\boldsymbol{\omega}}$ . We finally concatenate the encoded positions, orientations, and velocities in a single vector that is being fed as input to the neural field. The neural field is a 3-layer MLP with SiLU [36] activations in-between, and outputs a latent force field,  $\mathbf{f}(\mathbf{v}) = \mathrm{MLP}([\gamma(\mathbf{p}), \delta(\boldsymbol{\omega}), \zeta(\mathbf{u}))$ .
Dynamic fields The neural fields used to model the dynamic fields are conditioned on a latent vector representation  $\mathbf{z} \in \mathbb{R}^{D_z}$  that describes prior knowledge about the underlying field, and are defined as  $\mathbf{f}(\mathbf{v} \mid \mathbf{z})$ . In our case, the latent representation should "summarize" the input graph such that it isolates only global effects from the field. To that end, we employ a simple global spatio-temporal attention mechanism, similar to Li et al. [25], that aggregates the input system in a latent vector representation. First, we define object embeddings  $\mathbf{o}_i = \mathrm{GRU}\big(\mathbf{W}_g\mathbf{x}_i^{1:T}\big)$ , where  $\mathbf{W}_g$  is a matrix used to linearly transform the inputs, and GRU is the Gated Recurrent Unit [7]. We also define temporal embeddings  $\mathbf{t} = \mathrm{PE}(t)$ , where PE are positional encodings [49]. Using these embeddings, we augment the input as  $\mathbf{s}_i^t = [\mathbf{x}_i^t, \mathbf{o}_i] + \mathbf{t}$ . The aggregation is then defined as follows:
$$
\mathbf {z} = \sum_ {i, t} \operatorname {s o f t m a x} \left(f _ {a} \left(\mathbf {s} _ {i} ^ {t}\right)\right) \cdot f _ {b} \left(\mathbf {s} _ {i} ^ {t}\right), \tag {7}
$$
where  $f_{a}:\mathbb{R}^{D_{s}}\to \mathbb{R},f_{b}:\mathbb{R}^{D_{s}}\to \mathbb{R}^{D_{z}}$  are 2-layer MLPs with SiLU activations in-between.
After having obtained a latent vector representation  $\mathbf{z}$  that summarizes the input system, we condition the neural field using FiLM [33]. We include FiLM layers after the first two linear layers of the neural field. The exact details are presented in Appendix A.1.1.
3.3 Approximate equivariance with global-local coordinate frames
Equivariant neural networks cannot capture non-local information, such as global field effects. In this work, we explicitly aim to discover these fields and disentangle them from local object interactions. An alternative, or rather complementary approach, would be to directly combine global and local information, following the recently proposed notion of approximate equivariance [51]. Starting from LoCS [24], we can integrate global information and still operate in local coordinate frames
by defining an auxiliary node-object corresponding to the global coordinate frame, i.e. an object positioned at the origin, and oriented to match the x-axis.
Similar to all objects in the system, the full state of the origin node  $\mathcal{O}$  comprises the concatenation of its position and velocity,  $\mathbf{x}_{\mathcal{O}} = [\mathbf{p}_{\mathcal{O}}, \mathbf{u}_{\mathcal{O}}]$ . We use an "artificial" velocity that matches the  $x$ -axis in order to compute a non-degenerate frame. As such, we have  $\mathbf{x}_{\mathcal{O}} = [0, \hat{\mathbf{x}}]$ . The origin state can be expressed in the local coordinate frame of the  $i$ -th object similarly to Equation (1), as follows:
$$
\mathbf {v} _ {\mathcal {O} | i} ^ {t} = \mathbf {R} _ {i} ^ {t \top} \left[ \mathbf {p} _ {\mathcal {O}} ^ {t} - \mathbf {p} _ {i} ^ {t}, \mathbf {u} _ {\mathcal {O}} ^ {t} \right] = \mathbf {R} _ {i} ^ {t \top} \left[ - \mathbf {p} _ {i} ^ {t}, \mathbf {u} _ {\mathcal {O}} ^ {t} \right]. \tag {8}
$$
Since graph networks are permutation equivariant, we need to explicitly distinguish between the origin node and other nodes. We circumvent that by augmenting each object's state with the origin node information expressed in local coordinate frames, extending Equations (2) and (3) to
$$
\mathbf {h} _ {j, i} ^ {t} = f _ {e} \left(\left[ \mathbf {v} _ {j | i} ^ {t}, \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {v} _ {\mathcal {O} | i} ^ {t}} \right]\right), \tag {9}
$$
$$
\Delta \mathbf {x} _ {i | i} ^ {t + 1} = f _ {v} \left(g _ {v} \left(\left[ \mathbf {v} _ {i | i} ^ {t}, \boxed {\mathbf {v} _ {\mathcal {O} | i} ^ {t}} \right]\right) + \frac {1}{| \mathcal {N} (i) |} \sum_ {j \in \mathcal {N} (i)} \mathbf {h} _ {j, i} ^ {t}\right). \tag {10}
$$
This approach pushes the information in the node states, and removes the need to add the origin node to the actual graph. We term this method G-LoCS (Global-Local Coordinate FrameS). In practice, similar to Aether, we formulate G-LoCS as a variational autoencoder [21, 37] with latent edge types. The full details are presented in Appendix A.2. Finally, in practice, we integrate G-LoCS in Aether, since it can enhance the performance of our method.