Learning Abstract World Models for Value-preserving Planning with Options

• Unconvincing experiments: We are currently working to add MBRL baselines to our experiments. However, we think it's important to notice that HIRO and Director do skill discovery. We assume that skills are given and we leave for future work to use our framework in a loop of skill discovery.

• Missing references: Thank you for pointing out these references. We have now included a new paragraph in related works to include classical work in state aggregation theory in RL---that is definitely relevant to our work. However, it is worth pointing out that our proposed framework builds over probabilistic groundings that are more general than state aggregation: instead of partitioning the state space, we can have overlapping supports for abstract states. Therefore, a ground state might belong to more that one abstract state.

• Issues in technical details: We understand the concern on the relationship between dynamics-preserving abstractions and bisimulation (and similar), and we have updated our writing to help with this.

• However, we argue that the purpose of an abstract model for a set of skills is to be able to solve many new tasks with those skills by planning. Hence, the ground reward is not meant to represent a task but to represent the cost or reward associated with executing skills (e.g. energy cost to execute a particular motor skill). Therefore, we are not interested in preserving Markovinity with respect to the ground reward because there’s a limit to how much we can abstract away then. For instance, in the Antmaze experiment, if the ground reward is dependent on the joints' states then we wouldn't be able to abstract beyond the expert features. Therefore, we estimate the average cost of executing an option when it start at $s$ and ends in $s'$ (averaged over possible trajectories---see ground reward definition). When we move to the abstract level, given an abstract state we don't know exactly in what ground state we would be, hence, the best we can do is to estimate the average ground reward under the grounding function, and in turn we preserve the mean value with respect to the grounding function.

We thank you for pointing this out and we have updated our writing to reflect this.

• The dynamics preserving is, in fact, strong. However, this is similar to what current MBRL methods use to build models with latent variables. It means that $\phi(s)$ has to be a sufficient statistic to predict the next state (e.g., reconstruction based representation learning). In the worst case, in which everything in $s$ is necessary to predict the next state $s'$, then we would expect $\phi$ to reduce to the identity. However, with a structured set of options, we expect that some features of the state are controlled and others won't be (e.g. ant's position in the maze). Therefore, because we do not know which state trajectory was taken by the option to get from $s\rightarrow s'$, then some details of the state will become more and more unpredictable. For instance, we would need to model the trajectory of the ant's joints to be able to predict the final pose. Bu,t by having options that control the displacement, given that we preserve information about where the ant starts we can predict where it will end. We leverage this to abstract details away from our ground state.

• Discrete vs. Continuous: Although there are, indeed, discrete problems that are important, as previous work has tackled, not all problems we might have interest in solving are naturally discrete. For instance, the navigation problems in our experiments show that there is a continuity of $(x,y)$ positions to navigate. While discretization is definitely possible, it is not necessarily the most useful representation. Moreover, we also show that our approach is a strict generalization of the strong subgoal property that allows to build discrete models (as done for PDDL in previous work).

• Dataset : yes! to build a good model we need a good coverage of the state-option space. In practice, what we do is to initialize the dataset with data from trajectories resulting from executing a random policy and, then continuously, enhance the dataset with trajectories obtained by the learning agent. This does not mean the agent uses these samples to learn, they are only used to update the model and the agent still plans in abstract imagination.

Thank you to the authors for responding. I believe there are some issues with unconvincing experiments brought up by a few reviewers, which are as of yet unresolved. I'll keep my score the same but encourage the authors to submit a new draft of the work to another venue. The required changes are significant enough that they warrant another round of reviews, imo.

Thank you for your comments! Below the answers to your questions:

• Markov Abstractions: Markov abstractions deal with finding compression functions that preserve Markovianity at the primitive action level in both transition dynamics and rewards. That is, as much compression as possible at the primitive action level and task level. When we introduce temporal abstractions, we lose information by only observing (s, o, s’) (as oppose to the entire trajectory generated by the option execution). Therefore, some details of the state become less and less predictable. We leverage that to abstract beyond the point of Markovianity (as defined in the original MDP) and we preserve Markovinity at a higher time scale. However, in that process we lose information about the ground state (given an abstract state we cannot know exactly in what ground state we would be in) and, hence, we cannot fully predict the ground reward function. And we don’t want to preserve the ground reward predictability! because we want to use our abstract model for new tasks that we can solve with our options.

• How is the dynamics preserving abstraction defined in Def. 3.6 different from the Markovian abstractions defined in [Srivastava et al. 2016]? Thank you for pointing this reference out! I believe this shares the spirit of abstraction we are proposing in this paper: we can introduce impreciseness without being incorrect. Abstracting beyond the Markovianity of the ground MDP, introduces impreciseness on what ground state will be when an option is executed in the real environment, but preserves the relevant information to execute the next option. For instance, in the antmaze experiment, the model preserves knowledge of the ant’s position in the maze but cannot predict precisely with what pose the agent will reach such position.

Thank you for pointing all these issues out! We have updated the notation all across the paper and hopefully now it should be clearer. We believe this has helped improve a lot the framework's exposition!

1) and 2) We use $T(s’|s,o) = Pr(s’|s,o)$. (we have updated this in the notation to use only T). Moreover, we use the expected-length model of options. Hence, what we’re saying is that we execute the option as a black box, and observe how long it took ($\tau$) and where it ended at s’, and we model the option’s mean duration and the observed transition probabilities $Pr(s’|s, o)$ independently. Hence, the effect of the initiation sets and termination conditions are subsumed by the model of options.

