Modeling Fly Social Behavior with Agent-Centric Pose Forecasting

The gist: This work introduces an agent-centric framework for training autoregressive models of animal behavior by constraining deep networks to receive sensory inputs and produce motor outputs matching those of the animal being modeled, which is applicable to single animals and groups.

Agent-Centric Framework

The core idea is to build generative models where each animal is treated as a separate agent that inputs sensory information and outputs movements in the body’s frame The framework constrains the deep network’s inputs and outputs to match the animal’s so that the network and the animal solve the same problem This approach is designed to make fitting a hypothesis about how an animal generates its behavior rather than a fit to its trajectories.

Representational Machinery and Library

The framework requires managing many parallel representations of the same data, along with ML-specific transformations like discretization To handle this, the authors release a general-purpose library called AnimalPoseForecasting, which handles the representational machinery this approach requires: composable, invertible operations between world-frame, egocentric, and ML-friendly representations of pose and sensing. The library is organized around composable and invertible sequences of operations that connect parallel representations of the data.

Sensory Input and Output Representations

The input sensory representation is hand-crafted to limit the agent to what the animal can perceive, rather than giving it omniscient access to the full state of the world For interacting flies, visual information about other flies is represented as a depth map, and tactile information is summarized by the height of the domed ceiling at each keypoint and pairwise distances between body parts of the focal fly and others. The output representation choices are explored along three axes: the form of the prediction (continuous, discrete, or hybrid), the temporal frame of reference (absolute pose versus change in pose), and the parameterization of pose (raw 2D keypoints, egocentric keypoints, or a character-animation-inspired decomposition).

Autoregressive Rollout in Dynamic Environments

Generation is performed by applying next-step prediction iteratively In a dynamic, multi-agent environment, the next step’s input depends on the state of the world, which depends on the actions of all agents in the previous step This necessitates several operations beyond the network forward pass at each rollout step, including every agent’s predicted movement must be inverted back to a change in world-coordinate pose. The library's design ensures that the same operations that map raw pose into the model’s egocentric input space at training time are inverted, at rollout time, to map predictions back into world coordinates.

Evaluation Criteria and Key Findings

The paper evaluates models using a suite of complementary criteria, including Qualitative characterization, Quantitatively comparing the distributions of selected behavior features between real and simulated data, and training discriminators to distinguish real from simulated data. Experiments on courting male flies show that the resulting models capture properties at multiple scales, including fine-grained motor patterns such as the coordinated leg gait of walking, recognizable behavioral motifs such as chasing and wing extension. The analysis reveals that agent-centric representations enabled by our library is most important for accurate behavior modeling Furthermore, the work demonstrates that discretization is important to avoid the pitfall of forecasting no movement. The framework's adaptability is shown by applying it to a new domain like RatInABox.

Model Comparison Insights

When comparing variants, the REFERENCE and NO HANDCRAFTED POSE model variants appear to produce realistic fly behavior The worst-performing variant is the KEYPOINTS model, which is the only variant that does not use an egocentric representation of sensory input and motion output. The study concludes that the agent-centric representations enabled by our library is most important for accurate behavior modeling. The ultimate goal is to use such models as instruments for scientific discovery, interrogating their internal structure to reveal computations the animal itself implements.

1. Agent-centric modeling for animal behavior allows fitting a hypothesis about how the animal generates its behavior rather than a fit to its trajectories. This enables AI systems to model complex social behaviors by constraining the deep network to receive the same sensory input and produce the same kind of motor output as the animal it models.

2. The AnimalPoseForecasting library provides composable, invertible operations between world-frame, egocentric, and ML-friendly representations, which allows researchers to systematically compare design choices like agent-centric versus world-centric representations without rewriting rollout machinery.

3. The framework facilitates mechanistic interpretability by allowing the interrogation of trained models to extract structure in the phenomenon the model has learned to capture, potentially revealing patterns that scientists have not previously articulated.

Understanding animal behavior at an algorithmic level -- what animals attend to, how they form internal representations, and how this maps to action -- remains a central challenge in neuroethology. We explore how well large generative models capture properties of fly social behavior. We combine transformers with agent-centric representations. Models input egocentric sensory observations and output egocentric movements, mirroring biological constraints. Social behavior emerges from agents independently sensing and responding to one another. We demonstrate that this simple model captures many properties of fly behavior, outperforming the world-centric representation and achieving SOTA on MABe 2022-fly. We systematically evaluate properties of the representation important for capturing the behavior distribution through a series of lesioning experiments. These comparisons were enabled by a new, general-purpose library for animal-centric pose forecasting.