Learning Interactive World Model for Object-Centric Reinforcement Learning

AI summary

Overview

• Research area: Reinforcement learning and world models, specifically object-centric representation learning for model-based control.
• Technical level: Advanced. The paper assumes familiarity with POMDPs, variational inference, slot attention, and hierarchical RL.
• Scope: The paper proposes a world model that factorizes state into objects and into static versus dynamic attributes, explicitly models pairwise object interactions as sparse time-varying graphs, and uses those interactions as subgoals for a hierarchical policy on simulated robotic and embodied-AI benchmarks.

What This Paper Is About

Most object-centric reinforcement learning methods split the world into objects but leave the relations between those objects implicit, even though real dynamics are driven by collisions, stacking, containment, and physical forces. This paper asks which types and degrees of decomposition actually make latent representations useful for efficient and generalizable policy learning. The authors answer with FIOC-WM, a world model that learns objects, their static versus dynamic attributes, and an explicit interaction structure, then reuses the learned interactions as composable sub-skills for long-horizon control.

Key Contributions

1. A formal problem setting, the Factored Interactive Object-centric POMDP (FIOC-POMDP). Each object state is split into dynamic variables (position, velocity) and constant properties (color, mass, friction), the object's evolution is a self-transition plus a sum of interaction transitions over a sparse time-varying interaction graph, and observations and rewards follow factored generative processes.
2. The FIOC-WM framework (two-level factorization plus interaction learning). It combines object-level factorization, attribute-level static/dynamic factorization, and explicit interaction-graph learning inside a single world model trained from pixels using pre-trained visual encoders (DINO-v2, R3M).
3. An interaction-centric hierarchical policy. A high-level policy selects a sequence of target interaction graphs, and a low-level policy executes each selected interaction using model predictive control (MPC) or proximal policy optimization (PPO), with the low-level policy pre-trained during offline model learning.
4. Empirical evaluation across five benchmark families (SpritesWorld, OpenAI-Gym Fetch, Franka-Kitchen, iGibson, Libero) on reconstruction quality, interaction-structure recovery, disentanglement, single-task policy learning, and attribute/compositional/skill generalization, plus an ablation study on Libero.

Main Findings

