Efficient Exploration and Discriminative World Model Learning with an Object-Centric Abstraction

We thank the reviewer for the very thoughtful review. We are glad to find our work deals with an important problem in RL, with promising results, and is well written. We appreciate the reviewer’s astute summary that “the definition of a problem can be just as important as the method to solve it”.

Below, we hope to address some of the weaknesses and questions for further discussion.

Response to Weaknesses (W) and Questions (Q)

W1: 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.

The Ab-MDP’s choice of modeling sets of (item id, item attribute) gives a flexible abstraction framework, as items and attributes can be flexibly defined. In practice, items usually correspond to actual objects in the world. This is not an uncommon practice, for instance, various works from already work with item sets as a level of abstraction (e.g. [LOCA20, NCS23]).

We also point the reviewer to Appendix section B.1.4 for some potential example application of Ab-MDP to various robotics environments.

W2a: 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.

Our environments are procedure across episodes: for example, in the Minihack Freeze-Wand environment, one out of the 27 different types of wands are randomly spawned in each episode. The location of all objects and the agent are also random across episodes.

In the scenario that item attributes change within an episode (independent of the agent’s perturbation), we think re-planning (which MEAD does) would deal with this.

W2b: 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.

It is not necessary to learn all possible combinations for our model to work. Generally speaking, any model-based method needs to model the distribution $P(X_{t+1} | X_t, b_t)$ (i.e. put density over the entire state space $\mathcal{X}$). We make specific assumptions about this distribution (attribute change for a single item in $X_t$), which allow us to model this distribution using a discriminative model. We believe our learning efficiency (Figure 4), positive transfer (Figure 5a,b), and discriminative model ablation (Section 4.4.1) results demonstrate that compared to common baseline methods, we can in fact do better with having seen the same amount of (item id, item attribute) combinations.

W2c: 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?

We perform an ablation by training the forward model with a different number of optimizer steps between data collection (note we reset the weights of our model and re-fit via sampling from the full dataset after each data collection; for ablation on this see Figure 28c). We see in Figure 30 that as long as we train for enough steps, the accuracy stays high (Figure 30 left), and the planning performance is good (Figure 30 right). If the model is insufficiently fit then accuracy decreases and planning gets worse.

We further note that our model is not “overfitting” to the environment: this is evident by the positive transfer results in Figure 5a and 5b, where our model can zero- and few-shot transfer to unseen environments / objects.

---

[LOCA20] Locatello, Francesco, et al. "Object-centric learning with slot attention." Advances in neural information processing systems 33 (2020): 11525-11538.

[NCS23] Chang, Michael, et al. "Neural constraint satisfaction: Hierarchical abstraction for combinatorial generalization in object rearrangement." arXiv preprint arXiv:2303.11373 (2023).

I want to thank the authors for the detailed response!

I feel my concerns have been addressed properly, however I still feel that the approach to define hard coded item attributes for a specific environment seems not very generalisable.

I do appreciate the authors' example applications in Table 1, but then I wonder why not at least 1 of these environments was tested as well, which would strengthen the contribution. In this way, it would be clearer that the approach works solidly and working out the abstraction is a separate area of research that can be left for future work.

I maintain my point that I find a world-model based method for discrete states very interesting and think this is the main contribution I value the most of this work. I therefore raise my score to marginally accept this paper.

We thank the reviewer for the insightful review. We are glad to see that we’ve done a good job in addressing all claims and contributions we set out to answer, and that it is clearly presented.

Below, we hope to address some of the weaknesses and questions for further discussion.

Response to Weakness (W) and Questions (Q)

W1: Abstractions are not discovered. It would be nice if this was told earlier in the paper.

We have revised our introductory section to clarify this, please see L064.

W2: Ab-MDP cannot be interpreted as options. Ab-MDP may be interpreted as a semi-MDP.. It is the behavior b that can be interpreted as an option.

We agree the Ab-MDP may be interpreted as a semi-MDP, and behaviors can be interpreted as options. We believe our analysis in Appendix B.1.1 reflected this exactly. We updated our wording (L114, L434) to reflect our intention of interpreting the Ab-MDP under the Options Framework, which include both semi-MDPs and temporally extended actions.

W3: Seeds, legends and figures

We apologize for the confusion. We have clarified all points in our updated draft:

• We added a table with seeds in Appendix Section F.1, Table 6
• The triangle actually refers to all three methods: which all get a reward of 0 and are overlapped. We have updated Figure 4b to make this less confusing
• The black lines in Figure 5 are the same curves as Figure 4b (plotting just the mean for clarity). We have updated Figure 5 caption to clarify this.
• Figure 5a: the blue line starts high as MEAD shows good zero-shot transfer from the sandbox environment to the new environment. This is likely because the sandbox environment spawns all objects with some probability, thus the agent has already seen the object (albeit in a setting where many other objects are present), and MEAD can effectively use this knowledge to transfer to environments that require using individual objects. The reviewer may also find our PCA analysis of MEAD’s representation space in Appendix F.5 and Figure 32b interesting, where we observe similar object-attribute pairs clustered together after learning.

We have also fixed the typos, we are very grateful for the reviewer’s effort in pointing them out.

Q1: 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 ?

This is correct. More precisely, since behaviours are defined as changes to an item’s attribute $b = ( \alpha^{(i)}, \xi’ )$, executing the behaviour policy may either lead to the change specified by the behaviour (i.e. the i-th item attribute changing to $\xi’$) from the current state $X_t$ to $X_{t+1}$, or not (no change or an un-specified change). Executing a competent behavior policy results in $P(X’ | X, b)$ having high probability where $X’$ contains the specified change (Equation 1).

We thank the reviewer for the insightful feedback. We are very encouraged to find that our paper is clearly presented with a thorough respect for related works, and demonstrates strong empirical results that are extensively ablated and analyzed.

Below, we hope to address the questions.

Response to Questions (Q)

Q1: What are the likely explanations for differences in sample efficiency between MEAD and baselines?

We believe this is due to how MEAD models $P(X_{t+1} | X_t, b_t)$. Specifically, we make specific assumptions about this transition (attribute change for a single item in $X_t$, see Equations 2 and 3), which allow us to use a discriminative model. We argue that this is a useful assumption as at an abstract level, many things in the world do in fact change only locally at a time. We are also not the first to make such an assumption (for e.g. [NCS23]).

On the other hand, Dreamer makes more general assumptions and model $P(X_{t+1} | X_t, b_t)$ generatively with little restrictions put on $X_{t+1}$. We directly show the difference in this modeling choice in Ablation Section 4.4.1, and show visualizations of the learned embedding space within the discriminative model in Appendix F.5.

Q2: Useful to have descriptions of the baseline methods. This will help increase clarity as currently it is difficult to keep track of differences.

Many thanks for the recommendation. We have updated our draft to include both a short description of baselines at the start of Section 4 and an Appendix Table 5 to fully outline their main differences. We hope the reviewer finds this helpful for improving clarity.

Q3: Why does MEAD not suffer from distribution shift? Does this have to do with the discriminative model it learns or because of other assumptions.

The reviewer is correct, we believe the greater robustness to distribution shift is due to discrimination vs. generation. Specifically, the forward model in MEAD only needs to identify whether a change in a single item is possible, which is easier to generalize given minor changes in other items. On the contrary, “generative” forward models (e.g. Dreamer) needs to model a distribution over the full set of items and is much harder. This is indirectly shown in Figure 7 ablation.

Q4: 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?

Our paper focuses mainly on what you do once you have a set of competent low-level policies. To get such policies, one could use behaviour cloning, hand-coded controllers, RL, etc.. We demonstrate one example here (of using RL) to learn them. Further, while $N \times M$ is indeed large, policies that change attributes are often generalizable. For instance in robotics, for the attribute change of “standing_over” to “in_inventory” (i.e. a “pickup” policy), typically people will learn one grasping policy and condition it on the object, rather than learning N seperate policies. [OKR24].

In our experiments, the specific number of competent low level policies depends on the environment. For instance in the Freeze-Wand environment (27 different types of wand), over 370 unique goals were attempted over the course of low level training, while only around 60 unique goals were achieved at least once. Training required ~ 20 million frames and takes 3-4 hours with a performant goal-conditioned PPO implementation. This is more samples than what is needed to train MEAD, albeit this cost is incurred by all methods we compare that use low level neural policies. All in all, the purpose of this section was mainly to reassure that these policies do not need to be perfect or hard-coded for MEAD to work.

Q5: 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.

Please see Global response A.2.

---

[NCS23] Chang, Michael, et al. "Neural constraint satisfaction: Hierarchical abstraction for combinatorial generalization in object rearrangement." arXiv preprint arXiv:2303.11373 (2023).

[OKR24] Liu, Peiqi, et al. "Ok-robot: What really matters in integrating open-knowledge models for robotics." arXiv preprint arXiv:2401.12202 (2024).

We thank the reviewer for the insightful feedback. We are glad to see the reviewer think our method is novel and clearly presented, and that MEAD demonstrates strong results in exploration, transfer and generalization, and has the added benefit of being an interpretable world model.

