🤖 machine learning

Joint Learning of Hierarchical Neural Options and Abstract World Model

The paper introduces AgentOWL, a sample-efficient method that jointly learns an abstract world model and hierarchical neural options to enable AI agents to acquire and compose new skills with less data and better generalization than existing model-free approaches.

Original authors: Wasu Top Piriyakulkij, Wolfgang Lehrach, Kevin Ellis, Kevin Murphy

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Wasu Top Piriyakulkij, Wolfgang Lehrach, Kevin Ellis, Kevin Murphy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to make a cup of coffee. You don't want it to learn every single tiny movement from scratch every time—like how to move its finger 1 millimeter to the left, then 1 millimeter forward. That would take forever and require millions of tries. Instead, you want it to learn "skills" (like "grab the mug," "pour the water," "press the button") and then combine those skills to make coffee.

This paper introduces a new AI system called AgentOWL (Option and World model Learning Agent) that is really good at learning these skills quickly and then using them to solve new, harder problems.

Here is how it works, broken down into simple concepts:

1. The Problem: The "Trial and Error" Trap

Most current AI robots learn by trying things randomly until they succeed. If you want them to learn a long chain of skills (like making coffee), they have to try random combinations of movements billions of times. It's like trying to open a safe by guessing every possible combination of numbers one by one. It's slow and inefficient.

2. The Solution: Two Superpowers

AgentOWL solves this by combining two things: Hierarchical Options (Skills) and an Abstract World Model (A mental map).

A. The "Skill Ladder" (Hierarchical Options)

Think of this like a video game with a "Save Point" system.

Level 1: The robot learns a simple skill, like "Walk to the door." Once it masters this, it saves it as a single button press.
Level 2: Now, instead of learning "Walk to the door" again, the robot uses that saved button to learn a bigger skill, like "Go to the kitchen."
Level 3: It keeps stacking these. "Go to kitchen" + "Pick up mug" = "Get coffee."

This creates a deep ladder of skills. The robot doesn't have to relearn how to walk every time it wants to get coffee; it just uses the "Walk" skill it already knows.

B. The "Mental Map" (Abstract World Model)

This is the secret sauce. Instead of trying to predict every single pixel on the screen (which is like trying to predict the weather by looking at every single raindrop), AgentOWL builds a simplified mental map.

The Analogy: Imagine you are planning a road trip. You don't need to know the color of every house you pass. You just need to know: "If I take Highway 101, I will end up in San Francisco."
How AgentOWL does it: It uses a special "World Model" that predicts the outcome of a skill, not the tiny steps in between. It asks: "If I use my 'Walk to door' skill, will I end up in the kitchen?" It ignores the messy details of the journey and focuses on the result.

3. The "Dreamer" Strategy

Here is the clever part: AgentOWL practices in its dreams (the mental map) before it tries in the real world.

Imagination: It simulates thousands of scenarios in its head using its simplified map. It tries different combinations of skills to see which ones lead to the goal.
Reality Check: Once it finds a good plan in its head, it tries that specific plan in the real game.
Learning: If the real world is slightly different from the dream, it updates its mental map.

This is why it is so efficient. It doesn't waste time running into walls in the real world; it figures out the path in its head first.

4. The "Smart Assistant" (Using LLMs)

Sometimes the robot gets stuck because it doesn't know what skill to learn next to reach a new goal. To fix this, AgentOWL asks a Large Language Model (LLM) for help.

The Analogy: Imagine you are trying to build a complex Lego castle, but you are missing a specific piece. You ask a friend (the LLM): "Hey, I have a tower and a wall. What piece do I need to connect them?"
The LLM suggests a "bridge" skill. The robot then tries to learn that specific bridge skill. This helps the robot build its skill ladder much faster than if it had to guess randomly.

5. The Results: What Did They Prove?

