EqMotion: Equivariant Multi-agent Motion Prediction with Invariant Interaction Reasoning

In plain words

Several moving things, such as particles, the atoms of a molecule, the joints of a human body or people walking together, have to be predicted from their past positions. A good predictor should not care where the whole scene sits or which way it faces. If you rotate or shift the scene, the predicted motion should rotate or shift the same way. Most earlier models do not guarantee this, so they learn the same pattern again for every orientation of the data.

EqMotion keeps track of positions in a form that rotates along with the scene. It also keeps separate, unchanging summaries such as speed and turning angle, and it uses those summaries to decide which agents affect each other. The position updates are simple operations, such as weighting and clipping, that do not break the rotation rule. The authors argue that this makes the model more compact and easier to train, and that it works on four different kinds of data.

TL;DR

EqMotion is a graph-style motion predictor that is equivariant to Euclidean transformations: its outputs transform with the input, while its interaction reasoning is invariant §3.2 §4. It learns interaction categories from invariant quantities and updates positions with equivariant operations §4.2 §4.3. The authors report state-of-the-art results on particle, molecule, H3.6M skeleton and ETH-UCY pedestrian prediction, all author-reported §5.

Problem

Motion prediction asks for the future trajectories of several interacting agents given their past motion §3.1. The paper argues that the model should be equivariant to Euclidean transformations (translation, rotation and reflection) while the interaction relations stay invariant §1 §3.2. The authors say most existing methods do not guarantee this. Methods that process abstract feature vectors lose the geometric relations between agents §1. Equivariant networks that use spherical harmonics or equivariant message passing are mostly state-to-state, using one past time step to predict one future step, which limits temporal and interaction modelling §2.

Method

EqMotion has four parts: an initial feature layer, an invariant interaction reasoning module, L stacked feature-learning layers that update a geometric feature (equivariant) and a pattern feature (invariant), and an equivariant output layer §4 eq. 1-5. The network keeps two features per agent. The geometric feature G is a set of C coordinates per agent in 3D, and the pattern feature h is a D-dimensional vector per agent §4.1. The overall procedure is given in equations 1 to 5 eq. 1-5.

Initial features §4.1 eq. 1. The geometric feature is a learned linear map over the past time axis, applied to the past positions after subtracting their mean over all agents and times, and then adding the mean back §4.1. The pattern feature comes from a per-step speed sequence (the norm of the velocity) and a turning-angle sequence (the angle between consecutive velocities), which are concatenated and embedded by an MLP or an LSTM §4.1. Speed and angle are invariant to rotation and translation, so the pattern feature is invariant §4.1 app. A.

Invariant interaction reasoning §4.2 eq. 2. The module produces, for each ordered pair of agents, a categorical vector c_ij with K entries, which is a softmax over K categories with temperature tau §4.2. The inputs are the initial pattern features and the column-wise L2 distances between the agents' initial geometric features §4.2. The module runs one message-passing step that sums edge messages over the other agents, then applies a second MLP to each pair of updated features to get the category logits §4.2. The category is invariant because it uses only invariant inputs app. A, item 2.

Equivariant geometric feature learning §4.3 eq. 3. Each layer applies three operations in order:

1. Inner-agent attention over time. The attention weight for each of the C coordinates is a function of the pattern feature, and the coordinates are rescaled around the mean of all coordinates §4.3 eq. 6.
2. Inter-agent aggregation. Each ordered pair of agents gets aggregation weights e_ij in R^C. These are a category-weighted sum of K per-category MLPs of the pattern features and column-wise distances §4.3 eq. 7. Each agent is then moved by the weighted sum of its differences to the other agents §4.3 eq. 7. The graph is fully connected in the experiments §4.3. The authors use a local neighbour set only for point-cloud-like cases with many nodes, and this case is not tested in the paper §4.3.
3. Equivariant non-linearity. Learned query and key coordinates are formed from the centred geometry. For each coordinate, the query is kept when its inner product with the key is non-negative. Otherwise, the component of the query parallel to the key is removed §4.3. The mean is added back in both branches §4.3.

