GroupNet: Multiscale Hypergraph Neural Networks for Trajectory Prediction with Relational Reasoning

In plain words

Predicting where several moving agents go next is easier when you know which of them move together. Many earlier models look only at pairs of agents: does A affect B? This paper asks a broader question: which sets of agents act as one group, and how strongly does each group pull on its members?

GroupNet builds those groups from the agents' past movements. It tries several group sizes at once (for NBA, groups of 2, 5 and 11 agents), so it sees pairs, small clusters and groups of all agents. For each group it learns three things: how strong the group's influence is, which kind of interaction it is, and how that interaction changes an agent's motion. Nobody labels the true interaction types; the model learns them while it predicts.

The authors put this module inside a standard prediction model that samples several possible futures (a CVAE). It improves that model, and it also improves some existing predictors when plugged into them. Its main test is NBA basketball tracking data, with ETH-UCY and drone-video (SDD) data as further tests.

TL;DR

GroupNet is a trajectory-prediction module. It builds hyperedges (agent groups) over agents at several sizes, learns a strength, a soft category and a category-specific MLP for each group, and passes messages over them. Inside a CVAE it has the best NBA numbers in the paper and the best SDD ADE; on ETH-UCY it is close to the best rather than first tab. 2 tab. 3 tab. 4. The released code covers NBA only, and the appendix that holds several training values is not in the LaTeX source we have.

Problem

Future motion depends on the agent's own momentum, on unpredictable intent, and on the influence of other agents. The paper targets the last term §1. Earlier social models use pooling, attention or graphs, but mostly model pairwise interactions; graph models with explicit relations usually use fixed, fully connected topologies and infer only an interaction category, not its strength §1 §2. The task is to predict future positions of all N agents from past positions (the paper's notation: past length T_p, future length T_f) §3. The paper calls the auxiliary goal of inferring interactions without labels "unsupervised relational reasoning" §3.

Evaluation uses NBA SportVU basketball tracking (2.0 s observed, 4.0 s predicted), ETH-UCY and the Stanford Drone Dataset (SDD) §6.1 §6.3, plus synthetic physics simulations where true interactions are known §6.2.

Method

The model has two stages. First, a multiscale hypergraph is inferred from agent embeddings §4.1. Second, interaction embeddings are computed and passed back to the agents over that hypergraph §4.2. Third, GroupNet is used as the social encoder inside a CVAE prediction system, with a residual decoder and a three-part loss §5.

Affinity. Each agent's past gives an embedding q_i (the paper uses an MLP f_Q for it) §4.1. Embeddings are normalised, and the affinity A_ij is their cosine similarity eq. 1 §4.1.

Hyperedges at several scales. The paper sets a sequence of group sizes K^(s). Scale 0 links each node to the nodes it is most affine to, giving pairwise edges §4.1. For scales s ≥ 1, each agent i gets one hyperedge of size K^(s). The exact choice maximises the total affinity inside the group eq. 2. The paper solves this exactly (enumeration) when the number of agents is small, and greedily otherwise: start from agent i and repeatedly add the agent with the highest affinity to i §4.1. All scales use the same affinity matrix, and every scale keeps all N nodes; only the hyperedge sizes change §4.1.

# Hypergraph topology for one scene  [§4.1] [eq. 1] [eq. 2]
# Input:  queries (N, d) from the past encoder; scales = [K_1, ..., K_S]
# Output: one incidence matrix per scale; row e = members of hyperedge e

def build_hypergraph(queries, scales):
    # Normalise each agent embedding to unit length  [§4.1]
    unit = queries / norm(queries, axis=-1)            # (N, d)
    # Cosine affinity between all agent pairs  [§4.1] [eq. 1]
    affinity = unit @ transpose(unit)                  # (N, N)
    incidences = []
    for group_size in scales:
        # Hyperedge i: agent i plus group_size - 1 agents, added one at a
        # time by highest affinity to agent i (greedy form of eq. 2)
        # [§4.1]. This equals the top-k affinities of row i, self included.
        incidence = top_k_per_row(affinity, group_size)  # (N, N), 0/1
        incidences.append(incidence)
    return incidences

