Consciousness-Inspired Spatio-Temporal Abstractions for Better Generalization in Reinforcement Learning
Authors: Harry Zhao
Safa Alver
Harm van Seijen
Romain Laroche
Doina Precup
Yoshua Bengio
Venue: ICLR 2024 poster (2024)
Keywords: Reinforcement Learning
Planning
Neural Networks
Temporal Difference Learning
Generalization
Deep Reinforcement Learning
Primary area: reinforcement learning
TL;DR: Planning for better generalization by using abstraction in both space and time

Abstract: Inspired by human conscious planning, we propose Skipper, a model-based reinforcement learning framework utilizing spatio-temporal abstractions to generalize better in novel situations. It automatically decomposes the given task into smaller, more manageable subtasks, and thus enables sparse decision-making and focused computation on the relevant parts of the environment. The decomposition relies on the extraction of an abstracted proxy problem represented as a directed graph, in which vertices and edges are learned end-to-end from hindsight. Our theoretical analyses provide performance guarantees under appropriate assumptions and establish where our approach is expected to be helpful. Generalization-focused experiments validate Skipper’s significant advantage in zero-shot generalization, compared to some existing state-of-the-art hierarchical planning methods.

Published as a conference paper at ICLR 2024
CONSCIOUSNESS-INSPIRRR) SPATIO-TEMPORAL AB- STRACTIONS FOR BETTER GENERALIZATION IN REIN- FORCEMENT LEARNING
Mingde Zhao Safa Alver Harm van Seijen', Romain Laroche, Doina Precupl.3.5 Yoshua Bengio2.3
McGill University, Université de Montréal, Mila, 'Sony AI, GGoogle DeepMindd (mingde, zhao, sa safa. @ megill 11. ca, harm, vansei jenesony..n.oo.. com
romain. laroche@gmaal com, dprecupêcs megill ca, yoshua wwnggoemingenneo.oooooo..... quebec
ABSTRACT
Inspired by human conscious planning, we propose Skipper, a model-based rein- forcement learning framework utilizing spatio-temporal abstractions generalize better in novel situations. It automatically decomposes the given task into smaller, more manageable subtasks, and thus enables sparse decision-making and focused computation on the relevant parts of the environment. The decomposition relies on the extraction of an abstracted proxy problem represented as a directed graph, in which vertices and edges are learned end-to-end from hindsight. Our theoret- ical analyses provide performance guarantees under appropriate assumptions and establish where approach is expected to be helpful. Generalization -occuse experiments validate Skipper's significant advantage in zero-shot generalization, compared to some existing state-of-the-ar hierarchical planning methods.
INTRODUCTION
Attending to relevant aspects in both time and space, human conscious planning breaks down long- horizon tasks into more manageable steps, each of which can be narrowed down further. Stemming from consciousness in the first sense (C1) (Dehaene et al., 2020), this type of planning focuses attention on mostly the important decision points (Sutton et al. 1999) and relevant environmental factors linking the decision points (Tang et al.) 2020), thus operating abstractly both in time and in space. In contrast, existing Reinforcement Learning (RL) agents either operate solely based mostly relatively shortsighted
on intuition (model-free methods) are limited to reasoning plans (model-based methods) Kahneman (2017). The intrinsic limitations constrain the real-world application of RL under glass ceiling formed by challenges of generalization.
Building on our previous work on conscious planning (Zhao et al.) 2021), we take inspirations to develop planning agent that automatically decomposes the complex task hand into smaller subtasks, by constructing abstract "proxy" problems. A proxy problem is represented as a graph where 1) the vertices consist of states imagined by generative model, corresponding to sparse decision points; and 2) the edges, which define temporally eextedd transitions, are constructed by focusing on small amount of relevant information from the states, using an attention mechanism. Once proxy problem is constructed and the agent solves it to form a plan, each of the edges defines a new sub-problem, on which the agent will focus next. This divide-and-conquer strategy allows constructing partial solutions that generalize better to situations, while also giving the agent flexibility to construct abstractions necessary for the problem at hand. Our theoretical analyses establish guarantees on the quality of the solution to the overall problem.
We also examine empirically advantages of teaf--taniing---teee after using only few training tasks. We show through detailed controlled experiments that the proposed framework, named Skipper, in most cases performs significantly better in terms of zero- shot generalization, compared to the baselines and some state-of-the-art Hierarchical Planning (HP) methods (Nasiriany et al. 2019 Hafner et al. 2022).
generalization of our method
Work largely done during Mingde, Harm and Romain's 's time at Microsoft Research Montreal. Source code of experiments available pnttps://giithhe com/mi 1 a-aiiggg i a a pper

---