Below, we hope to address some of the weaknesses and questions for further discussion. We've added additional experiments and will refer to their exact locations in the updated pdf.

Response to Weakness (W) and Questions (Q)

W1: The approach inherently assumes access to an object-centric mapping and competent object-perturbing policies.

Regarding getting mapping: we acknowledge this limitation. Please also see Global Response A.1 for a further discussion on the object-centric mapping.

Regarding competent policies, we show in section 4.3.1 that we can learn them using reinforcement learning as long as we have access to the object-centric mapping (albeit less sample-efficiently). In practice it is also often easier to get sub-policies, such as by imitation learning [GUP19], than composing them together.

W2a and Q1: Details of training the object-centric mapping. How does MEAD depend on the quality of this mapping.

Per section 4.3.1, we learn the object centric mapping via supervised learning, assuming a labeled dataset of $(S, X)$ pairs can be collected. Given this dataset, we train an encoder which takes the low level (image-like) observation as input, and outputs a set of vectors. We can train this with supervised learning to predict the set of (item id, item attribute) categories given a low level observation. We have added more detailed training and architectural details for this in Appendix section F.2.2.

Additionally, we conducted an experiment by varying the training dataset size to investigate the effect of an inaccurate encoder map (M) on the downstream MEAD performance. We observe with less training data, the (held out) accuracy of the map decreases as expected (Figure 25, left). Nevertheless, we find that even with imperfect map, MEAD is surprisingly robust and able to solve the task even with ~60% prediction accuracy (Figure 25, right).

The point of the experiment is to demonstrate MEAD can work with a sufficiently well-learned mapping. As stated clearly, this work is not focused on unsupervised learning of such a mapping, although how to use such a mapping is not clear and the point of this current work.

W2b: More details on policies

Low level policies can be trained as goal-conditioned policies (conditioned on a specific behaviour), as described in section 4.3.1. We see from Appendix Section D.2.2 that for “imperfect” low level policies (i.e. do not 100% lead to the desired transition), MEAD is still robust. In fact, MEAD directly models the probability of policy success, and Dijkstra plans a path to maximize this probability.

W3: Not clear if PPO and Dreamer-v3 are applied to Ab-MDP or to the MDP setting… 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.

All methods (PPO and Dreamerv3) are applied to the Ab-MDP (as we state in L124), with the only exception being the black triangles / stars in Figure 4. We apologize if this point was not clear and have updated Figure 4b's legend. Specifically, the point of Figure 4b is three-fold:

• We wished to show that “vanilla” standard methods (e.g. PPO) can be applied to solve an Ab-MDP
• When solving the task in the Ab-MDP (as opposed to the low level MDP), regardless of method, we get a large efficiency gain (e.g. comparison between vanilla PPO in Ab-MDP in yellow, and state-of-the-art exploration methods + PPO in the low level MDP, as black stars and triangles)
• Finally, MEAD exploits the structure of the Ab-MDP better and is the most efficient method within the Ab-MDP

The point here is that Ab-MDP is a useful framework, and MEAD an efficient method to solve it. We do not claim MEAD as a general MDP solver (as is the case for e.g. PPO), and therefore do not compare MEAD with other methods in low level MDPs that PPO and Dreamer usually operate in.

Clarity

We are grateful for the clarity comments. We have updated the new pdf to reflect this (L163, L111, L320). We hope the reviewer find the changes suitable and we are open to additional comments.

---

[GUP19] Gupta, Abhishek, et al. "Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning." arXiv preprint arXiv:1910.11956 (2019).

We thank the reviewer for the insightful and positive review. We are glad to see the paper contains interesting insights for abstract-level models, is empirically very strong, with extensive comparisons and analysis.

Below, we address some of the weaknesses and questions for further discussion. We've added additional experiments and will refer to their exact locations in the updated pdf.

Response to Weaknesses (W) and Questions (Q)

W1: Environment where agent/object positions need to be modeled on the high level would not work

Please see global response A.2.

Q1: Exploration used in section 4.1.. Is there count-based reward used when running Dijkstra?

There is not. Dijkstra is used to search over a sparser graph with a sparse reward function. We do not explore in this phase, but purely exploit to plan toward a goal state.

Q2: Why not employ value iteration on the high-level instead of Dijkstra? This could deal with stochastic outcomes.

Value iteration (at least the standard version we are aware of) is too compute intensive. Consider an Ab-MDP with N=13 items and M=6 possible item attributes (as is the case in many MiniHack tasks), there are $M^N \approx 13 \text{ billion}$ possible abstract states $X$. Even though most state transitions are impossible, (naive) value iteration requires either iterating over or matrix-multiplying with all states, which is prohibitively expensive.