• Reconstruction quality (LPIPS, lower is better): FIOC scores 0.007 on Fetch versus 0.042 for Dreamer-V3, 0.039 for TD-MPC2, 0.026 for EIT, and 0.009 for DINO-WM. On Kitchen, FIOC scores 0.038 against Dreamer-V3 0.102, TD-MPC2 0.123, EIT 0.096, and DINO-WM 0.035. On Libero, FIOC scores 0.027 versus 0.089 (Dreamer-V3), 0.061 (TD-MPC2), 0.040 (EIT), and 0.035 (DINO-WM).
• Interaction learning: Evaluated with normalized Structured Hamming Distance (nSHD) between inferred and ground-truth interaction structures, FIOC consistently outperforms attention-based, post hoc interaction inference across all reported cases and shows a smaller empirical generalization gap (the gap between a bar's top and the shaded single-task region). Among the three FIOC variants, variational masks with categorical distributions perform best, particularly with a large number of objects.
• Disentanglement quality: Linear probing MSE against ground-truth attributes in SpritesWorld shows the learned static and dynamic representations capture color/shape and position/velocity respectively. Object-centric DINO features generally outperform vanilla DINO features on dynamic attributes, while object-centric clustering tends to degrade static attribute representations such as color and shape; the disentanglement module improves static attributes within each object.
• Policy learning and generalization (success rate): FIOC reaches 0.79 ± 0.13 on iGibson and 0.76 ± 0.14 on Libero for attribute generalization. For compositional generalization FIOC records 0.86 ± 0.10 on Push & Switch and 0.70 ± 0.09 on Libero. For skill generalization FIOC records 0.81 ± 0.06 on Push & Switch and 0.73 ± 0.06 on Franka Kitchen. On attribute generalization for Push & Switch, TD-MPC2 leads with 0.95 ± 0.02, ahead of EIT 0.92 ± 0.04, FIOC 0.91 ± 0.05, and Dreamer-V3 0.90 ± 0.07. EIT is the second-best baseline across the generalization tasks.
• Sample efficiency: Learning curves sampled every 100 time steps on iGibson and Libero show world models with object interactions (FIOC and EIT) converging faster than Dreamer-V3 and TD-MPC2; FIOC converges faster than EIT and reaches a higher final success rate on Libero.
• Offline RL comparison: Against DINO-WM, the FIOC variants using DINO-v2 and R3M pre-trained embeddings achieve superior performance in both single-task learning and generalization.
• Ablations on Libero: Removing the hierarchical policy causes the largest drop (0.81 to 0.58, down 0.23, on single task; 0.70 to 0.42, down 0.28, on compositional generalization). Removing interaction modeling drops performance to 0.63 (down 0.18) and 0.52 (down 0.18). Using random actions in offline learning gives 0.64 (down 0.17) and 0.48 (down 0.22). Removing the diversity term gives 0.62 (down 0.19) and 0.50 (down 0.20). Removing state factorization gives 0.77 (down 0.04) and 0.64 (down 0.06). Removing the pre-trained low-level policy gives 0.69 (down 0.12) and 0.59 (down 0.11).

Methodology in Plain English

The framework runs in two stages.

Stage 1 (offline model learning). Raw images are passed through pre-trained visual encoders, DINO-v2 or R3M. Slot attention clusters the resulting embeddings into a fixed set of object slots, following prior work. A variational auto-encoder maps each slot to a latent state, with two separate encoders extracting a static part and a dynamic part. The static part is regularized to stay constant over time (a temporal-consistency loss), and a contrastive loss separates static features across different slots so objects do not encode the same attributes. The dynamic part is modeled with a GRU and a variational prior conditioned on the previous state, the previous action, and the interaction graph. The interaction graph itself is a binary N-by-N adjacency matrix per timestep, learned either through variational masks or through conditional independence testing. Reconstruction, next-step prediction, KL, and reward losses are used to train the model. A low-level policy is also trained here by learning the inverse mapping from interactions to actions, using MPC or PPO.

Stage 2 (online hierarchical policy learning). A high-level policy selects a sequence of subgoal interaction graphs that transform the current graph step by step toward the task goal graph, with each subgoal graph differing from the previous one in a single interaction. The action space is narrowed by first choosing a small subset of anchor objects (typically one or two) and then picking a target object conditioned on the anchor, which reduces the combinatorial space. The high-level policy is optimized with the task reward plus a diversity reward defined as one over the square root of the number of graphs visited so far. The low-level policy, pre-trained offline, is fine-tuned online as new interaction data arrives.

Evaluation measures reconstruction with LPIPS, interaction recovery with normalized SHD, attribute disentanglement with linear probing MSE, and control with average success rate. Specific dataset sizes and training compute totals for each benchmark are not reported in the provided text.

Why This Matters

The paper argues that explicit, modular interaction learning, rather than object factorization alone, is the key inductive bias for robust control. It provides evidence that decomposing distant goals into sequences of pairwise object interactions improves sample efficiency and generalization, and it shows that a world model's learned interaction structure can be reused directly as a sub-skill vocabulary rather than being discarded after training.

Real-world applications suggested by the setting and benchmarks:

• Robotic manipulation: long-horizon pick-and-place, pushing, and switch-toggling tasks of the kind represented by OpenAI-Gym Fetch, Franka-Kitchen, and Libero.
• Household and embodied assistants: everyday household interaction tasks with many objects, as in iGibson.
• Simulated data generation and planning: using the learned model as a cheap surrogate for decision making instead of interacting with the real environment.
• Compositional generalization in deployment: recombining previously seen object attributes and skills into new task configurations, which the paper evaluates explicitly.

Industry relevance follows from the sample-efficiency claim: world models that recover interaction structure require less online interaction, and the offline pre-training of a reusable low-level interaction policy gives a natural route toward modular, transferable robot skills.

Future Directions

1. Removing the dependence on a pre-trained object-centric model for object discovery, which the authors list as an explicit limitation of FIOC-WM.
2. Generalizing interaction models beyond seen object categories; the authors state that current interaction models primarily generalize to object categories seen during training.
3. Extending the framework to real-world robotic settings, with the authors pointing to robot-learning foundation models as potential enablers.
4. Scaling the interaction-graph selection and scheduling problem. The high-level action space grows exponentially with the number of objects; the current work sidesteps this by focusing on a small anchor subset and a diversity reward, leaving open how to schedule interactions at larger scale.

Target Audience

Researchers and graduate students working on world models, model-based RL, and object-centric representation learning, particularly those interested in factored state spaces, causal or interaction structure discovery, and hierarchical skill decomposition. Practitioners building robot learning systems that need efficient long-horizon control and compositional generalization will also find the framing useful, though the paper's technical level is advanced and assumes fluency with variational inference, POMDP formalisms, and slot-based representation learning.

Authors’ abstract

Agents that understand objects and their interactions can learn policies that are more robust and transferable. However, most object-centric RL methods factor state by individual objects while leaving interactions implicit. We introduce the Factored Interactive Object-Centric World Model (FIOC-WM), a unified framework that learns structured representations of both objects and their interactions within a world model. FIOC-WM captures environment dynamics with disentangled and modular representations of object interactions, improving sample efficiency and generalization for policy learning. Concretely, FIOC-WM first learns object-centric latents and an interaction structure directly from pixels, leveraging pre-trained vision encoders. The learned world model then decomposes tasks into composable interaction primitives, and a hierarchical policy is trained on top: a high level selects the type and order of interactions, while a low level executes them. On simulated robotic and embodied-AI benchmarks, FIOC-WM improves policy-learning sample efficiency and generalization over world-model baselines, indicating that explicit, modular interaction learning is crucial for robust control.