3) We have corrected the notation. This was a typo! However, the option’s duration maps to a real because we are modeling the average duration when the option starts executing at s (over possible trajectories the option can take from $s\rightarrow s’$).

4) We have removed this confusing definition because of impreciseness. However, we meant to say that we only consider policies that can be learned in the abstract model. Because in the abstract model we don't want to model $s$ fully then we need a restricted class of policies that only uses the information available at the abstract level.

• Model-free vs. Model-based: It is worth pointing out that the comparison between DDQN on the ground MDP and DDQN in the abstract MDP shows that “losing” the information from a Markov state does not impede the DDQN agent finding a optimal policy (the DDQN is only ever trained with trajectories from the abstract model). This is because the retained information is the only relevant information. We, however, understand the concern about the evaluation between model-free and model-based comparisons. Currently, we are working on a Dreamer-like baseline that we will add to the draft asap.

• Environment choice: It is true that we can choose more intricate domains. However, I would argue that our Antmaze domain (even though not egocentric as the one referenced) shows something interesting: the state vector is an expert feature vector that is Markov by design. However, introducing skills for navigation allow us to build a model that abstract away (automatically) even more details and let the agent concentrate in the only important features for maze navigation (the ant's position in the maze).

• Dynamics-preserving vs. Bisimulation (Model-preserving abstractions): The main difference is that our dynamics-preserving abstraction does not impose conditions over the reward function. Our main desiderata is that we can re-use the abstract model to plan for different tasks and, therefore, the reward function of the ground MDP will be relevant to measure the cost (negative reward) of executing an option. This quantity, however, may be low-level trajectory-dependent and, hence, we might not be able to abstract our state beyond the Markov state of the ground MDP (e.g. we need to know the joints' states to compute the energy cost). Therefore, we decided to preserve only our ability to predict dynamics, and estimate the base reward of executing an option as the average reward accumulated (over the possible trajectories the option can generate). Hence, bisimulation can help reduce a high-dimensional observation (e.g. pixels) to a Markov state for the ground MDP. But to reduce our state even more, we use options and the dynamics-preserving condition.

• Missing references: We have added this references to our related works sections! thanks!

Minor questions:

• Yes! $\mathcal{L}^T_{\theta, \phi}$ affects the learning of $f_\phi$ (corrected this sentence in the paper, thanks!)
• There's not difference between figure 7 and 8 (removed).