Message passing per scale. Node embeddings start from q_i §4.2. In the node-to-hyperedge step, each hyperedge pools its members with learned weights w_j = F_w(v_j, sum of members), giving a collective embedding z_i §4.2. From z_i the model computes a strength r_i = sigmoid(F_r(z_i)) and a category vector c_i = softmax((F_c(z_i) + g) / τ), where g is Gumbel(0, 1) noise and τ is a temperature §4.2. The interaction embedding is then e_i = r_i · sum_l c_il · F_l(sum of members), with one MLP F_l per category eq. 3 §4.2. In the hyperedge-to-node step, each agent updates as v_i ← f_v([v_i, sum of the hyperedges containing i]) §4.2. The paper says the two steps repeat "several times" per scale, without a count §4.2. Scale-specific layers are not shared, and the final agent embedding concatenates all scales §4.2.

# One round of node-to-hyperedge and hyperedge-to-node  [§4.2] [eq. 3]
# Input:  agents (N, d); incidence (E, N), row = members of a hyperedge
# Output: updated agents (N, d)

def message_round(agents, incidence, category_fns, tau):
    # Sum of members for each hyperedge  [§4.2]
    member_sum = incidence @ agents                          # (E, d)
    # Member weights depend on the hyperedge's member sum  [§4.2]
    weights = mlp_w(agents[None], member_sum[:, None])       # (E, N, 1)
    weights = weights * incidence[..., None]                 # members only
    collective = sum(weights * agents[None], axis=1)         # (E, d)
    # Strength in (0, 1) and Gumbel-softmax category  [§4.2]
    strength = sigmoid(mlp_r(collective))                    # (E,)
    gumbel = sample_gumbel(shape=(E, num_categories))        # (E, L)
    category = softmax((mlp_c(collective) + gumbel) / tau)   # (E, L)
    # Interaction embedding e_i = r_i * sum_l c_il F_l(member sum)  [eq. 3]
    edge_emb = zeros(E, d)
    for l in range(num_categories):
        edge_emb += category[:, l, None] * category_fns[l](member_sum)
    edge_emb = strength[:, None] * edge_emb                  # (E, d)
    # Hyperedge-to-node: agent i gets the sum of its hyperedges  [§4.2]
    edge_sum = transpose(incidence) @ edge_emb               # (N, d)
    return f_v(concat(agents, edge_sum))                     # (N, d)

Prediction system (CVAE). Two GroupNet encoders produce V+ from the future and V- from the past §5. Two MLPs turn [V+, V-] into a Gaussian posterior with mean μ_q and standard deviation σ_q, and the latent Z is sampled from it §5. At test time, Z is drawn from the prior N(0, λI) §5. The encoder output is [Z, V-] §5.

Residual decoder. Two decoding blocks with the same structure. Each block has a GRU encoder over the past error and two MLP output heads §5. Block 1 reads the past and the encoder output; block 2 reads the past minus block 1's reconstruction. The predicted future and the reconstructed past are sums over the two blocks §5. The residual decoder is similar to the one of Cao et al. §5.

Loss. The total loss is L = L_elbo + L_rec + L_variety §5. The ELBO term is α times the squared future error plus β times a KL term to the prior. The reconstruction term is γ times the squared past error. The variety term is the minimum over K sampled futures of the squared error §5. The paper sets α, β, γ and λ to 1.0 §6.1.

# One training step of the CVAE prediction system  [§5] [§6.1]
# Input:  past (N, T_p, 2), future (N, T_f, 2) for one scene
# Output: scalar loss

