Entity Abstraction in Visual Model-Based Reinforcement Learning

Rishi Veerapaneni; John D. Co-Reyes; Michael Chang; Michael Janner; Chelsea Finn; Jiajun Wu; Joshua B. Tenenbaum; Sergey Levine

Source-linked AI summary

Entity Abstraction in Visual Model-Based Reinforcement Learning

Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua B. Tenenbaum, Sergey Levine

arXiv:1910.12827v5Oct 28, 2019cs.LGcs.CVcs.NEstat.ML

TL;DR

The paper asks whether entity-centric modeling improves generalization to unseen combinatorial physical tasks. It introduces OP3, which learns and grounds entity variables from raw visual observations for prediction and planning, and reports substantially better block-stacking accuracy than video prediction baselines. The discussion attributes this scope to reusable local entity functions, while noting the symmetry assumption underlying the abstraction.

Problem
The paper asks whether modeling entities and local interactions improves generalization to physical tasks with unseen object configurations and counts.
Method
OP3 combines object-centric perception, prediction, and planning with interactive inference that grounds unsupervised entity variables in raw visual observations.
Results
Two to three times greater accuracy than state-of-the-art video prediction models is reported for novel single- and multi-step block-stacking tasks.
Takeaways & Limitations
Entity abstraction enables reusable local functions for individual entities and entity pairs, supporting modeling of variable numbers of objects.
Takeaways & Limitations
OP3 assumes that multi-entity scene dynamics can be reduced to modeling one entity and its interactions, then applying the same function to every entity.

Abstract

from arXiv · show

This paper tests the hypothesis that modeling a scene in terms of entities and their local interactions, as opposed to modeling the scene globally, provides a significant benefit in generalizing to physical tasks in a combinatorial space the learner has not encountered before. We present object-centric perception, prediction, and planning (OP3), which to the best of our knowledge is the first fully probabilistic entity-centric dynamic latent variable framework for model-based reinforcement learning that acquires entity representations from raw visual observations without supervision and uses them to predict and plan. OP3 enforces entity-abstraction -- symmetric processing of each entity representation with the same locally-scoped function -- which enables it to scale to model different numbers and configurations of objects from those in training. Our approach to solving the key technical challenge of grounding these entity representations to actual objects in the environment is to frame this variable binding problem as an inference problem, and we develop an interactive inference algorithm that uses temporal continuity and interactive feedback to bind information about object properties to the entity variables. On block-stacking tasks, OP3 generalizes to novel block configurations and more objects than observed during training, outperforming an oracle model that assumes access to object supervision and achieving two to three times better accuracy than a state-of-the-art video prediction model that does not exhibit entity abstraction.

1 Introduction

The paper hypothesizes that modeling objects and their local interactions can improve generalization across combinatorial physical tasks. OP3 implements this idea through entity-centric perception, prediction, planning, and unsupervised interactive grounding.

Motivation: Compositional physical task spaces may grow combinatorially with object configurations and counts, motivating models of objects and their local interactions.This is the paper’s central hypothesis about generalization to novel physical tasks.
Entity abstraction: Entity abstraction applies the same locally scoped function symmetrically to each entity, separating each object’s modeling from computations involving other objects.The approach combines entity-centric functions with an interactive inference algorithm for grounding abstract variables in observations.
OP3: OP3 is a model-based reinforcement learner that predicts and plans over entity variables inferred from raw visual observations.Its entity variables are acquired without supervision on object identity.
Generalization: OP3 generalizes to variable numbers of objects and varied contexts without retraining by modeling entities locally rather than scenes globally.This follows from protecting each object’s dynamics and appearance modeling from computations involving other objects.
Results: OP3 achieves two to three times greater accuracy than state-of-the-art video prediction models on novel single- and multi-step block-stacking tasks.The comparison is reported as an empirical contribution of the paper.

2 Related Work

Related work often represents scenes globally, assumes fixed or supervised entities, or lacks mechanisms for temporal entity correspondence. OP3 instead combines symmetric entity processing with recurrent probabilistic inference for dynamic, unsupervised grounding.

Comparison: OP3 differs from related dynamic-scene methods through its fully probabilistic factorized latent-variable formulation and symmetric processing, including entity disambiguation.The caption contrasts this design with global or autoregressive/convolutional processing in other methods.
Representation learning: Global scene representations may be susceptible to the binding problem and may require data to learn that the same object transfers across contexts.The cited prior approaches do not encode the inductive biases associated with individual-object properties and inter-object relationships.
Unsupervised grounding: Many entity-based methods pre-specify identities or require supervision, segmentations, crops, or a simulator, while other unsupervised methods factorize scenes into pixel-level entities.The paper identifies grounding entities in raw observations as a crucial unresolved problem for autonomous learning and interaction.
Entity abstraction: OP3 processes each entity representation with the same function, enforcing local-property invariance to changes in global structure.This supports transferring models of a single entity across different global contexts.
Temporal modeling: Unlike Transporter, COBRA, and C-SWMs, OP3 can update entity beliefs with information from subsequent frames through recurrent probabilistic dynamics.The comparison concerns temporal refinement and correspondence between entity representations across frames.

