Efficient Exploration and Discriminative World Model Learning with an Object-Centric Abstraction
Authors: Anthony GX-Chen
Kenneth Marino
Rob Fergus
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
model based reinforcement learning
world model
exploration
hierarchy
Primary area: reinforcement learning
TL;DR: We develop a discriminative model-based RL method to efficiently learn a semantic-level world model

Abstract: In the face of difficult exploration problems in reinforcement learning, we study whether giving an agent an object-centric mapping (describing a set of items and their attributes) allow for more efficient learning. We found this problem is best solved hierarchically by modelling items at a higher level of state abstraction to pixels, and attribute change at a higher level of temporal abstraction to primitive actions. This abstraction simplifies the transition dynamic by making specific future states easier to predict. We make use of this to propose a fully model-based algorithm that learns a discriminative world model, plans to explore efficiently with only a count-based intrinsic reward, and can subsequently plan to reach any discovered (abstract) states.

We demonstrate the model's ability to (i) efficiently solve single tasks, (ii) transfer zero-shot and few-shot across item types and environments, and (iii) plan across long horizons. Across a suite of 2D crafting and MiniHack environments, we empirically show our model significantly out-performs state-of-the-art low-level methods (without abstraction), as well as performant model-free and model-based methods using the same abstraction. Finally, we show how to learn low level object-perturbing policies via reinforcement learning, and the object mapping itself by supervised learning.

Published as a conference paper at ICLR 2025
EFFICIENT EXPLORATION AND DISCRIMINATIVE
WORLD MODEL LEARNING WITH AN OBJECT-
CENTRIC ABSTRACTION
Anthony GX-Chen
Kenneth Marino
Rob Fergus
Center for Data Science
Google DeepMind The University of Utah
Dept. of Computer Science
New York University
New York University
anthony gx. chen@nyu. edu
ABSTRACT
In the face of difficult exploration problems in reinforcement learning, we study whether giving an agent an object-centric mapping (describing a set of items and their artributes) allow for more efficient learning. We found this problem is best solved hierarchically by modelling items at a higher level of state abstraction to pixels, and attribute change at higher level of temporal abstraction to primitive actions. This abstraction simplifies the transition dynamic by making specific future states easier to predict. We make of this propose a fully model-based algorithm that learns a discriminative world model, plans to explore efficiently with only a count-based intrinsic reward, and can subsequently plan reach any discovered (abstract) states.
We demonstrate the model's ability to (i) efficiently solve single tasks, (ii) transfer zero-shot and few-shot across item types and environments, and (iii) plan across long horizons. Across a suite of 2D crafting and MiniHack environments, we empirically show our model significantly out-performs state-of- tfae-oof-the-ar low-level methods (without abstraction), well as performant model-free and model- based methods using the same abstraction. Finally, we show how to learn low level object-perturbing policies via reinforcement learning, and the object mapping it- self by supervised learning.
INTRODUCTION
Since the inception of reinforcement learning (RL), the problems of exploration and world model learning have been active, open questions. RL requires agent to learn both basic motor abilities, and explore long sequences of interactions the environment. Currently, we have made tremendous progress problems of low-level control for short-horizon problems. With well defined single tasks, given enough samples (or demonstrations). there are well-developed ways of training low level policies reliably (Schulman al., 2017; Hafner al., 2023).
Perhaps a better way to explore complex environments is to treat the high-level exploration problem separately. We increasingly have the ability get semantically rich abstractions: via representation learning (Locatello et al., 2020), object segmentation (Kirillov et al., 2023), and visual-language models (Alayrac et al., 2022); even difficult control settings such roboties (Liu et al., 2024). In addition, with the recent explosion in the field of natural language processing, people are in- creasingly interested in addressing RL environments on a semantic level (Andreas et al., 2017: Jiang et al., 2019; Chen et al., 2021). Thus, we ask timely question: if the agent has good ability to perceive objects, what more can we do other than tabula rasa learning at the level of pixels and primitive motor actions?
In this work, we focus on exploration and world modelling at a semantic level. We explore a simple yet flexible abstraction: in addition to the agent's raw observations (e.g. pixels), it sees an abstract representation of set of items and their attributes (i.e. object states). We find that to navigate these two levels effectively, it is best construct proxy decision process (Ab-MDP) with state and temporal abstractions at the level of objects, with handful of competent object-perturbng low level policies. By construction, dynamics this semantic level becomes structured, and

---