def training_step(past, future):
    # Two GroupNet encoders: past and future  [§5]
    past_emb = past_encoder(past)                            # (N, d)
    future_emb = future_encoder(future)                      # (N, d)
    # Gaussian posterior from both embeddings  [§5]
    joint = concat(future_emb, past_emb)
    mu = mlp_mu(joint)                                       # (N, d_z)
    sigma = mlp_sigma(joint)                                 # (N, d_z)
    z = mu + sigma * randn_like(mu)                          # reparameterise
    # Residual decoder: future and past reconstruction  [§5]
    pred, recon = residual_decoder(past_emb, z, past)
    # Variety term: K futures from prior samples N(0, lam * I) [§5] [§6.1]
    # Paper: min over k of one norm; K is not stated for training
    # (code uses 20 [code: train_hyper_nba.py#L28]).
    best = infinity
    for k in range(K):
        z_k = sample_prior()                                 # N(0, lam I)
        pred_k, _ = residual_decoder(past_emb, z_k, past)
        best = min(best, squared_error(pred_k, future))
    # alpha = beta = gamma = 1.0  [§6.1]
    kl = kl_gaussian(mu, sigma, prior_var=lam)               # KL to N(0, lam I)
    loss = (squared_error(pred, future) + kl
            + squared_error(recon, past) + best)
    return loss

# Residual decoder: two blocks, each reads the remaining past error  [§5]
# Input:  past_emb (N, d), latent z (N, d_z), past (N, T_p, 2)

def residual_decoder(past_emb, z, past):
    context = concat(past_emb, z)                            # (N, d_c)
    recon_1, future_1 = block_1(context, past)
    recon_2, future_2 = block_2(context, past - recon_1)
    pred = future_1 + future_2                               # sum  [§5]
    recon = recon_1 + recon_2                                # sum  [§5]
    return pred, recon

# Inference: K futures per agent from prior samples  [§5] [§6.1] [tab. 2]
# K = 20 in the NBA table (minADE20 / minFDE20)

def predict(past, K):
    past_emb = past_encoder(past)
    futures = []
    for k in range(K):
        z = sample_prior()                                   # N(0, lam I)
        pred, _ = residual_decoder(past_emb, z, past)
        futures.append(pred)
    return futures                                           # K x (N, T_f, 2)

Figures are described only in words: the architecture is in fig. 2 and the prediction system in fig. 3; the relational-reasoning visualisations are fig. 4 to fig. 6; the NBA qualitative comparison is fig. 7.

Hyperparameters

Results

Author-reported. Setups: NBA with 2.0 s observed and 4.0 s predicted, metric minADE20/minFDE20 in meters §6.3 tab. 2. SDD with a 4.8 s horizon, metric in pixels tab. 3. ETH-UCY with five subsets and the average, metric in meters; the horizon is not stated in §6.3 tab. 4. Relational reasoning on synthetic data §6.2 tab. 1. Baselines are as named in each table. The paper's text says that "no single method achieves the best performances across all three datasets" §6.3.

Table 1: interaction-category accuracy, mean ± std over 5 runs (%) tab. 1

Table 2: NBA, minADE20 / minFDE20 (m) tab. 2. * = NMMP is the previous best.

The text gives these relative gains at 4.0 s: GroupNet+CVAE reduces minADE and minFDE against NMMP by 15.0% and 17.6%; GroupNet+NMMP reduces them by 6.0% and 12.2% §6.3. These match the table values.

Table 3: SDD, minADE20 / minFDE20 (pixels), 4.8 s tab. 3. * = PECNet is the previous best.

The text reports that GroupNet+CVAE lowers minADE20 by 6.5% against PECNet, and GroupNet+PECNet lowers minADE20/minFDE20 by 3.1%/3.4% §6.3. Its minFDE20 (16.11) is higher than PECNet's (15.88) in this table.

Table 4: ETH-UCY, minADE20 / minFDE20 (m) tab. 4. * = Trajectron++ is the previous best.

The text says GroupNet+Trajectron++ lowers the average minFDE20 by 7.3% against Trajectron++ and reaches a new best §6.3. It says that the CVAE variant with GroupNet is "comparable" with the best on ETH-UCY, not first tab. 4.

Table 5: ablation of the feature module on NBA, minADE20 / minFDE20 (m) tab. 5

Table 6: NBA, group scales, minADE20 / minFDE20 (m) tab. 6

