S TOCHASTIC P REDICTION OF M ULTI -A GENT I NTERACTIONS FROM P ARTIAL O BSERVATIONS Chen Sun Google Research Per Karlsson
Google Research
Jiajun Wu CSAIL Kevin Murphy Google Research
A BSTRACT
We present a method that learns to integrate temporal information, from a learned dynamics model, with ambiguous visual information, from a learned vision model, in the context of interacting agents.
Our method is based on a graph-structured variational recurrent neural network (Graph-VRNN), which is trained end-to-end to infer the current state of the (partially observed) world, as well as to forecast future states.
We show that our method outperforms various baselines on two sports datasets, one based on real basketball trajectories, and one generated by a soccer game engine
I NTRODUCTION
Imaging watching a soccer game on television.
At any given time, you can only see a subset of the players, and you may or may not be able to see the ball, yet you probably have some reasonable idea about where all the players currently are, even if they are not in the ﬁeld of view.
(For example, the goal keeper is probably close to the goal.)
Similarly, you cannot see the future, but you may still be able to predict where the “agents” (players and ball) will be, at least approximately.
Crucially, these problems are intertwined: we are able to predict future states by using a state dynamics model, but we can also use the same dynamics model to infer the current state of the world by extrapolating from the last time we saw each agent.
In this paper, we present a uniﬁed approach to state estimation and future forecasting for problems of this kind.
More precisely, we assume the observed data consists of a sequence of video frames, v 1: T , obtained from a stationary or moving camera.
The desired output is a (distribution over a) structured representation of the scene at each time step, p ( s t | v 1: t ) , as well as a forecast into the future, p ( s t +∆ | v 1: t ) , where s kt encodes the state (e.g., location) of the k ’th agent and s t = {
s kt } Kk =1 .
∗
The classical approach to this problem (see, e.g., Bar-Shalom et al. (2011)) is to use state-space models, such as Kalman ﬁlters, for tracking and forecasting, combined with heuristics, such as nearest neighbor, to perform data association (i.e., inferring the mapping from observations to latent objects).
Such generative approaches require a dynamical model for the states, p ( s t | s t − 1 ) , and a likelihood model for the pixels, p ( v t | s t ) .
These are then combined using Bayes’ rule.
However, it is hard to learn good generative model of pixels, and inverting such models is even harder.
By contrast, our approach is discriminative, and learns an inference network to compute the posterior belief state p ( s t | v 1: t ) directly.
In particular, our model combines ideas from graph networks, variational autoencoders, and
RNNs in a novel way, to create what we call a graph-structured variational recurrent neural network (Graph-VRNN).
We have tested our approach on two datasets: real basketball trajectories, rendered as a series of (partially observed) bird’s eye views of the court; and a simple simulated soccer game, rendered using a 3d graphics engine, and viewed from a simulated moving camera.
We show that our approach can infer the current state more accurately than other methods, and can also make more accurate future forecasts.
We also show that our method can vary its beliefs in a qualitatively sensible way.
For s t v t s t + z t + h t h t h t + v t s t z t z t + example, it “knows” the location of the not , since the not much (in our simulation).
Thus it learns to “see beyond the pixels”.
In summary, our main contribution is a uniﬁed way to do state estimation and future forecasting at the level of objects and relations directly from pixels using Graph-VRNN.
We believe our technique will have a variety of other applications beyond analysing sports videos, such as self-driving cars (where inferring and predicting the motion of pedestrians and vehicles is crucial), and human-robot interaction.
R ELATED WORK
In this section, we brieﬂy review some related work.
Graph-structured and stochastic RNNs.
There are several papers that combine standard RNNs (recurrent neural networks) with graph-structured representations, such as Chang et al.
Battaglia et al.
(2016); Sanchez-Gonzalez et al.
There are also several papers that extend standard RNNs by adding stochastic latent variables, notably the variational RNN approach of Chung et al.
(2015), as well as recent extensions, such Krishnan et al.
; Fraccaro et al.
; Karl et al.
; Goyal et al.
; Buesing et al.
However, as far as we know, combining VRNNs with graphs is novel.
Predicting future pixels from past pixels.
There are many papers that try to predict the future at the pixel level (see e.g., (Kitani et al., 2017) for a review).
Some use a single static stochastic variable, as in a conditional VAE (Kingma & Welling, 2014), which is then “decoded” into a sequence using an RNN, either using a VAE-style loss (Walker et al., 2016; Xue et al., 2016) or a GAN-style loss (Vondrick et al., 2016b; Mathieu et al., 2016).
More recent work, based on VRNNs, uses temporal stochastic variables, e.g., the SV2P model of Babaeizadeh et al.
(2018) and the SVGLP model of
Denton & Fergus (2018).
There are also various GAN-based approaches, such as the SAVP approach of Lee et al.
and the MoCoGAN approach of Tulyakov et al.
The recent sequential
AIR (attend, infer, repeat) model of Kosiorek et al.
(2018) uses a variable-sized, object-oriented latent state space rather than a ﬁxed-dimensional unstructured latent vector.
This is trained using a
VAE-style loss, where p ( v t | s t ) is an RNN image generator.
Forecasting future states from past states.
There are several papers that try to predict future states given past states.
In many cases the interaction structure between the agents is assumed to be known (e.g., using fully connected graphs or using a sparse graph derived from spatial proximity), as in the social LSTM (Alahi et al., 2016) and the social GAN (Gupta et al., 2018), or methods based on inverse optimal control (Kitani et al., 2017; Lee et al., 2017).
In the “neural relational inference” method of Kipf et al.
(2018), they infer the interaction graph from trajectories, treating it as a static latent variable in a VAE framework.
Ehrhardt et al.
(2017) proposed to predict a Gaussian distribution for the uncertain future states of a sphere.
It is also possible to use graph attention networks (Veli ˇ ckovi ´ c et al., 2018) for this task (see e.g., Hoshen (2017); Battaglia et al. [Figure]
Forecasting future states from past pixels.
Our main interest is predicting the hidden state of the world given noisy visual evidence.
(Vondrick et al., 2016a) train a conditional CNN to predict future feature vectors, and several papers (e.g., (Lerer et al., 2016; Mottaghi et al., 2016; Wu et al., 2017; van Steenkiste et al., 2018;
Fragkiadaki et al., 2016; Watters et al., 2017)) predict future object states.
Our work differs by allowing the vision and dynamics model to share and update a common belief space; also, our model is stochastic, not deterministic.
Our method is related to the “deep tracking” approach of (Dequaire et al., 2018).
However, they assume the input is a partially observed binary occupancy grid , and the output is a fully observed binary occupancy grid, whereas we predict the state of individual objects using a graph-structured model.
Our method is also related to the ”backprop Kalman ﬁltering” approach of (Haarnoja et al., 2016).
However, they assume the structure and dynamics of the latent state is known, and they use the Kalman ﬁlter equations to integrate information from the dynamical prior and (a nonlinear function of) the observations, whereas we learn the dynamics, and also learn how to integrate the two sources of information using attention.
M ETHODS 3. VRNN S
We start by reviewing the VRNN approach of Chung et al.
This model has three types of variable: the observed output x t , the stochastic VAE state z t , and the deterministic RNN hidden state h t , which summarizes x ≤ t and z ≤ t .
We deﬁne the following update equations:
[Equation]
where ϕ are neural networks (see section 3.4 for the details).
VRNNs are trained by maximizing the evidence lower bound (ELBO), which can be decomposed as follows:
[Equation]
We use Gaussians for the prior and posterior, so we can leverage the reparameterization trick to optimize this objective using SGD, as explained in Kingma & Welling (2014).
In practice, we scale the KL divergence term by β , which we anneal from 0 to 1, as in (Bowman et al., 2016).
Also, we start by using the ground truth values of x t − 1 (teacher forcing), and then we gradually switch to using samples from the model; this technique is known as scheduled sampling (Bengio et al., 2015), and helps the method converge.
3. A DDING GRAPH STRUCTURE
To create a structured, stochastic temporal model, we associate one VRNN with each agent.
The hidden states of the RNNs interact using a fully connected graph interaction network.
(This ensures the model is permutation invariant to the ordering of the agents.)
To predict the observable state s kt for agent k , we decode the hidden state vector h kt using an MLP with shared parameters.
The resulting model and operations are illustrated in ﬁg. 2.
3. C ONDITIONING ON IMAGES
We can create a conditional generative model of the form p ( s 1: t | v 1: t ) by using a VRNN (where x t = s t is the generated output) by making all the equations in section 3.1 depend on v t as input.
This is similar to a standard sequence-to-sequence model, but augmented to the graph setting, and with additional stochastic noise injected into the latent dynamics.
We choose to only condition the outputs s t on the visual input v t , rather than making the hidden state, h t , depend on v t .
The reason is that we want to be able to perform future prediction in the latent state space without having to condition on images.
In addition, we want the latent noise z t to reﬂect variation of the dynamics, not the appearance of the world.
See ﬁg. 1 for a simpliﬁed illustration of our model (ignoring the graph network component).
All parameters are shared across all agents except for the visual encoder; they have a shared feature extraction ”backbone”, but then use agent-speciﬁc features to identify the speciﬁc agent.
(Agents have the same appearance across videos.)
Thus the model learns to perform data association.
The key issue is how to deﬁne the output decoder, p ( s t | z t , h t − 1 , v t ) , in such a way that we properly combine the information from the current (partially observed) visual input, v t , with our past beliefs about the object states, h t − 1 , as well as any stochastic noise z t (used to capture any residual uncertainty).
In a generative model, this step would be performed with Bayes’ rule, combining the dynamical prior with the visual likelihood.
In our discriminative setting, we can learn how to weight the information sources using attention.
In particular, we deﬁne
[Equation]
(6) where ϕ DV is the visible decoder, ϕ DH is the hidden decoder, and the α i terms are attention weights, computed as follows:
[Equation]
(7) where i ∈ { V, H } .
This is similar to the gating mechanism in an LSTM, where we either pass through the current observations, or we pass through the prior hidden state, or some combination of the two.
If the visual input is uninformative (e.g., for future frames, we set v t +∆ = 0 ), the model will rely entirely on its dynamics model.
To capture uncertainty more explicitly, we also pass in s heat t into the ϕ Si functions that compute attention, where s heat t = softmax( ϕ dec ( z t , h t − 1 , v t )) is the set of “heatmaps” over object locations.
(This is illustrated by the s t − 1 → s t edges in ﬁg. 1.)
We also replace the sample z t with its sufﬁcient statistics, ( µ t , Σ t ) , computed using the state dependent prior, ϕ prior ( h t − 1 ) .
To encourage the model to learn to forecast future states, in addition to predicting the current state, we modify the above loss function to maximize a lower bound † on log p ( s 1: T +∆ | v 1: T ) , computed as follows:
[Equation]
where λ t = max( t/T, 1) is a weighting factor that trades off state estimation and forecasting losses,
N prior t = N ( z t | ϕ prior ( h t − 1 )) is the prior, N enc t =
N ( z t | ϕ enc ( h t − 1 , s t )) is the variational posterior, the expectation is over z 1: T , and where we set v τ =
if τ > t .
3. I MPLEMENTATION DETAILS
For the image encoder, we use convolutional neural networks with random initialization.
In the case when there are multiple frames per step, we combine the image network with an S3D (Xie et al., 2018) inspired network, which adds additional temporal convolutions on spatial feature maps.
For the recurrent model, we tried both vanilla RNNs and GRUs, and found that results were similar.
We report results using the GRU model.
For the graph network, we use relation networks (Santoro et al., 2017), i.e. a fully connected topology with equal edge weights.
During training, both state estimation and future prediction losses are important, and need to be weighted properly.
When the future prediction loss weight is too small, the dynamics model is not trained well (reﬂected by log-likelihood).
We ﬁnd that scaling the losses using an exponential discount factor for future predictions is a good heuristic.
We normalize the total weights for state estimation losses and future predictions losses to be the same.
R ESULTS
In this section, we show results on two datasets, comparing our full model to various baselines.
4. D ATASETS
Since we are interested in inferring and forecasting the state of a system composed of multiple interacting objects, based on visual evidence, analysing sports videos is a natural choice.
We focus on basketball and soccer, since these are popular games for which we can easily get data.
The states here are the identities of the objects (players and ball), and their ( x, y ) locations on the ground plane.
Although specialized solutions to state estimation for basketball and soccer already exist (e.g., based on wearable sensors (Sanguesa, 2017), multiple calibrated cameras (Manaﬁfard et al., 2017), or complex monocular vision pipelines (Rematas et al., 2018)), we are interested in seeing how far we can get with a pure learning based approach, since such a method will be more generally applicable.
Basketball.
We use the basketball data from Zhan et al.
Each example includes the trajectories of 11 agents (5 offensive players, 5 defensive players and 1 ball) for 50 steps.
We follow the standard setup from Zhan et al.
and just model the 5 offensive players, ignoring the defense.
However, we also consider the ball, since it has very different dynamics than the players.
Overall, there are 107,146 training and 13,845 test examples.
We generate bird’s-eye view images based on the trajectories (since we do not have access to the corresponding video footage of the game), where each agent is represented as a circle color-coded by its identity.
To simulate partial observation, we randomly remove one agent from the rendered image every 10 steps.
Some example rendered images can be found in ﬁg. 3b.
Soccer.
To evaluate performance in a more visually challenging environment, we consider soccer, where it is much harder to get a bird’s-eye view of the ﬁeld, due to its size.
Instead most soccer videos are recorded from a moving camera which shows a partial proﬁle view of the ﬁeld.
Since we do not have access to ground truth soccer trajectories, we decided to make our own soccer simulator, which we call Soccer World , using the Unity game engine.
The location of each player is determined by a hand-designed “AI”, which is a probabilistic decision tree based on the player’s current state, the states of other players and the ball.
The players can take actions such as kicking the ball and tackling.
The ball is driven by the players’ actions.
Each game lasts 5 minutes.
For each game, a player with the same id is assigned to a random position to play, except for the two goal keepers.
We render the players using off-the-shelf human models from Unity.
The identities of the players are color-coded on their shirts.
The camera tracks the ball, and shows a partial view of the ﬁeld at each frame.
See ﬁg. 3a for some visualizations.
We create 700 videos for training, and 300 videos for test.
We apply a random sliding window of 10 seconds to sample the training data.
The test videos are uniformly [Figure] segmented into 10-second clips, resulting in 9000 total test examples.
We plan to release the videos along with the game engine after publication of the paper.
4. E XPERIMENTAL SETUP
The ﬁrst task we evaluate is inferring the current state (which we deﬁne to be the 2d location) of the all the objects (players and ball).
To do this, we replace the discrete predictions with the corresponding real-valued coordinates by using weighted averaging.
We then compute normalized (cid:96) 2 distance between the ground truth locations and the predicted locations of all objects at each step.
The second task is predicting future states.
Since there are many possible futures, using (cid:96) 2 loss does not make sense, either for training or evaluation (c.f., (Mathieu et al., 2016)).
Instead we evaluate the negative log-likelihood on the discretized predictions, which can better capture the multi-modal nature of the problem.
(This is similar to the perplexity measure used to evaluate language models, except we apply it to each object separately and sum the results.)
For each task, we consider the following models: Visual only : standalone visual encoder without any recurrent neural network as backbone.
For future prediction, we follow Felsen et al.
(2017) and directly predict future locations from last observed visual features.
RNN : standard RNN, the hidden states are globally shared by all agents.
VRNN : standard VRNN.
Indep-RNN : one RNN per agent, each RNN runs independently without any interaction.
Social-RNN : one RNN per agent, agents interact via the pooling mechanism in Social GAN (Gupta et al., 2018).
Graph-RNN : one RNN per agent, with a graph interaction network.
Graph-VRNN : the full model, which adds stochasticity to
Graph-
Our visual encoder is based on ResNet-18 (He et al., 2016), we use the ﬁrst two blocks of ResNet to maintain spatial resolution, and then aggregate the feature map with max pooling.
The encoder is pre-trained on visible players, and then ﬁne-tuned for each baseline.
We ﬁnd this step important to stabilize training.
For the soccer data, we down-sample the video to 4 FPS, and treat 4 frames ( second) as one step.
We consider 10 steps in total, 6 observed, 4 unobserved.
We set the size of GRU hidden states to 128 for all baselines.
The state decoder is a 2-layer MLP.
For basketball data, we set every 5 frames as one step, and consider 10 steps as well.
The size of GRU hidden states is set to 128.
The parameters of the VRNN backbone and state decoders are shared across all agents, while each agent has its own visual encoder as “identiﬁer”.
For all experiments, we use the standard momentum optimizer.
The models are trained on 6 V100 GPUs with synchronous training with batch size ofper GPU, we train the model for 80K steps on soccer and 40K steps on basketball.
We use a linear learning rate warmup schedule for the ﬁrst 1K steps, followed by a cosine learning rate schedule.
[Table]
The hyper parameters, such as the base learning rate and the KL divergence weight β , are tuned on a hold-out validation set.
[Table] [Table]
4. Q UANTITATIVE RESULTS
Basketball.
The right half of Table 1 shows the average normalized (cid:96) 2 distance between the true and predicted location of all the agents for the basketball data.
(Error bars are not shown, but variation across trials is quite small.)
We see that the Graph-VRNN error generally goes down over time, as the system integrates evidence from multiple frames.
(Recall that one of the agents becomes “invisible” every 10 frames or 2 steps, so just looking at the current frame is insufﬁcient [Table]
The visual-only baseline has more or less constant error, which is expected since it does not utilize information from earlier observations.
All other methods outperform the visual-only baseline.
We can see that stochasticity helps (VRNN better than RNN, Graph-VRNN better than Graph-RNN), and
Graph-RNN is better than vanilla RNN.
As most of the players are visible in the basketball videos, we also report performance for visible and hidden (occluded) agents in Table 2.
As expected, the (cid:96)distances for visible agents are very low, since the localization task is trivial for this data.
When the agents are hidden, Graph-VRNN signiﬁcantly outperforms other baselines.
Again, we observe that stochasticity helps.
We also ﬁnd that graph network is a better interaction model than social pooling, both outperform Indep-RNN.
Soccer.
The left half of Table 1 shows the average normalized (cid:96) 2 distance between the true and predicted location of all the agents for the soccer data as a function of time.
The results are qualitatively similar to the basketball case, although in this setting the (cid:96) 2 distances are higher since the vision problem is much harder.
However, the gains from adding stochasticity to the Graph-RNN are smaller in this setting.
We believe the reason for this is that the dynamics of the agents in soccer world is much more predictable than in the basketball case, because our simulator is not very realistic.
(This will become more apparent in Section 4.4.)
Forecasting.
In this section, we assess the ability of the models to predict the future.
In particular, the input is the last 6 steps, and we forecast the next 4 steps.
Since the future is multimodal, we do not use (cid:96) 2 error, but instead we compute the log-likelihood of the discretized ground truth locations, normalized by the performance of random guessing § , which is a uniform prior over all locations.
The prior baseline trains the decoder based on constant zero input vectors, thus reﬂecting the prior marginal distribution over locations.
The visual only baseline predicts all 4 future steps based on last observed visual feature, as is in Felsen et al.
The results are shown in Table 3.
Not surprisingly, we see that modeling the interaction between the agents helps, and adding stochasticity to the latent dynamics also helps.
4. Q UALITATIVE RESULTS
This section shows some qualitative visualizations of the belief state of the system (i.e., the marginal distributions p ( s kt +∆ | v 1: t ) ) over time.
(Note that this is different from the beliefs over the internal states, p ( z t +∆ | v 1: t ) , which are uninterpretable.)
We consider the tracking case, where ∆ = 0 , and the forecasting case, where ∆ > 0 .
We include visualizations of sampled trajectories in the Appendix.
Basketball.
In Figure 4, we visualize the belief state for a single agent in the basketball dataset as a function of time.
More precisely, we visualize the output of the decoder,
p kt =
ϕ
dec ( v t , h t − 1 , z kt ) , as a heatmap, where z kt ∼ ϕ prior ( h t − 1 ) are different stochastic samples of the latent state drawn from the dynamical prior.
When we roll forward in time, we sample a speciﬁc value s kt ∼ p kt , and use this to update each h kt .
When we forecast the future, we set v t = 0 .
During tracking, when visual evidence is available, we see that the entropy of the belief state reduces over time, and is consistent across different draws.
However, when forecasting, we see increased entropy, reﬂecting uncertainty in the dynamics.
Furthermore, the trajectory of the belief state differs along different possible future scenarios (values of z kt ).
Figure 6 shows the examples trajectories for ﬁve attacking players generated by Graph-VRNN and baseline methods.
In Figure 6a, we can see that the samples generated by
Graph-VRNN are consistent with ground truth for the ﬁrst 5 to 6 steps (observed), and are diverse for the next 4 steps (not observed).
Figure 6b shows the trajectories generated for the same example by
Indep-RNN and vanilla RNN, which are cluttered or shaky.
Soccer.
Figure 5 (left) shows the belief states for the soccer domain for three different kinds of agents: a regular player (top row), the goal keeper (middle row), and the ball (bottom row).
For the player, we see that the initial belief state reﬂects the 4-4-2 pattern over possible initial locations of the players.
In frame 3, enough visual evidence has been accumulated to localize the player to one of two possible locations.
For the future prediction, we draw a single stochastic sample (to break symmetry), and visualize the resulting belief states.
We see that the model predicts the player will start moving horizontally (as depicted by the heatmap diffusing), since he is already at the top edge of the ﬁeld.
The goal keeper beliefs are always deterministic, since in the simulated game, the movement of the goalie is below the quantization threshold.
The ball is tracked reliably (since the camera is programmed to track the ball), but future forecasts start to diffuse.
Figure 5 (right) shows similar results for the basketball domain.
We see that the dynamics of the players and ball are much more complex than in our soccer simulator.
C ONCLUSIONS AND FUTURE WORK
We have presented a method that learns to integrate temporal information with partially observed visual evidence, based on graph-structured VRNNs, and shown that it outperforms various baselines on two simple datasets.
In the future, we would like to consider more challenging datasets, such as real sports videos.
We would also like to reduce the dependence on labeled data, perhaps by using some form of self-supervised learning.
Acknowledgements.
We thank Raymond Yeh and Bo Chang for helpful discussions.
A. S AMPLED SOCCER TRAJECTORIES
Figure A1 shows the sampled trajectories for Soccer World.
Unlike basketball, Soccer World has a moving camera with limited ﬁeld of view.
We observe that the trajectories for the ﬁrst few observed steps are quite shaky since only a few players have been observed.
We thus show the trajectories from t = 5 .
From ﬁg. A1, we can see that the trajectories generated by RNN are very shaky.
GraphVRNN generates much better trajectories, but some of the players are assigned incorrect identities, or incorrect locations.
We conjecture that this issue can be mitigated by providing longer visual inputs to the model, such that most of the players could be observed at some point of the videos.
A. S AMPLED BASKETBALL TRAJECTORIES
Figure A2 and Figure A3 provide more visualizations for different sampled trajectories by GraphVRNN, and comparison with baseline methods.
The observations are consistent as in Figure 6.