Invariant pattern feature learning §4.4 eq. 4. The pattern feature is updated by message passing. The edge message is an MLP of the two pattern features and the column-wise distances between the agents' geometric features. The messages are summed over the other agents and combined with the agent's own pattern feature through a second MLP §4.4. Messages use the layer's input geometry and pattern, so the geometry and pattern updates of one layer read the same previous-layer features eq. 3-4.

Equivariant output §4.5 eq. 5. After L layers, the prediction is a linear map of the centred geometric feature, plus the mean added back §4.5. The paper writes the weight as W_out in R^{T_f x D}, which does not match the shape of the geometric feature §4.5. See Pitfalls.

Theory §4.6 [Theorem 1]. Theorem 1 states that the initial geometric feature and the output are equivariant, that the initial pattern feature, the reasoning module's categories and the pattern updates are invariant, and that each geometric update is equivariant §4.6 [Theorem 1]. Corollary 1 combines these into equivariance of the whole network §4.6. The proofs are in the appendix app. A.

Optional operations described in the appendix are a discrete cosine transform of the time axis (applied before the network and inverted after it, with the mean removed and added back) app. B.1, and an extra term that adds an MLP of the speed to the geometric feature before the non-linearity app. B.2. For multi-prediction on pedestrians, the last feature layer and the output layer are repeated as 20 parallel heads, and the training loss takes the minimum over heads §5.4 app. C.

# EqMotion forward pass  [§4] [eq. 1-5]
# Input:  past_pos (M, T_p, 3), past_vel (M, T_p, 3)
# Output: future_pos (M, T_f, 3)

def eqmotion(past_pos, past_vel):
    # Pattern feature from speed and turning angle  [§4.1]
    speed = norm(past_vel, over xyz)            # (M, T_p)
    turn = angle(past_vel[t], past_vel[t - 1])  # (M, T_p), t = 0: see Pitfalls
    pattern = phi_init_h(concat(speed, turn))   # (M, D)
    # Geometric feature: linear over time, centred  [§4.1] [eq. 1]
    center = mean(past_pos, over agents and times)  # (3,)
    geo = W_init_g(past_pos - center) + center      # (M, C, 3)
    # Interaction categories from invariant inputs  [§4.2] [eq. 2]
    category = reason_interactions(geo, pattern)    # (M, M, K)
    for layer in range(L):
        # Pattern update reads the previous geometry  [§4.4] [eq. 4]
        new_pattern = pattern_update(geo, pattern)
        # Geometry update reads the previous pattern  [§4.3] [eq. 3]
        geo = geometry_update(geo, pattern, category)
        pattern = new_pattern
    # Linear readout over the coordinate axis  [§4.5] [eq. 5]
    center = mean(geo, over agents and coordinates)  # (3,)
    return W_out(geo - center) + center              # (M, T_f, 3)


def reason_interactions(geo, pattern):
    # Column-wise distances, invariant  [§4.2] [eq. 2]
    dist = norm(geo[:, None] - geo[None, :], over xyz)  # (M, M, C)
    # Broadcast the pattern pairs to (M, M, .) before concatenating
    edge_in = concat(pattern[:, None], pattern[None, :], dist)
    edge = phi_rm(edge_in)                       # (M, M, D')
    # Sum messages over the other agents  [§4.2]
    summed = sum over j of edge[:, j]            # (M, D')
    pattern2 = phi_rh(concat(summed, pattern))   # (M, D')
    edge_in2 = concat(pattern2[:, None], pattern2[None, :], dist)
    # Softmax over K categories with temperature tau  [§4.2]
    return softmax(phi_rc(edge_in2) / tau)       # (M, M, K)


