Source-linked AI summary
Relational Neural Expectation Maximization: Unsupervised Discovery of Objects and their Interactions
Sjoerd van Steenkiste, Michael Chang, Klaus Greff, Jürgen Schmidhuber
TL;DR
The paper asks how common-sense physical reasoning can be learned from visual observations without supervised object representations. It introduces R-NEM, which discovers objects with Neural Expectation Maximization and models pairwise interactions relationally; experiments report more accurate dynamics, occlusion handling, and generalization across object counts.
Problem
Existing compositional physical-reasoning methods rely on supervised object representations, while unsupervised pixel-space methods lack compositional object structure.
Method
R-NEM combines unsupervised Neural Expectation Maximization for object discovery with an interaction function that models pairwise effects between object representations.
Results
R-NEM captures physical dynamics more accurately than other methods, generalizes to different object counts, and predicts movement and collisions through complete occlusion.
Takeaways & Limitations
The method can serve as an approximate simulator that maintains object states and predicts interactions in occluded scenes.
Takeaways & Limitations
R-NEM requires injected noise or reduced capacity to encourage useful groupings, which may prevent detection of very small input regularities.
Abstract
from arXiv · showhide
Common-sense physical reasoning is an essential ingredient for any intelligent agent operating in the real-world. For example, it can be used to simulate the environment, or to infer the state of parts of the world that are currently unobserved. In order to match real-world conditions this causal knowledge must be learned without access to supervised data. To address this problem we present a novel method that learns to discover objects and model their physical interactions from raw visual images in a purely \emph{unsupervised} fashion. It incorporates prior knowledge about the compositional nature of human perception to factor interactions between object-pairs and learn efficiently. On videos of bouncing balls we show the superior modelling capabilities of our method compared to other unsupervised neural approaches that do not incorporate such prior knowledge. We demonstrate its ability to handle occlusion and show that it can extrapolate learned knowledge to scenes with different numbers of objects.
1 INTRODUCTION
The paper frames common-sense physical reasoning as essential for intelligent agents, but notes that existing compositional approaches typically require supervised object representations. R-NEM addresses this gap by learning objects and their interactions directly from raw images without supervision.
- Common-sense physical reasoning supports simulation and inference about currently unobserved parts of the world.
- Objects act as compositional primitives that let humans decompose scenes, represent relations, and reason about dynamics and interactions.
- Successful compositional approaches learn general physical dynamics between explicit object pairs, but rely on externally supplied object representations and supervision.
- Unsupervised pixel-space approaches lack compositional object representations, limiting efficient learning of functions over multiple entities and human-like generalization.
- R-NEM learns physical interactions between objects from raw visual images unsupervised and factors interactions between object pairs.Its relational mechanism is designed to support efficient learning and generalization to scenes with varying object counts without retraining.
2 METHOD
R-NEM combines Neural Expectation Maximization for unsupervised object discovery with a relational interaction function. The method uses compositional object representations and pairwise message-like effects to model object dynamics and interactions.
- 2.1 NEURAL EXPECTATION MAXIMIZATION: N-EM groups pixels into objects and learns a separate distributed representation for each object through differentiable iterative clustering.Each representation uses the same feature format and contains information about its corresponding visual object.
- 2.1 NEURAL EXPECTATION MAXIMIZATION: N-EM models images as mixtures of K components, with latent binary assignments indicating which component generated each pixel.
- 2.1 NEURAL EXPECTATION MAXIMIZATION: The E-step computes soft pixel-to-component assignments, while the generalized M-step updates object representations using gradient ascent.
- 2.1 NEURAL EXPECTATION MAXIMIZATION: Differentiable unrolling lets back-propagation train the image model, while the inter-cluster loss restricts each representation to pixels assigned to it.
- 2.2 RELATIONAL NEURAL EXPECTATION MAXIMIZATION: RNN-EM captures individual-object dynamics but cannot capture relations between objects, motivating R-NEM’s relational recurrence.
- 2.2 RELATIONAL NEURAL EXPECTATION MAXIMIZATION: The interaction function’s inductive bias balances interaction modeling with preservation of compositional object representations.Assuming no interactions recovers RNN-EM, whereas an unconstrained function over all representations can forfeit compositionality.
- 2.2 RELATIONAL NEURAL EXPECTATION MAXIMIZATION: R-NEM updates each object using pairwise effects from other objects, with shared embeddings, effect vectors, and attention coefficients selecting relevant interactions.The total interaction effect is a weighted sum of other objects’ effects.
3 RELATED WORK
Related work spans symbolic physics engines and state-to-state prediction, with object-based representations supporting generalization and simulation. R-NEM connects unsupervised visual learning with compositional object representations and relational interaction modeling.
- Symbolic approaches infer parameters of physics engines, restricting them to synthetic environments.
- State-to-state prediction methods use compact object descriptions as inputs or training targets, achieving strong generalization and simulation capabilities.
- Purely unsupervised state-to-state methods using raw visual inputs have not matched these capabilities.
- R-NEM’s interaction function is a message-passing network with neighborhood attention and permutation-equivariant set-function properties.
- N-EM discovers compositional object representations, while TAG, RTagger, and attention-based methods make coherent representations for physical interactions more difficult.
- Related research combines neural-network learnability with symbolic-program compositionality across physics, games, algorithms, visual understanding, and language.
4 EXPERIMENTS
R-NEM is evaluated on bouncing-ball dynamics, occlusion, long-horizon simulation, extrapolation to more objects, and visually complex gameplay. Across these settings, experiments assess prediction accuracy, compositionality, relational modeling, and generalization.
- Bouncing Balls: R-NEM learns compositional object representations from bouncing-ball sequences, grouping pixels belonging to individual balls into distinct components.The experiments use 64×64 binary images over 30 time-steps with four balls initialized using random positions, masses, and velocities.
- Bouncing Balls: An ARI score of 0.8 indicates that roughly 80% of test sequences model each ball with a single component.R-NEM also produces markedly lower losses than the other evaluated methods on four-ball sequences.
- Bouncing Balls: Increasing training components to K = 8 further reduces loss, particularly relational BCE, while improving grouping quality.This result is consistent with the proposed benefit of compositional object representations for modeling physical interactions.
- Extrapolating learned knowledge: R-NEM extrapolates better than competing methods from four-ball training sequences to test sequences containing 6–8 balls.R-NEM and RNN-EM scale markedly better than LSTM, whose relative performance worsens with additional balls.
- Attention: Attention coefficients become non-zero during collisions, indicating that components incorporate information from interacting context objects.Removing attention preserves four-ball performance but worsens extrapolation to environments with more balls.
- Simulation: R-NEM produces mostly accurate recursive simulations over ten predicted time-steps, while an RNN consistently deviates from ground truth.Its failures frequently occur when one component models more than one ball.
- Hidden Factors: R-NEM maintains object states through complete occlusion and accurately predicts a ball’s reappearance after an invisible curtain.The experiment evaluates object modeling and interactions when the object is hidden.
- Space Invaders: On Space Invaders sequences, R-NEM accurately models visually complex frames while groups represent high-level entities with similar movement patterns.The model is trained on 84 × 84 binarized images over 25 gameplay time-steps using four groups and the agent’s action.
5 DISCUSSION AND CONCLUSION
R-NEM uses object-centered inductive biases to learn physical reasoning from visual observations without supervision. It models dynamics across environments, handles occlusion, generalizes to different object counts, and remains limited by grouping, prediction, and training constraints.
- R-NEM incorporates inductive biases about objects and interactions through clustering and an interaction function, learning their specific dynamics from visual observations.
- R-NEM captures physical dynamics more accurately than other methods and generalizes to environments with different numbers of objects.
- R-NEM can predict object movement and collisions during complete occlusion, providing an approximate simulator with object-specific persistence.
- Dynamic grouping would require top-down agent feedback, while injected noise can prevent detection of very small input regularities.
- Interactions among groups make more than ten components unstable, and the multiple objectives can make R-NEM and RNN-EM challenging to train.
- R-NEM assumes object behavior follows inferable rules, and its second loss term hinders modeling complex varying backgrounds.
A EXPERIMENT DETAILS
The experiments use fixed optimization and dataset settings, evaluate grouping against ground truth with ARI, and apply validation-based early stopping.
- All experiments use ADAM with default parameters, batch size 64, and 50 000 training, 10 000 validation, and 10 000 test inputs.
- Grouping quality is measured with Adjusted Rand Index against ground truth while ignoring background and overlap regions.
- Training uses early stopping when validation loss has not improved for 10 epochs.
A.1 BOUNCING BALLS
The bouncing-balls experiments use binary image sequences with heterogeneous balls and elastic collisions. R-NEM combines a recurrent encoder-decoder with a learned interaction function and is evaluated against related recurrent baselines and variants.
- The bouncing-balls data contains 64 × 64 binary-image sequences over 30 time-steps with two ball types differing in mass and radius.One ball is six times heavier and 1.25 times larger in radius than the other.
- The encoder-decoder uses convolutional layers with a recurrent bottleneck, followed by fully connected and upsampling convolutional decoding layers.
- Upsampling followed by ordinary convolution replaces transposed convolutions to avoid frequency artifacts, without layer normalization on the recurrent connection.
- R-NEM is trained with noisy inputs and next-step prediction, using a Bernoulli pixel prior and blocking gradients through γ.The input uses bitflip noise with p = 0.2.
- The interaction function is built from encoding, embedding, effect, and attention multilayer perceptrons.
- Deeper architectures did not produce significant improvement, while comparison and extrapolation experiments use K = 5 and K = 8 components respectively.
- The RNN baseline uses K = 1, and the no-attention R-NEM variant removes MLPatt so that α:,: = 1.
- For simulation, thresholded maxk ψ values at 0.1 replace unavailable ground-truth assignments; the occlusion dataset uses three equal-mass balls and K = 5.
A.2 SPACE INVADERS
The Space Invaders experiment uses binarized visual sequences generated from a pretrained DQN and adapts the architecture and noise level to the Atari observations and actions.
- Space Invaders uses a different encoder and decoder because its images are 84 × 84.
- The listed architecture includes fully connected, recurrent, and decoding layers with layer normalization.
- R-NEM concatenates an embedding of the agent’s action to the hidden state at each time-step.The action embedding uses a single-layer, 10-unit ReLU MLP.
- Training uses K = 4 and input noise 0.02 to preserve tiny elements such as bullets occupying 1–2 pixels.