Dijkstra avoids this by only locally searching from the initial states based on possible state transitions until a goal state is found (we do not search a transition if its success probability is low), and is efficient when the transitions are sufficiently sparse. In the case transition is not sparse, MCTS can be used (as is done for the exploration phase).

Q3: 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 comparison 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.

We ran an experiment examining the effect of different train ratio and decided on 1024 being the best choice for training from scratch (this also happened to be the highest ratio used in the original paper), see Figure 35a. We ran additional ablation for the transfer experiment and found that the train ratio had little effect on the initial learning speed, albeit a too-high ratio may destabilize training (Figure 35b). Details for this are in Appendix G.2.

We also have the MuZero training ratio ablation in Figure 36, described further in Appendix G.3.

All in all, we did not find that pushing the train ratio higher consistently helps with greater sample efficiency. Also note none of the settings are more sample-efficient than MEAD.

Response to Comments

We are deeply thankful for the writing suggestions! We have incorporated all changes in the attached updated pdf, as well as adding the new related work.

Dijkstra & Exploration

Many thanks to the reviewer for clarifying the question. We run our method MEAD in two separate experimental "stages":

• Pure exploration phase: we use the count-based intrinsic reward, and we use MCTS with the (partially learned) model to do search to find states with high intrinsic reward (Illustrated in Figure 3a). This encourages the agent to visit the state space widely and generate data. We use this data to train the model, and over time the model becomes more and more accurate, which in turn further helps with exploration. There is no notion of extrinsic reward or goal states here.

• Evaluation / exploitation phase: given a model (trained up to some data budget from the pure exploration phase), we evaluate it by using Dijkstra to try to find a goal state (Figure 3b). It is entirely possible that if the goal state has not been visited, the transitions to the goal state will be deemed impossible by the model, and Dijkstra will not find the correct path to the goal state. This is what happens when the exploration phase has not been run for long enough (e.g. Figure 4b, with low env frames the episode reward is low). Also note we do not add the data generated in the eval phase into the buffer used for training the model.

---

High Frequency Changing Attributes

if you have variables that are changing with high frequency, e.g. something like a position (can be discrete) the efficiency of the approach suffers as the ab-MDP becomes big

Indeed, this is likely to be true and we will add this as a limitation. In general whether or not attributes like this is present would depend on the abstraction map $M$ itself, and the hope with (and general point of) abstractions is that things should happen at a slow(er) temporal resolution at a higher level. Nevertheless, we agree the Ab-MDP is likely less suitable for directly modelling these lower level attributes -- it's inductive biases make it less general for all problems, but highly efficient for a (in our opinion still flexible and useful) subset of problems.

---

We hope this addresses the reviewer's questions. Please let us know if you have more questions, and we apologize for this delayed response.

[XMEM22] Cheng, Ho Kei, and Alexander G. Schwing. "Xmem: Long-term video object segmentation with an atkinson-shiffrin memory model." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.

[SAM23] Kirillov, Alexander, et al. "Segment anything." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[ZHU23] Zhu, Yifeng, et al. "Learning generalizable manipulation policies with object-centric 3d representations." arXiv preprint arXiv:2310.14386 (2023).

[QIA24] Qian, Jianing, et al. "Task-Oriented Hierarchical Object Decomposition for Visuomotor Control." arXiv preprint arXiv:2411.01284 (2024).

[SHI24] Shi, Junyao, et al. "Plug-and-play object-centric representations from “what” and “where” foundation models." ICRA. 2024.

[XIE24] Xie, William, et al. "DeliGrasp: Inferring Object Properties with LLMs for Adaptive Grasp Policies." 8th Annual Conference on Robot Learning. 2024.

[LOCA20] Locatello, Francesco, et al. "Object-centric learning with slot attention." Advances in neural information processing systems 33 (2020): 11525-11538.

[NCS23] Chang, Michael, et al. "Neural constraint satisfaction: Hierarchical abstraction for combinatorial generalization in object rearrangement." arXiv preprint arXiv:2303.11373 (2023).

[TRT21] Janner, Michael, Qiyang Li, and Sergey Levine. "Offline reinforcement learning as one big sequence modeling problem." Advances in neural information processing systems 34 (2021): 1273-1286.

[VQA17] Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." Advances in neural information processing systems 30 (2017).

[VQB24] Lee, Seungjae, et al. "Behavior generation with latent actions." arXiv preprint arXiv:2403.03181 (2024).