Efficiently Scanning and Resampling Spatio-Temporal Tasks with Irregular Observations

## Abstract

Abstract Various works have aimed at combining the inference efficiency of recurrent models and training parallelism of multi-head attention for sequence modeling. However, most of these works focus on tasks with fixed-dimension observation spaces, such as individual tokens in language modeling or pixels in image completion. To handle an observation space of varying size, we propose a novel algorithm that alternates between cross-attention between a 2D latent state and observation, and a discounted cumulative sum over the sequence dimension to efficiently accumulate historical information. We find this resampling cycle is critical for performance. To evaluate efficient sequence modeling in this domain, we introduce two multi-agent intention tasks: simulated agents chasing bouncing particles and micromanagement analysis in professional StarCraft II games. Our algorithm achieves comparable accuracy with a lower parameter count, faster training and inference compared to existing methods.

## 1 Introduction

Spatio-temporal modelling tasks with complex unstructured or semi-structured state and observation spaces can be identified in a variety of domains. Designing deep learning algorithms to excel at these tasks often requires deliberate handling of both the accumulation of knowledge from historical observations, and the summarization of current observations, which can be computationally expensive. This is a particular concern in domains such as motion prediction and behaviour modeling, which have real-time requirements and often need to be performed using relatively low compute resources available on an edge devices. Recurrency is a popular paradigm in deep learning as it naturally maps onto these problems, which are sequential and causal in nature. While iterative processing of data with recursion can be efficiently performed in $\mathcal{O}(1)$ with respect to the input sequence, transformer and convolution methods introduce compute and memory complexity correlations with the input sequence length of $\mathcal{O}(N^{2})$ and $\mathcal{O}(N)$ respectively. However, transformer and convolution algorithms are more parallelizable at training time, efficiently utilizing hardware and invoking backward propagation paths that aren’t correlated with the sequence length. This leads to a trade-off between compute efficiency at training or inference time.

Proposals to address this trade-off can be grouped into various categories, each of which aim to take advantage of inference efficiency of recurrence and parallelization at training time. One avenue aims to address the $\mathcal{O}(N^{2})$ space and compute complexity that multi-head attention invokes over the sequence by introducing variations with linearized or amortized attention calculations Sun et al. (); Peng et al. (); Katharopoulos et al. (). State space models (SSMs) Dao & Gu (); Gu et al. () are presented as a compelling alternative to transformers with demonstrated efficacy in a variety of long range dependency tasks and language modeling. SSMs can be formulated as a convolution for training time paralellism or as a recurrent model for efficient inference over long sequences (Gu et al. ()). Fu et al. () propose to directly learn an unrolled state space model with an efficient long-convolution framework. In each case, these methods are evaluated on sequences with a fixed dimensional observation space $\mathcal{O}\in\mathbb{R}^{d}$, such as tokenized text, image pixels or audio spectrogram data. Focusing on these sequences may leave a blind-spot in tasks with an irregular observation spaces $\mathcal{O}\in\mathbb{R}^{d(t)}$ such those in as multi-agent interactions, where the number of agents and links between these may vary over time.

In this work, we investigate a range of encoding and state space modelling approaches for these settings, and propose a novel and efficient algorithm that utilizes a 2D latent state and alternates between input sampling, and accumulating historical information as a weighted sum (inclusive scan). This weighted sum can be performed efficiently in parallel on a GPU with an inclusive-scan Merrill & Garland () during training, and incrementally during inference. We show that the resampling cycle is more effective than a continued self-attention block, or not alternating between accumulation and processing. This method natively supports a two dimensional latent state. We evaluate our method against a transformer encoder (Vaswani et al., ), several RNN models (Chung et al., ; Hochreiter, ), and Mamba2 (Dao & Gu, ).