The researchers tested AgentOWL on three very hard video games (Atari games like Montezuma's Revenge and Pitfall) where the robot has to explore a maze and find specific objects.

Faster Learning: AgentOWL learned more skills using less data (fewer game attempts) than other standard AI methods.
Solving Hard Puzzles: Other AI methods gave up on the hardest levels because the path was too long and complex. AgentOWL succeeded because it could break the long path into small, manageable "skills" and plan them out in its head.
Zero-Shot Generalization: This is the coolest part. The researchers moved the robot to a brand-new starting position it had never seen before. Because AgentOWL understood the logic of the world (the map) and had a library of skills, it could immediately figure out how to get to the goal without any new training. It was like moving a chess player to a new board; they didn't need to relearn how to move the pieces, they just applied their strategy to the new setup.

Summary

AgentOWL is like a student who doesn't just memorize answers but learns concepts.

It breaks big problems into small, reusable skills.
It builds a simplified mental map to predict what those skills do.
It practices in its head to find the best path before acting.
It asks a smart assistant for help when it gets stuck.

The result is an AI that learns new, complex tasks much faster and can adapt to new situations without needing to be retrained from scratch.

Technical Summary: Joint Learning of Hierarchical Neural Options and Abstract World Model (AgentOWL)

1. Problem Statement

The paper addresses the challenge of enabling decision-making agents to cumulatively acquire new skills by composing existing ones. While the options framework formalizes this as learning a hierarchy of policies (options) that achieve increasingly difficult goals, standard model-free hierarchical reinforcement learning (RL) approaches suffer from poor sample efficiency. As the number of acquired skills grows, the effective action space expands, making policy learning increasingly data-intensive and intractable for new goals.

The core problem is a tradeoff: learning new skills quickly versus the number of skills already acquired. Standard model-free methods require exponentially more samples as learning progresses. The authors propose that model-based RL can resolve this by using a world model to plan and rule out options before real-world execution. However, existing world modeling approaches often lack the data efficiency required to learn complex dynamics from limited samples, or they fail to provide the necessary temporal and state abstractions to make planning tractable.

2. Methodology: AgentOWL

The authors propose AgentOWL (Option and World model Learning Agent), a system that jointly learns a set of hierarchical neural options and an abstract world model in a sample-efficient manner.

2.1 Abstract World Modeling

The system employs a Product-of-Experts (PoE-World) framework to learn a world model that abstracts over both states and time.

State Abstraction: Instead of predicting raw pixel transitions, the model predicts future abstract states defined by goal predicates $f(s) = (g_1(s), g_2(s), \dots)$ . This reduces the dimensionality of the prediction task.
Symbolic Representation: The world model is represented as a product of experts, where each expert is a short symbolic program (Python snippet) generated by a Large Language Model (LLM). These programs model independent causal mechanisms.
Frame Axiom Prior: To maintain sample efficiency, the model incorporates a "frame axiom prior." This biases the model to believe that an option $o_i$ primarily changes its corresponding feature $f_i$ and does not change other features $f_j$ ( $j \neq i$ ). This is implemented via Maximum A Posteriori (MAP) estimation with Gaussian priors, significantly reducing the data required to learn the transition dynamics.
Reconstruction: Since the abstract model predicts $p(f' | s)$ , a kernel density estimator (weighting function) is used heuristically to sample full states $s'$ from abstract states $f'$ , allowing the agent to interact with a simulated environment without learning a full parametric generative model.

2.2 Joint Learning of Hierarchical Options

AgentOWL iteratively trains to master a sequence of goals $\{g_1, \dots, g_n\}$ by learning corresponding options $\{o_1, \dots, o_n\}$ .

LLM-Based Sub-Option Hypothesis: When an agent needs to achieve a new target goal $g$ but lacks a direct option, an LLM is prompted to hypothesize a new sub-option $o_{h \to g}$ . The LLM selects a precondition $h$ (a previously learned sub-goal) from the existing set of options $\Omega$ that could serve as a stepping stone. This new option and its hypothetical world model are added to the system.
Model-Based Exploration: The agent learns a policy $\pi_{wm}$ within the abstract world model using Deep Q-Learning (DQN). This policy serves as an exploration guide for the real-world policy $\pi_{real}$ . The combined policy is defined as $\pi(a|s) = (1-\epsilon)\pi_{real}(a|s) + \epsilon\pi_{wm}(a|s)$ . The exploration rate $\epsilon$ is annealed from 1 to 0, allowing the agent to transition from model-based planning to fully model-free learning as the real-world policy stabilizes.
Stable Hierarchical Training: To mitigate the non-stationarity caused by training lower-level options (which changes the environment seen by higher-level options), the system employs a stabilization mechanism. Data from episodes is only added to the replay buffer for training if all sub-options executed in that episode are "stable" (i.e., they have been trained on sufficient samples $n > n_{threshold}$ and have a high goal completion rate $\delta > \delta_{threshold}$ ).

3. Key Contributions

AgentOWL Framework: A novel architecture that jointly learns hierarchical neural options and an abstract world model, enabling sample-efficient skill acquisition.
Abstract World Model with Symbolic Programs: The integration of PoE-World with LLM-generated symbolic programs allows for learning complex, temporally abstract dynamics from very little data, overcoming the "one-step trap" of low-level models.
LLM-Driven Option Discovery: A method for using LLMs to hypothesize new sub-options (preconditions) based on existing skills, effectively automating the curriculum of skill acquisition without requiring manual engineering of sub-goals.
Stable Hierarchical DQN: A training procedure that stabilizes the learning of deep option hierarchies by filtering data based on the stability of sub-options, addressing the non-stationarity problem inherent in hierarchical RL.

4. Experimental Results

The method was evaluated on three object-centric Atari (OCAtari) games known for hard exploration: Montezuma's Revenge, Pitfall, and Private Eye.

Skill Acquisition Efficiency: AgentOWL mastered a significantly higher fraction of options using fewer environment steps compared to baselines, including Rainbow DQN, Goal-conditioned DQN, and Hierarchical DQN (without the world model). While baselines plateaued at a low number of mastered skills, AgentOWL continued to acquire complex skills.
Generalization to Hard Goals: Qualitative analysis showed that baselines failed to acquire options for difficult goals (e.g., reaching specific rooms in Pitfall) because they could not discover the necessary long action sequences. AgentOWL successfully planned these sequences abstractly and executed them hierarchically.
Implicit Learning: Experiments demonstrated that AgentOWL implicitly improves sub-options even when only rewarded for a higher-level goal. Sub-options useful for the target goal showed increased success rates, while irrelevant ones did not.
Zero-Shot Generalization: In a novel starting state experiment, AgentOWL could compose existing options to achieve previously mastered goals without additional training, provided it was given a single new option to navigate from the new start to the original start state. Baselines lacking the abstract world model and hierarchical composition failed to adapt zero-shot.

5. Significance and Claims

The paper claims that option learning and world model learning are deeply intertwined. By learning an abstract world model that captures the effects of options, the agent can plan efficiently, which in turn facilitates the discovery and refinement of new skills.

The authors emphasize that AgentOWL demonstrates:

Sample Efficiency: Achieving complex skill hierarchies with significantly less data than model-free baselines.
Compositional Generalization: The ability to combine learned skills to solve novel problems (zero-shot generalization) and to implicitly learn sub-skills required for higher-level goals.
Scalability of Planning: The use of abstract, temporally extended models allows for planning over long horizons that would be intractable for low-level models.

The paper concludes that while current limitations exist (such as the need for ordered goal sequences and symbolic inputs), the approach highlights a promising direction for building agents that can cumulatively expand their knowledge and capabilities through the interplay of skill acquisition and world modeling.