The paper says performance rises with more scales at first, then levels off §6.4. The relational-reasoning results are qualitative beyond Table 1: the learnt affinity and hyperedges follow the physical groups §6.2 fig. 4; the fig. 5 shows example trajectories for the three interaction types; the learnt strength grows with the charge on the fixed particle fig. 6.

Paper vs code

Code: MediaBrain-SJTU/GroupNet at commit ad3f20f50ed8cf05bad19189d29bf66223f1f642 (MIT). The link is on the paper's first page p. 1, which matches the criterion for an official repository. Only NBA code is in the repository: train_hyper_nba.py, test_nba.py, model/, data/dataloader_nba.py. No code is present for SDD, ETH-UCY or the physics experiments.

The README says the released pretrained checkpoint (saved_models/nba/pretrain.p) performs better than the paper's numbers and follows a fix to a loss-scaling mistake README.md README.md. We did not train or evaluate, so none of the numbers above are reproduced here.

• Scale 0. §4.1 says each node links to its most affine nodes §4.1; §6.1 says all other agents §6.1. The code builds an all-ones matrix, so every agent links to every agent, self included MS_HGNN_batch.py.
• Hyperedge search. The code takes the top-k affinities per row, self included MS_HGNN_batch.py. That equals the paper's greedy step. The exact enumeration (init_adj_attention_listall) is in the code but disabled by self.listall = False MS_HGNN_batch.py MS_HGNN_batch.py. If the scale equals the agent count, the code returns one hyperedge over all agents, not one per agent MS_HGNN_batch.py.
• Scales. §6.1 gives 2, 5, 11 for NBA §6.1. The code default is [5, 11] train_hyper_nba.py, and the all-pairs module runs separately. Table 6's best row is 1, 2, 5, 11 tab. 6. The paper does not say which module is the scale "2".
• Affinity encoder. The paper uses an MLP f_Q §4.1 and "3 layers" for all MLPs §6.1. The code uses linear layers over time-flattened features, plus a hard-coded agent-role one-hot: agents 0–4 get one flag, 5–9 another, and 10 a third GroupNet_nba.py GroupNet_nba.py. The future encoder also hard-codes 11 agents GroupNet_nba.py. The MLP helper gives one hidden layer as two linear layers MS_HGNN_batch.py.
• Message rounds. "Several" rounds §4.2; the code runs one round per scale GroupNet_nba.py MS_HGNN_batch.py.
• Weights and normalisation. The code applies a softmax over the members of each hyperedge to the attention weights MS_HGNN_batch.py. It divides aggregated messages by the agent count MS_HGNN_batch.py MS_HGNN_batch.py. Neither appears in the paper.
• Strength and category. The code multiplies the category distribution by the strength (factor * distribution) before the category MLPs MS_HGNN_batch.py. That is the paper's r · c product eq. 3. The Gumbel temperature is 0.5 MS_HGNN_batch.py, and the number of categories is 6 or 10 MS_HGNN_batch.py MS_HGNN_batch.py.
• Gumbel noise at test. model.eval() is called test_nba.py, but no module checks self.training. The Gumbel noise is drawn on every forward pass MS_HGNN_batch.py, so the category and therefore the interaction embedding are random at test time.
• Posterior parameters. The paper gives a standard deviation σ_q §5. The code splits its output into a mean and a log-variance, and forms σ = exp(0.5 · logvar) GroupNet_nba.py.
• Prior. The code's prior is N(0, I) with learn_prior off by default, matching λ = 1 GroupNet_nba.py train_hyper_nba.py.
• Loss normalisation. The paper's L has no normalisation §5. The code sums the prediction and reconstruction errors over agents, then divides by the batch size and by T_f = 10 GroupNet_nba.py. The KL is summed, then divided by batch × agents GroupNet_nba.py. The variety term is a mean over agents with no division by T_f GroupNet_nba.py. The three terms are therefore not weighted as §5's formula reads.
• KL floor. The code clamps the KL below at 2.0 with clamp_min_, which gives no KL gradient below that value GroupNet_nba.py train_hyper_nba.py. The paper does not mention this.
• Variety samples. The code draws the K samples from the prior and takes the minimum per agent GroupNet_nba.py. §5 writes one minimum over one norm §5.
• Output frame. The future is an offset from the last observed position GroupNet_nba.py. The past reconstruction is trained against absolute positions GroupNet_nba.py. The paper gives the future directly §5.
• Metric labels. The paper gives columns 1.0 s to 4.0 s tab. 2. The code reports "ADE 1.0s" as the mean of the ADE at 0.8 s and 1.2 s, and "ADE 3.0s" as the mean at 2.8 s and 3.2 s test_nba.py. The paper does not say how those two columns are computed.
• Data. The loader divides coordinates by 94/28 dataloader_nba.py and takes the first 32,500 training and 12,500 test trajectories dataloader_nba.py dataloader_nba.py. The paper states none of these.