Published as a conference paper at ICLR 2024
2 PRELIMINARIES
Reinforcement Learning & Problem Setting. An RL agent interacts with an environment via a sequence of actions to maximize its cumulative reward. The interaction is usually modeled Markov Decision Process (MDP) M = A, P. R, d. 7), where S and A are the set of states and actions, P P x A DS is the state transition function, R. S x A x S - R is the reward function, d : S + Dist(S) is the initial state distribution, and Y E [0, 1] is the discount factor. The agent needs to learn a policy : S + Dist(A) that maximizes the value of the states, i.e. the expected discounted cumulative reward E,, P>
Th of R(S,, Ar, 1 ~ d]. where Ti denotes the time step at which the episode terminates. A value estimator Q : S x A + IR can be used to guide the search for a good policy. However, real-wort problems are typically partially observable, meaning that at each time step t, instead of states, the agent receives an observation rz+1 X, where X' is the observation space. The agent then needs to infer the state from the sequence of observations, usually with state encoder.
One important goal of RL is to achieve high (generalization) performance on evaluation tasks after learning from limited number of training tasks, where the evaluation and training distributions may differ; for instance, policy for a robot may need to be trained in a simulated environment for safety reasons, but would need be deployed on physical device, a setting called sim2real. Discrepancy between task distributions is often recognized as major reason why RL agents are yet to be applied pervasively in the real world (Igl al., 2019). To address this issue, in this paper, agents are trained on small set of fixed training tasks, then evaluated in unseen tasks, where there are environmental variations, but the core strategies needed to finish the task remain consistent. To generalize well, the agents need to build learned skills which capture the consistent knowledge across tasks.
Deep Model-based RL. Deep model-based RL uses predictive or generative models to guide the search for policies (iilver et al., 2017). In terms of generalization, rich models, expressed by Neural Networks (NNs), may capture generalizable information and infer latent causal structure. Back- ground planning agents e.g., Dreamer (Hafner et al., 2023)
model data generator to im- prove the value estimators and policies, which executes in background without directly engaging in the environment (Sutton T991). These agents do not improve on the trained policy decision time. In contrast, decision-time planning agents e.g., MuZero (Schrittwieser et af. 2020) and PlaNet (Hafner al. (2019) actively use models at decision time to make better decisions. Recently, [Alver & Precup (022) suggests that the latter approach provides better generalization, aligning with ob- servations from cognitive behaviors (Momennejad et al. 2017).
Options& & Goal-Conditioned RL. Temporal abstraction allows agents to use sub-policies, and to model the environment over extended time scales, achieve both better generalization and the divide and conquer of larger problems. Options and their models provide formalism for temporal abstraction in RI. (Sutton et al., 1999]. Each option consists of an initiation condition, a policy, and a termination condition. For any set of options defined on an MDP, the decision process that selects only among those options, executing each to termination, is a Semi-MDP (SMDP) (Sutton) let al. 1999 Puterman 2014), consisting of the set of states S, the set of options O, and for each state-option pair, an expected return, and a joint distribution of the next state and transit time. In this paper, we focus
conlicond options, where the initiation set covers the whole state space S. Each such option is a tuple o = (T, 3), where т : S + (ist(AA) is the (intra-)option policy and B S + {0, 13 indicates when goal state is reached. Hindsight Experience Replay (HER) (Andrychowicz et al. 2017) is often used to train goal-conditioned options by sampling a transition (Te, 0r, Tt+1, 1+) together with an additional observation の from the same trajectory, which is re-labelled as a "goal".
3 Skipper: SPATIALLY & TEMPORALLY ABSTRACT PLANNING
In this section, we describe the main ingredients of Skipper a framework that formulates a proxy problem for given task, solves this problem, and then proceeds to "fill in" the details of the plan.
3.1
PROXY PROBLEMS
Proxy problems are finite graphs constructed at decision-time time, whose vertices are states and whose directed edges estimate transitions between the vertices, shown in Fig. [
We call the states
2

---

Published as a conference paper at ICLR 2024
selected to be vertices of the proxy problems checkpoints, to differentiate from other uninvolved states. The current state is always included as one of the vertices. The checkpoints are proposed by generative model and represent some states that the agent might experience in the current episode, often denoted as 5 in this paper. Each edge is annotated with estimates of the cumulative discount and reward associated with the transition between the connected checkpoints; these estimates learned over the relevant aspects of the environment and depend on the agent's 's capability. As the low-Ievel policy implementing checkpoint transitions improves, the edges strengthen. Planning in proxy problem is temporally abstract, since the checkpoints act as sparse decision points. Estimating each checkpoint transition is spatially abstract, as an option corresponding to such task would base its decisions only some aspects of the environment state (Bengio 2017: Konidaris & Barto 2009). to improve generalization as well as computational efficiency (Zhao et al. (0221)
A proxy problem can he viewed as a deterministic SMDP, where each directed edge is implemented as a checkpoint-conditioned op- tion. It can be fully described by the discount and reward matrices, I and V*, where 72/ and Dij are defined as:
1 = E, [ |So Si, STI 53]]
" = E - RR 85. ST. =
By planning with 1 and V* , e.g. using SMDP value iteration (Sut- lon et al.) 1999), we can solve the proxy problem, and form a jumpy plan to travel between states in the original problem. If the proxy problems can be estimated well, the obtained solution will be of good quality, as established in the following theorem:
Proxy Problem on a World Navigation Task: the MDP of the origi- nal problem is in gray and the terminal states marked with squares. An agent needs to get from the (filled red) position. to the goal (filled green). Distant goals be reached by lever- aging proxy problem with 12 checkpoints (outlined orange).
Figure 1:
Let Ir be the SMDP policy (high-level) and be the low-level policy. Let V and 1 denote learned estimates of the SMDP model. If the estimation accuracy satisfies:
Theorem
[0 by 109 < 6 (1 - 7)
og < Umax < (1 -7)mmax
and Vi,j-
(3)
Then, the estimated value of the composite uuom (s) accurate up to error terms linear in cy and eyy
1-1
(pon(s) = - 0x ( O
II
9x (1
Upon + Ev + E. Umax + o(co + G))
R+1
1-y (1
k=O
E=
where 0z (8:(88)) = 07 and im (5,(s,) = 97 and Umax denotes the maximum value.
The theorem indicates that once the agent achieves high accuracy estimation of the model for the proxy problem and a near-optimal lower-level policy w, it converges toward optimal performance (proof in Appendix [D.2]. The theorem also makes no assumption difficult to learn good π for far away targets. Despite the theorem's is generality, in the experiments, we limit ourselves navigation tasks with sparse rewards for reaching goals, where the goals are included as permament vertices in the proxy problems. This is a case where the accuracy assumption can be met non- trivially,, i.e., while avoiding degenerate proxy problems whose edges involve no rewards. Following Thm. I we train estimators for e and Yrr and refer to this as edge estimation.
since it would likely be
3.2 DESIGN CHOICES
To implement planning over proxy problems, Skipper embraces the following design choices:
Decision-time planning is employed due to its ability to improve the policy in novel situations;
Spatio-temporal abstraction: temporal abstraction breaks down the given task into smaller ones, while spatial abstraction over the state features improves local learning and generalization;
We use "spatial abstraction" to denote specifically the hehavior of constraining decision-making to the relevant environmental factors during an option. Please check Section 4 for discussions and more details.
3

---

Published as a conference paper at ICLR 2024
Higher quality proxies: we introduce pruning techniques improve the quality of proxy problems;
Learning end-to-end from hindsight, off-policy: to maximize sample efficiency and the case of training. we propose to auxiliary (off-) joff-) methods for edge estimation, and learn context- conditioned checkpoint generation, both from hindsight experience replay;
Delusion suppression:
we propose delusion suppression technique to minimize the behavior of chasing non-existent outcomes. This is done by exposing edge estimation to imagined targets that would otherwise not exist in experience.
3.3 PROBLEM 1: EDGE ESTIMATION
First, we discuss how to estimate the edges of the proxy problem, given a set of already generated checkpoints. Inspired by conscious information processing in brains (Dehaene al. 2020) and existing approach in Sylvain et al. 2019), we introduce a local perceptive field selector, d, consisting of attention bottleneck that (soft- (soft-)selects Jselects the top-k local segments of the full state (e.g. a feature map by typical convolutional encoder); all segments of the state compete for the & attention slots, Lee. irrelevant aspects of state are discouraged or discarded, to form a partial state representation (Mott et al.) 2019 Tang et al., 2020; Zhao et al., 2021] Alver & Precup, 2023). We provide an example in Fig. (see purple parts). Through o, the auxiliary estimators, to be discussed soon, force the bottleneck mechanism to promote aspects relevant to the local estimation of connections between the checkpoints. The rewards and discounts are then estimated from the partial state o (S), conditioned on the agent's policy.
3.3.1 BASIS FOR CONNECTIONS: CHECKPOINT-ACHIEVE POLICY
The low-level policy maximizes an intrinsic reward, s.4. the target checkpoint S can be reached. The choice of intrinsic reward flexible; for example, one could use reward of +1 when S within small radius of 8 according to some distance metric, or use reward-respecting intrinsic rewards that enable more sophisticated behaviors, as in (Sutton et al. 2022). In the following, for simplicity, we will denote the checkpoint-achhieemnee condition with equality: S 8
3.3.2 ESTIMATE CONNECTIONS
We learn the connection estimates with auxiliary reward signals that are designed to be not task- specific (Zhao et al. 2019). These estimates are learned using C51-style distributional RL, where the output of each estimator takes the form of a histogram over scalar support (Dabney et al.) 2018).
Cumulative Reward. The cumulative discounted task reward ufy learned by policy evaluation on an auxiliary reward that is the same as the original task reward everywhere except when reaching the target. Given hindsight sample (r,, at, F++1, F1+1, 2 and the corresponding encoded sample
(st, de, Tr+1, St+1, s0), we train Vn with KL-divergence as follows:
[[(s; ae, 8t+1) + 70x (o(((1)))))))))))))))))))))))))) ar+ (o if 81+1 # 8
(+ ()) RR(s; at; 8t+1)
(4)
if 8 8
where (s) is the spatially-abstractedd from the full state s and 4441 x-1o((s+1)) (s)))
Cumulative Distances Discounts. With C51 and uniform supports, the cumulative discount lead- ing to &e under unfortunately more difficult learn than Vg, since the prediction would be heavily skewed towards 1 if s 1. Yet, we can instead effectively estimate cumulative (truncated) distances (or trajectory length) under æ. Such distances can be learned with policy evaluation, where the auxiliary reward is +1 on every transition, except at the targets:
+ D= (a((111)) (1)))
s1+1 7 8
(x(a(s)), (s)o(s9)) (r(((ss).ooll((())
1
if st+1 = S
if 81+1 is terminal and s1+1 # s
where 4+41 ((11(((411)) The cumulative discount is then recovered by replacing the support of the output distance histogram with the corresponding discounts. Additionally, the learned distance is used to prune unwanted checkpoints simplify the proxy problem, well prune far- fetched edges. The details of pruning will be presented shortly.
Please refer the Appendix D.I for the properties of the learning rules for 0z and fm.
4

---

Published as a conference paper at ICLR 2024
spatial abstraction
eNewtoon
checkpoint policy
Hargeto
seleet
reward estimator V
partial ofs)
discount estimator
edges
splitter
partial desc."
fuser
Aprunn
full state
z(6)
wv
evertices
context cls)
checkpoint generatar
observation
state encoding
teniporal abstraction
proxy problem planning
Figure 2: Skipper Framework: 1) Partial states consist of few local fields, soft-selected via top-k attention
(Gupta al. 2021). Skipper's edge estimations and low-level behaviors a are based on the partial states. 2) The checkpoint generator learns by splitting the full state into context and partial descriptions, and fusing them to reconstruct the input. It imagines checkpoints by sampling partial descriptions and combining them with the episodic contexts; 3) We prune the vertices and edges of the denser graphs to extract sparser proxy problems. Once a plan is formed, the immediate checkpoint target used condition the policy. In the proxy problem example, blue edges are estimated to be bidirectional and red edges have the other direction pruned.
3.4 PROBLEM 2: VERTEX GENERATION
The checkpoint generator aims to directly imagine the possible future states without needing to know how exactly the agent might reach them nor worrying about if they are reachable. The feasibility of checkpoint transitions will be abstracted by the connection estimates instead.
To make the checkpoint generator generalize well across diverse tasks, while still being able to capture the underlying causal mechanisms in the environment (a challenge for existing model-based methods) (Zhagg et al. 2020), we propose that the checkpoint generator learns to split the state representation into two parts: episodic context and partial description. In navigation problem, for example, as in Fig. 2 a context could be a representation of the map of gridworld, and the partial description be the 2D-coordinates of the agent's location. In different contexts, the same partial description could correspond to very different states. Yet, within the same context, we should be able to recover the same state given the same partial description.
As shown in Fig. 2 this information split is achieved using two functions: the splitter Egz, which maps the input state S into representation of context c(S) and a partial description =(S), well as the fuser + which, when applied to the input (c, =), recovers S. In order to achieve consistent context extraction across states in the same episode, at training time, we force the context to be extracted from other states in the same episode, instead of the input.
We sample in hindsight a diverse distribution of target encoded (full) states S given any current S.. Hence, we make the generator a conditional Variational AutoEncoder ( (VAE) (Sohn et al. 2015) which learns a distribution p(S (CSs)) Σ, P(S (C(S),, 2 2))(=)((S)) where C(S,) is the extracted context from S, and are the partial descriptions. We train the generator by minimizing the evidence lower bound on (St, 50) pairs chosen with HER.
Similarly to Hafner et al. 2023). we constrain the partial description bundle of binary variables and train them with the straight- gradients (Bengio al. 2013). These binary latents can be easily sampled or composed for generation. Compared models such as that in Director (Hafner cet al. 2022). which generates intermediate goals given the on on-policy trajectory, ours can generate and handle more diverse distribution of states, beneficial for planning in novel scenarios.
3.4.1 PRUNING
In this paper,we we limit ourselves only to checkpoints from a return- conditional generation model, leaving the question of how to improve the quality of the generated checkpoints for future work. Without learning, the proxy problem can be improved by making more sparse, and making the proxy problem vertices more evenly spread in state space. To achieve this, we propose prun- ing algorithm based on k-medoids clustering (Kaffman & Rousseeuw 1990), which only requires pairwise distance estimates between states. During proxy problem construction, we first sample larger number of checkpoints, and then cluster them and select the centers (which are always real states instead of imaginary weighted sums of state representations).
5

---

Published as a conference paper at ICLR 2024
Notably, for sparse reward tasks, the generator cannot guarantee the presence of the rewarding checkpoints in the proposed proxy problem. We could remedy this by explicitly learning the gener- ation of the rewarding states with another conditional generator. These rewarding states should be kept as vertices (immune from pruning).
In addition pruning the vertices, we also prune the edges according to distance threshold, i.e., all edges with estimated distance over the threshold are deleted from the complete graph of the pruned vertices. This biases potential plans towards shorter-length, smaller-scale sub-problems, as far-away checkpoints are difficult for π to achieve, trading optimality for robustness.
3.4.2 SAFETY & DELUSION CONTROL
Model- HRL agents can be prone blindly optimizing for objectives without understanding the consequences (Langosco et al. 2022, Paolo et al. (022). We propose a technique to suppress delusions by exposing edge estimation to potentially delusional targets that do not exist in the expe- rience replay buffer. Details and examples are provided in the Appendix.
4 RELATED WORKS & DISCUSSIONS
Temporal Abstraction. Resembling attention in human consciousness, choosing checkpoint tar- get selection towards certain decision points in the dimension of time, i.e. form of temporal abstraction. Constraining options, Skipper learns the options targeting certain "outcomes", which dodges the difficulties of option collapse (Bacon et al.] (2017) and option outcome modelling by de- sign. The constraints indeed shift the difficulties to generator learning (Silver & Ciosek] 2012: TTngg Salakhutdinov 2019). We expect this to entail benefits where states are easy to learn and gener- ate, and in stochastic environments where the outcomes of unconstrained options are difficult to learn. Constraining options was also investigated in Sharma et al. (2019) in an unsupervised setting.
Spatial Abstraction is different from "state abstraction" (Sacerdoti 1974 Knoblock 1994). which evolved to be a more general concept that embraces mainly the aspect of state aggregation, iec. state space partitioning (Li et al., 2006). Spatial abstraction, defined to capture the behavior of attention in conscious planning in the spatial dimensions, focuses on the within-state partial selection of the environmental state for decision-making. It corresponds naturally to the intuition that state represen- tations should contain useful aspects of the environment, while not all aspects are useful for a par- ticular intent. Efforts toward spatial abstraction are traceable to early hand-coded proof-of-concepts proposed in e.g. Dietterich| (2000). Until only recently, attention mechanisms had primarily been used to construct state representations in model-fres agents for sample efficiency purposes, with- out the focus on generalization (Mott et al. 2019: Manchin et al. 2019 Tang et al. 2020). In Fu jet al. (202T): Zadaianchuk al. (2020); Shah et al. (2021), 3 more recent model-based approaches, spatial abstractions are attempted to remove visual distractors. Concurrently, emphasizing on gener- alization, our previous work (Zhao et al. 202T) used spatially-abstract partial states in decision-time planning. We proposed an attention bottleneck to dynamically select subset of environmental en- tities during the atomic-step forward simulation, without explicit goals provided as in Zadaianchuk let al. (2020). Skipper's checkpoint transition is step-up from our old approach, where we now show that spatial abstraction, an overlooked missing flavor, is as crucial for longer-term planning as temporal abstraction (Konidaris & Barto 2009).
Task Abstraction via Goal Composition The early work McGovern & Barto (2001) suggested to
use bottleneck states subgoals abstract given tasks into manageable steps. Nair al. (2018):
Florensa et al.) 2018) use generative model to imagine subgoals while AhileeEysenbache et al. (2019 search directly on the experience replay. In Kim et al. (2021), promising states explore are generated and selected with thortest-path algorithms. Similar ideas have been attempted for guided exploration (Erraqabi et al., 202T. Kuikarni al. (016). Similar to Hatner et al. (2022), Czechowski et al. (2021) generate fixed-steps ahead subgoals for reasoning tasks, while Bagaria et al. (2021) augments the search graph by states reached fixed-steps ahead. Nasiriany et al. (2019): Xie et al. (2020); Shi
et al. (2022) employ CEM plan a chain of subgoals towards the task goal (Rubinstein) 1997b).
Skipper utilizes proxy problems to abstract the given tasks via spatio-tempora abstractions (Bagaria et al. 2021). Checkpoints
be seen as sub-goals that generalize the notion of "landmarks' or "waypoints" Suttoon et al. (1999); Diett)ercihttt (2000) (Svinov et al. (2018). Zhang et al. (202T) used latent landmark graphs as high-level guidance, where the landmarks sparsified with weighted

---

Published as a conference paper at ICLR 2024
sums in the latent space to compose subgoals. In comparison, checkpoint pruning selects subset of generated states, which is less prone to issues created by weighted sums.
Planning Estimates. Zhang et al. (2021) propose a distance estimate with explicit regression. With TDMs (Pong et al. 2018). LEAP (Nasiraay et al. 2019) embraces a sparse intrinsic reward based distances to the goal. Contrasting with distance estimates, there is empirical ev- idence of TDMs' compatibility with stochasticity and terminal states. Notably, [Eyeebbach et al.] (2019) employs a similar distance learning scheme learn the shortest path distance between states found in the experience replay; while our estimators learn the distance conditioned on evolving policies. Such aspect was also investigated in Nachum al. (2018).
Decision- Time HP Methods. Besides LEAP (Nasiriany ct al. 2019), decision time planning with evolutionary algorithms was investigated in Nair & Finn 2000
Hafner et al. (2019).
EXPERIMENTS
As introduced in Sec. 2 our first goal test the zero-shot generalization ability of trained agents. To fully understand the results, it is necessary to have precise control of the difficulty of the training and evaluation tasks. Also, to validate if the empirical performance of our agents matches the formal analyses (Thm. D. we need to know how close to the (optimal) ground truth our edge estimations and checkpoint policies are. These goals lead to the need for environments whose ground truth information (optimal policies, true distances between checkpoints, etc) can be computed. Thus, we base our experimental setting on the SaniGrid-Baby-ll framework (Chevalier- Chevaaler-Boiivert et al.) 2018b Hui et al. (2200). Specifically, we build the experiments used in previous works (Zhao et al. [021] Alver & Precup/ 2022): the agent needs to navigate to the goal from its initial state gridworlds tilled with terminal lava traps generated randomly according to difficulty parameter, which controls their density. During evaluation, the agent is always spawned at the opposite side from the goals. During training, the agent's is position uniformly initialized speed up training. We provide results for non- unifornfoo training initialization in the Appendix.
These fully observable tasks prioritize on the challenge of reasoning over causal mechanisms over learning representations from complicated observations, which is not the focus of this work. Across all experiments, we sample training tasks from an environment distribution of difficulty 0.4: each cell in the field has probability 0.4 to be filled with lava while guaranteeing a path from the initial position to the goal. The evaluation tasks are sampled from a gradient of OOD difficulties 0.25, 0.35, 0.45 and 0.55, where the training difficulty acts as a median. To step up the long(er) term generalization difficulty compared to existing work, we conduct experiments done on large, 12 x 12 maze sizes, (see the visualization in Fig 2). The agents trained for 1.5 x 106 interactions.
We compare Skipper against two state-of of-the-art Hierarchical Planning (HP) methods: LEAP
(Nasiriayy et al. 2019) and Director (Hafner et al.) 2022). The comparative results include:
Skipper-once: Skipper agent that generates one proxy problem the start of the episode, and the replanning (choosing a checkpoint target based on the existing proxy problem) only triggers quick re-selection of the immediate checkpoint target;
Skipper-regen: A Skipper agent that re-generates a proxy problem when replanning is triggered;
modelfree: A model-free baseline agent sharing the same base architecture with the Skipper vari- ants prioritized distributional Double DQN (Dabney et al. 2018 Van Hasselt et al. 2016);
Director: A tuned Director agent (Hafner et al. 2022 fed with simplified visual inputs. Since Director discards trajectories that are not long enough for training purposes, we make sure that the same amount of training data gathered for the other agents;
LEAP: A re-implemented LEAP for discrete action spaces. Due to low performance, we replaced the VAE and the distance learning mechanisms with our counterparts. We waived the interaction costs for its generator pretraining stage, only showing the second stage of RI pretrainin..
Please refer the Appendix for more details and insights on these agents.

---

Published as a conference paper at ICLR 2024
5.1 GENERALIZATION PERFORMANCE
Fig. 3 shows how the agents' generalization performance evolves during training. These results are obtained with 50 fixed sampled training tasks (different 50s for each seed), representative configuration of different numbers of training tasks including {1, 25, 50, 100, oc B whose results are in the Appendix. In Fig- a), we observe how well an agent performs on its training tasks. an agent performs well here but badly in b), c). d) and e), e.g. the modelfree baseline, then we suspect that it overfitted on training tasks, likely indicating a reliance on memorization (Cobbe et al. 2020).
advantage in the generalization performance of the Skipper agents throughout training. We have also included significance tests and power analyses (Colas et al.) 2018 Patterson et al. 2023 in the Appendix, together with results for other training configurationss The regen variant exhibits dominating performance over all others. This is likely due to the frequent reconstruction of the graph makes the agent less prone to being trapped in low-quality proxy prob- lem and provides extra adaptability in novel scenarios (more discussions in the Appendix). During training, Skippers behave less optimally than expected, despite the strong generalization on evalu- ation tasks. As our ablation results and theoretical analyses consistently show, such a phenomenon composite outcome of inaccuracies both in the proxy problem and the checkpoint policy. One major symptom of an inaccurate proxy problem is that the agent would chase delusional targets. We address this behavior with the delusion suppression technique, to be discussed in the Appendix.
We observe a paatistically-ssinniiiei
3.0
Skipper-once
Skipper-regen
LEAP
Director
modeifree
0.8
100000
104
0.2
00
0.0
0.5
3.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
3,0
3.5
0.0
0.5
100
1.5
0.0
0.5
1,0
training stopu
training dagg
training stepe
training Mags
teatning stepa
(a) training. diff 0.4
1OOO,, diff 0.25 (c) OOD, diff 0.35
(d) OOD, diff 0.45
(e) OOD, diff 0.55
Figure 3: Generalization Performance of Agents During Training: the x-axes correspond to training
progress, while the aligned y-axes represent the success rate of episodes (optimal is 1.0). Each agent is trained with 50 tasks. Each data point is the average success rate over 20 evaluation episodes, and cach error bar (95% confidence interval) processed from 20 independent seed runs. Training tasks performance is shown in (a) while OOD evaluation performance shown (b). (c). (d), (e).
Better than the modelfree baseline, LEAP obtains reasonable generalization performance, despite the extra budget it needs for pretraining. In the Appendix, we show that LEAP benefits largely from the delusion suppression technique. This indicates that optimizing for path the latent space may be prone to errors caused by delusional subgoals. Lastly, we see that the Director agents suffer in these experiments despite their good performance in the single environment experimental settings reported by Hafner et al. (022). We present additional experiments in the Appendix to show that Director is ill-suited for generalization-fouur single environment configurations, but its performance deteriorates fast with more training tasks. This indicates poor scalability in terms of generalization, a limitation to its application in real-world scenarios.
settings: Director still performs well in
5.2 SCALABILITY OF GENERALIZATION PERFORMANCE
Like Cobbe al. (220)) investigate the scalability of the agents' generalization abilities across different numbers of training tasks. To this end, in Fig. 4 we present the results of the agents' final evaluation performance after training over different numbers of training tasks.
With more training tasks, Skippers and the baseline show consistent improvements in generaliza- tion performance. While both LEAP and Director behave similarly as in the previous subsection, notably, the modelfree baseline can reach similar performance as Skipper, but only when trained on a different task in each episode, which is generally infeasible in the real world beyond simulation.
oo training tasks mean that an agent is trained on different task for each episode. In reality, this may lead prohihitive costs in creating the training environment.
8

---

Published as a conference paper at ICLR 2024
Skipper-regen
LEAP
Director
modelfree
6.8
(aal tae
6.6
02
Skipper-onoe
60
100
23 tnadt
(a] training, diff 0.4
(b) OOD, diff 0.25
(c) OOD, diff 0.35
(d) OOD, diff 0.45
(e) OOD, diff 0.55
Figure 4: Generalization Performance of Agents on Different Numbers Training Tasks: each data point corresponding error bar (95% confidence interval) are based on final performance from 20 independent seed runs. Training task performance is shown in (a) while OOD performance is shown in (b), (c), (d), (c).
Notably, the Skipper agents as well as the adapted LEAP behave poorly during training when being trained only task, the split of context and partial information cannot be achieved. Training invalidates purpose of proposed cettraizzatinn-fcccsse checkpoint generator.
task
5.3 ABLATION & SENSITIVITY STUDIES
In the Appendix, we present ablation results confirming the effectiveness of delusion suppression, k-medoids pruning and the effectiveness of spatial abstraction via the local perception field. We also provide sensitivity study for the number of checkpoints in each proxy problem.
5.4 SUMMARY OF EXPERIMENTS
Within the scope of the experiments, we conclude that Skipper provides benefits for generalization; And it achieves better generalization when exposed to more training tasks;
From the content presented the Appendix, we deduce additionally that:
Spatial abstraction based the local perception field is crucial for the scalability of the agents;
Skipper performs well by reliably decomposing the given tasks, and achieving the sub-tasks bustly. Its performance bottlenecked by the accuracy of the estimated proxy problems as well as the checkpoint policies, which correspond to goal generalization and capability generalization, spectively, identified in Langosco et al. (2022). This matches well with our theory. The proposed delusion suppression technique (in Appendix) is effective in suppressing plans with non-existent checkpoints as targets, thereby increasing the accuracy of the proxy problems;
LEAP fails to generalize well within its original form and can generalize better when combined with the ideas proposed in this paper; Director may generalize better only in domains where long and informative trajectory collection is possible;
We verified empirically that, as expected, Skipper is compatible with stochasticity.
CONCLUSION & FUTURE WORK
Building on previous work on spatial abstraction (Zhao et al. 202I), we proposed, analyzed and validated Skipper, which generalizes its learned skills better than the compared methods, due to its combinned spatio-temporal abstractions. A major unsatisfactory aspect of this work is that we generated checkpoints at random by sampling the partial description space. Despite the pruning mechanisms, the generated checkpoints, and thus temporal abstraction, do not prioritize the pre- dictable, important states that matter form meaningful plan (Simşek & Barto 2004). We would like to continue investigating the possibilities along this line. Additionally, we would like to explore other environments where the accuracy assumption (in Thm. I can meaningfully hold, i.e. beyond sparse reward cases.

---

Published as a conference paper at ICLR 2024
7 REPRODUCIBILITY STATEMENT
The results presented in the experiments are felluc-eproducible with the open-sourced repository https////gthuubb..ommmmmm SkipeS
ACKNOWIEDGEMENTS
Mingde is grateful for the financial support of the FRQ (Fonds de recherche du Québec) and the col- laborative efforts during his time at Microsoft Research Montreal. We want to thank Riashat Islam, Ziyu Wang for the discussions in the initial stages of the project. We appreciate the computational resources allocated to us from Mila and McGill university.
10

---

Published as a conference paper at ICLR 2024
REFERENCES
Safa Alver and Doina Precup. Understanding decision-tinie time vs. background planning in model-base reinforcement learning. arXiv preprint arXiv:2206. 08442, 2022.
Safa Alver and Doina Precup. Minimal vatue-equivalent partial models for scalable and robust plan- ning in lifelong reinforcement learning. In Sarath Chandar, Razvan Pascanu, Hanie Sedghi, and Doina Precup (eds.), Proceedings of The 2nd Conference on Lifelong Learning Agents, volume 232 of Proceedings of Machine Learning Research, pp. 548-567. PMLR, 22-25 Aug 2023. URL https::/proceediing
wtw.s/ http///////aaaaaaaaaaaa html
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojciech Zaremba. Hindsight experience re- play. Advances in neural information processing systems, 30, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.064400, 2016.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
Akhil Bagaria, Jason K Senthil, and George Konidaris. Skill discovery for exploration and planning using deep skill graphs. In International Conference on Machine Learning, pp. 521-531. PMLR, 2021.
Yoshua Bengio. The consciousness prior. arXiv, 1709.08568, 2017. http:://rxiv.org/// 1709. 08569]
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint:1308.3433 2013..
Maxime Chevalier Boisserr,, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. International Conference on Learning Representations, 2018a.
fttt:::/aaxiv..rgee haagg//b////// / 11810 082722
Maxime Cheveeier-ooisvert,, Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for openai gym. GitHuu repository, 2018b. https://github.. com max mecb/
Sym-minigrid
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pp. 2048- 2056. PMLR, 2020.
Cédric Colas, Olivier Sigaud, and Pierre- Yves Oudeyer. How many random seeds? statistical power analysis in deep reinforcement learning experiments, 2018.
Özgür Şimşek and Andrew G. Barto. Using relative novelty to identify useful temporal abstrac- tions in reinforcement learning. In Proceedings of the Twenty-Firs/ International Conference on Machine Learning. ICML '04, pp. 95, New York, NY, USA, 2004. Association for Computing Machinery. ISBN 158 1 138385.
1015353
Konrad Czechowski, Tomasz Odrzygózdž. Marek Zbysiński, Michat Zawalski. Krzysztof Olejnik, Yuhuai Wu, Fukasz Kucinski, and Piotr Miloš. Subgoal search for complex reasoning tasks. arXiv preprint:2108. /204, 2021.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 32, 2018
Stanislas Dehaene, Hakwan Lau, and Sid Kouider. What consciousness, and could machines have it? Science, 358, 2020. URL Stttps://sclence.cceeeeeeee wwtp:///eeeeeaaaccoooooooooooooe org/ wrtp://ooo/oooooo..e....

---

Published as a conference paper at ICLR 2024
Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decompo- sition. Journal of artificial intelligence research, 13:227-303, 2000.
Akram Erraqabi, Mingde Zhao, Marlos Machado, Yoshua Bengio, Sainbayar Sukhbaatar, Ludovic Denoyer, and Alessandro Lazaric. Exploration-driver representation learning in reinforcement learning. In /CML 2021 Workshop Unsupervised Reinforcement Learning, 2021.
Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine Search on the replay buffer: Bridging planning and reinforcement learning. In Н. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neu- ral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL
httpo::/prrcceeinsss.eeeeiii hwtp://ss...iiii..ccccc.cooooe httpteerrf.leeppppeeeeeeea Filess/pper/2019/E1ee Sc48Ef18e0a4//aaattalao 2oofffleeccs/aaaftddeeeeeeeeeee2 Paper, pdf
Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. In International conference on machine learning, pp. 1515-1528.
PMLR, 2018.
Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference Machine Learning, pp. 3480-3491, PMLR, 2021.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods, 2018.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. arXiv preprint:2 1ttpprin::21000689900 2021.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 25655. PMLR, 2019.
Danijar Hafner, Kuang-Huu Huei Lee, Ian Fischer, and Pieter Abbeel. Deep hierarchical planning from pixels. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Ad- vances in Neural Information Processing Systems, 2022. httptppp:://peeeeeeeee..... Forum?id=w2k634 5y9_d
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint:2301.0 14044, 2023.
Chad Hogg, Ugur Kuter, and Héctor Muñoz Avila. Learning hierarchical task networks for non- deterministic planning domains. In Proceedings of the 21st International Joint Conference on Artificial Intelligence, IICAI'09, pp. 1708-1714, San Francisco, CA, USA, 2009, Morgan Kauf- mann Publishers Inc
David Yu-Tung Hui, Maxime Chevalier- Chevalier-Boisvert,, Dzmitry Bahdanau, and Yoshua Bengio. Babyai 1.1, 2020.
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek. Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization reinforcement learning with selective noise injection and information bottleneck. Advances in neural information processing systems, 32, 2019.
Daniel Kahneman. Thinking, fast and slow. 2017.
Leonard Kaufman and Peter Rousseeuw. Partitioning Around Medoids (Program PAM), chapter 2, 125. pp. 68-125. John Wiley & Sons, Ltd, 1990. URL /tRLhhtss:///tt htttt:::://onnnnelbbarrry www.bbrrrrww................rree com/
doi htl/abb//1n /10 10200/90477733333 6801 ch2
Junsu Kim, Younggyo Sco, and Jinwoo Shin. Landmark- -uided subgoal generation in hierarchical reinforcement learning. arXiv 100 2021.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv
preprint: 1412.6990, 2014.
12

---

Published as a conference paper at ICLR 2024
Craig A Knoblock. Automatically generating abstractions for planning. Artificial intelligence, 68 (2):243-302, 1994.
George Dimitri Konidaris and Andrew G Barto. Efficient skill learning using abstraction selection. In IJCAI, volume 9, pp. 1107-1112, 2009.
Tejas D Kulkarni, Karthik Narasimhan, Ardavan Saeedi, and Josh Tenenbaum. Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
Lauro Langosco Di Langosco, Jack Koch, Lee D Sharkey, Jacob Pfau, and David Krueger. Goal mis- generalization in deep reinforcement learning. In International Conference on Machine Learn- ing, volume 162, pp. 12004-12019, 17-23 Jul 2022. URL
html
Lihong Li, Thomas J Walsh, and Michael L Littman. Towards unified theory of state abstraction for mdps. AI&M, 1(2):3, 2006.
Anthony Manchin, Ehsan Abbasnejad, and Anton Van Den Hengel. Reinforcement learning with at- tention that works: A self-supervised approach. In Neural Information Processing: 26th Interna- lional Conference, ICONIP 2019, Sydney, NSW, Australia, December 12-15, 2019, Proceedings, Part V 26, pp. 223-230. Springer, 2019.
Amy McGovern and Andrew G Barto. Automatic discovery of subgoals in reinforcement learning using diverse density. 2001.
Ida Momennejad, Evan M Russek, Jin H Cheong, Matthew M Botvinick, Nathaniel Douglass Daw, and Samuel J Gershman. The successor representation in human reinforcement learning. Nature human behaviour, 19))66806622, 2017.
Alexander Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, and Danilo Jimenez Rezende. Towards interpretable reinforcement learning using attention augmented agents. Advances in neural information processing systems, 32, 2019.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. Advances in neural information processing systems, 31, 2018.
Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. arXiv 11077int:1807.4479999 1807.047777 2018.
Suraj Nair and Chelsea Finn. Hierarchical foresight: Self-supervised learning of long honizoorizon tasks via visual subgoal generation. In International Conference on Learning Representations, 2020.
URL https:///peereeiieeeee net/ htt.://wwwwweeeeeeeeeeoooaaaaaaa
Soroush Nasiriany, Vitchyr Pong, Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. Advances in Neural Information Processing Systems, 32, 2019.
Giuseppe Paolo, Jonas Konzalez-Billandon,, Albert Thomas, and Balázs Kégl. Guided safe shooting: model based reinforcement learning with safety constraints, 2022.
Andrew Patterson, Samuel Neumann, Martha White, and Adam White. Empirical design in rein- forcement learning, 2023.
Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal difference models: Model- free deep for model-based control. arXiv preprint: 1 1802.09081, 2018.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John
Wiley & Sons, 2014.
Reuven Y. Rubinstein. Optimization of computer simulation models with rare events. Euro-
pean Journal of Operational Research, -911:89-112, 1997a. ISSN 0377-2217. doi: https: HlbbuummgJ0..0111000077221100000
URL ht https: www.wciiineeiieeettttt com/
sci astiancee/rtiiiaaaal///777777713
13

---

Published as a conference paper at ICLR 2024
Reuven Y Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1):89-112, 1997b.
Earl D Sacerdoti. Planning in a hierarchy of abstraction spaces. Artificial intelligence, 152):155-135, 1974.
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametric topological memory for navigation. arXiv preprint w:Xiv:/83.000655 2018.
Julian Schrittwieser, 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 leamned model. Nature, 588(7839):604 609, 2020.
Dhruv Shah, Peng Xu, Yao Lu, Ted Xiao, Alexander Toshev, Sergey Levine, and Brian Ichter. Value function spaces: Skill-centric state abstractions for long-horizon reasoning. arXiv preprint 1xXiv:2/11.0// 189, 2021.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics- aware unsupervised discovery of skills. arXiv preprint arXiv: 190700007 2019.
Lucy Xiaoyang Shi, Joseph Lim, and Youngwoon Lee. Skill-based model- reinforcement learning. arXiv preprint arXiv:2207.07550 2022.
David Silver and Kamil Ciosek. Compositional planning using optimal option models. arXiv preprint arXiv: 6473, 2012.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. Nature, 550(7676))354-335 2017.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models.
Cortes, Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015. URL http:::/prooceeeeogs wew.rsseeeurrppp www.eeewwwws....eeeeeeeeeeee cc/paper http://aperrfflsesppeeeeeee///// files/paper/2015/ File/ 2elee/dd5522495555 /aapaa//0000eaaa0000777113333330 Paperr.pppppppppppp.
Richard S. Sutton. Dyna, integrated architecture for learning, planning, and reacting. SIGART
Bulletin, 2(4):160-163, 1991. URL httL:ttttt.//bt.11/////1
Richard S. Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial Intelligence, 11((1)181-211, 1999. k
URL
Richard S Sutton, Marlos C Machado, G Zacharias Holland, David szepesvari Finbarr Timbers, Brian Tanner, and Adam White. Reward-respecting subtasks for model-based reinforcement learning, arXiv preprint:2202.034000 2022.
Tristan Sylvain, Linda Petrini, and Devon Hjelm. Locality and compositionality in zero-shot learn- ing. arXiv preprint arXiv:1912.12173 2019.
Charlie Tang and Russ R Salakhutdinov. Multiple futures prediction. Advances in neural information processing systems, 32, 2019.
Yujin Tang, Duong Nguyen, and David Ha. Neuroevolution of self-interpretaale agents. In Proceed- ings of the 2020 Genetic and Evolutionary Computation Conference, pp. 424. 2020.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in
neural information processing systems, 30, 2017.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-
learning. Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
Kevin Xie, Homanga Bharadhwaj. Danijar Hafner, Animesh Garg, and Florian Shkurti. Latent skill planning for exploration and transfer. arXiv preprint arXiv:2011.1389) 2020.
14