To test the efficacy of these training and inference efficient algorithms on tasks with more complex observation spaces, we use two multi-agent interaction benchmarks. The first is a “gymnasium” style simulation that involves agents chasing randomly assigned particles(^1^11Environment and path-planning algorithm derived from [here](https://www.doc.ic.ac.uk/%C2%A0ajd/Robotics/RoboticsResources/planningmultirobot.py).). The second benchmark is based on StarCraft II (SC2), a real-time strategy video game, where we extract instances where the players are in combat. Each of these tasks involves a multidimensional time-varying observation space, where the model must analyse the motion of agents to identify their intended goal.
In summary our contributions are as follows:

- •
The introduction of two multi-agent interaction challenges to better evaluate sequence modeling algorithms with irregular observation spaces.
- •
A novel algorithm to efficiently address sequence modelling tasks with irregular observation spaces. We find that this algorithm achieves comparable accuracy to alternatives with a lower parameter count and improved throughput in training and inference regimes.
- •
The finding that resampling the observation space conditioned on accumulated sequence data is particularly beneficial in sequence modelling tasks.

## 2 Background

Efficient and parallelisable sequence modeling is of great interest to the research community, as attaining efficient utilisation of parallel computation at training time while performing inference efficiently are desirable attributes of sequence modeling algorithms. Previously, a trade-off had to be considered between using a recurrent neural network and a window-based algorithm such as a transformer or convolutional neural network. Recurrent methods utilize a hidden state to retain information between steps in the sequence, enabling $\mathcal{O}(1)$ compute and memory complexity with respect to the length of the sequence when performing incremental inference. However, incremental processing results in under-utilization of compute hardware and long back-propagation chains dependent on the sequence length, slowing the training process. Alternatively, transformer or convolution based methods are highly parallelisable, taking advantage of modern hardware, and the back-propagation chain is related to the depth of the model, rather than the length of the sequence. This results in greater training throughput. This comes at the cost of lack of incremental processing for inference. At inference time, a window of data is processed causing runtime compute and memory complexity of $\mathcal{O}(L)$ and $\mathcal{O}(L^{2})$ for convolution and transformer methods respectively.

Spatio-temporal tasks can be modeled as a state space $\mathcal{S}$ with an observation space $\mathcal{O}$ that evolves over time according to dynamics model $\mathcal{D}$. In some domains, $\mathcal{O}$ is a variably-sized set, ${\mathbb{O}}_{t}$, that changes between environment instances, or over the duration of the sequence. This creates challenges in both concisely and effectively compressing a sequence of variably sized observations into some latent representation of the state of the system $\mathcal{S}_{t}^{\prime}:={\mathbb{O}}_{0\dots T}$. This can be exacerbated if the temporal duration of this task is indefinite as $T\in\mathbb{N}^{+}$.

As a particularly pertinent example of this class of problem, StarCraft II (SC2) is a real-time strategy game where players build an economy and military in order to defeat the opposing team. At a macro scale, SC2 follows a typical rock-paper-scissors approach where there are counter-strategies that can be employed against a given player. However, micromanagement in SC2 games is also a contributing factor to the performance of a player, with novices performing a few dozen actions per minute, whereas the typical professional completes hundreds of actions per minute. These actions are typically taken on variably sized sets of units, buildings, targets or objectives. Analysing the micromanagement of a player can potentially give insights into strategy used.

## 3 Environments

We describe the environments that are used to evaluate the efficacy of the algorithms first, as this will provide context for the design motivations in Section 4. Each consist of two sets of agents which interact with one another with some objective. We use these environments to test whether models are capable of capturing relevant information from observations of a complex environment, in order to infer some properties from it. The predominant task is a target assignment problem, where an agent from one set is (possibly) assigned to target another. This task is similar to other intention recognition tasks which are pertinent in domains such as human robot interaction (Felip et al., ; Ahmad et al., ), that aim to infer and predict the objective of agents though observation. Each environment has an additional task to demonstrate generalisation of the architecture to multiple problems problems, beyond the target assignment problem.

Figure: (a) Chasing-Targets gymnasium environment. Robots (marked with a trajectory trail) are randomly allocated targets (blue dots) to chase. When targets are reached (Green dots), robots are randomly assigned new targets.
Refer to caption: extracted/5919246/images/gym-example.png

### 3.1 Chasing Targets

The chasing targets environment involves an arbitrary number of two-wheeled robots that chase an arbitrary number of particles bouncing around the environment, while trying to avoid collisions with one another. Each robot is randomly assigned a target particle at the beginning of the simulation and initialized in a random stationary pose. When the robot reaches the particle, it is randomly assigned to a new particle. Robots are controlled with a simple cost function to select the best control inputs that minimizes the robot’s distance to the target’s projected position and a penalty term if a collision is forecast to occur with another robot. Invalid control inputs, such as exceeding a speed limit are disallowed. Particles are initialised at random positions in the arena with a random initial velocity sampled from a normal distribution and clamped to a maximum range. A visualisation of the environment is depicted in Figure 1(a).

The main modelling objective in this task is to determine the intended target (particle) for each of the robots in the scene. We consider two variations of this task. In the first, the observation space consists of the current pose of the robots and particles. A random number of robots and particles are sampled from a uniform distribution for each training and validation sample. For the second task variation, the observation space only includes the robots, and the model must predict where particles are located in the scene. This is challenging as not only the position of the particles are unknown, but the cardinality must also be estimated. This task is formulated as an occupancy problem, the model needs to estimate the likelihood that a position in the environment is occupied by a particle.

### 3.2 StarCraft II

The StarCraft II data used is sampled from tournaments hosted between 2019 and 2023. To focus on battle sequences where players are micro-managing their units in combat, we target parts of the game when damage dealt or received by a player exceeds a threshold. Non-overlapping sequences of a fixed size are created at these instances. We note that observation data is irregularly sampled, hence the time duration of each sequence will vary. This domain has some subtleties compared to Chasing-Targets (Section 3.1) as units are not always assigned to an enemy unit. They can either be idle, or assigned to a position to move to. Hence, we introduce null option for the assignment problem. Furthermore, we introduce a secondary task, where we are required to estimate if a unit has been given a target position command and the location of that target position. The assignment and the position estimation problem are jointly learned and performed by each model.

The observation space for StarCraft II includes a terrain height-map and the position and properties (health, damage, etc.) of player and enemy units. The observation space is restricted to a region-of-interest (ROI) of a fixed size, see Appendix A.1 for detail on how the ROI is calculated. An example of an extracted ROI with unit data and their assignments is depicted in Figure 1(b). Units outside of the ROI are truncated, resulting in a time-varying set of units in the environment as they enter, exit or die in combat. The observability rules of the game apply from a player’s point-of-view. Enemy units will be hidden in the fog of war or when an obfuscation ability is used, for example Zerg players can “burrow” units into the ground.

Figure: (a) BERT Encoder

Both environments above are good examples of sequential modelling and prediction problems involving irregular, multidimensional time varying observation spaces. Below, we describe a sequential encoding and dynamics modelling strategy suitable for these tasks.

## 4 Method

We evaluate several methods to encode spatio-temporal features into a series of latent states. The decoder for each of the tasks is fixed in order to isolate the contribution of the encoding methodology.

### 4.1 Spatio-Temporal Encoding

This section outlines how the latent representation of the scene is constructed, using the StarCraft task above as an example. The scene observation is tokenized identically for each of these encoders. First, the $(x,y,\theta)$ pose of the agents are sinusoidally encoded. Units from StarCraft II have additional information included such as health and a learned embedding associated with the type of unit. From this variably sized set of tokens, ${\mathbb{O}}^{\prime}=\{{\bm{o}}\}_{n_{o}},{\bm{o}}\in\mathbb{R}^{d_{o}}$, sampled from the observation, we must summarize a fixed set ${\mathbb{L}}=\{{\bm{l}}\}_{n_{l}},{\bm{l}}\in\mathbb{R}^{d_{l}}$ suitable for latent dynamics modelling.

We use two methods for summarizing ${\mathbb{O}}^{\prime}$ into ${\mathbb{L}}^{\prime}$, as illustrated in Figure 2, each of which involves Multi-head Attention between ${\mathbb{O}}^{\prime}$ and ${\mathbb{L}}$. The first is a BERT-style (Devlin, ) transformer encoder. Latents ${\mathbb{L}}$ are concatenated to ${\mathbb{O}}^{\prime}$ , and act as the “[CLS]” tokens of the BERT encoder. The encoded “[CLS]” tokens, ${\mathbb{L}}^{\prime}$, are used as a fixed-size representation of the variably sized observation data. The second method uses a block of cross-attention layers (X-Attn), ${\mathbb{L}}$ to query key-value pairs generated from ${\mathbb{O}}^{\prime}$ to transfer relevant information. We denote these encoders as $\mathbf{Enc}$.

Without loss of generality, consider the case where we have two observation sources. In the tasks above, these correspond to the two teams of agents. Since there are two distinct observation sources, ${\mathbb{O}}_{p}$ and ${\mathbb{O}}_{e}$, we test a variety of methods to determine an effective method for combining both into ${\mathbb{L}}^{\prime}$. The “Fused” method, Figure 3(a), adds a learned embedding per source and then concatenates the sources together for encoding. This method has fewer parameters than the alternative methods as there is only one $\mathbf{Enc}$. This also enables flexibility in gathering the optimal amount of information from each type of source. However, this comes at the cost of a larger attention matrix inside $\mathbf{Enc}$ with $\mathcal{O}((N_{p}+N_{e})^{2})$. The “Piece-wise” method, Figure 3(b), encodes ${\mathbb{O}}_{p}$ and ${\mathbb{O}}_{e}$ separately with half of the latent state ${\mathbb{L}}$ used for each observation source. This method has the benefit of a smaller attention matrix $\mathcal{O}(N_{p}^{2}+N_{e}^{2})$, but enforces an equal weighting of information between the two sources, which is potentially sub-optimal. “Sequental” processes ${\mathbb{O}}_{p}$ then ${\mathbb{O}}_{e}$, as depicted in Figure 3(c). This renders a smaller attention matrix, but reduces parallelism and increases model depth.

Figure: (a) Fused Encoding

#### 4.1.1 Scan Encoder

The key design objective of this spatio-temporal encoder is to efficiently aggregate historical information with a parallel algorithm, and to use that accumulated knowledge to sample the current observation. Furthermore, we use a set of tokens to represent our hidden state, rather than a single vector that is common to most algorithms. The motivation behind this is that attention mechanisms can be then utilized effectively with this hidden state, to either inject or extract information from this state. To achieve this, we use a weighted sum of an initially encoded input data from previous steps in the sequence. The accumulation can be performed efficiently in parallel over the sequence dimension with an inclusive scan operation Merrill & Garland (). The proposed inclusive-scan algorithm uses the cross-attention encoder (Fig. 2(b)) with the sequential algorithm (Fig. 3(c)) to create a recursive query driven sequential modelling approach. Inclusive-scan is performed on the updated variables to accumulate temporal information. This allows variables to accumulate context, so that when an observation is re-queried, it is conditioned on an accumulated history. This process is depicted in Figure 4. At inference time, we can calculate the next accumulated encoding with a scaled copy of the previous time-step ${\mathbb{L}}_{xt}^{\prime}=\frac{1}{\gamma}{\mathbb{L}}_{x(t-1)}^{\prime}+{
\mathbb{L}}_{xt}$, resulting in a computational complexity of $\mathcal{O}(1)$ with respect to the length of the sequence. We show in Section 5 that cycling between cross-attention and inclusive-scan is superior to a block of cross-attention layers with an inclusive-scan at the end or sampling the input once, and cycling between inclusive-scan and self-attention. We include the number of cycles as nomenclature for this scan encoder in Section 5, for example Scan $4\times$ is a scan encoder with four cycles.

Figure: Figure 4: The inclusive-scan encoder alternates between sampling the observation based on some latent variable ${\mathbb{L}}_{x}$ and accumulating a weighted sum ${\mathbb{L}}^{\prime}$. A learned set of parameters ${\mathbb{L}}_{0}$ is used as the query in the first block. $\gamma\geq 1$

An important subtlety is that the inclusive-scan is weighted such that the historical contribution decays over time with $\gamma\geq 1$. This ensures that the accumulation does not diverge in magnitude. We show in Section 5 that the scan with $\gamma=2$ outperforms $\gamma=1$ (a simple cumulative sum). A PyTorch extension was written to efficiently perform the forward and backward method of discounted inclusive-scan both on CPU and on GPU with CUDA(^2^22[https://github.com/5had3z/torch-discounted-cumsum-nd](https://github.com/5had3z/torch-discounted-cumsum-nd)).

#### 4.1.2 Baseline Encoders

The spatio-temporal transformer (S.T.) groups tokens from the same time-step to create observation chunks at each time-step. A learned position embedding of the absolute time-step is added to the tokens. We use a transformer encoder with causal masking to process the entire sequence. We evaluate both the BERT and cross-attention encoders on the agent assignment challenge.

We evaluate each of the recurrent neural networks included in PyTorch: RNN, GRU and LSTM. We use BERT to encode the observation data and a set of recurrent modules to process each latent encoding. We find that learned initialization performs better than zero initialization with an LSTM and therefore also use learned initialization for other module structures (Appendix A.3).

We use Mamba2 as a point of comparison against modern SSMs. The cross-attention encoder is used to encode observation data, and like the recurrent neural networks above, we use a set of Mamba2 modules to process each latent encoding. We use 64 for the state dimension, 4 for the convolution dimension and an expansion factor of 2 as the parameters for Mamba2.

The height-map is additional context provided in the SC2 task. ResNet-18 is used as the feature extractor, however we replace the final $1\times 1$ adapative-average-pool and fully-connected layers with a $4\times 4$ adaptive-average-pool layer to create a grid of features. Sinusoidal position embeddings are added to the feature grid and then flattened, creating the set of contextual tokens. These tokens are appended to the extracted spatio-temporal features before being passed to the task decoder.

### 4.2 Task Decoding

In this section, we detail how the decoding strategy used for each benchmark task. The foundation of each task decoder is a cross-attention operation to gather information from the latent state.

Goal Assignment aims to find the correlation between agents and targets. To encode the agent, we first project the agent feature vector with a linear layer, then use cross-attention to transfer information from the latent state that is useful to associate agents and targets. The targets are also projected with a linear layer. For the StarCraft II challenge, a learned null token is appended to the set of projected targets. Finally, the correlation between agents and targets is determined using a dot-product. Soft-max is applied along the target dimension to calculate a final weighting between the agents and targets. We note that this is a one-to-many problem, many agents can have the same target, but each agent only has one target (including the null assignment).

We utilize the negative log likelihood (NLL) as a categorical loss between the predicted assignment and the ground truth assignment. However, there are some complications due to padding introduced to make regular tensor shapes from the variable agents and targets in the scene. Invalid agents can simply be omitted from the loss function with a mask. To handle invalid targets, we group agents with the same number of targets in the scene and truncate the tensor to this number. This allows us to apply NLL loss while avoiding gradients from invalid targets and agents. In the StarCraft II task, there is a significant imbalance between agents with a target and those without. A weighting factor of $0.05$ is applied to emphasise non-null assignment losses.

A Movement Target Assignment task completes the motion action space that units in StarCraft II can take. Here, the objective is to estimate the coordinates on the map where the unit is moving to. This plays an important part of the game as manoeuvring units into advantageous positions such as the high ground or choke-points is an important strategy in game-play. We evaluate several methods of estimating the target position with both regression and categorical techniques including: global cartesian, relative cartesian and relative polar. Furthermore, a logit is emitted to estimate the likelihood that the unit is following a position command and position estimate is valid.

An L2 loss is used for regression-based position estimation and HL-Gauss (Imani & White, ) is used for categorical-based estimation. We include additional logic in the HL-Gauss loss to correct angle wrapping for polar coordinate estimation, further detail is provided in Appendix A.2. Binary cross-entropy loss is used for the valid position command logit. Position commands are under-represented in the dataset compared to target commands. To increase the amount of position data, we create faux position commands from the coordinates of target units. This does not effect the valid position command logit training.

Hidden Target Estimation in Chasing-Targets is formulated as an occupancy prediction problem. An array of coordinates is sinusoidally encoded to query the spatio-temporal encoding the likelihood of hidden target occupation. This is performed with cross-attention between the positional query and the spatio-temporal encoding. A linear layer transforms this result to a single logit. Focal loss Lin et al. () is used for occupancy prediction to address class imbalance between unoccupied and occupied pixels. This rendering method can create an image of any resolution as position queries can be continuous, however we use a fixed grid for simplicity.

## 5 Results

We used PyTorch to train our models. Unless otherwise specified, experiments used a batch size of 64, the AdamW optimizer with a learning rate of $1e^{-4}$ and a polynomial schedule with power 0.9 and gradient clipping of 0.1(^3^33Link to source and model configurations will be provided later).

### 5.1 Chasing Targets

Each instance of the environment is randomly generated each time, sampling the number of robots and targets from a uniform distribution, and randomly placing them on the field. The field is a $4\times 4$m grid and the maximum velocity of the agents is 0.5m/s. The first 10 iterations of the simulation are skipped to remove the domain gap between the behaviour of the robots after random initialization, and the steady state chasing and switching targets. The duration of the simulation is 41 iterations.

#### 5.1.1 Target Assignment

We train the target assignment challenge for $\approx 47$k iterations and use a latent state $\mathcal{L}\in\mathbb{R}^{8\times 128}$. The number of robots in each simulation is sampled from $\mathcal{U}(8,15)$ and the number of targets $\mathcal{U}(3,6)$. Figure 5 shows the complex trade-space between training throughput, memory usage, and the parameter count of the model. The LSTM is the fastest model to train with a step time of 17.7ms, has the second lowest parameter count and third lowest memory consumption. The most accurate model, the spatio-temporal (S.T) transformer with X-Attn encoder has the second shortest step time, moderate memory consumption and parameter count. The Scan $\times 4$ encoder has the least number of parameters by a large margin, however suffers from lower accuracy in this task compared to S.T and Mamba2. Mamba2 performs similarly to S.T with the lowest memory consumption during training.

Figure: Figure 5: Average Top 1 assignment accuracy over the simulation sequence with training cost. The top left corner is ideal in each scenario. The Scan encoder linearly scales in performance and cost from ^1 to ^2 by introducing an additional self-attention layer in each recursion block.

We notice a trend of assignment accuracy decay after a peak earlier in the sequence in Figure 6. This is more pronounced in the better performing S.T and Mamba2 models, converging to approximately the same performance as the Scan $4\times$ and LSTM by the end of the sequence.

Figure: Figure 6: Assignment accuracy comparison over a sequence.

To understand the contributions of each of the design decisions in the Scan encoder, we remove individual components from a baseline model. If a cumulative sum is performed without weighting, accuracy decays towards the end of the sequence, and in general underperforms the baseline $\gamma=2$. Reducing the size of the latent state to $\mathbb{R}^{1\times 128}$ has significant performance consequences as not enough historical information is retained over the sequence. This performs almost as poorly as not performing a scan at all, the worst performing experiment conducted.

Figure: Figure 7: Contribution ablation for a Scan $2\times$ encoder over a sequence.

#### 5.1.2 Hidden Target Estimation

While the S.T with cross attention performs well in target assignment accuracy, it struggles with hidden target estimation. We train the hidden target estimation challenge for $\approx 94$k iterations with a batch size of 32 and latent state $\mathcal{L}\in\mathbb{R}^{16\times 128}$. The number of robots in each simulation is sampled from $\mathcal{U}(4,12)$ and the number of targets from $\mathcal{U}(2,5)$. Figure 8 shows that the Scan and Mamba encoders perform equivalently, whereas the transformer encoder fails on this task.

Figure: Figure 8: Model prediction AUC over the hidden target estimation sequence.

Figure 9 illustrates a prediction sequence from this task. Although numerically Scan and Mamba produce low AUC scores in absolute terms, they are often correctly able to infer target locations in the scene.

Figure: (a) 0
Refer to caption: extracted/5919246/images/uncover-seq/0.png

### 5.2 StarCraft II

We train each model for $\approx 205$k iterations with a battle sequence length of 30 and use latent state $\mathcal{L}\in\mathbb{R}^{16\times 256}$. StarCraft II is more challenging than Chasing-Targets, evident by the significantly lower top-1 assignment score for non-null assignments in Table 1. There are substantially more units on the field, which also vary in cardinality as they enter, exit and die in combat. Furthermore, in dense combat scenarios it is potentially more challenging for the model to precisely assign targets. We find that relative cartesian categorical representation of position performs the best and is used for the following experiments. An ablation of position representation can be found in Appendix A.2.

There is change in ranking of training cost from Chasing-Targets to SC2, as the duration of the sequence is shorter by 25%, and the task results in a difference in the interplay between efficient sampling of the observation, and temporal accumulation. As the number of units on the screen is significantly greater than Chasing-Targets, quickly compressing the observation space with the shallowest model is more important. The baseline encoders use two block repetitions for encoding the observation, compared to Scan encoder which has encoding blocks correlated with the number of repetitions, four in the case of Scan $4\times$. A larger latent state is used as the complexity of the scene and task has increased. All encoders increase in parameter count with the token dimension. Adding more tokens to the latent state increases in number of LSTM and Mamba2 modules.

**Table 1: Average performance over the SC2 battle sequence and training cost.**
| Encoder | Top1 | Top1 | Top5 | MSE | F1 | # M | Mem | Train |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| +null | Param. | GB | it/sec |  |  |  |  |  |
| LSTM | 0.271 | 0.899 | 0.656 | 0.200 | 0.793 | 22.1 | 17.5 | 9.00 |
| S.T w/ X-Attn | 0.264 | 0.906 | 0.656 | 0.191 | 0.809 | 17.9 | 9.61 | 15.1 |
| Mamba2 | 0.188 | 0.899 | 0.622 | 0.191 | 0.811 | 29.1 | 9.11 | 14.5 |
| Scan $4\times$ | 0.198 | 0.916 | 0.626 | 0.173 | 0.841 | 21.6 | 14.0 | 9.65 |

## 6 Discussion and Limitations

Our benchmark sequences are performed over relatively short timespans. Accuracy over the Chasing-Targets challenge saturates early into this cycle, however we note accuracy decay in some encoders (Fig. 6). SC2 accuracy continues to increase until the 30 step cut-off. Micromanagement tasks may not elicit long range dependencies to the extent of short term ones, with players performing hundreds of actions per minute. There may be additional value in the development and benchmarking of the models above on tasks with longer term temporal dependencies and irregular and complex observation spaces.

## 7 Related work

State Space Models (SSMs) are ubiquitous in sequential modelling tasks. In control and dynamics modelling, they have been applied to learn locally linear latent dynamics models conditioned on exogenous inputs (Watter et al., ; Banijamali et al., ; Jaques et al., ; Levine et al., ; Karl et al., ; Han et al., ). More recently, these models have become popular with endogenous inputs. Here, SSMs parameterize a sequential process with endogenous input ${\bm{u}}(t)$ to model the output process ${\bm{y}}(t)$ with hidden state ${\bm{x}}(t)$,

$$ $\displaystyle\dot{{\bm{x}}}(t)={\bm{A}}{\bm{x}}(t)+{\bm{B}}{\bm{u}}(t)$ (1) $\displaystyle{\bm{y}}(t)={\bm{C}}{\bm{x}}(t)+{\bm{D}}{\bm{u}}(t).$ (2) $$

Parameters $({\bm{A}},{\bm{B}},\dots)$ can be approximated in a discrete form $(\bar{{\bm{A}}},\bar{{\bm{B}}},\dots)$, appropriate for a sequence-to-sequence mapping, such that they can be learned with gradient descent. However, these matrices must be deliberately structured for effectiveness, usually guided by HiPPO theory Gu et al. (). The key desirability of SSMs is re-parameterization as a convolution to enable parallel computation for training,

$$ $\displaystyle{\bm{y}}_{t}=\sum_{k=0}^{t}\bar{{\bm{C}}}\bar{{\bm{A}}}^{k}\bar{{ \bm{B}}}{\bm{u}}_{t}+\bar{{\bm{D}}}{\bm{u}}_{t}$ (3) $\displaystyle{\bm{y}}=\bar{{\bm{K}}}*{\bm{u}}+\bar{{\bm{D}}}{\bm{u}},$ (4) $$

where ${\bm{K}}=(\bar{{\bm{C}}}\bar{{\bm{A}}}^{i}\bar{{\bm{B}}})_{i=0\dots L-1}$ for sequence length $L$. Rather than learning the SSM matrices, Fu et al. () propose learning $\bar{{\bm{K}}}$ directly with various regularization techniques and an efficient long convolution implementation to enable feasibility. Gu & Dao () introduce the notion of SSM matrices also being time dependent, similar to an LSTM, adding mechanisms for selectively propagating or forgetting state information dependent on the current token. This is also accompanied by a specialised implementation for performance. In follow up work, Dao & Gu () decomposes the SSM into diagonal and low-rank blocks for computational efficiency.

Sub-quadratic Transformer Attention has been proposed with various techniques. Katharopoulos et al. () re-parameterise multi-head self-attention with a non-negative feature map (i.e. $\text{elu}(x)+1$) applied to $\mathbf{K},\mathbf{Q}$ separately to attain linear time and memory complexity, and enable incremental processing in scenarios with causal masking. RWKV (Peng et al., ) utilizes learned position biases, similar to (Zhai et al., ), rather than queries from that position. Sun et al. () follows a similar mechanism to our proposed method where a hidden state $\mathcal{S}$ is propagated with a decay factor $\gamma$ and updated with the current input, $\mathcal{S}_{n}=\gamma\mathcal{S}_{n-1}+\mathcal{X}_{n}$. However, our cyclic scanning and updating conditions the input on the accumulated history, before adding to the time-decaying latent state, as illustrated in Figure 4.

Irregular Spatio-Temporal Tasks can be found in many domains. Multi-object tracking is a popular task in computer vision, which aims to associate detected objects between frames. Recursive processing is popular in end-to-end learning of the multi-object tracking problem (Meinhardt et al., ; Zhang et al., ; Zhu et al., ). This methodology is suitable for training short-term tracking algorithms, appropriate for many benchmarks (Dendorfer, ; Dave et al., ). However, over a longer horizon, for example in long-term Re-identification tasks, recursive methods scale poorly, hence further innovation is required. Longer horizon methods (Qin et al., ) are typically are not trained end-to-end like for this reason. Graph neural networks (GNNs) (Battaglia et al., ) are popular architectures for variable dimensional observations or irregularly structured data due to their permutation invariance. These models generally use zero-padding to handle irregular sized observations and rely on fixed graph structures for global context, although some approaches attempt to model the evolution of the graph over time (Pareja et al., ; Manessi et al., ). Unfortunately, these approaches lack the scalability of state space models. They are also difficult to frame in an incremental manner for inference, GNNs typically process the full sequence. For example, Gao et al. () construct a graph representation of a traffic scene with agent tracklets, lanes and crosswalks to produce a future trajectory estimation for the agents. GNNs such as these lack incremental inference, at each time-step the spatio-temporal graph is modified with the new observation, and re-processed in its entirety.

## 8 Conclusion

This work introduces a sequential modelling approach for irregular observation sources of varying dimensionality and cardinality. We investigate the use of various encoders to pack these observations into a fixed dimensional latent, and corresponding sequence modelling approaches. We propose a novel algorithm that alternates between cross attention and a weighted inclusive scan to accumulate context over time and show that this performs on par with existing sequence modelling approaches on two new challenging benchmarks, with comparably higher training and inference speeds and a lower parameter count. A key benefit of the proposed approach is that it can naturally be tweaked to allow a broad range of performance/compute tradeoffs (see Appendix A.4 for more extensive experiments in this vein) as required for a given application.

## 9 Reproducibility

Model source code and experiment configuration files will be made public later. Chasing targets can be installed with pip chasing-targets-gym, source code is at [https://github.com/5had3z/chasing-targets-gym](https://github.com/5had3z/chasing-targets-gym). StarCraft II uses the sc2-serializer framework, [https://github.com/5had3z/sc2-serializer](https://github.com/5had3z/sc2-serializer), the tournaments dataset was used.

## Appendix A Appendix

### A.1 StarCraft II ROI Calculation

The objective of the ROI is to contain the main action happening on the screen, namely the units in combat. To achieve this, we gather the coordinates of all the units with target or motion commands over the sequence. K-Means clustering is performed on the positions and sorted by the number of members in the clusters. If the limits centroids of the K-means clusters can fit within the desired ROI size, then the mean of the centroids is used for the center of the ROI. If this condition is not fulfilled, the smallest cluster is removed. This is done iteratively until the condition is satisfied. This calculation can be performed live in the dataloader during training, or calculated once and saved a key associated with the sample to be re-indexed during training.

### A.2 StarCraft II Unit Positioning

HL-Gauss (Imani & White, ) loss is used for categorical position encoding. Instead of a one-hot encoding label, a gaussian kernel is used as the categorical target with $\mu$ equal to the position and $\sigma$, which we set to $0.2$. The appropriate number of bins used for the prediction is correlated to the $\sigma$ chosen (Farebrother et al., ), we use 20. We decode the categorical output of the model by integrating over the product of the soft-max output and the value each bin represents. We perform an ablation study with the large Scan $6\times$ encoder, shown in Table 2, to determine the best performing position learning schema. Based on these results, we use a relative cartesian coordinate system for unit position estimation, encoded with a categorical distribution.

**Table 2: Position representation ablation performed with Scan $6\times$ encoder. Here Cart. is cartesian, Rel. is relative position from the unit, Glbl. is global position in the ROI and Cat. is categorical position encoding.**
| Rep. | Frame | Encoding | MSE |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| 0 | 5 | 15 | 25 |  |  |  |
| Cart. | Rel. | Scalar | 0.300 | 0.258 | 0.253 | 0.251 |
| Cart. | Glbl. | Scalar | 0.299 | 0.259 | 0.254 | 0.241 |
| Polar | Rel. | Cat. | 0.258 | 0.221 | 0.216 | 0.216 |
| Cart. | Rel. | Cat. | 0.248 | 0.211 | 0.207 | 0.206 |
| Cart. | Glbl. | Cat. | 0.254 | 0.220 | 0.218 | 0.219 |

In this section we discuss how we handle wrapping $\theta$ when using a polar coordinate system with a categorical representation. When creating the categorical target, we double the template range to $[-4\pi,4\pi]$. Once the kernel is projected onto the template, we add the reflection of the additional range to the opposite side in the normal range, i.e. we reflect and add the bins from $[-4\pi,-2\pi]$ to $[0,2\pi]$, and repeat for the the opposite tail. To handle angle wrapping during decoding, we find the argmax category and rotate the categories so the argmax is centered. With this new categorical range centered on the argmax, the values each bin represents is recalculated to the new range. The standard decoding algorithm is used to yield the final result.

### A.3 Recurrent Model Comparison

We evaluated all the available recurrent neural network modules available in PyTorch, and briefly compare a learned or zero initialization for the LSTM. The standard RNN performs poorly whilst the LSTM and GRU are similar. Learned or zero initialization for the LSTM also has minimal impact on the performance of the model on the Chasing-Targets benchmark, with the learned initialisation having a slightly greater peak, but zeros climbing to higher accuracy faster. Each of these models use the X-Attn encoder with piece-wise encoding algorithm.

Figure: Figure 10: Comparison between recurrent models and hidden state initialization where Z=zeros and L=learned.

### A.4 Chasing Targets Detailed Results

Table 3 provides a larger set of results including various encoder (BERT, X-Attn) and encoding algorithm (Fused, Piece-wise) variants for the S.T, Recurrent and Mamba2 encoders. We also include slight variations of the Scan encoder with removed self-attention layers to trade accuracy and efficiency (indicated by ^1) or replace cross-attention with the observation in the repeat layers, with self attention blocks (indicated by ^2).

**Table 3: Algorithm performance, training and evaluation cost of tested algorithms on Chasing-Targets. ^1 A self-attention layer is skipped in the encoder block. ^2 The input is only sampled once.**
| Encoder | Top1 Accuracy | # M | Mem. | Train | Infer. |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 0 | 5 | 10 | 40 | Param. | GB | it/sec | it/sec |  |
| RNN-L,Fused,BERT | 0.359 | 0.529 | 0.592 | 0.633 | 0.76 | 1.82 | 60.7 | 89.4 |
| GRU-L,Fuse,BERT | 0.372 | 0.652 | 0.706 | 0.711 | 1.29 | 1.88 | 58.8 | 86.6 |
| LSTM-L,Fuse,BERT | 0.375 | 0.646 | 0.703 | 0.703 | 1.56 | 1.89 | 56.4 | 81.4 |
| S.T,PW,BERT | 0.377 | 0.688 | 0.739 | 0.719 | 2.10 | 2.81 | 35.3 | 55.2 |
| S.T,PW,X-Attn | 0.374 | 0.692 | 0.745 | 0.716 | 2.10 | 1.86 | 42.7 | 52.1 |
| Mamba2,PW,X-Attn | 0.381 | 0.673 | 0.724 | 0.703 | 2.79 | 1.04 | 30.9 | 31.6 |
| Mamba2,Fuse,BERT | 0.386 | 0.667 | 0.719 | 0.713 | 2.79 | 1.95 | 35.4 | 61.0 |
| Scan $1\times$ | 0.358 | 0.514 | 0.557 | 0.600 | 0.41 | 0.74 | 77.7 | 92.5 |
| Scan $2\times$ | 0.334 | 0.658 | 0.688 | 0.712 | 0.65 | 1.13 | 49.7 | 54.8 |
| Scan $4\times^{1}$ | 0.399 | 0.629 | 0.668 | 0.677 | 1.31 | 1.82 | 30.4 | 33.3 |
| Scan $4\times^{2}$ | 0.368 | 0.661 | 0.696 | 0.705 | 0.91 | 1.47 | 39.7 | 41.3 |
| Scan $4\times$ | 0.364 | 0.680 | 0.717 | 0.704 | 1.31 | 1.95 | 29.2 | 30.6 |
| Scan $6\times$ | 0.403 | 0.696 | 0.723 | 0.703 | 1.91 | 2.79 | 20.7 | 21.7 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See Weaknesses.
soundness: 2
strengths: The paper investigates the use of various encoders to encode irregular observations into a fixed dimensional latent, and corresponding sequence modelling approaches.
summary: The paper discusses a approach to spatio-temporal modelling for tasks with irregular observation spaces. Experiments are conducted on multi-agent scenario's like robotic target-chasing simulations and StarCraft II gameplay analysis.


Motivation: The paper aruges that current architectures like Transformers and LSTMs works well in scenarios with fixed observation spaces but struggle with tasks where observation spaces vary over time, such as multi-agent interactions. The paper aims to combine the efficiency of recurrent models with the parallelism of attention mechanisms, focusing on varying-dimension observation spaces.

Proposed Method:  The paper introduce an algorithm using a 2D latent state, alternating between cross-attention on observations and a discounted cumulative sum over time, capturing long context.  

Experiments: Two benchmarks are used for evaluation: (a) Chasing Targets: Agents pursuing dynamic targets, evaluating the model's ability to infer intended goals in real-time; (b) StarCraft II: A environment with complex, temporally varying observations, where models infer combat unit assignments and predict unit movement.

Results: The proposed model achieves competitive accuracy with fewer parameters and faster inference compared to standard methods (e.g., transformers). The "resampling cycle," improves sequence modeling accuracy by conditioning observations on accumulated sequence data.

===================

After rebuttal: I've read the review by other reviewers as well as the rebuttal by the authors.

This statement by the authors summarizes the contribution of the paper well.
"Our results primarily set out to explore the trade-offs between efficiency and performance in the irregular observation case. This means results should be interpreted given the compute budget available to a practitioner. If only performance matters, practitioners should select a temporal transformer with sequential cross attention (X-Attn Sequential) for observation encoding. For efficiency and model parsimony, an inclusive scan should be selected with X-Attn Sequential."

It will be very helpful to revise the introduction and abstract accordingly.
weaknesses: - The model's efficiency enables its use in real-time applications with irregular observation space, but studying the effect on longer sequence durations as well as diversifying the tasks (i.e., conduct experiments on more tasks) would further validate its utility in long-term planning tasks.
- There's some work in the literature like Temporal Latent Bottleneck (https://arxiv.org/abs/2205.14794) and Block Recurrent Transformers (https://arxiv.org/abs/2203.07852); which combines strengths of attention and recurrence.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: I’m not very familiar with this topic, so my questions are mainly about the application settings.

1.	The observations of Chasing-target and StarCraft2 are all low-dimensional. How to extend your method to more realistic and complex scenarios, e.g., vision-based navigation tasks?

2.	How can this method be applied to longer-horizon tasks?

3.	I would like to see more ablation studies on the selection of hyperparameter $\gamma$.
soundness: 4
strengths: The proposed method, especially integrating inclusive scan algorithm, is novel and is an effective way to accumulate history.

The experiment results are good. The proposed method improves the performance and reduce the computation overhead.
summary: This paper investigates how to combine the inference efficiency of recurrent models and the training parallelism of multi-head attention transformers, especially for dealing with the varying size observation space. The authors propose a novel algorithm to cycle between cross-attention and inclusive scan to efficiently accumulate historical information. They evaluate their method on two benchmarks, Chasing-Targets gymnasium and StarCraft2.
weaknesses: Please refer to the Questions part.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Section 3.2 Where exactly does the StarCraft data come from? It would be nice to have a detailed description of the data collection and formatting process in the Appendix
2. How different is this method of encoding observations from previous work on StarCraft (regardless of the exact task, the encoding stage should be similar).
3. What is the "learned initialization" in Section 4.1.2
4. Is there some reason why baselines couldn't be matched on some axis? (parameter count is probably easiest)
5. What I understand in Section 3.2 is that the StarCraft task is a 3-class (idle, move to position, move to enemy) classification problem given sequential context, but still this doesn't exactly specify the setup for me, could you clarify what exactly is the task?
6. Section 7 Given the related work (especially Sun et al. 2023) is there some reason why none of these were used as baselines?
7. The limitations seem a bit unfortunate, but also foreseeable, is there some reason that the authors specifically designed their own benchmarks instead of working with something longer-term?
soundness: 3
strengths: - The contribution of the task based on StarCraft can be an interesting domain for future work to study
- The writing is in general fairly clear to follow
summary: The paper proposes a method for sequential prediction tasks that can handle observations of varying sizes at different times. It also simultaneously introduces two intention-prediction tasks in which to test this method. Experiments are performed against a few baselines and components of the method are ablated.
weaknesses: I should start off by pointing out that I'm not precisely acquainted with this literature, so that will affect my judgment.

## Major points
1. One large issue I see with the paper is that I'm not totally certain why these particular tasks are good ones in order to explore modeling arbitrary-sized observation spaces. I think StarCraft as a domain makes sense, but there has been much prior work in StarCraft (which is not cited in the paper) that has used it in a different way, and I don't understand why the decision was made to couple "intention-prediction" tasks with the method. It seems that some domains like sequential decision making in healthcare, or some kind of variable-sized text prediction task would be equally fine as a benchmark domain. This is further complicated by the relatively unclear definition of the exact input and output for each task in Section 3. Given the motivation of these tasks is that it shows up often in interactions, it would seem that a space like Shared Autonomy might provide some more common benchmarks.

2. Another big issue that I see is that the paper does not make a strong enough case for the choice of baselines. I would expect there to be accompanying work in the domain previously with which to compare against, but because 4.1.2 does not cite any prior work, it's unclear to me that these choices are well-founded. Certainly I'm familiar with Mamba2, but the rest is unclear. Also the choice of encoders for the RNN methods feels somewhat unjustified.

3. The last major issues I see is that the results in Figure 5 and Table 1 feel somewhat inconclusive. Because nothing is controlled (either parameter count, memory, or training time) between methods, and there is no clear winner across metrics, it's difficult to see the benefit. When coupling these mixed results with the choice of new bespoke domains, it's even harder to understand the longer-term contribution. I don't think it's the case that a new method needs to be always better than prior work, but the fact that it does not appear to have particular advantages in computation that was claimed as a contribution ("efficiently address sequence modeling" at the end of Section 1), makes it hard to understand the contribution.

## Specific points

1. Section 4.1.2 add a citation to Mamba2
2. Section 4.2 it would be nice to have the problem written out mathematically as I find the current description a bit hard to follow
3. Section 5 is missing comment on hyperparameter tuning, especially for baselines
4. Fig 7 should be referenced in 5.1.1

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1) Could you clarify the contributions of the paper in terms of the models introduced, notably wrt to SSMs. I read the related works, but they did not do a good job connecting to what was done in prior sections.
2) Why is it important that the benchmarks include PvP / PvE games?
3) Did you ablate over decoders when you decided on the final scheme to compare methods? Is it possible that some results are decoder dependent (ie there's a relationship between decoder and encoder results).
4) Could you explain in more detail the motivation of the scan methods? What do you think is going on that makes this method work? What does this have to do with variable observation spaces?
5) What is Lx in 4.1.1? What does the first subscript mean in L_xt?
6) Does the model ever output observation predictions? You mention "sample the current observation" but it's unclear what you mean.
7) What do you mean by "learned initialization" in 4.1.2?
8) What is the consequence of those PvE encoders from Figure 3 on the results?
soundness: 3
strengths: Overall the paper is interesting and approaches a very important problem I think in a comprehensive way, covering some benchmarks that seem well motivated to me as well as some methodologies for dealing with the variable observation problem. The experiments I think are clear enough and the conclusions drawn I think are sound, demonstrating that some methods perform better than others.
summary: The paper tackles the issues of spatio-temporal modeling in the context of variable observation spaces, such as when there are variable objects across episodes / scenes, etc or when the object number varies within an episode (such as when entities enter and leave an arena). The authors introduce a few baseline tasks representative of this problem and cover various attention-based encoding and temporal aggregation schemes, some of which are advertised as being novel (such as some of the encoder methods). They perform various analysis on different methodologies, such as task performance, memory, speed, etc on these tasks which highlight some of the benefits of various methods (such as the scan-based one).
weaknesses: There are places where the writing is good, and there are places where the writing and motivation is not so good. The introduction is very vague and lacks citations on some critical claims in the paper that make it difficult for readers who are not very familiar with the domain to understand well. The benchmarks are introduced well enough, but it's unclear why we need multi-player problems to address the core problem of variable observation spaces. Critically though when the methods come around the organization of the paper suffers, as it becomes extremely difficult to understand the motivation of different design choices, such as the scan operation (e.g., why was the scheme in figure 4 motivated, other than post-hoc performance?). The methods in Figure 3 seem interesting, but I had difficulty finding their consequence in the paper later on.

(Many of these and questions below were addressed sufficiently in the rebuttal and revision)

Meta-review: The paper presents an interesting approach to sequence modeling with irregular observation spaces, exploring novel scan operations and cross-attention techniques for handling variable-dimensional observations. While the work demonstrates potential in multi-agent tasks and introduces innovative modeling strategies, it struggles with comprehensive validation and generalizability. Reviewers noted significant challenges in methodology, including unclear motivation and insufficient differentiation from existing approaches. The benchmarks, primarily based on custom multi-agent tasks, do not definitively prove the method's broader applicability, and performance improvements remain marginal. Despite the authors' efforts to address concerns through additional experiments and clarifications, fundamental limitations persist. The incremental nature of the contribution, combined with presentation challenges and limited experimental rigor, suggests that while the work shows promise, it is not yet sufficiently developed for publication at a top-tier conference.

Decision: Reject