SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Relax, it doesn’t matter how you get there: A new self-supervised approach for multi-timescale behavior analysis
Mehdi Azabou* Georgia Tech
Michael J. Mendelson  
Georgia Tech
Nauman Ahad  
Georgia Tech
Maks Sorokin  
Georgia Tech
Shantanu Thakoor DeepMind
Carolina Urzay  
Georgia Tech
Eva L. Dyer* Georgia Tech
Unconstrained and natural behavior consists of dynamics that are complex and unpredictable, especially when trying to predict what will happen multiple steps into the future. While some success has been found in building representations of animal behavior under constrained or simplified task-based conditions, many of these models cannot be applied to free and naturalistic settings where behavior becomes increasingly hard to model. In this work, we develop a multi-task representation learning model for animal behavior that combines two novel components: (i) an action-prediction objective that aims to predict the distribution of actions over future timesteps, and (ii) a multi-scale architecture that builds separate latent spaces to accommodate short- and long-term dynamics. After demonstrating the ability of the method to build representations of both local and global dynamics in robots in varying environments and terrains, we apply our method to the MABe 2022 Multi-Agent Behavior challenge, where our model ranks first overall on both mice and fly benchmarks. In all of these cases, we show that our model can build representations that capture the many different factors that drive behavior and solve a wide range of downstream tasks.
Behavior is shaped by various factors operating across different timescales. Immediate motivations can drive moment-to-moment interactions, while long-term experiences or even the time of day can influence behavior on broader scales. Analyzing these dynamics, particularly in complex and naturalistic contexts (34; 28), has now become a critical component in many modern studies in neuroscience (28), cognitive science, and in social behavior and decision making (38; 52; 56; 5). Additionally, monitoring and tracking systems now allow for modeling of multi-agent interactions (27; 11; 16) and social behaviors (52; 1), providing valuable insights into dynamics across many individuals.
In order to learn latent factors that may influence behavioral patterns, a promising solution is to build models of behavior in a unsupervised manner (25; 68). Unsupervised models are of particular interest in this domain as it becomes hard to identify complex behaviors which can be composed of many "syllables" of movement (65), and are thus hard and tedious to annotate. Recent work in this direction build such representations using generative modeling and reconstruction-based objectives,
*Contact: {mazabou,evadyer}@gatech.edu. Project page and code: https://multiscale-behavior.github.io/
Figure 1: Overview of our approach. Bootstrap Across Multiple Scales (BAMS) uses two temporal convolutional networks with two latent spaces, each with their own receptive field sizes. The model is trained on a novel learning objective that consists in predicting future action distributions instead of future action sequences. In this figure, we use sample data from the MABe dataset, only a subset of the channels are shown.
typically by performing open loop (58; 13; 10) or closed loop (13) prediction of observations or actions multiple timesteps into the future.
However, when using a reconstruction or prediction objective to analyze behavior, future actions become hard to predict and models can start to become myopic, focusing only on short-range interactions in the data (65). To circumvent this overly local learning of dynamics, there have been a number of efforts to build models of long-term behavioral style (41; 6), where instance-level learning methods are used to extract a single representation for an entire sequence. However, these models then lose their ability to provide time-varying representations that capture the dynamic nature of different behaviors. It is still an outstanding challenge to build representations that can capture both short-term behavioral dynamics along with longer-term trends and global structure.
In this work, we develop a new self-supervised approach for learning multiscale representations of behavior. Our method consists of two core innovations: (i) a novel action-prediction approach that aims to predict the distribution of actions over future timesteps, without modeling exactly when each action is taken, and (ii) a novel multi-scale architecture that builds separate latent spaces to accommodate short- and long-term dynamics. We combine both of these innovations and show that our approach can capture both long-term and short-term attributes of behavior and work flexibly to solve a variety of different downstream tasks.
To test our approach, we utilize behavioral datasets that contain multiple tasks that vary in complexity and contain distinct multi-timescale dynamics. Using NVIDIA's Isaac Gym (37), we generate a synthetic dataset of the multi-limb kinematics from quadruped robots by varying the robot's morphological properties and the environment's terrain type and difficulty. Using this robot behavior data, we demonstrate that our method can effectively build dynamical models of behavior that accurately elicit both the robot and environment properties, without any explicit training signal encouraging the learning of this information.
Having established that our approach can successfully predict the complex behavior of an artificial creature, we apply it to two multi-agent behavior benchmarks (61) and challenges with multiple tasks that vary in their frame-level (local) vs. sequence-level (global) labels and properties. On the mouse triplet benchmark, we rank first overall on the leaderboard  ${}^{1}$  (averaged across 13 tasks),first on all of the 4 global tasks,and are in the top-3 on all the 9 frame-level local subtasks. In one of the global tasks (decoding the strain of the mouse), we achieve impressive performance over the other methods, with a  ${10}\%$  gap over the next best performing method. On the fruit fly groups benchmark, we also rank first overall  ${}^{2}$  (averaged across 50 tasks),and outperform other methods on both average frame-level and sequence-level subtasks. Our results demonstrate that our approach can provide representations that can be used to decode meaningful information from behavior that spans many timescales (longer approach interactions, grooming, etc.) as well as global attributes like the time of day or the strain of the mouse.
1 aicrowd.com/challenges/multi-agent-behavior-challenge-2022/problems/mabe-2022-mouse-triplets/leaderboards?post_challenge=true
2 aicrowd.com/challenges/multi-agent-behavior-challenge-2022/problems/mabe-2022-fruit-fly-groups/leaderboards?post_challenge=true
The contributions of this work include:
- A self-supervised framework that learns representations in two separate long-term and short-term embedding spaces. Bootstrapping is performed within each timescale, using a latent predictive loss across positive views only, and hence can process much longer sequences than contrastive methods that require negative views.
- A novel prediction task for behavioral analysis and cloning called HoA (histogram of actions), which aims to predict the future distribution of keypoints instead of the precise ordering of future states. We use an efficient implementation of the 2-Wasserstein divergence as a measure of distributional fit between the true movement distribution and the predicted movements.
- New state-of-the-art performance: When applied to the Multi-agent (MABe) Benchmark, our model is ranked #1 overall and achieves top scores on all of the sequence-level (global) tasks. In aggregate, we achieve the top rank on the leaderboard.
- A procedurally generated dataset of walking quadruped robots with ground truth annotations and multiple frame- and sequence-level tasks. We believe that this dataset can provide a robust benchmark for future work in this direction.
Our approach addresses two critical challenges of modeling naturalistic behavior (Figure 1). In Section 2.2, we introduce a novel distributional-relaxation of the reconstruction-based learning objective. In Section 2.3, we introduce an architecture and self-supervised learning objectives that support the learning of behavior at different levels of temporal granularity.
We assume a fixed dataset of  $D$  trajectories, each comprised of a sequence of observations  $\mathbf{x}_t$  and/or actions  $\mathbf{y}_t$ . Where actions are not explicitly provided, in many cases we can infer actions based on the difference between consecutive observations. Our goal is to learn, for each timestep, behavioral representations  $\mathbf{z}_t$  that capture both global-information such as the strain of the mouse or the time of day, as well as temporally-localised representations such as the activity each mouse is engaged in at a given point in time. As obtaining labeled datasets for realistically-useful scales of agent population and diversity of behavior is impractical, we aim to learn representations in an unsupervised manner.
2.2 Histogram of Actions (HoA): A novel objective for predicting future
Modeling behavior dynamics can be done by training a model to predict future actions. This reconstruction-based objective becomes challenging when behavior is complex and non-stereotyped. Let us consider the example of a mouse scanning the room, rotating its head from one side to the other. It is possible to extrapolate the trajectory of the head over a few milliseconds, but prediction quickly becomes impossible, not because this particular behavior is complex but because any temporal misalignment in the prediction leads to increasing errors.
We propose to predict the distribution of future actions rather than their sequence. The motivation behind this distributional-relaxation of the reconstruction objective lies in blurring the exact temporal unfolding of the actions while preserving their behavioral fingerprint.
Predicting histograms of future actions. Let  $\mathbf{y}_t \in \mathbb{R}^N$  be the action vector at time  $t$ . Each feature in the action vector can, for example, represent the linear velocity of a joint or the angular velocity of the head. Given observations  $[\mathbf{x}_0, \dots, \mathbf{x}_t]$  of the behavior at timesteps 0 through  $t$ , the objective is to predict the distribution of future actions over the next  $L$  timesteps. For each  $i$ -th element of the action vector, we compute a one-dimensional normalized histogram of the values it takes between timesteps  $t + 1$  and  $t + L$ . We pre-partition the space of action values into  $K$  equally spaced bins, resulting in a  $K$ -dimensional histogram that we denotes as  $\mathbf{h}_{t,i}$ , for all keypoints  $1 \leq i \leq N$ .
We introduce a predictor  $g$  that, given the extracted representation  $\mathbf{z}_t$ , predicts all feature-wise histograms of future actions. The predictor is a multi-layer perceptron (MLP) with an output space in  $\mathbb{R}^{N\times K}$ . The output is split into  $N$  vectors, which are normalized using the softmax operator. We obtain  $[\hat{\mathbf{h}}_{t,1},\dots ,\hat{\mathbf{h}}_{t,n}]$ , each estimating the histogram of the  $i$ -th action feature following timestep  $t$ .
$\mathbf{EMD}^2$  loss for histograms. To measure the loss between the predicted and target histograms, we use the Discrete Wasserstein distance, also known as the Earth Mover's Distance (EMD). This distance is obtained by solving an optimal transport problem that consists in moving mass from one distribution to the other while incurring the lowest transport cost. In our case, the cost of moving mass from one bin to another is equal to the number of steps between the two bins.
Because our histogram has equally-sized bins, the EMD is equivalent to the Mallows distance which has a closed-form solution (31; 20). In particular we use  $\mathrm{EMD}^2$  which has been shown to be easier to optimize and converge faster (53). The loss is defined as follows:
$$
\mathcal {D} _ {\mathrm {E M D} ^ {2}} \left(\mathbf {h} _ {t, i}, \hat {\mathbf {h}} _ {t, i}\right) = \sum_ {k = 1} ^ {K} \left(\mathrm {C D F} _ {k} \left(\mathbf {h} _ {t, i}\right) - \mathrm {C D F} _ {k} \left(\widehat {\mathbf {h}} _ {t, i}\right)\right) ^ {2}, \tag {1}
$$
where  $\mathrm{CDF}_k(\mathbf{h})$  is the  $k$ -th element of the cumulative distribution function of  $\mathbf{h}$ .
The total loss is obtained by summing over all features of the action vector, which leaves us with the following loss at time  $t$ :
$$
\mathcal {L} _ {t} = \sum_ {i = 1} ^ {N} \mathcal {D} _ {\mathrm {E M D} ^ {2}} \left(\mathbf {h} _ {t, i}, \widehat {\mathbf {h}} _ {t, i}\right). \tag {2}
$$
2.3 Multi-timescale bootstrapping in a temporally-diverse architecture
In order to form richer and multi-scale representations of behavior, we use a self-supervised learning objective. We introduce a new approach using latent predictive losses to build representations across different scales while preserving the granularity in each. We achieve this by explicitly separating the short-term and long-term representations, then bootstrapping within each representation space. This approach enables us to learn from otherwise incompatible representation learning objectives (55; 62).
2.3.1 Two latent spaces are better than one
Our goal is to capture and separate short-term and long-term dynamics in two different spaces. We use the Temporal Convolutional Network (TCN) (4) as our building block. The TCN produces a representation at time  $t$  that only depends on the past observations (46).
We design an architecture that separates the different timescales by using two TCN encoders: A short-term encoder,  $f_{s}$ , that captures short-term dynamics and targets momentary behaviors such as drinking, running or chasing; A long-term encoder,  $f_{l}$ , that captures long-term dynamics and targets longstanding factors that modulate behavior (strain of mouse, time of day). Architecturally, the difference between the two is that we increase the number of layers and use larger dilation rates (8) for the long-term encoder, thus effectively covering a larger receptive field (more history) in the input sequence. All feature embeddings extracted by the TCNs are concatenated, to produce the final embedding,  $\mathbf{z}_{t} = \mathrm{concat}[\mathbf{z}_{t}^{s},\mathbf{z}_{t}^{l}]$ .
2.3.2 Bootstrapping Across Multiple Scales
We draw inspiration from recent work (17; 50; 18) that uses latent bootstrapping to learn a latent space where "positive" views are mapped close to each other. Unlike other contrastive methods (69), bootstrapping does not require negative examples.
In the context of temporal representation learning, a common assumption is that points that are nearby in time are positive views of each other and can be constrained to lie nearby in the latent space (69; 3). In our case, we can bootstrap and find positive views at both the short-term and also at a more long-term scale, as illustrated in Figure 2.
Bootstrapping short-term representations. We randomly select samples, both future or past, that are within a small window  $\Delta$  of the current timestep  $t$ . In other words,  $\delta \in [-\Delta, \Delta]$ . Bootstrapping
involves using a shallow network to predict the representation of one view from the other (17). We
Figure 2: Visualization of the short-term and long-term windows used to build multi-scale similarity. Positives are selected within a window. The window is small for short-term embeddings and can be as large as the entire sequence for long-term embeddings.
use a predictor  $q_{s}$  that takes in the short-term embedding  $\mathbf{z}_t^s$  and learns to regress  $\mathbf{z}_{t + \delta}^{s}$  using the loss:
$$
\mathcal {L} _ {r, s h o r t} = \left\| \frac {q _ {s} \left(\mathbf {z} _ {t} ^ {s}\right)}{\left\| q _ {s} \left(\mathbf {z} _ {t} ^ {s}\right) \right\| _ {2}} - \operatorname {s g} \left[ \frac {\mathbf {z} _ {t + \delta} ^ {s}}{\left\| \mathbf {z} _ {t + \delta} ^ {s} \right\| _ {2}} \right] \right\| _ {2} ^ {2}, \tag {3}
$$
where  $\mathrm{sg}[\cdot ]$  denotes the stop gradient operator. Unlike (17), we do not use an exponential moving average of the model, but simply increase the learning rate of the predictor as in (50).
Bootstrapping long-term representations. For long-term behavior embeddings which should be stable at the level of a sequence, we sample any other time point in the same sequence, i.e.  $t' \in [0, T]$ . We use a similar setup for the long-term behavior embedding, where predictor  $q_l$  is trained over longer time periods or in the limit, over the entire sequence.
$$
\mathcal {L} _ {r, l o n g} = \left\| \frac {q _ {l} \left(\mathbf {z} _ {t} ^ {l}\right)}{\| q _ {l} \left(\mathbf {z} _ {t} ^ {l}\right) \| _ {2}} - \operatorname {s g} \left[ \frac {\mathbf {z} _ {t ^ {\prime}} ^ {l}}{\| \mathbf {z} _ {t ^ {\prime}} ^ {l} \| _ {2}} \right] \right\| _ {2} ^ {2} \tag {4}
$$
Finally, we optimize the proposed multi-task architecture with a combined loss:
$$
\mathcal {L} = \mathcal {L} _ {t} + \alpha \left(\mathcal {L} _ {r, s h o r t} + \mathcal {L} _ {r, l o n g}\right) \tag {5}
$$
where  $\alpha$  is a scalar that is used to weigh the contribution of the short- and long-term contrastive losses. In practice, we find that we simply need to choose  $\alpha$  that re-scales the contrastive losses to the same order of magnitude as the HoA prediction loss.
3.1 Simulated Quadrupeds Experiment
3.1.1 A synthetic dataset of simulated legged robots
To test our model's ability to separate behavioral factors that vary in complexity and contain distinct multi-timescale dynamics, we introduce a new dataset generated from a heterogeneous population of quadrupeds traversing different terrains. Simulation enables access to information that is generally inaccessible or hard to acquire in a real-world setting and provides ground-truth information about the agent and the world state.
Agents. We use advanced robotic systems (33) that imitates 4-legged creatures capable of various locomotion skills. These robots are trained to walk on challenging terrains using reinforcement learning (51). We use two robots that differ by their morphology, ANYmal B and ANYmal C. To create heterogeneity in the population, we randomize the body mass of the robot as well as the target traversal velocity. We track a set of 24 proprioceptive features including linear and angular velocities of the robots' joints.
Procedurally generated environments. Using NVIDIA's Isaac Gym (33) simulation environment, we procedurally generate maps composed of multiple segments of different terrains types (Figure 3). We consider five different terrains including flat surfaces, pits, hills, ascending and descending stairs. We also vary the roughness and slope of the terrain to control the difficulty of terrain traversal.
Figure 3: Quadrupeds walking on procedurally generated map. (A) Illustration showing the two robots (ANYmal B and ANYmal C) walking on the procedurally generated map, with segments of different terrain types. As the robots traverse different terrains, the velocity of their joints is tracked and visualized. (B) Long-term embeddings learned by BAMS. On the left, we overlay the labels for the type of robot and on the right we overlay the velocity of the robot; we observe a clear organization of the latents in terms of their velocity and robot type.
Model	Sequence-level Tasks	Frame-level Tasks
Robot Type* (↑)	Linear velocity (↓)	Terrain type* (↑)	Terrain slope (↓)	Terrain difficulty (↓)
PCA	99.72	0.069	08.83	0.037	0.790
TCN	99.93	0.102	33.03	0.037	0.080
BAMS	99.96	0.038	39.89	0.033	0.078
L short-term	100.0	0.094	34.86	0.036	0.079
L long-term	99.88	0.020	32.39	0.036	0.078
Table 1: Linear readouts of robot behavior. For each task, we report the linear decoding performance on sequence-level and frame-level tasks. Tasks marked with * are classification tasks for which the F1-score is reported, for the remaining tasks, we report the mean-squared error.
Experimental setup. We collect 5182 trajectories of robots walking through terrains. We record for 3 minutes at a frequency of  $50\mathrm{Hz}$ . For evaluation, we split the dataset into train and test sets (80/20 split) and use multi-task probes that correspond to different long-term and short-term behavioral factors. More details can be found in Appendix A.
3.1.2 Results
Results in Table 1 suggest that our model performs well on these diverse prediction tasks. A major advantage of our method is the separation of the short-term and long-term dynamics, which enables us to more clearly identify the multi-scale factors. While some of the tasks are represented best in the mixed model, we find that the linear velocity is more decodable in the long-term embedding, while terrain type is more decodable in the short-term embedding. We further analyze the formed representations by visualizing the embeddings in the different spaces. In Figure 3-B, we visualize the long-term embedding space and find that our model is able to capture the main factors of variance in the dataset, corresponding to the robot type and the velocity at which the robots are moving. This suggests that in the absence of labels, the learned embedding can provide valuable insights into how the recorded population is distributed without the need for annotations. In Appendix A, we visualize the extracted embeddings of a single sequence over time. We find that the long-term embedding is more stable and smooth, while the short term embeddings reveal different blocks of behavior that change more frequently.
3.2 Experiments on Mouse Triplet Dataset
3.2.1 Experimental Setup and Tasks
Dataset description. The mouse triplet dataset (61) is part of the Multi-Agent Behavior Challenge (MABe 2022), hosted at CVPR 2022. This large-scale dataset was introduced to address the lack of standardized benchmarks for representation learning of animal behavior. It consists of a set of trajectories from three mice interacting in an open-field arena. A total of 5336 one-minute clips, recorded from a top-view camera, were curated and processed to track twelve anatomically defined keypoints on each mouse, as shown in Figure 4.
Figure 4: Multi-Agent Behavior (MABe) - Mouse Triplets Challenge. (A) Keypoint tracking approaches are used to extract keypoints from many positions on the mouse body in a video. (B) Methods are evaluated across 13 different tasks.
As part of this benchmark, a set of 13 common behavior analysis tasks were identified, and are used to evaluate the performance of representation learning methods. Over the course of these sequences, the mice might exhibit individual and social behaviors. Some might unfold at the frame level, like chasing or being chased, others at the sequence level, like light cycles affecting the behavior of the mice or mouse strains that inherently differentiate behavior.
Integrating features across multiple animals. We process the trajectory data to extract 36 features characterizing each mouse individually, including head orientation, body velocity and joint angles. We construct the short-term TCN and long-term TCN encoders to have representation of size 32 each, and receptive fields approximated to be 60 and 1200 frames respectively. We compute the histogram of actions over 1 second, i.e.  $L = 30$  and use  $K = 32$  bins. We build representations for each mouse independently, which means that at time  $t$ , and for each frame  $t$ , we produce embeddings  $\mathbf{z}_{t,1}$ ,  $\mathbf{z}_{t,2}$