

\subsection{Multi-agent trajectory modelling}

Our problem description closely follows \citet{alcorn2021baller2vec}, whom we quote here:

\begin{quote}
Let $A = \{1, 2, \dots, B\}$ be a set indexing $B$ agents and $P = \{p_{1}, p_{2}, \dots, p_{K}\} \subset A$ be the $K$ agents involved in a particular sequence.
[Let] $C_{t} = \{(x_{t,1}, y_{t,1}), (x_{t,2}, y_{t,2}), \dots, (x_{t,K}, y_{t,K})\}$ [be] an unordered set of $K$ coordinate pairs such that $(x_{t,k}, y_{t,k})$ are the coordinates for agent $p_{k}$ at time step $t$.
The ordered sequence of sets of coordinates $\mathcal{C} = (C_{1}, C_{2}, \dots, C_{T})$, together with $P$, thus defines the trajectories for the $K$ agents over $T$ time steps.
\end{quote}

In multi-agent trajectory modeling, the goal is to model a joint probability of the form:

\[
     p(\Delta x_{t,1}, \Delta x_{t, 2}, \dots, \Delta x_{t, K} | x_{1:t,1}, x_{1:t,2}, \dots, x_{1:t,K})
\]

\noindent
i.e., the joint probability of the $K$ agents' trajectories $\Delta x_{t,k}$ at time step $t$ given the agents' location histories $x_{1:t,k}$.
We note here that the common practice of \textit{simultaneously} predicting the trajectories for all of the agents at a specific time step is not required by theory.
Using the chain rule of probability, the joint probability of the agents' trajectories can be factorized as, e.g.:

\[
    p(\Delta x_{t,1}, \Delta x_{t, 2}, \dots, \Delta x_{t, K}) = p(\Delta x_{t,1}) p(\Delta x_{t, 2} | \Delta x_{t,1}) \dots p(\Delta x_{t, K} | \Delta x_{t,1}, \Delta x_{t, 2}, \dots, \Delta x_{t, K - 1})
\]

where we omit the conditional historical trajectories for brevity.
As a result, it is perfectly acceptable to generate trajectories agent-wise, using the previously generated trajectories as additional conditioning information when generating the trajectories for later agents (see Figure \ref{fig:legal_dep}).

\begin{figure}[ht]
\centering
\subfigure{\includegraphics[width=\columnwidth]{legal}}
\caption{
At inference time, a model is not \textit{required} to \textit{simultaneously} generate the trajectories for all of the agents at a specific time step.
An alternative strategy is to allow the model to generate the agents' trajectories one at a time, and let the model use the previously generated trajectories to inform the trajectories it generates for the remaining agents.
}
\label{fig:legal_dep}
\end{figure}

\subsection{\texttt{baller2vec}{} is a (conditional) generative model.}\label{sec:baller2vec_generative}

\texttt{baller2vec}{} is a recently described \textit{multi-entity} Transformer that can model sequences of \textit{sets} (the underlying data structure for multi-agent spatiotemporal systems), as opposed to sequences of individual inputs (like words in a sentence).
When used to model the game of basketball, the input at each time step for \texttt{baller2vec}{} is a set of feature vectors where each feature vector contains information about the identity and location of a player on the court.
\texttt{baller2vec}{} maps each input feature vector to an output feature vector, which is then used to ``classify'' the binned trajectory for that specific player at that specific time step.

Here, we provide a probabilistic interpretation of \texttt{baller2vec}{}, which establishes the theoretical grounds for using the chain rule to generate trajectories agent-wise at each time step in \texttt{baller2vec++}{}.
Without loss of generality, we only consider one-dimensional trajectories for a single agent here.
To briefly summarize, the outputs of the softmax function over the $n$ binned trajectories in \texttt{baller2vec}{} can be interpreted as mixture proportions for a mixture of uniform distributions with predetermined bounds that partition the Euclidean trajectory space.
Further, because $x_{t + 1} = x_{t} + \Delta x_{t}$, \texttt{baller2vec}{} is in fact a conditional generative model that assigns a probability to a sequence of trajectories given the initial position of the agent, i.e., $p(\Delta x_{1}, \Delta x_{2}, \dots, \Delta x_{T} | x_{1})$.
Using the chain rule, we decompose the joint probability of the trajectories as:

\[
    p(\Delta x_{1}, \Delta x_{2}, \dots, \Delta x_{T}) = p(\Delta x_{1}) p(\Delta x_{2} | \Delta x_{1}) \dots p(\Delta x_{T} | \Delta x_{1}, \Delta x_{2}, \dots, \Delta x_{T - 1})
\]

\noindent
to reflect their temporal structure (we omit the conditional initial position term for brevity).
Therefore, new trajectories can be generated from \texttt{baller2vec}{} with the following procedure (see Figure \ref{fig:graphical_model}):

\begin{enumerate}
    \item First, sample one of the $n$ different mixture components using the mixture proportions output from the classifier $f$ (i.e., \texttt{baller2vec}{}) conditioned on the agent's current position, i.e., $i \sim \text{Categorical}(\pi_{1}, \pi_{2}, \dots, \pi_{n})$ where $[\pi_{1}, \pi_{2}, \dots, \pi_{n}] = f(x_{t})$.
    \item Next, sample a trajectory from the uniform distribution associated with the sampled component, i.e., $\Delta x_{t} \sim \mathcal{U}(a_{i}, b_{i})$.
    \item Finally, add the sampled trajectory to the agent's input position to generate the agent's position at the start of the next time step, i.e., $x_{t+1} = x_{t} + \Delta x_{t}$.
\end{enumerate}

\begin{wrapfigure}{r}{0.5\linewidth}
\vskip -0.2in
\includegraphics[width=\linewidth]{graphical_model_tight}
\caption{
\texttt{baller2vec}{} can be viewed as a conditional generative model that assigns a probability to a sequence of trajectories given the initial positions of the agents.
Here, we show a graphical model depiction of a \texttt{baller2vec}{} model that generates a sequence of one-dimensional trajectories for a single agent.
Given the initial position of the agent (the circle containing $x_{1}$), one of $n$ different uniform distributions (the square containing $i_{1}$) is sampled using the mixture proportions ($\pi_{i}$) output by \texttt{baller2vec}{} ($f$).
The agent's trajectory (the diamond containing $\Delta x_{1}$) is then sampled from the selected uniform distribution, which has bounds $-\infty < a_{i} < b_{i} < \infty$.
At the start of the next time step, the agent's position is $x_{2} = x_{1} + \Delta x_{1}$.
Maximizing the likelihood of \texttt{baller2vec}{} as a classifier over the binned trajectories is thus equivalent to maximizing its likelihood when assuming the trajectories are generated from a mixture of uniform distributions that partition the Euclidean trajectory space (see Section \ref{sec:baller2vec_generative} for details).
}
\label{fig:graphical_model}
\end{wrapfigure}
Let $[\Delta x_{min}, \Delta x_{max})$ be an interval on the real line such that any trajectory $\Delta x < \Delta x_{min}$ or $\Delta x \geq \Delta x_{max}$ has zero density (i.e., such trajectories are humanly impossible).
Let $\{[a_{i}, b_{i})\}_{i=1}^{n}$ be a set of $n$ intervals that partition the interval $[\Delta x_{min}, \Delta x_{max})$ into $n$ bins, i.e., $\cup_{i=1}^{n} [a_{i}, b_{i}) = [\Delta x_{min}, \Delta x_{max})$ and $i \neq j \implies [a_{i}, b_{i}) \cap [a_{j}, b_{j}) = \emptyset$.
Recall that the probability density function (PDF) for a uniform distribution with bounds $-\infty < a < b < \infty$ is:

\[
    p(\Delta x) = 
    \begin{cases}
        \frac{1}{b - a} & \text{for } \Delta x \in [a,b)  \\
        0               & \text{otherwise}
    \end{cases}
\]

\noindent
Letting $c_{i} = \frac{1}{b_{i} - a_{i}}$, the PDF for a mixture of uniforms with these bounds is thus:

\begin{equation}\label{eq:mix_uniforms}
    p(\Delta x) = \sum_{i=1}^{n} \pi_{i} \mathcal{U}(\Delta x; a_{i}, b_{i}) = \sum_{i=1}^{n} \pi_{i} c_{i}
\end{equation}