Published as a conference paper at ICLR 2025
make use of this to design a discriminative model-based RL method that learns to predicts if specific object perturbation will be successful. Our method, Model- Exploration of abstracted Attribute Dynamics (MEAD), is fit with simple objective to model semantic level transitions, resulting in stable model improvements with no need for auxiliary objectives. It is fully online: the agent plans to explore without extrinsic reward learn about the world it inhabits using count- based objective. The same model can be used at any point to plan to solve downstream tasks when given a goal function, without the need train policy or value functions.
The main results of this work investigates decision making at the semantic level as defined by the Ab-MDP, assuming access an object centric map and object perturbing policies. We empirically evaluate our method in a set of 2D crafting games, and MiniHack environments that are known to be very difficult to explore (Henaff et al., 2022). We then show how to learn components of the Ab- MDP when it is not given: by reinforcement learning object perturbation policies given an object mapping, and supervised learning of the object centric mapping itself. We leave the unsupervised discovery of the object centric mapping complementary and future work.
The main contributions of this work are:
We define the Ab-MDP: simple human-understandaable hierarchical MDP. The structure this abstraction allows for pre-trained or from-scratch visual encoders to used. also allows method (MEAD)
discriminative objective for world-mode learning.
We introduce our model-learning method (MEAD). MEAD uses stable discriminative ob- jective model item attribute changes, which is more efficient learning than generative- based baselines and also allows for efficient inference-tima planning.
We empirically show MEAD extracts an interpretable world model, accurately plans over long horizons, reaches competitive final performance all levels with greater sample effi- ciency to strong baselines, and transfers well to new environments.
We demonstrate how to learn parts of the Ab-MDP when is not given: by reinforcement learning object-perturbing policies, and supervised learning the object-centric mapping-
2 PROBLEM SETTING
We begin by considering a (reward free) "base" Markov Decision Process (MDP) as the tu- ple A,
A, P). with state space S, primitive action space A, and transition probability function S x Ax > [0, 1] (Puterman, 1994). Below we construct the "abstracted" space.
2.1 ABSTRACT STATES AND BEHAVIOURS
We give the agent deterministic object-centric mapping M: S - X' from low level observations to abstract state X E X. The abstract state is a discrete of items, each having form (item identity, item attribute). Denote the i-th item in this (ordered) set zz(1), each ((() consists of an item's identity (+) (e.g. "potion"), and attribute g(() (e.g. "in inventory"). In practice, each X is represented as x (diden + daw)) matrix containing N items (including empty items) each having diden dimensional identity embedding and datr dimensional attribute embedding.
(,()g(()
Denote a space B of object perturbation policies which we call behaviours. Each behaviour has an associated mi + A. behaviour b = (a(*)e describes a single item ( and desired new attribute. For N objects and attributes, there are N x single item behaviours. Not all behaviours changes (6)) to have g' (due to the change being impossible from current X, or having a bad mτ). behaviour competent if executing behaviour policy m, from state S, M(S) = X, result in arriving at state S', M(S') = X', (, E X', within ie steps with high probability. See Appendix B.I for more details. Executing competent behaviours lead to predictable changes.
For the main experiments we assume access to (learned or given) abstract states and behaviours and build on top. We address learning both in Section 4.3.
2.2 ABSTRACTED ITEM-ATTRIBUTE MDP
We now define an Abstracted Item-Attribut MDP using the abstract state X' and behaviour B spaces, which we refer to in brief as Ab-MDP. The abstract mapping M provides an state abstraction. We
Here we define single item hehaviours. Behaviours also he multi items, Appendix 1.
2

---

Published as a conference paper at ICLR 2025
further define the Ab- to have coarser temporal resolution than the low level MDP: abstract transition only occurs when an abstract state changes, or if the state stays the same but & low level steps clapses (see Definition B.1 for more details).
Specifically, the Ab-MDP is the tuple (X, T), with abstract states X. abstract behaviours B, and transition probability function T X x B x X - [0. 1]. T(X|X, b) describes the probability of being in next abstract state X', when starting from X, executing behaviour b, and waiting until an abstract transition occurs. Figure provides an example of an abstract transition and associated low level steps. The AB-MDP can be interpreted via the Options Framework (Sutton et al., 1999; Precup, 2000), which we discuss in Appendix B.I.I.
There is an intuitive relationship between the competence of a behaviour and the transition prob- ability function T. For instance, executing an incompetent behaviour from abstract state X will likely lead to the next abstract state being itself, X X,, or some hard-to-predicc distribution over next states. Executing competent behaviour lead to a predictable next state that contains the proposed attribute change with high probability.
any case, we can treat the Ab-MDP as a regular MDP use standard RL algorithms to solve This can be viewed solving a proxy problem of the low level base MDP. For the remainder of this paper, methods are trained exclusively within the Ab-MDP unless stated otherwise.
Item Ids
Item Attributes
id
attr
ABSTRACTED MDP
id
attr
stairs
player
standing on
potion
inventory
Symbols
Map
Map
Map
Map
abstract by behavi-
bT.
b
bb policy
01
022
01 primitive
action
"Welcome
NetHack"
LOW LEVEL MDP
"You see here potion"
An example state transition in an Ab-MDP (here defined MiniHack). Abstract states are sets of (item identity, item attribute), and behaviour by have structure (item, new attribute). An abstract state can correspond to multiple low level states, and an abstract behaviour multiple primi- tive actions. We provide legends for the item identities and attributes illustrated (left rectangles).
Figure 1:
METHODS
While any standard RL methods can be used to solve the Ab-MDP, develop method to make better use the item-attribute structure and item-perte cem--erturbing behaviours. Specifically, we design a model-based method that efficiently explores and learn accurate full world model. Our method, Model-based Exploration of abstracted Attribute Dynamics (MEAD) can be broken down into three components: (i) A forward model fo (X, b). (ii) A reward function r r(X, b), (iii) A planner which maximizes the reward function.
Our forward model takes in the current abstract state X and behaviour b, and predicts the probability of success of the change proposed by the abstract behaviour. We describe the procedure for training the forward model in Section 3.1.
We use a different reward function and planner at training versus inference time. During training, the agent learns an accurate (abstract) world model without rewards. We function Tintr (X, b) that encourages infrequently visited states, and plan using Monte-Carlo Tree Search (MCTS) toward these states. Section 3.2 describes the exploration phase in detail.
intrinsic reward
At inference time, our aim is to reach a particular abstract state. We define our reward function rgoaal as 1 for the goal abstract states and for all others. Because our forward model has been trained to accurately estimate state transitions, we simply run Dijkstra's 's algorithm using fo and rgaal to maximize expected reward. Section 3.3 describes this.
3

---

Published as a conference paper at ICLR 2025
3.1 FORWARD MODEL
A forward model takes in a current abstract state and behaviour to model the distribution over next states, X p(X(+1|X, b,). Typically, this distribution has support over the full state space X. We observe that by construction, Ab-MDP considers competent policies as ones that change single item-atteet -tt--atributes. Thus, we assume after an abstract transition, either (i) single item's 's attribute changes according to by, or (ii) the abstract state is unchanged. It is of course possible for multiple items' attributes to change at once, or for an unexpected item's 's attribute to change- we show with a competent set of policies and re-planning
model is robust to this in practice in Appendix B.2.2.
Discriminative modelling Given an abstract transition (Xe, be, 1), we consider this tran- sition successful if the next abstract state con- tains the attribute change proposed by the be- haviour: (a()) {') Xe+1 for by (a - Executing a competent policy has a high prob- ability of leading to successful transition.
next st probabitty
X
fo
WF
P
b
X
We model this success probability directly, ((,,,b)) = Pr(a' E Xx+1(X,,bb),
The forward model fo predicts the prob- ability p that the behaviour b, is successful from state X.. The next state distribution is modelled as binary distribution (Equation 3).
Figure 2:
(1)
where x' = (a'))"))
We refer to this way of modelling discrimi- native world modelling, which models the probability that the state change specified by a behaviour successful. As behaviours are by construction item perturbations, we can view this as asking question about allowable item changes ("can change this item to have this attribute?"). This is different from generative world modelling which models Pr(X [+X Xx by) without any restrictions on possible future states (i.e. "what are all possible outcomes of my action in this state"). We find this inductive bias leads to more data efficient learning and multi-step planning (Section 4.4.1).
Forward Prediction
We can model the next state after successful transition as,
X b( = X (z(0)) U (1 . for any b = (a ..)).)))..
Then, given success probability function q (Equation 1), model the next state distribution as: Pr(X++1XX,,b) b.) [(X,, with probability g(X,, bz), Xx
(3)
with probability 1 - (Xe, b.).
See Figure for an example. As mentioned above, our model robust to cases where multiple items change through re-plamning (Appendix B.2.2).
Model Learning We use model function class fo : x B + [o, 1] with trainable parameters 0, and fit to approximate success probability (Equation 1). We can estimate the empirical success probability from a dataset of trajectories simply by counting, and fit a binary cross entropy loss. We found this to be a stable, simple-to-optimize objective. More details in Appendix C.I.
3.2 PLANNING FOR EXPLORATION
We want the agent to explore indefinitely to discover how the world works and represent it in its model fe. We use a count-based intrinsic reward introduced Zhang et al. (2018) which encourages the agent to uniformly cover all nate-behhaviouss (details in Appendix C.3):
Tintr (X,)) b) = VT/ (N(X,b) +eT),
(4)
where N(X, b) is a count of the number of times (X, b) is visited, T' is the total number of (abstract) time-steps, and 0.001 a smoothing constant.
We use Monte-Carlo Tree Search (MCTS) to find behaviours to reach states that maximizes Equa- tion 4. This is done within the partially learned model's imagination and allow the agent to plan toward novel states multiple steps in the future. The agent takes the first behaviour proposed. Fur- ther, since behaviours do not always succeed, one needs multiple state-visitations to estimate q(X X, b) well. We found MCTS exploration discovers more novel transition than exploring randomly (Sec- tion 4.4.2).
We use set notations to denote removal the item x ()) = () from the set X, and replacing it with item (0) g').
4

---

Published as a conference paper at ICLR 2025
bg 0.982
by
0.967
X
0.90
Model
0,982
tarsitinn
0.001
transiionn
(a) Task agnostic training using MCTS and based reward explore and learn the world model
(h) Dijkstra planning eval time maximize success probability of reaching any goal state
Figure 3: Planning within model's imagination to both explore, and to reach any goal state.
3.3 PLANNING FOR GOAL
Once the agent has sufficiently explored the environment and well-fitted, the model fo usually con- tains a small set of high success probability edges between abstract states (see Figure 3b; most item- attribute changes are generally impossible). We can use this model fo to plan to reach any desired (abstract) world state. As our model function that tells it when it is in a desired state: Tgoal (X) 1 if X is goal state, and 0 otherwise.
trained to maximize external rewards, we simply give it
We use Dijkstra's algorithm which finds the shortest path in weighted graph. We turn our world model into a graph weighted by success probabilities, using a negative log transform:
d(X, ((X,( b)) = log fo (X,b), with A(X, b) defined in Equation 2.
(5)
Thus, finding the shortest path between current and goal state returns an abstract behaviour se- quences that maximizes the probability of successfully reaching the goal.
RESULTS
We use two sets of of environment. One 2D crafting, which requires agent craft objects follow- ing a Minecraft-like crafting tree with each crafting stage requiring multiple low level actions (Chen et al., 2020). Second is MiniHack (Samvelyan et al., 2021), which contains a series of extremely difficult exploration games (Henaff al., 2022). We detail the environments further in Appendix D.
For baselines, we use strong model-free (asynch PPO, (Schulman et al., 2017: Petrenko et al., 2020)) and model-based (Dreamer-v3 (Hafner et al., 2023; Sekar et al., 2020) and MuZero (Schrittwieser et al., 2020; Werner Duvaud, 2019)) baselines. Dreamer-v3 learns a generative world model and plans using model-free method in imagination, while MuZero learns generative world model and plans using MCTS. We outline different baseline methods in greater details and do additional com- parisons in Appendix E. All the methods are trained in the Ab-MDP. We report low level steps and plot the 95% confidence interval. See Appendix G for other experimental details.
4.1 LEARNING FROM SCRATCH SINGLE EN ENVIRONMENTS
Figure shows the from-scratch Ab-MDP training results for strong basclines and our method. We see that Dreamer-v3, a of-ee-of-theaar model-based method, is highly sample efficient. MuZero, another model-based method utilizing online planning, is less efficient. PPO, model-free method, stably optimizes the reward, but at data budget of an additional two orders of magnitude (hun- dreds of thousands versus 10 million). Our method learns the Ab-MDP with an additional order of magnitude sample efficiency improvement compared the already efficient Dreamer. We also refer the reader to Figure 3b (which is clean illustration of the same data presented in Appendix F-6) for an example of the interpretable world model we can extract from MEAD.
For reference, plot the performance for SOTA algorithms in the low level MDP (without Ab- MDP) as triangles and stars. particular in MiniHack tasks (Fig.4b), model-free methods with random and global exploration bonuses (IMPALA, RND, ICM) perform poorly even after 50 million steps, and methods using episodic exploration bonus perform better (E3B). Nonetheless, all of them are less efficient than the least efficient method trained on the Ab-MDP.
These results show that Ab-MDP helps simplify the problem, with all methods reaching reasonable performances-2 feat not achievable by most methods trained in the low level MiniHack MDP
5

---

Published as a conference paper at ICLR 2025
Levitale-Bools
Freez-Hoom
Moan Epsadd Rotum
Wean Titodde Retum
1.00
8.73
0.75
6:2
0.00
101
307
301
10
100
10'
Model-Bessed (ours)
Freeze-Wand
Freeze Randdm-anndm
Model-Based
Nowe (piedde Retur
3,5
10
(Dreamer-3)
8.75
(ours)
Model-Based (Dresmenva)
Moan Episode
0.3
Model-Based
(MuZero)
62-
Modal-Fre
Model-free (PPO)
E38 (low levell
Model free, Low level (PPO)
190
1001 32 20
IMPALA (low
30
201 10"
10
RND (low levell
Erv Frammes Scalel
Iry Fremmes Scale)
Frammes (Log Scile)
(low
(a) 2D crafting games
(b) MiniHack skill games
Figure 4: Results in games trained from scratch in Ab-MDP. Triangles and stars denote low-level only methods. For MiniHack, we also show the final performance of state-of-the-ar exploration methods in the low level MDP (star/triaggle), from Henaff et al. (2022). X-axis on log scale.
(Henaff al., 2022). However, also that the problem is still non-trivial, with the model-free baseline still requiring millions of frames to achieve success. Our method reaches the same peak performance the model-free method while using fewer samples than the model-based baseline.
4.2
TRANSFER AND COMPOSITIONS
Mean Epteooee Retum
Model based lours, pretrainingl
Model based lours, count-based pretrain)
0.5
Dreamery33 Inc pretran)
PPO DE SEaTAr
Dreamory3
.6.5
Moan
Icount: based pretrain!
Dreameru3
300000 End Framess
203000
idisagreement pretrain)
Dreamenuj3
20005 40000
Frames
20030
46000
trewand-based pretrain)
pretrain ning ng Treate Rendomi
peetreiigg in Preaee Randomi
(a) Sandbox
(b) Freeze-Random transfer
(c) Compositional environment.
Figure 5: Transfer experiments and a difficult compositional environment. Dotted lines in Sa and 5b are the same curves from Figure 4b, with just the means shown for clarity.
Next, since the Ab-MDP abstraction models object identities and shared attributes, we evaluate the abstract level models can transfer object knowledge to new settings. Here exclusively the MiniHack environments due to their rich repertoire of object types and higher difficulty. Note we each transfer experiments in multiple environments; the result is similar, we report a single illustrative environment in Figure 5 for brevity and provide full results in the Appendix Figure 16.
Sandbox transfer We first place the agent in a "sandbox" environment where all objects are spawned with some probability (Appendix D.2.1). The agent is allowed to interact with the vironment with only an exploration reward (no task-specific reward) for 200k frames. We then fine-tune the agent in each of the environments. For model, continue train So within the new environment. For Dreamer-v3, we transfer only the world model (encoder, decoder and dynam- ies function) and re-initialize the policy and reward functions." Results for all five environments follow similar trends; we show representative example in Figure 5a, with full results in Figure 16.
Surprisingly, Dreamer-v3 trained with both count- and disagreement-bene based intrinsic rewards fails to get any meaningful rewards in the allocated fine-tuning budget (Figure 5a, red & orange curves), although training for longer seems to show a some reward indicating negative transfer (see Ap- lower per-
pendix 1.1.2). Similarly, the model-free PPO policy shows negative transfer-reachinn formance with the same data budget as compared training from scratch (results Appendix F. since x-axis in different order of magnitude). A possible explanation is while the agent observes all object types, the data distribution dominated by seeing more objects on average in the training
Thee point of this experiment is to evaluate world knowledge transfer. Nevertheless, fine-tuning the ploratory reward function and policy also achieves no hetter performance.

---

Published as a conference paper at ICLR 2025
environment than in the evaluation environments, thus presenting a distribution shift. In contrast, our model exhibits good zero-shot performance, and reaches optimal performance efficiently.
To help the Dreamer-v3 baseline further, give it reward function to "use" the items pre-training sandbox environment. This is a kind of privileged information which biases the model representation to better model the evaluation-time task. We observe that this variant of Dreamer- v3 exhibits positive transfer, hinting at Dreamer's ' representation reliance on the reward function. Nevertheless, does not exhibit the same degree of zero-shot or fine-tuning capability as our model.
in the
Transfer across object classes We further evaluate the ability of the models transfer to new object types. For example, would a model learn to use potion better, after it has learned to use a freeze wand, since the abstract level interactions of these objects are similar (the agent can stand on top of these objects, pick it up, and use them)? We take 200k frame checkpoints in the F Freeze Random environment and train them the other environments for an additional 50k frames.
We again show one representative comparison in Figure 5b and present the full eeeviirommenss result in Figure 16b. We see our model exhibits reasonable zero-shot performance and improves further through fine-tuning. Similar to the sandbox environment, we observe that exploration-based training of both PPO (results plotted Appendix F.1.1 due to x-axis being in different order of magnitude) and Dreamer-v3 exhibits negative transfer in our setting, with both methods eventually achieving better rewards but slower than training from scratch. We again pre- trait Dreamer-v3 with privileged reward information that matches the downstream task, but still requires generalization to novel objects. We see that in general this way of training dreamer resulted in positive transfer - learss to reach higher reward with fewer samples compared to training from scratch.
is somewhat surprising that our model transfer so well zero-shot to completely new objects in Levitate-Boots and Potiian. We analyze this further in Appendix F5:: for (item id, item attribute) vectors with unseen ids, we find the model internally cluster them by object attributes, which generalizes from Freeze objects to Levitate objects.
Compositional planning environment Finally, we designed an environment similar the 5 dif- ficult Minihack exploration levels, but with the required planning depth twice as long. This is an even more challenging task where the agent needs to pick up both a freeze and levitate object, make itself levitate (using the levitate object). then cast freezing spell to terminate the episode. Further, the levitate spell must cast after all the objects have been picked up (since the agent cannot pickup objects while levitating), but before the freeze spell is cast. This is challenging environment requiring precise sequences of abstract behaviours.
We observe in Figure 5c that when pre-trained on Freeze objects, model trains more stably in this environment (but no longer shows good zero-shot results), succeeding majority of times. Our from scratch model solves this task with better rewards. Our method also learns highly complex yet still interpretable world graph, with a simplified version shown in Figure 34. contrast, all other methods, with and without pre-training, do not find the correct solution all, but instead stuck at a 0 reward local minima.
4.3
LEARNING OBJECT PERTURBING POLICIES AND OBJECT-CENTRIC MAP
Behaviour Success
Episode Reward
Encoder Accuracy
Episode Reward
Averoge Success Pro
1.0
a..
9.4
0.5
58
0.50
9.3
0.0
Crafted Policies
56-
Learned Neural
02
-0.5
0.00
Encoder
21
Learned
Policies
62
-0.25
Ground Truth Encoder
0.0
5.
-0.50
60000
Model
Low Frames 167
Low
(a) Learning B
(b) Perf. with learned B
(c) Performance with learned object mapping
Figure 6: Performance using learned Ab-MDP components. All error bars denote 95 confidence interval of mean with exception of encoder accuracy which shows I standard deviation.
*PPO and extended Dreamer-v3 curves in Appendix F.I.
Previous environments need minimum of three abstract transitions complete; this least six.
7

---

Published as a conference paper at ICLR 2025
Up to this point have assumed the existence of an object-centric mapping and competent be- haviours (Sec. 2.1), which can be combined with our method for efficient exploration and world model learning. We now demonstrate how one might learn both components.
4.3.1 1 LEARNING LOW LEVEL POLICIES
We observe low level policies can be learned with RL if an object-centric map exists. Specifi- cally, reward the agent for achieving abstract transition successfully: from X,, if executing behaviour b = (a(*) 5') until an abstract transition results in X containing item (a(*) g²), the as- sociated policy Th gets a reward +1 (else-1) -1). We train low level policies to replace all behaviours in the MiniHack single task environments, and observe the success probability can be stably opti- mized with PPO (Fig.6a; note it does not reach as it is averaged over all attempted behaviours, some are for impossible transitions). Solving Ab-MDP with learned behaviours perform similarly (Fig.6b). The reported 2D Crafting results (Section 4.1) already use learned behaviours. Details in Appendix F.2.
4.3.2 LEARNING OBJECT-CENTRE MAP
we wish to also learn the object-centric mapping, we can do so using supervised learning. We generate a dataset of 100k transitions using random policy in MiniHack's Freeze- envi- ronment, and use the hand-crafted mapping to provide ground truth abstract state labels. We train neural net to produce the abstract state X from S. We compare planning performance using the learned object-centric mapping in Figure 6c and observe it reaches similar performance as the ground truth mapping, even though the encoder is not perfect in predicting the abstract states. We discuss alternative way of getting this mapping in Section 6.
4.4 ABLATIONS
0.8
We ablate major design choices below, with additional ab- lations in Appendix F.3.
Epicode Reward
0.4
4.4.1 GENERATIVE Vs. DISCRIMINATIVE MODELS
0.2
A common way of learning a forward model is to model the next state generative-ly, i.e. R441 ~ f(X,, by), with support over the entire state space X. MEAD, we opted to constrain the possible future either be the same the current time-step, or one where there is a single item change (Fig. 2). This has the distdvantage of being less flexible, but the advantage of being casier to model. To isolate only the effect of discriminative modelling, we design an experiment where we fit models to the same expert dataset, with the only difference being the model type (fit either predict success probability in the discriminative case, or to predict the distribution over next item-atrribute sets), while keeping all other variables constant. Indeed, given the same (expert) data budget, discriminative model learns better model for Ab-MDP planning in the lower data regime (Figure 7, experimental details in Appendix F.3.
Discrimin- ative
-0.2
-0.4
Generative
5000
10600
Dataset Size (Transitions)
Figure 7: Planning performance with discriminative vs. generative models
4.4.2 COUNT-BASED EXPLORATION
Numbar Discovered Object-Attribute Trans
We observe using the count-based intrinsic reward (Equa- tion 4) with MCTS was important for exploration. Figure 8 shows count-based exploration discovering many more valid object attribute transitions compared to a randomly exploring agent. Correspondingly. this leads to a model fo that is better at achieving goals.
80-
60-
MCTS (depth=1) MCTS (d (depth=4)
20
Random Exploration
4.4.3 PARAMETRIC VS. NON-PARAMETRIC MODELS
6oóno 120000
Instead of fitting discriminative model success probability, we can directly use the empirical count-based estimate from a dataset (a non parametric model). This is done learn the model in Zhang et al. (2018). We compare this choice in Ap- pendix F.1.3 and show our parametric model has better sam- ple efficiency, demonstrating the generalization benefits of our parametric forward model.
Env Frames
Figure 8: Effect of count-based ex- ploration in the Freeze-Wart en- vironment. MCTS discovers more unique valid (item identity, item tribute, new attribute) transitions.
8

---

Published as a conference paper at ICLR 2025
5 RELATED WORK
Hierarchical RL Our Ab-MDP is a hierarchical state and temporal abstraction formalism and can be interpreted under the Options Framework (Appendix B.I). Within the Options Framework (Sut- ton et al., 1999; Precup, 2000), Option discovery remains an open and important question. Existing methods typically do option discovery in a "bottom-up" fashion by learning structures from their training environments (Bacon et al., 2016; Machado et al., 2017; Ramesh et al., 2019; Machado et al., 2021). In contrast, our Ab-MDP uses "top-down" prior knowledge in the form of object- perception, and define options that perturb objects' attributes.
Planning with abstract representations Abstraction is a fundamental concept (Abel, 2022) and we select a subset of most relevant works. The Ab-MDP can be viewed as more structured MDP formulation in which learning and planning is simplified. This relates to the Factored MDP (Boutilier et al., 1995; 1999; Guestrin et al., 2003; Degris & Sigaud, 2013), which models the con- ditional independence structure between factors in a factorized state space as Bayesian Networks. also relates to OO-MDP (Diuk et al., 2008; Keramati et al., 2018), which describes state as a set of object and object states which interact through relations. We discuss their mathematical relation- ships in Appendix 1.3. Also related is MaxQ and AMDP (Dietterich, 1998; Gopalan et al., 2017), 7), which provides a decomposition of base MDP into multi-level sub-MDP hierarchies to plan more efficiently. Transition dynamics can also be defined as symbolic domain specific language (Mao al.. 2023). Above works require pre-speeiffyng additional interaction rules, functions, and object relations to simplify planning, whereas we only assume the ability to see objects and their attributes, and learn the forward model as a neural network. The CEE-US model of Sancaktar et al. (2022) use ensembles of graph neural networks to exploit object-centric state representation which factorizes into entity and agent information, and show strong model learning and zero-shot generalization ca- pabilities in robot manipulation domain. DAC-MDP (Shrestha et al., 2020) clusters low level states into "core states" to solve by value iteration, with a focus on getting good policies in an offline RL setting rather than abstract world model learning. More distantly related is Veerapaneni et al. (2019) which focus more on learning good object-centric representations rather than effective usage of an abstract level. Nasiriany al. (2019) tackles sub-policy learning and planning with abstraction, focusing on implicitly representing abstractions rather than using semantically meaningful ones.
Zhou et al. (2022); Haramati et al. (2024); Chang al. (2023) are works that leverage object-centric representations with a set of discrete, factorized descriptions-a beneficial inductive bias Ab-MDP shares (i.e. item identities and attributes). In particular, Chang et al. (2023) uses slot attention to learn factorized "entity types" and "entity states", albeit relying on more restrictive assumptions about entity dynamics. Haramati al. (2024) extracts an object-centric representation using Deep Latent Particles, which extracts latent "particles" describing pixel-based objects in terms of their pixel location, size, etc. Both works offer potential ways of learning the object-centric map M. We discuss their relationship with Ab-MDP further in Appendix B.1.3 and compare against the model- learning method from Chang et al. (2023) in Appendix F.1.3.
Most similar is Zhang et al. (2018), which does model-based planning over abstract attributes. We differ in having a more expressive abstraction framework: modelling a flexible set of objects and shareable attributes. rather than a fixed, binary set of attributes only. We also use parametric model to more efficiently learn and is able to transfer object information to new environments. We compare against their non-parametric approach and show better sample efficiency in Section 4.4.3.
RL in semantic spaces As language models have grown in interest and sophistication, there has been an increasing interest in semantic and linguistic abstractions in RL. Early works such as An- dreas et al. (2017) looked at handcrafted abstract spaces such as we do here to show that creating abstractions allows for better learning and exploration. Other works such as Jiang et al. (2019); Chen et al. (2021) specifically use natural language as the abstract space. In others, the language space abstraction used as a way of guiding exploration (Tam et al., 2022; Guo al., 2023).
Model-based RL Our work is related to other work in Model- based Reinforcement Learning (Sut- ton, 1990; Kaiser et al., 2019; Doya et al., 2002). In particular we compare to the Dreamer line of work (Hafner et al., 2019; 2023; Sekar et al., 2020; Mendonca et al., 2021; Hafner et al., 2022). These methods were developed for a pixels setting and employ a two step learning strategy: learn- ing generative forward model to predicts the next states X', then learning policy in imagination. We compare directly to this in our experiments, where we show that our method (which employs a discriminative forward model) learns more quickly and transfers better in our setting.
9

---

Published as a conference paper at ICLR 2025
Exploration in RL Exploration is a well-studied problem in RL. For more extensive literature review, Amin et al. (2021) gives a comprehensive picture. Most relevant to here are count-ban based methods such Raileanu & Rocktäschel (2020); Zhang et al. (2021). Our work takes inspiration from the hard exploration problems from Henaff et al. (2022). There are also many disagreement- based methods including Sekar al. (2020) and Mendonca al. (2021); we compare directly the former for our Dreamer benchmarks with pre-training.
6 DISCUSSION
In this work introduce the Ab-MDP, which uses human understandable object-atrribtte relations to build an abstract MDP. This enables us to learn discriminitive model-based planning method (MEAD) which explores and learns this semantic world model without extrinsic rewards.
Limitations The major limitation of the work is that the construction of Ab-MDP requires the ex- istence of an object-centric mapping. While we show how these components can be learned in Section 4.3, learning the object mapping still requires a labelled dataset, and method to automat- ically discover object-centric representations in the form items and attributes remains an open question. Promisingly, the field as whole is moving in a direction where abstractions are becoming more easily obtainable. For instance, object-centric abstract maps can be acquired using segmenta- tion models (Kirillov et al., 2023) and visual language models (Alayrac et al., 2022). Similarly, the short-horizon (low-level) manipulation policies can be acquired more efficiently through imitation learning. For example, it has been found in robotics that simply using off-the-shelf models for per- ception and skill policies is possible (Liu et al., 2024). Finally, we currently focus only on discrete attributes, meaning the Ab-MDP cannot precisely capture continuously varying quantities (e.g. po- sition and velocity). However, this is only a restriction in the absiract level, and does not impose any fundamental restriction on the kind of learnable policies in the base MDP. other words, Ab-MDP models discrete, abstract relationships present in the base MDP, while low-level policies remain free utilize continuous features and output continuous actions in the base MDP drive attribute changes for the Ab-MDP.
Overall, as foundational models advance, the field is progressively moving toward working with dis- crete, semantically meaningful representations. Concurrently, reinforcement learning provides rig- orous and principled frameworks for designing intelligent agents, encompassing exploration, plan- ning, and hierarchical behaviours. We hope this work demonstrates productive synthesis of these perspectives, and takes a step forward in the development of more intelligent autonomous systems.
ACKNOWLEDGMENTS
This work is supported by ONR MURI #N00014-22-1-2777 ONR #N00014--1-1-27558, and the National Science Foundation under NSF Award 1922658. AGC is supported by the Natural Sciences and Engineering Research Council of Canada (NSERC), PGSD3--59927--20 2021. Cette recherche a été financée par le Conseil de recherches sciences naturelles et génie du Canada (CRSNG), We are grateful for insightful discussions with Rajesh Ranganath. We thank Doina Precup, Arun Ahuja, Ben Evans, Siddhant Haldar, Nikhil Bhattasali, Ulyana Piterbarg, Nikhil Dongyan Lin, and Mandana Samiei for their helpful feedback on earlier drafts of this paper. Finally, we thank the anonymous reviewers for making this work stronger in many ways.
PGSD3-559278-200
REFERENCES
David Abel. A theory of abstraction in reinforcement learning. ArXiv, abs/2203.00397, 2022. URL https://apl www.nticsccooooo....eee. www.oo//oaao... Corpporrpsssi..c 222232139.
Jean- Bean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lene, Arthur Mensch, Katie Millican, Malcolm Reynolds, Roman Ring, Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob Menick, Se- bastian Borgeaud, Andy Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira. Oriol Vinyals, Andrew Zisserman, and Karen Simonyan. Flamingo: vi- sual language model for few-shot learning. ArXiv, abs/2204. 14198, 2022. URL https: //api. semant cachola wwwp:llaaa.rrororooe//ee 2484777744
10

---

Published as a conference paper at ICLR 2025
Susan Amin, Maziar Gomrokchi, Harsh Satija, Herke Hoof, and Doina Precup. A survey of exploration methods reinforcement learning. ArXiv, abs/2109.00157, 2021. URL https: //api. www.sseeaannnnolaaa...e.. www.laar.rrrroooooeroeee wttp:o/rrrsS:::3333333.....
Jacob Andreas, Dan Klein, and Sergey Levine. Modular multitask reinforcement learning with policy sketches. In JCML, pp. 166-175, JMLR. org, 2017.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. ArXiv, tess/600.0011444 2016. URL https://api semant csccol wwwpchllla.rrr///oooooooooe Orgg/CorpusDD 6627476.
Andrew G. Barto and Sridhar Mahadevan. Recent advances in hierarchical reinforcement
learning. Discrete Event Dynamic Systems, 13:41-77, 2003. URL https: //api. commantcschop wtt.:orro/oooott::3BB
Craig Boutilier, Richard Dearden, Moisés Goldszmidt, et. Exploiting structure in policy construc- tion. In I/CAI, volume 14, pp. 1104-1113, 1113, 1995.
Craig Boutilier, Thomas Dean, and Steve Hanks. Decision-theoretic planning: Structural assump- tions and computational leverage. Journal of Artificial Intelligence Research, 11:1-94, 1999.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2019.
Michael Chang, Alyssa I. Dayan, Franziska Meier, Thomas I. Griffiths, Sergey Levine, and Amy Zhang. Neural constraint satisfaction: Hierarchical abstraction for combinatorial generalization in object rearrangement. arXiv preprint arXiv:2303. 2023.
Valerie Chen, Abhinav Kumar Gupta, and Kenneth Marino. Ask your humans: Using human structions to improve generalization in reinforcement learning. ArXiv, abs/201 abs/201 1.00517, 2020. URL https: / api. semantiiinn icccoolaccccccrr.. org/ Corpus 226226934.
Valerie Chen, Abhinav Gupta, and Kenneth Marino. Ask your humans: Using human instructions to improve generalization in reinforcement learning. ICLR, 2021.
Thomas Degris and Olivier Sigaud. Factored markov decision processes. Markov Decision Pro- cesses Artificial Intelligence, pp. 126, 2013.
Thomas G. Dictterich. The maxq method for hierarchical reinforcement learning. In International Conference on Machine Learning, 1998. URL https://api. wwmanticsccooaa........c.c.c.... oorgg IerpusID:0568555
Carlos Diuk, Andre Cohen, and Michael L. Littman. An object-oriented representation for efficient
reinforcement learning. In International Conference on Machine Learning, 2008. URL https:
//api. semant irtnmnanticccolaa.....
207168200.
Kenji Doya, Kazuyuki Samejima, Ken ichi Katagiri, and Mitsuo Kawato. Multiple model-based
reinforcement learning. Neural Computation, 14:1347-1369, 2002. URL htttps://ap. :1458994444.3
semanticscholes wrw.a.rorr..ooooooe...eee.oooer. Cht///orppssDDD it
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. International conference on machine learning. pp. 1407-1416. PMLR, 2018.
Nakul Gopalan, Marie desJardins, Michael L. Littman, James MacGlashan, S. Squire, Stefanie Tellex, John Winder, and Lawson L. S. Wong. Planning with abstract markov decision pro- cesses. In International Conference on Automated Planning and Scheduling, 2017. URL https://aph wem.nticsccoo ar org/ Cotp otr.ei/ccsooo.o.oooo/oooooooo
25501330.
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algo- rithms for factored mdps. Journal of Artificial Intelligence Research, 19:399-461 2003.

---

Published as a conference paper at ICLR 2025
Zhourui Guo, Meng Yao, Yang Yu, and Qiyue Yin. Improve the efficiency of deep reinforcement learning through semantic exploration guided by natural language. ArXiv, 0bss2309.11155. 2023. URL https: /app...../// semant
org/ Ootg/CorpusI......oooooo 2 62084426.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving horizon tasks via imitation and reinforcement learning. arXiv preprint 0rXiv:1910.11956. 2019.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to con- trol: Learning behaviors by latent imagination. ArXiv, abs/1912.01603, 2019. URL https:
//api. aeppi.immmniicccolaaa org/ Caapg//orpuusDi 208547755.
Danijar Hafner, Kuang-Huei Lee, Ian S. Fischer, and P. Abbeel. Deep hierarchical planning from org/
pixels. ArXiv, 0bs/2206.04114, 2022. URL https:/ /api. CorpusID:249538577
Danijar Hafner, J. Pašukonis, Jimmy Ba, and Timothy P. Lillicrap.
Mastering diverse do-
mains through world models. ArXiv, 0a/2301.04100 2023. URL setps://aphe
semant orrcchosl...rr//ooosssssis
Dan Haramati, Tal Daniel, and Aviv Tamar. Entity-centric reinforcement learning for object manip- ulation from pixels. arXiv preprint arXiv:2404.01220, 2024.
Mikael Henaff, Roberta Raileanu, Minqi Jiang, and Tim Rocktäschel. Exploration via elliptical episodic bonuses. Advances in Neural Information Processing Systems, 35:37631-37644 2022.
Yiding Jiang, Shixiang Shane Gu, Kevin P Murphy, and Chelsea Finn. Language as an abstraction for hierarchical deep reinforcement learning. In NeurIPS, pp. 9414-9426, 2019.
Lukasz Kaiser, Mohammad Babacizadch. Piotr Milos, Blazej Osinski, Roy H. Campbell, K. Czechowski, D. Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiud- din, Ryan Sepassi, G. Tucker, and Henryk Michalewski. Model-based reinforcement learning for atari. ArXiv, abs/1903.00374, 2019. URL https: /aap://appaaa.. CorpusID: 62856232.
Ramtin Keramati, Jay Whang, Patrick Cho, and Emma Brunskill. Fast exploration with simplified models and approximately optimistic planning in model based reinforcement learning. arXiv: Artificial Intelligence, 2018. URL https: http:///p......sssseiseeooe semanticscho..o org/ CorpusID: 53781531.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross B. Girshick. Segment anything- sion (ICCV), pp. 3992-4003,
2023 IEEE/CVF International Conference on Computer Vi- URL https: /ttp://app.... www.nticscoolar.co.ooooo
2023.
Heinrich Küttler, Nantas Nardelli, Thibaut Lavril, Marco Selvatici, Viswanath Sivakumar, Tim Rocktäschel, and Edward Grefenstette. TorchBeast: PyTorch Platform for Dis- tributed RL. arXiv preprint arXiv:1910003552, 2019. URL https: httts://ggthhbs com/ Facebokresearch/toreee. wwwpseearrrccctttbbeesssee
Heinrich Küttler, Nantas Nardelli, Alexander H. Miller, Roberta Raileanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The NetHack Learning Environment. In Proceedings of the Conference on Neural Information Processing Systems (NeurIPS), 2020.
Seungjae Lee, Yibin Wang, Haritheja Etukuru, H Jin Kim, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Behavior generation with latent actions. arXiv preprint arXiv:2403.0318/ 2024.
Peigi Liu, Yaswanth Orru, Chris Paxton, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Ok- robot: What really matters in integrating open- knonllnoweedge models for robotics. arXiv preprint arXiv:2401. 12202, 2024.
12

---

Published as a conference paper at ICLR 2025
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centrie learning with slot atten- tion. Advances in neural information processing systems, 53:11525-11538 2020.
Marlos C. Machado, Marc G. Bellemare, and Michael Bowling. A laplacian framework for option discovery in reinforcement learning. In International Conference on Machine Learning. 2017.
URL https:://a/pp. http:////aap..ssessscsscsccoo
org/ htte//CrpusDDD: 7795421.
Marlos C. Machado, André Barreto, and Doina Precup. Temporal abstraction in reinforcement learning with the successor representation. J. Mach. Learn. Res., 24:80:1-80:69, 2021. URL
atttps:/aape www.stecsannnnooo.eeeeee http:s///ppp.seenniiccchaaa org/ :tt.:/oo......o..ooooosaa:a::::.
Jiayuan Mao, Tomas Lozano-Perez, Joshua B. Tenenbaum, and Leslie Pack Kaelbling. Pdsketch:
Integrated planning domain programming and learning. In NewrIPS, volume Abs/2203.05501, 2023. URL httph:t///// http:////pp.smmmtttccsciiiie www.nticschooeeeeeeee http://cccccccalooooooooreee www.://osoo/oooo.::..... :
Russell Mendonca, Oleh Rybkin, Kostas Danilidis, Danijar Hafner, and Deepak Pathak. Dis- covering and achieving goals via world models. ArXiv, abs/2110.09514, 2021. URL https: www.ssssmaaticcieieeeeeee
org/ wwwpoooppg..../.....o. 239016113911444
Soroush Nasiriany, Vitchyr H. Pong, Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. ArXiv, 0ass/9911.08455 2019. URL https: /tts:///pppp /api. semanticscholl org/
CorpusID:20277300
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. 2778-2787. PMLR, 2017.
Jeffrey Pennington, Richard Socher, and Christopher D. Manning. GloVe: Global vectors for word representation. In EMNLP, 2014.
Aleksei Petrenko, Zhehui Huang, Tushar Kumar, Gaurav S. Sukhatme, and Vladlen Koltun. Sam- ple factory: Egocentric 3d control from pixels at 100000 FPS with asynchronous reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Re- search, pp. 7652-766 7662. PMLR, 2020. URL http:/ htttt::/pproceeiine mlr. prwss//isss/////////e/. petrenko20a.htie
Doina Precup. Temporal abstraction in reinforcement learning. Amherst, 2000.
University of Massachusetts
Martin L. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 1994.
Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for procedurally-generatee environments. ArXiv, abs/2002. 12292, 2020. URL. https://aag
cemanticschen ar. wwt.ssccchccc.c.aeeeeeeoeeeeeeee :o
Rahul Ramesh, Manan Tomar, and Balaraman Ravindran. Successor options: An option discovery
framework for reinforcement learning. ArXiv, abs/1905.0574. 2019. URL https://api:
gemanticscholadt wttp ID: 1D:31111277.
Mikayel Samvelyan, Robert Kirk, Vitaly Kurin, Jack Parker-Holder, Minqi Jiang, Eric Hambro, Fabio Petroni, Heinrich Kuttler, Edward Grefenstette, and Tim Rocktäschel. Minihack the planet: A sandbox for open-ended reinforcement learning research. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021. URL https: / //openzeview. net. f for corummm f www.dddssswwweeeem
Cansu Sancaktar, Sebastian Blaes, and Georg Martius. Curious exploration via structured world models yields zero-shot object manipulation. Advances in Neural Information Processing Sys- fems, 35:24170-24183 2022.
13

---

Published as a conference paper at ICLR 2025
Julian Schrittweser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-60 609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proxi- mal policy optimization algorithms. ArXiv, abs/1707.00347, 2017. URL https://api.
semanticschol org/ azg/CorpusID- 28695052.
Ramanan Sekar, Oleh Rybkin, Kostas Danillidis, P. Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International Conference on Machine Leurning, 2020. URL https::/aap..eemaaae.e..e wemanticschoole wtt.:oorg///oooooo 216559511.
Aayam Shrestha, Stefan Lee, Prasad Tadepalli, and Alan Fern. Deepaveragers: Offlie reinforce- ment learning by solving derived non-parametric mdps. ArXiv, abs/2010.08891, 2020. URL
https: /appi. semanticscholo ar. org/CorpusIDD :2:20003::222
Richard S. Sutton. Dyna, an integrated architecture for learning. planning, and reacting. SIGART Bull., 2:160-163, 1990. URL https: /app. semanticscho.. org/CorpusID:
207162288.
Richard S. Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: frame- work for temporal abstraction in reinforcement learning. Artif. Intell., 112:181-211, 1999. URL
https: / //pp.... semant wesschclsch.rr. org/Corpus 7
Allison C. Tam, Neil C. Rabinowitz, Andrew Kyle Lampinen, Nicholas A. Roy, Stephanie C. Y. Chan, DJ Strouse, Jane X. Wang, Andrea Banino, and Felix Hill. Semantic exploration from language abstractions and pretrained representations. ArXiv, abs/2204.05080, 2022. URL aittps://ap. scschoic www.or///oooooo Ctrppo :
Vaswani. Attention all you need. Advances in Neural Information Processing Systems, 2017.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua B. Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. ArXiv, abs/1910. 12827, 2019. URL. https://///..c...tt..ttt /api. appi... semant cschola a .orgg
:orpusID: 200Mm004:244488555
Aurèle Hainaut Werner Duvaud. Muzero general: Open reimplementation of muzero. https: wtt.uuu..oo//eeeeeddrreraee
2019.
Amy Zhang, Sainbayar Sukhbaatar, Adam Lerer, Arthur Szlam, and Rob Fergus. Composable planning with attributes. In International Conference on Machine Learning, pp. 5842-5851.
PMLR, 2018.
Tianjun Zhang, Huazhe Xu, Xiaolong Wang, Yi Wu, Kurt Keutzer, Joseph E. Gonzalez, and Yuan- dong Tian. Noveld: A simple yet effective exploration criterion. In Neural Information Processing Systems, 2021. URL https::///aat wwwanticscholaattttooooecccccaae org Corpps 245877021.
Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for com- positional generalization in control. arXiv preprint arXiv:2203.05966 2022.
14

---

Published as a conference paper at ICLR 2025
A BROADER IMPACT
Our contributions are fundamental reinforcement learning algorithms, and we hope our work will contribute to the goal of developing generally intelligent systems. However, do not focus applications in this work, and substantial additional work will be required to apply our methods to real-world settings.
Regarding compute resources, we use an internal clusters with Nvidia A100 and H100 GPUs. All experiments use most one GPU and for more than days. Running baselines along with randomization of seeds requires multiple GPUs at once.
B PROBLEM SETTING DETAILS
B. ABSTRACTED MDP
Here we re-iterate our framework and discuss it in greater details. We first define a low level (reward- free) Markov Decision Process (MDP, Puterman (1994)) as the tuple (S, A, P), with (low level) state space S, primitive action space A, and transition probability function P 8 x A x S - [0, 1]. Suppose the agent interacts with the environments at discrete time-steps u = 0, 1, 2, .... then P(su+1|$u a.) deseribes the probability of transitioning to state Su+1 E S after one time-step when choosing action Qu E A in the current low level state Su-
To construct an Ab-MDP, level: M : S > X. That is, any low level state S E S maps onto some abstract state X E X (multiple S can and do map the same X).
start with surjective deterministic mapping M from low to abstract
Define behaviours B. Each b (a!9, E B describes single item id (1) and desired new attribute, along with an associated low level policy nr : S + A which tries to set the specified item to have the attribute. A behaviour can be competent incompetent. Executing competent behaviours lead to the specified diem-attribll changes with high probability.
A behaviour may be incompetent because (i) the proposed attribute change is impossible within the rules of the world (e.g. item is not a craft-able object) (ii) the attribute change is not possible from the current state (e.g. do not have the correct raw ingredients in inventory), or (iii) the behaviour is not sufficiently trained to perform the specified change.
We build simple abstraction that allow us learn purely at the abstract level while still being able to influence actions at the low level. Notice the surjective map M give us state abstraction. We build temporal abstraction by defining abstract state transitions as happening at a slower time-scale than the low level time-step u.
Definition B.1. An abstract state transition has occurred (low level) time-step if: (i) the abstract state has changed, M(ss) ) * M(su-1): or (ii) the abstract state has not changed in the past steps, M(ss.) = M(ss-1)= = ... ((s-kk).
Definition B.2. An (max-k) Ab-MDP is defined as the tuple (N, Tx), with abstract states X abstract behaviours B. and transition probability function Tr X' x B X [0, 1]. (X' b) describes the probability of being in abstract state X' when starting from X and executing behaviour b until an abstract state transition occurs (Definition B.
In the main text, we write Ab-MDP and denote the transition function as T for brevity (instead of writing max-k Ab-MDP and transition function T&). & was chosen for baseline methods as worked slightly better. MEAD works just as well with k = 8 and l: = 16.
B.1.1 INTERPRETATION VIA OPTIONS FRAMEWORK
natural to interpret our hicrarchical framework within the Options framework (Sutton et al., 1999). An option consists of policy and an initiation set of states I C S. used to denote the set of all possible historical states.
Ω x A - [0, 1), a termination function B Ω [0. 1),
Our set of behaviour b B is set of options. Each behaviour specifies behaviour policy To. which deterministically terminate when an abstract state transition occurs (Definition B. and they can be initiated from all states (I = S).
15

---

Published as a conference paper at ICLR 2025
The point of the Options Framework and temporal abstraction generally is to break up a monolithic problem into re-usable sub-problems. How one breaks up a problem is an open question. Some popular previous approaches include finding bottleneck states and using the successor representation to discover well-connected states. We can view our approach as another way of breaking up the problem when prior knowledge in the form of objects is available: we break up the problem the level of single item-attribute changes (with options defined as policies that can change a single item's 's attribute).
B.1.2 A NOTE ON SEMI-MDI MDP
The decision process that selects amongst options is a semi-MDP (Theorem 1, (Sutton et al., 1999)). The general way to model semi-MDP dynamic requires considering the amount of low level timesteps that passes between each abstract level decision 1999; Barto & Mahadevan, 2003). Instead, we chose to simply treat the abstracted level as having fixed time interval in between steps. This is chosen for simplicity as our focus is with learning efficiency on the abstract level, and this formulation allow us to directly apply efficient algorithms developed solve MDPs. We also see empirically that given our M and B this choice does not lead to less optimal policy compared to algorithms that directly solve the low level MDP.
random variable (Sutton et al.,
B.1.3 RELATIONSHIP TO OTHER MDP FORMALISM
Our Ab-MDP related to Factored MDPs (Boutilier et al., 1995; 1999; Degris & Sigaud, 2013). Factored MDPs model states as decomposable into a set of factors X, {xx(1,, .x(2),... ...,, with each factor having domain ( This allows the transition probability function to be decomposed into Dynamic Bayesian Networks (DBNs) between the factors at two successive time-steps t 11 and t, with the factors as the nodes of the DBN. Concretely, this decomposes the transition probability function for each action into product of conditional probabilities,
Pa (x((x(X111)
II Pa(at Parent((r)) (3)
where (f)) denotes the factors in X,-1 which are the parent of x ( In the most generic form, each action a E A induces a different DBN. If the DBN is known or well-approximated, the Factored MDP can be efficiently solved (Guestrin al., 2003).
Similarly, the Ab-MDP considers a decomposition of state into items (analogous to factors), and our method MEAD shares the philosophy of efficient utilization of a factorized state space. While one could interpret Ab-MDP as kind of Factored MDP with discrete domains, the two differs in their emphasis on transition vs. actions. Factored MDPs defines arbitrary actions, each inducing a potentially different conditional independence structure in the transition dynamics. The Ab-MDP instead defines a set of next state changes as the actions, without explicitly specifying any inde- pendence structures in its transitions dynamics. Our method MEAD then makes use of this action representation to directly model the probability of item-attribus changes, while implicitly learning the transition dynamics in the weights of the neural network.
Our Ab-MDP also related to the Object-oriented MDP (00-MDP) of Diuk et al. (2008), and shares similar structure of deseribing set of "objects" and their attributes separately. In O0- MDP, an environment consists of a set of objects O = {or..... ...,, each belonging to one of c classes {Ci,...: C.), and each having set of (class) altributes Att(C) {C.a1, .C.a,,...CC. (C.a,,...CC...) An OO-MDP "object state" -otte"-ostate-iss a value assignment to all of its attributes (fa), H, and the O0-MDP "state" is the union of all of its object states:
14=10t-ttate
Our Ab-MDP similarly contains a set of items, X = (2x(1) ..... ...,,,(())) (N)), with each item consisting of an item identity and an attribute, (0 = (1) g()). We can view our Ab-MDP a soft generalization of O0-MDP by modelling:
The OO-MDP object (and its associated class) as an Ab-MDP item identity; ( = Os, The OO-MDP attribute the Ab-MDP item attribute: g()) (C.a),-., Og C.
To model state transitions, OO-MDP makes additional assumptions that objects interact through their class identities. Specifically, set of relations determine if two objects interact, with resulting
16

---

Published as a conference paper at ICLR 2025
effects that is used to update the object state. This requires the designer to additionally provide set of accurate relations functions and effects, with the benefit of facilitating within generalization. In contrast, Ab-MDP makes additional assumption about object class, relations, or interactions between objects. We instead directly models (probabilistically) whether an effect is possible. This is done in fully data-driven way through neural function f X' x B [0, 1] and empirically is shown to generalize well across item-attribute variations. Item interactions are implicitly represented by the attention mechanism over the set of input items.
The neural constraint satisfaction (NCS) approach of Chang et al. (2023) introduces an entity- sel state representation, where each entity h(6) is decomposable into type 209) and entity state g(o) Our Ab-MDP can naturally be mapped onto their entity-state abstraction by setting
Ab-MDP item identity as the NCS type (a = 2 and the Ab-MDP item attribute as the NCS entity state (g(i) ).))) Note the Ab-MDP "behaviour" (i.e. abstract action) is equivalent to the NCS "action".
NCS makes additional assumptions that attribute change is dependent only the behaviour: c(t)
St+1 /(5(") .b). On the other hand, Ab-MDP models more general relationships e(1)
be) (e.g. Equation 1). The NCS assumption is restrictive in that it would fail to capture cases where an item's 's attribute change depends what the item's identity is, and/or the attributes of other items.
f(a, 55 a....a, ....
(1) (1)
(M) <(M)
The Deep Latent Particles (DLP) approach used in Haramati al. (2024) learns latent space containing set of K particles:
{(2p, 2s. 2d, 21, (4224), 2f); ) K 2p R describes the particle's (x.y) position on an image, 2, E R² describes its (x.y) scale, za E R describe its "depth". 2, R its transparency, and 2 E R' the visual features around the particle. We can cast this into the Ab-MDP by having K items whose item identities is the particle's index, ( = i, and item attribute a 6 I- dimensional vector describing the particle features; c(i) = (Zp: Ræs Ed; zt; zf); DLP can be one way of getting abstraction from pixels, albeit it is restricted to settings where the state is image-based.
1.1.4 EXAMPLE APPLICATIONS
We provide further examples possible Ab-MDP constructions for few common settings. Note we provide one example abstraction per domain, though multiple types of Ab-MDP abstractions can theory be constructed depending on the choice of the item set and attribute set.
Environment
Description
Example abstract states
How to get abstract behaviours
Kitchen (Gupta
A robot arm is in a MuJoCo kitchen with multiple things to do such as turning on lights, turning on ket- tle, opening oven / mi- crowave / cabinet, etc. A robot arm pushes cubes on table their desired positions
Item set: (light, oven, mi- crowave, cabinet). At- tribute set: {open, closed]
Imitating expert play data in Gupta et (2019) which already produce trajectories that lead to abstract state changes
et al., 2019)
Cube pushing (Haramati et al.. 2024)
Item set: [red cube, blue cube, green cube, ...). At- tribute set: {at goal posi- tion, not goal position) Item set: [drawer, bag, oven, box, bread, can}. At- tribute set: {open, closed, on table, in bag, in drawer, in oven)
RL to push a single block goal at time
Real world robot (Lee et al., 2024)
A Stretch robot in a kitchen, moving ob- jects in and out of con- tainers
Human imitation learning data con- taining trajectories between abstract state changes
Table I: Example Ab-MDP's applied to various domains
17

---

Published as a conference paper at ICLR 2025
B.1.5 MULTI-ITEM BEHAVIOURS
In the Ab-MDP, we consider abstract states as an (ordered) set of items X = (x(11, xx(2) ,..... where each item consist of identity and attribute x(6i (a (())). A behaviour is description of which item(s) to change (Section 2.1). For instance, the single-item behaviour = {(ol)) {')) specifies a change of the i-th item to a new attribute E'. We can similarly specify multi-item be- haviours, such as a two-item behaviour = ((()))")) (a),) (a))) This behaviour is competent after an abstract state change, the i-th item takes attribute ç and the j-th item takes attribute s" with high probability.
For N objects, attributes, and 7 items changes consider in behaviours,
Number of behaviours = N x m'
For instance, for single item change behaviours (7 = 1). there are (1) x m = x m behaviours. For two-item change behaviours (I 2), there are (2) m behaviours, and In the ex- treme, one can choose model a all possible item attribute changes N), which results in m behaviours that is exponential in the number of items considered.
B.2 ANALYSIS OF ASSUMPTIONS
B.2.1 THE GENERALITY OF AB-MDP
We make the observation that any sparse reward RL task can be described by a Ab-MDP with single item,
Remark B.3, For any MDP with a reward function R: S + R. where,
R(S) 1 if S G and otherwise,
for S E 9 terminal goal state(s), we can deseribe this as an Ab-MDP with a single item a and binary attributes (E1,52). The abstract map M S > X can be constructed in the following way,
M(S) = [((a, if
(8)
(1), otherwise.
The proxy problem of transition from X = [((a(( {1) to x= ((a, E2)} solves the sparse reward task the base MDP and a competent behaviour for this transition is good policy is the base task.
The point of this is merely say that Ab-MDP does not restrict the RL problem in particular way- Of course, such a construction an Ab-MDP is not interesting and learning a competent behaviour for this Ab-MDP is easier than learning behaviour that maximizes the reward function in the base MDP.
The point of Ab-MDP, however, is that task can be broken-up in interesting way, e.g. in case through richer set of items and attributes. We demonstrate in this work that an object-centric item-attribute set encoding result in set of intuitively base behaviours and a world model that interesting.
B.2.2 DISCRIMINATIVE MODELLING WITH SINGLE ITEM CHANGE
In our forward model, we made the assumption to only model single item's attribute change (Sec- tion 3.1 and Equation 3), and that if the attribute does not change then the abstract state X stays the same as X, We analyze this assumption more deeply in this section.
Robustness through re-planning We first note that in practice, it is often the case that multiple item's attribute changes. For instance, in the Crafting environments, crafting an item often con- sumes the raw ingredient used to craft it. Thus, plan that considers multiple items being crafted (attribute change from absent to in inventory) will inaccurately predict that the raw ingredi- ent stays in the inventory (attribute of raw ingredient stays the same, staying as in inventory) rather than changing to absent. We illustrate the plan generate at each abstract step in Figure 9, and show that despite this being the case, the plan still generates optimal action sequence; and that at each time-step. re- ce-planning with the newly obserbved abstract state updates the incorrect abstract prediction be more accurate.
18

---

Published as a conference paper at ICLR 2025
True States
Plan
Plan at
t=1
Plan =
1=2
Figure 9: An illustrative example following the logic of the 2D Crafting game. The objective of the game is to get planks by chopping trees (which changes the plank's ss attribute from absent to in inventory, and the tree's attribute to from in world d to absent as the resource depleted), then craft wooden stairs using the plank (which depletes the plank resources to absent and creates wooden stairs in inventory). Comparing the true state trajectory with the imagined plan at t=0, notice that the imagined future is incorrect: modelling only single item changes do not capture the fact that resources are depleted, but only new resources are created. Nevertheless, at 1=1, the model observes the abstract state and can therefore re-plan with the most up-to-date information which corrects for previously incorrect predictions. A case like this happens in practice within the 2D Crafting games.
An adversarial failure Here artificially design adversarial setting which is not countered in our evaluation environments, but would in theory result in sub-optimal plans. With reference Figure 10a, construct an environment whose goal is construct both wooden stair and a wooden door. Both require a single unit of wooden plank to craft. The attribute here is the quantity of the item in the inventory. Suppose further that once we start crafting, we can no longer collect additional wooden planks.
The correct sequence of behaviour would first collect another unit of wood, before crafting wooden stairs and wooden door. Planning with discriminative single item change model would result in sub-optimal plan, as it fails consider that crafting a wooden stair result in the wooden plank being consumed, and therefore cannot subsequently craft the wooden door (Fig. 10a, plan at t=)).
Note that this failure case is limitation of doing discriminative modelling with single item changes (Section 3.1), not the Ab-MDP framework. Doing generative modelling (considering all possible X X' given X,, b,) would be able to learn the correct relationship- at the cost of data ineffi- ciency (Section 4.4.1). 1). Nevertheless, we discuss how to fix it below with discriminative modelling.
Discriminative modelling of multi-object changes One simple way of generating plans with multi-item dependencies over time is to consider multi-item behaviours (Section B.1.5). With ref- erence to Figure 10b, we notice that simply modelling two item changes can perfectly capture the dynamics within this task and generate the optimal plan. This comes at cost of a greater number of behaviours to consider at each step of planning (E.g. for 9 behaviours to consider for single-item changes, and 27 behaviours to consider for two-items changes).
items, m attributes, there
plan to gener- Generally speaking, with small change in the representation of the problem, ate the optimal sequence of abstract behaviours. This can also happen through changing the abstract mapping-es e.g. by modelling items along with the resources needed to create them. Overall, view
19

---

Published as a conference paper at ICLR 2025
Correct Trajectory
B
Taken Trajectory
Goal
Plan (t=0)
Deal trmretion
(auctsso))
Beat Inarection (lullct)
Imagine transtion (sucosss)
Plan (t=1)
Imsegine tranatiion (falldd
(a) Planning with single item-attribute changes fails this adversarial task.
Taken Trajectory
Goal
Plan (t=0)
Plan (t=1)
Plan (t=2)
(b) Planning with double item-attribute changes models attribute dynamics perfectly.
Figure 10: Single and double item- item-attribute change planning
this as a problem constructing the right abstraction for the problem which, although important, is not the main focus of this work and nor fundamental issue with our modelling choice.
C MEAD: EXTENDED METHODS
Here we provide more details about our methods
C.1 MODEL FITTING
Given dataset of (abstract) trajectories collected via the exploration strategy outlined in Section 3.2, D {X1,, dy, X2, bz, ):....== we can estimate the empirical success probability g of all ob- served item-attribute pairs by counting their occurrences:
q(X,)) 6) ~ b) = Number of (Xe, X(+1) transitions where (a()) E X4+1 +e
(9)
Number of (Xe, be) + 2e
20

---

Published as a conference paper at ICLR 2025
where is small smoothing constant (usually 1e-3). This be efficiently implemented as a hash map, via hashing (Xx, by) as keys, and the number of successful and total transitions (starting from (X,,6,) b,)) as values.
To train our discriminative model fo &' x B + [0, 1], we minimize binary cross entropy loss with p the target,
8* arg min E (p(X,b) log (ff(XX,)) (1 p(X, log (1 fa(
(10)
where the expectation estimated by uniformly sampling unique (Xx, b) pairs from dataset.
Appendix C.2 contains architectural details for fa- All in all, found that sampling method and loss function lead to a stable, himple-o-optimizze objective. Moreover, it is possible to directly do planning, is done in Zhang al. (2018). However, show in Section F.1.3 that is more efficient to use our parametric modelling approach as it generalizes better to new objects.
Algorithm Example Data Collection and Training loop for MEAD
Require: Model fg
repeat
Collect n frames of data in environment using model fo to plan to explore
Add data to dataset D
Reset model weights while accuracy requirement not met do if local minima then
Reset model weights 0
end if
Sample from D and minimize Equation 10, for m optimizer steps end while until time runs out
We provide the pseudocode for training MEAD in Algorithm 1. The accuracy requirement running average binary classification accuracy (typically by checking if the difference is accuracy is close to zero (taken as exponentially smoothed average be insensitive enough gradient updates between data collection steps the precise setting of checking for accuracy and local minima does not matter much.
0.95). The local minima is evaluated
stochasticity). We find in Figure 28 that the most important setting is run
C.2 MODEL ARCHITECTURE
de
dy
drah
K
linear
QK' Vdla
IR (logits)
dideo datt
ds
dot prod
key-query-rolus
layor norm
layer applied eiemen-wiiee
projections (linear)
Figure 11: Model Architecture
We a small key-query attention architecture that takes in sets of vectors produce binary prediction (Figure 11). Each item at()) (g(())())) is represented by its item identity and item attribute embeddings. An Ab-MDP state containing a set of m items represented as a m x (daten + dam) matrix. The behaviour = (a (() describes an item identity and the new attribute to try to change to, and is diden dar dimensional vector. The set of items and the behaviour is passed
21

---

Published as a conference paper at ICLR 2025
element- wise through small MLP (specifically 2 layer ReLU network with 128 units in its hidden layer) to produce a set of m x demb item embeddings, and a Memb-dimensione behaviour embedding (demh 128). The item embeddings are linearly projected into key and value matrices K and V both with shape m x dR, and the behaviour embedding is linearly projected into a x dy query vector (dx 256). We compute the scaled dot product attention (Vaswani, 2017):
(QK'
Attention (Q, K, V) softmax
Vdz V
(11)
The output of the attention operation is dg vector. Finally, it is linearly projected to scalar, representing the un-normalized logit for binary prediction.
C.3 EXPLORATION
We reward function that encourages indefinitely exploring all states and behaviours,
T
Tints (X. ((X,6)+ +
(12)
where N(X, b) is count of the number of times (X, b) is visited, T is the total number of abstract time-steps, and = 0.001 is a smoothing constant.
Intuitively, state with many visits (large N(X, b)) will have lower reward, while T prevents the intrinsic reward from shrinking to zero (so the agent explores indefinitely). Indeed, in the bandits setting, this reward function is maximized when all states are visited uniformly (Zhang al., 2018), We emphasize that all of our model training is done with this intrinsic reward, with no task-specific information injected training
time.
We use MCTS to find high (intrinsic) reward states. MCTS iteratively expands a search tree and bal- ances between exploring less-visited states (using an upper confidence tree) and exploiting reward- ing states. 6 In our case, MCTS generates child nodes of abstract state X by proposing behaviours b and checking (XX, b) 0.5 (Figure 3a). We also introduce a small degree of randomness so all behaviours have a small chance of being selected.
We modify the MCTS algorithm to return the the maximum Nintr (X, b) encountered along each path during the back-propagaaaee propagation stage and set fix expansion depth without a simulation stage from the leaf node. This done to prevent loops and encourage the agent to go toward frontier states with the least amount of visitations.
D ENVIRONMENTS
D.1 2D CRAFTING
We adopt three crafting environments proposed in Chen et al. (2020), which contain Minecraft like crafting logic. The agent needs to traverse grid world, collect resources, and build increasingly more complex items by collecting and crafting ingredients (e.g. getting axe, chopping tree to get wood, before being able make wooden planks and furniture that depend on planks). Primitive ac- tions consist of movement in the four cardinal directions (N, E, W), toggle (of switch to open doors), grab (c.g. pickaxe), mine (c.g. ores), craft (at crafting bench of specific items). The low level observation space consist of dictionary containing: (i) an (5. 5, 300) matrix describing 5 x world with each state being described by an 300-dim word embedding. (ii) an 300-dim inven- tory embedding describing things in the inventory, and (iii) 300-dim goal embedding describing in words the task to complete. We describe each environment's 's crafting sequences required to finish each task in Table
'This exploration is purely within the model's 's imagination, rather than the real environment. Balancing exploration and exploitation is still necessary exhaustive search in the abstract space is highly inefficient intractable.
22

---

Published as a conference paper at ICLR 2025
tex: (CMPTT)
Iast (1K
Inc: "TTckass":
Inc (*Plckas"" "Cathlomme"'
les: ('Thckas""
Example episode in example 2D Crafting environment, figure taken from Chen et al. (2020). To solve this task, the agent needs to pickup the pickaxe, mine the cobblestone stash to get cobblestone, then the acquired cobblestone craft cobblestone stairs. The items must be acquired in that order as each step requires having the item from the previous step (with an additional obstacle of being able cross the door only if it has key).
Figure 12:
Environment Name
Crafting Task Description
Steps
Pickup pickaxe, mine gold
Steps
Pickup pickaxe, mine diamond orc, craft diamond boots Pickup axe, chop tree to get wood, craft wooden plank from wood, then craft either wooden stairs wooden door from plank
Steps
Table 2: 2D Crafting environments.
Abstraction States We build simple abstraction to map from the low level observation to abstract representation. We take as items all objects in the world, in the inventory, and potentially craft-able objects. We assign each item one the the attributes in Table 3.
We take each item identity as its 300-dim word embedding provided from the environment, and each attribute as a one-hot embedding, repeated also be 300-dim. Each item in the item set is described as 600-dim vector.
Attribute
Deseription
TN WORLD
An item the world but not in the inventory.
IN. INVENTORY
An item is in the player's inventory.
ABSENT
An item neither in the inventory in the world.
Table 3: 2D Crafting attributes.
D.2 MINIHACK
We also use the MiniHack environment (Samvelyan et al., 2021), which is built on the Nethack Learning Environment (Küttler et al., 2020), We use five hard-exploration games from Henaff et al. (2022) requiring interaction with different types of objects. The base game has multi-modal ob- servation space involving image, text, inventory, etc. Solving the game from the base observation is extremely difficult since each environment has multiple variants of object (e.g. Freeze- Wand has 27 different types of wands), and randomly generated object locations. Indeed, Henaff et al. (2022) found that performant RL methods such as IMPALA (Espeholt et al., 2018; Küttler et al., 2019) and all tested global exploration methods (ICM (Pathak et al., 2017) and RND (Burda et al., 2019)) fail to solve these environments.
The base environment contains dictionary observation space, containing 21 79 grid of glyphs (integer denoting one of 5976 nethack glyph types), an 27-dimensional vector denoting agent's stats, 256 dim vector for message, and another vector denoting agent's inventory. The environments are:
"Mini "MinflaccceeiiiaaeePPPoooorrreed
"Mini - Tee
23

---

Published as a conference paper at ICLR 2025
"MiniHack Beaecck-Frreeee Freeze Randde- - Restricted
Generically, the Nethack Learning Environment contains over 100 actions; whereas the "restricted" set of games reduce this space. The possible primitive actions include the 8 cardinal actions (N, E, S, W, NE, SE, SW, NW), PICKUP, ZAP, WEAR, APPLY, QUAFF, PUTON, FIRE, READ, and RUSH. It may be even less dependeing on the specific environment.
Abstract States We hand-construct a map that maps from the multi-modal observations to a set of (object identities, object attribute) vectors. We do this by taking the glyph and i inventory observation modalities and filtering out all glyphs containing blank spaces, floor, and walls. To get the object attributes, We take the remaining glyphs and mark them being "in world" if they are in g glyph, and "in inventory" if they were from the inventory modality. Then, if the agent object was beside object and took a primitive action move towards it, mark the object as "standing on". If we observe levitating message (e.g. "up, up and away" from the me message modality, we set the agent object have "levitate" attribute. Alternatively, if the low level episode terminates along with the last observation being either levitate message or a "freeze spell" message (e.g. "the frozen bolt bounces"), then we set the corresponding levitate or freeze object to have "used" status. Finally, if an object disappears from the low level observation that is not being stood on, we set it to have "removed" attribute. We summarize this in Table 4.
Attribute Name
Meaning
How to Infer
in world standing on
The object in the world Object in world, and the
Object is in the glyph observation
Object was in the glyph observation and message says "you see here..."
player standing on it Object in inventory
in inventory levitating used
Object in the inventory observation
Agent is levitating An object is used
Message says a levitate message
A freeze or levitate message appears, and episode terminates The object is not longer present
removed
Object disappears from world
Table 4: Object Attributes and how to infer them.
As for object identities, they are simply the glyph number given to them in MiniHack (e.g. 329 for the player agent).
Embeddings We then embed the object identities and attributes before providing them to the agents. For object identities, we take each object's 's glyph number, and map it to the NetHack text de- scription for that glyph (which we can extract from the screen descript ions modality in the low level observation). For each glyph's 's text description (which is short sentence), we pre-process the sentence to remove stop words, then take each word's GloVe (Pennington et al., 2014) embed- ding and averaging over the sentence. We use the -00-dimensional Wikipedia 2014 Gigaword 5 Glo Ve word vector embeddings from https: //tp//ni...../////////////////// stanford. odu/ppojecte glove/ (6B tokens, 400k vocab, uncased). The 300-dim vectors were chosen found them to have naturally good separation between object types in preliminary analysis.
We embed the object attributes egories and repeated to fill 300 dimension). This is done to make sure the input space is equally represented by the identity and attribute embeddings.
-i0edimensional one-hot vectors (i.e. one-hot between six cat-
D.2.1 SANDBOX ENVIRONMENT
We design a sandbox environment where the agent can observe and interact with all items it may encounter in one the five evaluation environments. In the sandbox, each of these objects has an 0.5 chance (independently) of spawning at a random location: levitate boots, levitate rings. levitate potions, freeze horn, freeze wands. Further there is 0.5 change that a second of the same object type will spawn. The agent interact with any of the objects for up 250 low level frames, or the episode terminates once the agent sets any of the object's attribute to "used". Figure 14c shows an instantiation of an episode.
24

---

Published as a conference paper at ICLR 2025
atir embeds
attr
embeds
X
1
600
Figure 13: Embedding
(a) Levitate-Potion
(b) Freeze-Random
(c) Sandbox
Figure 14: Example Minihack environments. Only cropped version of the screen is rendered for presentation clarity. The full observation space contains larger screen, agent's inventory, in in-game stats, and in-game messages (not shown here).
D.2.2 Low LEVEL POLICY ERRORS
The learned and hand-defined low level policies are not perfect. We see that depending on the initial abstract state, the low level policy succeeds average around 90% of the time, but can be as low as 10% for particular states (Figure 15).
world)
goal final state
FREEZE HORNS inventory)
levitating
used
FREEZE HORNS world)
standing, to_top in jnventory
FREEZE, HORNS (standing_on) on top)
world levitating
FREEZE WANDS (in, world)
FREEZE_WANDS
FREEZE, WANDS (standing_on) cattinggnooooooooooo___
used
LEVITATE BOOTS world)
standing_on _n__to in j_inventtory
LEVITATE_ BOOTT (standing_on, top)
in world
LEVITATE POTIONS world)
to 0.0
LEVITATE_ POTIONS (standing_oo) on top)
60
STAIRS world)
STAIRS_UP UP (standing_on, top)
0.0
0.2
0.4
0.6
0.8
1.0
success probability
Figure 15: Low-Level policy success rates for an environment
Some common failure cases for these low-level policies are:
1. To item in the inventory, the policy needs select which item (e.g. drink potion the policy needs to first do primitive action "quaff", then primitive action that selects the inventory object to drink, etc.). When there are multiple objects in the inventory it just picks random object, meaning it might not pick the object specified by the behaviour
2. The "go stand on" policy is coded to move in a straight line towards an object. Since we measure policy "success" whenever the abstract state changes, if the agent walks
25

---

Published as a conference paper at ICLR 2025
another object route, the Ab-MDP will register abstract state change and deem the transition unsuccessful
3. To successfully cast a freeze spell, the NetHack game needs to return a message saying the spell has bounced off wall. To knowledge, this is stochastic- we got far figuring out that the spell needs to travel sufficient distance to bounce, but even so we were unable produce low level policy that results in spell bounce 100% of the time. This succeeds 80% of the time
E BASELINE METHODS
For baselines, we use an environment wrapper which abstracts away the low level observations into only the abstract level observations and behaviours (as described in Section 2.1). We similarly equip the baseline policies with transformer-basee encoders, with input masking in places where the input sequence longer than the number of objects in the set as well as positional embedding, which we found in preliminary experiments helped with training. For action spaces, the agent generates discrete integer action which maps onto an object index in the observation set as well as the object attribute to change to.
Specifically, we use Dreamer-v3 a performant moodel-based baseline (Hafner al., 2023), MuZero as another model-basee baseline (Schrittwieser et al., 2020; Werner Duvaud, 2019), and PPO model free baseline (Schulman et al., 2017; Petrenko et al., 2020). For exploration-reward transfer disagreement-based model for Dreamer intrinsic reward (Sekar et al., 2020), as well as using the same count- based reward to train all models.
experiments,
Finally, also compare against other methods that uses structured MDP, such the attribute planner Zhang al. (2018), and NCS in Chang et al. (2023). These results can found in Section R.1.3.
We outline the main differences between the different model-learning methods Table 5.
State
Repre-
Function Approx
Planning
sentation
Dreamer-V3 (Hafner
Feature vector
Neural Net (genera-
Model free (in imag- ination)
al., 2023)
tive) Neural Net (genera- tive)
MuZero (Schrit-
Feature vector
MCTS
twieser al., 2020)
DAC-MDP (Shrestha
Feature vector
k-NN
Value iteration (on k- NN "cores")
al., 2020)
Attribute Planner
Tabular
No approximation
Dijkstra
(Zhang al., 2018)
DLP-EIT (Haramati
Sets of vec-
Neural Net (genera- tive)
Model free
et al., 2024)
tors
NCS (Chang et al.,
Sets of {a,E)
K K-means (on just Greedy tribute vectors è Neural net (discrimi- Dijkstra native)
2023)
vectors
MEAD
Sets of (a.E)
vectors
Ablation
Sets of {) vectors
K-means (on just at- Dijkstras tribute vectors )
Table 5: Model-based methods
F EXTENDED RESULTS
F.1 EXTENDED MAIN RESULTS
All runs, extended and main, are run with a minimum of independent seeds. We outline the specific number of seeds for each method and each main experiment in Table 6.
26

---

Published as a conference paper at ICLR 2025
Experiment Group
Method
Seeds
MEAD
4
Learning from scratch, 2D crafting games (Figure 4a)
Dreamer- V3
PPO
3
MEAD
5
Learning from seratch, Minihack skill games (Figure 4b)
MuZero
3
Dreamer-Va V3
4
PPO
3
Transfer from Sandbox pre-training (Figure 5a)
MEAD
3
Dreamer- V3 (count disagreement
3/3/3
/ reward)
Transfer from
MEAD
3
Freeze-Random pretraining (Figure 5b)
Dreamer-V3 (count disagreement reward)
3/3/3
Compositional Environment (Figure 5c)
MEAD (from scratch / pretrained)
3/3
Dreamer-V3 (from scratch / pre- trained)
3/3
Table 6: Number of seeds used our main results.
To get the performance of the various algorithms in the low level MDP plotted them in Fig- ure 4b, we extracted them from Figure 13 Henaff et al. (2022), using the WebPlorDigitizer app (https: apps. automeria i.. ht//wd///
Levltate-Bocts
Levitate-Potion
Mean matum
6:75-
0.75
Moan Epicade netum
Levitate-Potion
0.0
0.0
-4.0
-0.00
-0.5
-05
15166 20003 JUdE 48au2 sal0
2866 2o602 Ba0 ao002 Salao
20200 so60
Freeze-Homm
Freeze-Wand
Freeze-Hor
Freeze-Wand
3.00
1.00
Moun Episoda Buturn
2.0
1.0
6.5
0,5
623
0.0
0.0
-2.2
6.5
05-
-4.75
-0.75
10300 20100 J002 40020 30000
woico 20100 2000E 40000 30020
Erv Frames (Unear Scale)
Env Frames (Linear Scale)
Freare-Bandom
Model based (burs. no pretrain!
(ours, no pretraining)
Dreamery3
(count-based pretraining)
8.79-
Model based tours,
Model based (ours,
Dreamenv3
Wean Eokaode
count-based pretrinnigg!
count-based pretraining)
(disagreement pretraining)
Dresmer pretrain)
Dreamer-y3
Dreamery3 (reward-based pretraining)
8.00
Dreameriy3
(no pretraining)
based pretraining)
-250
Dreamenu3
disagresmert pretraining)
(b) Transfer to new object types. Pretrained in Freeze Random environment where Freeze Horn and Freeze Wand objects appear with 0.5 prohabilities, fine-tuning remaining environments. (Top row): fine-tuning performance where new object type fully unseen, (bottom row): fine- tuning performance for seen objects but appearing in different frequencies.
1000e 20003 30006 48909 5500
Dreamery33
Eny Tramme (Linear
(reward-based pretraining!
(a) Sandbox transfer performance. Episode return from fine-tuning model weights from pre- training a sandbox environment containing multiple ob- jects instances and types.
Figure Full transfer experiments
F.I.I
PPO ALL RUNS
We plot all PPO runs with different pre-training settings in Figure 18, Pre-training resulted only negative transfer compared to pre-training.
27

---

Published as a conference paper at ICLR 2025
Paromaaic mode; coutt baard craiion NCTS
Non param abric
court baned expleration NCTS
paramabicc randon expicrattor
Figure 17: Training from scratch with different model types
Frezze
Freeze
Schodd
90+
-MA-
+
-EA
Court based Prut sining in Frocze-Random Random
Countbased Pretraining in Playground
No Protraining
Figure 18: PPO all runs.
F.1.2 DREAMER
We show extended transfer results for Dreamer. We see in Figures 19 and 20 that when trained for much longer, Dreamer with both disagreement- and count-based pre-training does start to get more rewards. However, we also observe the reward curves with pre-training is much worse than the reward curves without pre-training (black dotted lines), indicating negative transfer across environment and object types for these methods.
Boots
Levitate Pot
Prezze Wand
Random
Tarma 1.3100
Drea maned peetrennng)
Driamma 192 (courlo baaddd-adddddd prezzin
Dreamer-33 [dlaggsseeer pratraiting!
Figure 19:
Longer Dreamer pretraining runs in the sandbox environment to show eval reward does increase with more samples, albeit pretraining seems result in negative transfer.
F.1.3 OTHER MODEL LEARNING APPROACHES
We plot comprehensively the non non- coraaarettrrc model results in Figure 21. This model uses table to keep track of transition successful probabilities between abstract states, used by the attribute planner in Zhang al. (2018). We show direct comparison between methods in Figure 17.
We further compare against an approach using K-means clustering on the attributes, and labeling edges (between the k-means "cores") with the last observed action that successfully led to tran- sition. This way of model-learning along with greedy planning was used by the NCS (neural con- straint satisfaction) method of Chang et al. (2023). Specifically NCS uses k-means clustering with greedy planning, which found to not work for any of our environments. A hybrid variant using k-means clustering with Dijkstra planning works MiniHack (albeit less well than MEAD), and fails completely in 2D crafting.
28

---

Published as a conference paper at ICLR 2025
Levitokt Boots
Levibate Pation
Freeze Hom
Freeze Word
2009 1000 L100 www.00000o.o calte a0000o [Jhaer Scale!
2000 EOOW LUBER ROPVERRRRROO
5000
20085
23000
m2 protraninn)
Dreamenug (count based pretranngg ngg
Dreamenug [dsagremment pretralsing)
Figure 20: Longer Dreamer pretraining runs in the Freeze Random environment to show eval reward does increase with more samples, ableit pretraining seems to result in negative transfer.
Exbood Tatuan
0.3
03
0.0-
00855
J0085
20000
200R
10000
3000
21000
30,00 00000 16813
20000 M000 R0.0
Non-Pamm MCTS Expleration
Non-Paramy Rendom Explorstion
(a) Mean episode return
Prsuee fum
PrecocWurd
Presos-Tancom
Dravee 3
58
25
8000
BT hramio
CreThunes
Dey Pransor
Dr hammo
Poram. MCTS Exploration
Random Exploration
(b) Number of unique (Object, Attribute, New Attribute) transitions discovered
Figure 21: Non-parametric agent. Notice the difference in X-axis scale for Freeze- Freeze-Wanddae and Freeze- Random environments.
F.2 Low LEVEL LEARNING
F.2. LEARNING Low LEVEL POLICIES
We can reinforcement learn low level behaviours. From some X,, we propose behaviour (a(9, E') and execute т. until abstract transition. If the resulting X contains item (a()) E'), we reward Ty with a reward of +1 (else -1).
While this can be implemented in many ways, in practice we goal conditioned RL set-up, where a single model free policy network takes in both the low level state and the behaviour embedding produce primitive action: -S B - A. The low level state is the environment's raw observations while the behaviour embedding is the same as the one given to the abstract level model. wrapper over the original environment which provides the policy with reward if (a <') E X441, and -1 otherwise. We use APPO the policy optimization algorithm (Petrenko et al., 2020). The point of this is merely to show that we can learn the set of behaviours neurally and building an Ab-MDP with them still works.
This is implemented
The full MiniHack low level policy learning result is shown in Figure 24a. The performance of MEAD in Ab-MDP with learned behaviours is shown in Figure 24b. Note that the reported 21) Crafting results in main text Section 4.1 already use neurally learned behaviours.
F.2.2 LEARNING OBJECT-CENTRE MAP
We train an object centric map which takes the low level observation and predict the abstract state: X, = (S, S,). The encoder map M, uses a convolutional encoder for MiniHack low level observations (specifically, the low level encoder from Henaff et al. (2022)), and learned embeddings
29

---

Published as a conference paper at ICLR 2025
Mumber Stecoveeed tonat
Noan sode Natum
Aopaci buto Tunittos
-33-
-14
Enu Frames
The Trannee
bir Prames
Non-Param, MCTS Explorstion
Parametrii, MCTS Exploratice
Counts MCTS (depth=11
Counts MCTS (depthed))
Exploration
(depth-2)
(a) Training from scratch with parametric (ours) and non-parametric model types in the Freeze-Wand en- vironment. (Left) Episade return. (Right) Number unique valid (item identity, item attribute, new at- tribute) transitions discovered.
(b) Effect of count-based intrinsic reward and MCTS in the Freeze-Wand enviromment. (Left) Episode return. (Right) Number of unique valid (item identity, item attribute, new attribute) transi- tions discovered by the agent.
Figure 22: Model ablations
Env: 2D Crafting, steps
Env: Minihack Freeze Wand
MMM
epeode reward
0.0
Mean ephode mwaard
0.6
0.4
0.0
Dijkstra MEAD (parametric madil, Dijstra)
MEAD (parametric model. Dijkstra) kmeans with Dijistra
-0.2
0.2
NCS model (Kmeans, greedy planningl
0.0
20000
40000
60000
80000 100000
20000
40000
Environment. steps
Environment stops
(a) Training from scratch in 2D crafting
(b) Training from scratch in MiniHack.
Figure 23: Performance of the k-means clustering on attributes method as used in NCS (Chang et al., 2023). NCS specifically uses k-means clustering with greedy planning, although also test vari- ant using k-means clustering with Dijkstra planning. Legends specify the function approximation and planning methods used.
for each discrete actions. After concatenating all encoded inputs, the latent is projected through two 2-layer feedforward nets, one producing categorical distribution over item identities, and another categorical distribution over item attributes. The output of this map two tensors each with shape (item set size x num categories). We can train this via a categorical cross-entropy loss to predict the correct identity and attribute categories for all items.
We collect a dataset of 100k transitions using a random policy in the MiniHack Freeze-Horn envi- ronment, and use the hand-crafted mapping provide ground truth abstract state labels. We train using Adams optimizer with batch size 4096 and learning rate 3e-4.
We investigate the effect of training with different sub-sets of the data both the encoder accuracy and the subsequent planning performance in Figure 25.
F.3 EXTENDED ABLATION RESULTS
F.3.1 GENERATIVE DISCRIMINATIVE MODELLING
We design an experiment where we isolate the difference made by the generative vs. discriminative modelling, which we can view as constraint over the support of the predicted distribution for X
We collect a dataset using a hand-crafted expert policy in MiniHack's with a degree of action noise added to increase coverage. This is done to remove the exploration problem: model only needs to learn to the trajectories well and it should contain within its world model the correct transitions required to solve the task. We fit two models:
Freeze-Home environment,
A discriminative model, which learns to predict success probability (Equation 1) as de- scribed by the main text.
30

---

Published as a conference paper at ICLR 2025
Freeze-Random
Freeze-Wand
Leviaae-Pation
0.6
16
0.-
6.2-
32
0.2
0.5
3.3 Eny frames
0.0
28
03
0.5
2.0
0.0
0.5
2.0
2.0
6.2
2.3
Env Fromess
1n7
Emw frammss
347
Env Framess
Env Frames
(a) Low level model free RT. to optimize behaviour success probabilities
dom
1.5
30-
-1.0
-10
Crafted
-00
L60 mod
Frames
Trames
Emn
Inu Frames
(b) MEAD performance in Ab-MDP with learned crafted behaviours
Figure 24: Learning low level behaviours within MiniHack
Encoder Accuracies
Performance with Different Encoders
0.8
Prop data
0.01
0.4
0.05
0.5
30000 Enviranment Staps
Proportion dota trained on
Figure 25: Effect of training on different proportions of the 100k dataset on encoder accuracy (left) and MEAD performance (right) in Mi Faeehacce--eeoorr.
A generative model, which learns predict the distribution over X.
The two models use the same architecture to encode the previous abstract state and behaviours. differs only in its output head: the discriminative model outputs a binary prediction, while the generative model outputs categorical prediction for each item. We observe in Figure 26 that both model classes optimize the loss.
Discriminative
Generative
10a
100
5000
100
5000
0.4
1000
10000
1000
10000
Bin Entrrpy
Cross Entropy
0.3
10-4
0.2
0.1
100
200
3.00
400
500
100
200
300
400
500
Optimization
Optimization
Figure 26: Loss curves for training generative and discriminative mdoels fixed expert dataset. Colors denote different number of dataset transition size.
To do planning, we run Dijkstra, with the edge weight as the negative log Pr(X'|X, b, 8) (where 0 denotes the generative discriminative model). There two ways can gen- erate the neighbours for the generative model: either by proposing the modal next state, ie. X arg maxx: Pr(X' X,,, 0), or by expected behaviour change X' A(X,b). Using
31

---

Published as a conference paper at ICLR 2025
pected behaviour change does not work in planning (the model never proposes behaviour sequences that lead to solving the task), thus we report results by using the modal next state. The results are reported in main text Figure 7, noting the error bars denote standard error of mean.
F.3.2 MODEL FITTING
We found that running sufficient number of model fitting steps in between data collection steps was important. Fully re-setting the weights before data fitting is also helpful.
F.3.3 GLOVE SEMANTIC EMBEDDING
0.84
059-
94
04-
04-
.233
02
Onr hameec
Emn Thames
Era Frames
Eve Pronnes
CT Frumoo
GloWe wwbredding stject name (Troren (neights)
Random doject wiberddingg ing (Ira na ghiti)
Figure 27: Training from scratch with frozen GloVe embeddings, or trainable randomly initialized weights
We investigate the impact having the object identity embeddings semantically meaningful versus random one-hot embedding in MiniHack. By default, all experiments in the main text take as input object identity embedding using GloVe Pennington et al. (2014) embedding, which are kept frozen throughout training. We ablate this by initializing the object identity embedding matrix to be randomly initialized from a standard Gaussian, (0, 1), and allow them to be updated by gradients.
Figure 27 shows the results in all five games. We observe that having word embedding helps bit in environment where there are many types of the same objects (i.e. Levitate- Potion and Freeze-Wand), This is unsurprising wand" and "metal wand" would have similar averaged embedding). Abeit they still performant overall. This hints at the abstraction being useful mainly due to its (object identity, object attribute) structure.
the text embedding between objects would be rather similar (e.g. "wooden
F.4 MODEL UPDATES
We show effects of number of model update steps during data collection in Figure 28, 29.&31.
Macki MUiCe pxps
6.0
0.3
2500
0.5-
0.2
02
Treshddi to cilict MON EER
0.44
12095
20065 40900
1900e
1900 2900 1900
4900
9900 0900
2000
45900
66899
100900
(a) Effect of model update steps (b) Effect model threshold logic (c) Weight reset strategies before prior being allowed to collect fitting on new data more data
between data collection steps.
Figure 28: Model training hyperparameters training in Freeze Wand environment.
32

---

Published as a conference paper at ICLR 2025
-44
TAMARGY TIIOOEE
ON Premes
(a) MCTS hyperparameter cf- fects, planning to find the max re- warding state using current para- metric model.
Dxfaver
(b) Using the empirical dataset probability for MCTS planning (i.c., non-parametric estimate)
Figure 29: MCTS exploration hyperparameters in the Freeze Wand environment.
Model Binary Accuracy
Mean Episode Reward
1.0
accuracy
0.9
Optim steps
0.5
0.8
500
Meaa Saade
500
1000
1000
0.7
0.5
20000
40000
10000
20000
30000
40000
Environment steps
Environment steps
Figure 30: Ablating just the effect of model optimization steps in Mini - -reezee-eeeee-ee Wand env. Model weights are reset after each data collection step (of 2.5k frames), then trained for the specified number steps before the next data collection step. With sufficient optimization steps, high model accuracy and planning performance is achieved.
F.5 MODEL EMBEDDING
We investigate the vector representation of different (object identity, object attribute) pairs in our model at the start of training and in the end. Results are shown in Figure 32. We observe vector cluster mainly along object attributes in the start of training (due to the orthogonal object embed- ding initialization). After training, vectors cluster by object attributes For instance, the STAIRS_UE object in Figure 32a are not clustered by attributes along the first two principle components in the Freeze Random environment, likely due their usage being very different as compared to the other usable objects. We also observe separation by object-identitie inside each of the object attribute clusters after training, but not before.
identity-dependent way.
Again referring to Figure 32 (a), we notice that unseen objects (not seen in the Freeze Random environment-beet it is the "levitate" objects) largely fall into clusters by attributes.
F.6 INTERNAL. WORLD GRAPHS
We also show the model's internal world model in highly interpretable way: through running a graph search algorithm (c.g. breadth first search) along transitions the model predicts is high probability, and reconstructing the model's internal graph.
We do this for two environments, showing only high probability transitions, and disallow visualiza- tion of loops in the graphs. The world model graph for Freeze-Random is in Figure 33, Note this is the same model which we illustrated in brief main text Figure 3b. The world model graph for 200k frames checkpoint in the eevittt-tthnn-Feee environment is in Figure 34, which is much more complex even with only the high probability edges.
33

---

Published as a conference paper at ICLR 2025
Dutaa Ungerte
my Mest
2 30009
22-
2052
-344
EnyTruma
Eny Prrmes
(a) Eval performance as function of data collection
(b) Wallclock time for different data collection frequencies
Figure 31: Data collection frequeney in the Freeze Wand environment.
G TRAINING DETAILS AND HYPERPARAMETERS
G.1 OUR MODEL
We describe our default parameters in Table 7.
Hyperparameter
Value
Optimizer
Adam
Learning rate
lc-4
Optimization
Adam betas
(0.9, 0.999)
Adam eps
1e-8
Max batch size
2048
Additional frames before training
2500
Data Collection
Weight reset before training
reset all weights
Minimum optimizer steps before data collection
2500
Random goal selection probability
0.2
Num simulations
16
MCTS Exploration
Max search depth
Probability for valid edge
0.5
Dijkstra max iters
100
Eval Planning
Dijkstra low probability cut-off
0.1
Table 7: Default training parameters for our model
G.2 DREAMER-V3
We describe the default parameter used the Dreamer-v3 baseline. We use PyTorch im- plementation of Dreamer: https: / huu..uu.... www.coo//ooieee /0mM/NM5122 htt////eeeeeeeeorreeeeee
which has the
similar training curves reported initially in Dreamer-v3. Unless otherwise mentioned, "XL" world model setting with the highest amount of training iteration, which had the best reported sample efficiency in Hafner al. (2023), and was confirmed by preliminary experiments (Fig- ure 35). We start with parameters as faithful to the reported ones in Hafner et al. (2023) as possible was found work across wide range of settings, then sweep over subset of hyperparameters and report the best ones (which we use) in Table 8.
We perform parameter sweep over subset and report the best ones (which we use) in Otherwise we keep the parameters as faithful to the reported ones in Hafner et al. (2023) as possible as it was found to work across a wide range of settings.
G.3 MUZERO
We use the MuZero code base from Werner Duvaud (2019). We sweep over learning rates (1x 10-5, 1 x 10-4, 1 x 10-3), encoding size (32, 64, 128, 256, 512), number of simulations (30, 90, 100,
34

---

Published as a conference paper at ICLR 2025
Env Frames:: 0
Env Frames: 200000
10
acdhute
By Objjctt IScntRies
POTIING
FREEZ,, HORNS
FROEZE, INAVES
ACENTSS
-0.5
1.0
Offect
stand on top
used
-0.5
0.0
0.5
1.0
10
PC
PC1
(a) Training Freeze Random environment. Per dimension variance ratio explained frame: (0.24735658, 0.22067907):
200k frames: 10.49951893 0.35679522)
Eny Frames:
Env Frames: 200000
By Object kieniiies
2.5
LEviTer BOOTS
0.9-
FREEZE HORRS
ponoos
FRAEZE INADS
ACENTS
5.0
5.0
@
By Giject Actr Butee
2.5
hluood
solo
-2.5
Inchaain
uood
-5.0
(b) Training in the Sanabox Playground environment. Per dimension variance ratio explained at O frame: (0.3139481. 0.23430094); 200k frames: (0.5275301, 0.23016657)
Figure 32: Principle components analysis (PCA) the (object identity, object attribute) vector representation in the model right before the key-query-value projection layer. PCA projects from 128 dimensions to 2 orthogonal dimensions of maximum variance. We show the representation at two points in training: frames and 200k frames (columns). We label each (object identity, object attribute) vector by its identity (top row) and by its attribute (bottom row) separately.
180]. using prioritized experience replay (or not), batch size {128, 2048). TD steps {4, 10, 20, 50, 100). and the delay in self-play simulations (allowing for more training steps in between data collections) of {0, 2, 8). We use the best hyperparameters and report them in Table 9.
We observe small effect of different self play delay lengths performing better in different games. We choose a delay of 0 as it performed the best in most games. Nevertheless, we report the perfor- mance of different self play delay lengths in Figure 36 for completeness.
G.4 PPO
We use the Sample-Factory Factory implementation of PPO (Petrenko et al., 2020). We run parameter sweep over subset and report the best parameters which we use in Table 10.
35

---

Published as a conference paper at ICLR 2025
AGENTS world. FREEZE WANDS world.
FREEZE WANDS: in world d standing on tog: 0.99
STAIRS UP:
in world -> standing on top: 0.992
STAIRS world.
AGENTS in world. STAIRS UP standing
FREEZE WANDS standing
FREEZE WANDS:
standing toy in inventory, 0.997
FREEZE WANDS inventory. AGENTS world,
STAMS UP: In world i> standing on top: 0.982 No Inventary <> used: 0.988
FREEZE WANDS;
FREEZE WANDS in inventory. world. standing top.
STAIRS UP in world,
Figure 33: Extracted internal world model with particular starting state in Freeze-Randone by starting traversal from a given initial state. Only the most relevant objects are included for clarity and only transitions with success probability above I are included. Note that each node has 40-50 edges, although we only visualize the high probability ones.
Hyperparameter
Value
Optimizer
Adam
Model
le-4
Actor Ir
3e-5
Critic Ir
3e-5
Optimization
Batch size
32
Batch length
64
Train ratio
1024
Replay capacity
106
GRU recurrent units
4096
Model Units
Dense hidden units
1024
MLP layers
5
Planning
Imagination horizon
15
Exploration (disagreement)
Ensemble size
10
Table 8: Dreamer parameters
Hyperparameter
Value
Learning rate init
1 x 10
Encoding size
512
Batch size
2048
Num simulations
100
TD steps
20
Prioritized experience replay
False
Table 9: MuZero parameters
36

---

Published as a conference paper at ICLR 2025
Figure 34: Extracted internal world model by starting traversal from a given initial state. Only the most relevant objects are included for clarity. The plan required to solve the Levitate-then-Free task involves going from the top (root) state to the bottom left state.
37

---

Published as a conference paper at ICLR 2025
1.0
0.75
Moan episode reward
0.50
0.25
Train ratio
Train radi
256
512
0.00
1024
0.25
1024
2048
-0.5
2048
4096
-0.50
120000160000
15000
30000
Environment steps
Environment steps
(a) Training from scratch (no pre- training)
(b) Fine-tuning after reward-based pre-
training
Figure 35: Effect of Dreamer-v3 training ratio (amount of training to data collec- tion) performance, for training from scratch in MLniHHcck-rrezee-Waa (left), and for fine-tuning in MiniHack-Levitate -Boots from reward-based pretraining in Mini Fack-Freezee-Raneee
Los Lato So0ts
Los Kato Putlon
Fresee Har
Freeze Wond
Freque Randemm
35
0,50
0.22
32-
+0.50
259000 43008 Eav Trumes
200095 4300E9 Day Tarmsss
200090 400003 Trammss
200890 469900
389000
409000
EX0
Figure 36: Muzero runs the Minihack environments, with different self-play delay lengths
Hyperparameter
Value
Number of workers
Number of env per worker
10
Data Collection
Rollout length
128
Recurrence length
Optimizer
Adam
Model Ir
lc-4
Optimization
Reward scale
1.0
Batch size
4096
Shuffle minibatches
False
Table 10: PPO Parameters
38


Reviews:
[{"review_id": "Gwz3gx9Kmi", "reviewer": "Reviewer_eBxo", "summary": "The paper proposes a method for solving particular object-centric abstract MDPs with discrimitive world models, count-based exploration and planning / shortest path methods.\nThe paper shows very strong empirical results on two problem domains, Minihack and Crafting. The paper also presents results on learning low-level behaviors and object-centric representation (in a supervised way).\nOverall the paper contains intersting insights into model-based RL when applied to higher level abstraction.", "questions": "1. I did not understand the exploration used in section 4.1:\nYou introduced the count-based exploration before, but that does not seem to be used here? The task return is when running Dijkstra, right? Is there an exploration phase here?\n1. Can the authors comment on why they did not employ value iteration on the high-level instead of Dijkstra. That could deal with stochastic outcomes. \n1. For the Dreamer and MuZero baseline: did you run enough / more updates per collected data in the transfer setting? If you reset the policy it needs to be learned again. Dreamer needs to run many updates using the dream-phase to rebuilt it. In this sense the comprison is a but unfair, as your are using run-time planning. I suggest to run Dreamer with a larger number of dream-updates with the new reward function before interacting with the environment.\n\n## Comments:\n- Very little detail about the training/architecture of the world model in the main paper, I suggest to at least mention the architecture and the loss in the main paper.\n- The paper needs some careful proof reading:\n    - Abstract: the last sentence : to reinforce learn and to supervise learn is an uncommon use of the words\n    - first sentence of intro: have been major roadblocks: I think this is not correct. They have been topics of active research and many very effective solutions have been proposed. So I suggest to change that sentence\n    - l086: eaching havin form\n    - l143: more fully\n    - l196: trainbale\n    - l251: plot 95 confidence interval\n    - l270: paragraph is missing details on how exploration is done etc.\n    - l313/314, 331/332, 404/405\n    \n- In the related work, you might want to mention this paper:\nCurious Exploration via Structured World Models Yields Zero-Shot Object Manipulation\nhttps://openreview.net/forum?id=NnuYZ1el24C\nthat also shows that structured representations, world-model structure, and planning for exploration yield strong models and zero-shot capabilities (shown in manipulation).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Overall, clearly presented\n- very efficient method for solving high-level problems of a particular structure (Ab-MDP)\n- study of different transfer aspects\n- extensive comparison to strong method and impressive results\n- comprehensive analysis with lots of additional details in the appendix", "weaknesses": "- The method needs quite strong assumptions on the high-level MDP, which are that usually only one element is changing or nothing changes. The authors also argue that replanning fixes some amount of violation of that assumption, but an environment where agent/object positions need to be modeled on the high level would not work from my understanding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pzPEHG9MRt", "reviewer": "Reviewer_74S3", "summary": "The proposed using an abstract model of the world to plan with, and presents a model learning algorithm along with it. The main selling point of these abstractions is that they are object-centric. They respect specific useful tools in the environments, and curate the agent’s learning to be around them. These objects and their attributes are also human-interpretable.\nThey show how this can lead to better performance than existing model-based methods, and perform an empirical study into the components responsible for this, along with listing the settings where it does work.", "questions": "The relationship between the competence of a behaviour and the transition matrix was not made clear. Is the paragraph trying to say: a competent behaviour will have most/all of the probability mass concentrated on some other state s.t. X_{t+1} != X_t ?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Answers the questions they presented in the abstract and introduction.\n* The Methods section is succinct and points the reader towards the relevant experiments, based on the topic / interest.\n* Contributions are all listed usefully (and addressed)", "weaknesses": "* Abstractions are not discovered. It would be nice if this was told earlier in the paper.\n* The Abstract MDP proposed here and the options framework are two different mathematical objects. The Ab-MDP cannot be interpreted as options. Under the options framework, the Ab-MDP _may_ be interpreted as a semi-MDP. As mentioned in the appendix, it is the behavior, b, can be interpreted as an option.\n* In Section 4:\n     - Paper doesn't explicitly mention how many seeds were used in the experiment. The appendix hints at least 3, but it would be good to know how many were used for each plot. This can contribute to the context with which the reader interpret the results, whilst also providing the paper a context with which to substantiate its claims.\n     - Figure 4b: Which triangle is IMPALA, RND and ICM in each plot?\n     - Figure 5: Why do none of the black lines have error bars? Are they single seed? If so, that would be useful to specify in the caption.\n     - Figure 5a: Why does the blue line start high? Is this just a plotting artefact?\n* [Very Minor] Last sentence of abstract: change wording to talk about “reinforcement learning” and “supervised learning”\n* [Very Minor] Typo: line 325: “representatin”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SCnSfmUBAR", "reviewer": "Reviewer_ccnG", "summary": "This paper proposes a method for efficient exploration and world model learning in reinforcement learning (RL) by leveraging an object-centric abstraction. The authors define an Abstracted Item-Attribute Markov Decision Process (Ab-MDP), where abstract states are sets of items and their attributes, and actions are abstract behaviors corresponding to object-perturbing policies. They introduce MEAD (Model-based Exploration of abstracted Attribute Dynamics), a fully model-based algorithm that learns a discriminative world model predicting the success probability of abstract behaviors. MEAD uses count-based intrinsic rewards and Monte-Carlo Tree Search (MCTS) for efficient exploration and plans using Dijkstra's algorithm to reach goal states. The method is evaluated on a suite of 2D crafting and MiniHack environments, demonstrating significant improvements in sample efficiency and performance over state-of-the-art baselines. The authors also show that MEAD transfers well to new environments and tasks and discuss how to learn the object-centric mapping and object-perturbing policies when they are not provided.", "questions": "1. Could you provide more details on how the object-centric mapping \\( M \\) is learned in practice, especially in complex environments and the environments where such mappings are not readily available? How does the performance of MEAD depend on the quality of this mapping?\n\n2. What are the computational requirements of MEAD compared to the baselines? For example, how does the planning  (MCTS and Dijkstra's algorithm used) affect it and is it feasible in more complex environments with larger state space and action space?\n\n3. How does your method compare to the methods closer to yours, such as the ones mentioned in your Related Work section?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Novelty: The approach proposed improves exploration and world model learning in RL by utilizing object-centric abstractions and discriminative modeling.\n2. Efficient Exploration: MEAD adopts a discriminative world model and countbased intrinsic reward for efficient exploration.\n3. Experimental Results: Plots and explanations are detailed and clear.\n4. Transfer and Generalization: The paper shows MEAD can transfer zero-shot and few-shot to new environments and object types indicating of its generalization ability.\n5. Interpretable World Model: Object-centric abstraction naturally can lead to an interpretable world model which can be helpful for understanding agent behavior.", "weaknesses": "1. Assumption of Given Abstraction: The approach inherently assumes access to an object-centric mapping and competent object-perturbing policies. This limits the applicability in scenarios where getting these mappings can be cumbersome.\n\n2. Abstraction Learning: The authors discuss learning the object-centric mapping and the policies. However, more details and empirical validation would strengthen this section.\n\n3. Experimens: It is not clear if the experiments using PPO and Dreamer-v3 are applied to Ab-MDP or to the MDP setting. It is mentioned that these methods are used in Ab-MDP, but there also is \"we also show the final performance of state-of-the-art exploration  methods in the low level MDP\" in the caption for Figure 4.\nAdditionally, there is no comparison with any of the methods that are closer to the proposed method (a good number of them are mentioned in the Related Work section). Also, for the comparison with PPO and Dreamer-v3, the proposed method is not evaluated in the settings and environments in which the baseline methods are mostly evaluated for their performance. Therefore, it is unclear how the proposed method is generalizable over more complex environments and the comparison do not sound quite fair.\n\n4. Clarity: Certain parts of the paper and the overall flow of information can be improved and clarified. For example, the paragraph starting at line 161 and ending in line 167 could be revised to convey the message more clearly. Additionally, the explanation of T(X′|X, b) in lines 109-111 could be rephrased for better clarity. Here are some errors I noticed: in line 316, \"the agent observe\" should be corrected to \"the agent observes,\" and in line 331, \"themin\" should be separated as \"them in.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9nXmZXIU9R", "reviewer": "Reviewer_UYL3", "summary": "This work presents MEAD, a model-based planning method that learns a semantic reward-free world model in an abstracted MDP. The learned model is then combined with standard search algorithms to generate goal-reaching trajectories.  The abstracted MDP comprises of an object-centric mapping, where a state consists of a discrete set if items, each associated with a set of attributes. Notably, this work focuses on learning a discriminative world model, which predicts the probability of reaching a particular state, given an action (defined in the paper as behaviour) and current state. Results on 2-D crafting and Minihack show that MEAD significantly outperforms existing baselines.", "questions": "1. The underlying assumptions of Section 4.1 need to be clarified further. If Dreamer and other baseline policies also assume access to these behaviors and operate in the same abstract MDP, then what are the most likely for the differences in sample efficiency. Inference for MEAD is done using Dijkstra's, how is it done for the other methods? Some analysis of the results in *Section 4.1* would be helpful . \n\n2. In general, it would be useful to have 2 sentence descriptions of the baseline methods or a table that shows the high-level differences between MEAD and baselines (access to behavior policies, abstract world model, inference-time strategy, etc.,). This will help increase clarity as currently it is a bit difficult to keep track of the differences. \n\n3. In the transfer experiments, Dreamer and other baseline struggles, and the authors hypothesize that this may be due to a distribution shift. Why does MEAD not suffer from a similar distribution shift? Does this have to do with the discriminative model it learns or because of other assumptions in the work. \n\n4. Section 4.3.1 discusses learning all low-level behaviours. However, I imagine that in most environments, the number of behaviors is extremely large ( $N*m$) and most of the them are incompetent. How much overload does this actually add in practice? What are the total number of low-level policies learned and what fraction of the behaviors are competent (success probability above some threshold)? What are the cumulative training steps (and wall-clock hours) spent in learning these behaviors, and how do they compare to the metrics for the world-model training? \n\n5. How do you expect this method to perform when attributes are continuous (eg., speed, position, etc.,)? Continuous attributes can be commonly found in many robotic domains, and this work relies on having a discrete set of states.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Clear writing and presentation**: The paper is well-written and generally easy to follow. It provides the right intuition and effectively builds up motivation where needed. The related work thoroughly covers existing work in the area.  \n2. **Strong and Extensive Results**: The proposed method performs strongly against competitive baselines and it can have significant advantages in settings where the assumptions of the work are satisfied. The paper also has a comprehensive set of experiments and analysis to support its claims. Notably, the paper has numerous ablations to understand the importance of the different components of the proposed method, which was enjoyable to read.", "weaknesses": "1. **Limited Applicability**: The presented method assumes access to object-centric maps, low-level behavior policies and discrete space for attributes. Many RL environments do not provide access to all of the above, and while it is possible to learn these policies/maps, it might require a lot of overhead. \n2. **More clarity in results**: The paper does not do a good job of introducing the baselines and only describes them superficially in the main paper. More detailed descriptions will help increase clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WB942EK7y9", "reviewer": "Reviewer_5fJW", "summary": "Summary: The authors propose a method to perform exploration and planning in an abstract MDP which they denote as Ab-MDP. \nThis assumes the  a set of abstract states is given, where the abstract state consists of an item id and a specific attribute for the object. The authors define as a competent behaviour, when the corresponding attribute of an item changes in an admissible manner within k steps. The authors train a forward model that predicts the probability i.e. the next \npossible state from the competent behaviour. For finding the correct set of behaviours, MCTS together with a count based intrinsic reward is used. During inference time, the behaviours can then simply be found\nby applying the Dijkstra algorithm on the probabilities of the competent behaviours put out by the forward model. The authors show that their model trained on the Ab-MDP outperforms methods trained on a low-level MDP as well \nas known model-based methods such as Dreamer on the Ab-MDP. \n\nOverall, I find the proposed abstraction an interesting and necessary view for exploration and policy learning in reinforcement learning. However, I feel that this work makes very strong assumptions on the needed abstractions for the proposed Ab-MDP, being tailored to the tested environments. \nOne thing that would mitigate this is a discussion on how the proposed item id, item attribute abstraction can be a general representation for other problems. For instance the attributes (in world, standing on, in inventory) may be meaningless in other environments or be a level of abstraction \nThat is too coarse to learn a good policy. How do we choose a good level of abstraction? Or what are desirable attributes that may allow for learning a good policy in most cases? I think this is important to make the Ab-MDP formalism more generalisable.", "questions": "- How do you train Dreamer-v3 on the proposed environments? I would assume performance to be worse for the proposed Ab-MDP since for instance the latent space bottleneck in Dreamer is tailored to continuous and not discrete states.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Learning from more complex and abstract behaviours is an important problem in reinforcement learning, especially since the current focus so far as the authors state is mainly on problems that can be solved without much planning.\n- The authors propose a straightforward architecture to explore and plan from discrete sets, i.e., abstract states showing they perform better than models that are not tailored to this abstract MDP.\n- The paper is well written and the provided illustrations help in making the method understandable. I also found the ablations to be insightful justifying the components of the method.\n- I think the results are promising, showing that the definition of a problem can be just as important as the underlying methods to solve it.", "weaknesses": "- The paper makes a lot of assumptions about the nature of how attributes and behaviours of the states are encoded. While a thorough explanation is provided for the given environments, it seems unfeasible to do this exhaustively for more complex environments.  I feel there needs to be a more general analysis of what constitutes an item attribute, I.e. a good abstraction level for an Ab-MDP.  In my view, when defining a new framework such as a variation of an MDP, it should define a general way to see the MDP that either is a useful abstraction or is tailored to a specific set of problem. While I see that the authors attempt this, it seems more that the authors restrict the MDP in a specific way to work with the given environments. \n\n- At inference time, the authors rely on Dijkstra to plan towards a specific state with their trained forward model. While this seems feasible for static environments, it seems less feasible for procedural environments, where, e.g., the position of objects could change. At the moment it seems to me that with the current environments the forward model has to learn all possible combinations of item id and item attributes. Could you illustrate how performance degrades w.r.t to how much the forward model is trained, i.e., how much the encoder overfits to the environments?\n\n- In the end, the authors perform a lot of manual feature engineering to make both environments fit in the Ab-MDP paradigm. In this way the contribution becomes really about learning a good model for this feature engineered solution. Although I do value the attempt of a world model paradigm that is based on discrete states (see positives).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)