Table of Contents
Fetching ...

Joint Learning of Hierarchical Neural Options and Abstract World Model

Wasu Top Piriyakulkij, Wolfgang Lehrach, Kevin Ellis, Kevin Murphy

TL;DR

AgentOWL addresses the challenge of scalable skill acquisition by jointly learning an abstract, temporally extended world model and a hierarchy of neural options. It combines a PoE-World-based abstract world model with LLM-assisted sub-goal proposals and model-based exploration to improve sample efficiency on hard object-centric Atari tasks. The approach yields more mastered options with fewer environment steps, enables zero-shot generalization to novel starting states, and demonstrates implicit sub-option refinement through world-model guided planning. These results highlight the value of tightly integrating abstract planning with hierarchical skill discovery for data-efficient, long-horizon decision making.

Abstract

Building agents that can perform new skills by composing existing skills is a long-standing goal of AI agent research. Towards this end, we investigate how to efficiently acquire a sequence of skills, formalized as hierarchical neural options. However, existing model-free hierarchical reinforcement algorithms need a lot of data. We propose a novel method, which we call AgentOWL (Option and World model Learning Agent), that jointly learns -- in a sample efficient way -- an abstract world model (abstracting across both states and time) and a set of hierarchical neural options. We show, on a subset of Object-Centric Atari games, that our method can learn more skills using much less data than baseline methods.

Joint Learning of Hierarchical Neural Options and Abstract World Model

TL;DR

AgentOWL addresses the challenge of scalable skill acquisition by jointly learning an abstract, temporally extended world model and a hierarchy of neural options. It combines a PoE-World-based abstract world model with LLM-assisted sub-goal proposals and model-based exploration to improve sample efficiency on hard object-centric Atari tasks. The approach yields more mastered options with fewer environment steps, enables zero-shot generalization to novel starting states, and demonstrates implicit sub-option refinement through world-model guided planning. These results highlight the value of tightly integrating abstract planning with hierarchical skill discovery for data-efficient, long-horizon decision making.

Abstract

Building agents that can perform new skills by composing existing skills is a long-standing goal of AI agent research. Towards this end, we investigate how to efficiently acquire a sequence of skills, formalized as hierarchical neural options. However, existing model-free hierarchical reinforcement algorithms need a lot of data. We propose a novel method, which we call AgentOWL (Option and World model Learning Agent), that jointly learns -- in a sample efficient way -- an abstract world model (abstracting across both states and time) and a set of hierarchical neural options. We show, on a subset of Object-Centric Atari games, that our method can learn more skills using much less data than baseline methods.
Paper Structure (36 sections, 6 equations, 9 figures, 7 tables)

This paper contains 36 sections, 6 equations, 9 figures, 7 tables.

Figures (9)

  • Figure 1: Illustration of hierarchical planning and execution in AgentOWL where the goal is to go to the left platform in the left room. Left: AgentOWL goes through possible plans and successfully makes a short plan (two high level steps) in its abstract world model to reach the goal. Right: AgentOWL executes a hierarchical sequence of options. (The hierarchical structure is represented by the indentation.)
  • Figure 2: Illustration of how AgentOWL learns its abstract world model (left) and hierarchical options (right). Left: given a dataset of option transitions, we learn an abstract world model using (an extension of) the method of piriyakulkij2025poeworld. Specifically, each expert is generated using LLM code synthesis, and the weight for each expert (denoted $\theta_i$) is learned using gradient descent on the likelihood objective. Right: AgentOWL learns a new option to achieve a new goal by leveraging previously acquired sub-options. Specifically, given the goal, we ask an LLM to hypothesize new sub-options (building on the already learned ones) that might help achieve the goal. The agent trains to master the highest-level option, $o_g$, which implicitly train these new sub-options, and simultaneously update our abstract world model to identify sub-options with high goal completion rates. Eventually, the agent masters the target option $o_g$ by composing good sub-options. Detailed pseudocode can be founded in \ref{['fig:pseudocode']}.
  • Figure 3: Fraction of options mastered vs number of environment steps for the three OCAtari's games we test on: Montezuma's Revenge, Pitfall, and Private Eye. Option is acquired once its success rate for the recent episodes reaches threshold $\delta=0.5$.
  • Figure 4: Screenshots of 4 rooms of Pitfall stitched together. Player starts in Room 1 (rightmost) and can traverse to other rooms through the sides of the screen. Goals that only AgentOWL masters within 5M environment steps (\ref{['fig:main_results']} middle) are shown in red borders.
  • Figure 5: Ablation study: Removing core components of AgentOWL lead to fewer options being acquired and/or more data from the environment being needed. Setting $n_{threshold} = 0$ means stabilization for hierarchical DQN is not implemented.
  • ...and 4 more figures