---

Published as a conference paper at ICLR 2024
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations. arXiv preprint arXiv:2011.14881, 2020.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. International Conference on Machine Learning, pp. 1121--11224 PMLR, 2020.
Lunjun Zhang, Ge Yang. and Bradly C Stadie. World model as graph: Learning latent landmarks for planning. International Conference on Machine Learning, pp. 12611 -12620. PMLR, 2021.
Mingde Zhao, Sitao Luan, lan Porada, Xiao- Chang, and Doina Precup. Meta-learning state-based eligibility traces for sample-efficient policy evaluation. arXiv preprint 1rXiv:/904.11433 2019.
Mingde Zhao, Zhen Liu, Sitao Luan, Shuyuan Zhang, Doina Precup, and Yoshua Bengio. A consciousness -insiesed planning agent for model-based reinforcement learning. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.). Advances in Neural Information Processing Systems, volume 34, pp. 1569-1581. Curran Associates, Inc., 2021. URL http:pss//pppoeeeeeee.....eooooo CC/ pappa files/paper/2021/
fi F122:/22251111170000666666666633
15

---

Published as a conference paper at ICLR 2024
A APPENDIX
Please the following to quickly navigate your points of interest.
Weaknesses Limitations (Sec.
Skipper Algorithmic Details (Sec. pseudocodes, k-medoids based pruning, delusion sup- pression
Theoretical Analyses (Sec. D.I]: detailed proofs, discussions
Implementation Details (Sec. E for Skipper, LEAP and Director
More Experiments (Sec. experimental results that cannot be presented in the main paper due page limit
Ablation Tests & Sensitivity Analyses (Sec.
B WEAKNESSES & LIMITATIONS
We would like to expand the discussions on the limitations to the current form Skipper, as well as the design choices that we seek to improve in the future:
We generate future checkpoints at random by sampling the partial description space. Despite the post-processing such as pruning, the generated checkpoints do not prioritize on the predictable, important states that matter the most to form a meaningful long-term plan.
The current implementation intended for pixel input fully-observable tasks with discrete state and action spaces. Such minimalistic form is because wish to isolate the unwanted challenges from other factors that are not closely related to the idea of this work, as well as to make the agent generalist possible. Skipper is naturally compatible with continuous actions spaces and the only thing we will need to do is to replace the haseline agent with a compatible one such as TD3 (Fujimoto al. 2018); the other hand, for continuous state spaces, the identification of the achievement of checkpoint becomes tricky. This is due to the fact that a strict identity between either must adopt distance measure for approximate state equivalence, or rely on the equivalence of the partial descriptions (which is adopted the current implementation). We intentionally designed the partial descrip- tions to be in the form of bundles of binary variables, so that this comparison could be done fast and trivially for any forms of the state space; for partial observability, despite that no recurrent mechanism has been incorporated in the current implementation. the framework is not incompat- ible. To implement that, we will need to augment the state encoder with recurrent or memory mechanisms and need to make the checkpoint generator directly work over the learned state representations. We acknowledge that future work is needed to verify Skipper's performance on the popular partially-observat benchmark suites, which requires the incorporation of compo- nents to handle partial observability as well as scaling up the architectures for more expressive power;
the current state and the target checkpoint may ever established,
We do not know the precise boundaries of the motivating theory on proxy problems, since it only indicates performance guarantees on the condition of estimation accuracy, which turn does not correspond trivially to set of well-defined problems. We are eager to explore, outside the scope sparse-reward navigation, how this approach can used to facilitate better generalization, and at the same time, try find more powerful theories that guide us better;
C Skipper's ALGORITHMIC DETAILS
C.1 OVERALL Skipper FRAMEWORK (PsEUDO-Con)
The pseudocode of Skipper is provided in Alg. together with the hyperparameters used in our
implementation.
16

---

Published as a conference paper at ICLR 2024
Algorithm/ 1: Skipper with Random Checkpoints (implementation choice in purple)
for each episode do
start of the subroutine to construct the proxy problem
generate more than necessary (32) checkpoints by sampling from the partial descriptions given the extracted context from the initial state;
(k: 12) 12)-mddoid pruning upon estimated distances among all checkpoints; // prune vertices use estimators to annotate the edges hetween the nodes (including a terminal state estimator correct the estimates);
prune edges that are too far-fetched according to distance estimations (threshold set to be 8, same replan interval); // prune edges
|| end of the subroutine to construct the proxy problem for each agent-environment interaction step until termination of episode do if decided to explore (DQN-style annealing e-greedy) then
L take random action;
else
abstract problem just constructed or checkpoint / timeout reached (2 8 steps
since last planned] then
[OPTIONAL RE-GENERATION] call the subroutine above for
Skipper-regen;
run value iteration (for 5 iterations) on proxy problem, select the target checkpoint;
follow the action suggested by the checkpoint-achhevinng policy;
time to train (every actions) then
sample hindsight transitions and train checkpoint-achievine policy, estimators (including a teriminal state estimator) and checkpoint generator;
[OPTIONAL DELUSION CONTROL]: train estimators using generated
checkpoints;
save interaction into the trajectory experience replay; convert trajectory into HER samples (relabel random states as additional goals);
17

---

Published as a conference paper at ICLR 2024
C.2 M-MEDOSDS BASED PRUNING
We present the pseudocode of the modified k-medoids algorithm for pruning overcrowded check- points in Alg. 2 Note that the presented pseudocode is optimized for readers' understanding, while the actual implementation is parallelized. The changes upon the original k-medoids algorithm is marked in purple, which implement forced preservation of data points: when k-medoids is called after the unpruned graph is constructed, Sy is set to be the set containing the goal state only. This is intended to span more uniformly in the state space with checkpoints, while preserving the goal.
Let the estimated distance matrix be D, where each element d.j represents the estimated trajectory length takes for и fulfill the transition from checkpoint cannot handle infinite distances (e.g. from terminal state to another state), the distance matrix D is truncated, and then we take the elementwise minimum between the truncated and D to preserve the way distances. The matrix containing the elementwise minimums would be the input of the pruning algorithm.
checkpoint j. Since k-medoids
Algorithm 2: Checkpoint Pruming with k-medoids
Data: X {œi, ( 22,,.... (nn) (state indices), D (estimated distance matrix), S. (states that
must be kept), k (#checkpoints to keep)
Result: So {M1 M2.... ........ Mx (checkpoints kept)
Initialize S = {M1, M2. ..,,... M&} randomly from X
make sure S, C So
repeat
Assign each data point n; to the nearest medoid Mj, forming clusters C1 C2...... Cku foreach medoid M, do
Calculate the cost J of M, as the sum of distances between Mj and the data points in Czi
Find the medoid Mj with the lowest cost Jyi
M, changes then
make sure Sy C So
Replace M, with the data point in Cj that minimizes the total cost;
until Convergence (no cost improvement);
C.3 DELUSION SUPPRESSION
RL agents are prone to blindly optimizing for an intrinsic objective without fully understanding the consequences of its actions. Particularly in model-based RL or in Hierarchical RL (HRL), there is a significant risk posed by the agents trying to achieve delusional future states that do not exist or beyond the safety constraints. With use of learned generative model, in Skipper and other HP frameworks, such risk is almost inevitable, because of uncontrollable generalization effects.
Generalization abilities of the generative models double-edged sword. The agent would take advantage of its potentials to propose novel checkpoints to improve its behavior, but is also at risk of wanting achieve non-existent unknown consequences. In Skipper. checkpoints imagined by the generative model could correspond to non-existent "states" that would lead to delusional edge estimates and therefore confuse planning. For instance, arbitrarily sampling partial descriptions may result in a delusional state where the agent is in cell that can never be reached from the initial states. Since such states do not exist in the experience replay. the estimators will have not learned how to handle them appropriately when encountered in the generated proxy problem during decision time. We present resulting failure mode in Fig. 5
To address such concerns, we propose optional auxiliary training procedure that makes the agent stay further away from delusional checkpoints. Due to the favorable properties the update rules of Dr (in fact, V, as well), all we have to do is to replace the hindsight-sampled target states with generated checkpoints, which contain non non-existent states. Then, the auxiliary rewards will all converge to the minimum in terms of favorability the non-existent states. This is implemented trivially by adding loss to the original training loss for the distance estimator, which we give a 0.25 scaling for stability.
18

---

Published as a conference paper at ICLR 2024
Figure Example of Failure Caused by Delusions: we illustrate instance of chasing delusional checkpoint one of our experimen- tal runs by Skipper. The distance (discount) estimator, probably due to the ill-generalization, estimates that the delusional checkpoint (yellow) very close to every other state. resulting plan was that the agent thought it could reach any far-away checkpoints using the delusional state to form shortcut: the goal was least 17 stops away would be reached 2.2.
D=1.17
delusional checkpoint
dorance ccror (delusion)
50
suppression
su0ces Ddte
0.6
0.6
0.6
0.4
0.4
0.4
o sisppression
10
0.2
0.3
1.0
1.5
0.0
0.5
1.0
155
0.0
0.5
1.0
1.5
0.0
0.5
10
15
training
(4) generalization performance (b) dislance estimation emoor (L1) (c) frequency of selecting delu- sional targets
(d) larget optimulity
training
for delusional edges
Figure 6: Performance of Skipper-once with the proposed Delusion Suppression Technique: each curve and corresponding error bar (95% CT) are processed from 20 independent seed runs. a) the performance across training tasks is shown. A more optimal performance when delusions suppressed; During training interactions, the in estimated (truncated) distance from delusional targets are significantly reduced target is reduced to almost negligible during the whole training process; d) The optimality of target checkpoint during training can be improved by the suppression. Each agent is trained with 50 environments and each curve is processed from 20 independent seed runs.
be achieved with Skipper-once in training tasks,
technique;
frequency of selecting delusional
Algorithm 3: Delusion Suppression
// This whole code block should be injected into the training loop if used generate using the checkpoint generator, from the sampled batch of encoded states, the target states (to overwrite those relabelled in the HER) i.e. replace (st, a, Fr41 S1+1, 5 with (st, Gz, T $141, 2 where oe are generated from the context of st train the distance estimator as if these are sampled from the HER
19

---

Published as a conference paper at ICLR 2024
We provide analytic results and related discussion for Skipper-once agents trained with the pro- posed delusion suppression technique on 50 training tasks in Fig. 6 The delusion suppression technique is not enabled by default because it not introduced in the main manuscript due to the page limits.
The delusion suppression technique can also be used to help understand the failure modes of LEAP Sec. E.2.4
D THEORETICAL ANALYSES
D.1 UPDATE RULES FOR EDGE ESTIMATION
First, we want to show that the update rules proposed in the main paper indeed estimate the desired cumulative discount and reward.
The low-level checkpoint-achiewv policy x is trained with an intrinsic reward to reach target state O The cumulative reward and cumulative discount estimated by applying policy evaluation given on the two sets of auxiliary reward signals, respectively.
For the cumulative discounted reward random variable:
(*(8s; ar () R(st, ae, Se11) + (S Act1ls0)
(5)
R(St: Ar,
(6)
where St+1 ~ p(-|st,aa)) az), A411 T-|1S++1, 20)) and with V= 1411/ 0
S
We overload the notation as follows: V ))) V (s, o) with ~ π. (-|ss, 8))
The cumulative discount random variable denotes the event that the trajectory did not terminate before reaching the target 8
(St, A118 = I (S (44111))144111))) = = })
(7)
(8)
where Tı denotes the timestep when the trajectory terminates, and with ( A411 (
if S.+1 = の and r (S++1, Ar+1 ()) 0 if S.+1 + の is terminal. We overload the notation as follows: I (s.
( A, s 0) with A 7-1SS+1,, so).
Note that, for the sake of simplicity, we take here the view that the terminality of states is determin- istic, but this not reductive as any state with a stochastic terminality can be split into two identical states: one that is deterministically non-terminal and the other that is deterministically terminal. Note also that we could adopt the view that the discount factor the constant probability of the trajectory to not terminate.
D.2 PERFORMANCE BOUND
We are going to denote the expected cumulative discounted reward, a.k.a. the state-action value with 9T - Eg [V], and let år. be our estimate for it. We are also going consider the state value 2x ()) = E. T[ass, の (s, a) and its estimate 0z Similarly, we denote the expected cu- mulative discount with 7 [T' and its estimate with 9..
We are in the presence of hierarchical policy. The high level policy consists in (potentially) stochastically picking a sequence of checkpoints. The low-level policy is implemented by which is assumed to be given and fixed for the moment. The composite policy μ is non-Markovian: it depends both on the current state and the current checkpoint goal. So there no notion of state value, except when arrive at checkpoint, i.e. when high level action (checkpoint selection) needs to be chosen.
Proceeding further, we adopt the view where the discounts way to represent the hazard of the environment: 1 - y is the probability of sudden trajectory termination. In this view, V, denotes the (undiscounted: there is discounting) expected sum of reward before reaching the next
20

---

Published as a conference paper at ICLR 2024
checkpoint, and more interestingly "n denotes the binomial random variable of non-termination during the transition to the selected checkpoint.
Making the following assumption that the trajectory terminates almost surely when reaching the goal, i.c Y (3(, = O, Vs;, the gain V can be written:
Vo ((S + I(S
Σ V(S k=0 Y=0
I r S
(9)
where Sk+1 ~ ((-1S)). where V
<+1) is the gain obtained during the path between Sk and +++) 1) is either 0 or depending whether the trajectory terminated reached where S K++ and [ S
S6+1 If we consider μ deterministic planning routine over the checkpoints, then the action space of μ boils down to a list of checkpoints {sā su. $i .... s =gg) Thanks to the Marko- vian property in checkpoints, have independence between V and r,,. therefore for the expected value o rT, we have:
(pon(Su) = Exoo SV ssll =
Σ (s
Il ==0
(10)
km
Having obtained the ground truth value, in the following, we are going to consider the estimates which may have small error terms:
[s)(s) 0z 0s)] < Umuxa < (1 Umax
and hx((s) fx(5) < < (1-))²
Vs. (11)
We are looking for performance bound, and assume without loss of generality that the reward function is non-negative, the values are guaranteed to be non-negative well. We provide an upper bound:
Uuon (8)
(s (
(12)
k=0
t=
(u
1 ( + C, Umax
II
(Ya (8) (+11) +
(13)
k=0
1=0
R-1
< Unoo (s) +
c Umax
II
(7
0 + + Σ
(vr (x (6 + t,, ke,n o(ee + +e,)
km
1=0
km
(14)
D
2/oo (5) + Umax Σ *k cy Umax
Σ k=O
ky + +o(c
(15)
Unoo (5) + Ev Umux
+ Umux
+ +o(exte)
(16)
(1 7
21

---

Published as a conference paper at ICLR 2024
Similarly, derive lower bound:
Uuom (s)
(+1
(17)
k=0
(Ux(s (sk
1
Ev Umax
II (s (5 (111 EE)
(18)
A-1
Unoo (8)
cp Umax
II
( (s 8 i+1 (4) Σ
(wr (> 1x21))+)))))))))))) - -v ken + olco + ())
h=
1=0
k=
(19)
00
06
Unon (s)
Umax
Σ
nh ey Umnx Σ
ly + o +e)
(20)
U/rr (s)
6 Umax
6y/xmax + +G) 7)-
(21)
We may therefore conclude that öfr equals Upor up to accuracy of SuPan + Eu 1-y (1-7)² + o + c,)-
Note that the requirement for the reward function to be positive is only a cheap technical trick to ensure we bound in the right direction of cy errors in the discounting, but that the theorem would still stand if were not the case.
D.3 No ASSUMPTION OPTIMALITY
If the low-level policy is perfect, then the best high-level policy is to choose directly the goal as target"] Our approach assumes that it would be difficult to learn effectively a æ when the target too far, and that would rather Therefore, we'll never want to make any optimality assumption on otherwise our approach is pointless. These theories we have initiated makes assumption on æ.
a proxy construct path with shorter-distance transitions.
The Theorem provides guarantees on the solution to the overall problem. The quality of the solution depends on both the quality of the estimates (distances/discounts, rewards) and the quality of the policy, as the theorem guarantees accuracy to the solution of the overall problem given current policy, which should evolve towards optimal during training. This means bad policy with good estimation will lead to an accurate yet bad overall solution. No matter the quality of the policy, with a bad estimation, it will result in a poor estimate of solutions. Only a near-optimal policy and good estimation will lead to near-optimal solution.
E IMPLEMENTATION DETAILS FOR EXPERIMENTS
E. Skipper
E.I.1 TRAINING
The agent is based on a distributional prioritized double DQN. All the trainable parameters are optimized with Adam at rate of 2.5 x 10-4 (Kingma & Ba 2014), with gradient clipping by value (maximum absolute value 1.0). The priorities for experience replay sampling are equal to the per-sample training loss.
E.1. FULL STATE ENCODER
The full-state encoder is two layered residual block (with kernel size 3 and doubled intermediate channels) combined with the 16-dimensional bag-of-words embedder of Baby Al (Hui et al.) 2020).
A triangular incquality can be shown that with perfect and a perfect estimate of and Yr. the perfor- mance will always be minimized by selecting sf 8g.
22

---

Published as a conference paper at ICLR 2024
E.1.3 PARTIAL STATE SELECTOR (SPATIAL ABSTRACTION)
The selector a is implemented with one-head (not multiheaded,, therefore the output linear transfor- mation of the default multihead attention implementation in PyTorch is disabled.) top-4 attention, with each local perceptive field of size 8 x 8 cells. Layer normalization before and after the spatial abstraction.
(Ba al. 2016 is used
E.1.4 ESTIMATORS
The estimators, which operate on the partial states, are 3-layered MLPs with 256 hidden units.
An additional estimator for termination is learned, which instead of taking a pair of partial states input, takes only one, and is learned to classify terminal states with cross-entropy loss. The estimated distance from terminal states to other states would be overwritten with oc. The internal y for intrinsic reward of æ is 0.95, while the task y is 0.99
D
3 28 29 30 31
The estimators use C51 distributional TD learning (Dbnney et al. 2018). That is, the estimators output histograms (softmax over vector outputs) instead of scalars. We regress the histogram to- wards the targets, where these targets are skewed histograms of scalar values, towards which KL-divergence is used to train. At the output, there are 16 bins for each histogram estimation (value for policy, reward, distance).
Figure Estimating Distribu- tions of Discount and Distance with the Same Histogram: by transplanting the support with the corresponding discount values, the distribution of the cumulative dis- count can be inferred.
1.5. RECOVERING DISCOUNTS FROM DISTANCES
We recover the distribution of the cumulative discount by replac- ing the support of the discretized truncated distances with the cor- responding discounts, as illustrated in Fig. / This addresses the problem of [][rP] * ,E!P), as the probability of having trajec- tory length of under policy from state sy so the same as a trajectory having discount 7
E. CHECKPOINT GENERATOR
Despite Skipper is designed to have the generator work on state level, that is, it should take learned state representations inputs and have state representations outputs, in experiments, the generator actually operates on observation inputs and outputs. This is because of the preferred compactness of the observations and the equivalence full states under full observability in our experiments.
The context extractor & 32-dimensional BabyAl BOW embedder. It encodes input observa- tion into representation of the episodic context.
The partial description extractor & made of 32-dimensional Baby BOW embedder, followed by 3 aforementioned residual blocks with 3 x 3 convolutions (doubling the feature dimension every time) in between, ended by global maxpool and final linear projection to the latent weights. The partial descriptions are bundles of binary latents, which could represent at most 64 "kinds" of checkpoints. Inspired by VQ-VAE (Van Den Oord et al. 2017). we use the argmax of the latent weights as partial descriptions, instead of sampling according to the softmax-ed weights. This en- ables easy comparison of current state to the checkpoints each state deterministically corresponds to one partial description. We identify reaching target checkpoint the partial description of the current state matches that of the target.
the partial description space, because
The fusing function first projects linearly the partial descriptions to a 128-dimensional space and then uses deconvolution to recover an output which shares the same size as the encoded context. Finally, residual block is used, followed by final 1x1 convolution that downscales the concate- nation of context together with the deconv's 'dd partial description into a 2D weight map. The agent's location is taken be the argmax of this weight map.
23

---

Published as a conference paper at ICLR 2024
The whole checkpoint generator is trained end-to-end with standard VAE loss. That is the sum of a KL-divergence for the agent's location, and the entropy of partial descriptions, weighted by 2.5 x 10-4, suggested in https: / //gttttth///// com/ sample losses in the batches are not weighted for training according to priority from the experience replay.
-VEE Note that the per-
We want to mention that one does not want to generate non-goal terminal states as checkpoints, could also seek to train on reversed (5', St)) pairs. In this case, the checkpoints to reconstruct will never be terminal.
E.1.7 HER
Each experienced transition is further duplicated into 4 hindsight transitions at the end of each episode. Each of these transitions combined with a randomly sampled observation from the same trajectory as the relabelled "goal". The size of the hindsight buffer is extended to 4 times that of the baseline that does not learn from hindsight accordingly, that is, 4 x 10
E. PLANNING
As introduced, we value iteration over options (Sutton al. 1999) plan the proxy problem represented as an SMDP. We use the matrix form C = Rexs + IVV, where R and I' are the estimated edge matrices for cumulative rewards, respectively. Note that this notation is different from the ones we used in the manuscript. The checkpoint value V. initialized as all-zero, is taken on the maximum of Q along the checkpoint target (the actions for μ dimension. When planning is initiated during decision time, the value iteration step is called 5 times. We do not run until convergence since with low-quality estimates during the early stages of the learning, this would be waste of time. The edges from the current state towards other states onc-directional, and the self-loops are also removed. This means the first column as well as the diagonal elements of R and I' are all zeros. Besides pruning edges based on the distance threshold, as introduced in the main paper, the terminal estimator also used to prune the matrices R and T: the rows corresponding to the terminal states are all zeros.
always set to be
The only difference between the two variants, i.e. Skipper-once and Skipper-regen is that the latter variant would discard the previously constructed proxy problem and construct a new one every time the planning is triggered. This introduces more computational effort while lowering the chance that the agent gets "trapped" in a bad proxy problem that cannot form effective plans to achieve the goal. If such a situation occurs with Skipper-regen, as long as the agent does not terminate the episode prematurely, a new proxy problem will be generated to hopefully address the issue. Empirically, as we have demonstrated in the experiments, such variant in the planning behavior results in generally significant improvements in terms of generalization abilities at the cost of extra computation.
E. HYPERPARAMETER TUNING
Some hyperparameters introduced by Skipper can be located in the pseudocode in Alg. !
Timeout and Pruning Threshold Intuitively, we tied the timeout to be equal to the distance pruning threshold. The timeout kicks in when the agent thinks a checkpoint can be achieved within e.g. steps, but already spent 8 steps yet still could not achieve it.
we fully used the advantage of our experiments on DP-solvable tasks: with a snapshot of the agent during its training, we can sample many ( starting state, target state) pairs and calculate the ground truth distance between the pair, as well as the failure rate of reaching from the starting state to the target state given the current policy n, then plot them as the x and y values respectively for visualization. We found such curves to evolve from high failure rate at the beginning, to monotonically increasing curve, where at small true distances, the failure rates are near zero. We picked 8 because the curve starts to grow explosively when the true distances are more than 9.
This leads to how we tuned the pruning (distance) threshold:
for k-medoids We tuned this by running a sensitivity analysis on Skipper agents with different k's, whose results are presented previously in this Appendix.
24

---

Published as a conference paper at ICLR 2024
Additionally, we prune from 32 checkpoints because 32 checkpoints could achieve (visually) good coverage of the state space as well as its friendliness to NVIDIA accelerators.
Size of local Perception Field We used local perception field of size 8 because our bascline model- free agent would be able to solve and generalize well within speaking, our spatial abstraction breaks down the overall tasks into 8 sub-tasks, which the policy could comfortably solve.
tasks, but not larger. Roughly
Model-free Baseline Architecture The baseline architecture (distributional, Double DQN) was heavily influenced by the architecture used in the previous work (Zhao et al. (021). which demon- strated success on similar but smaller-scale experiments (8 x 8). "The difference is that while then we used computationally heavy components such as transformer layers on set-based representa- tion, we replaced them with simpler and effective local perception component. We validated our model-free baseline performance on the tasks proposed in Zhao et al. (202T).
E.2 LEAP
E.2.1 ADAPTATION FOR DISCRETE ACTION SPACES
The LEAP bascline has been implemented from scratch for our experiments, since the original open-sourced implementation" 4 was not compatible with environments with discrete action spaces. LEAP's training involves two pretraining stages, that are, generator pretraining and distance estima- tor pretraining, which were originally named the VAE and RL pretrainings. Despite our best effort, that is to be covered in detail, we found that LEAP was unable to get reasonable performance in its original form after rebasing on a discrete model-free RL baseline.
E.2.2 REPLACING THE MODEL
We tried to identify the reasons why the generalization performance of the adapted LEAP was unsat- isfactory: we found that the original VAE used in LEAP is not capable to handle even few training tasks, let alone generalize well to the evaluation tasks. Even by combining the idea of the context / partial description split (still with continuous latents), during decision time, the planning results given by the evolutionary algorithm (Cross Entropy Method, CEM, Rubinstein (1997a)) almost al- ways produce delusional plans that are catastrophic in terms of performance. This was why switched into LEAP the same conditional generator we proposed in the paper, and adapted CEM accordingly, due to the change from continuous latents to discrete.
We also did not find that using the pretrained VAE representation as the state representation during the second stage helped the agent's performance, the paper claimed. In fact, the adapted LEAP variant could only achieve decent performance after learning a state representation from scratch in the RL pretraining phase. Adopting Skipper's splitting generator also disables such choice.
E.2.3 REPLACING TDM
The original distance estimator based on Temporal Difference Models (TDM) also does not show capable performance in estimating the length of trajectories, even with the help of ground truth distance function (calculated with DP). Therefore, we switched to learning the distance estimates with our proposed method. Our distance estimator is not sensitive to the sub-goal time budget as TDM and is hence more versatile in environments like that was used in the main paper, where the trajectory length of each checkpoint transition could highly vary. Like for Skipper, an additional terminal estimator has been learned to make LEAP plamning compatible with the terminal lava states. Note that this LEAP variant was trained on the same sampling scheme with HER as in Skipper.
The introduced distance estimator, as well as the accompanying full-state encoder, are of the same architecture, hyperparameters, and training method as those used in Skipper. The number of in- termediate subgoals for LEAP planning is tuned to be 3, which close to how many intermediate checkpoints Skipper typically needs to reach before finishing the tasks. The CEM called with iterations for each plan construction, with a population size of 128 and an elite population of size 16. We found no significant improvement in enlarging the search budget other than additional wall
httpp:://github.ccom
25

---

Published as a conference paper at ICLR 2024
time. The new initialization of the population is by sampling c-mean of the clite population (the binary partial descriptions), where 0.01 to prevent the loss of diversity. Because of the very expensive cost of using CEM decision time and its low return of investment in terms of gen- eralization performance, during the RI. pretraining phase, the agent performs random walks over uniformly random initial states collect experience.
E.2.4 FAILURE MODE: DELUSIONAL. PLANS
Interestingly, we find that major reason why LEAP does not generalize well is that it often gen- erates delusional plans that lead to catastrophic subgoal transitions. This is likely because of its blind optimization in the latent space towards shorter path plans: any paths with delusional shorter distances would be prefered.. We present the results with LEAP combined with proposed delu- sion suppression technique in Fig. g We find that the adapted LEAP agent, with our generator, our distance estimator, and the delusion suppression technique, is actually able achieve significantly better generalization performance.
1.0
6.6
64
LEAP
LEAP LEAP-suppress
01
0.5
2.0
0.5
1.0
tmiining steps
taaiiing steps OOD,
tminiing steps
baining stegs
toaling
(a) training,
(c) OOD, 0.35
(d) OOD, 0.45
(e)
0.55
Figure 8: Comparative Results of LEEP with and without the delusion suppression technique: the results are obtained with 50 training tasks. The results are obtained from 20 independent seed runs.
E.3 DIRECTOR
E.3.1 ADAPTATION
We based our experiments of Director (Hafner et al.] 2022) on the publicly available code chttps://githhn
com dan! /ar/ director) released by the authors. Except for a few changes in the parameters, which are depicted in Tab. [ we have used the default con- figuration provided for Atari environments. Note that as the Director version in which the worker receives no task rewards performed worse in our tasks, we have used the version in which the worker receives scaled task rewards (referred as "Director (worker task reward)" in Hafner et al. (2022)). This agent has also been shown to perform better across various domains in Hafner al. (2022).
Encoder. Unlike Skipper and LEAP agents, the Director agent receives as input a simplified RGB image of the current state of the environment (see Fig. 99. This is because we found that Director performed better with its original architecture, which was designed for image-based ob- servations. We removed all textures to simplify the RGB observations.
Figure 9: An exam-
ple for simplified observa- tions for Director.
E.3.2 FAILURE MODES: BAD GENERALIZATION SENSITIVE TO SHORT TRAJECTORIES
We investigated why Director is unable to achieve good training perfor- mance(Fig. 3) As Director was designed to be trained solely on environments where is able to collect long trajectories to train good enough recurrent world model (Hafner al.) 2022), hypothesized that Director may perform better in domains where it is able to interact with the envi- ronment through longer trajectories by having better recurrent world models (i.e., the agent does not
Training Performance.
26

---

Published as a conference paper at ICLR 2024
Table 1: The changed parameters and their values in the config file of the Director agent.
Parameter
Value
replay_Size Liize
2M
replay_chunk
12
imag.horizon horizon
w.V_skill_duratie
4
train skill duration
worker
{extr: 0.5, expl: 0.0, goal: 1.0)
sticky gray
False False
immediately die as result interacting with specific objects the environment). To test this, we experimented with variants of the used tasks, where the lava cells replaced with wall cells, the agent does not die upon trying to move towards them (we refer to this environment as the "walled" environment). The corresponding results 50 training tasks are depicted in Fig. TO As can be seen, the Director agent indeed performs better within the training tasks than in the environments with lava.
Generalization Performance. We also investigated why Director is unable to achieve good gener- alization (Fig. B). As Director trains its policies solely from the imagined trajectories predicted by its learned world model, we believe that the low generalization performance is due to Director being unable to learn a good enough world model that generalizes to the evaluation tasks. The generaliza- tion performances in both the "walled"" and regular environments, depicted in Fig. [OO indeed support this argument. Similar to what we did in the main paper, we also present experimental results for how the generalization performance changes with the number of training environments. Results in Fig. II show that the number of training environments has little effect on its poor generalization performance.
140
Director-lava
Director-wall
0.8
042000
04
0.2
0.0
0.5
3.0
1.5
0.
0.%
1.0
1.5
80
0.5
3.0
3.5
0n
0.
1.5
80
115
1.0
Inaitinn stopo
traising Aaga
Inaining stepe
training Mtgs
Inaining stogo
training, dif' 0.4
OOD, diff 0.25
OOD, dil 0.35
(d) OOD, dilf 0.45
(e) OOD, dift 0.55
Figure 10: Comparative Results Director on Environments with Lavas and on those Walls: results are obtained with 50 training tasks. The results for Director-lava (same as in the main paper) are obtained from 20 independent seed runs.
3.0
Dull
66
6.4
Director-lacw
60
50
(a) training. diff' 0.4
(b) OOD, diff 0.25
(c) OOD, diff 0.35
(d) OOD, dift 0.45
(e) OOD, 0.55
Figure 11: Generalization Performance of Agents on Different Numbers of Training Tasks (while Direc- tor the walled environments): besides Director, each data point and corresponding bar (95% confidence interval) are processed from the final performance from 20 independent seed Director-wall results are obtained from runs.
27

---

Published as a conference paper at ICLR 2024
F EXPERIMENTAL RESULTS (CONT.)
We present the experimental results that the main paper could not hold due to the page limit.
F.1
SKipper-ONCE SCALABILITY
We present the performance of Skipper-once different numbers of training tasks in Fig. 12
7.2
66
644
ent
envs
enos inf eves
5 entos
50 enos
0.
1/0
1.5
traiting
training
(a) training,
OOD,
(c) 0.35
(d) OOD,
(e) OOD, 0.55
Figure 12: Generalization Performance of Skipper-once on different numbers of training tasks: each enror har (95% confidence interval) is obtained from 20 independent seed runs.
F.2 Skipper-REGEN SCALABILITY
We present the performance of Skipper-regen on different numbers of training tasks in Fig.
1.0
940000 224
06
644
eno
enar
enes inf erns
5
50 enoe
0.5
training
truining
(a) training.
OOD,
(c) OOD, 0.35
(6) OOD,
(e) OOD, 0.55
Figure 13: Performance of Skipper-regen on different numbers of training tasks: each error har (95% confidence interval) is obtained from 20 independent seed runs.
F.3 MODELFREE BASELINE SCALABILITY
We present the performance of the modelfree baseline different numbers of training tasks in Fig.
14
3.0
25 enzus
100 enos
0.8
50 entus
inf envs
040000 30le
0.2
5 entos
0.0
0.5
3.0
1.5
0.0
0.5
1.0
1.5
08
0.5
3,0
3.5 0.
05
10
05
18
traming steps
traming degs
tramiin steps
tnining dags
tramiin deps
training. diff 0.4
OOD, diff 0.25
OOD, diff 0.35
(d) OOD, diff 0.45
(e) OOD, 0.53
Figure 14: Generalization Performance the modelfree baseline on different numbers of training tasks: each error bar (95% confidence interval) is obtained from 20 independent seed runs.
28

---

Published as a conference paper at ICLR 2024
F.4 LEAP SCALABILITY
We present the performance of the adapted LEAP baseline on different numbers of training tasks in Fig. 15
10
100000 280
046
eNP
ends
exes
50
inf
00
1.5
(a) training. 0.4
OOD, (c) OOD, 0.33
(d) OOD, 0.45
(e) OOD, 0.55
Generalization Performance of the LEAP baseline on different numbers of training tasks: each error bar (95% confidence interval) is obtained from 20 independent seed runs.
Figure 15:
F,5 DIRECTOR SCALABILITY
present the performance of the adapted Director baseline on different numbers of training tasks
in Fig. 16
emt
25 enos
100 enos
68
50
infenss
62
0.0
0.0
0.5
3.0
1.5
0.0
0.5
1.0
1.5
0.
0.5
3,0
tramiing strps
taaning deps
trumiin steps
taiting sheps
teattin steps
(a) training, diff 0.4
(b) OOD, diff 0.25
(6) OOD. diff 0.35
(d) OOD, diff 0.45
(e) O0D, diff 0.55
Figure 16: Generalization Performance of the Director baseline on different numbers of training tasks: each error bar (95% confidence interval) is obtained from 20 independent seed runs.
F.6 GENERALIZZTTON PERFORMANCE ON DIFFERENT NUMBERS OF TRAINING TASKS
The performance of all agents on all training configurations i.e. different numbers of training tasks, are presented Fig. Fig.
Fig.
20 Fig. 21 and Fig. 22
3.2
Skipper-onee
Skipper-regen
LEAP
Director
modeiffre
066
0000000
0.
1/0
1.5
traiting
tnainin
(a) training.
OOD, (e) 0.35
(8) OOD, 0.45
(e) 0.55
Figure 17: Generalization Performance of the Agents when trained with training task: each error bar (95% confidence interval) is obtained from 20 independent seed runs.
29

---

1.0 Published as a conference paper at ICLR 2024
1.0
Skipper-once
Seepeer-regen
LEAP
Director
moodelfree
941200 2315
044
01
0.5
2.0
2.5
0.
0.5
1.0
1.5
1.0
tminning steps
taikiig strgs
tminting strpe
taining stege
teatting steps
training.
OOD, diff
(d)
(e) 0.33
Figure 18: Generalization Performance of the Agents when trained with 5 training tasks: each error bar (95% confidence interval) is obtained from 20 independent seed runs.
1.0
Skipper-oner
Skipper-regen
LEAP
Director
modelffee
5400000 2800
156
104
0.2
010
0.5
7.0
0.0
0.5
1.0
tmiining
taaiting steps
traininng
Craining stops
Hoalinng
(a) training
OOD,
(d)
(e)
Figure 19: Generalization Performance of the Agents when trained with 25 training tasks: each error bar (95% confidence interval) is obtained from 20 independent seed runs.
1.0
Skipper-onne
Skipper-rager
LEAP
Director
madeiffre
0.8
943000 28t
04
0.0
0.5
3.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
3.
3.5
0.0
0.5
0.5 textning steps (e) OOD, 0.35
traning rteps
traming deps
teiming steps
traiiing dhge
training, diff 0.4
OOD, diff 0.
OOD, 0.35
(d) OOD, 0.45
Figure 20: Generalization Performance the Agents when trained with 50 training tasks (same as in the main paper): each error bar (95% confidence interval) is obtained from 20 independent seed runs.
1.0
10200 288
0/6
0.2
Skipper-onoe
Skipper-regen
LEAP
Director
modelfree
0.
0.5
2.0
0.
0.5
1.0
1.5
06
0.5
1,3
tmining rtpp
taining deps
tmining steps
taining stops
teitning steps
training, dift 0.4
OOD, diff
OOD,
OOD.
(e) 0.55
Figure 21: Generalization Performance of the Agents when trained with 100 training tasks: each emor har (95% confidence interval) is obtained from 20 independent seed runs.
30

---

Published as a conference paper at ICLR 2024
1.0
500 2886
556
64
Skipper-onoe
Skipper-rggen
LEAP
Director
modelfree
60
steps
(a) training. 0.4
OOD, diff 0.25
(c) OOD, 0.35
(d) OOD, 0.45
(e) OOD, 0.53
Figure 22: Generalization Performance of the Agents when trained with co training tasks (a new task
each training episode): each error bar (95% confidence interval) is obtained from 20 independent seed runs.
F.6.1 STATISTICAL SIGNIFICANCE POWER ANALYSES
Besides visually observing generally non-overlapping confidence intervals, present the pairwise
t-test results of Skipper-once and Skipper-regen against the compared methods. In addition, the advantage is significant, perform power analyses to determine if the number of seed (20) was enough to make the significance claim. These results are shown in Tab. 2 and Tab. respectively.
Table 2: Skipper-once others: significance & power
method ltask difficulty
0.25
0.35
0.45
0.55
leap
22
NO
NO
NO
train envs
director
15
11
22
11
baseline
NO
38
36
NO
leap
28
NO
NO
NO
train envs
director
NO
NO
NO
22
baseline
11
8
10
12
leap
15
13
11
7
25 train envs
director
2
2
2
baseline
2
2
leap director
17
11
train envs
2
2
2
bascline
2
2
2
2
leap director
15
10
7
9
100 train envs
2
2
2
2
basclinc
2
2
2
2
leap director
32
5
7
3
inf train
2
2
2
2
baseline
NO
NO
NO
NO
threshold: 0.05.
Effect size set be the difference of the means of the compared pairs Colas al. 2018
Colls hold if results NOT significant insufficient seeds for statistical power. significant cases, minimum number seeds statistical
provided.
As we can observe from the tables, generally there is significant evidence of generalization advan-
tage in Skipper variants compared to the other methods, especially when the number of training environments are between 25 to 100. Additionally, as expected, Skipper-regen displays more dom- inating performance compared to that of Snipper-once.
31

---

Published as a conference paper at ICLR 2024
Table Skipper-regen v.s, others: significance power
method (task difficulty
0.25
0.35
0.45
0.55
leap
32
NO
NO
NO
train envs
director
16
13
23
10
bascline
NO
NO
NO
NO
leap
NO
NO
NO
NO
train
director
33
NO
NO
NO
baseline
6
8
4
leap
10
7
5
4
train envs
director
2
baseline
2
2
2
leap
6
50 train envs
director
2
2
2
2
baseline
2
2
2
2
leap
7
3
3
2
100 train envs
director
2
2
2
2
baseline
2
2
2
2
leap
15
3
2
2
inf train
director
2
2
2
2
baseline
NO
NO
35
threshold: 0.05.
Efflect size set to be the difference of the means of the compured pairs. Coolas al. 2018).
Colls bold if results NOT significant insufficient seeds for statistical power. significant cases, minimum number soods statistical power provided.
G ABLATION & SENSITIVITY
G.1 VALIDATION EFFECTIVENESS STOCHASTIC ENVIRONMENTS
We present the performance of the agents in stochastic variants of the used environment. Specifi- cally, with probability 0.1, each action is changed into a random action. We present the 50-training tasks performance evolution in Fig. 23 The results validate the compatibility of our agents with stochasticity in environmental dynamics. Notably, the performance the baseline deteriorated to worse than even Director with the injected stochasticity. The compatibility of Hierarchical RL frameworks to stochasticity has been investigated in Hogg et al. (2099).
Skipper-onae
Skipper-regen
LEAP
Director
moddelfree
0.8
62
69
tmiining steps
strss
teaining steps
tsaining steps
textning dieps
(4) training, dift' 0.4
(b) OOD. diff 0.25
(c) OOD, diff 0.35
(d) OOD, difl 0.45
(e) OOD, dift 0.55
Figure 23: Generalization Performance of agents in stochastic environments: e-greddy style randomness added to each primitive action with O.L. Each agent trained environments and each curve processed from 20 independent seed runs.
G.2 A ABLATION FOR SPATIAL ABSTRACTION
We present in Fig. 24 the ablation results the spatial abstraction component with Skipper-once agent, trained with 50 tasks. The alternative component of the attention- bottleneck, which is without the spatial abstraction, is MI.P flattened full state. The results confirm significant
32

---

Published as a conference paper at ICLR 2024
advantage in terms of generalization performance well as sample efficiency in training, introduced by spatial abstraction.
1.0
se spat. abst.
280
1400000
w/ sput, abst.
00
0.5
1.5
(a) training.
OOD, diff (c) OOD, 0.33 (d) OOD, 0.45
(e) OOD, 0.53
Figure 24: Ablation for Spatial Abstraction on Skipper-once agent: cach agent is trained with 50 cnviron- ments and each curve is processed from 20 independent seed runs.
G.3 ACCURACY OF PROXY PROBLEMS CHECKPOINT POLICIES
We present in Fig. 25 the ablation results the accuracy of proxy problems well as the check- point policies the Skipper-once agents, trained with 50 tasks. The ground truths are computed via DP the optimal policies, which also suggested by DP. Concurring with theoretical analyses, the results indicate that the performance of Skipper is determined (bottlenecked) by the accuracy of the proxy problem estimation
the high-level and the optimality of the checkpoint policy on the lower level. Specifically, the curves for the generalization performance across training tasks, in (a) of 25] indicate that the lower than expected performance is composite outcome of eтrоrs in the two levels. In the next part, we address major misbehavior of inaccurate proxy problem estimation chasing delusional targets.
1.0
6.8
066
10200
64
Skipper-once
optimal plan
optimal policy
000
0.5
3.0
15
on
0%
10
1.5
nC
1.0
15
on
0.
1n
1.5
80
16
1.5
Inaininn slepo
Leainiing daga
Inaining stepe
trairing
Maps
Inaining stgg
training. 0.4
OOD, diff 0.
OOD, 0.35
(d) OOD, 0.45
(e) OOD, 0.35
Figure 25: Skipper-once Empirical Performance v.s. ground truths: both the optimal policy and optimal plan variants are assisted by DP. The default deterministic setting induces the fact that combining optimal policy and optimal plan results in 1.0 success rate. The figures suggest that the learned agent limited by errors both the proxy problem estimation and the checkpoint policy Each agent trained with 50 environments and each curve's processed from 20 independent seed runs.
G.4 TRAINING INITIALIZATION: UNIFORM V.S. SAME AS EVALUATION
We compare the agents' performance with and without uniform initial state distribution. The non- uniform starting state distributions introduce additional difficulties in terms of exploration. In Pre- sented in Fig. 26 these results are obtained from training on 50 tasks.
conclude that given similar computational budget, using non-uniform initialization only slows down the learning curves without introducing significant changes to our conclusions, and thus we use the ones with uniform initialization for presentation in the main paper.
G.5 ABLATION: VERTEX PRUNING
As mentioned previously, each proxy problem in the experiments are reduced from 32 vertices to 12 with such techniques. We compare the performance curves of the used configuration against
33

---

Published as a conference paper at ICLR 2024
1.0
Shipper-mgen
LEAP
Director
modelffre
500000 2886
566
64
Skipper-once
60
30
(a) training. 0.4
(b) OOD, diff 0.25
(c) OOD, 0.35
(d) OOD, 0.45
(e) OOD, 0.53
Figure 26:
Comparative Results on 50 training tasks without uniform initial state distribution: each curve processed from 20 independent seed runs.
baseline that generates 12-vertex proxy problems without pruning. We present in Fig. 27 these ablation results on the component of k-medoids checkpoint pruning. We observe that the pruning only increases the generalization but also the stability of performance.
1.0
se prume
166
64
w//
0.
000
0.5
2.0
0.
tmiining
taaiting steps
training
Uniiiing stgg
molning
(a) training.
OOD,
(c) OOD, 0.35
(d) OOD, 0.45
(e) OOD, 0.55
Figure 27: Ablation Results on 50 training tasks for k-medoids pruning: cach curve is processed from 20 independent seed runs.
G.6 SENSITIVITY: NUMBER OF VERTICES
We provide a sensitivity analysis the number of checkpoints (number of vertices) in each proxy problem. We present the results of Skipper-once on 50 training tasks with different numbers of post-pruning checkpoints (all reduced from 32 by pruning). in Fig. 28 From the results, we can see that long as the number of checkpoints is above 6, Skipper exhibits good performance. We therefore chose 12, the one with a rather small computation cost, as the default hyperparameter.
ckpts
24 ckpls
32 ckpts
0.8
066
02
dkpts
0.
stope
daga
Inaining stego
Laaining stgg
Inaining steps
(a) training diff 0.4
(b) OOD, diff 0.25
(c) OOD, diff 0.35
(4) OOD, diff 0.45
(e) OOD, 0.55
Figure 28: Sensitivity of Skipper-once the number of checkpoints in each proxy problem: each agent trained with 50 environments. All curves processed from 20 independent seed
G.7 ABLATION: PLANNING OVER PROXY PROBLEMS
We provide additional results for the readers to intuitively understand the effectiveness of plan- ning over proxy problems. This is done by comparing the results of Skipper-once with baseline
34

---

Published as a conference paper at ICLR 2024
Skipper-ger goal that blindly selects the task goal as its target all the time. We present the results based Concurring with our vision on temporal abstraction, we can see
on 50 training tasks in Fig. 29 that solving more manageable sub- aub-problems leads to faster convergence. The Skipper-goal variant catches up later when the policy slowly improves to be capable of solving longer distance navigation.
2.0
Skipper-goel
64
140000 eat
6.6
62
Skipper-oncce
00
0.3
1.0
Shaining stepo
traiing dags
Iuiining stepo
tratning Mags
Inaining stepe
(a) training, dift 0.4
(b) OOD, diff 0.25
(c) OOD, diff 0.35
(d) OOD, dil 0.45
(e) O0D, dift 0.55
Figure 29: Effectiveness of Proxy Problem based Planning: each agent is trained with 50 environments and each curve processed from 20 independent seed runs.
35


Reviews:
[{"review_id": "GJN5eHsqGe", "reviewer": "Reviewer_muKb", "summary": "The paper proposes a new model-based reinforcement learning agent, Skipper, with the explicit aim to generalize its learned \"skills\" to novel situations. Skipper creates a \"proxy problem\" for a given task. Proxy problems are finite graphs created at decision time, whose vertices (\"checkpoints\") are states and whose directed edges are temporally abstract transitions between these states. The vertices are proposed by a generative model (which splits the state into two parts: a context (which is a general description of the environment), and partial description (which is the agent's current state) , while the estimates of the edges (holding time and expected reward) are learned while interacting with the environment. The training is the sum of the losses of the vertex generation and edge estimation.\n\nThe algorithm is evaluated on the MiniGrid-BabyAI framework against the modelfree RL agent that is the basis of the architecture and two relevant Hierarchical Planning methods (LEAP and Director). The evaluation metric is generalization across different instantiations of the environment: subsequent difficulties mean more lava states in the grid. It is demonstrated that Skipper makes use of the context latent representation, as well as the estimated edges to transfer learned policies faster than the baselines.", "questions": "Could you please define spatial abstraction and highlight how it is different from state abstraction? (ideally with a citation) I had a look at all cited papers in the relevant subsection in Section 5, but non of them contained the term \"spatial abstraction\". Machin et al. 2019 used the term \"spatial (and temporal) information\", but it was not defined.\n\n“Existing RL agents either operate solely based on intuition (model-free methods) or are limited to reasoning over mostly shortsighted plans (model-based methods)” could you please back this rather strong statement with a citation (especially the first part).\n\n\"Model-based HRL agents can be prone to blindly optimizing for objectives without understanding the consequences.\" - again, at least a citation would be in order for such a claim.\n\n“For any set of options defined on an MDP, the decision process that selects only among those options, executing each to termination, is a Semi-MDP (SMDP, Puterman 2014)” Could you please point to the part in the book that defines SMDPs as such? \n\n“Taking inspiration from conscious information processing in brains [...]” - could you please cite the works whose findings inspired the introduced local perceptive field selector?\n\n3.3.2, the update rule for the Cumulative Reward: could you please explains where is the KL-divergence in that update rule? To me it just looks like each $v(s,a)$ is updated with the TD-target $(r(s,a,s') + \\gamma v(s,a))$?\n\n\"The following estimates are learned with using distributional RL, where the output of each estimator takes the form of a histogram over scalar support (Dabney et al., 2018)\"  - I don't understand what this sentence is trying to say, and the linked paper did not help. Could you please elaborate?\n\nHow were hyperparameters tuned for the algorithm (the algorithmic specific ones, such as timeout, pruning threshold, k for k-medoid, etc. as well as the ones basis distributed DDQN) and the various baselines for this task?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Originality: splitting the state to context and partial information the way it is proposed and using it for transfer across slightly altered tasks seems intuitive yet to the best of my knowledge original.\n\nQuality: experiments ran over 20 seeds, which is substantially more than the usual 5 and is generally recommended the absolute minimum for a pilot study. The plots are well annotated, captioned with error bars included. \n\nClarity: the paper is well-written and easily understood in most parts.\n\nSignificance: the results as presented appear to be significantly better than the baselines. The proposed method coupling state (or \"spatial\") and temporal abstraction and using it for transfer across tasks is an interesting avenue for further research.", "weaknesses": "As mentioned above, the paper has many strengths, however, it is unfortunately not devoid of weaknesses. There are certain parts of the paper that I did not understand, hence I asked several questions in the next section. Furthermore, there is very little information about the way experiments were carried out (with regards to hyperparameter tuning for the proposed algorithm and the baselines on the various environments). These details are required to assess the validity of the claims and to assure reproducibility.\n\nOnce my doubts are cleared, I will consider raising my score accordingly.\n\nI listed the experiments ran with 20 seeds as strength and it is. However, in order to claim significance, it is suggested that after the pilot study with 20 seeds, statistical power analysis should be performed to determine the necessary sample size [1][2]. \n\n[1] How Many Random Seeds? Statistical Power Analysis in Deep Reinforcement Learning Experiments, Cedric Colas, Olivier Sigaud, and Pierre-Yves Oudeyer, 2018\n[2] Empirical Design in Reinforcement Learning, Andrew Pattersion, Samuel Neumann, Martha White, Adam White, 2023\n\nI recommend either adding a statistical power analysis to the paper or softening the claims a little bit.\n\nThe following typos/style issues did not affect the score:\n\nThe abbreviation RL should be defined as reinforcement learning (RL) the first time it is used., even though it is clear for the audience what it stands for, it is good practice.\n\ns⊙  is defined much later than when first introduced (in Theorem 1) which is confusing.\n\nSection 2, 2nd paragraph, also Section 4 1s paragraph: no need to capitalize the word following a semicolon (;) or end the previous subsentence with a full stop and the capitalize the start of the sentence.\n\nSection 2\n\n1st paragraph: Q: |\\mathcal{S}| \\times |\\mathcal{A}|  (| | missing for S)\n\n2nd paragraph: \"evaluated _on_ unseen tasks, _in which_ there are environment variations, but the core strategies need to finish the task remain [...]\" (not remains)\n\n3rd paragraph: \"Recent work suggests that this approach\" -> the latter approach\n\nFigure 2 should be moved at least one page down, there are concepts in it that are explained much later.\n\n\"The following estimates are learned with using distributional RL\" - with or using\n\n3.3\n\nthe expression \"on top of\" occurs twice and I find it rather ambiguous. Could you use some other expression, e.g. with or after, or anything else that conveys what you are trying to say.\n\n3.3.2\n\n3rd paragraph: \"Similarly to $V_{\\pi}$, we would want to know the cumulative discount leading to the target s⊙ under $\\pi$.\" - There is no need for this sentence. Just start the next sentence as \"The cumulative discount leading to the target s⊙ under $\\pi$ is difficult to learn...\"\n\n3.4\n\nsecond paragraph: \"tall order\" - this is slang, consider changing it to something like \"difficult task\"\n\n5\n\n\"Prior to LEAP (Nasiriany et al., 2019), path planning with evolutionary algorithms was investigated by Nair & Finn (2019); Hafner et al. (2022); Mendonca et al. (2021) propose world models to assist temporally abstract background planning.\"\n\nNot clear where the first sentence ends and the second begins. Use a full stop and an 'and' where applicable.", "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": "sTiATrkaOG", "reviewer": "Reviewer_iydj", "summary": "The paper proposes a method for hierarchical planning by learning spatio-temporal abstraction. The approach is tested for zero-shot generalization.", "questions": "Included with the limitations.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper supports the importance of spatial and temporal abstractions together.\n\n- The approach to minimize plans with non-existent checkpoints is interesting.", "weaknesses": "- In my opinion, the paper lacks an overall description of the algorithm which largely limits the understanding of the paper. \n\n- A proxy problem is estimated using the models by a conditional generative model that randomly samples checkpoints. These checkpoints are only pruned, not improved and thus the approach seems to be heavily dependent on the plan found earlier even if it can not be achieved at the low level. The delusion suppression approach only minimizes the non-existent checkpoints. Are there any guarantees on whether a checkpoint is achievable? \n\n- There is much research done on learning temporal and/or state abstractions automatically for planning and learning [1,2,3,4]. It is not clear how the proposed work is different or novel compared to these different lines of research. The related work needs a thorough justification of differences with different categories of work.\n\n- The proposed approach seems applicable to grid-like environments only as also included in the experiments.\n\nReferences:\n\n[1] Hogg, C., Kuter, U. and Munoz-Avila, H., 2009, July. Learning Hierarchical Task Networks for Nondeterministic Planning Domains. In IJCAI (pp. 1708-1714).\n\n[2] Dadvar, M., Nayyar, R.K. and Srivastava, S., 2023, July. Conditional abstraction trees for sample-efficient reinforcement learning. In Uncertainty in Artificial Intelligence (pp. 485-495). PMLR.\n\n[3] Ravindran, B., AC, I. and IIL, R., Hierarchical Reinforcement Learning using Spatio-Temporal Abstractions and Deep Neural Networks.\n\n[4] McGovern, A. and Barto, A.G., 2001. Automatic discovery of subgoals in reinforcement learning using diverse density.", "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": "OyQe1nJiRd", "reviewer": "Reviewer_3j7M", "summary": "This paper presents Skipper, a model-based reinforcement learning approach that leverages both spatial and temporal abstraction to perform well on problems with sparse rewards. The abstraction takes the form of what the authors refer to as a \"proxy problem\": a directed graph in which vertices and edges correspond to states and the ability to reliability traverse between them (respectively). These proxy problems are estimated at decision time from learned models, trained end-to-end via hindsight. The authors show the ability of Skipper to perform well on a navigation-oriented gridworld problem [via the MiniGrid-BabyAI framework] and in particular show how Skipper generalizes quickly given only a handful of training tasks, outperforming both model-based and model-free baselines in its ability to generalize.", "questions": "I would like the authors to address what I identify as the two core weaknesses in my comments above. In summary:\n\n1. Could the authors comment on the (implicit) assumptions Skipper makes about the task/environment? What decisions about Skipper's design (e.g., using an attention mechanism over the observations to propose vertices for proxy problem estimation) are core to Skipper's functionality or can be changed without changing the overall approach? On what types of problems is Skipper useful (or not)?\n2. Relatedly, a small summary of Limitations or Weaknesses of the proposed approach would help clarify when Skipper should be used, and thus strengthen the paper. If the answers to my question 1 were in the form of such a section, I think it could be effective and helpful for improving clarity.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "This is a well written paper. It is easy to follow despite the challenging subject material that presents a well-scoped approach to tackling model-based RL for gridworld-like problems. The paper is generally good and, between the main body of the paper and the expansive additional experiments and details in the Appendix, the work shows the effectiveness of Skipper. Using hindsight to train Skipper is well-motivated. The baseline algorithms against which Skipper is compared are appropriate, and it is clear that Skipper outperforms them under multiple different circumstances. In particular, the generalization-focused experiments are interesting and convincing, and show the benefits of the decision-time planning approach that Skipper takes to planning.", "weaknesses": "There are two (related) core weakness of the paper: (1) it seems that Skipper is designed primarily with gridworld environments in mind and seems to be relatively limited in what problems it is suited to solve and (2) the paper is missing a clear \"Weaknesses\" or \"Limitations\" section, which would help the reader understand when Skipper will be appropriate.\n\nThe experiments in the paper focus somewhat narrowly on a single family of gridworld environments. While this is not on its own a problem, there are questions regarding how applicable Skipper is to tasks beyond those shown in the paper. Specifically, Skipper's vertex generation module consists of an attention mechanism over the observation. For the experiments presented in the manuscript, the observations all take the form of overhead views of a map of the environment and so directly correspond to the state space. This means that Skipper is particularly well-suited to work on environments such as these, since vertex proposal works directly on the observation, and (unless I am mistaken in my understanding) arguably Skipper is over-engineered to work on environments such as these.\n\nThe appeal of the other model-based reinforcement learning approaches mentioned in the paper, including MuZero and Dreamer (against which Skipper is compared), is that they work in fairly general observation spaces, both capable of taking as input RGB images and learning a model from those. While the paper on the whole generally does an effective job of communicating what assumptions the system makes and how it works, the construction of Skipper seems to implicitly assume a certain type of state or observation representation that limits the applicability of the work in general. If the authors could clarify this point and make explicit these so-far implicit conditions, it would greatly strengthen the paper.\n\nRelatedly, while the paper spends considerable time discussing the ways in which the proposed Skipper approach is effective, it lacks a clear Weaknesses or Limitations section. Such a section, and comments regarding which problems or domains Skipper is best suited, would greatly aid in understanding when Skipper is appropriate. If possible, I would appreciate hearing the authors' thoughts on the size of map or problem that Skipper would be expected to work; I do not think additional experiments are necessary or appropriate, though I am interested to understand if Skipper would succeed on larger occupancy grids and navigation tasks.", "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": "NM0XDYPcPu", "reviewer": "Reviewer_qHoq", "summary": "The paper introduces a planning agent called Skipper that overcomes the limitations of existing reinforcement learning (RL) agents. \n\nRL agents can either rely on intuition or only plan for short-term tasks. Skipper, on the other hand, uses human-like thinking to break down complex tasks into smaller parts called \"proxy\" problems. \n\nThese problems are represented as graphs, with decision points as dots and connections between them as lines. Skipper solves each problem and then focuses on the connections to solve the next part. \n\nThis approach helps Skipper handle new situations better and create useful shortcuts. The paper also shows that Skipper performs better than other methods in terms of adapting to new situations, even with limited training.", "questions": "What are the potential future directions for research and development in the field of planning agents like Skipper? Are there any specific improvements or enhancements that you plan to explore in your future work? For example, how does this method influence the research of more complex environments?\n\nCan you provide more information about the theoretical guarantees of the quality of the solution to the overall problem? What are the key insights or findings from the theoretical analysis? For example, what are  the limitations when applying Skipper in the continuous control environments.\n\nThe authors mentioned that non-uniform starting state distributions introduce additional difficulties in terms of exploration and therefore globally slow down the learning process. Is this caused by the intrinsic limitations of the Skipper method?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Skipper employs both temporal and spatial abstraction. Temporal abstraction breaks down complex tasks into smaller sub-tasks, while spatial abstraction allows the agent to focus on relevant aspects of the environment, improving learning efficiency and generalization.\n\nDecision-time planning enhances the agent's ability to handle novel situations. By formulating a simplified proxy problem and planning within it, the agent can generate plans in real time, adapting to changing environments or goals.\n\nSkipper uses proxy problems, represented as finite graphs, to create a simplified version of the task. The use of checkpoints as vertices in the proxy problems allows for a sparse representation, reducing computational complexity. \n\nThe experiments are conducted in gridworld environments with varying levels of difficulty, providing a realistic and challenging testbed for the agents. The use of MiniGrid-BabyAI framework ensures a well-defined and computationally feasible evaluation environment.", "weaknesses": "The experiments are conducted in simulated grid world environments, which might not fully represent the complexities of real-world scenarios. The findings might not directly translate to applications in more intricate, dynamic, and unstructured environments.\n\nIn Section 4: the paper assumes fully observable environments, which might not hold in many real-world applications where agents have limited or partial observability. The performance of Skipper in partially observable environments is not explored, limiting its practical applicability.\n\nThe paper introduces different variants of Skipper. While this provides a comprehensive analysis, it also complicates the understanding of the core method. A clearer presentation of the differences and justifications for each variant would enhance the paper's 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": {}}]

Author rebuttal:

Decision: Accept (poster)