def geometry_update(geo, pattern, category):
    # Inner-agent attention over the C coordinates  [§4.3] [eq. 6]
    weights = phi_att(pattern)                       # (M, C)
    center = mean(geo, over agents and coordinates)  # (3,)
    geo = weights[..., None] * (geo - center) + center
    # Inter-agent aggregation weighted by categories  [§4.3] [eq. 7]
    dist = norm(geo[:, None] - geo[None, :], over xyz)  # (M, M, C)
    edge_in = concat(pattern[:, None], pattern[None, :], dist)
    weights = sum over k of category[..., k] * phi_e_k(edge_in)  # (M, M, C)
    diff = geo[:, None] - geo[None, :]               # (M, M, C, 3)
    geo = geo + sum over j of weights[..., None] * diff  # (M, C, 3)
    # Equivariant non-linearity: project queries against keys  [§4.3]
    center = mean(geo, over agents and coordinates)
    centered = geo - center
    queries = W_Q(centered)                          # (M, C, 3)
    keys = W_K(centered)                             # (M, C, 3)
    dot = sum(queries * keys, over xyz)              # (M, C)
    unit = keys / norm(keys, over xyz)
    projected = queries - sum(queries * unit, over xyz) * unit
    return where(dot >= 0, queries, projected) + center


def pattern_update(geo, pattern):
    # Column-wise distances, invariant  [§4.4]
    dist = norm(geo[:, None] - geo[None, :], over xyz)  # (M, M, C)
    edge_in = concat(pattern[:, None], pattern[None, :], dist)
    messages = phi_m(edge_in)                    # (M, M, D)
    # Sum over neighbours; fully connected in the experiments  [§4.3] [§4.4]
    summed = sum over j of messages[:, j]        # (M, D)
    return phi_h(concat(pattern, summed))        # (M, D)

Notes on the pseudocode: "sum over j" runs over all other agents, and the pairwise tensors include the self pair, which contributes zero to the geometric difference. Shapes are written for one batch element. The pseudocode uses the paper's equations. The official code differs in several places, listed under Paper vs code.

Hyperparameters

Values are from the paper unless marked "code". "Not stated" means the paper and this card's sources do not say.

Results

Author-reported throughout. Each table lists the setup and the baselines the paper compares against.

Particle dynamics: interaction recognition accuracy and consistency (%, mean ± std over 5 runs) tab. 1. Setup: 5 particles, Springs or Charged simulation §5.1. Consistency is the share of identical recognitions over 20 random Euclidean transforms §5.1.

Future prediction on the Charged setting is shown as a displacement-error curve over time, where the paper reports EqMotion as best §5.1 fig. 3. No numbers are given for it in the text.

Molecule dynamics, MD17: ADE / FDE (x 10^-2, lower is better) tab. 4. Setup: 10 frames in, 10 out, one model per molecule §5.2 app. D.1.2. The text reports average reductions of 34.2% (ADE) and 30.1% (FDE) across the four molecules §5.2.

ETH-UCY pedestrians: ADE / FDE (m) tab. 5. Setup: 8 frames observed, 12 predicted, leave-one-out over five subsets §5.4 app. D.1.4. Reported by the paper as best-of-20 for the multi-prediction block.

Deterministic (one prediction per input) tab. 5:

Multi-prediction (best of 20) tab. 5:

The text reports FDE reductions of 10.4% (deterministic, against MemoNet's average) and 7.9% (multi-prediction, against MID's average) §5.4.

H3.6M skeleton, short-term: average MPJPE (mm) over 15 actions tab. 2. Setup: 22 joints, prediction at 80, 160, 320 and 400 ms §5.3. Per-action numbers are in tab. 9, which is not transcribed here. The text reports an average reduction of 8.6% §5.3.

H3.6M skeleton, long-term: average MPJPE (mm) over 8 actions tab. 3. Setup: prediction at 560 ms and 1000 ms §5.3. The text reports an average reduction of 3.8% §5.3.

Ablations on H3.6M short-term, average MPJPE tab. 6 tab. 7 tab. 8.

The paper also reports, in a figure, that EqMotion uses about 5% of the training data and still matches some baselines trained on all of it §5.5 fig. 5. It reports a model size under 30% of the others' sizes §5.5 fig. 6. These are shown in plots that this card does not transcribe.

Headline claims in the abstract and introduction §1: improvements of 24.0 / 30.1 / 8.6 / 9.2% on the four tasks, with state of the art on all four. See Pitfalls for which metric each figure refers to.

Paper vs code