Pitfalls

• Appendix missing. §6.1 defers dataset details and training settings to an appendix §6.1. Our LaTeX source has no appendix, so epochs, batch size, the LR decay factor and the hidden widths are not stated, and the code is the only source for them. Check the arXiv PDF before relying on any of them.
• H orientation. The paper's incidence H^(s) is N × E, with agents as rows §4.1. The code's incidence is E × N, with hyperedges as rows MS_HGNN_batch.py. Mixing the two conventions transposes the messages silently.
• Self in every hyperedge. Affinities are cosine similarities of unit vectors, so A_ii = 1 is the maximum. Top-k and greedy selection therefore always include agent i in its own hyperedge.
• Scale equal to N. The paper has one hyperedge per agent, all containing every agent §4.1. The code returns one hyperedge of shape 1 × N MS_HGNN_batch.py. A node then receives one copy of the message in the code and N copies in the paper's version, which changes the magnitude.
• Stochastic inference. Category sampling stays on at test time MS_HGNN_batch.py. Two calls on the same input can give different interaction embeddings. Fix a seed, or average over calls, if you need determinism.
• Hard-coded agents. The role one-hot and the agent count of 11 are specific to the NBA layout (5 + 5 + ball) GroupNet_nba.py. Other agent counts or orderings will not work without changes.
• Absolute coordinates. Inputs are absolute positions GroupNet_nba.py, so the model is not translation invariant. Use the same court frame as training.
• Loss weights are not the paper's. See the loss-normalisation item above. A faithful reimplementation of §5's formula will weight the terms differently from the released code.
• Best-of-K cost. Training decodes 21 futures per agent: one from the posterior sample and 20 from prior samples, in two decoder calls GroupNet_nba.py GroupNet_nba.py. Inference runs it K = 20 times.
• Scale and node count. The method is tested on 11 agents over 2.0 s of history and 4.0 s of prediction. The conclusion lists larger agent counts, longer windows and time-varying agent counts as future work §7. The paper gives no runtime or memory numbers.
• Stated vs code scales. The NBA scales in the text (2, 5, 11) are not the code default (5, 11). Use the code's value only if you want to match the release.
• Results are author-reported. The NBA table is consistent with its text, and the SDD and ETH-UCY gains in the text match the tables. Nothing here is reproduced.

Lineage

• Builds on: the CVAE framework §5, which the paper cites as Sohn et al. (2015) §5; the residual decoder from Cao et al. (2021) §5; the variety loss from Social-GAN §5; the Gumbel-softmax relaxation of Maddison et al. §4.2 (arXiv:1611.00712 per the paper's bibliography); Adam (arXiv:1412.6980 per the bibliography) §6.1. Relational predecessors are NRI, dNRI and EvolveGraph §2.
• Compared against: Social-LSTM, Social-GAN, Social-STGCNN, STGAT, STAR, PECNet, NMMP and Trajectron++ on the real datasets tab. 2 tab. 3 tab. 4; NRI on the category task tab. 1; MLP, NMMP and non-local attention in the ablation tab. 5.
• Superseded by: not known from this wiki. The code README mentions an extension, DynGroupNet, which we did not check README.md.