# AXIOM: Learning to Play Games in Minutes with Expanding Object-Centric Models

arXiv: https://arxiv.org/abs/2505.24784 · HTML: https://mlsanity.com/paper/2505.24784/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 25 of its 53 references that are in the map.
- First posted on arXiv 30 May 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 25 of 25)

- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [Multi-Object Representation Learning with Iterative Variational Inference](https://mlsanity.com/paper/1903.00450/) (2019)
- [Deep Reinforcement Learning: An Overview](https://mlsanity.com/paper/1701.07274/) (2017)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Expectation Propagation for approximate Bayesian inference](https://mlsanity.com/paper/1301.2294/) (2013)
- [Bigger, Better, Faster: Human-level Atari with human-level efficiency](https://mlsanity.com/paper/2305.19452/) (2023)
- [Provable Compositional Generalization for Object-Centric Learning](https://mlsanity.com/paper/2310.05327/) (2023)
- [EfficientZero V2: Mastering Discrete and Continuous Control with Limited Data](https://mlsanity.com/paper/2403.00564/) (2024)
- [Variational Inference MPC for Bayesian Model-based Reinforcement Learning](https://mlsanity.com/paper/1907.04202/) (2019)
- [Recurrent switching linear dynamical systems](https://mlsanity.com/paper/1610.08466/) (2016)
- [Dependent Multinomial Models Made Easy: Stick Breaking with the P'olya-Gamma Augmentation](https://mlsanity.com/paper/1506.05843/) (2015)
- [Bayesian model reduction](https://mlsanity.com/paper/1805.07092/) (2018)
- [Slot State Space Models](https://mlsanity.com/paper/2406.12272/) (2024)
- [From pixels to planning: scale-free active inference](https://mlsanity.com/paper/2407.20292/) (2024)
- [Object-Centric Learning with Slot Mixture Module](https://mlsanity.com/paper/2311.04640/) (2023)
- [SPARTAN: A Sparse Transformer World Model Attending to What Matters](https://mlsanity.com/paper/2411.06890/) (2024)
- [Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learning](https://mlsanity.com/paper/2501.16443/) (2025)
- [Supervised structure learning](https://mlsanity.com/paper/2311.10300/) (2023)
- [Slot Structured World Models](https://mlsanity.com/paper/2402.03326/) (2024)
- [Dynamic Markov Blanket Detection for Macroscopic Physics Discovery](https://mlsanity.com/paper/2502.21217/) (2025)
- [Gradient-free variational learning with conditional mixture networks](https://mlsanity.com/paper/2408.16429/) (2024)
- [Structure learning with Temporal Gaussian Mixture for model-based Reinforcement Learning](https://mlsanity.com/paper/2411.11511/) (2024)

## Built on by (most cited first, 3 of 3)

- [Graph-Based Exploration for ARC-AGI-3 Interactive Reasoning Tasks](https://mlsanity.com/paper/2512.24156/) (2025)
- [Natural Building Blocks for Structured World Models: Theory, Evidence, and Scaling](https://mlsanity.com/paper/2511.02091/) (2025)
- [Soft Geometric Inductive Bias for Object Centric Dynamics](https://mlsanity.com/paper/2512.15493/) (2025)

## Abstract (arXiv)

Current deep reinforcement learning (DRL) approaches achieve state-of-the-art performance in various domains, but struggle with data efficiency compared to human learning, which leverages core priors about objects and their interactions. Active inference offers a principled framework for integrating sensory information with prior knowledge to learn a world model and quantify the uncertainty of its own beliefs and predictions. However, active inference models are usually crafted for a single task with bespoke knowledge, so they lack the domain flexibility typical of DRL approaches. To bridge this gap, we propose a novel architecture that integrates a minimal yet expressive set of core priors about object-centric dynamics and interactions to accelerate learning in low-data regimes. The resulting approach, which we call AXIOM, combines the usual data efficiency and interpretability of Bayesian approaches with the across-task generalization usually associated with DRL. AXIOM represents scenes as compositions of objects, whose dynamics are modeled as piecewise linear trajectories that capture sparse object-object interactions. The structure of the generative model is expanded online by growing and learning mixture models from single events and periodically refined through Bayesian model reduction to induce generalization. AXIOM masters various games within only 10,000 interaction steps, with both a small number of parameters compared to DRL, and without the computational expense of gradient-based optimization.