3 Problem Formulation

The problem formulation represents a visual scene with latent variables for its objects and seeks a predictive model that supports planning for novel compositional manipulation tasks. Learning is framed probabilistically through observation, dynamics, and recognition distributions optimized with an ELBO.

Latent representation: OP3 represents a scene’s object states with a set of latent random variables H1:K rather than one latent variable for the whole scene.The paper distinguishes physical objects from model entities representing their states.
Generative model: The model uses X for image observations and A for agent actions, with observation and dynamics distributions shared across timesteps.The goal is to model observations and latent entities generated over action sequences.
Generalization target: The target is to solve novel compositional object-manipulation problems, such as building block towers after training only on how blocks fall.This formulation emphasizes generalization from observed interactions to new tasks.
Prediction and planning: Planning requires approximating the posterior predictive distribution of observations d steps into the future given past observations and actions.The paper casts this approximation as a variational inference problem.
Learning objective: The ELBO trains recognition states that reconstruct observations through G and predict future entity states through D.The objective links perceptual reconstruction with predictive dynamics modeling.

4 Object-Centric Perception, Prediction, and Planning (OP3)

OP3 builds perception, dynamics, inference, and planning around symmetrically processed entity variables rather than a global scene representation. Its models use local entity and pairwise interactions, while interactive inference grounds latent variables to objects from raw visual observations.

Entity abstraction: Entity abstraction applies the same entity-centric function to every entity, transferring learned modeling of one generic entity across scene instances.This symmetry is encoded through map(f, H1:K), which broadcasts the function to all entity variables.
Object-centric perception: The observation model represents images as compositions of entity-generated sub-images weighted by masks encoding relative depth.The mixture model combines per-entity components at each pixel using weights derived from entities’ depth from the camera.
Object-centric prediction: The dynamics model applies the same function d to each entity, combining individual dynamics, action effects, and interactions with other entities.A pairwise function doo is reused across entity pairs, reducing variable configurations to generic entity and entity-pair operations.
Interactive inference: Variable binding is framed as posterior inference, because a single parallel recognition pass cannot assign different objects to entity variables.Iterative inference refines posterior parameters using generative-model feedback, while interactive inference adds temporal continuity and actions to disambiguate objects.
Object-centric planning: Grounded entity pointers support object-centric actions by mapping an entity identifier to pick coordinates, enabling (entity id, place xy) actions without a scripted policy.This extends OP3 beyond prediction and planning in entity space to selecting which object to manipulate.

5 Experiments

OP3 is evaluated on simulated block-stacking tasks requiring combinatorial generalization, single-step and multi-step planning, and unsupervised object grounding. It outperforms the compared baselines and uses interactive inference to represent coherent objects in simulated and real-world scenes.

5.1 Combinatorial Generalization without Object Supervision: OP3 trains only on predicting falling objects, then generalizes to novel block-stacking tasks with two to three times better accuracy than a state-of-the-art video prediction model.The training set contains up to five objects, while test scenes contain up to nine objects in unseen structures.
5.1 Combinatorial Generalization without Object Supervision: OP3 achieves better block-tower accuracy than O2P2 despite using raw images without ground-truth object-identity supervision.O2P2 uses image segmentations, whereas OP3 uses only raw images as input.
5.2 Multi-Step Planning: Entity pointers let OP3 bias pick locations toward objects, improving multi-step planning over random pick locations in sparse scenes.Random pick locations rarely intersect objects, whereas entity pointers provide object locations without supervision.
5.2 Multi-Step Planning: In multi-step planning, OP3 achieves 1.5-2 times SAVP’s accuracy with pick-location actions and 2-3 times SAVP’s accuracy with entity-based actions.The entity-based action space is (entity id, place xy), while the standard space is (pick xy, place xy).
5.2 Multi-Step Planning: Interactive inference uses temporal continuity and action feedback to disambiguate objects and refine entity representations during training.OP3 alternates refinement and prediction steps, producing coherent object decompositions and retaining information about partially occluded objects.
5.3 Real World Evaluation: On cluttered real-world videos, OP3 obtains more accurate object segmentations than vanilla IODINE when movement is needed to disambiguate objects.OP3 separates entities such as a towel, robotic arm, and container edges after observing the effects of arm movement.