Code is at commit 5aec2e0b61c511fa93a24138dd90da59a089084b (MIT). The repository has no deterministic entry point for ETH-UCY: main_eth_diverse.py only builds the 20-head model model_t.py. The paper's deterministic ETH numbers tab. 5 therefore cannot be reproduced from the official scripts.

1. Inner-agent attention has a residual term in code. Code returns the rescaled coordinates plus the input coordinates gcl_t.py. Paper's equation has no such term eq. 6 §4.3. The same residual appears in the n-body layer gcl_t.py.
2. Inter-agent aggregation differs in code. The category-weighted factors are summed and then passed through an extra MLP (factor_mlp) before the coordinate update gcl_t.py. Paper's aggregation has no such MLP eq. 7. The H3.6M and ETH layers sum only the first L-2 categories: range(category_num-2) gcl_t.py gcl_t.py. With L = 4, categories 3 and 4 get no weight. The n-body layer sums all K categories gcl_t.py.
3. Velocity term. The H3.6M layer adds a learned linear map of the velocity vectors to the coordinates in every layer gcl_t.py. The paper's main equations do not include this. Its optional section describes an MLP of the speed magnitude, which is a different term app. B.2.
4. DCT is always on for H3.6M. apply_dct is set to True model_t.py and applied to positions and velocities model_t.py. The paper calls it optional app. B.1.
5. Learned agent token. The H3.6M script forces add_agent_token = True main_h36m.py. The model then concatenates a learned per-agent embedding of size 22 to the pattern feature model_t.py model_t.py. The paper does not describe this.
6. Reasoning inputs. Code computes the interaction categories from coordinates and velocities together, so the distance covers 2C channels model_t.py model_t.py. It also runs a learned MLP on the distances before the edge MLP model_t.py. The paper uses only G^(0), with C channels §4.2. The temperature is 1 in code model_t.py; the paper leaves tau unstated.
7. Epsilon in the non-linearity. Code divides by k_norm_sq + 1e-4 gcl_t.py. The paper's formula has no epsilon §4.3.
8. Activation. The paper says MLPs use ReLU app. D.2. Code passes the EqMotion default nn.SiLU() into every layer model_t.py model_t.py.
9. Initial pattern feature. The paper's embedding is an MLP or LSTM §4.1. Code uses one linear layer for the speed sequence and one for the angle sequence, each of width D/2, then concatenates them model_t.py model_t.py. For the first time step, code sets the velocity to its own previous value, so the first angle is 0 model_t.py.
10. Output weight shape. The paper writes W_out in R^{T_f x D} §4.5. Code maps the C coordinates to T_f steps with a linear layer over the coordinate axis model_t.py model_t.py. The code's shape is the one that fits the geometric feature.
11. Training loss. The paper states no loss for the single-prediction settings. Code uses the mean over agents and frames of the L2 norm per joint main_h36m.py main_nbody.py. Multi-prediction on ETH uses the minimum over 20 heads of the mean L2 norm, not squared main_eth_diverse.py. The paper gives the squared version app. C.
12. ETH loss and metrics use the ego agent only. Default training uses agent index 0 only main_eth_diverse.py. The evaluation also reads agent 0 only main_eth_diverse.py. Full-agent supervision is behind --supervise_all main_eth_diverse.py main_eth_diverse.py.
13. Time-weighted loss for 25-frame H3.6M. When the output length is 25, the code weights step t by 25 / (1 + 4t/25) main_h36m.py main_h36m.py. The paper does not describe this weighting.
14. Molecule channels. Paper: C = 64 app. D.2. Code: hid_channel=32 main_md17.py.
15. Batch sizes. Paper: 50 for particle and molecule app. D.2. Code defaults: 100 main_nbody.py main_md17.py.
16. Pedestrian learning rates. Paper: ETH 8e-4, Zara1 5e-4 app. D.2. Code: all subsets except Hotel use 1e-3 main_eth_diverse.py.
17. Zara1 channels. Paper: C = 64 for all pedestrian subsets app. D.2. Code: 128 for Zara1 main_eth_diverse.py.
18. Pedestrian decay and epochs. Paper: per-subset decay rates and 50 epochs app. D.2. Code: one decay of 0.8, only with --apply_decay main_eth_diverse.py, and 60 epochs by default main_eth_diverse.py. The ETH command in the README sets neither README.md.
19. H3.6M long-term epochs and decay. Paper: 100 epochs app. D.2. Code: 80 by default main_h36m.py; the README long-term command does not set epochs and does set --apply_decay README.md.
20. Molecule frame sampling. Paper: sampling rate 20 and gap 400 app. D.1.2. Code: delta_frame = 50 main_md17.py. I did not trace how this value is used, so the two are not reconciled.
21. Seeds. The H3.6M script picks a random seed when none is given main_h36m.py. Results are not reproducible by default.