\noindent
where $p(\Delta x)$ is the density assigned to $\Delta x$ by the mixture, $\pi_{i}$ is the mixture proportion for the mixture component indexed by $i$ (i.e., $0 \leq \pi_{i} \leq 1$ and $\sum \pi_{i} = 1$), and $\mathcal{U}(\Delta x; a_{i}, b_{i})$ is the density assigned to $\Delta x$ by the uniform distribution with bounds $-\infty < a_{i} < b_{i} < \infty$.
Because the bounds of the uniform distributions partition $[\Delta x_{min}, \Delta x_{max})$, Equation \eqref{eq:mix_uniforms} reduces to:

\[
    p(\Delta x) = \pi_{i'} c_{i'}
\]

\noindent
where $\Delta x \in [a_{i'}, b_{i'})$ (because the other uniform distributions will assign a density of zero to $\Delta x$).
The likelihood for data $D$ (with $|D| = N$) is then:

\[
    \mathcal{L}(D) = \prod_{j=1}^{N} p(\Delta x_{j}) = \prod_{j=1}^{N} \pi_{j,i'} c_{j,i'}
\]

\noindent
where $\pi_{j,i'}$ is the mixture proportion assigned to the component with  $\Delta x_{j} \in [a_{i'}, b_{i'})$ and $c_{j,i'}$ is the associated density.
Taking the negative logarithm of the likelihood gives:

\begin{equation}\label{eq:mixture_nll}
    -\ln(\mathcal{L}(D)) = -\sum_{j=1}^{N} \ln(\pi_{j,i'}) -\sum_{j=1}^{N} \ln(c_{j,i'})
\end{equation}

\noindent
Because the bounds are fixed, the second summation is a constant, and Equation \eqref{eq:mixture_nll} becomes:

\[
    -\ln(\mathcal{L}(D)) = -\sum_{j=1}^{N} \ln(\pi_{j,i'}) + C
\]

\noindent
where $C = -\sum_{j=1}^{N} \ln(c_{j,i'})$.
Therefore, minimizing the loss of \texttt{baller2vec}{} as a classifier of binned trajectories is equivalent to minimizing the loss of the model when assuming the trajectories are generated from a mixture of uniform distributions as specified in Equation \eqref{eq:mix_uniforms}.


\subsection{Toy dataset}
Each training sample was initialized with the agents starting at $(-1, 0)$ and $(1, 0)$ on a grid in random order (i.e., the first agent could be placed to either the left or the right of the origin).
At each time step, one of nine actions (corresponding to the 3$\times$3 grid surrounding the agent) was sampled from a uniform distribution, and each of the agents was translated along this trajectory.
This process was repeated for 20 time steps (see Figure \ref{fig:toy}(a) for a sample).

\subsection{Basketball dataset}

We used the same National Basketball Association (NBA) dataset\footnote{\url{https://github.com/linouk23/NBA-Player-Movements}} employed by \citet{alcorn2021baller2vec}, whom we paraphrase here:

\begin{quote}
    The NBA dataset consists of trajectories from 631 games from the 2015-2016 season, which were split into 569/30/32 training/validation/test games, respectively.
    During training, each sequence was sampled using the following procedure: (1) randomly select a training game, (2) randomly select a starting time from the game, (3) take the following four seconds of data and downsample it to 5 Hz from the original 25 Hz, and then (4) randomly (with a probability of 0.5) rotate the court 180\degree.
    This sampling procedure gave us access to on the order of $\sim$82 million different (albeit overlapping) training sequences.
    For both the validation and test sets, $\sim$1,000 different, \textit{non-overlapping} sequences were selected for evaluation by dividing each game into $\lceil \frac{1,000}{N} \rceil$ non-overlapping chunks (where $N$ is the number of games), and using the starting four seconds from each chunk as the evaluation sequence.
\end{quote}

\subsection{Model}\label{sec:model}

Our \texttt{baller2vec++}{} and \texttt{baller2vec}{} models for the basketball dataset closely followed \cite{alcorn2021baller2vec}, and so largely resemble the original Transformer architecture \cite{vaswani2017attention}.
Specifically, the Transformer settings were: $d_{\text{model}} = 512$ (the dimension of the input and output of each Transformer layer), eight attention heads, $d_{\textrm{ff}} = 2048$ (the dimension of the inner feedforward layers), six layers, no dropout, and no positional encoding.
Each MLP (i.e., $g_{z}$, $g_{u}$, and $g_{r}$) had 128, 256, and 512 nodes in its three layers, respectively, and a ReLU nonlinearity following each of the first two layers.
The player embeddings \cite{alcorn20182vec} had 20 dimensions, and $h_{t,k}$ was a binary variable indicating the side of the frontcourt for player $p_{k}$ (i.e., the direction of his team's hoop) at time step $t$.
Lastly, the $11 \text{ ft} \times 11 \text{ ft}$ 2D Euclidean trajectory space was binned into 121 $1 \text{ ft} \times 1 \text{ ft}$ squares.

We used the Adam optimizer \cite{kingma2014adam} with an initial learning rate of $10^{-6}$, $\beta_{1} = 0.9$, $\beta_{2} = 0.999$, and $\epsilon = 10^{-9}$ to update the model parameters, of which there were $\sim$19 million.
The learning rate was reduced to $10^{-7}$ after 20 epochs of the validation loss not improving.
Models were implemented in PyTorch and trained on a single NVIDIA GTX 1080 Ti GPU for $\sim$650 epochs (seven days) where each epoch consisted of 20,000 training samples, and the validation set was used for early stopping.

For the toy dataset, we used scaled down versions of the basketball models with $d_{\text{model}} = 128$, four attention heads, $d_{\textrm{ff}} = 512$, and two layers in the Transformer.
Additionally, each MLP had two layers with 64 and 128 nodes, respectively.
The models were trained for 50 epochs of 500 samples per epoch ($\sim$10.5 minutes) using a single learning rate of $10^{-5}$.

\subsection{\texttt{baller2vec++}{} can effectively model coordinated agents in both simulated and real settings}\label{sec:effective}

For the toy dataset, the training loss for \texttt{baller2vec}{} converged to $\sim$2.2 $\approx -\ln(\frac{1}{9})$, i.e., the model was simply independently guessing the trajectories for both agents at every time step.
In contrast, the training loss for \texttt{baller2vec++}{} converged to $\sim$1.1 $\approx -\ln(\frac{1}{9}) \div 2$, which is the expected loss for a model that perfectly learns the deterministic relationship between the agents' trajectories (because the prediction for the second agent will always contribute $-\ln(1.0) = 0$ to the loss).
\begin{wrapfigure}{r}{0.7\linewidth}
\centering
\subfigure[Training sample.]{\includegraphics[width=0.2\columnwidth]{train_toy}}
\hskip 0.2cm
\subfigure[\texttt{baller2vec}{}.]{\includegraphics[width=0.2\columnwidth]{b2v_toy}}
\hskip 0.2cm
\subfigure[\texttt{baller2vec++}{}.]{\includegraphics[width=0.2\columnwidth]{b2v++_toy}}
\caption{When trained on a dataset of perfectly coordinated agent trajectories (a),  the trajectories generated by \texttt{baller2vec}{} are completely \textit{uncoordinated} (b) while the trajectories generated by \texttt{baller2vec++}{} are perfectly coordinated (c).
Animated versions can be found in the code repository.
}
\vskip -0.1in
\label{fig:toy}
\end{wrapfigure}
When generating trajectories with \texttt{baller2vec}{}, the agents are completely uncoordinated, with each agent following an independent random walk around the grid (Figure \ref{fig:toy}(b)).
In contrast, trajectories generated by \texttt{baller2vec++}{} display the same coordinated agent behavior as the training data (Figure \ref{fig:toy}(c)).

For the basketball dataset, \texttt{baller2vec++}{} achieved an average NLL of 0.472 on the test set, 8.9\% better than the average NLL for \texttt{baller2vec}{} (0.518) (see Figure \ref{fig:basketball} for trajectories generated by \texttt{baller2vec++}{} and \texttt{baller2vec}{}).
As was observed in \cite{alcorn2021baller2vec}, the trajectory bin distributions for \texttt{baller2vec}{} become much more certain after observing a portion of the sequence (Figure \ref{fig:plots}), which suggests \texttt{baller2vec}{} may be inferring some of the shared hidden variables (e.g., plays) influencing the players.
If that hypothesis was true, the performance gap between \texttt{baller2vec++}{} and \texttt{baller2vec}{} should be largest at the beginning of the sequence (before any shared hidden variables can be inferred by \texttt{baller2vec}{}).
Indeed, the average NLL for \texttt{baller2vec++}{} in the \textit{first time step} of each test set sequence (1.567) is 16.1\% better than the average NLL for \texttt{baller2vec}{} (1.869), while the average NLL for \texttt{baller2vec++}{} in the \textit{last time step} of each test set sequence (0.420) is only 9.7\% better than the average NLL for \texttt{baller2vec}{} (0.465) (see Figure \ref{fig:plots}).

\begin{figure}[h]
\centering
\subfigure{\includegraphics[width=\columnwidth]{plots}}
\caption{\textbf{Left}: when modeling the trajectories of professional basketball players, the performance gap between \texttt{baller2vec++}{} and \texttt{baller2vec}{} is largest at the beginning of the sequence, before shared unobserved variables can be inferred by \texttt{baller2vec}{}.
Each bar indicates a model's average NLL over the entire test set for that particular time step.
For full sequences, \texttt{baller2vec++}{} outperforms \texttt{baller2vec}{} by 8.9\%.
\textbf{Right}: the joint probability assigned to a sequence by \texttt{baller2vec++}{} is approximately permutation invariant with respect to the order of the agents.
For each point, its $x$ value indicates \texttt{baller2vec++}{}'s average NLL for a test set sequence using the \textit{original} order of the agents in the sequence, while its $y$ value indicates \texttt{baller2vec++}{}'s average NLL \textit{for the same sequence} with the order of the agents \textit{shuffled}.
The shuffled average NLLs are highly correlated with their corresponding unshuffled average NLLs.
}
\vskip -0.2in
\label{fig:plots}
\end{figure}

\subsection{\texttt{baller2vec++}{} makes better predictions when conditioned on concurrent trajectory information from other agents}\label{sec:conditioning}

Implicit in much of our discussion has been the intuition that providing a model with additional (relevant) information will improve its performance.
To empirically test this conjecture, we compared the performance of \texttt{baller2vec++}{} when predicting the trajectory of a specific basketball player placed in the \textit{first position} of the player order (i.e., when $k = 1$) vs. predicting the trajectory for that \textit{same player} placed in the \textit{last position} (i.e., when $k = 10$).
Specifically, for each player in each test sequence, we calculated the NLL of the player's trajectory in the first time step\footnote{Because, as previously discussed, the benefits of \texttt{baller2vec++}{} were most pronounced in the first time step.} with the player in the \textit{first position} of the player order.
Next, we moved the player to the \textit{last position} of the player order, and then randomly shuffled the remaining nine players 10 times, calculating the NLL for the player in the last position each time.
Finally, we calculated the average percent change in the last position NLLs relative to their corresponding first position NLLs.
On average, moving a player from the first to the last position improved the NLL for the player's trajectory by 14.6\%.

\subsection{The joint probability assigned to a sequence by \texttt{baller2vec++}{} is approximately permutation invariant with respect to the order of the agents}\label{sec:permutation}

To determine whether or not \texttt{baller2vec++}{} respects the fact that any ordering of a chain rule decomposition of a joint probability produces the same value, we measured how much the average NLL for each test sequence in the basketball dataset varied when the order of the agents changed.
Specifically, for each test set sequence, we shuffled the order of the agents 10 times.
Then, for each permuted sequence, we calculated the percent error\footnote{See: \url{https://en.wikipedia.org/wiki/Approximation_error\#Formal_Definition}.} in the average NLL relative to the original, \textit{unshuffled} sequence.
Across all test sequences, the average percent error was only $\pm1.5$\%.
Further, as can be seen in Figure \ref{fig:plots}, the shuffled average NLLs are highly correlated with their corresponding unshuffled average NLLs (Pearson correlation coefficient = 0.997), i.e., the joint probability assigned to a sequence by \texttt{baller2vec++}{} is approximately permutation invariant with respect to the order of the agents.

\section{Introduction and Related Work}
\input{Introduction}

\section{Background}\label{sec:background}
\input{Background}

\section{Model Architecture}\label{sec:baller2vec++}
\input{Architecture}

\section{Experiments}
\input{Experiments}

\section{Results}
\input{Results}

\section{Conclusion and Future Work}\label{sec:conclusion}
\input{Conclusion}

\section{Author Contributions}
MAA conceived and implemented the architecture, designed and ran the experiments, and wrote the manuscript.
AN partially funded MAA, provided the GPUs for the experiments, and commented on the manuscript.

\section{Acknowledgements}
We would like to thank Jan Van Haaren for his helpful suggestions on how to improve the manuscript.