6 Discussion

The discussion presents OP3 as a probabilistic bridge between entity-level symbolic reasoning and continuous visual data. Its design combines locally shared entity functions with inference-based grounding and predictive rollouts.

6 Discussion: Entity abstraction factorizes scenes into local entity states processed symmetrically by the same entity-centric function.This differs from global latent-state models and fixed, permutation-sensitive entity processing.
6 Discussion: Framing entity variables as latent random states turns binding abstract variables to concrete objects into an inference problem.The model refines posterior estimates from raw visual observations rather than receiving object-identity supervision.
6 Discussion: OP3’s predictive posterior rollout extends the dynamics model beyond the last observation and uses the observation model to predict future observations.The associated objective adds future reconstruction terms to the standard ELBO terms for observed timesteps.
6 Discussion: Interactive inference updates entity posteriors using observations, previous states, actions, and learned dynamics predictions.The procedure can be trained end-to-end with an ELBO at every timestep and adapted to multi-step rollouts.

E Cost Function

The cost function compares candidate and goal configurations through entity-level masked subimages and a color-sensitive pixel criterion. This decomposition better supports block-tower planning than a cost based on whole images.

E Cost Function: OP3 represents each entity with a masked subimage and decomposes configuration cost into distances between entity states.The masked subimage is defined as ˆI(Hk) := m(Hk) · g(X | Hk).
E Cost Function: Single-step planning uses L2 distance between corresponding masked subimages to compare entity states.This cost is used in the first block-stacking environment.
E Cost Function: Multi-step planning uses a color-sensitive distance that avoids assigning low cost when a block of the wrong color occupies another block’s goal position.The criterion compares overlapping object pixels and penalizes mismatched colors.

F Architecture and Hyperparameter Details

OP3 uses shared neural components to model entity dynamics, action effects, object interactions, and observations. The implementation includes recurrent posterior refinement, mixed deterministic-stochastic latents, and end-to-end training with ELBO objectives.

F Architecture and Hyperparameter Details: The latent representation combines a deterministic component of size Rd with a stochastic component of size Rs.The paper reports Rs = 64 and R = 128 for the listed architecture settings.
F Architecture and Hyperparameter Details: The refinement network receives image-sized likelihood, mask, gradient, and coordinate inputs before recurrently updating posterior parameters.Posterior parameters and their gradients are concatenated with convolutional features and passed to an LSTM.
F Architecture and Hyperparameter Details: The dynamics model applies the same action-conditioned function to each entity, using other entities to model interactions.Action and object effects are represented by shared functions across entities and entity pairs.
F Architecture and Hyperparameter Details: The dynamics network outputs posterior parameters that are split into deterministic and stochastic components by separate networks.The component networks are fdet and fsto, and the functions are parameterized by single-layer MLPs.

G.1 Single-Step Block-Stacking

The single-step block-stacking dataset contains trajectories pairing action-conditioned before images with after images, requiring prediction of each block’s steady-state outcome.

The dataset contains 60,000 trajectories of before-and-after 64x64 images.
Each action selects a block shape, color, position, and orientation before dropping it.
The model predicts the steady-state result after the dropped block settles.

G.2 Multi-Step Block-Stacking

The multi-step task evaluates block manipulation across sequences of pick-and-place actions, using entity-based actions to address sparse successful pick locations and visualizing iterative inference and prediction.

Dataset: The dataset contains 10,000 trajectories from environments with two differently colored blocks, each comprising five 64x64 frames.
Dataset: Actions include randomly relocating blocks, stacking blocks on other blocks, and random pick-and-place locations.
Qualitative results: OP3 predicts steady-state effects while capturing gravity and interactions with other objects.
Qualitative results: Rollouts refine a single input image across iterations before predicting outcomes from the learned hidden states.
Action spaces: Entity Pick Place selects a latent entity to move before choosing its destination, avoiding uniformly sampled pick locations with only 2% success probability.

H Ablations

Ablations show that entity-centric structure matters: removing symmetry or decomposing costs globally harms generalization, while entity-factorized costs retain the model’s advantage.

Model ablations: Unfactorized and no-weight-sharing models are unable to sufficiently generalize on the block-stacking task.
Cost ablations: Using an unfactorized cost significantly underperforms decomposing the cost per entity’s masked subimage.
Model ablations: The no-weight-sharing model did not converge during training.
Interpretability: The paper does not explicitly explore interpretability, although entity-factorized representations are presented as interpretable by construction.