Pitfalls

• Headline percentages are not all traceable. The abstract lists 24.0 / 30.1 / 8.6 / 9.2% without naming the metric. The text gives 8.6% for H3.6M short-term and 3.8% for long-term §5.3, 34.2% (ADE) and 30.1% (FDE) for molecules §5.2, and 10.4% / 7.9% (FDE) for pedestrians §5.4. The 24.0% for particle dynamics and the 9.2% for pedestrians have no matching sentence in the body. I computed the average reduction over the four horizons in tab. 2 against SPGSN from the printed numbers as about 8.8%, which differs from the 8.6% the paper states. The exact averaging or baseline is not stated.
• Multi-prediction numbers are best-of-20. The minimum over 20 heads is an oracle selection on the test split. Compare them only with other best-of-N results, not with single-prediction numbers §5.4 tab. 5.
• Ego-only evaluation on ETH. Default ADE/FDE uses agent 0 only main_eth_diverse.py. Full-agent scores need --supervise_all during training and a separate evaluation. The paper does not say which agents are scored.
• Output shape ambiguity. W_out is written as T_f x D, but the geometric feature is C x 3 §4.5. The code uses C to T_f model_t.py. Implement the code's version or the output will not fit the paper's definition.
• Temperature and time-zero angle are unstated. tau is not given §4.2. The first turning angle is undefined in the paper; code sets it to 0 model_t.py.
• acos at +-1. The turning angle uses acos of a clamped cosine. Its derivative is unbounded at +-1, so near-parallel consecutive velocities can give unstable gradients. Code adds 1e-6 to the norms, not to the arccos argument model_t.py.
• Quadratic cost in agents. The fully connected graph gives pairwise tensors of size M x M x C x 3 in every layer §4.3 gcl_t.py. The paper says local neighbourhoods are used for large point sets but does not test them §4.3.
• Equivariance is proven, not tested. Theorem 1 covers the layers as written §4.6. The code adds residuals, an extra MLP and a learned agent token that the proof does not cover [Paper vs code, items 1-5]. Test rotation and reflection invariance of the code you ship before relying on the property.
• Eval split details. The paper gives leave-one-out for ETH-UCY app. D.1.4 and a subject split for H3.6M app. D.1.3. It does not say which checkpoint is reported. The code saves best-on-test checkpoints main_h36m.py, which select on the test split.
• Reported ablation numbers are single numbers. The ablation tables have no variance and no seed tab. 6 tab. 7 tab. 8.
• Human-motion data preparation. Joint and clip selection come from the baseline paradigm the paper cites §5.3. Other preprocessing is in the code, not the paper.

Lineage

• Positioned against: EGNN (Satorras et al.), the equivariant message passing that the authors say handles only state-to-state prediction, which EqMotion extends to sequences §2.
• Compared against: Corr. (path) and Corr. (LSTM), EGNN, NRI and dNRI on particle reasoning tab. 1; Radial Field, TFN, SE(3)-Transformer, EGNN, LSTM, S-LSTM, NRI, NMMP and GroupNet on MD17 tab. 4; S-LSTM, SGAN, Traj++, TransF, MemoNet, NMMP, PECNet, AgentFormer, GroupNet, MID and GP-Graph on ETH-UCY tab. 5; Res-sup., Traj-GCN, DMGNN, MSRGCN, PGBIG and SPGSN on H3.6M tab. 2 tab. 3.
• Superseded by: not known from this wiki.